From 7cc96f675a91aba27d4ac1d155bb75d0939ca47b Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 12 Aug 2026 16:54:43 -0400 Subject: [PATCH 01/34] feat(bigquery): add zero-copy queryArrow API for Arrow VectorSchemaRoot streaming --- java-bigquery/google-cloud-bigquery/pom.xml | 9 + .../cloud/bigquery/ArrowQueryResult.java | 60 +++ .../cloud/bigquery/ArrowQueryResultImpl.java | 329 ++++++++++++++++ .../com/google/cloud/bigquery/BigQuery.java | 52 +++ .../google/cloud/bigquery/BigQueryImpl.java | 311 +++++++++++++++ .../cloud/bigquery/QueryRequestInfo.java | 14 +- .../cloud/bigquery/ArrowQueryResultTest.java | 354 ++++++++++++++++++ .../cloud/bigquery/BigQueryImplTest.java | 36 ++ .../cloud/bigquery/it/ITBigQueryTest.java | 47 +++ 9 files changed, 1210 insertions(+), 2 deletions(-) create mode 100644 java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResult.java create mode 100644 java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java create mode 100644 java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryResultTest.java diff --git a/java-bigquery/google-cloud-bigquery/pom.xml b/java-bigquery/google-cloud-bigquery/pom.xml index 765a2e650ee2..3f186b68624e 100644 --- a/java-bigquery/google-cloud-bigquery/pom.xml +++ b/java-bigquery/google-cloud-bigquery/pom.xml @@ -122,6 +122,15 @@ arrow-memory-netty + + com.google.api + gax-grpc + + + io.grpc + grpc-api + + com.google.errorprone error_prone_annotations 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..13c7f2e11bb8 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResult.java @@ -0,0 +1,60 @@ +/* + * 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. + * + *

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..6637f1483fbf --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java @@ -0,0 +1,329 @@ +/* + * 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 java.io.IOException; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +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 final Object lock = new Object(); + private boolean closed = false; + private boolean iteratorCreated = false; + private ServerStream serverStream; + + ArrowQueryResultImpl( + Object arrowSchema, + JobId jobId, + long totalRows, + byte[] initialRecordBatchBytes, + String streamName, + BigQueryReadClient readClient) { + this( + arrowSchema, + jobId, + /* queryId= */ null, + /* jobCreationReason= */ null, + totalRows, + initialRecordBatchBytes, + streamName, + readClient); + } + + ArrowQueryResultImpl( + Object arrowSchema, + JobId jobId, + String queryId, + JobCreationReason jobCreationReason, + long totalRows, + byte[] initialRecordBatchBytes, + String streamName, + BigQueryReadClient readClient) { + if (arrowSchema instanceof Schema) { + this.arrowSchema = (Schema) arrowSchema; + } else { + this.arrowSchema = null; + } + 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) { + this.allocator = ArrowDeserializer.createChildAllocator("ArrowQueryResult"); + List vectors = ArrowPojoUtils.createVectors(this.arrowSchema, this.allocator); + this.root = new VectorSchemaRoot(vectors); + this.loader = new VectorLoader(this.root); + } else { + this.allocator = null; + this.root = null; + this.loader = null; + } + } + + static ArrowQueryResultImpl fromReadSession( + ReadSession readSession, JobId jobId, BigQueryReadClient readClient) { + Schema pojoSchema = null; + if (readSession.hasArrowSchema()) { + try { + pojoSchema = + (Schema) + ArrowDeserializer.deserializeSchema( + readSession.getArrowSchema().getSerializedSchema().toByteArray()); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow schema from ReadSession", e); + } + } + String streamName = + readSession.getStreamsCount() > 0 ? readSession.getStreams(0).getName() : null; + 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() { + synchronized (lock) { + checkNotClosed(); + if (iteratorCreated) { + throw new IllegalStateException("ArrowQueryResult can only be iterated once"); + } + iteratorCreated = true; + return new VectorBatchIterator(); + } + } + + @Override + public void close() { + synchronized (lock) { + if (closed) { + return; + } + closed = true; + Throwable firstException = null; + + if (serverStream != null) { + try { + serverStream.cancel(); + } catch (Throwable t) { + firstException = 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 != null) { + throw new RuntimeException("Failed to close Arrow resources", firstException); + } + } + } + + private void checkNotClosed() { + if (closed) { + throw new IllegalStateException("ArrowQueryResult has already been closed"); + } + } + + private final class VectorBatchIterator implements Iterator { + private boolean yieldedInitialBatch = false; + private Iterator streamIterator = null; + private boolean streamInitialized = false; + private long totalRowsYielded = 0; + + @Override + public boolean hasNext() { + synchronized (lock) { + if (closed) { + return false; + } + if (!yieldedInitialBatch + && initialRecordBatchBytes != null + && initialRecordBatchBytes.length > 0) { + return true; + } + ensureStreamInitialized(); + if (streamIterator == null) { + return false; + } + return streamIterator.hasNext(); + } + } + + @Override + public VectorSchemaRoot next() { + synchronized (lock) { + checkNotClosed(); + + // 1. Yield initial batch from REST response if present + if (!yieldedInitialBatch + && initialRecordBatchBytes != null + && initialRecordBatchBytes.length > 0) { + yieldedInitialBatch = true; + try { + loadBatchBytes(initialRecordBatchBytes); + totalRowsYielded += root.getRowCount(); + return root; + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load initial Arrow record batch", e); + } + } + yieldedInitialBatch = true; + + // 2. Stream subsequent batches from gRPC + ensureStreamInitialized(); + if (streamIterator == null || !streamIterator.hasNext()) { + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } + + while (streamIterator.hasNext()) { + ReadRowsResponse response = streamIterator.next(); + if (response.hasArrowRecordBatch()) { + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = + response.getArrowRecordBatch(); + try { + loadBatchBytes(batch.getSerializedRecordBatch().toByteArray()); + totalRowsYielded += root.getRowCount(); + return root; + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); + } + } + } + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } + } + + private void ensureStreamInitialized() { + if (streamInitialized) { + return; + } + streamInitialized = true; + if (totalRows >= 0 && totalRowsYielded >= totalRows && yieldedInitialBatch) { + return; + } + if (streamName != null && readClient != null) { + ReadRowsRequest request = + ReadRowsRequest.newBuilder() + .setReadStream(streamName) + .setOffset(totalRowsYielded) + .build(); + serverStream = readClient.readRowsCallable().call(request); + streamIterator = serverStream.iterator(); + } + } + + private void loadBatchBytes(byte[] bytes) throws IOException { + try (ByteArrayReadableSeekableByteChannel byteChannel = + new ByteArrayReadableSeekableByteChannel(bytes); + ReadChannel readChannel = new ReadChannel(byteChannel); + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch(readChannel, allocator)) { + if (deserializedBatch != null) { + loader.load(deserializedBatch); + } + } + } + } +} diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index 9fca8b042100..7ca564912c43 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java @@ -1639,6 +1639,58 @@ TableResult query(QueryJobConfiguration configuration, JobOption... options) TableResult query(QueryJobConfiguration configuration, JobId jobId, JobOption... options) throws InterruptedException, JobException; + /** + * [Beta] Runs the query associated with the request and returns an {@link + * ArrowQueryResult} yielding Apache Arrow {@code VectorSchemaRoot} batches directly for zero-copy + * vector access. + * + *

Callers must manage off-heap native memory by closing the returned {@link ArrowQueryResult} + * (e.g. via a {@code try-with-resources} block). + * + *

