From 82693bf941b6dfb950a215e2bed3a593bd74210d Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 16:15:57 -0400 Subject: [PATCH 01/11] feat(bigquery): accelerate row-based query() with Arrow wire format --- .../google/cloud/bigquery/BigQueryImpl.java | 105 ++++-- .../cloud/bigquery/BigQueryImplTest.java | 347 +++++++++++++++++- 2 files changed, 421 insertions(+), 31 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 b0b5756c6866..df6e87426d12 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 @@ -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; @@ -76,6 +77,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; @@ -2114,7 +2116,7 @@ public TableDataList call() throws IOException { } } - private static Iterable transformTableData( + private static List transformTableData( Iterable tableDataPb, final Schema schema, boolean useInt64Timestamps) { return ImmutableList.copyOf( Iterables.transform( @@ -2353,8 +2355,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) { @@ -2382,25 +2397,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 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 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) @@ -2417,16 +2482,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) @@ -2452,11 +2514,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 = @@ -2511,6 +2568,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) { 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 cb67ac54aa4a..0d57d4cb3546 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 @@ -37,6 +37,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; import com.google.api.client.googleapis.json.GoogleJsonError; import com.google.api.client.googleapis.json.GoogleJsonResponseException; @@ -45,6 +46,8 @@ import com.google.api.gax.paging.Page; import com.google.api.gax.retrying.ResultRetryAlgorithm; import com.google.api.gax.retrying.TimedAttemptSettings; +import com.google.api.gax.rpc.ServerStream; +import com.google.api.gax.rpc.ServerStreamingCallable; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; import com.google.api.services.bigquery.model.JobConfigurationQuery; @@ -70,6 +73,9 @@ import com.google.cloud.bigquery.spi.BigQueryRpcFactory; 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.ReadRowsRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; import com.google.common.base.Function; import com.google.common.base.Supplier; import com.google.common.collect.ImmutableList; @@ -77,13 +83,27 @@ import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; +import com.google.common.io.BaseEncoding; +import com.google.protobuf.ByteString; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; import java.io.IOException; import java.math.BigInteger; import java.net.ConnectException; import java.net.UnknownHostException; +import java.nio.channels.Channels; import java.util.Collections; import java.util.List; import java.util.Map; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.ipc.WriteChannel; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.ipc.message.MessageSerializer; +import org.apache.arrow.vector.types.pojo.ArrowType; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -2905,41 +2925,348 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException } @Test - void testQueryThrowsWhenArrowResultsFormat() { + void testQueryArrowDefaultsToJobCreationOptional() throws IOException, InterruptedException { + QueryJobConfiguration config = QueryJobConfiguration.newBuilder("SELECT 1").build(); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-optional-1") + .setJobComplete(true) + .setTotalRows(java.math.BigInteger.ZERO); + + ArgumentCaptor requestPbCapture = ArgumentCaptor.forClass(QueryRequest.class); + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) + .thenReturn(queryResponsePb); + + bigquery = options.getService(); + ArrowQueryResult result = bigquery.queryArrow(config); + assertNotNull(result); + assertEquals("q-optional-1", result.getQueryId()); + assertNull(result.getJobId()); + + QueryRequest requestPb = requestPbCapture.getValue(); + assertEquals("JOB_CREATION_OPTIONAL", requestPb.getJobCreationMode()); + assertEquals("ARROW", requestPb.getQueryResultsFormat()); + } + + @Test + void testQueryArrowResultsFormatUnsupportedConfiguration() { QueryJobConfiguration config = QueryJobConfiguration.newBuilder("SELECT 1") .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setDestinationTable(TableId.of("dataset", "table")) .build(); bigquery = options.getService(); IllegalArgumentException exception = assertThrows(IllegalArgumentException.class, () -> bigquery.query(config)); - assertTrue(exception.getMessage().contains("Use queryArrow() instead")); + assertTrue( + exception + .getMessage() + .contains("Arrow results format is only supported for fast query path execution")); } @Test - void testQueryArrowDefaultsToJobCreationOptional() throws IOException, InterruptedException { - QueryJobConfiguration config = QueryJobConfiguration.newBuilder("SELECT 1").build(); + void testQueryWithArrowFormatFastPath() throws IOException, InterruptedException { + 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(); + } + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = new com.google.api.services.bigquery.model.QueryResponse() - .setQueryId("q-optional-1") + .setQueryId("q-arrow-1") .setJobComplete(true) - .setTotalRows(java.math.BigInteger.ZERO); + .setTotalRows(BigInteger.ONE) + .setArrowSchema( + new com.google.api.services.bigquery.model.ArrowSchema() + .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes))); ArgumentCaptor requestPbCapture = ArgumentCaptor.forClass(QueryRequest.class); when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) .thenReturn(queryResponsePb); + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT 1 as id") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); bigquery = options.getService(); - ArrowQueryResult result = bigquery.queryArrow(config); + TableResult result = bigquery.query(config); assertNotNull(result); - assertEquals("q-optional-1", result.getQueryId()); - assertNull(result.getJobId()); + assertEquals("q-arrow-1", result.getQueryId()); + assertNotNull(result.getSchema()); + assertEquals(1, result.getSchema().getFields().size()); + assertEquals("id", result.getSchema().getFields().get(0).getName()); QueryRequest requestPb = requestPbCapture.getValue(); - assertEquals("JOB_CREATION_OPTIONAL", requestPb.getJobCreationMode()); assertEquals("ARROW", requestPb.getQueryResultsFormat()); } + @Test + void testQueryWithArrowFormatMultiplePages() throws IOException, InterruptedException { + 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(); + } + + // Prepare page 2 Arrow batch for streaming + byte[] page2BatchBytes; + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + BigIntVector idVector = new BigIntVector("id", allocator); + idVector.allocateNew(1); + idVector.set(0, 2L); + 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); + page2BatchBytes = out.toByteArray(); + } + } finally { + idVector.close(); + } + } + + 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-multipage") + .setJobComplete(true) + .setJobReference(queryJob.toPb()) + .setTotalRows(BigInteger.valueOf(2L)) + .setPageToken("1") + .setArrowSchema( + new com.google.api.services.bigquery.model.ArrowSchema() + .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes))); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) + .thenReturn(queryResponsePb); + + // Mock BigQueryReadClient for page 2 + @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); + + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(page2BatchBytes)) + .build(); + ReadRowsResponse streamResponse = + ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build(); + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator()); + + BigQueryReadClient mockReadClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + when(mockReadClient.readRowsCallable()).thenReturn(mockCallable); + + bigquery = options.getService(); + ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); + TableResult result = bigquery.query(config); + assertNotNull(result); + assertEquals("q-arrow-multipage", result.getQueryId()); + assertTrue(result.hasNextPage()); + assertEquals("1", result.getNextPageToken()); + + Page page2 = result.getNextPage(); + assertNotNull(page2); + List page2Rows = ImmutableList.copyOf(page2.getValues()); + assertEquals(1, page2Rows.size()); + assertEquals("2", page2Rows.get(0).get(0).getStringValue()); + + verify(mockCallable).call(any(ReadRowsRequest.class)); + } + + @Test + void testQueryWithArrowFormatMultiplePagesWithMaxResults() + throws IOException, InterruptedException { + 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(); + } + + // Prepare page 2 Arrow batch for streaming with 2 rows + byte[] page2BatchBytes; + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + BigIntVector idVector = new BigIntVector("id", allocator); + idVector.allocateNew(2); + idVector.set(0, 2L); + idVector.set(1, 3L); + idVector.setValueCount(2); + 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); + page2BatchBytes = out.toByteArray(); + } + } finally { + idVector.close(); + } + } + + 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-multipage-maxresults") + .setJobComplete(true) + .setJobReference(queryJob.toPb()) + .setTotalRows(BigInteger.valueOf(3L)) + .setPageToken("1") + .setArrowSchema( + new com.google.api.services.bigquery.model.ArrowSchema() + .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes))); + + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class))) + .thenReturn(queryResponsePb); + + // Mock BigQueryReadClient for page 2 + @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); + + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(page2BatchBytes)) + .build(); + ReadRowsResponse streamResponse = + ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build(); + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator()); + + BigQueryReadClient mockReadClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + when(mockReadClient.readRowsCallable()).thenReturn(mockCallable); + + bigquery = options.getService(); + ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient); + + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setMaxResults(2L) + .build(); + TableResult result = bigquery.query(config); + assertNotNull(result); + assertTrue(result.hasNextPage()); + + Page page2 = result.getNextPage(); + assertNotNull(page2); + List page2Rows = ImmutableList.copyOf(page2.getValues()); + // Since maxResults is 2 and initialRowOffset is 1, page2 should only contain 1 row even though + // stream returned 2 rows + assertEquals(1, page2Rows.size()); + assertEquals("2", page2Rows.get(0).get(0).getStringValue()); + // Since totalRowsReturned == maxResults, hasNextPage must be false + assertFalse(page2.hasNextPage()); + assertNull(page2.getNextPage()); + + // When maxResults is 1, initialRowOffset (1) already reaches maxResults, so hasNextPage is + // false immediately + QueryJobConfiguration configMax1 = + QueryJobConfiguration.newBuilder("SELECT id FROM test") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setMaxResults(1L) + .build(); + TableResult resultMax1 = bigquery.query(configMax1); + assertNotNull(resultMax1); + assertFalse(resultMax1.hasNextPage()); + assertNull(resultMax1.getNextPage()); + } + + @Test + void testArrowQueryPageFetcherSerialization() throws Exception { + 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(); + } + + 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-multipage-ser") + .setJobComplete(true) + .setJobReference(queryJob.toPb()) + .setTotalRows(BigInteger.valueOf(2L)) + .setPageToken("1") + .setArrowSchema( + new com.google.api.services.bigquery.model.ArrowSchema() + .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes))); + + 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(); + TableResult result = bigquery.query(config); + assertNotNull(result); + assertTrue(result.hasNextPage()); + assertEquals("1", result.getNextPageToken()); + + // Serialize and deserialize TableResult + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (java.io.ObjectOutputStream oos = new java.io.ObjectOutputStream(baos)) { + oos.writeObject(result); + } + + TableResult deserializedResult; + try (java.io.ObjectInputStream ois = + new java.io.ObjectInputStream(new ByteArrayInputStream(baos.toByteArray()))) { + deserializedResult = (TableResult) ois.readObject(); + } + + assertNotNull(deserializedResult); + assertEquals(result.getSchema(), deserializedResult.getSchema()); + assertEquals(result.getTotalRows(), deserializedResult.getTotalRows()); + assertEquals(result.getQueryId(), deserializedResult.getQueryId()); + assertEquals("1", deserializedResult.getNextPageToken()); + assertTrue(deserializedResult.hasNextPage()); + } + @Test void testGetQueryResults() throws IOException { JobId queryJob = JobId.of(JOB); From 56eaebb8da0be877bcf0af2a80afcf4d92ca163a Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 18:43:29 -0400 Subject: [PATCH 02/11] fix(bigquery): validate serialized Arrow schema exists in query response --- .../google/cloud/bigquery/BigQueryImpl.java | 3 +++ .../cloud/bigquery/BigQueryImplTest.java | 24 +++++++++++++++++++ 2 files changed, 27 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 df6e87426d12..90e893a694b7 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 @@ -2360,6 +2360,9 @@ public com.google.api.services.bigquery.model.QueryResponse call() byte[] arrowSchemaBytes = null; if (results.getJobComplete() && (results.getSchema() != null || isArrow)) { if (isArrow) { + if (results.getArrowSchema().getSerializedSchema() == null) { + throw new BigQueryException(0, "Arrow schema is missing from the response"); + } arrowSchemaBytes = results.getArrowSchema().decodeSerializedSchema(); try { arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes); 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 0d57d4cb3546..c7864c3fdedc 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 @@ -3267,6 +3267,30 @@ void testArrowQueryPageFetcherSerialization() throws Exception { assertTrue(deserializedResult.hasNextPage()); } + @Test + void testQueryWithArrowFormatMissingSerializedSchema() { + 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-missing-schema") + .setJobComplete(true) + .setJobReference(queryJob.toPb()) + .setTotalRows(BigInteger.valueOf(2L)) + .setArrowSchema(new com.google.api.services.bigquery.model.ArrowSchema()); + + 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 schema is missing from the response")); + } + @Test void testGetQueryResults() throws IOException { JobId queryJob = JobId.of(JOB); From 9293b9dc4afd4b0270686d0f8570c6a00c8adc1b Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 18:48:11 -0400 Subject: [PATCH 03/11] fix(bigquery): enforce maxResults on first page and check jobComplete for hasMorePages --- .../google/cloud/bigquery/BigQueryImpl.java | 36 +++++++++---------- 1 file changed, 16 insertions(+), 20 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 90e893a694b7..3fa406e4180f 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 @@ -2402,31 +2402,27 @@ public com.google.api.services.bigquery.model.QueryResponse call() Collection 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(); + 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 = - results.getRows() != null - ? transformTableData( - results.getRows(), - schema, - getOptions().getDataFormatOptions().useInt64Timestamp()) - : ImmutableList.of(); + transformTableData( + results.getRows(), schema, getOptions().getDataFormatOptions().useInt64Timestamp()); + } + + if (content.getMaxResults() != null && firstPageRows.size() > content.getMaxResults()) { + firstPageRows = + ImmutableList.copyOf(Iterables.limit(firstPageRows, content.getMaxResults().intValue())); } - boolean hasMorePages = results.getPageToken() != null; + boolean hasMorePages = results.getPageToken() != null && results.getJobComplete(); long initialRowOffset = 0L; if (hasMorePages && isArrow) { Long parsedOffset = Longs.tryParse(results.getPageToken()); From e55dcfaeed47a672ed1e483147ff95dc193d2a45 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 18:50:39 -0400 Subject: [PATCH 04/11] fix(bigquery): guard against null ArrowRecordBatch and simplify BigInteger reference --- .../google/cloud/bigquery/BigQueryImpl.java | 21 ++++++++++++------- .../cloud/bigquery/BigQueryImplTest.java | 2 +- 2 files changed, 14 insertions(+), 9 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 3fa406e4180f..2147b203d565 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 @@ -2402,14 +2402,19 @@ public com.google.api.services.bigquery.model.QueryResponse call() Collection firstPageRows; if (isArrow) { - try { - firstPageRows = - ArrowDeserializer.deserializeRecordBatch( - results.getArrowRecordBatch().decodeSerializedRecordBatch(), - schema, - arrowSchemaPojo); - } catch (IOException e) { - throw new BigQueryException(0, "Failed to deserialize Arrow record batch", e); + 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); + } } } else { firstPageRows = 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 c7864c3fdedc..f4f9995ca8d7 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 @@ -2931,7 +2931,7 @@ void testQueryArrowDefaultsToJobCreationOptional() throws IOException, Interrupt new com.google.api.services.bigquery.model.QueryResponse() .setQueryId("q-optional-1") .setJobComplete(true) - .setTotalRows(java.math.BigInteger.ZERO); + .setTotalRows(BigInteger.ZERO); ArgumentCaptor requestPbCapture = ArgumentCaptor.forClass(QueryRequest.class); when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) From 9c6629b46e3ffc40ad75f7fc86c4b62928310ff0 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 18:54:47 -0400 Subject: [PATCH 05/11] fix(bigquery): declare throws Exception on testQueryWithArrowFormatMissingSerializedSchema --- .../test/java/com/google/cloud/bigquery/BigQueryImplTest.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 f4f9995ca8d7..9428456748e8 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 @@ -3268,7 +3268,7 @@ void testArrowQueryPageFetcherSerialization() throws Exception { } @Test - void testQueryWithArrowFormatMissingSerializedSchema() { + void testQueryWithArrowFormatMissingSerializedSchema() 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() From d2b49c9157b0eca3f783db457ee88bcc15eee96e Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 18:57:07 -0400 Subject: [PATCH 06/11] fix(bigquery): guard against null results.getRows() when transforming table data --- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 2147b203d565..9541ca819e55 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 @@ -2419,7 +2419,9 @@ public com.google.api.services.bigquery.model.QueryResponse call() } else { firstPageRows = transformTableData( - results.getRows(), schema, getOptions().getDataFormatOptions().useInt64Timestamp()); + results.getRows() != null ? results.getRows() : ImmutableList.of(), + schema, + getOptions().getDataFormatOptions().useInt64Timestamp()); } if (content.getMaxResults() != null && firstPageRows.size() > content.getMaxResults()) { From 400d7c3b8bdf67e283adce8b81548b922a63efbe Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 20:12:58 -0400 Subject: [PATCH 07/11] chore(bigquery): isolate Arrow fast query logic into queryRpcArrow --- .../google/cloud/bigquery/BigQueryImpl.java | 231 +++++++++++++----- 1 file changed, 170 insertions(+), 61 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 9541ca819e55..74420d808570 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 @@ -2355,24 +2355,8 @@ public com.google.api.services.bigquery.model.QueryResponse call() long numRows; Schema schema; - 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) { - if (results.getArrowSchema().getSerializedSchema() == null) { - throw new BigQueryException(0, "Arrow schema is missing from the response"); - } - 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.getJobComplete() && results.getSchema() != null) { + schema = Schema.fromPb(results.getSchema()); if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) { numRows = 0L; } else if (results.getNumDmlAffectedRows() != null) { @@ -2400,28 +2384,158 @@ public com.google.api.services.bigquery.model.QueryResponse call() SessionInfo sessionInfo = results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null; - Collection firstPageRows; - if (isArrow) { - 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); - } + if (results.getPageToken() != null) { + JobId jobId = JobId.fromPb(results.getJobReference()); + String cursor = results.getPageToken(); + 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() != null ? results.getRows() : ImmutableList.of(), + schema, + getOptions().getDataFormatOptions().useInt64Timestamp()))) + .setJobId(jobId) + .setQueryId(results.getQueryId()) + .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason())) + .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L) + .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, + transformTableData( + results.getRows() != null ? results.getRows() : ImmutableList.of(), + schema, + getOptions().getDataFormatOptions().useInt64Timestamp()))) + // 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) + .setStatementType(statementType) + .setTotalBytesBilled(totalBytesBilled) + .setTotalBytesProcessed(totalBytesProcessed) + .setTotalSlotMs(totalSlotMs) + .setNumDmlAffectedRows(numDmlAffectedRows) + .setSessionInfo(sessionInfo) + .build(); + } + + 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( + new Callable() { + @Override + public com.google.api.services.bigquery.model.QueryResponse call() + throws IOException { + return 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 bigQueryErrors = + Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION); + throw new BigQueryException(bigQueryErrors); + } + + if (!Boolean.TRUE.equals(results.getJobComplete()) || results.getArrowSchema() == null) { + 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"); + } + + 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 { - firstPageRows = - transformTableData( - results.getRows() != null ? results.getRows() : ImmutableList.of(), - schema, - getOptions().getDataFormatOptions().useInt64Timestamp()); + 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; + + Collection 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); + } } if (content.getMaxResults() != null && firstPageRows.size() > content.getMaxResults()) { @@ -2429,9 +2543,10 @@ public com.google.api.services.bigquery.model.QueryResponse call() ImmutableList.copyOf(Iterables.limit(firstPageRows, content.getMaxResults().intValue())); } - boolean hasMorePages = results.getPageToken() != null && results.getJobComplete(); + boolean hasMorePages = + results.getPageToken() != null && Boolean.TRUE.equals(results.getJobComplete()); long initialRowOffset = 0L; - if (hasMorePages && isArrow) { + if (hasMorePages) { Long parsedOffset = Longs.tryParse(results.getPageToken()); initialRowOffset = parsedOffset != null ? parsedOffset : firstPageRows.size(); if (content.getMaxResults() != null @@ -2444,30 +2559,21 @@ public com.google.api.services.bigquery.model.QueryResponse call() if (hasMorePages) { JobId jobId = JobId.fromPb(results.getJobReference()); String cursor = results.getPageToken(); - - NextPageFetcher 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)); - } + NextPageFetcher pageFetcher = + new ArrowQueryPageFetcher( + jobId, + schema, + arrowSchemaBytes, + arrowSchemaPojo, + getOptions(), + initialRowOffset, + content.getMaxResults(), + optionMap(options)); return TableResult.newBuilder() .setSchema(schema) .setTotalRows(numRows) - .setPageNoSchema( - new PageImpl<>( - // fetch next pages of results - pageFetcher, cursor, firstPageRows)) + .setPageNoSchema(new PageImpl<>(pageFetcher, cursor, firstPageRows)) .setJobId(jobId) .setQueryId(results.getQueryId()) .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason())) @@ -2480,6 +2586,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() .setSessionInfo(sessionInfo) .build(); } + // only 1 page of result return TableResult.newBuilder() .setSchema(schema) @@ -2489,7 +2596,6 @@ public com.google.api.services.bigquery.model.QueryResponse call() new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)), null, firstPageRows)) - // Return the JobID of the successful job .setJobId( results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null) .setQueryId(results.getQueryId()) @@ -2572,6 +2678,9 @@ && getOptions().getOpenTelemetryTracer() != null) { content.setTimeoutMs(timeoutMs); } + if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) { + return queryRpcArrow(projectId, content, options); + } return queryRpc(projectId, content, options); } From 5765fcdd9b3a70aba242d68791b1e14dbbdd1354 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 20:16:30 -0400 Subject: [PATCH 08/11] fix(bigquery): guard against null jobReference in queryRpcArrow --- .../main/java/com/google/cloud/bigquery/BigQueryImpl.java | 8 ++++++++ 1 file changed, 8 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 74420d808570..c8c7326f583d 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 @@ -2485,6 +2485,10 @@ public com.google.api.services.bigquery.model.QueryResponse call() } 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); } @@ -2557,6 +2561,10 @@ public com.google.api.services.bigquery.model.QueryResponse call() } 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 pageFetcher = From 84cea659ef9fde9cf966f4df9bc46df02516fa0d Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 20:20:26 -0400 Subject: [PATCH 09/11] fix(bigquery): simplify Callable to lambda and remove redundant jobComplete check --- .../java/com/google/cloud/bigquery/BigQueryImpl.java | 11 ++--------- 1 file changed, 2 insertions(+), 9 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 c8c7326f583d..cb6dc4f998bc 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 @@ -2457,13 +2457,7 @@ && getOptions().getOpenTelemetryTracer() != null) { try (Scope queryRpcScope = queryRpc != null ? queryRpc.makeCurrent() : null) { results = BigQueryRetryHelper.runWithRetries( - new Callable() { - @Override - public com.google.api.services.bigquery.model.QueryResponse call() - throws IOException { - return bigQueryRpc.queryRpcSkipExceptionTranslation(projectId, content); - } - }, + () -> bigQueryRpc.queryRpcSkipExceptionTranslation(projectId, content), getOptions().getRetrySettings(), getOptions().getResultRetryAlgorithm(), getOptions().getClock(), @@ -2547,8 +2541,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() ImmutableList.copyOf(Iterables.limit(firstPageRows, content.getMaxResults().intValue())); } - boolean hasMorePages = - results.getPageToken() != null && Boolean.TRUE.equals(results.getJobComplete()); + boolean hasMorePages = results.getPageToken() != null; long initialRowOffset = 0L; if (hasMorePages) { Long parsedOffset = Longs.tryParse(results.getPageToken()); From 051bbde7222c1477dcbcd9b4424f759fbf36395f Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 20:24:05 -0400 Subject: [PATCH 10/11] docs(bigquery): document queryRpcArrow execution and pagination flow --- .../google/cloud/bigquery/BigQueryImpl.java | 19 +++++++++++++++++++ 1 file changed, 19 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 cb6dc4f998bc..681f2bc0b8fe 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 @@ -2438,6 +2438,19 @@ public com.google.api.services.bigquery.model.QueryResponse call() .build(); } + /** + * Executes a fast-path query RPC request expecting Arrow-formatted wire response. + * + *

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 { @@ -2478,6 +2491,7 @@ && getOptions().getOpenTelemetryTracer() != null) { 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( @@ -2491,6 +2505,7 @@ && getOptions().getOpenTelemetryTracer() != 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 { @@ -2520,6 +2535,7 @@ && getOptions().getOpenTelemetryTracer() != null) { SessionInfo sessionInfo = results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null; + // Deserialize first page of rows from the Arrow record batch (if present). Collection firstPageRows; if (results.getArrowRecordBatch() == null || results.getArrowRecordBatch().getSerializedRecordBatch() == null) { @@ -2536,11 +2552,13 @@ && getOptions().getOpenTelemetryTracer() != null) { } } + // 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) { @@ -2553,6 +2571,7 @@ && getOptions().getOpenTelemetryTracer() != null) { } } + // Multi-page results: configure ArrowQueryPageFetcher for subsequent tabledata.list calls. if (hasMorePages) { if (results.getJobReference() == null) { throw new BigQueryException( From e5fde792d15cecccc551c402f7e3a3831f0ca3b4 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Tue, 15 Sep 2026 10:57:49 -0400 Subject: [PATCH 11/11] chore(bigquery): change firstPageRows to List and remove unused import --- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 3 +-- 1 file changed, 1 insertion(+), 2 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 681f2bc0b8fe..07b8bef1c9d3 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 @@ -77,7 +77,6 @@ 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; @@ -2536,7 +2535,7 @@ && getOptions().getOpenTelemetryTracer() != null) { results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null; // Deserialize first page of rows from the Arrow record batch (if present). - Collection firstPageRows; + List firstPageRows; if (results.getArrowRecordBatch() == null || results.getArrowRecordBatch().getSerializedRecordBatch() == null) { firstPageRows = ImmutableList.of();