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 648153872a3e..b0b5756c6866 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 @@ -27,6 +27,8 @@ import com.google.api.gax.core.NoCredentialsProvider; import com.google.api.gax.paging.Page; import com.google.api.gax.rpc.HeaderProvider; +import com.google.api.gax.rpc.NoHeaderProvider; +import com.google.api.gax.rpc.ServerStream; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; import com.google.api.services.bigquery.model.ProjectList; @@ -53,6 +55,8 @@ import com.google.cloud.bigquery.storage.v1.BigQueryReadSettings; import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest; import com.google.cloud.bigquery.storage.v1.DataFormat; +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.common.annotations.VisibleForTesting; import com.google.common.base.Function; @@ -70,10 +74,13 @@ import io.opentelemetry.context.Scope; import java.io.IOException; import java.net.URI; +import java.util.ArrayDeque; import java.util.ArrayList; import java.util.Collections; +import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Queue; import java.util.concurrent.Callable; import java.util.concurrent.locks.ReentrantLock; import java.util.regex.Matcher; @@ -280,9 +287,180 @@ public Page getNextPage() { } } + /** + * NextPageFetcher implementation for queries returning results in Arrow format. Reads subsequent + * pages from the job's default gRPC storage read stream. + * + *

Note: Neither {@link Page} nor {@link TableResult} implements {@link AutoCloseable}. The + * underlying gRPC stream is automatically canceled and resources released when iteration reaches + * the end (or maximum results requested) or when an error occurs. Callers that do not iterate to + * completion rely on server-side stream timeouts and garbage collection to release stream + * resources. + */ + static class ArrowQueryPageFetcher implements NextPageFetcher { + private static final long serialVersionUID = 1L; + private static final long DEFAULT_PAGE_SIZE = 10000L; + + private final JobId jobId; + private final Schema schema; + private final byte[] arrowSchemaBytes; + private final BigQueryOptions serviceOptions; + private final long maxResults; + private final Map optionsMap; + + private transient org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo; + private transient BigQueryReadClient bqReadClient; + private transient ServerStream stream; + private transient Iterator streamIterator; + private transient Queue buffer = new ArrayDeque<>(); + private long totalRowsReturned = 0L; + private boolean streamClosed = false; + + ArrowQueryPageFetcher( + JobId jobId, + Schema schema, + byte[] arrowSchemaBytes, + org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo, + BigQueryOptions serviceOptions, + long initialRowOffset, + Long maxResults, + Map optionsMap) { + this.jobId = jobId; + this.schema = schema; + this.arrowSchemaBytes = arrowSchemaBytes; + this.arrowSchemaPojo = arrowSchemaPojo; + this.serviceOptions = serviceOptions; + this.totalRowsReturned = initialRowOffset; + this.maxResults = maxResults != null ? maxResults : Long.MAX_VALUE; + this.optionsMap = optionsMap; + } + + @Override + public Page getNextPage() { + if (buffer == null) { + buffer = new ArrayDeque<>(); + } + if (streamClosed || totalRowsReturned >= maxResults) { + closeClient(); + return null; + } + + Number optionPageSize = + optionsMap != null ? (Number) optionsMap.get(BigQueryRpc.Option.MAX_RESULTS) : null; + long pageSize = + optionPageSize != null && optionPageSize.longValue() > 0 + ? optionPageSize.longValue() + : DEFAULT_PAGE_SIZE; + List rowBatch = new ArrayList<>((int) Math.min(pageSize, 10000L)); + + try { + String location = + jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation(); + if (location == null) { + throw new BigQueryException( + 0, "Location must be specified to read Arrow rows from storage stream"); + } + + if (bqReadClient == null) { + BigQuery service = serviceOptions.getService(); + if (service instanceof BigQueryImpl) { + BigQueryImpl impl = (BigQueryImpl) service; + bqReadClient = impl.getBigQueryReadClient(); + } else { + throw new IllegalStateException( + "Arrow query result pagination requires an instance of BigQueryImpl to manage BigQueryReadClient lifecycle"); + } + } + + if (streamIterator == null) { + String streamName = + String.format( + "projects/%s/locations/%s/jobs/%s/streams/_default", + jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(), + location, + jobId.getJob()); + + ReadRowsRequest readRowsRequest = + ReadRowsRequest.newBuilder() + .setReadStream(streamName) + .setOffset(totalRowsReturned) + .build(); + + stream = bqReadClient.readRowsCallable().call(readRowsRequest); + streamIterator = stream.iterator(); + } + + if (arrowSchemaPojo == null && arrowSchemaBytes != null) { + arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes); + } + + boolean hasMore = + ArrowDeserializer.loadArrowRows( + streamIterator, + arrowSchemaPojo, + schema, + rowBatch, + buffer, + pageSize, + totalRowsReturned, + maxResults); + + if (rowBatch.isEmpty()) { + streamClosed = true; + closeClient(); + return null; + } + + totalRowsReturned += rowBatch.size(); + + String nextPageToken = null; + if (hasMore && totalRowsReturned < maxResults) { + nextPageToken = String.valueOf(totalRowsReturned); + } else { + streamClosed = true; + closeClient(); + } + + return new PageImpl<>(this, nextPageToken, rowBatch); + + } catch (BigQueryException e) { + streamClosed = true; + closeClient(); + throw e; + } catch (Exception e) { + streamClosed = true; + closeClient(); + throw new BigQueryException(0, "Failed to read Arrow rows from storage stream", e); + } + } + + private void closeClient() { + if (stream != null) { + try { + stream.cancel(); + } catch (Exception e) { + // Ignore cancellation exceptions + } + } + bqReadClient = null; + streamIterator = null; + stream = null; + } + } + private final ReentrantLock readClientLock = new ReentrantLock(); private transient BigQueryReadClient bqReadClient; + @VisibleForTesting + void setBigQueryReadClient(BigQueryReadClient client) { + readClientLock.lock(); + try { + this.bqReadClient = client; + } finally { + readClientLock.unlock(); + } + } + /** * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming * Arrow query results, reusing credentials and channel configuration from this {@link @@ -324,7 +502,7 @@ private static void configureReadSettings( } else { settingsBuilder.setCredentialsProvider(NoCredentialsProvider.create()); } - HeaderProvider headerProvider = options.getMergedHeaderProvider(null); + HeaderProvider headerProvider = options.getMergedHeaderProvider(new NoHeaderProvider()); if (headerProvider != null) { settingsBuilder.setHeaderProvider(headerProvider); } diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryPageFetcherTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryPageFetcherTest.java new file mode 100644 index 000000000000..b572946fef7d --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryPageFetcherTest.java @@ -0,0 +1,349 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; + +import com.google.api.gax.paging.Page; +import com.google.api.gax.rpc.ServerStream; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.cloud.bigquery.spi.v2.BigQueryRpc; +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.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.protobuf.ByteString; +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.io.ObjectInputStream; +import java.io.ObjectOutputStream; +import java.nio.channels.Channels; +import java.util.Collections; +import java.util.List; +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.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class ArrowQueryPageFetcherTest { + + private static final String PROJECT = "test-project"; + private static final String LOCATION = "US"; + private static final String JOB = "test-job"; + + private BufferAllocator allocator; + private org.apache.arrow.vector.types.pojo.Schema arrowSchema; + private Schema bqSchema; + private byte[] schemaBytes; + + @BeforeEach + void setUp() throws IOException { + allocator = new RootAllocator(Long.MAX_VALUE); + 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)))); + bqSchema = ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema); + + try (ByteArrayOutputStream out = new ByteArrayOutputStream()) { + MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema); + schemaBytes = out.toByteArray(); + } + } + + @AfterEach + void tearDown() { + allocator.close(); + } + + private byte[] createBatchBytes(List values) throws IOException { + BigIntVector idVector = new BigIntVector("id", allocator); + idVector.allocateNew(values.size()); + for (int i = 0; i < values.size(); i++) { + idVector.set(i, values.get(i)); + } + idVector.setValueCount(values.size()); + + 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); + return out.toByteArray(); + } + } finally { + idVector.close(); + } + } + + private BigQueryReadClient createMockReadClient( + ServerStreamingCallable mockCallable) { + BigQueryReadClient mockClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + when(mockClient.readRowsCallable()).thenReturn(mockCallable); + return mockClient; + } + + @Test + void testGetNextPage_singlePage() throws IOException { + byte[] batchBytes = createBatchBytes(ImmutableList.of(10L, 20L)); + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(batchBytes)) + .build(); + ReadRowsResponse response = + 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(response).iterator()); + + BigQueryReadClient mockReadClient = createMockReadClient(mockCallable); + + BigQueryOptions options = + BigQueryOptions.newBuilder().setProjectId(PROJECT).setLocation(LOCATION).build(); + BigQuery service = options.getService(); + ((BigQueryImpl) service).setBigQueryReadClient(mockReadClient); + + JobId jobId = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + BigQueryImpl.ArrowQueryPageFetcher fetcher = + new BigQueryImpl.ArrowQueryPageFetcher( + jobId, + bqSchema, + schemaBytes, + arrowSchema, + options, + /* initialRowOffset= */ 0L, + /* maxResults= */ 10L, + Collections.emptyMap()); + + Page page = fetcher.getNextPage(); + assertNotNull(page); + List rows = ImmutableList.copyOf(page.getValues()); + assertEquals(2, rows.size()); + assertEquals("10", rows.get(0).get(0).getStringValue()); + assertEquals("20", rows.get(1).get(0).getStringValue()); + assertNull(page.getNextPageToken()); + assertFalse(page.hasNextPage()); + + verify(mockCallable).call(any(ReadRowsRequest.class)); + } + + @Test + void testGetNextPage_multiplePagesWithPageSizeOption() throws IOException { + byte[] batch1Bytes = createBatchBytes(ImmutableList.of(1L, 2L)); + byte[] batch2Bytes = createBatchBytes(ImmutableList.of(3L, 4L)); + + ReadRowsResponse resp1 = + ReadRowsResponse.newBuilder() + .setArrowRecordBatch( + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(batch1Bytes)) + .build()) + .build(); + ReadRowsResponse resp2 = + ReadRowsResponse.newBuilder() + .setArrowRecordBatch( + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(batch2Bytes)) + .build()) + .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(resp1, resp2).iterator()); + + BigQueryReadClient mockReadClient = createMockReadClient(mockCallable); + + BigQueryOptions options = + BigQueryOptions.newBuilder().setProjectId(PROJECT).setLocation(LOCATION).build(); + BigQuery service = options.getService(); + ((BigQueryImpl) service).setBigQueryReadClient(mockReadClient); + + JobId jobId = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + BigQueryImpl.ArrowQueryPageFetcher fetcher = + new BigQueryImpl.ArrowQueryPageFetcher( + jobId, + bqSchema, + schemaBytes, + arrowSchema, + options, + /* initialRowOffset= */ 0L, + /* maxResults= */ 10L, + ImmutableMap.of(BigQueryRpc.Option.MAX_RESULTS, 2L)); + + Page page1 = fetcher.getNextPage(); + assertNotNull(page1); + List page1Rows = ImmutableList.copyOf(page1.getValues()); + assertEquals(2, page1Rows.size()); + assertEquals("1", page1Rows.get(0).get(0).getStringValue()); + assertEquals("2", page1Rows.get(1).get(0).getStringValue()); + assertTrue(page1.hasNextPage()); + assertEquals("2", page1.getNextPageToken()); + + Page page2 = page1.getNextPage(); + assertNotNull(page2); + List page2Rows = ImmutableList.copyOf(page2.getValues()); + assertEquals(2, page2Rows.size()); + assertEquals("3", page2Rows.get(0).get(0).getStringValue()); + assertEquals("4", page2Rows.get(1).get(0).getStringValue()); + assertFalse(page2.hasNextPage()); + assertNull(page2.getNextPageToken()); + assertNull(page2.getNextPage()); + } + + @Test + void testGetNextPage_respectsMaxResults() throws IOException { + byte[] batchBytes = createBatchBytes(ImmutableList.of(1L, 2L, 3L)); + ReadRowsResponse resp = + ReadRowsResponse.newBuilder() + .setArrowRecordBatch( + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(batchBytes)) + .build()) + .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(resp).iterator()); + + BigQueryReadClient mockReadClient = createMockReadClient(mockCallable); + + BigQueryOptions options = + BigQueryOptions.newBuilder().setProjectId(PROJECT).setLocation(LOCATION).build(); + BigQuery service = options.getService(); + ((BigQueryImpl) service).setBigQueryReadClient(mockReadClient); + + JobId jobId = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + // initialRowOffset = 1, maxResults = 2 -> only 1 more row allowed + BigQueryImpl.ArrowQueryPageFetcher fetcher = + new BigQueryImpl.ArrowQueryPageFetcher( + jobId, + bqSchema, + schemaBytes, + arrowSchema, + options, + /* initialRowOffset= */ 1L, + /* maxResults= */ 2L, + Collections.emptyMap()); + + Page page = fetcher.getNextPage(); + assertNotNull(page); + List rows = ImmutableList.copyOf(page.getValues()); + assertEquals(1, rows.size()); + assertEquals("1", rows.get(0).get(0).getStringValue()); + assertFalse(page.hasNextPage()); + assertNull(page.getNextPageToken()); + assertNull(page.getNextPage()); + } + + @Test + void testGetNextPage_missingLocationThrowsException() { + BigQueryOptions options = BigQueryOptions.newBuilder().setProjectId(PROJECT).build(); + JobId jobId = JobId.of(PROJECT, JOB); // No location set on JobId or options + + BigQueryImpl.ArrowQueryPageFetcher fetcher = + new BigQueryImpl.ArrowQueryPageFetcher( + jobId, + bqSchema, + schemaBytes, + arrowSchema, + options, + /* initialRowOffset= */ 0L, + /* maxResults= */ 10L, + Collections.emptyMap()); + + BigQueryException thrown = assertThrows(BigQueryException.class, fetcher::getNextPage); + assertTrue(thrown.getMessage().contains("Location must be specified")); + } + + @Test + void testSerialization() throws Exception { + BigQueryOptions options = + BigQueryOptions.newBuilder().setProjectId(PROJECT).setLocation(LOCATION).build(); + JobId jobId = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build(); + + BigQueryImpl.ArrowQueryPageFetcher fetcher = + new BigQueryImpl.ArrowQueryPageFetcher( + jobId, + bqSchema, + schemaBytes, + arrowSchema, + options, + /* initialRowOffset= */ 10L, + /* maxResults= */ 10L, + Collections.emptyMap()); + + ByteArrayOutputStream baos = new ByteArrayOutputStream(); + try (ObjectOutputStream oos = new ObjectOutputStream(baos)) { + oos.writeObject(fetcher); + } + + BigQueryImpl.ArrowQueryPageFetcher deserializedFetcher; + try (ObjectInputStream ois = + new ObjectInputStream(new ByteArrayInputStream(baos.toByteArray()))) { + deserializedFetcher = (BigQueryImpl.ArrowQueryPageFetcher) ois.readObject(); + } + + assertNotNull(deserializedFetcher); + java.lang.reflect.Field bufferField = + BigQueryImpl.ArrowQueryPageFetcher.class.getDeclaredField("buffer"); + bufferField.setAccessible(true); + assertNull(bufferField.get(deserializedFetcher)); + + // Calling getNextPage() lazily initializes the transient buffer without throwing NPE + assertNull(deserializedFetcher.getNextPage()); + assertNotNull(bufferField.get(deserializedFetcher)); + } +}