Prerequisite: Requires the BigQuery Storage Read API ({@code + * bigquerystorage.googleapis.com}) to be enabled on your GCP project. + * + * @param configuration the query configuration + * @param options query options + * @return an {@link ArrowQueryResult} streaming Arrow vectors + * @throws BigQueryException upon failure + * @throws InterruptedException if the current thread gets interrupted while waiting for the query + * to complete + * @throws JobException if the job completes unsuccessfully + */ + @BetaApi + default ArrowQueryResult queryArrow(QueryJobConfiguration configuration, JobOption... options) + throws InterruptedException, JobException { + throw new UnsupportedOperationException("queryArrow is not implemented"); + } + + /** + * [Beta] Runs the query associated with the request, using the given JobId, and returns an + * {@link ArrowQueryResult} yielding Apache Arrow {@code VectorSchemaRoot} batches directly for + * zero-copy vector access. + * + *

Callers must manage off-heap native memory by closing the returned {@link ArrowQueryResult} + * (e.g. via a {@code try-with-resources} block). + * + *

Prerequisite: Requires the BigQuery Storage Read API ({@code + * bigquerystorage.googleapis.com}) to be enabled on your GCP project. + * + * @param configuration the query configuration + * @param jobId the job ID to use + * @param options query options + * @return an {@link ArrowQueryResult} streaming Arrow vectors + * @throws BigQueryException upon failure + * @throws InterruptedException if the current thread gets interrupted while waiting for the query + * to complete + * @throws JobException if the job completes unsuccessfully + */ + @BetaApi + default ArrowQueryResult queryArrow( + QueryJobConfiguration configuration, JobId jobId, JobOption... options) + throws InterruptedException, JobException { + throw new UnsupportedOperationException("queryArrow is not implemented"); + } + /** * Starts the query associated with the request, using the given JobId. It returns either * TableResult for quick queries or Job object for long-running queries. 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 da4b11e676dd..3271bdc5da29 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 @@ -18,10 +18,12 @@ import static com.google.cloud.bigquery.PolicyHelper.convertFromApiPolicy; import static com.google.cloud.bigquery.PolicyHelper.convertToApiPolicy; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; import static java.net.HttpURLConnection.HTTP_NOT_FOUND; import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; +import com.google.api.gax.core.FixedCredentialsProvider; import com.google.api.gax.paging.Page; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; @@ -45,6 +47,11 @@ import com.google.cloud.bigquery.JobStatistics.SessionInfo; 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.BigQueryReadSettings; +import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest; +import com.google.cloud.bigquery.storage.v1.DataFormat; +import com.google.cloud.bigquery.storage.v1.ReadSession; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.Strings; @@ -54,14 +61,19 @@ import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; +import com.google.common.net.HostAndPort; +import io.grpc.ManagedChannelBuilder; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; import io.opentelemetry.context.Scope; import java.io.IOException; +import java.net.URI; import java.util.ArrayList; +import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.locks.ReentrantLock; import java.util.regex.Matcher; import java.util.regex.Pattern; import org.checkerframework.checker.nullness.qual.NonNull; @@ -266,6 +278,76 @@ public Page getNextPage() { } } + private final ReentrantLock readClientLock = new ReentrantLock(); + private transient BigQueryReadClient bqReadClient; + + /** + * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming + * Arrow query results, reusing credentials and channel configuration from this {@link + * BigQueryImpl}. + * + * @return the active BigQueryReadClient instance + * @throws IOException if initializing the storage read client fails + */ + BigQueryReadClient getBigQueryReadClient() throws IOException { + readClientLock.lock(); + try { + if (bqReadClient == null) { + BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder(); + configureReadSettings(settingsBuilder, getOptions()); + bqReadClient = BigQueryReadClient.create(settingsBuilder.build()); + } + return bqReadClient; + } finally { + readClientLock.unlock(); + } + } + + /** + * Configures a {@link BigQueryReadSettings.Builder} with credentials, universe domain, custom + * endpoint, and transport settings mapped from the given {@link BigQueryOptions}. + * + * @param settingsBuilder the builder to configure + * @param options the source BigQueryOptions + */ + private static void configureReadSettings( + BigQueryReadSettings.Builder settingsBuilder, BigQueryOptions options) { + if (options.getCredentials() != null) { + settingsBuilder.setCredentialsProvider( + FixedCredentialsProvider.create(options.getCredentials())); + } + if (options.getUniverseDomain() != null) { + settingsBuilder.setUniverseDomain(options.getUniverseDomain()); + } + if (options.getHost() != null) { + String host = options.getHost(); + String target = host; + if (target.contains("://")) { + target = URI.create(target).getAuthority(); + } + HostAndPort hostAndPort = HostAndPort.fromString(target); + String endpointHost = hostAndPort.getHost(); + if (endpointHost.contains("bigquery.googleapis.com")) { + endpointHost = + endpointHost.replace("bigquery.googleapis.com", "bigquerystorage.googleapis.com"); + } else if (endpointHost.contains("bigquery.private.googleapis.com")) { + endpointHost = + endpointHost.replace( + "bigquery.private.googleapis.com", "bigquerystorage.private.googleapis.com"); + } else if (endpointHost.startsWith("bigquery.")) { + endpointHost = endpointHost.replaceFirst("^bigquery\\.", "bigquerystorage."); + } + int port = hostAndPort.getPortOrDefault(443); + settingsBuilder.setEndpoint(endpointHost + ":" + port); + if (endpointHost.contains("localhost") || endpointHost.contains("127.0.0.1")) { + settingsBuilder.setTransportChannelProvider( + BigQueryReadSettings.defaultGrpcTransportProviderBuilder() + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .build()); + } + } + } + private final HttpBigQueryRpc bigQueryRpc; private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG = @@ -2178,6 +2260,11 @@ 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 = @@ -2232,6 +2319,7 @@ && getOptions().getOpenTelemetryTracer() != null) { return queryRpc(projectId, content, options); } + return create(JobInfo.of(jobId, configuration), options); } finally { if (querySpan != null) { @@ -2240,6 +2328,229 @@ && getOptions().getOpenTelemetryTracer() != null) { } } + @Override + public ArrowQueryResult queryArrow(QueryJobConfiguration configuration, JobOption... options) + throws InterruptedException, JobException { + return queryArrow(configuration, (JobId) null, options); + } + + @Override + public ArrowQueryResult queryArrow( + QueryJobConfiguration configuration, JobId jobId, JobOption... options) + throws InterruptedException, JobException { + return queryArrowWithTimeout(configuration, jobId, null, options); + } + + private ArrowQueryResult queryArrowWithTimeout( + QueryJobConfiguration configuration, JobId jobId, Long timeoutMs, JobOption... options) + throws InterruptedException, JobException { + checkNotNull(configuration, "configuration cannot be null"); + Span querySpan = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + querySpan = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.queryArrowWithTimeout") + .setAllAttributes(jobId != null ? jobId.getOtelAttributes() : Attributes.empty()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope queryScope = querySpan != null ? querySpan.makeCurrent() : null) { + QueryJobConfiguration arrowConfig = configuration; + if (arrowConfig.getQueryResultsFormat() != QueryResultsFormat.ARROW) { + arrowConfig = + configuration.toBuilder().setQueryResultsFormat(QueryResultsFormat.ARROW).build(); + } + if (arrowConfig.getJobCreationMode() == null) { + arrowConfig = + arrowConfig.toBuilder() + .setJobCreationMode(QueryJobConfiguration.JobCreationMode.JOB_CREATION_OPTIONAL) + .build(); + } + + QueryRequestInfo requestInfo = + new QueryRequestInfo(arrowConfig, getOptions().getDataFormatOptions()); + + boolean useFastPath = + requestInfo.isFastQuerySupported() + && arrowConfig.getDestinationTable() == null + && (jobId == null || jobId.getJob() == null); + + if (useFastPath) { + String projectId = + jobId != null && jobId.getProject() != null + ? jobId.getProject() + : getOptions().getProjectId(); + QueryRequest content = requestInfo.toPb(); + if (jobId != null && jobId.getLocation() != null) { + content.setLocation(jobId.getLocation()); + } else if (getOptions().getLocation() != null) { + content.setLocation(getOptions().getLocation()); + } + if (timeoutMs != null) { + content.setTimeoutMs(timeoutMs); + } + + Map optionsMap = optionMap(options); + com.google.api.services.bigquery.model.QueryResponse results; + try { + 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); + } + + if (results.getErrors() != null) { + List bigQueryErrors = + Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION); + throw new BigQueryException(bigQueryErrors); + } + + JobId actualJobId = + results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : jobId; + + Object arrowSchema = null; + if (results.getArrowSchema() != null) { + try { + arrowSchema = + ArrowDeserializer.deserializeSchema( + results.getArrowSchema().decodeSerializedSchema()); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e); + } + } + + long numRows = -1L; + if (results.getNumDmlAffectedRows() != null) { + numRows = results.getNumDmlAffectedRows(); + } else if (results.getTotalRows() != null) { + numRows = results.getTotalRows().longValue(); + } + + byte[] initialBatchBytes = null; + if (results.getArrowRecordBatch() != null + && results.getArrowRecordBatch().getSerializedRecordBatch() != null) { + initialBatchBytes = results.getArrowRecordBatch().decodeSerializedRecordBatch(); + } + + String streamName = null; + if (actualJobId != null && actualJobId.getJob() != null) { + String jobProject = + actualJobId.getProject() != null ? actualJobId.getProject() : projectId; + String jobLocation = + actualJobId.getLocation() != null + ? actualJobId.getLocation() + : (content.getLocation() != null + ? content.getLocation() + : getOptions().getLocation()); + if (jobLocation != null) { + streamName = + String.format( + "projects/%s/locations/%s/jobs/%s/streams/_default", + jobProject, jobLocation, actualJobId.getJob()); + } + } + + BigQueryReadClient client; + try { + client = getBigQueryReadClient(); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + } + + JobCreationReason jobCreationReason = + results.getJobCreationReason() != null + ? JobCreationReason.fromPb(results.getJobCreationReason()) + : null; + + return new ArrowQueryResultImpl( + arrowSchema, + actualJobId, + results.getQueryId(), + jobCreationReason, + numRows, + initialBatchBytes, + streamName, + client); + } else { + // Fallback path: jobs.insert + BigQuery Storage Read API + Job job = create(JobInfo.of(jobId, arrowConfig), options); + Job completedJob; + try { + completedJob = job.waitFor(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw e; + } + + if (completedJob.getStatus().getError() != null) { + throw new BigQueryException( + Collections.singletonList(completedJob.getStatus().getError())); + } + + TableId destinationTable = null; + if (completedJob.getConfiguration() instanceof QueryJobConfiguration) { + destinationTable = + ((QueryJobConfiguration) completedJob.getConfiguration()).getDestinationTable(); + } + if (destinationTable == null) { + destinationTable = arrowConfig.getDestinationTable(); + } + if (destinationTable == null) { + throw new BigQueryException(0, "Unable to resolve destination table for fallback query"); + } + + String destProject = + destinationTable.getProject() != null + ? destinationTable.getProject() + : (jobId != null && jobId.getProject() != null + ? jobId.getProject() + : getOptions().getProjectId()); + String parent = String.format("projects/%s", destProject); + String srcTable = + String.format( + "projects/%s/datasets/%s/tables/%s", + destProject, destinationTable.getDataset(), destinationTable.getTable()); + + BigQueryReadClient client; + try { + client = getBigQueryReadClient(); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + } + + CreateReadSessionRequest request = + CreateReadSessionRequest.newBuilder() + .setParent(parent) + .setReadSession( + ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) + .setMaxStreamCount(1) + .build(); + ReadSession readSession = client.createReadSession(request); + + return ArrowQueryResultImpl.fromReadSession(readSession, completedJob.getJobId(), client); + } + } finally { + if (querySpan != null) { + querySpan.end(); + } + } + } + @Override public QueryResponse getQueryResults(JobId jobId, QueryResultsOption... options) { Map optionsMap = optionMap(options); diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java index c224bed5cc58..14d2c65fe78a 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java @@ -46,6 +46,8 @@ final class QueryRequestInfo { private final DataFormatOptions formatOptions; private final String reservation; private final Long jobTimeoutMs; + private final QueryResultsFormat queryResultsFormat; + private final ArrowSerializationOptions arrowSerializationOptions; QueryRequestInfo( QueryJobConfiguration config, com.google.cloud.bigquery.DataFormatOptions dataFormatOptions) { @@ -63,9 +65,11 @@ final class QueryRequestInfo { this.useLegacySql = config.useLegacySql(); this.useQueryCache = config.useQueryCache(); this.jobCreationMode = config.getJobCreationMode(); - this.formatOptions = dataFormatOptions.toPb(); + this.formatOptions = dataFormatOptions != null ? dataFormatOptions.toPb() : null; this.reservation = config.getReservation(); this.jobTimeoutMs = config.getJobTimeoutMs(); + this.queryResultsFormat = config.getQueryResultsFormat(); + this.arrowSerializationOptions = config.getArrowSerializationOptions(); } /** @@ -142,6 +146,12 @@ QueryRequest toPb() { if (jobTimeoutMs != null) { request.setJobTimeoutMs(jobTimeoutMs); } + if (queryResultsFormat != null) { + request.setQueryResultsFormat(queryResultsFormat.toString()); + } + if (arrowSerializationOptions != null) { + request.setArrowSerializationOptions(arrowSerializationOptions.toPb()); + } return request; } @@ -161,7 +171,7 @@ public String toString() { .add("useQueryCache", useQueryCache) .add("useLegacySql", useLegacySql) .add("jobCreationMode", jobCreationMode) - .add("formatOptions", formatOptions.getUseInt64Timestamp()) + .add("formatOptions", formatOptions != null ? formatOptions.getUseInt64Timestamp() : null) .add("reservation", reservation) .add("jobTimeoutMs", jobTimeoutMs) .toString(); 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..e62155663478 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryResultTest.java @@ -0,0 +1,354 @@ +/* + * 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 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.BigQueryReadSettings; +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.cloud.bigquery.storage.v1.stub.EnhancedBigQueryReadStub; +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); + EnhancedBigQueryReadStub mockStub = mock(EnhancedBigQueryReadStub.class); + BigQueryReadSettings mockSettings = mock(BigQueryReadSettings.class); + try { + java.lang.reflect.Field settingsField = BigQueryReadClient.class.getDeclaredField("settings"); + settingsField.setAccessible(true); + settingsField.set(mockClient, mockSettings); + + java.lang.reflect.Field stubField = BigQueryReadClient.class.getDeclaredField("stub"); + stubField.setAccessible(true); + stubField.set(mockClient, mockStub); + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + when(mockStub.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, 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"), 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"), 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); + + @SuppressWarnings("unchecked") + ServerStream mockServerStream = mock(ServerStream.class); + 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"), 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); + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(response).iterator()); + + @SuppressWarnings("unchecked") + ServerStreamingCallable mockCallable = + mock(ServerStreamingCallable.class); + 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/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 9a398e74a67d..cb67ac54aa4a 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 @@ -2904,6 +2904,42 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException assertEquals((Long) 1000L, requestPb.getTimeoutMs()); } + @Test + void testQueryThrowsWhenArrowResultsFormat() { + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT 1") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); + bigquery = options.getService(); + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> bigquery.query(config)); + assertTrue(exception.getMessage().contains("Use queryArrow() instead")); + } + + @Test + 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 testGetQueryResults() throws IOException { JobId queryJob = JobId.of(JOB); diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 9744eebc2d81..f90a72bd9795 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -45,6 +45,7 @@ import com.google.cloud.bigquery.Acl.DatasetAclEntity; import com.google.cloud.bigquery.Acl.Expr; import com.google.cloud.bigquery.Acl.User; +import com.google.cloud.bigquery.ArrowQueryResult; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.DatasetField; import com.google.cloud.bigquery.BigQuery.DatasetListOption; @@ -118,6 +119,7 @@ import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode; import com.google.cloud.bigquery.QueryJobConfiguration.Priority; import com.google.cloud.bigquery.QueryParameterValue; +import com.google.cloud.bigquery.QueryResultsFormat; import com.google.cloud.bigquery.Range; import com.google.cloud.bigquery.RangePartitioning; import com.google.cloud.bigquery.Routine; @@ -203,6 +205,7 @@ import java.util.concurrent.TimeoutException; import java.util.logging.Level; import java.util.logging.Logger; +import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -7497,6 +7500,50 @@ void testQueryWithTimeout() throws InterruptedException { assertTrue(millis < 1_000_000 * 2); } + @Test + void testQueryResultsFormatArrow() throws InterruptedException { + RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); + BigQuery bigQuery = bigqueryHelper.getOptions().getService(); + String query = "SELECT 1 as id, 'hello' as name, TIMESTAMP('2026-08-10T12:00:00Z') as ts"; + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder(query) + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL) + .build(); + try (ArrowQueryResult result = bigQuery.queryArrow(config)) { + assertNotNull(result); + int batchCount = 0; + long totalRows = 0; + for (VectorSchemaRoot root : result) { + batchCount++; + totalRows += root.getRowCount(); + assertEquals(1, root.getRowCount()); + } + assertTrue(batchCount > 0); + assertEquals(1, totalRows); + } + } + + @Test + void testQueryResultsFormatArrowMultiPage() throws InterruptedException { + RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); + BigQuery bigQuery = bigqueryHelper.getOptions().getService(); + String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x"; + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder(query) + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL) + .build(); + try (ArrowQueryResult result = bigQuery.queryArrow(config)) { + assertNotNull(result); + long totalRows = 0; + for (VectorSchemaRoot root : result) { + totalRows += root.getRowCount(); + } + assertEquals(15000, totalRows); + } + } + @Test void testUniverseDomainWithInvalidUniverseDomain() { RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); From 64d1729e623fb50ede80643485541c92857642b4 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 10:47:48 -0400 Subject: [PATCH 02/34] refactor(bigquery): remove unused 6-argument constructor in ArrowQueryResultImpl --- .../cloud/bigquery/ArrowQueryResultImpl.java | 18 --------- .../cloud/bigquery/ArrowQueryResultTest.java | 39 +++++++++++++++++-- 2 files changed, 35 insertions(+), 22 deletions(-) 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 index 6637f1483fbf..c41e11f97088 100644 --- 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 @@ -59,24 +59,6 @@ class ArrowQueryResultImpl implements ArrowQueryResult { private boolean iteratorCreated = false; private ServerStream serverStream; - ArrowQueryResultImpl( - Object arrowSchema, - JobId jobId, - long totalRows, - byte[] initialRecordBatchBytes, - String streamName, - BigQueryReadClient readClient) { - this( - arrowSchema, - jobId, - /* queryId= */ null, - /* jobCreationReason= */ null, - totalRows, - initialRecordBatchBytes, - streamName, - readClient); - } - ArrowQueryResultImpl( Object arrowSchema, JobId jobId, 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 index e62155663478..7407bc0f61b8 100644 --- 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 @@ -148,7 +148,15 @@ void testSingleBatchIteration() throws IOException { JobId jobId = JobId.of("test-project", "job_123"); try (ArrowQueryResult result = - new ArrowQueryResultImpl(arrowSchema, jobId, 2L, batchBytes, null, null)) { + new ArrowQueryResultImpl( + arrowSchema, + jobId, + /* queryId= */ null, + /* jobCreationReason= */ null, + 2L, + batchBytes, + null, + null)) { assertEquals(arrowSchema, result.getArrowSchema()); assertEquals(jobId, result.getJobId()); assertEquals(2L, result.getTotalRows()); @@ -172,7 +180,15 @@ void testIteratorCannotBeCreatedTwice() throws IOException { byte[] batchBytes = createTestBatchBytes(ImmutableList.of(1L), ImmutableList.of("Alice")); try (ArrowQueryResult result = - new ArrowQueryResultImpl(arrowSchema, JobId.of("j1"), 1L, batchBytes, null, null)) { + 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); @@ -185,7 +201,15 @@ void testCloseIsIdempotentAndReleasesResources() throws IOException { byte[] batchBytes = createTestBatchBytes(ImmutableList.of(1L), ImmutableList.of("Alice")); ArrowQueryResult result = - new ArrowQueryResultImpl(arrowSchema, JobId.of("j1"), 1L, batchBytes, null, null); + new ArrowQueryResultImpl( + arrowSchema, + JobId.of("j1"), + /* queryId= */ null, + /* jobCreationReason= */ null, + 1L, + batchBytes, + null, + null); result.close(); result.close(); assertThrows(IllegalStateException.class, result::iterator); @@ -221,7 +245,14 @@ void testMultiBatchStreaming() throws IOException { String streamName = "projects/p/locations/l/jobs/j/streams/_default"; try (ArrowQueryResult result = new ArrowQueryResultImpl( - arrowSchema, JobId.of("j"), 5L, initialBatchBytes, streamName, mockClient)) { + arrowSchema, + JobId.of("j"), + /* queryId= */ null, + /* jobCreationReason= */ null, + 5L, + initialBatchBytes, + streamName, + mockClient)) { Iterator it = result.iterator(); // Batch 1 (initial REST response) From 7c48bfd1fd78fd71952682758b73f1a9c4b07fc0 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 11:35:29 -0400 Subject: [PATCH 03/34] refactor(bigquery): use strongly-typed Arrow Schema and VectorSchemaRoot.create in ArrowQueryResultImpl --- .../cloud/bigquery/ArrowQueryResultImpl.java | 18 +++++------------- .../google/cloud/bigquery/BigQueryImpl.java | 2 +- 2 files changed, 6 insertions(+), 14 deletions(-) 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 index c41e11f97088..bc841515753a 100644 --- 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 @@ -23,10 +23,8 @@ import com.google.cloud.bigquery.storage.v1.ReadSession; import java.io.IOException; import java.util.Iterator; -import java.util.List; import java.util.NoSuchElementException; import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.ipc.ReadChannel; @@ -60,7 +58,7 @@ class ArrowQueryResultImpl implements ArrowQueryResult { private ServerStream serverStream; ArrowQueryResultImpl( - Object arrowSchema, + Schema arrowSchema, JobId jobId, String queryId, JobCreationReason jobCreationReason, @@ -68,11 +66,7 @@ class ArrowQueryResultImpl implements ArrowQueryResult { byte[] initialRecordBatchBytes, String streamName, BigQueryReadClient readClient) { - if (arrowSchema instanceof Schema) { - this.arrowSchema = (Schema) arrowSchema; - } else { - this.arrowSchema = null; - } + this.arrowSchema = arrowSchema; this.jobId = jobId; this.queryId = queryId; this.jobCreationReason = jobCreationReason; @@ -83,8 +77,7 @@ class ArrowQueryResultImpl implements ArrowQueryResult { if (this.arrowSchema != null) { this.allocator = ArrowDeserializer.createChildAllocator("ArrowQueryResult"); - List vectors = ArrowPojoUtils.createVectors(this.arrowSchema, this.allocator); - this.root = new VectorSchemaRoot(vectors); + this.root = VectorSchemaRoot.create(this.arrowSchema, this.allocator); this.loader = new VectorLoader(this.root); } else { this.allocator = null; @@ -99,9 +92,8 @@ static ArrowQueryResultImpl fromReadSession( if (readSession.hasArrowSchema()) { try { pojoSchema = - (Schema) - ArrowDeserializer.deserializeSchema( - readSession.getArrowSchema().getSerializedSchema().toByteArray()); + ArrowDeserializer.deserializeSchema( + readSession.getArrowSchema().getSerializedSchema().toByteArray()); } catch (IOException e) { throw new BigQueryException(0, "Failed to deserialize Arrow schema from ReadSession", e); } 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 3271bdc5da29..24c148c6c398 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 @@ -2423,7 +2423,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() JobId actualJobId = results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : jobId; - Object arrowSchema = null; + org.apache.arrow.vector.types.pojo.Schema arrowSchema = null; if (results.getArrowSchema() != null) { try { arrowSchema = From 79d9cbdb473d5890309cb9e776dc26e51cfa333f Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 11:52:24 -0400 Subject: [PATCH 04/34] style(bigquery): remove extraneous empty line in BigQueryImpl --- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 1 - 1 file changed, 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 24c148c6c398..afc7254ed6d7 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 @@ -2319,7 +2319,6 @@ && getOptions().getOpenTelemetryTracer() != null) { return queryRpc(projectId, content, options); } - return create(JobInfo.of(jobId, configuration), options); } finally { if (querySpan != null) { From 359b64474042b84fff4326fb3fe938244fa978c0 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 12:27:13 -0400 Subject: [PATCH 05/34] fix(bigquery): resolve Mockito Java 8 JSpecify compatibility and optional client initialization in Arrow query --- .../google/cloud/bigquery/BigQueryImpl.java | 15 +++++++++----- .../cloud/bigquery/ArrowQueryResultTest.java | 20 ++++++++++++------- 2 files changed, 23 insertions(+), 12 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 afc7254ed6d7..22df13db2b5e 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 @@ -24,6 +24,7 @@ import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; import com.google.api.gax.core.FixedCredentialsProvider; +import com.google.api.gax.core.NoCredentialsProvider; import com.google.api.gax.paging.Page; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; @@ -315,6 +316,8 @@ private static void configureReadSettings( if (options.getCredentials() != null) { settingsBuilder.setCredentialsProvider( FixedCredentialsProvider.create(options.getCredentials())); + } else { + settingsBuilder.setCredentialsProvider(NoCredentialsProvider.create()); } if (options.getUniverseDomain() != null) { settingsBuilder.setUniverseDomain(options.getUniverseDomain()); @@ -2464,11 +2467,13 @@ public com.google.api.services.bigquery.model.QueryResponse call() } } - BigQueryReadClient client; - try { - client = getBigQueryReadClient(); - } catch (IOException e) { - throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + BigQueryReadClient client = null; + if (streamName != null) { + try { + client = getBigQueryReadClient(); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + } } JobCreationReason jobCreationReason = 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 index 7407bc0f61b8..e8f65b9948c8 100644 --- 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 @@ -25,6 +25,7 @@ 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; @@ -122,9 +123,12 @@ private byte[] createTestBatchBytes(List ids, List names) throws I private BigQueryReadClient createMockReadClient( ServerStreamingCallable mockCallable) { - BigQueryReadClient mockClient = mock(BigQueryReadClient.class); - EnhancedBigQueryReadStub mockStub = mock(EnhancedBigQueryReadStub.class); - BigQueryReadSettings mockSettings = mock(BigQueryReadSettings.class); + BigQueryReadClient mockClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + EnhancedBigQueryReadStub mockStub = + mock(EnhancedBigQueryReadStub.class, withSettings().withoutAnnotations()); + BigQueryReadSettings mockSettings = + mock(BigQueryReadSettings.class, withSettings().withoutAnnotations()); try { java.lang.reflect.Field settingsField = BigQueryReadClient.class.getDeclaredField("settings"); settingsField.setAccessible(true); @@ -225,10 +229,11 @@ void testMultiBatchStreaming() throws IOException { @SuppressWarnings("unchecked") ServerStreamingCallable mockCallable = - mock(ServerStreamingCallable.class); + mock(ServerStreamingCallable.class, withSettings().withoutAnnotations()); @SuppressWarnings("unchecked") - ServerStream mockServerStream = mock(ServerStream.class); + ServerStream mockServerStream = + mock(ServerStream.class, withSettings().withoutAnnotations()); when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream); com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = @@ -354,12 +359,13 @@ void testFromReadSessionFallback() throws IOException { .build(); @SuppressWarnings("unchecked") - ServerStream mockServerStream = mock(ServerStream.class); + ServerStream mockServerStream = + mock(ServerStream.class, withSettings().withoutAnnotations()); when(mockServerStream.iterator()).thenReturn(ImmutableList.of(response).iterator()); @SuppressWarnings("unchecked") ServerStreamingCallable mockCallable = - mock(ServerStreamingCallable.class); + mock(ServerStreamingCallable.class, withSettings().withoutAnnotations()); when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream); BigQueryReadClient mockReadClient = createMockReadClient(mockCallable); From d291158152318ce05971f17cebb6852b8f6b8e3e Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:02:08 -0400 Subject: [PATCH 06/34] fix(bigquery): rethrow JVM Error in close and translate stream/storage exceptions to BigQueryException --- .../cloud/bigquery/ArrowQueryResultImpl.java | 52 ++++++++++++------- .../google/cloud/bigquery/BigQueryImpl.java | 7 ++- 2 files changed, 38 insertions(+), 21 deletions(-) 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 index bc841515753a..4c9b11102361 100644 --- 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 @@ -188,6 +188,8 @@ public void close() { } 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); } @@ -217,11 +219,15 @@ public boolean hasNext() { && initialRecordBatchBytes.length > 0) { return true; } - ensureStreamInitialized(); - if (streamIterator == null) { - return false; + try { + ensureStreamInitialized(); + if (streamIterator == null) { + return false; + } + return streamIterator.hasNext(); + } catch (Exception e) { + throw new BigQueryException(0, "Error reading from Arrow stream", e); } - return streamIterator.hasNext(); } } @@ -246,26 +252,32 @@ public VectorSchemaRoot next() { yieldedInitialBatch = true; // 2. Stream subsequent batches from gRPC - ensureStreamInitialized(); - if (streamIterator == null || !streamIterator.hasNext()) { - throw new NoSuchElementException("No more Arrow batches available in query stream."); - } + try { + ensureStreamInitialized(); + if (streamIterator == null || !streamIterator.hasNext()) { + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } - while (streamIterator.hasNext()) { - ReadRowsResponse response = streamIterator.next(); - if (response.hasArrowRecordBatch()) { - com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = - response.getArrowRecordBatch(); - try { - loadBatchBytes(batch.getSerializedRecordBatch().toByteArray()); - totalRowsYielded += root.getRowCount(); - return root; - } catch (IOException e) { - throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); + while (streamIterator.hasNext()) { + ReadRowsResponse response = streamIterator.next(); + if (response.hasArrowRecordBatch()) { + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = + response.getArrowRecordBatch(); + try { + loadBatchBytes(batch.getSerializedRecordBatch().toByteArray()); + totalRowsYielded += root.getRowCount(); + return root; + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); + } } } + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } catch (NoSuchElementException | BigQueryException e) { + throw e; + } catch (Exception e) { + throw new BigQueryException(0, "Error reading from Arrow stream", e); } - throw new NoSuchElementException("No more Arrow batches available in query stream."); } } 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 22df13db2b5e..21a5c12ce158 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 @@ -2544,7 +2544,12 @@ public com.google.api.services.bigquery.model.QueryResponse call() ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) .setMaxStreamCount(1) .build(); - ReadSession readSession = client.createReadSession(request); + ReadSession readSession; + try { + readSession = client.createReadSession(request); + } catch (Exception e) { + throw new BigQueryException(0, "Failed to create ReadSession for fallback query", e); + } return ArrowQueryResultImpl.fromReadSession(readSession, completedJob.getJobId(), client); } From 1098b0c9984d6d8cab6881c68d310a484beff168 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:08:27 -0400 Subject: [PATCH 07/34] feat(bigquery): wait for incomplete fast-path query and stream ByteString zero-copy --- .../cloud/bigquery/ArrowQueryResultImpl.java | 45 ++++++++++++++----- .../google/cloud/bigquery/BigQueryImpl.java | 15 +++++++ 2 files changed, 49 insertions(+), 11 deletions(-) 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 index 4c9b11102361..2aff67b195d8 100644 --- 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 @@ -242,7 +242,7 @@ public VectorSchemaRoot next() { && initialRecordBatchBytes.length > 0) { yieldedInitialBatch = true; try { - loadBatchBytes(initialRecordBatchBytes); + loadBatch(initialRecordBatchBytes); totalRowsYielded += root.getRowCount(); return root; } catch (IOException e) { @@ -264,7 +264,7 @@ public VectorSchemaRoot next() { com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = response.getArrowRecordBatch(); try { - loadBatchBytes(batch.getSerializedRecordBatch().toByteArray()); + loadBatch(batch.getSerializedRecordBatch()); totalRowsYielded += root.getRowCount(); return root; } catch (IOException e) { @@ -289,18 +289,29 @@ private void ensureStreamInitialized() { if (totalRows >= 0 && totalRowsYielded >= totalRows && yieldedInitialBatch) { return; } - if (streamName != null && readClient != null) { - ReadRowsRequest request = - ReadRowsRequest.newBuilder() - .setReadStream(streamName) - .setOffset(totalRowsYielded) - .build(); - serverStream = readClient.readRowsCallable().call(request); - streamIterator = serverStream.iterator(); + 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 + + ")"); + } + return; } + ReadRowsRequest request = + ReadRowsRequest.newBuilder() + .setReadStream(streamName) + .setOffset(totalRowsYielded) + .build(); + serverStream = readClient.readRowsCallable().call(request); + streamIterator = serverStream.iterator(); } - private void loadBatchBytes(byte[] bytes) throws IOException { + private void loadBatch(byte[] bytes) throws IOException { try (ByteArrayReadableSeekableByteChannel byteChannel = new ByteArrayReadableSeekableByteChannel(bytes); ReadChannel readChannel = new ReadChannel(byteChannel); @@ -311,5 +322,17 @@ private void loadBatchBytes(byte[] bytes) throws IOException { } } } + + private void loadBatch(com.google.protobuf.ByteString byteString) throws IOException { + try (java.nio.channels.ReadableByteChannel channel = + java.nio.channels.Channels.newChannel(byteString.newInput()); + ReadChannel readChannel = new ReadChannel(channel); + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch(readChannel, allocator)) { + if (deserializedBatch != null) { + loader.load(deserializedBatch); + } + } + } } } 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 21a5c12ce158..8729df1d73f9 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 @@ -2347,6 +2347,7 @@ private ArrowQueryResult queryArrowWithTimeout( QueryJobConfiguration configuration, JobId jobId, Long timeoutMs, JobOption... options) throws InterruptedException, JobException { checkNotNull(configuration, "configuration cannot be null"); + Job.checkNotDryRun(configuration, "queryArrow"); Span querySpan = null; if (getOptions().isOpenTelemetryTracingEnabled() && getOptions().getOpenTelemetryTracer() != null) { @@ -2425,6 +2426,20 @@ public com.google.api.services.bigquery.model.QueryResponse call() JobId actualJobId = results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : jobId; + if (results.getJobComplete() != null && !results.getJobComplete()) { + if (actualJobId == null) { + throw new BigQueryException( + 0, "Query is incomplete but no job reference was returned."); + } + Job job = getJob(actualJobId); + if (job != null) { + job = job.waitFor(); + if (job.getStatus().getError() != null) { + throw new BigQueryException(Collections.singletonList(job.getStatus().getError())); + } + } + } + org.apache.arrow.vector.types.pojo.Schema arrowSchema = null; if (results.getArrowSchema() != null) { try { From 531eb5ec722a336f2a1962c4da68e863a72c2e3e Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:13:11 -0400 Subject: [PATCH 08/34] fix(bigquery): read results via Storage Read API when fast-path query completes asynchronously --- .../google/cloud/bigquery/BigQueryImpl.java | 40 +++++++++++++++++++ 1 file changed, 40 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 8729df1d73f9..d488e7964e2d 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 @@ -2437,6 +2437,46 @@ public com.google.api.services.bigquery.model.QueryResponse call() if (job.getStatus().getError() != null) { throw new BigQueryException(Collections.singletonList(job.getStatus().getError())); } + TableId destinationTable = null; + if (job.getConfiguration() instanceof QueryJobConfiguration) { + destinationTable = + ((QueryJobConfiguration) job.getConfiguration()).getDestinationTable(); + } + if (destinationTable == null) { + throw new BigQueryException( + 0, "Unable to resolve destination table for completed query"); + } + String destProject = + destinationTable.getProject() != null + ? destinationTable.getProject() + : (jobId != null && jobId.getProject() != null + ? jobId.getProject() + : getOptions().getProjectId()); + String parent = String.format("projects/%s", destProject); + String srcTable = + String.format( + "projects/%s/datasets/%s/tables/%s", + destProject, destinationTable.getDataset(), destinationTable.getTable()); + BigQueryReadClient client; + try { + client = getBigQueryReadClient(); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + } + CreateReadSessionRequest request = + CreateReadSessionRequest.newBuilder() + .setParent(parent) + .setReadSession( + ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) + .setMaxStreamCount(1) + .build(); + ReadSession readSession; + try { + readSession = client.createReadSession(request); + } catch (Exception e) { + throw new BigQueryException(0, "Failed to create ReadSession for completed query", e); + } + return ArrowQueryResultImpl.fromReadSession(readSession, job.getJobId(), client); } } From 8765958e46063367d510a66830b20b1447d37bff Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:18:26 -0400 Subject: [PATCH 09/34] fix(bigquery): address review comments on null checks, error handling, and concurrency --- .../cloud/bigquery/ArrowQueryResultImpl.java | 46 ++++++--- .../google/cloud/bigquery/BigQueryImpl.java | 99 ++++++++++--------- 2 files changed, 89 insertions(+), 56 deletions(-) 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 index 2aff67b195d8..019c52e19e81 100644 --- 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 @@ -24,6 +24,7 @@ import java.io.IOException; 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; @@ -52,7 +53,7 @@ class ArrowQueryResultImpl implements ArrowQueryResult { private final VectorSchemaRoot root; private final VectorLoader loader; - private final Object lock = new Object(); + private final ReentrantLock lock = new ReentrantLock(); private boolean closed = false; private boolean iteratorCreated = false; private ServerStream serverStream; @@ -80,6 +81,11 @@ class ArrowQueryResultImpl implements ArrowQueryResult { this.root = VectorSchemaRoot.create(this.arrowSchema, this.allocator); this.loader = new VectorLoader(this.root); } 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; @@ -98,8 +104,12 @@ static ArrowQueryResultImpl fromReadSession( throw new BigQueryException(0, "Failed to deserialize Arrow schema from ReadSession", e); } } - String streamName = - readSession.getStreamsCount() > 0 ? readSession.getStreams(0).getName() : null; + + String streamName = null; + if (readSession.getStreamsCount() > 0) { + streamName = readSession.getStreams(0).getName(); + } + return new ArrowQueryResultImpl( pojoSchema, jobId, @@ -138,19 +148,23 @@ public long getTotalRows() { @Override public Iterator iterator() { - synchronized (lock) { + 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() { - synchronized (lock) { + lock.lock(); + try { if (closed) { return; } @@ -193,6 +207,8 @@ public void close() { } else if (firstException != null) { throw new RuntimeException("Failed to close Arrow resources", firstException); } + } finally { + lock.unlock(); } } @@ -210,7 +226,8 @@ private final class VectorBatchIterator implements Iterator { @Override public boolean hasNext() { - synchronized (lock) { + lock.lock(); + try { if (closed) { return false; } @@ -228,12 +245,15 @@ public boolean hasNext() { } catch (Exception e) { throw new BigQueryException(0, "Error reading from Arrow stream", e); } + } finally { + lock.unlock(); } } @Override public VectorSchemaRoot next() { - synchronized (lock) { + lock.lock(); + try { checkNotClosed(); // 1. Yield initial batch from REST response if present @@ -278,6 +298,8 @@ public VectorSchemaRoot next() { } catch (Exception e) { throw new BigQueryException(0, "Error reading from Arrow stream", e); } + } finally { + lock.unlock(); } } @@ -317,9 +339,10 @@ private void loadBatch(byte[] bytes) throws IOException { ReadChannel readChannel = new ReadChannel(byteChannel); ArrowRecordBatch deserializedBatch = MessageSerializer.deserializeRecordBatch(readChannel, allocator)) { - if (deserializedBatch != null) { - loader.load(deserializedBatch); + if (deserializedBatch == null) { + throw new IOException("Unexpected end of stream when deserializing ArrowRecordBatch"); } + loader.load(deserializedBatch); } } @@ -329,9 +352,10 @@ private void loadBatch(com.google.protobuf.ByteString byteString) throws IOExcep ReadChannel readChannel = new ReadChannel(channel); ArrowRecordBatch deserializedBatch = MessageSerializer.deserializeRecordBatch(readChannel, allocator)) { - if (deserializedBatch != null) { - loader.load(deserializedBatch); + if (deserializedBatch == null) { + throw new IOException("Unexpected end of stream when deserializing ArrowRecordBatch"); } + loader.load(deserializedBatch); } } } 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 d488e7964e2d..8195f5520747 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 @@ -2432,52 +2432,57 @@ public com.google.api.services.bigquery.model.QueryResponse call() 0, "Query is incomplete but no job reference was returned."); } Job job = getJob(actualJobId); - if (job != null) { - job = job.waitFor(); - if (job.getStatus().getError() != null) { - throw new BigQueryException(Collections.singletonList(job.getStatus().getError())); - } - TableId destinationTable = null; - if (job.getConfiguration() instanceof QueryJobConfiguration) { - destinationTable = - ((QueryJobConfiguration) job.getConfiguration()).getDestinationTable(); - } - if (destinationTable == null) { - throw new BigQueryException( - 0, "Unable to resolve destination table for completed query"); - } - String destProject = - destinationTable.getProject() != null - ? destinationTable.getProject() - : (jobId != null && jobId.getProject() != null - ? jobId.getProject() - : getOptions().getProjectId()); - String parent = String.format("projects/%s", destProject); - String srcTable = - String.format( - "projects/%s/datasets/%s/tables/%s", - destProject, destinationTable.getDataset(), destinationTable.getTable()); - BigQueryReadClient client; - try { - client = getBigQueryReadClient(); - } catch (IOException e) { - throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); - } - CreateReadSessionRequest request = - CreateReadSessionRequest.newBuilder() - .setParent(parent) - .setReadSession( - ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) - .setMaxStreamCount(1) - .build(); - ReadSession readSession; - try { - readSession = client.createReadSession(request); - } catch (Exception e) { - throw new BigQueryException(0, "Failed to create ReadSession for completed query", e); - } - return ArrowQueryResultImpl.fromReadSession(readSession, job.getJobId(), client); + if (job == null) { + throw new BigQueryException( + 0, "Query is incomplete and job could not be retrieved: " + actualJobId); + } + job = job.waitFor(); + if (job == null) { + throw new BigQueryException(0, "Job no longer exists or could not be retrieved."); + } + if (job.getStatus().getError() != null) { + throw new BigQueryException(Collections.singletonList(job.getStatus().getError())); } + TableId destinationTable = null; + if (job.getConfiguration() instanceof QueryJobConfiguration) { + destinationTable = + ((QueryJobConfiguration) job.getConfiguration()).getDestinationTable(); + } + if (destinationTable == null) { + throw new BigQueryException( + 0, "Unable to resolve destination table for completed query"); + } + String destProject = + destinationTable.getProject() != null + ? destinationTable.getProject() + : (jobId != null && jobId.getProject() != null + ? jobId.getProject() + : getOptions().getProjectId()); + String parent = String.format("projects/%s", destProject); + String srcTable = + String.format( + "projects/%s/datasets/%s/tables/%s", + destProject, destinationTable.getDataset(), destinationTable.getTable()); + BigQueryReadClient client; + try { + client = getBigQueryReadClient(); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + } + CreateReadSessionRequest request = + CreateReadSessionRequest.newBuilder() + .setParent(parent) + .setReadSession( + ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) + .setMaxStreamCount(1) + .build(); + ReadSession readSession; + try { + readSession = client.createReadSession(request); + } catch (Exception e) { + throw new BigQueryException(0, "Failed to create ReadSession for completed query", e); + } + return ArrowQueryResultImpl.fromReadSession(readSession, job.getJobId(), client); } org.apache.arrow.vector.types.pojo.Schema arrowSchema = null; @@ -2556,6 +2561,10 @@ public com.google.api.services.bigquery.model.QueryResponse call() throw e; } + if (completedJob == null) { + throw new BigQueryException(0, "Job no longer exists or could not be retrieved."); + } + if (completedJob.getStatus().getError() != null) { throw new BigQueryException( Collections.singletonList(completedJob.getStatus().getError())); From 4aa070bd45e02e80b455a63bc239517bc0a70ed5 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:30:29 -0400 Subject: [PATCH 10/34] fix(bigquery): close BigQueryReadClient, propagate headers, and document single mutated vector root --- .../cloud/bigquery/ArrowQueryResult.java | 6 ++++ .../cloud/bigquery/ArrowQueryResultImpl.java | 28 +++++++++++++++++-- .../com/google/cloud/bigquery/BigQuery.java | 11 +++++++- .../google/cloud/bigquery/BigQueryImpl.java | 17 +++++++++-- 4 files changed, 56 insertions(+), 6 deletions(-) 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 index 13c7f2e11bb8..6d4efa893fdb 100644 --- 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 @@ -24,6 +24,12 @@ * [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. 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 index 019c52e19e81..aeb8d6ce38ea 100644 --- 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 @@ -77,9 +77,31 @@ class ArrowQueryResultImpl implements ArrowQueryResult { this.readClient = readClient; if (this.arrowSchema != null) { - this.allocator = ArrowDeserializer.createChildAllocator("ArrowQueryResult"); - this.root = VectorSchemaRoot.create(this.arrowSchema, this.allocator); - this.loader = new VectorLoader(this.root); + 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) { diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index 7ca564912c43..8ff5e6bbf8e6 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java @@ -40,7 +40,16 @@ * * @see Google Cloud BigQuery */ -public interface BigQuery extends Service { +public interface BigQuery extends Service, AutoCloseable { + + /** + * Closes any background resources and active streaming clients (such as {@code + * BigQueryReadClient}) managed by this BigQuery service instance. + * + * @throws Exception if closing underlying resources fails + */ + @Override + void close() throws Exception; /** * Fields of a BigQuery Dataset resource. 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 8195f5520747..85bccadae6c6 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 @@ -319,6 +319,9 @@ private static void configureReadSettings( } else { settingsBuilder.setCredentialsProvider(NoCredentialsProvider.create()); } + if (options.getMergedHeaderProvider(null) != null) { + settingsBuilder.setHeaderProvider(options.getMergedHeaderProvider(null)); + } if (options.getUniverseDomain() != null) { settingsBuilder.setUniverseDomain(options.getUniverseDomain()); } @@ -2394,8 +2397,6 @@ && getOptions().getOpenTelemetryTracer() != null) { if (timeoutMs != null) { content.setTimeoutMs(timeoutMs); } - - Map optionsMap = optionMap(options); com.google.api.services.bigquery.model.QueryResponse results; try { results = @@ -2917,4 +2918,16 @@ private static boolean isRetryErrorCodeHttpNotFound(BigQueryRetryHelperException } return false; } + + @Override + public void close() throws Exception { + readClientLock.lock(); + try { + if (bqReadClient != null) { + bqReadClient.close(); + } + } finally { + readClientLock.unlock(); + } + } } From 76b6294fb7871bc6d85c116756e2d69208d1f6f4 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:35:01 -0400 Subject: [PATCH 11/34] fix(bigquery): remove checked Exception from BigQuery.close signature --- .../src/main/java/com/google/cloud/bigquery/BigQuery.java | 4 +--- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index 8ff5e6bbf8e6..c37179baf0bf 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java @@ -45,11 +45,9 @@ public interface BigQuery extends Service, AutoCloseable { /** * Closes any background resources and active streaming clients (such as {@code * BigQueryReadClient}) managed by this BigQuery service instance. - * - * @throws Exception if closing underlying resources fails */ @Override - void close() throws Exception; + void close(); /** * Fields of a BigQuery Dataset resource. 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 85bccadae6c6..8267eab214f2 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 @@ -2920,7 +2920,7 @@ private static boolean isRetryErrorCodeHttpNotFound(BigQueryRetryHelperException } @Override - public void close() throws Exception { + public void close() { readClientLock.lock(); try { if (bqReadClient != null) { From e23176f088f303fdbc3b278334a1f8fa10cd52a1 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:39:48 -0400 Subject: [PATCH 12/34] fix(bigquery): narrow lock scope in VectorBatchIterator to allow concurrent cancellation --- .../cloud/bigquery/ArrowQueryResultImpl.java | 213 ++++++++++++------ 1 file changed, 146 insertions(+), 67 deletions(-) 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 index aeb8d6ce38ea..93332a81c7ba 100644 --- 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 @@ -246,113 +246,192 @@ private final class VectorBatchIterator implements Iterator { private boolean streamInitialized = false; private long totalRowsYielded = 0; + private boolean isClosed() { + lock.lock(); + try { + return closed; + } finally { + lock.unlock(); + } + } + + private boolean hasInitialBatchToYield() { + lock.lock(); + try { + return !yieldedInitialBatch + && initialRecordBatchBytes != null + && initialRecordBatchBytes.length > 0; + } finally { + lock.unlock(); + } + } + + private Iterator getStreamIterator() { + lock.lock(); + try { + return streamIterator; + } finally { + lock.unlock(); + } + } + @Override public boolean hasNext() { - lock.lock(); + if (isClosed()) { + return false; + } + if (hasInitialBatchToYield()) { + return true; + } try { - if (closed) { + ensureStreamInitialized(); + Iterator iterator = getStreamIterator(); + if (iterator == null) { return false; } - if (!yieldedInitialBatch - && initialRecordBatchBytes != null - && initialRecordBatchBytes.length > 0) { - return true; - } - try { - ensureStreamInitialized(); - if (streamIterator == null) { - return false; - } - return streamIterator.hasNext(); - } catch (Exception e) { - throw new BigQueryException(0, "Error reading from Arrow stream", e); + return iterator.hasNext(); + } catch (Exception e) { + if (isClosed()) { + return false; } - } finally { - lock.unlock(); + 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 { - checkNotClosed(); - - // 1. Yield initial batch from REST response if present if (!yieldedInitialBatch && initialRecordBatchBytes != null && initialRecordBatchBytes.length > 0) { yieldedInitialBatch = true; + initialBytes = initialRecordBatchBytes; + } else { + yieldedInitialBatch = true; + } + } finally { + lock.unlock(); + } + + if (initialBytes != null) { + try { + lock.lock(); try { - loadBatch(initialRecordBatchBytes); + checkNotClosed(); + loadBatch(initialBytes); totalRowsYielded += root.getRowCount(); return root; - } catch (IOException e) { - throw new BigQueryException(0, "Failed to load initial Arrow record batch", e); + } finally { + lock.unlock(); } + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load initial Arrow record batch", e); } - yieldedInitialBatch = true; + } - // 2. Stream subsequent batches from gRPC - try { - ensureStreamInitialized(); - if (streamIterator == null || !streamIterator.hasNext()) { - throw new NoSuchElementException("No more Arrow batches available in query stream."); - } + // 2. Stream subsequent batches from gRPC + try { + ensureStreamInitialized(); + Iterator iterator = getStreamIterator(); + if (iterator == null || !iterator.hasNext()) { + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } - while (streamIterator.hasNext()) { - ReadRowsResponse response = streamIterator.next(); - if (response.hasArrowRecordBatch()) { - com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = - response.getArrowRecordBatch(); + while (iterator.hasNext()) { + ReadRowsResponse response = iterator.next(); + if (response.hasArrowRecordBatch()) { + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = + response.getArrowRecordBatch(); + try { + lock.lock(); try { + checkNotClosed(); loadBatch(batch.getSerializedRecordBatch()); totalRowsYielded += root.getRowCount(); return root; - } catch (IOException e) { - throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); + } finally { + lock.unlock(); } + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); } } - throw new NoSuchElementException("No more Arrow batches available in query stream."); - } catch (NoSuchElementException | BigQueryException e) { - throw e; - } catch (Exception e) { - throw new BigQueryException(0, "Error reading from Arrow stream", e); } - } finally { - lock.unlock(); + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } 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); } } private void ensureStreamInitialized() { - if (streamInitialized) { - return; + lock.lock(); + try { + if (streamInitialized) { + return; + } + if (closed) { + return; + } + if (totalRows >= 0 && totalRowsYielded >= totalRows && yieldedInitialBatch) { + 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; + } + } finally { + lock.unlock(); } - streamInitialized = true; - if (totalRows >= 0 && totalRowsYielded >= totalRows && yieldedInitialBatch) { - return; + + ReadRowsRequest request; + lock.lock(); + try { + request = + ReadRowsRequest.newBuilder() + .setReadStream(streamName) + .setOffset(totalRowsYielded) + .build(); + } finally { + lock.unlock(); } - 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 - + ")"); + + ServerStream stream = readClient.readRowsCallable().call(request); + + lock.lock(); + try { + if (closed) { + stream.cancel(); + return; } - return; + serverStream = stream; + streamIterator = stream.iterator(); + streamInitialized = true; + } finally { + lock.unlock(); } - ReadRowsRequest request = - ReadRowsRequest.newBuilder() - .setReadStream(streamName) - .setOffset(totalRowsYielded) - .build(); - serverStream = readClient.readRowsCallable().call(request); - streamIterator = serverStream.iterator(); } private void loadBatch(byte[] bytes) throws IOException { From 49066581711588448a61fae7ecb45e911516e999 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:44:13 -0400 Subject: [PATCH 13/34] fix(bigquery): set bqReadClient = null upon close --- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 1 + 1 file changed, 1 insertion(+) 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 8267eab214f2..06f15bc90413 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 @@ -2925,6 +2925,7 @@ public void close() { try { if (bqReadClient != null) { bqReadClient.close(); + bqReadClient = null; } } finally { readClientLock.unlock(); From c687221a7c3c907255f3e6189893f43a9a4bc85a Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:48:48 -0400 Subject: [PATCH 14/34] fix(bigquery): prefetch and peek Arrow record batches to preserve Iterator contract --- .../cloud/bigquery/ArrowQueryResultImpl.java | 70 +++++++++++++------ 1 file changed, 47 insertions(+), 23 deletions(-) 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 index 93332a81c7ba..08aedb0e18ca 100644 --- 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 @@ -246,6 +246,8 @@ private final class VectorBatchIterator implements Iterator { private boolean streamInitialized = false; private long totalRowsYielded = 0; + private ReadRowsResponse peekedResponse = null; + private boolean isClosed() { lock.lock(); try { @@ -283,13 +285,23 @@ public boolean hasNext() { if (hasInitialBatchToYield()) { return true; } + if (peekedResponse != null) { + return true; + } try { ensureStreamInitialized(); Iterator iterator = getStreamIterator(); if (iterator == null) { return false; } - return iterator.hasNext(); + while (iterator.hasNext()) { + ReadRowsResponse response = iterator.next(); + if (response.hasArrowRecordBatch()) { + peekedResponse = response; + return true; + } + } + return false; } catch (Exception e) { if (isClosed()) { return false; @@ -338,33 +350,45 @@ public VectorSchemaRoot next() { // 2. Stream subsequent batches from gRPC try { - ensureStreamInitialized(); - Iterator iterator = getStreamIterator(); - if (iterator == null || !iterator.hasNext()) { + 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."); } - while (iterator.hasNext()) { - ReadRowsResponse response = iterator.next(); - if (response.hasArrowRecordBatch()) { - com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = - response.getArrowRecordBatch(); - try { - lock.lock(); - try { - checkNotClosed(); - loadBatch(batch.getSerializedRecordBatch()); - totalRowsYielded += root.getRowCount(); - return root; - } finally { - lock.unlock(); - } - } catch (IOException e) { - throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); - } + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = + targetResponse.getArrowRecordBatch(); + try { + lock.lock(); + try { + checkNotClosed(); + loadBatch(batch.getSerializedRecordBatch()); + totalRowsYielded += root.getRowCount(); + return root; + } finally { + lock.unlock(); } + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); } - throw new NoSuchElementException("No more Arrow batches available in query stream."); } catch (NoSuchElementException | BigQueryException e) { throw e; } catch (Exception e) { From e4c72f436167322405e8c4537e7b38df7128092c Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:52:13 -0400 Subject: [PATCH 15/34] fix(bigquery): provide default no-op implementation for BigQuery.close --- .../src/main/java/com/google/cloud/bigquery/BigQuery.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index c37179baf0bf..61cbc7b6d8dc 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java @@ -47,7 +47,7 @@ public interface BigQuery extends Service, AutoCloseable { * BigQueryReadClient}) managed by this BigQuery service instance. */ @Override - void close(); + default void close() {} /** * Fields of a BigQuery Dataset resource. From 8d49d2ad546d883b58cca7aede3357a7495e0a98 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 21:57:34 -0400 Subject: [PATCH 16/34] fix(bigquery): manage ArrowRecordBatch lifecycle to prevent premature deallocation --- .../cloud/bigquery/ArrowQueryResultImpl.java | 47 +++++++++++++++---- 1 file changed, 37 insertions(+), 10 deletions(-) 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 index 08aedb0e18ca..f5e8c6e88908 100644 --- 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 @@ -52,6 +52,7 @@ class ArrowQueryResultImpl implements ArrowQueryResult { 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; @@ -200,6 +201,18 @@ public void close() { 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(); @@ -240,6 +253,20 @@ private void checkNotClosed() { } } + void loadBatch(ArrowRecordBatch newBatch) { + lock.lock(); + try { + checkNotClosed(); + if (currentRecordBatch != null) { + currentRecordBatch.close(); + } + currentRecordBatch = newBatch; + loader.load(currentRecordBatch); + } finally { + lock.unlock(); + } + } + private final class VectorBatchIterator implements Iterator { private boolean yieldedInitialBatch = false; private Iterator streamIterator = null; @@ -334,10 +361,10 @@ public VectorSchemaRoot next() { if (initialBytes != null) { try { + loadBatch(initialBytes); lock.lock(); try { checkNotClosed(); - loadBatch(initialBytes); totalRowsYielded += root.getRowCount(); return root; } finally { @@ -377,10 +404,10 @@ public VectorSchemaRoot next() { com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = targetResponse.getArrowRecordBatch(); try { + loadBatch(batch.getSerializedRecordBatch()); lock.lock(); try { checkNotClosed(); - loadBatch(batch.getSerializedRecordBatch()); totalRowsYielded += root.getRowCount(); return root; } finally { @@ -461,26 +488,26 @@ private void ensureStreamInitialized() { private void loadBatch(byte[] bytes) throws IOException { try (ByteArrayReadableSeekableByteChannel byteChannel = new ByteArrayReadableSeekableByteChannel(bytes); - ReadChannel readChannel = new ReadChannel(byteChannel); - ArrowRecordBatch deserializedBatch = - MessageSerializer.deserializeRecordBatch(readChannel, allocator)) { + ReadChannel readChannel = new ReadChannel(byteChannel)) { + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch(readChannel, allocator); if (deserializedBatch == null) { throw new IOException("Unexpected end of stream when deserializing ArrowRecordBatch"); } - loader.load(deserializedBatch); + ArrowQueryResultImpl.this.loadBatch(deserializedBatch); } } private void loadBatch(com.google.protobuf.ByteString byteString) throws IOException { try (java.nio.channels.ReadableByteChannel channel = java.nio.channels.Channels.newChannel(byteString.newInput()); - ReadChannel readChannel = new ReadChannel(channel); - ArrowRecordBatch deserializedBatch = - MessageSerializer.deserializeRecordBatch(readChannel, allocator)) { + ReadChannel readChannel = new ReadChannel(channel)) { + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch(readChannel, allocator); if (deserializedBatch == null) { throw new IOException("Unexpected end of stream when deserializing ArrowRecordBatch"); } - loader.load(deserializedBatch); + ArrowQueryResultImpl.this.loadBatch(deserializedBatch); } } } From 090f08b53dd7b5413599e2422d109c62b4494662 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 22:02:14 -0400 Subject: [PATCH 17/34] fix(bigquery): ensure exception safe cleanup of ArrowRecordBatch and remove redundant lock --- .../cloud/bigquery/ArrowQueryResultImpl.java | 36 ++++++++++++------- 1 file changed, 23 insertions(+), 13 deletions(-) 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 index f5e8c6e88908..ca4b8a4a16ee 100644 --- 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 @@ -457,17 +457,11 @@ private void ensureStreamInitialized() { lock.unlock(); } - ReadRowsRequest request; - lock.lock(); - try { - request = - ReadRowsRequest.newBuilder() - .setReadStream(streamName) - .setOffset(totalRowsYielded) - .build(); - } finally { - lock.unlock(); - } + ReadRowsRequest request = + ReadRowsRequest.newBuilder() + .setReadStream(streamName) + .setOffset(totalRowsYielded) + .build(); ServerStream stream = readClient.readRowsCallable().call(request); @@ -494,7 +488,15 @@ private void loadBatch(byte[] bytes) throws IOException { if (deserializedBatch == null) { throw new IOException("Unexpected end of stream when deserializing ArrowRecordBatch"); } - ArrowQueryResultImpl.this.loadBatch(deserializedBatch); + boolean loaded = false; + try { + ArrowQueryResultImpl.this.loadBatch(deserializedBatch); + loaded = true; + } finally { + if (!loaded) { + deserializedBatch.close(); + } + } } } @@ -507,7 +509,15 @@ private void loadBatch(com.google.protobuf.ByteString byteString) throws IOExcep if (deserializedBatch == null) { throw new IOException("Unexpected end of stream when deserializing ArrowRecordBatch"); } - ArrowQueryResultImpl.this.loadBatch(deserializedBatch); + boolean loaded = false; + try { + ArrowQueryResultImpl.this.loadBatch(deserializedBatch); + loaded = true; + } finally { + if (!loaded) { + deserializedBatch.close(); + } + } } } } From 8702031f82b7048e2e15986e55da600781a93204 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 22:14:47 -0400 Subject: [PATCH 18/34] fix(bigquery): mock BigQueryReadClient directly without reflection --- .../cloud/bigquery/ArrowQueryResultTest.java | 19 +------------------ .../org.mockito.plugins.MockMaker | 1 + 2 files changed, 2 insertions(+), 18 deletions(-) create mode 100644 java-bigquery/google-cloud-bigquery/src/test/resources/mockito-extensions/org.mockito.plugins.MockMaker 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 index e8f65b9948c8..17a08d7651ae 100644 --- 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 @@ -31,12 +31,10 @@ 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.BigQueryReadSettings; 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.cloud.bigquery.storage.v1.stub.EnhancedBigQueryReadStub; import com.google.common.collect.ImmutableList; import com.google.protobuf.ByteString; import java.io.ByteArrayOutputStream; @@ -125,22 +123,7 @@ private BigQueryReadClient createMockReadClient( ServerStreamingCallable mockCallable) { BigQueryReadClient mockClient = mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); - EnhancedBigQueryReadStub mockStub = - mock(EnhancedBigQueryReadStub.class, withSettings().withoutAnnotations()); - BigQueryReadSettings mockSettings = - mock(BigQueryReadSettings.class, withSettings().withoutAnnotations()); - try { - java.lang.reflect.Field settingsField = BigQueryReadClient.class.getDeclaredField("settings"); - settingsField.setAccessible(true); - settingsField.set(mockClient, mockSettings); - - java.lang.reflect.Field stubField = BigQueryReadClient.class.getDeclaredField("stub"); - stubField.setAccessible(true); - stubField.set(mockClient, mockStub); - } catch (ReflectiveOperationException e) { - throw new RuntimeException(e); - } - when(mockStub.readRowsCallable()).thenReturn(mockCallable); + when(mockClient.readRowsCallable()).thenReturn(mockCallable); return mockClient; } 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 From 012c7171a7cca715b0934a8fe123e7618dbf1a8c Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 22:19:21 -0400 Subject: [PATCH 19/34] fix(bigquery): address review comments on batch loading and project id resolution --- .../com/google/cloud/bigquery/ArrowQueryResultImpl.java | 2 +- .../main/java/com/google/cloud/bigquery/BigQueryImpl.java | 8 ++++---- 2 files changed, 5 insertions(+), 5 deletions(-) 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 index ca4b8a4a16ee..7fb4048fcedb 100644 --- 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 @@ -257,11 +257,11 @@ void loadBatch(ArrowRecordBatch newBatch) { lock.lock(); try { checkNotClosed(); + loader.load(newBatch); if (currentRecordBatch != null) { currentRecordBatch.close(); } currentRecordBatch = newBatch; - loader.load(currentRecordBatch); } finally { lock.unlock(); } 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 06f15bc90413..a6cf70469f47 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 @@ -2456,8 +2456,8 @@ public com.google.api.services.bigquery.model.QueryResponse call() String destProject = destinationTable.getProject() != null ? destinationTable.getProject() - : (jobId != null && jobId.getProject() != null - ? jobId.getProject() + : (actualJobId.getProject() != null + ? actualJobId.getProject() : getOptions().getProjectId()); String parent = String.format("projects/%s", destProject); String srcTable = @@ -2586,8 +2586,8 @@ public com.google.api.services.bigquery.model.QueryResponse call() String destProject = destinationTable.getProject() != null ? destinationTable.getProject() - : (jobId != null && jobId.getProject() != null - ? jobId.getProject() + : (completedJob.getJobId().getProject() != null + ? completedJob.getJobId().getProject() : getOptions().getProjectId()); String parent = String.format("projects/%s", destProject); String srcTable = From 6c65fe4207f957c07fd684b1bffad0365793b370 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 22:23:49 -0400 Subject: [PATCH 20/34] fix(bigquery): protect Arrow batch deserialization under lock and ensure offset visibility --- .../cloud/bigquery/ArrowQueryResultImpl.java | 79 +++++++++++-------- 1 file changed, 45 insertions(+), 34 deletions(-) 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 index 7fb4048fcedb..2154d0d0dac4 100644 --- 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 @@ -427,6 +427,7 @@ public VectorSchemaRoot next() { } private void ensureStreamInitialized() { + long offset; lock.lock(); try { if (streamInitialized) { @@ -453,15 +454,13 @@ private void ensureStreamInitialized() { streamInitialized = true; return; } + offset = totalRowsYielded; } finally { lock.unlock(); } ReadRowsRequest request = - ReadRowsRequest.newBuilder() - .setReadStream(streamName) - .setOffset(totalRowsYielded) - .build(); + ReadRowsRequest.newBuilder().setReadStream(streamName).setOffset(offset).build(); ServerStream stream = readClient.readRowsCallable().call(request); @@ -480,44 +479,56 @@ private void ensureStreamInitialized() { } private void loadBatch(byte[] bytes) throws IOException { - try (ByteArrayReadableSeekableByteChannel byteChannel = - new ByteArrayReadableSeekableByteChannel(bytes); - ReadChannel readChannel = new ReadChannel(byteChannel)) { - 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(); + lock.lock(); + try { + checkNotClosed(); + try (ByteArrayReadableSeekableByteChannel byteChannel = + new ByteArrayReadableSeekableByteChannel(bytes); + ReadChannel readChannel = new ReadChannel(byteChannel)) { + 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(); + } } } + } finally { + lock.unlock(); } } private void loadBatch(com.google.protobuf.ByteString byteString) throws IOException { - try (java.nio.channels.ReadableByteChannel channel = - java.nio.channels.Channels.newChannel(byteString.newInput()); - 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(); + lock.lock(); + try { + checkNotClosed(); + try (java.nio.channels.ReadableByteChannel channel = + java.nio.channels.Channels.newChannel(byteString.newInput()); + 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(); + } } } + } finally { + lock.unlock(); } } } From 11ac9a52ccb8a45fe676ccc8e4c1e70c7f07ac7f Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 22:30:39 -0400 Subject: [PATCH 21/34] fix(bigquery): import ByteString and channel types instead of using FQCNs --- .../com/google/cloud/bigquery/ArrowQueryResultImpl.java | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) 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 index 2154d0d0dac4..cce092414284 100644 --- 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 @@ -21,7 +21,10 @@ 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; @@ -505,12 +508,11 @@ private void loadBatch(byte[] bytes) throws IOException { } } - private void loadBatch(com.google.protobuf.ByteString byteString) throws IOException { + private void loadBatch(ByteString byteString) throws IOException { lock.lock(); try { checkNotClosed(); - try (java.nio.channels.ReadableByteChannel channel = - java.nio.channels.Channels.newChannel(byteString.newInput()); + try (ReadableByteChannel channel = Channels.newChannel(byteString.newInput()); ReadChannel readChannel = new ReadChannel(channel)) { ArrowRecordBatch deserializedBatch = MessageSerializer.deserializeRecordBatch(readChannel, allocator); From fcf19c08e0c77cc02553a982539d787448c67f3b Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 22:35:07 -0400 Subject: [PATCH 22/34] fix(bigquery): ensure new batch is tracked before closing old batch --- .../com/google/cloud/bigquery/ArrowQueryResultImpl.java | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 index cce092414284..810729556d06 100644 --- 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 @@ -261,10 +261,11 @@ void loadBatch(ArrowRecordBatch newBatch) { try { checkNotClosed(); loader.load(newBatch); - if (currentRecordBatch != null) { - currentRecordBatch.close(); - } + ArrowRecordBatch oldBatch = currentRecordBatch; currentRecordBatch = newBatch; + if (oldBatch != null) { + oldBatch.close(); + } } finally { lock.unlock(); } From 0e6063544508d52352de9de58cff77b6296aa80c Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 22:40:23 -0400 Subject: [PATCH 23/34] fix(bigquery): hold lock across stream initialization in VectorBatchIterator --- .../cloud/bigquery/ArrowQueryResultImpl.java | 25 ++++--------------- 1 file changed, 5 insertions(+), 20 deletions(-) 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 index 810729556d06..e099e06a0a03 100644 --- 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 @@ -431,13 +431,9 @@ public VectorSchemaRoot next() { } private void ensureStreamInitialized() { - long offset; lock.lock(); try { - if (streamInitialized) { - return; - } - if (closed) { + if (streamInitialized || closed) { return; } if (totalRows >= 0 && totalRowsYielded >= totalRows && yieldedInitialBatch) { @@ -458,22 +454,11 @@ private void ensureStreamInitialized() { streamInitialized = true; return; } - offset = totalRowsYielded; - } finally { - lock.unlock(); - } - - ReadRowsRequest request = - ReadRowsRequest.newBuilder().setReadStream(streamName).setOffset(offset).build(); + long offset = totalRowsYielded; + ReadRowsRequest request = + ReadRowsRequest.newBuilder().setReadStream(streamName).setOffset(offset).build(); - ServerStream stream = readClient.readRowsCallable().call(request); - - lock.lock(); - try { - if (closed) { - stream.cancel(); - return; - } + ServerStream stream = readClient.readRowsCallable().call(request); serverStream = stream; streamIterator = stream.iterator(); streamInitialized = true; From 2e556e2ba4584fd91d867bb1f46a59d073892ed3 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Fri, 11 Sep 2026 22:44:57 -0400 Subject: [PATCH 24/34] fix(bigquery): rethrow BigQueryException directly in hasNext to avoid redundant wrapping --- .../java/com/google/cloud/bigquery/ArrowQueryResultImpl.java | 5 +++++ 1 file changed, 5 insertions(+) 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 index e099e06a0a03..0f4d1a0b06ed 100644 --- 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 @@ -333,6 +333,11 @@ public boolean hasNext() { } } return false; + } catch (BigQueryException e) { + if (isClosed()) { + return false; + } + throw e; } catch (Exception e) { if (isClosed()) { return false; From 8767fef066498abd6dd4925f9bbef8c2e411a9b9 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Sat, 12 Sep 2026 07:32:54 -0400 Subject: [PATCH 25/34] fix(bigquery): consolidate lock acquisition in VectorBatchIterator --- .../cloud/bigquery/ArrowQueryResultImpl.java | 20 ++++--------------- 1 file changed, 4 insertions(+), 16 deletions(-) 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 index 0f4d1a0b06ed..5584a02e6b9f 100644 --- 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 @@ -371,14 +371,7 @@ public VectorSchemaRoot next() { if (initialBytes != null) { try { loadBatch(initialBytes); - lock.lock(); - try { - checkNotClosed(); - totalRowsYielded += root.getRowCount(); - return root; - } finally { - lock.unlock(); - } + return root; } catch (IOException e) { throw new BigQueryException(0, "Failed to load initial Arrow record batch", e); } @@ -414,14 +407,7 @@ public VectorSchemaRoot next() { targetResponse.getArrowRecordBatch(); try { loadBatch(batch.getSerializedRecordBatch()); - lock.lock(); - try { - checkNotClosed(); - totalRowsYielded += root.getRowCount(); - return root; - } finally { - lock.unlock(); - } + return root; } catch (IOException e) { throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); } @@ -494,6 +480,7 @@ private void loadBatch(byte[] bytes) throws IOException { } } } + totalRowsYielded += root.getRowCount(); } finally { lock.unlock(); } @@ -520,6 +507,7 @@ private void loadBatch(ByteString byteString) throws IOException { } } } + totalRowsYielded += root.getRowCount(); } finally { lock.unlock(); } From 7ff06d1cfd40d244ff7adc772699b13e4711c668 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Sat, 12 Sep 2026 07:37:08 -0400 Subject: [PATCH 26/34] fix(bigquery): perform blocking gRPC call outside lock during stream initialization --- .../cloud/bigquery/ArrowQueryResultImpl.java | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) 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 index 5584a02e6b9f..84d88f245e3a 100644 --- 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 @@ -422,12 +422,17 @@ public VectorSchemaRoot next() { } private void ensureStreamInitialized() { + ReadRowsRequest request; lock.lock(); try { if (streamInitialized || closed) { return; } - if (totalRows >= 0 && totalRowsYielded >= totalRows && yieldedInitialBatch) { + if (totalRows >= 0 + && totalRowsYielded >= totalRows + && (yieldedInitialBatch + || initialRecordBatchBytes == null + || initialRecordBatchBytes.length == 0)) { streamInitialized = true; return; } @@ -446,10 +451,23 @@ private void ensureStreamInitialized() { return; } long offset = totalRowsYielded; - ReadRowsRequest request = - ReadRowsRequest.newBuilder().setReadStream(streamName).setOffset(offset).build(); + request = ReadRowsRequest.newBuilder().setReadStream(streamName).setOffset(offset).build(); + } finally { + lock.unlock(); + } - ServerStream stream = readClient.readRowsCallable().call(request); + 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; From 75498096dc539b42950e722be2dcad6d4cc18f54 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Sat, 12 Sep 2026 07:44:58 -0400 Subject: [PATCH 27/34] fix(bigquery): translate IOException to BigQueryException inside getBigQueryReadClient --- .../google/cloud/bigquery/BigQueryImpl.java | 30 +++++++------------ 1 file changed, 10 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 a6cf70469f47..15da52edcbda 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 @@ -288,15 +288,19 @@ public Page getNextPage() { * BigQueryImpl}. * * @return the active BigQueryReadClient instance - * @throws IOException if initializing the storage read client fails + * @throws BigQueryException if initializing the storage read client fails */ - BigQueryReadClient getBigQueryReadClient() throws IOException { + BigQueryReadClient getBigQueryReadClient() { readClientLock.lock(); try { if (bqReadClient == null) { BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder(); configureReadSettings(settingsBuilder, getOptions()); - bqReadClient = BigQueryReadClient.create(settingsBuilder.build()); + try { + bqReadClient = BigQueryReadClient.create(settingsBuilder.build()); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + } } return bqReadClient; } finally { @@ -2464,12 +2468,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() String.format( "projects/%s/datasets/%s/tables/%s", destProject, destinationTable.getDataset(), destinationTable.getTable()); - BigQueryReadClient client; - try { - client = getBigQueryReadClient(); - } catch (IOException e) { - throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); - } + BigQueryReadClient client = getBigQueryReadClient(); CreateReadSessionRequest request = CreateReadSessionRequest.newBuilder() .setParent(parent) @@ -2530,11 +2529,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() BigQueryReadClient client = null; if (streamName != null) { - try { - client = getBigQueryReadClient(); - } catch (IOException e) { - throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); - } + client = getBigQueryReadClient(); } JobCreationReason jobCreationReason = @@ -2595,12 +2590,7 @@ public com.google.api.services.bigquery.model.QueryResponse call() "projects/%s/datasets/%s/tables/%s", destProject, destinationTable.getDataset(), destinationTable.getTable()); - BigQueryReadClient client; - try { - client = getBigQueryReadClient(); - } catch (IOException e) { - throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); - } + BigQueryReadClient client = getBigQueryReadClient(); CreateReadSessionRequest request = CreateReadSessionRequest.newBuilder() From 7b170dfb53e74d25f251d4cc0378a8a8f525df06 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Sat, 12 Sep 2026 07:48:43 -0400 Subject: [PATCH 28/34] fix(bigquery): simplify anonymous Callable to lambda in queryArrowWithTimeout --- .../main/java/com/google/cloud/bigquery/BigQueryImpl.java | 8 +------- 1 file changed, 1 insertion(+), 7 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 15da52edcbda..394b29ba8c43 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 @@ -2405,13 +2405,7 @@ && getOptions().getOpenTelemetryTracer() != null) { try { 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(), From 6d600efc25cb9883b46674a0919074e218c0d9fd Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 10:38:02 -0400 Subject: [PATCH 29/34] fix(bigquery): preserve existing Service contract by removing AutoCloseable from BigQuery --- .../java/com/google/cloud/bigquery/BigQuery.java | 9 +-------- .../com/google/cloud/bigquery/BigQueryImpl.java | 13 ------------- 2 files changed, 1 insertion(+), 21 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index 61cbc7b6d8dc..7ca564912c43 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java @@ -40,14 +40,7 @@ * * @see Google Cloud BigQuery */ -public interface BigQuery extends Service, AutoCloseable { - - /** - * Closes any background resources and active streaming clients (such as {@code - * BigQueryReadClient}) managed by this BigQuery service instance. - */ - @Override - default void close() {} +public interface BigQuery extends Service { /** * Fields of a BigQuery Dataset resource. 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 394b29ba8c43..24e9b9213e65 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 @@ -2902,17 +2902,4 @@ private static boolean isRetryErrorCodeHttpNotFound(BigQueryRetryHelperException } return false; } - - @Override - public void close() { - readClientLock.lock(); - try { - if (bqReadClient != null) { - bqReadClient.close(); - bqReadClient = null; - } - } finally { - readClientLock.unlock(); - } - } } From c8575961ddf093111e0b5ed4fa76a43e48681b12 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 12:42:55 -0400 Subject: [PATCH 30/34] refactor(bigquery): isolate Arrow result model and streaming iterator for PR 1 --- .../com/google/cloud/bigquery/BigQuery.java | 52 --- .../google/cloud/bigquery/BigQueryImpl.java | 369 ------------------ .../cloud/bigquery/QueryRequestInfo.java | 14 +- .../cloud/bigquery/BigQueryImplTest.java | 36 -- .../cloud/bigquery/it/ITBigQueryTest.java | 47 --- 5 files changed, 2 insertions(+), 516 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index 7ca564912c43..9fca8b042100 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java @@ -1639,58 +1639,6 @@ TableResult query(QueryJobConfiguration configuration, JobOption... options) TableResult query(QueryJobConfiguration configuration, JobId jobId, JobOption... options) throws InterruptedException, JobException; - /** - * [Beta] Runs the query associated with the request and returns an {@link - * ArrowQueryResult} yielding Apache Arrow {@code VectorSchemaRoot} batches directly for zero-copy - * vector access. - * - *

Callers must manage off-heap native memory by closing the returned {@link ArrowQueryResult} - * (e.g. via a {@code try-with-resources} block). - * - *

Prerequisite: Requires the BigQuery Storage Read API ({@code - * bigquerystorage.googleapis.com}) to be enabled on your GCP project. - * - * @param configuration the query configuration - * @param options query options - * @return an {@link ArrowQueryResult} streaming Arrow vectors - * @throws BigQueryException upon failure - * @throws InterruptedException if the current thread gets interrupted while waiting for the query - * to complete - * @throws JobException if the job completes unsuccessfully - */ - @BetaApi - default ArrowQueryResult queryArrow(QueryJobConfiguration configuration, JobOption... options) - throws InterruptedException, JobException { - throw new UnsupportedOperationException("queryArrow is not implemented"); - } - - /** - * [Beta] Runs the query associated with the request, using the given JobId, and returns an - * {@link ArrowQueryResult} yielding Apache Arrow {@code VectorSchemaRoot} batches directly for - * zero-copy vector access. - * - *

Callers must manage off-heap native memory by closing the returned {@link ArrowQueryResult} - * (e.g. via a {@code try-with-resources} block). - * - *

Prerequisite: Requires the BigQuery Storage Read API ({@code - * bigquerystorage.googleapis.com}) to be enabled on your GCP project. - * - * @param configuration the query configuration - * @param jobId the job ID to use - * @param options query options - * @return an {@link ArrowQueryResult} streaming Arrow vectors - * @throws BigQueryException upon failure - * @throws InterruptedException if the current thread gets interrupted while waiting for the query - * to complete - * @throws JobException if the job completes unsuccessfully - */ - @BetaApi - default ArrowQueryResult queryArrow( - QueryJobConfiguration configuration, JobId jobId, JobOption... options) - throws InterruptedException, JobException { - throw new UnsupportedOperationException("queryArrow is not implemented"); - } - /** * Starts the query associated with the request, using the given JobId. It returns either * TableResult for quick queries or Job object for long-running queries. 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 24e9b9213e65..da4b11e676dd 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 @@ -18,13 +18,10 @@ import static com.google.cloud.bigquery.PolicyHelper.convertFromApiPolicy; import static com.google.cloud.bigquery.PolicyHelper.convertToApiPolicy; import static com.google.common.base.Preconditions.checkArgument; -import static com.google.common.base.Preconditions.checkNotNull; import static java.net.HttpURLConnection.HTTP_NOT_FOUND; import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; -import com.google.api.gax.core.FixedCredentialsProvider; -import com.google.api.gax.core.NoCredentialsProvider; import com.google.api.gax.paging.Page; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; @@ -48,11 +45,6 @@ import com.google.cloud.bigquery.JobStatistics.SessionInfo; 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.BigQueryReadSettings; -import com.google.cloud.bigquery.storage.v1.CreateReadSessionRequest; -import com.google.cloud.bigquery.storage.v1.DataFormat; -import com.google.cloud.bigquery.storage.v1.ReadSession; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.Strings; @@ -62,19 +54,14 @@ import com.google.common.collect.Iterables; import com.google.common.collect.Lists; import com.google.common.collect.Maps; -import com.google.common.net.HostAndPort; -import io.grpc.ManagedChannelBuilder; import io.opentelemetry.api.common.Attributes; import io.opentelemetry.api.trace.Span; import io.opentelemetry.context.Scope; import java.io.IOException; -import java.net.URI; import java.util.ArrayList; -import java.util.Collections; import java.util.List; import java.util.Map; import java.util.concurrent.Callable; -import java.util.concurrent.locks.ReentrantLock; import java.util.regex.Matcher; import java.util.regex.Pattern; import org.checkerframework.checker.nullness.qual.NonNull; @@ -279,85 +266,6 @@ public Page getNextPage() { } } - private final ReentrantLock readClientLock = new ReentrantLock(); - private transient BigQueryReadClient bqReadClient; - - /** - * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming - * Arrow query results, reusing credentials and channel configuration from this {@link - * BigQueryImpl}. - * - * @return the active BigQueryReadClient instance - * @throws BigQueryException if initializing the storage read client fails - */ - BigQueryReadClient getBigQueryReadClient() { - readClientLock.lock(); - try { - if (bqReadClient == null) { - BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder(); - configureReadSettings(settingsBuilder, getOptions()); - try { - bqReadClient = BigQueryReadClient.create(settingsBuilder.build()); - } catch (IOException e) { - throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); - } - } - return bqReadClient; - } finally { - readClientLock.unlock(); - } - } - - /** - * Configures a {@link BigQueryReadSettings.Builder} with credentials, universe domain, custom - * endpoint, and transport settings mapped from the given {@link BigQueryOptions}. - * - * @param settingsBuilder the builder to configure - * @param options the source BigQueryOptions - */ - private static void configureReadSettings( - BigQueryReadSettings.Builder settingsBuilder, BigQueryOptions options) { - if (options.getCredentials() != null) { - settingsBuilder.setCredentialsProvider( - FixedCredentialsProvider.create(options.getCredentials())); - } else { - settingsBuilder.setCredentialsProvider(NoCredentialsProvider.create()); - } - if (options.getMergedHeaderProvider(null) != null) { - settingsBuilder.setHeaderProvider(options.getMergedHeaderProvider(null)); - } - if (options.getUniverseDomain() != null) { - settingsBuilder.setUniverseDomain(options.getUniverseDomain()); - } - if (options.getHost() != null) { - String host = options.getHost(); - String target = host; - if (target.contains("://")) { - target = URI.create(target).getAuthority(); - } - HostAndPort hostAndPort = HostAndPort.fromString(target); - String endpointHost = hostAndPort.getHost(); - if (endpointHost.contains("bigquery.googleapis.com")) { - endpointHost = - endpointHost.replace("bigquery.googleapis.com", "bigquerystorage.googleapis.com"); - } else if (endpointHost.contains("bigquery.private.googleapis.com")) { - endpointHost = - endpointHost.replace( - "bigquery.private.googleapis.com", "bigquerystorage.private.googleapis.com"); - } else if (endpointHost.startsWith("bigquery.")) { - endpointHost = endpointHost.replaceFirst("^bigquery\\.", "bigquerystorage."); - } - int port = hostAndPort.getPortOrDefault(443); - settingsBuilder.setEndpoint(endpointHost + ":" + port); - if (endpointHost.contains("localhost") || endpointHost.contains("127.0.0.1")) { - settingsBuilder.setTransportChannelProvider( - BigQueryReadSettings.defaultGrpcTransportProviderBuilder() - .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) - .build()); - } - } - } - private final HttpBigQueryRpc bigQueryRpc; private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG = @@ -2270,11 +2178,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 = @@ -2337,278 +2240,6 @@ && getOptions().getOpenTelemetryTracer() != null) { } } - @Override - public ArrowQueryResult queryArrow(QueryJobConfiguration configuration, JobOption... options) - throws InterruptedException, JobException { - return queryArrow(configuration, (JobId) null, options); - } - - @Override - public ArrowQueryResult queryArrow( - QueryJobConfiguration configuration, JobId jobId, JobOption... options) - throws InterruptedException, JobException { - return queryArrowWithTimeout(configuration, jobId, null, options); - } - - private ArrowQueryResult queryArrowWithTimeout( - QueryJobConfiguration configuration, JobId jobId, Long timeoutMs, JobOption... options) - throws InterruptedException, JobException { - checkNotNull(configuration, "configuration cannot be null"); - Job.checkNotDryRun(configuration, "queryArrow"); - Span querySpan = null; - if (getOptions().isOpenTelemetryTracingEnabled() - && getOptions().getOpenTelemetryTracer() != null) { - querySpan = - getOptions() - .getOpenTelemetryTracer() - .spanBuilder("com.google.cloud.bigquery.BigQuery.queryArrowWithTimeout") - .setAllAttributes(jobId != null ? jobId.getOtelAttributes() : Attributes.empty()) - .setAllAttributes(otelAttributesFromOptions(options)) - .startSpan(); - } - try (Scope queryScope = querySpan != null ? querySpan.makeCurrent() : null) { - QueryJobConfiguration arrowConfig = configuration; - if (arrowConfig.getQueryResultsFormat() != QueryResultsFormat.ARROW) { - arrowConfig = - configuration.toBuilder().setQueryResultsFormat(QueryResultsFormat.ARROW).build(); - } - if (arrowConfig.getJobCreationMode() == null) { - arrowConfig = - arrowConfig.toBuilder() - .setJobCreationMode(QueryJobConfiguration.JobCreationMode.JOB_CREATION_OPTIONAL) - .build(); - } - - QueryRequestInfo requestInfo = - new QueryRequestInfo(arrowConfig, getOptions().getDataFormatOptions()); - - boolean useFastPath = - requestInfo.isFastQuerySupported() - && arrowConfig.getDestinationTable() == null - && (jobId == null || jobId.getJob() == null); - - if (useFastPath) { - String projectId = - jobId != null && jobId.getProject() != null - ? jobId.getProject() - : getOptions().getProjectId(); - QueryRequest content = requestInfo.toPb(); - if (jobId != null && jobId.getLocation() != null) { - content.setLocation(jobId.getLocation()); - } else if (getOptions().getLocation() != null) { - content.setLocation(getOptions().getLocation()); - } - if (timeoutMs != null) { - content.setTimeoutMs(timeoutMs); - } - com.google.api.services.bigquery.model.QueryResponse results; - try { - results = - BigQueryRetryHelper.runWithRetries( - () -> 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); - } - - if (results.getErrors() != null) { - List bigQueryErrors = - Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION); - throw new BigQueryException(bigQueryErrors); - } - - JobId actualJobId = - results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : jobId; - - if (results.getJobComplete() != null && !results.getJobComplete()) { - if (actualJobId == null) { - throw new BigQueryException( - 0, "Query is incomplete but no job reference was returned."); - } - Job job = getJob(actualJobId); - if (job == null) { - throw new BigQueryException( - 0, "Query is incomplete and job could not be retrieved: " + actualJobId); - } - job = job.waitFor(); - if (job == null) { - throw new BigQueryException(0, "Job no longer exists or could not be retrieved."); - } - if (job.getStatus().getError() != null) { - throw new BigQueryException(Collections.singletonList(job.getStatus().getError())); - } - TableId destinationTable = null; - if (job.getConfiguration() instanceof QueryJobConfiguration) { - destinationTable = - ((QueryJobConfiguration) job.getConfiguration()).getDestinationTable(); - } - if (destinationTable == null) { - throw new BigQueryException( - 0, "Unable to resolve destination table for completed query"); - } - String destProject = - destinationTable.getProject() != null - ? destinationTable.getProject() - : (actualJobId.getProject() != null - ? actualJobId.getProject() - : getOptions().getProjectId()); - String parent = String.format("projects/%s", destProject); - String srcTable = - String.format( - "projects/%s/datasets/%s/tables/%s", - destProject, destinationTable.getDataset(), destinationTable.getTable()); - BigQueryReadClient client = getBigQueryReadClient(); - CreateReadSessionRequest request = - CreateReadSessionRequest.newBuilder() - .setParent(parent) - .setReadSession( - ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) - .setMaxStreamCount(1) - .build(); - ReadSession readSession; - try { - readSession = client.createReadSession(request); - } catch (Exception e) { - throw new BigQueryException(0, "Failed to create ReadSession for completed query", e); - } - return ArrowQueryResultImpl.fromReadSession(readSession, job.getJobId(), client); - } - - org.apache.arrow.vector.types.pojo.Schema arrowSchema = null; - if (results.getArrowSchema() != null) { - try { - arrowSchema = - ArrowDeserializer.deserializeSchema( - results.getArrowSchema().decodeSerializedSchema()); - } catch (IOException e) { - throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e); - } - } - - long numRows = -1L; - if (results.getNumDmlAffectedRows() != null) { - numRows = results.getNumDmlAffectedRows(); - } else if (results.getTotalRows() != null) { - numRows = results.getTotalRows().longValue(); - } - - byte[] initialBatchBytes = null; - if (results.getArrowRecordBatch() != null - && results.getArrowRecordBatch().getSerializedRecordBatch() != null) { - initialBatchBytes = results.getArrowRecordBatch().decodeSerializedRecordBatch(); - } - - String streamName = null; - if (actualJobId != null && actualJobId.getJob() != null) { - String jobProject = - actualJobId.getProject() != null ? actualJobId.getProject() : projectId; - String jobLocation = - actualJobId.getLocation() != null - ? actualJobId.getLocation() - : (content.getLocation() != null - ? content.getLocation() - : getOptions().getLocation()); - if (jobLocation != null) { - streamName = - String.format( - "projects/%s/locations/%s/jobs/%s/streams/_default", - jobProject, jobLocation, actualJobId.getJob()); - } - } - - BigQueryReadClient client = null; - if (streamName != null) { - client = getBigQueryReadClient(); - } - - JobCreationReason jobCreationReason = - results.getJobCreationReason() != null - ? JobCreationReason.fromPb(results.getJobCreationReason()) - : null; - - return new ArrowQueryResultImpl( - arrowSchema, - actualJobId, - results.getQueryId(), - jobCreationReason, - numRows, - initialBatchBytes, - streamName, - client); - } else { - // Fallback path: jobs.insert + BigQuery Storage Read API - Job job = create(JobInfo.of(jobId, arrowConfig), options); - Job completedJob; - try { - completedJob = job.waitFor(); - } catch (InterruptedException e) { - Thread.currentThread().interrupt(); - throw e; - } - - if (completedJob == null) { - throw new BigQueryException(0, "Job no longer exists or could not be retrieved."); - } - - if (completedJob.getStatus().getError() != null) { - throw new BigQueryException( - Collections.singletonList(completedJob.getStatus().getError())); - } - - TableId destinationTable = null; - if (completedJob.getConfiguration() instanceof QueryJobConfiguration) { - destinationTable = - ((QueryJobConfiguration) completedJob.getConfiguration()).getDestinationTable(); - } - if (destinationTable == null) { - destinationTable = arrowConfig.getDestinationTable(); - } - if (destinationTable == null) { - throw new BigQueryException(0, "Unable to resolve destination table for fallback query"); - } - - String destProject = - destinationTable.getProject() != null - ? destinationTable.getProject() - : (completedJob.getJobId().getProject() != null - ? completedJob.getJobId().getProject() - : getOptions().getProjectId()); - String parent = String.format("projects/%s", destProject); - String srcTable = - String.format( - "projects/%s/datasets/%s/tables/%s", - destProject, destinationTable.getDataset(), destinationTable.getTable()); - - BigQueryReadClient client = getBigQueryReadClient(); - - CreateReadSessionRequest request = - CreateReadSessionRequest.newBuilder() - .setParent(parent) - .setReadSession( - ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) - .setMaxStreamCount(1) - .build(); - ReadSession readSession; - try { - readSession = client.createReadSession(request); - } catch (Exception e) { - throw new BigQueryException(0, "Failed to create ReadSession for fallback query", e); - } - - return ArrowQueryResultImpl.fromReadSession(readSession, completedJob.getJobId(), client); - } - } finally { - if (querySpan != null) { - querySpan.end(); - } - } - } - @Override public QueryResponse getQueryResults(JobId jobId, QueryResultsOption... options) { Map optionsMap = optionMap(options); diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java index 14d2c65fe78a..c224bed5cc58 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java @@ -46,8 +46,6 @@ final class QueryRequestInfo { private final DataFormatOptions formatOptions; private final String reservation; private final Long jobTimeoutMs; - private final QueryResultsFormat queryResultsFormat; - private final ArrowSerializationOptions arrowSerializationOptions; QueryRequestInfo( QueryJobConfiguration config, com.google.cloud.bigquery.DataFormatOptions dataFormatOptions) { @@ -65,11 +63,9 @@ final class QueryRequestInfo { this.useLegacySql = config.useLegacySql(); this.useQueryCache = config.useQueryCache(); this.jobCreationMode = config.getJobCreationMode(); - this.formatOptions = dataFormatOptions != null ? dataFormatOptions.toPb() : null; + this.formatOptions = dataFormatOptions.toPb(); this.reservation = config.getReservation(); this.jobTimeoutMs = config.getJobTimeoutMs(); - this.queryResultsFormat = config.getQueryResultsFormat(); - this.arrowSerializationOptions = config.getArrowSerializationOptions(); } /** @@ -146,12 +142,6 @@ QueryRequest toPb() { if (jobTimeoutMs != null) { request.setJobTimeoutMs(jobTimeoutMs); } - if (queryResultsFormat != null) { - request.setQueryResultsFormat(queryResultsFormat.toString()); - } - if (arrowSerializationOptions != null) { - request.setArrowSerializationOptions(arrowSerializationOptions.toPb()); - } return request; } @@ -171,7 +161,7 @@ public String toString() { .add("useQueryCache", useQueryCache) .add("useLegacySql", useLegacySql) .add("jobCreationMode", jobCreationMode) - .add("formatOptions", formatOptions != null ? formatOptions.getUseInt64Timestamp() : null) + .add("formatOptions", formatOptions.getUseInt64Timestamp()) .add("reservation", reservation) .add("jobTimeoutMs", jobTimeoutMs) .toString(); 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..9a398e74a67d 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 @@ -2904,42 +2904,6 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException assertEquals((Long) 1000L, requestPb.getTimeoutMs()); } - @Test - void testQueryThrowsWhenArrowResultsFormat() { - QueryJobConfiguration config = - QueryJobConfiguration.newBuilder("SELECT 1") - .setQueryResultsFormat(QueryResultsFormat.ARROW) - .build(); - bigquery = options.getService(); - IllegalArgumentException exception = - assertThrows(IllegalArgumentException.class, () -> bigquery.query(config)); - assertTrue(exception.getMessage().contains("Use queryArrow() instead")); - } - - @Test - 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 testGetQueryResults() throws IOException { JobId queryJob = JobId.of(JOB); diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index f90a72bd9795..9744eebc2d81 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java @@ -45,7 +45,6 @@ import com.google.cloud.bigquery.Acl.DatasetAclEntity; import com.google.cloud.bigquery.Acl.Expr; import com.google.cloud.bigquery.Acl.User; -import com.google.cloud.bigquery.ArrowQueryResult; import com.google.cloud.bigquery.BigQuery; import com.google.cloud.bigquery.BigQuery.DatasetField; import com.google.cloud.bigquery.BigQuery.DatasetListOption; @@ -119,7 +118,6 @@ import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode; import com.google.cloud.bigquery.QueryJobConfiguration.Priority; import com.google.cloud.bigquery.QueryParameterValue; -import com.google.cloud.bigquery.QueryResultsFormat; import com.google.cloud.bigquery.Range; import com.google.cloud.bigquery.RangePartitioning; import com.google.cloud.bigquery.Routine; @@ -205,7 +203,6 @@ import java.util.concurrent.TimeoutException; import java.util.logging.Level; import java.util.logging.Logger; -import org.apache.arrow.vector.VectorSchemaRoot; import org.junit.jupiter.api.AfterAll; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -7500,50 +7497,6 @@ void testQueryWithTimeout() throws InterruptedException { assertTrue(millis < 1_000_000 * 2); } - @Test - void testQueryResultsFormatArrow() throws InterruptedException { - RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); - BigQuery bigQuery = bigqueryHelper.getOptions().getService(); - String query = "SELECT 1 as id, 'hello' as name, TIMESTAMP('2026-08-10T12:00:00Z') as ts"; - QueryJobConfiguration config = - QueryJobConfiguration.newBuilder(query) - .setQueryResultsFormat(QueryResultsFormat.ARROW) - .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL) - .build(); - try (ArrowQueryResult result = bigQuery.queryArrow(config)) { - assertNotNull(result); - int batchCount = 0; - long totalRows = 0; - for (VectorSchemaRoot root : result) { - batchCount++; - totalRows += root.getRowCount(); - assertEquals(1, root.getRowCount()); - } - assertTrue(batchCount > 0); - assertEquals(1, totalRows); - } - } - - @Test - void testQueryResultsFormatArrowMultiPage() throws InterruptedException { - RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); - BigQuery bigQuery = bigqueryHelper.getOptions().getService(); - String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x"; - QueryJobConfiguration config = - QueryJobConfiguration.newBuilder(query) - .setQueryResultsFormat(QueryResultsFormat.ARROW) - .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL) - .build(); - try (ArrowQueryResult result = bigQuery.queryArrow(config)) { - assertNotNull(result); - long totalRows = 0; - for (VectorSchemaRoot root : result) { - totalRows += root.getRowCount(); - } - assertEquals(15000, totalRows); - } - } - @Test void testUniverseDomainWithInvalidUniverseDomain() { RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create(); From dc0289907ffc70c3f13421a867ebe00fc98a1c8d Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 13:39:46 -0400 Subject: [PATCH 31/34] docs(bigquery): add Javadoc comments to non-override methods in ArrowQueryResultImpl --- .../cloud/bigquery/ArrowQueryResultImpl.java | 72 +++++++++++++++++++ 1 file changed, 72 insertions(+) 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 index 84d88f245e3a..b2e79aa9b4ea 100644 --- 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 @@ -62,6 +62,20 @@ class ArrowQueryResultImpl implements ArrowQueryResult { 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, @@ -118,6 +132,16 @@ class ArrowQueryResultImpl implements ArrowQueryResult { } } + /** + * 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; @@ -250,12 +274,23 @@ public void close() { } } + /** + * 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 { @@ -279,6 +314,11 @@ private final class VectorBatchIterator implements Iterator { 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 { @@ -288,6 +328,12 @@ private boolean isClosed() { } } + /** + * 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 { @@ -299,6 +345,12 @@ private boolean hasInitialBatchToYield() { } } + /** + * 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 { @@ -421,6 +473,13 @@ public VectorSchemaRoot next() { } } + /** + * 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(); @@ -476,6 +535,12 @@ private void ensureStreamInitialized() { } } + /** + * 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 { lock.lock(); try { @@ -504,6 +569,13 @@ private void loadBatch(byte[] bytes) throws IOException { } } + /** + * 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 { lock.lock(); try { From 3065daf17743ac038f1713ddb0fa2db14decd0a2 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 13:51:11 -0400 Subject: [PATCH 32/34] refactor(bigquery): deduplicate loadBatch in ArrowQueryResultImpl --- .../cloud/bigquery/ArrowQueryResultImpl.java | 27 +------------------ 1 file changed, 1 insertion(+), 26 deletions(-) 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 index b2e79aa9b4ea..1a2d2c897ad7 100644 --- 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 @@ -35,7 +35,6 @@ 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 @@ -542,31 +541,7 @@ private void ensureStreamInitialized() { * @throws IOException if deserialization fails */ private void loadBatch(byte[] bytes) throws IOException { - lock.lock(); - try { - checkNotClosed(); - try (ByteArrayReadableSeekableByteChannel byteChannel = - new ByteArrayReadableSeekableByteChannel(bytes); - ReadChannel readChannel = new ReadChannel(byteChannel)) { - 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(); - } + loadBatch(ByteString.copyFrom(bytes)); } /** From 39425776f982a39abda0a478a51c9d4676f48a23 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 14:01:16 -0400 Subject: [PATCH 33/34] perf(bigquery): ensure zero-copy batch loading via ReadableByteChannel --- .../cloud/bigquery/ArrowQueryResultImpl.java | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) 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 index 1a2d2c897ad7..9fd92fb4d1f4 100644 --- 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 @@ -35,6 +35,7 @@ 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 @@ -541,7 +542,7 @@ private void ensureStreamInitialized() { * @throws IOException if deserialization fails */ private void loadBatch(byte[] bytes) throws IOException { - loadBatch(ByteString.copyFrom(bytes)); + loadBatch(new ByteArrayReadableSeekableByteChannel(bytes)); } /** @@ -552,11 +553,21 @@ private void loadBatch(byte[] bytes) throws IOException { * @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 (ReadableByteChannel channel = Channels.newChannel(byteString.newInput()); - ReadChannel readChannel = new ReadChannel(channel)) { + try (ReadChannel readChannel = new ReadChannel(channel)) { ArrowRecordBatch deserializedBatch = MessageSerializer.deserializeRecordBatch(readChannel, allocator); if (deserializedBatch == null) { From 1b283dbcd1053df7f28bb2ef674e433c6e4a0d6d Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 14:55:36 -0400 Subject: [PATCH 34/34] fix(bigquery): remove unused gax-grpc and grpc-api dependencies from pom.xml --- java-bigquery/google-cloud-bigquery/pom.xml | 9 --------- 1 file changed, 9 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/pom.xml b/java-bigquery/google-cloud-bigquery/pom.xml index 3f186b68624e..765a2e650ee2 100644 --- a/java-bigquery/google-cloud-bigquery/pom.xml +++ b/java-bigquery/google-cloud-bigquery/pom.xml @@ -122,15 +122,6 @@ arrow-memory-netty - - com.google.api - gax-grpc - - - io.grpc - grpc-api - - com.google.errorprone error_prone_annotations