From d031b760273af80fa910f4814dd3c4c02ab73dc1 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 21:24:50 -0400 Subject: [PATCH 1/4] feat(bigquery): add Storage Read API slow-path fallback for row-based query() --- .../google/cloud/bigquery/BigQueryImpl.java | 241 ++++++++++++++++-- .../cloud/bigquery/BigQueryImplTest.java | 130 +++++++++- 2 files changed, 337 insertions(+), 34 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index a6098f27f5ad..7cd09add57b2 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -308,6 +308,7 @@ static class ArrowQueryPageFetcher implements NextPageFetcher { 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; @@ -331,7 +332,30 @@ static class ArrowQueryPageFetcher implements NextPageFetcher { long initialRowOffset, Long maxResults, Map 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 optionsMap) { this.jobId = jobId; + this.customStreamName = customStreamName; this.schema = schema; this.arrowSchemaBytes = arrowSchemaBytes; this.arrowSchemaPojo = arrowSchemaPojo; @@ -360,21 +384,6 @@ public Page getNextPage() { List rowBatch = new ArrayList<>((int) Math.min(pageSize, 10000L)); try { - // Resolve job location in order: JobId location -> BigQueryOptions location -> "global" - // default. - // The Storage Read API stream resource name requires a location component (e.g. - // projects/{project}/locations/{location}/jobs/{job}/streams/_default). If no specific - // location - // was provided on the job or service options, defaulting to "global" allows queries created - // without an explicit location to still stream results without failing. - String location = jobId.getLocation(); - if (location == null) { - location = serviceOptions.getLocation(); - } - if (location == null) { - location = "global"; - } - if (streamIterator == null) { if (bqReadClient == null) { BigQuery service = serviceOptions.getService(); @@ -390,13 +399,32 @@ public Page getNextPage() { } } - // Construct the default stream path for reading job query results via Storage Read API. - 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 { + // Resolve job location in order: JobId location -> BigQueryOptions location -> "global" + // default. + // The Storage Read API stream resource name requires a location component (e.g. + // projects/{project}/locations/{location}/jobs/{job}/streams/_default). If no specific + // location was provided on the job or service options, defaulting to "global" allows + // queries created without an explicit location to still stream results without failing. + String location = jobId.getLocation(); + if (location == null) { + location = serviceOptions.getLocation(); + } + if (location == null) { + location = "global"; + } + + // Construct the default stream path for reading job query results via Storage Read API. + 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() @@ -2741,8 +2769,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); @@ -3022,6 +3049,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 firstPage = pageFetcher.getNextPage(); + List 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 optionsMap = optionMap(options); diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index b4b4081bfe89..9cefda8587b5 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/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 mockCallable = + mock(ServerStreamingCallable.class, withSettings().withoutAnnotations()); + @SuppressWarnings("unchecked") + ServerStream 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 From 80059a6dadd6bc983f6024bfa9f94782a5690285 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 22:19:34 -0400 Subject: [PATCH 2/4] refactor(bigquery): simplify empty PageImpl and use readSession estimatedRowCount in slow-path fallback --- .../google/cloud/bigquery/BigQueryImpl.java | 23 ++++--------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 7cd09add57b2..7c6720113aad 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -3148,11 +3148,7 @@ private TableResult queryFallbackArrow( return TableResult.newBuilder() .setSchema(schema) .setTotalRows(numDmlAffectedRows != null ? numDmlAffectedRows : 0L) - .setPageNoSchema( - new PageImpl<>( - new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)), - null, - ImmutableList.of())) + .setPageNoSchema(new PageImpl<>(null, null, ImmutableList.of())) .setJobId(completedJob.getJobId()) .setRowsInPage(0L) .setStatementType(statementType) @@ -3181,29 +3177,18 @@ private TableResult queryFallbackArrow( 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() + : (readSession.getEstimatedRowCount() > 0 + ? readSession.getEstimatedRowCount() : 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())) + firstPage != null ? firstPage : new PageImpl<>(null, null, ImmutableList.of())) .setJobId(completedJob.getJobId()) .setRowsInPage(rowsInPage) .setStatementType(statementType) From f8f1a8ce4832ff8fdb85c1404f2945732fee0a44 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Thu, 17 Sep 2026 21:08:50 -0400 Subject: [PATCH 3/4] docs(bigquery): document execution stages in queryFallbackArrow --- .../java/com/google/cloud/bigquery/BigQueryImpl.java | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 7c6720113aad..c96cc776f284 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -3064,6 +3064,7 @@ private ArrowQueryResult createArrowQueryResultFromTable( private TableResult queryFallbackArrow( JobId jobId, QueryJobConfiguration configuration, JobOption... options) throws InterruptedException { + // Submit the query job via jobs.insert and poll until completion. Job job = create(JobInfo.of(jobId, configuration), options); Job completedJob = job.waitFor(); @@ -3075,6 +3076,7 @@ private TableResult queryFallbackArrow( throw new BigQueryException(Collections.singletonList(completedJob.getStatus().getError())); } + // Resolve the query's destination table where the completed job wrote its results. TableId destinationTable = null; if (completedJob.getConfiguration() instanceof QueryJobConfiguration) { destinationTable = @@ -3087,6 +3089,7 @@ private TableResult queryFallbackArrow( throw new BigQueryException(0, "Unable to resolve destination table for fallback query"); } + // Extract query execution statistics from the completed job metadata. JobStatistics.QueryStatistics stats = completedJob.getStatistics() instanceof JobStatistics.QueryStatistics ? (JobStatistics.QueryStatistics) completedJob.getStatistics() @@ -3099,6 +3102,7 @@ private TableResult queryFallbackArrow( Long numDmlAffectedRows = stats != null ? stats.getNumDmlAffectedRows() : null; SessionInfo sessionInfo = stats != null ? stats.getSessionInfo() : null; + // Create a Storage Read API ReadSession targeting the destination table in Arrow format. String destProject = destinationTable.getProject() != null ? destinationTable.getProject() @@ -3126,6 +3130,7 @@ private TableResult queryFallbackArrow( throw new BigQueryException(0, "Failed to create ReadSession for fallback query", e); } + // Deserialize the Arrow schema and convert to BigQuery Schema for TableResult metadata. org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo = null; byte[] arrowSchemaBytes = null; if (readSession.hasArrowSchema()) { @@ -3144,6 +3149,8 @@ private TableResult queryFallbackArrow( String streamName = readSession.getStreamsCount() > 0 ? readSession.getStreams(0).getName() : null; + // If the destination table has no data streams (e.g. DDL/DML statements or empty results), + // return an empty TableResult populated with execution statistics. if (streamName == null) { return TableResult.newBuilder() .setSchema(schema) @@ -3160,6 +3167,7 @@ private TableResult queryFallbackArrow( .build(); } + // Initialize the page fetcher targeting the ReadSession stream to load the first page of rows. ArrowQueryPageFetcher pageFetcher = new ArrowQueryPageFetcher( completedJob.getJobId(), @@ -3177,6 +3185,8 @@ private TableResult queryFallbackArrow( firstPage != null ? ImmutableList.copyOf(firstPage.getValues()) : ImmutableList.of(); long rowsInPage = (long) firstPageRows.size(); + // Determine total row count: prefer DML affected rows, fallback to estimated row count, + // or actual first page rows returned. long totalRows = numDmlAffectedRows != null ? numDmlAffectedRows @@ -3184,6 +3194,7 @@ private TableResult queryFallbackArrow( ? readSession.getEstimatedRowCount() : rowsInPage); + // Assemble and return the complete TableResult. return TableResult.newBuilder() .setSchema(schema) .setTotalRows(totalRows) From fd83fae1427499b849f2221c245b5e714e228570 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 18 Sep 2026 10:44:52 -0400 Subject: [PATCH 4/4] feat(bigquery): support mid-flight slow-path transition for incomplete Arrow queries --- .../google/cloud/bigquery/BigQueryImpl.java | 75 ++++++--- .../cloud/bigquery/BigQueryImplTest.java | 152 +++++++++++++++++- 2 files changed, 201 insertions(+), 26 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index c96cc776f284..013a16de3748 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -2554,10 +2554,23 @@ && getOptions().getOpenTelemetryTracer() != null) { throw new BigQueryException(bigQueryErrors); } - // If query is incomplete, Arrow format for slow query path is not yet supported. + // If the query is incomplete (took longer than the fast-path timeout), transition + // to the slow path: retrieve the created job, wait for completion, and read the + // results using the Storage Read API in Arrow format. if (!Boolean.TRUE.equals(results.getJobComplete())) { - throw new UnsupportedOperationException( - "Arrow results format for slow query path execution is not yet supported."); + if (results.getJobReference() == null) { + throw new BigQueryException( + 0, "Query is incomplete, but no job reference was returned to await completion."); + } + JobId jobId = JobId.fromPb(results.getJobReference()); + Job job = getJob(jobId, options); + if (job == null) { + throw new BigQueryException(0, "Job no longer exists or could not be retrieved: " + jobId); + } + Job completedJob = job.waitFor(); + Long maxResults = + content.getMaxResults() != null ? content.getMaxResults().longValue() : null; + return readArrowTableResultFromJob(completedJob, maxResults, null, options); } // If query is complete but Arrow schema is missing or invalid, fallback to standard JSON @@ -3050,24 +3063,18 @@ private ArrowQueryResult createArrowQueryResultFromTable( } /** - * 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}. + * Reads query result rows from a completed query job's destination table using the Storage Read + * API in Arrow format, returning a populated {@link TableResult}. * - * @param jobId the job ID, or {@code null} - * @param configuration the query job configuration + * @param completedJob the completed query job + * @param maxResults maximum results requested, or {@code null} + * @param fallbackDestinationTable fallback destination table if not present on the job * @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 + * @return a {@link TableResult} containing decoded rows and execution metadata + * @throws BigQueryException if job failed or ReadSession creation fails */ - private TableResult queryFallbackArrow( - JobId jobId, QueryJobConfiguration configuration, JobOption... options) - throws InterruptedException { - // Submit the query job via jobs.insert and poll until completion. - Job job = create(JobInfo.of(jobId, configuration), options); - Job completedJob = job.waitFor(); - + private TableResult readArrowTableResultFromJob( + Job completedJob, Long maxResults, TableId fallbackDestinationTable, JobOption... options) { if (completedJob == null) { throw new BigQueryException(0, "Job no longer exists or could not be retrieved."); } @@ -3076,17 +3083,17 @@ private TableResult queryFallbackArrow( throw new BigQueryException(Collections.singletonList(completedJob.getStatus().getError())); } - // Resolve the query's destination table where the completed job wrote its results. + // Resolve the destination table containing the query results. TableId destinationTable = null; if (completedJob.getConfiguration() instanceof QueryJobConfiguration) { destinationTable = ((QueryJobConfiguration) completedJob.getConfiguration()).getDestinationTable(); } if (destinationTable == null) { - destinationTable = configuration.getDestinationTable(); + destinationTable = fallbackDestinationTable; } if (destinationTable == null) { - throw new BigQueryException(0, "Unable to resolve destination table for fallback query"); + throw new BigQueryException(0, "Unable to resolve destination table for query job"); } // Extract query execution statistics from the completed job metadata. @@ -3127,7 +3134,7 @@ private TableResult queryFallbackArrow( try { readSession = client.createReadSession(request); } catch (Exception e) { - throw new BigQueryException(0, "Failed to create ReadSession for fallback query", e); + throw new BigQueryException(0, "Failed to create ReadSession for query job", e); } // Deserialize the Arrow schema and convert to BigQuery Schema for TableResult metadata. @@ -3177,7 +3184,7 @@ private TableResult queryFallbackArrow( arrowSchemaPojo, getOptions(), 0L, - configuration.getMaxResults(), + maxResults, optionMap(options)); Page firstPage = pageFetcher.getNextPage(); @@ -3211,6 +3218,28 @@ private TableResult queryFallbackArrow( .build(); } + /** + * 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 { + // Submit the query job via jobs.insert and poll until completion. + Job job = create(JobInfo.of(jobId, configuration), options); + Job completedJob = job.waitFor(); + return readArrowTableResultFromJob( + completedJob, configuration.getMaxResults(), configuration.getDestinationTable(), options); + } + @Override public QueryResponse getQueryResults(JobId jobId, QueryResultsOption... options) { Map optionsMap = optionMap(options); diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 9cefda8587b5..530b5bddaaf1 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -3406,20 +3406,166 @@ void testQueryWithArrowFormatIncompleteJob() throws Exception { .setJobComplete(false) .setJobReference(queryJob.toPb()); + 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.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) .thenReturn(queryResponsePb); + 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 mockCallable = + mock(ServerStreamingCallable.class, withSettings().withoutAnnotations()); + @SuppressWarnings("unchecked") + ServerStream 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 id FROM test") .setQueryResultsFormat(QueryResultsFormat.ARROW) .build(); - UnsupportedOperationException e = - assertThrows(UnsupportedOperationException.class, () -> bigquery.query(config)); + + TableResult result = bigquery.query(config); + + assertNotNull(result); + assertEquals(1, Iterables.size(result.iterateAll())); + FieldValueList row = result.iterateAll().iterator().next(); + assertEquals(42L, row.get("id").getLongValue()); + assertEquals(50L, result.getTotalSlotMs().longValue()); + assertEquals(100L, result.getTotalBytesBilled().longValue()); + assertEquals(200L, result.getTotalBytesProcessed().longValue()); + assertEquals(StatementType.SELECT, result.getStatementType()); + } + + @Test + void testQueryWithArrowFormatIncompleteJobMissingJobReference() throws Exception { + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-arrow-incomplete-no-job") + .setJobComplete(false); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) + .thenReturn(queryResponsePb); + + bigquery = options.getService(); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); + + BigQueryException e = assertThrows(BigQueryException.class, () -> bigquery.query(config)); assertTrue( e.getMessage() - .contains("Arrow results format for slow query path execution is not yet supported.")); + .contains( + "Query is incomplete, but no job reference was returned to await completion.")); + } + + @Test + void testQueryWithArrowFormatIncompleteJobJobNotFound() throws Exception { + JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-arrow-incomplete") + .setJobComplete(false) + .setJobReference(queryJob.toPb()); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) + .thenReturn(queryResponsePb); + when(bigqueryRpcMock.getJobSkipExceptionTranslation(eq(PROJECT), eq(JOB), any(), any())) + .thenThrow(new BigQueryException(404, "Not Found")); + + bigquery = options.getService(); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); + + BigQueryException e = assertThrows(BigQueryException.class, () -> bigquery.query(config)); + assertTrue(e.getMessage().contains("Job no longer exists or could not be retrieved")); } @Test