feat(bigquery): accelerate row-based query() with Arrow wire format

Author: jinseopkim0

java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java

@@ -67,6 +67,7 @@
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.net.HostAndPort;
+import com.google.common.primitives.Longs;
 import io.grpc.ManagedChannelBuilder;
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.api.trace.Span;
@@ -75,6 +76,7 @@
 import java.net.URI;
 import java.util.ArrayDeque;
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
@@ -2110,7 +2112,7 @@ public TableDataList call() throws IOException {
     }
   }
 
-  private static Iterable<FieldValueList> transformTableData(
+  private static List<FieldValueList> transformTableData(
       Iterable<TableRow> tableDataPb, final Schema schema, boolean useInt64Timestamps) {
     return ImmutableList.copyOf(
         Iterables.transform(
@@ -2349,8 +2351,21 @@ public com.google.api.services.bigquery.model.QueryResponse call()
 
     long numRows;
     Schema schema;
-    if (results.getJobComplete() && results.getSchema() != null) {
-      schema = Schema.fromPb(results.getSchema());
+    boolean isArrow = results.getArrowSchema() != null;
+    org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo = null;
+    byte[] arrowSchemaBytes = null;
+    if (results.getJobComplete() && (results.getSchema() != null || isArrow)) {
+      if (isArrow) {
+        arrowSchemaBytes = results.getArrowSchema().decodeSerializedSchema();
+        try {
+          arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes);
+        } catch (IOException e) {
+          throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e);
+        }
+        schema = ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchemaPojo);
+      } else {
+        schema = Schema.fromPb(results.getSchema());
+      }
       if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) {
         numRows = 0L;
       } else if (results.getNumDmlAffectedRows() != null) {
@@ -2378,25 +2393,75 @@ public com.google.api.services.bigquery.model.QueryResponse call()
     SessionInfo sessionInfo =
         results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null;
 
-    if (results.getPageToken() != null) {
+    Collection<FieldValueList> firstPageRows;
+    if (isArrow) {
+      if (results.getArrowRecordBatch() != null
+          && results.getArrowRecordBatch().getSerializedRecordBatch() != null) {
+        try {
+          firstPageRows =
+              ArrowDeserializer.deserializeRecordBatch(
+                  results.getArrowRecordBatch().decodeSerializedRecordBatch(),
+                  schema,
+                  arrowSchemaPojo);
+        } catch (IOException e) {
+          throw new BigQueryException(0, "Failed to deserialize Arrow record batch", e);
+        }
+      } else {
+        firstPageRows = ImmutableList.of();
+      }
+    } else {
+      firstPageRows =
+          results.getRows() != null
+              ? transformTableData(
+                  results.getRows(),
+                  schema,
+                  getOptions().getDataFormatOptions().useInt64Timestamp())
+              : ImmutableList.of();
+    }
+
+    boolean hasMorePages = results.getPageToken() != null;
+    long initialRowOffset = 0L;
+    if (hasMorePages && isArrow) {
+      Long parsedOffset = Longs.tryParse(results.getPageToken());
+      initialRowOffset = parsedOffset != null ? parsedOffset : firstPageRows.size();
+      if (content.getMaxResults() != null
+          && (initialRowOffset >= content.getMaxResults()
+              || firstPageRows.size() >= content.getMaxResults())) {
+        hasMorePages = false;
+      }
+    }
+
+    if (hasMorePages) {
       JobId jobId = JobId.fromPb(results.getJobReference());
       String cursor = results.getPageToken();
+
+      NextPageFetcher<FieldValueList> pageFetcher;
+      if (isArrow) {
+        pageFetcher =
+            new ArrowQueryPageFetcher(
+                jobId,
+                schema,
+                arrowSchemaBytes,
+                arrowSchemaPojo,
+                getOptions(),
+                initialRowOffset,
+                content.getMaxResults(),
+                optionMap(options));
+      } else {
+        pageFetcher = new QueryPageFetcher(jobId, schema, getOptions(), cursor, optionMap(options));
+      }
+
       return TableResult.newBuilder()
           .setSchema(schema)
           .setTotalRows(numRows)
           .setPageNoSchema(
               new PageImpl<>(
                   // fetch next pages of results
-                  new QueryPageFetcher(jobId, schema, getOptions(), cursor, optionMap(options)),
-                  cursor,
-                  transformTableData(
-                      results.getRows(),
-                      schema,
-                      getOptions().getDataFormatOptions().useInt64Timestamp())))
+                  pageFetcher, cursor, firstPageRows))
           .setJobId(jobId)
           .setQueryId(results.getQueryId())
           .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
-          .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
+          .setRowsInPage((long) firstPageRows.size())
           .setStatementType(statementType)
           .setTotalBytesBilled(totalBytesBilled)
           .setTotalBytesProcessed(totalBytesProcessed)
@@ -2413,16 +2478,13 @@ public com.google.api.services.bigquery.model.QueryResponse call()
             new PageImpl<>(
                 new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
                 null,
-                transformTableData(
-                    results.getRows(),
-                    schema,
-                    getOptions().getDataFormatOptions().useInt64Timestamp())))
+                firstPageRows))
         // Return the JobID of the successful job
         .setJobId(
             results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null)
         .setQueryId(results.getQueryId())
         .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
-        .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
+        .setRowsInPage((long) firstPageRows.size())
         .setStatementType(statementType)
         .setTotalBytesBilled(totalBytesBilled)
         .setTotalBytesProcessed(totalBytesProcessed)
@@ -2448,11 +2510,6 @@ public Object queryWithTimeout(
       throws InterruptedException, JobException {
     Job.checkNotDryRun(configuration, "query");
 
-    if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) {
-      throw new IllegalArgumentException(
-          "QueryResultsFormat.ARROW is not supported with query(). Use queryArrow() instead.");
-    }
-
     // If JobCreationMode is not explicitly set, update it with default value;
     if (configuration.getJobCreationMode() == null) {
       configuration =
@@ -2507,6 +2564,12 @@ && getOptions().getOpenTelemetryTracer() != null) {
 
         return queryRpc(projectId, content, options);
       }
+
+      if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) {
+        throw new IllegalArgumentException(
+            "Arrow results format is only supported for fast query path execution (e.g. no destination table, no custom clustering, etc.).");
+      }
+
       return create(JobInfo.of(jobId, configuration), options);
     } finally {
       if (querySpan != null) {