feat(bigquery): add Storage Read API slow-path fallback for row-based query()

Author: jinseopkim0

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

@@ -303,6 +303,7 @@ static class ArrowQueryPageFetcher implements NextPageFetcher<FieldValueList> {
     private static final long DEFAULT_PAGE_SIZE = 10000L;
 
     private final JobId jobId;
+    private final String customStreamName;
     private final Schema schema;
     private final byte[] arrowSchemaBytes;
     private final BigQueryOptions serviceOptions;
@@ -326,7 +327,30 @@ static class ArrowQueryPageFetcher implements NextPageFetcher<FieldValueList> {
         long initialRowOffset,
         Long maxResults,
         Map<BigQueryRpc.Option, ?> optionsMap) {
+      this(
+          jobId,
+          null,
+          schema,
+          arrowSchemaBytes,
+          arrowSchemaPojo,
+          serviceOptions,
+          initialRowOffset,
+          maxResults,
+          optionsMap);
+    }
+
+    ArrowQueryPageFetcher(
+        JobId jobId,
+        String customStreamName,
+        Schema schema,
+        byte[] arrowSchemaBytes,
+        org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo,
+        BigQueryOptions serviceOptions,
+        long initialRowOffset,
+        Long maxResults,
+        Map<BigQueryRpc.Option, ?> optionsMap) {
       this.jobId = jobId;
+      this.customStreamName = customStreamName;
       this.schema = schema;
       this.arrowSchemaBytes = arrowSchemaBytes;
       this.arrowSchemaPojo = arrowSchemaPojo;
@@ -355,13 +379,6 @@ public Page<FieldValueList> getNextPage() {
       List<FieldValueList> rowBatch = new ArrayList<>((int) Math.min(pageSize, 10000L));
 
       try {
-        String location =
-            jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation();
-        if (location == null) {
-          throw new BigQueryException(
-              0, "Location must be specified to read Arrow rows from storage stream");
-        }
-
         if (bqReadClient == null) {
           BigQuery service = serviceOptions.getService();
           if (service instanceof BigQueryImpl) {
@@ -374,12 +391,23 @@ public Page<FieldValueList> getNextPage() {
         }
 
         if (streamIterator == null) {
-          String streamName =
-              String.format(
-                  "projects/%s/locations/%s/jobs/%s/streams/_default",
-                  jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(),
-                  location,
-                  jobId.getJob());
+          String streamName;
+          if (customStreamName != null) {
+            streamName = customStreamName;
+          } else {
+            String location =
+                jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation();
+            if (location == null) {
+              throw new BigQueryException(
+                  0, "Location must be specified to read Arrow rows from storage stream");
+            }
+            streamName =
+                String.format(
+                    "projects/%s/locations/%s/jobs/%s/streams/_default",
+                    jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(),
+                    location,
+                    jobId.getJob());
+          }
 
           ReadRowsRequest readRowsRequest =
               ReadRowsRequest.newBuilder()
@@ -2710,8 +2738,7 @@ && getOptions().getOpenTelemetryTracer() != null) {
       }
 
       if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) {
-        throw new UnsupportedOperationException(
-            "Arrow results format for slow query path execution is not yet supported.");
+        return queryFallbackArrow(jobId, configuration, options);
       }
 
       return create(JobInfo.of(jobId, configuration), options);
@@ -2991,6 +3018,172 @@ private ArrowQueryResult createArrowQueryResultFromTable(
     return ArrowQueryResultImpl.fromReadSession(readSession, jobId, client);
   }
 
+  /**
+   * Executes a slow-path query job using {@code jobs.insert}, awaits its completion, and streams
+   * the result rows via the BigQuery Storage Read API in Arrow format, wrapping the decoded rows in
+   * a {@link TableResult}.
+   *
+   * @param jobId the job ID, or {@code null}
+   * @param configuration the query job configuration
+   * @param options query job options
+   * @return a {@link TableResult} containing the decoded rows and job execution metadata
+   * @throws InterruptedException if interrupted while awaiting job completion
+   * @throws BigQueryException if job execution or ReadSession creation fails
+   */
+  private TableResult queryFallbackArrow(
+      JobId jobId, QueryJobConfiguration configuration, JobOption... options)
+      throws InterruptedException {
+    Job job = create(JobInfo.of(jobId, configuration), options);
+    Job completedJob = job.waitFor();
+
+    if (completedJob == null) {
+      throw new BigQueryException(0, "Job no longer exists or could not be retrieved.");
+    }
+
+    if (completedJob.getStatus().getError() != null) {
+      throw new BigQueryException(Collections.singletonList(completedJob.getStatus().getError()));
+    }
+
+    TableId destinationTable = null;
+    if (completedJob.getConfiguration() instanceof QueryJobConfiguration) {
+      destinationTable =
+          ((QueryJobConfiguration) completedJob.getConfiguration()).getDestinationTable();
+    }
+    if (destinationTable == null) {
+      destinationTable = configuration.getDestinationTable();
+    }
+    if (destinationTable == null) {
+      throw new BigQueryException(0, "Unable to resolve destination table for fallback query");
+    }
+
+    JobStatistics.QueryStatistics stats =
+        completedJob.getStatistics() instanceof JobStatistics.QueryStatistics
+            ? (JobStatistics.QueryStatistics) completedJob.getStatistics()
+            : null;
+
+    StatementType statementType = stats != null ? stats.getStatementType() : null;
+    Long totalBytesBilled = stats != null ? stats.getTotalBytesBilled() : null;
+    Long totalBytesProcessed = stats != null ? stats.getTotalBytesProcessed() : null;
+    Long totalSlotMs = stats != null ? stats.getTotalSlotMs() : null;
+    Long numDmlAffectedRows = stats != null ? stats.getNumDmlAffectedRows() : null;
+    SessionInfo sessionInfo = stats != null ? stats.getSessionInfo() : null;
+
+    String destProject =
+        destinationTable.getProject() != null
+            ? destinationTable.getProject()
+            : (completedJob.getJobId() != null && completedJob.getJobId().getProject() != null
+                ? completedJob.getJobId().getProject()
+                : getOptions().getProjectId());
+    String parent = String.format("projects/%s", destProject);
+    String srcTable =
+        String.format(
+            "projects/%s/datasets/%s/tables/%s",
+            destProject, destinationTable.getDataset(), destinationTable.getTable());
+
+    BigQueryReadClient client = getBigQueryReadClient();
+    CreateReadSessionRequest request =
+        CreateReadSessionRequest.newBuilder()
+            .setParent(parent)
+            .setReadSession(
+                ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW))
+            .setMaxStreamCount(1)
+            .build();
+    ReadSession readSession;
+    try {
+      readSession = client.createReadSession(request);
+    } catch (Exception e) {
+      throw new BigQueryException(0, "Failed to create ReadSession for fallback query", e);
+    }
+
+    org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo = null;
+    byte[] arrowSchemaBytes = null;
+    if (readSession.hasArrowSchema()) {
+      arrowSchemaBytes = readSession.getArrowSchema().getSerializedSchema().toByteArray();
+      try {
+        arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes);
+      } catch (IOException e) {
+        throw new BigQueryException(0, "Failed to deserialize Arrow schema from ReadSession", e);
+      }
+    }
+    Schema schema =
+        arrowSchemaPojo != null
+            ? ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchemaPojo)
+            : (stats != null ? stats.getSchema() : null);
+
+    String streamName =
+        readSession.getStreamsCount() > 0 ? readSession.getStreams(0).getName() : null;
+
+    if (streamName == null) {
+      return TableResult.newBuilder()
+          .setSchema(schema)
+          .setTotalRows(numDmlAffectedRows != null ? numDmlAffectedRows : 0L)
+          .setPageNoSchema(
+              new PageImpl<>(
+                  new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
+                  null,
+                  ImmutableList.of()))
+          .setJobId(completedJob.getJobId())
+          .setRowsInPage(0L)
+          .setStatementType(statementType)
+          .setTotalBytesBilled(totalBytesBilled)
+          .setTotalBytesProcessed(totalBytesProcessed)
+          .setTotalSlotMs(totalSlotMs)
+          .setNumDmlAffectedRows(numDmlAffectedRows)
+          .setSessionInfo(sessionInfo)
+          .build();
+    }
+
+    ArrowQueryPageFetcher pageFetcher =
+        new ArrowQueryPageFetcher(
+            completedJob.getJobId(),
+            streamName,
+            schema,
+            arrowSchemaBytes,
+            arrowSchemaPojo,
+            getOptions(),
+            0L,
+            configuration.getMaxResults(),
+            optionMap(options));
+
+    Page<FieldValueList> firstPage = pageFetcher.getNextPage();
+    List<FieldValueList> firstPageRows =
+        firstPage != null ? ImmutableList.copyOf(firstPage.getValues()) : ImmutableList.of();
+    long rowsInPage = (long) firstPageRows.size();
+
+    Table destTable = null;
+    try {
+      destTable = getTable(destinationTable);
+    } catch (Exception e) {
+      // Non-fatal table lookup failure
+    }
+    long totalRows =
+        numDmlAffectedRows != null
+            ? numDmlAffectedRows
+            : (destTable != null && destTable.getNumRows() != null
+                ? destTable.getNumRows().longValue()
+                : rowsInPage);
+
+    return TableResult.newBuilder()
+        .setSchema(schema)
+        .setTotalRows(totalRows)
+        .setPageNoSchema(
+            firstPage != null
+                ? firstPage
+                : new PageImpl<>(
+                    new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
+                    null,
+                    ImmutableList.of()))
+        .setJobId(completedJob.getJobId())
+        .setRowsInPage(rowsInPage)
+        .setStatementType(statementType)
+        .setTotalBytesBilled(totalBytesBilled)
+        .setTotalBytesProcessed(totalBytesProcessed)
+        .setTotalSlotMs(totalSlotMs)
+        .setNumDmlAffectedRows(numDmlAffectedRows)
+        .setSessionInfo(sessionInfo)
+        .build();
+  }
+
   @Override
   public QueryResponse getQueryResults(JobId jobId, QueryResultsOption... options) {
     Map<BigQueryRpc.Option, ?> optionsMap = optionMap(options);