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

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

@@ -68,6 +68,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;
@@ -2114,7 +2115,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(
@@ -2394,7 +2395,7 @@ public com.google.api.services.bigquery.model.QueryResponse call()
                   new QueryPageFetcher(jobId, schema, getOptions(), cursor, optionMap(options)),
                   cursor,
                   transformTableData(
-                      results.getRows(),
+                      results.getRows() != null ? results.getRows() : ImmutableList.of(),
                       schema,
                       getOptions().getDataFormatOptions().useInt64Timestamp())))
           .setJobId(jobId)
@@ -2418,7 +2419,7 @@ public com.google.api.services.bigquery.model.QueryResponse call()
                 new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
                 null,
                 transformTableData(
-                    results.getRows(),
+                    results.getRows() != null ? results.getRows() : ImmutableList.of(),
                     schema,
                     getOptions().getDataFormatOptions().useInt64Timestamp())))
         // Return the JobID of the successful job
@@ -2436,6 +2437,198 @@ public com.google.api.services.bigquery.model.QueryResponse call()
         .build();
   }
 
+  /**
+   * Executes a fast-path query RPC request expecting Arrow-formatted wire response.
+   *
+   * <p>Deserializes the returned Arrow IPC schema and record batch into standard {@link
+   * TableResult} row-based representations, configuring {@link ArrowQueryPageFetcher} for
+   * subsequent pages when pagination tokens are returned.
+   *
+   * @param projectId project ID in which to execute the query
+   * @param content query request content PB
+   * @param options job options
+   * @return either a {@link Job} if incomplete/fallback is required, or {@link TableResult}
+   * @throws InterruptedException if interrupted while awaiting RPC execution
+   */
+  private Object queryRpcArrow(
+      final String projectId, final QueryRequest content, JobOption... options)
+      throws InterruptedException {
+    com.google.api.services.bigquery.model.QueryResponse results;
+    Span queryRpc = null;
+    if (getOptions().isOpenTelemetryTracingEnabled()
+        && getOptions().getOpenTelemetryTracer() != null) {
+      queryRpc =
+          getOptions()
+              .getOpenTelemetryTracer()
+              .spanBuilder("com.google.cloud.bigquery.BigQuery.queryRpc")
+              .setAttribute("bq.query.project_id", projectId)
+              .setAllAttributes(otelAttributesFromQueryRequest(content))
+              .setAllAttributes(otelAttributesFromOptions(options))
+              .startSpan();
+    }
+    try (Scope queryRpcScope = queryRpc != null ? queryRpc.makeCurrent() : null) {
+      results =
+          BigQueryRetryHelper.runWithRetries(
+              () -> bigQueryRpc.queryRpcSkipExceptionTranslation(projectId, content),
+              getOptions().getRetrySettings(),
+              getOptions().getResultRetryAlgorithm(),
+              getOptions().getClock(),
+              DEFAULT_RETRY_CONFIG,
+              getOptions().isOpenTelemetryTracingEnabled(),
+              getOptions().getOpenTelemetryTracer());
+    } catch (BigQueryRetryHelper.BigQueryRetryHelperException e) {
+      throw BigQueryException.translateAndThrow(e);
+    } finally {
+      if (queryRpc != null) {
+        queryRpc.end();
+      }
+    }
+
+    if (results.getErrors() != null) {
+      List<BigQueryError> bigQueryErrors =
+          Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION);
+      throw new BigQueryException(bigQueryErrors);
+    }
+
+    // If query is incomplete or did not return an Arrow schema, fallback to fetching the Job.
+    if (!Boolean.TRUE.equals(results.getJobComplete()) || results.getArrowSchema() == null) {
+      if (results.getJobReference() == null) {
+        throw new BigQueryException(
+            0, "Job is incomplete or Arrow schema is missing, but no job reference was returned.");
+      }
+      JobId jobId = JobId.fromPb(results.getJobReference());
+      return getJob(jobId, options);
+    }
+
+    if (results.getArrowSchema().getSerializedSchema() == null) {
+      throw new BigQueryException(0, "Arrow schema is missing from the response");
+    }
+
+    // Deserialize Arrow IPC schema and convert to veneer BigQuery Schema.
+    byte[] arrowSchemaBytes = results.getArrowSchema().decodeSerializedSchema();
+    org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo;
+    try {
+      arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes);
+    } catch (IOException e) {
+      throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e);
+    }
+    Schema schema = ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchemaPojo);
+
+    long numRows;
+    if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) {
+      numRows = 0L;
+    } else if (results.getNumDmlAffectedRows() != null) {
+      numRows = results.getNumDmlAffectedRows();
+    } else {
+      numRows = results.getTotalRows().longValue();
+    }
+
+    StatementType statementType =
+        results.getStatementType() != null
+            ? StatementType.valueOf(results.getStatementType())
+            : null;
+    Long totalBytesBilled = results.getTotalBytesBilled();
+    Long totalBytesProcessed = results.getTotalBytesProcessed();
+    Long totalSlotMs = results.getTotalSlotMs();
+    Long numDmlAffectedRows = results.getNumDmlAffectedRows();
+    SessionInfo sessionInfo =
+        results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null;
+
+    // Deserialize first page of rows from the Arrow record batch (if present).
+    List<FieldValueList> firstPageRows;
+    if (results.getArrowRecordBatch() == null
+        || results.getArrowRecordBatch().getSerializedRecordBatch() == null) {
+      firstPageRows = ImmutableList.of();
+    } else {
+      try {
+        firstPageRows =
+            ArrowDeserializer.deserializeRecordBatch(
+                results.getArrowRecordBatch().decodeSerializedRecordBatch(),
+                schema,
+                arrowSchemaPojo);
+      } catch (IOException e) {
+        throw new BigQueryException(0, "Failed to deserialize Arrow record batch", e);
+      }
+    }
+
+    // Enforce maxResults limit on the first page if requested.
+    if (content.getMaxResults() != null && firstPageRows.size() > content.getMaxResults()) {
+      firstPageRows =
+          ImmutableList.copyOf(Iterables.limit(firstPageRows, content.getMaxResults().intValue()));
+    }
+
+    // Calculate row offset and determine if subsequent pages exist.
+    boolean hasMorePages = results.getPageToken() != null;
+    long initialRowOffset = 0L;
+    if (hasMorePages) {
+      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;
+      }
+    }
+
+    // Multi-page results: configure ArrowQueryPageFetcher for subsequent tabledata.list calls.
+    if (hasMorePages) {
+      if (results.getJobReference() == null) {
+        throw new BigQueryException(
+            0, "More pages exist, but no job reference was returned to fetch them.");
+      }
+      JobId jobId = JobId.fromPb(results.getJobReference());
+      String cursor = results.getPageToken();
+      NextPageFetcher<FieldValueList> pageFetcher =
+          new ArrowQueryPageFetcher(
+              jobId,
+              schema,
+              arrowSchemaBytes,
+              arrowSchemaPojo,
+              getOptions(),
+              initialRowOffset,
+              content.getMaxResults(),
+              optionMap(options));
+
+      return TableResult.newBuilder()
+          .setSchema(schema)
+          .setTotalRows(numRows)
+          .setPageNoSchema(new PageImpl<>(pageFetcher, cursor, firstPageRows))
+          .setJobId(jobId)
+          .setQueryId(results.getQueryId())
+          .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
+          .setRowsInPage((long) firstPageRows.size())
+          .setStatementType(statementType)
+          .setTotalBytesBilled(totalBytesBilled)
+          .setTotalBytesProcessed(totalBytesProcessed)
+          .setTotalSlotMs(totalSlotMs)
+          .setNumDmlAffectedRows(numDmlAffectedRows)
+          .setSessionInfo(sessionInfo)
+          .build();
+    }
+
+    // only 1 page of result
+    return TableResult.newBuilder()
+        .setSchema(schema)
+        .setTotalRows(numRows)
+        .setPageNoSchema(
+            new PageImpl<>(
+                new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
+                null,
+                firstPageRows))
+        .setJobId(
+            results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null)
+        .setQueryId(results.getQueryId())
+        .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
+        .setRowsInPage((long) firstPageRows.size())
+        .setStatementType(statementType)
+        .setTotalBytesBilled(totalBytesBilled)
+        .setTotalBytesProcessed(totalBytesProcessed)
+        .setTotalSlotMs(totalSlotMs)
+        .setNumDmlAffectedRows(numDmlAffectedRows)
+        .setSessionInfo(sessionInfo)
+        .build();
+  }
+
   @Override
   public TableResult query(QueryJobConfiguration configuration, JobId jobId, JobOption... options)
       throws InterruptedException, JobException {
@@ -2452,11 +2645,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 =
@@ -2509,8 +2697,17 @@ && getOptions().getOpenTelemetryTracer() != null) {
           content.setTimeoutMs(timeoutMs);
         }
 
+        if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) {
+          return queryRpcArrow(projectId, content, options);
+        }
         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) {