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

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()
@@ -2713,8 +2741,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);
@@ -2994,6 +3021,157 @@ 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<>(null, 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();
+
+    long totalRows =
+        numDmlAffectedRows != null
+            ? numDmlAffectedRows
+            : (readSession.getEstimatedRowCount() > 0
+                ? readSession.getEstimatedRowCount()
+                : rowsInPage);
+
+    return TableResult.newBuilder()
+        .setSchema(schema)
+        .setTotalRows(totalRows)
+        .setPageNoSchema(
+            firstPage != null ? firstPage : new PageImpl<>(null, 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);

java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java

@@ -74,8 +74,11 @@
 import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
 import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc;
 import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
+import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest;
 import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
 import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
+import com.google.cloud.bigquery.storage.v1.ReadSession;
+import com.google.cloud.bigquery.storage.v1.ReadStream;
 import com.google.common.base.Function;
 import com.google.common.base.Supplier;
 import com.google.common.collect.ImmutableList;
@@ -2949,19 +2952,126 @@ void testQueryArrowDefaultsToJobCreationOptional() throws IOException, Interrupt
   }
 
   @Test
-  void testQueryArrowResultsFormatUnsupportedConfiguration() {
+  void testQueryWithArrowFormatSlowPathFallback() throws Exception {
+    JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
+    com.google.api.services.bigquery.model.JobStatus jobStatus =
+        new com.google.api.services.bigquery.model.JobStatus().setState("DONE");
+
+    com.google.api.services.bigquery.model.Job jobResponsePb =
+        new com.google.api.services.bigquery.model.Job()
+            .setJobReference(queryJob.toPb())
+            .setStatus(jobStatus)
+            .setConfiguration(
+                new com.google.api.services.bigquery.model.JobConfiguration()
+                    .setQuery(
+                        new JobConfigurationQuery()
+                            .setQuery("SELECT id FROM test")
+                            .setDestinationTable(TABLE_ID.toPb())))
+            .setStatistics(
+                new com.google.api.services.bigquery.model.JobStatistics()
+                    .setTotalSlotMs(50L)
+                    .setQuery(
+                        new com.google.api.services.bigquery.model.JobStatistics2()
+                            .setStatementType("SELECT")
+                            .setTotalBytesBilled(100L)
+                            .setTotalBytesProcessed(200L)));
+
+    when(bigqueryRpcMock.createSkipExceptionTranslation(
+            any(com.google.api.services.bigquery.model.Job.class), any()))
+        .thenReturn(jobResponsePb);
+    when(bigqueryRpcMock.getJobSkipExceptionTranslation(eq(PROJECT), eq(JOB), any(), any()))
+        .thenReturn(jobResponsePb);
+    when(bigqueryRpcMock.getQueryResultsSkipExceptionTranslation(
+            eq(PROJECT), eq(JOB), any(), any()))
+        .thenReturn(
+            new GetQueryResultsResponse().setJobComplete(true).setJobReference(queryJob.toPb()));
+
+    org.apache.arrow.vector.types.pojo.Schema arrowSchema =
+        new org.apache.arrow.vector.types.pojo.Schema(
+            ImmutableList.of(
+                org.apache.arrow.vector.types.pojo.Field.nullable(
+                    "id", new ArrowType.Int(64, true))));
+
+    byte[] schemaBytes;
+    try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+      MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema);
+      schemaBytes = out.toByteArray();
+    }
+
+    byte[] batchBytes;
+    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+      BigIntVector idVector = new BigIntVector("id", allocator);
+      idVector.allocateNew(1);
+      idVector.set(0, 42L);
+      idVector.setValueCount(1);
+      try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(idVector))) {
+        VectorUnloader unloader = new VectorUnloader(root);
+        try (ArrowRecordBatch recordBatch = unloader.getRecordBatch();
+            ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+          WriteChannel channel = new WriteChannel(Channels.newChannel(out));
+          MessageSerializer.serialize(channel, recordBatch);
+          batchBytes = out.toByteArray();
+        }
+      } finally {
+        idVector.close();
+      }
+    }
+
+    com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch =
+        com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder()
+            .setSerializedRecordBatch(ByteString.copyFrom(batchBytes))
+            .build();
+    ReadRowsResponse streamResponse =
+        ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build();
+
+    @SuppressWarnings("unchecked")
+    ServerStreamingCallable<ReadRowsRequest, ReadRowsResponse> mockCallable =
+        mock(ServerStreamingCallable.class, withSettings().withoutAnnotations());
+    @SuppressWarnings("unchecked")
+    ServerStream<ReadRowsResponse> mockServerStream =
+        mock(ServerStream.class, withSettings().withoutAnnotations());
+    when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream);
+    when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator());
+
+    ReadSession readSession =
+        ReadSession.newBuilder()
+            .setName("projects/" + PROJECT + "/locations/" + LOCATION + "/sessions/session-1")
+            .setArrowSchema(
+                com.google.cloud.bigquery.storage.v1.ArrowSchema.newBuilder()
+                    .setSerializedSchema(ByteString.copyFrom(schemaBytes)))
+            .addStreams(ReadStream.newBuilder().setName("stream-1"))
+            .build();
+
+    BigQueryReadClient mockReadClient =
+        mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
+    when(mockReadClient.createReadSession(any(CreateReadSessionRequest.class)))
+        .thenReturn(readSession);
+    when(mockReadClient.readRowsCallable()).thenReturn(mockCallable);
+
+    bigquery = options.getService();
+    ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient);
+
     QueryJobConfiguration config =
-        QueryJobConfiguration.newBuilder("SELECT 1")
+        QueryJobConfiguration.newBuilder("SELECT id FROM test")
             .setQueryResultsFormat(QueryResultsFormat.ARROW)
-            .setDestinationTable(TableId.of("dataset", "table"))
+            .setDestinationTable(TABLE_ID)
             .build();
-    bigquery = options.getService();
-    UnsupportedOperationException exception =
-        assertThrows(UnsupportedOperationException.class, () -> bigquery.query(config));
-    assertTrue(
-        exception
-            .getMessage()
-            .contains("Arrow results format for slow query path execution is not yet supported."));
+
+    TableResult result = bigquery.query(config);
+    assertNotNull(result);
+    assertEquals(1, Iterables.size(result.getValues()));
+    assertEquals(queryJob, result.getJobId());
+    assertEquals("42", result.getValues().iterator().next().get(0).getStringValue());
+    assertEquals(StatementType.SELECT, result.getStatementType());
+    assertEquals(100L, result.getTotalBytesBilled());
+    assertEquals(200L, result.getTotalBytesProcessed());
+    assertEquals(50L, result.getTotalSlotMs());
+
+    verify(bigqueryRpcMock)
+        .createSkipExceptionTranslation(
+            any(com.google.api.services.bigquery.model.Job.class), any());
+    verify(mockReadClient).createReadSession(any(CreateReadSessionRequest.class));
+    verify(mockCallable).call(any(ReadRowsRequest.class));
   }
 
   @Test