diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResult.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResult.java new file mode 100644 index 000000000000..6d4efa893fdb --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResult.java @@ -0,0 +1,66 @@ +/* + * 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 com.google.api.core.BetaApi; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; + +/** + * [Beta] A query result container providing zero-copy access to Apache Arrow {@link + * VectorSchemaRoot} batches. + * + *

Important Usage Warning: The {@link VectorSchemaRoot} returned by the iterator is a + * single, shared, mutated instance across iterations. Data in the root is only valid during the + * current iteration step and will be overwritten or cleared on the next call to {@link + * java.util.Iterator#next()}. Callers needing data across iteration steps must copy the data out of + * the vectors before advancing the iterator. + * + *

Implementations manage direct off-heap native memory buffers. Callers must invoke {@link + * #close()} (idiomatically via a {@code try-with-resources} block) to ensure native allocations and + * underlying gRPC streaming channels are deterministically released. + */ +@BetaApi +public interface ArrowQueryResult extends AutoCloseable, Iterable { + + /** Returns the Apache Arrow schema of the result vectors. */ + Schema getArrowSchema(); + + /** + * Returns the job ID associated with the query execution, or {@code null} if no job was created + * (e.g. when optional job creation was used). + */ + JobId getJobId(); + + /** Returns the query ID associated with the query execution, or {@code null} if unavailable. */ + String getQueryId(); + + /** + * Returns the reason a job was created when optional job creation was requested, or {@code null} + * if no job was created or if the query ran via the fallback path. + */ + JobCreationReason getJobCreationReason(); + + /** Returns the total number of rows across all batches if known, or {@code -1} if unknown. */ + long getTotalRows(); + + /** + * Releases underlying direct off-heap memory allocations and closes any active stream channels. + */ + @Override + void close(); +} diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java new file mode 100644 index 000000000000..9fd92fb4d1f4 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java @@ -0,0 +1,592 @@ +/* + * 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 com.google.api.gax.rpc.ServerStream; +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.cloud.bigquery.storage.v1.ReadSession; +import com.google.protobuf.ByteString; +import java.io.IOException; +import java.nio.channels.Channels; +import java.nio.channels.ReadableByteChannel; +import java.util.Iterator; +import java.util.NoSuchElementException; +import java.util.concurrent.locks.ReentrantLock; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.VectorLoader; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ReadChannel; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.ipc.message.MessageSerializer; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; + +/** + * Implementation of {@link ArrowQueryResult} that provides zero-copy streaming of Apache Arrow + * {@link VectorSchemaRoot} batches across initial REST response and subsequent gRPC stream. + */ +class ArrowQueryResultImpl implements ArrowQueryResult { + + private final Schema arrowSchema; + private final JobId jobId; + private final String queryId; + private final JobCreationReason jobCreationReason; + private final long totalRows; + private final byte[] initialRecordBatchBytes; + private final String streamName; + private final BigQueryReadClient readClient; + + private final BufferAllocator allocator; + private final VectorSchemaRoot root; + private final VectorLoader loader; + private ArrowRecordBatch currentRecordBatch; + + private final ReentrantLock lock = new ReentrantLock(); + private boolean closed = false; + private boolean iteratorCreated = false; + private ServerStream serverStream; + + /** + * Constructs an {@link ArrowQueryResultImpl}. + * + * @param arrowSchema the Arrow schema describing column types, or {@code null} if empty + * @param jobId the ID of the query job + * @param queryId the ID of the fast-path query execution + * @param jobCreationReason the reason why a job was created + * @param totalRows the total number of rows returned by the query, or -1 if unknown + * @param initialRecordBatchBytes serialized Arrow record batch bytes from the REST response + * @param streamName the Storage Read API stream name for reading subsequent rows + * @param readClient the {@link BigQueryReadClient} for streaming rows via gRPC + * @throws IllegalArgumentException if {@code arrowSchema} is null but query data or stream is + * present + */ + ArrowQueryResultImpl( + Schema arrowSchema, + JobId jobId, + String queryId, + JobCreationReason jobCreationReason, + long totalRows, + byte[] initialRecordBatchBytes, + String streamName, + BigQueryReadClient readClient) { + this.arrowSchema = arrowSchema; + this.jobId = jobId; + this.queryId = queryId; + this.jobCreationReason = jobCreationReason; + this.totalRows = totalRows; + this.initialRecordBatchBytes = initialRecordBatchBytes; + this.streamName = streamName; + this.readClient = readClient; + + if (this.arrowSchema != null) { + BufferAllocator alloc = null; + VectorSchemaRoot vRoot = null; + try { + alloc = ArrowDeserializer.createChildAllocator("ArrowQueryResult"); + vRoot = VectorSchemaRoot.create(this.arrowSchema, alloc); + this.loader = new VectorLoader(vRoot); + this.allocator = alloc; + this.root = vRoot; + } catch (Throwable t) { + if (vRoot != null) { + try { + vRoot.close(); + } catch (Throwable suppressed) { + t.addSuppressed(suppressed); + } + } + if (alloc != null) { + try { + alloc.close(); + } catch (Throwable suppressed) { + t.addSuppressed(suppressed); + } + } + throw t; + } + } else { + if ((initialRecordBatchBytes != null && initialRecordBatchBytes.length > 0) + || streamName != null) { + throw new IllegalArgumentException( + "Arrow schema cannot be null when query results or streams are present."); + } + this.allocator = null; + this.root = null; + this.loader = null; + } + } + + /** + * Constructs an {@link ArrowQueryResultImpl} from a BigQuery Storage Read API {@link + * ReadSession}. + * + * @param readSession the read session containing the Arrow schema and stream names + * @param jobId the ID of the associated BigQuery query job + * @param readClient the client used to stream rows from the read session + * @return a new {@link ArrowQueryResultImpl} instance + * @throws BigQueryException if deserializing the Arrow schema from the session fails + */ + static ArrowQueryResultImpl fromReadSession( + ReadSession readSession, JobId jobId, BigQueryReadClient readClient) { + Schema pojoSchema = null; + if (readSession.hasArrowSchema()) { + try { + pojoSchema = + ArrowDeserializer.deserializeSchema( + readSession.getArrowSchema().getSerializedSchema().toByteArray()); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow schema from ReadSession", e); + } + } + + String streamName = null; + if (readSession.getStreamsCount() > 0) { + streamName = readSession.getStreams(0).getName(); + } + + return new ArrowQueryResultImpl( + pojoSchema, + jobId, + /* queryId= */ null, + /* jobCreationReason= */ null, + /* totalRows= */ -1L, + /* initialRecordBatchBytes= */ null, + streamName, + readClient); + } + + @Override + public Schema getArrowSchema() { + return arrowSchema; + } + + @Override + public JobId getJobId() { + return jobId; + } + + @Override + public String getQueryId() { + return queryId; + } + + @Override + public JobCreationReason getJobCreationReason() { + return jobCreationReason; + } + + @Override + public long getTotalRows() { + return totalRows; + } + + @Override + public Iterator iterator() { + lock.lock(); + try { + checkNotClosed(); + if (iteratorCreated) { + throw new IllegalStateException("ArrowQueryResult can only be iterated once"); + } + iteratorCreated = true; + return new VectorBatchIterator(); + } finally { + lock.unlock(); + } + } + + @Override + public void close() { + lock.lock(); + try { + if (closed) { + return; + } + closed = true; + Throwable firstException = null; + + if (serverStream != null) { + try { + serverStream.cancel(); + } catch (Throwable t) { + firstException = t; + } + } + if (currentRecordBatch != null) { + try { + currentRecordBatch.close(); + currentRecordBatch = null; + } catch (Throwable t) { + if (firstException == null) { + firstException = t; + } else { + firstException.addSuppressed(t); + } + } + } + if (root != null) { + try { + root.close(); + } catch (Throwable t) { + if (firstException == null) { + firstException = t; + } else { + firstException.addSuppressed(t); + } + } + } + if (allocator != null) { + try { + allocator.close(); + } catch (Throwable t) { + if (firstException == null) { + firstException = t; + } else { + firstException.addSuppressed(t); + } + } + } + if (firstException instanceof RuntimeException) { + throw (RuntimeException) firstException; + } else if (firstException instanceof Error) { + throw (Error) firstException; + } else if (firstException != null) { + throw new RuntimeException("Failed to close Arrow resources", firstException); + } + } finally { + lock.unlock(); + } + } + + /** + * Asserts that this query result instance has not been closed. + * + * @throws IllegalStateException if the query result has already been closed + */ + private void checkNotClosed() { + if (closed) { + throw new IllegalStateException("ArrowQueryResult has already been closed"); + } + } + + /** + * Loads an {@link ArrowRecordBatch} into the underlying {@link VectorSchemaRoot} and releases the + * previously loaded batch to prevent memory leaks. + * + * @param newBatch the Arrow record batch to load + */ + void loadBatch(ArrowRecordBatch newBatch) { + lock.lock(); + try { + checkNotClosed(); + loader.load(newBatch); + ArrowRecordBatch oldBatch = currentRecordBatch; + currentRecordBatch = newBatch; + if (oldBatch != null) { + oldBatch.close(); + } + } finally { + lock.unlock(); + } + } + + private final class VectorBatchIterator implements Iterator { + private boolean yieldedInitialBatch = false; + private Iterator streamIterator = null; + private boolean streamInitialized = false; + private long totalRowsYielded = 0; + + private ReadRowsResponse peekedResponse = null; + + /** + * Checks whether the enclosing query result has been closed. + * + * @return {@code true} if closed, {@code false} otherwise + */ + private boolean isClosed() { + lock.lock(); + try { + return closed; + } finally { + lock.unlock(); + } + } + + /** + * Checks whether the initial Arrow batch from the query response is pending and unconsumed. + * + * @return {@code true} if an initial batch is present and not yet yielded, {@code false} + * otherwise + */ + private boolean hasInitialBatchToYield() { + lock.lock(); + try { + return !yieldedInitialBatch + && initialRecordBatchBytes != null + && initialRecordBatchBytes.length > 0; + } finally { + lock.unlock(); + } + } + + /** + * Retrieves the active gRPC stream iterator. + * + * @return the {@link Iterator} of {@link ReadRowsResponse} messages, or {@code null} if not + * initialized + */ + private Iterator getStreamIterator() { + lock.lock(); + try { + return streamIterator; + } finally { + lock.unlock(); + } + } + + @Override + public boolean hasNext() { + if (isClosed()) { + return false; + } + if (hasInitialBatchToYield()) { + return true; + } + if (peekedResponse != null) { + return true; + } + try { + ensureStreamInitialized(); + Iterator iterator = getStreamIterator(); + if (iterator == null) { + return false; + } + while (iterator.hasNext()) { + ReadRowsResponse response = iterator.next(); + if (response.hasArrowRecordBatch()) { + peekedResponse = response; + return true; + } + } + return false; + } catch (BigQueryException e) { + if (isClosed()) { + return false; + } + throw e; + } catch (Exception e) { + if (isClosed()) { + return false; + } + throw new BigQueryException(0, "Error reading from Arrow stream", e); + } + } + + @Override + public VectorSchemaRoot next() { + if (isClosed()) { + throw new NoSuchElementException("ArrowQueryResult has already been closed"); + } + + // 1. Yield initial batch from REST response if present + byte[] initialBytes = null; + lock.lock(); + try { + if (!yieldedInitialBatch + && initialRecordBatchBytes != null + && initialRecordBatchBytes.length > 0) { + yieldedInitialBatch = true; + initialBytes = initialRecordBatchBytes; + } else { + yieldedInitialBatch = true; + } + } finally { + lock.unlock(); + } + + if (initialBytes != null) { + try { + loadBatch(initialBytes); + return root; + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load initial Arrow record batch", e); + } + } + + // 2. Stream subsequent batches from gRPC + try { + ReadRowsResponse targetResponse = null; + if (peekedResponse != null) { + targetResponse = peekedResponse; + peekedResponse = null; + } else { + ensureStreamInitialized(); + Iterator iterator = getStreamIterator(); + if (iterator == null || !iterator.hasNext()) { + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } + + while (iterator.hasNext()) { + ReadRowsResponse response = iterator.next(); + if (response.hasArrowRecordBatch()) { + targetResponse = response; + break; + } + } + } + + if (targetResponse == null || !targetResponse.hasArrowRecordBatch()) { + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } + + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = + targetResponse.getArrowRecordBatch(); + try { + loadBatch(batch.getSerializedRecordBatch()); + return root; + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); + } + } catch (NoSuchElementException | BigQueryException e) { + throw e; + } catch (Exception e) { + if (isClosed()) { + throw new NoSuchElementException("Query stream was closed."); + } + throw new BigQueryException(0, "Error reading from Arrow stream", e); + } + } + + /** + * Initializes the gRPC {@code ReadRows} stream via {@link BigQueryReadClient} if more rows + * remain to be consumed and the stream has not yet been started. + * + * @throws BigQueryException if stream initialization fails or required stream parameters are + * missing + */ + private void ensureStreamInitialized() { + ReadRowsRequest request; + lock.lock(); + try { + if (streamInitialized || closed) { + return; + } + if (totalRows >= 0 + && totalRowsYielded >= totalRows + && (yieldedInitialBatch + || initialRecordBatchBytes == null + || initialRecordBatchBytes.length == 0)) { + streamInitialized = true; + return; + } + if (streamName == null || readClient == null) { + if (totalRows > 0 && totalRowsYielded < totalRows) { + throw new BigQueryException( + 0, + "Cannot stream query results: stream name or read client is missing, " + + "but there are more rows to read (totalRows=" + + totalRows + + ", yielded=" + + totalRowsYielded + + ")"); + } + streamInitialized = true; + return; + } + long offset = totalRowsYielded; + request = ReadRowsRequest.newBuilder().setReadStream(streamName).setOffset(offset).build(); + } finally { + lock.unlock(); + } + + ServerStream stream = readClient.readRowsCallable().call(request); + + lock.lock(); + try { + if (closed || streamInitialized) { + try { + stream.cancel(); + } catch (Throwable t) { + // ignore + } + return; + } + serverStream = stream; + streamIterator = stream.iterator(); + streamInitialized = true; + } finally { + lock.unlock(); + } + } + + /** + * Deserializes an Arrow record batch from raw bytes and loads it into the root vector. + * + * @param bytes serialized Arrow record batch bytes + * @throws IOException if deserialization fails + */ + private void loadBatch(byte[] bytes) throws IOException { + loadBatch(new ByteArrayReadableSeekableByteChannel(bytes)); + } + + /** + * Deserializes an Arrow record batch from a protobuf {@link ByteString} and loads it into the + * root vector. + * + * @param byteString serialized Arrow record batch bytes as a {@link ByteString} + * @throws IOException if deserialization fails + */ + private void loadBatch(ByteString byteString) throws IOException { + loadBatch(Channels.newChannel(byteString.newInput())); + } + + /** + * Deserializes an Arrow record batch from a {@link ReadableByteChannel} and loads it into the + * root vector. + * + * @param channel readable byte channel providing serialized Arrow record batch bytes + * @throws IOException if deserialization fails + */ + private void loadBatch(ReadableByteChannel channel) throws IOException { + lock.lock(); + try { + checkNotClosed(); + try (ReadChannel readChannel = new ReadChannel(channel)) { + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch(readChannel, allocator); + if (deserializedBatch == null) { + throw new IOException("Unexpected end of stream when deserializing ArrowRecordBatch"); + } + boolean loaded = false; + try { + ArrowQueryResultImpl.this.loadBatch(deserializedBatch); + loaded = true; + } finally { + if (!loaded) { + deserializedBatch.close(); + } + } + } + totalRowsYielded += root.getRowCount(); + } finally { + lock.unlock(); + } + } + } +} diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryResultTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryResultTest.java new file mode 100644 index 000000000000..17a08d7651ae --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryResultTest.java @@ -0,0 +1,374 @@ +/* + * 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.when; +import static org.mockito.Mockito.withSettings; + +import com.google.api.gax.rpc.ServerStream; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.cloud.bigquery.storage.v1.ArrowSchema; +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.cloud.bigquery.storage.v1.ReadSession; +import com.google.cloud.bigquery.storage.v1.ReadStream; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.channels.Channels; +import java.nio.charset.StandardCharsets; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.VarCharVector; +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.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class ArrowQueryResultTest { + + private BufferAllocator testAllocator; + + @BeforeEach + void setUp() { + testAllocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + void tearDown() { + testAllocator.close(); + } + + private Schema createTestArrowSchema() { + Field idField = + new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), ImmutableList.of()); + Field nameField = + new Field("name", FieldType.nullable(new ArrowType.Utf8()), ImmutableList.of()); + return new Schema(ImmutableList.of(idField, nameField)); + } + + private byte[] createTestBatchBytes(List ids, List names) throws IOException { + BigIntVector idVector = new BigIntVector("id", testAllocator); + idVector.allocateNew(ids.size()); + for (int i = 0; i < ids.size(); i++) { + if (ids.get(i) != null) { + idVector.set(i, ids.get(i)); + } else { + idVector.setNull(i); + } + } + idVector.setValueCount(ids.size()); + + VarCharVector nameVector = new VarCharVector("name", testAllocator); + nameVector.allocateNew(names.size()); + for (int i = 0; i < names.size(); i++) { + if (names.get(i) != null) { + nameVector.set(i, names.get(i).getBytes(StandardCharsets.UTF_8)); + } else { + nameVector.setNull(i); + } + } + nameVector.setValueCount(names.size()); + + VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(idVector, nameVector)); + VectorUnloader unloader = new VectorUnloader(root); + ArrowRecordBatch recordBatch = unloader.getRecordBatch(); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + WriteChannel channel = new WriteChannel(Channels.newChannel(out)); + MessageSerializer.serialize(channel, recordBatch); + + recordBatch.close(); + root.close(); + return out.toByteArray(); + } + + private BigQueryReadClient createMockReadClient( + ServerStreamingCallable mockCallable) { + BigQueryReadClient mockClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + when(mockClient.readRowsCallable()).thenReturn(mockCallable); + return mockClient; + } + + @Test + void testSingleBatchIteration() throws IOException { + Schema arrowSchema = createTestArrowSchema(); + byte[] batchBytes = + createTestBatchBytes(ImmutableList.of(1L, 2L), ImmutableList.of("Alice", "Bob")); + JobId jobId = JobId.of("test-project", "job_123"); + + try (ArrowQueryResult result = + new ArrowQueryResultImpl( + arrowSchema, + jobId, + /* queryId= */ null, + /* jobCreationReason= */ null, + 2L, + batchBytes, + null, + null)) { + assertEquals(arrowSchema, result.getArrowSchema()); + assertEquals(jobId, result.getJobId()); + assertEquals(2L, result.getTotalRows()); + + Iterator it = result.iterator(); + assertTrue(it.hasNext()); + VectorSchemaRoot root = it.next(); + assertEquals(2, root.getRowCount()); + BigIntVector idVec = (BigIntVector) root.getVector("id"); + assertEquals(1L, idVec.get(0)); + assertEquals(2L, idVec.get(1)); + + assertFalse(it.hasNext()); + assertThrows(NoSuchElementException.class, it::next); + } + } + + @Test + void testIteratorCannotBeCreatedTwice() throws IOException { + Schema arrowSchema = createTestArrowSchema(); + byte[] batchBytes = createTestBatchBytes(ImmutableList.of(1L), ImmutableList.of("Alice")); + + try (ArrowQueryResult result = + new ArrowQueryResultImpl( + arrowSchema, + JobId.of("j1"), + /* queryId= */ null, + /* jobCreationReason= */ null, + 1L, + batchBytes, + null, + null)) { + Iterator it1 = result.iterator(); + assertNotNull(it1); + assertThrows(IllegalStateException.class, result::iterator); + } + } + + @Test + void testCloseIsIdempotentAndReleasesResources() throws IOException { + Schema arrowSchema = createTestArrowSchema(); + byte[] batchBytes = createTestBatchBytes(ImmutableList.of(1L), ImmutableList.of("Alice")); + + ArrowQueryResult result = + new ArrowQueryResultImpl( + arrowSchema, + JobId.of("j1"), + /* queryId= */ null, + /* jobCreationReason= */ null, + 1L, + batchBytes, + null, + null); + result.close(); + result.close(); + assertThrows(IllegalStateException.class, result::iterator); + } + + @Test + void testMultiBatchStreaming() throws IOException { + Schema arrowSchema = createTestArrowSchema(); + byte[] initialBatchBytes = + createTestBatchBytes(ImmutableList.of(1L, 2L), ImmutableList.of("A", "B")); + byte[] streamingBatchBytes = + createTestBatchBytes(ImmutableList.of(3L, 4L, 5L), ImmutableList.of("C", "D", "E")); + + @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(streamingBatchBytes)) + .build(); + ReadRowsResponse streamResponse = + ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build(); + + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator()); + + BigQueryReadClient mockClient = createMockReadClient(mockCallable); + + String streamName = "projects/p/locations/l/jobs/j/streams/_default"; + try (ArrowQueryResult result = + new ArrowQueryResultImpl( + arrowSchema, + JobId.of("j"), + /* queryId= */ null, + /* jobCreationReason= */ null, + 5L, + initialBatchBytes, + streamName, + mockClient)) { + Iterator it = result.iterator(); + + // Batch 1 (initial REST response) + assertTrue(it.hasNext()); + VectorSchemaRoot root1 = it.next(); + assertEquals(2, root1.getRowCount()); + + // Batch 2 (streaming gRPC response) + assertTrue(it.hasNext()); + VectorSchemaRoot root2 = it.next(); + assertEquals(3, root2.getRowCount()); + + assertFalse(it.hasNext()); + assertThrows(NoSuchElementException.class, it::next); + } + } + + @Test + void testQueryIdAndJobCreationReason() throws IOException { + Schema schema = createTestArrowSchema(); + byte[] batch1 = createTestBatchBytes(ImmutableList.of(1L), ImmutableList.of("Alice")); + JobId jobId = JobId.of("p", "j"); + String queryId = "query-12345"; + JobCreationReason reason = + JobCreationReason.fromPb( + new com.google.api.services.bigquery.model.JobCreationReason().setCode("REQUESTED")); + + ArrowQueryResultImpl result = + new ArrowQueryResultImpl(schema, jobId, queryId, reason, 1L, batch1, null, null); + + assertEquals(queryId, result.getQueryId()); + assertNotNull(result.getJobCreationReason()); + assertEquals(JobCreationReason.Code.REQUESTED, result.getJobCreationReason().getCode()); + assertEquals(jobId, result.getJobId()); + + try (ArrowQueryResult res = result) { + Iterator it = result.iterator(); + assertTrue(it.hasNext()); + VectorSchemaRoot root = it.next(); + assertEquals(1, root.getRowCount()); + assertFalse(it.hasNext()); + } + } + + @Test + void testStatelessSingleBatchIterationWithoutJobId() throws IOException { + Schema schema = createTestArrowSchema(); + byte[] batch1 = + createTestBatchBytes(ImmutableList.of(1L, 2L), ImmutableList.of("Alice", "Bob")); + + ArrowQueryResultImpl result = + new ArrowQueryResultImpl( + schema, + /* jobId= */ null, + /* queryId= */ "stateless-q-1", + /* jobCreationReason= */ null, + 2L, + batch1, + /* streamName= */ null, + /* readClient= */ null); + + assertNull(result.getJobId()); + assertEquals("stateless-q-1", result.getQueryId()); + assertNull(result.getJobCreationReason()); + + try (ArrowQueryResult res = result) { + Iterator it = result.iterator(); + assertTrue(it.hasNext()); + VectorSchemaRoot root = it.next(); + assertEquals(2, root.getRowCount()); + assertFalse(it.hasNext()); + } + } + + @Test + void testFromReadSessionFallback() throws IOException { + Schema schema = createTestArrowSchema(); + byte[] batch1 = createTestBatchBytes(ImmutableList.of(10L), ImmutableList.of("FallbackUser")); + + ByteArrayOutputStream schemaOut = new ByteArrayOutputStream(); + WriteChannel schemaChannel = new WriteChannel(Channels.newChannel(schemaOut)); + MessageSerializer.serialize(schemaChannel, schema); + + ArrowSchema arrowSchemaPb = + ArrowSchema.newBuilder() + .setSerializedSchema(ByteString.copyFrom(schemaOut.toByteArray())) + .build(); + ReadStream streamPb = + ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/str1").build(); + ReadSession readSession = + ReadSession.newBuilder().setArrowSchema(arrowSchemaPb).addStreams(streamPb).build(); + + ReadRowsResponse response = + ReadRowsResponse.newBuilder() + .setArrowRecordBatch( + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(batch1)) + .build()) + .build(); + + @SuppressWarnings("unchecked") + ServerStream mockServerStream = + mock(ServerStream.class, withSettings().withoutAnnotations()); + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(response).iterator()); + + @SuppressWarnings("unchecked") + ServerStreamingCallable mockCallable = + mock(ServerStreamingCallable.class, withSettings().withoutAnnotations()); + when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream); + + BigQueryReadClient mockReadClient = createMockReadClient(mockCallable); + + JobId jobId = JobId.of("p", "fallback-job"); + ArrowQueryResultImpl result = + ArrowQueryResultImpl.fromReadSession(readSession, jobId, mockReadClient); + + assertEquals(jobId, result.getJobId()); + assertNull(result.getQueryId()); + assertNull(result.getJobCreationReason()); + + try (ArrowQueryResult res = result) { + Iterator it = result.iterator(); + assertTrue(it.hasNext()); + VectorSchemaRoot root = it.next(); + assertEquals(1, root.getRowCount()); + BigIntVector idVector = (BigIntVector) root.getVector("id"); + assertEquals(10L, idVector.get(0)); + assertFalse(it.hasNext()); + } + } +} diff --git a/java-bigquery/google-cloud-bigquery/src/test/resources/mockito-extensions/org.mockito.plugins.MockMaker b/java-bigquery/google-cloud-bigquery/src/test/resources/mockito-extensions/org.mockito.plugins.MockMaker new file mode 100644 index 000000000000..1f0955d450f0 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/resources/mockito-extensions/org.mockito.plugins.MockMaker @@ -0,0 +1 @@ +mock-maker-inline