From ddda891537b15de7afd55d91fe8df371b28c628b Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 15:36:47 -0400 Subject: [PATCH 01/12] feat(bigquery): add zero-copy queryArrow API for Arrow VectorSchemaRoot streaming --- java-bigquery/google-cloud-bigquery/pom.xml | 9 + .../com/google/cloud/bigquery/BigQuery.java | 52 +++ .../google/cloud/bigquery/BigQueryImpl.java | 371 ++++++++++++++++++ .../cloud/bigquery/QueryRequestInfo.java | 14 +- .../cloud/bigquery/BigQueryImplTest.java | 36 ++ 5 files changed, 480 insertions(+), 2 deletions(-) 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/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..fd999540f4c1 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,13 @@ 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; @@ -45,6 +48,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 +62,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 +279,85 @@ 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 = @@ -2178,6 +2270,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 = @@ -2240,6 +2337,280 @@ && 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); + } + + /** + * Executes a query in Arrow format with an optional execution timeout. + * + * @param configuration query job configuration + * @param jobId job identifier, or {@code null} + * @param timeoutMs query timeout in milliseconds, or {@code null} + * @param options query job options + * @return an {@link ArrowQueryResult} for streaming results + * @throws InterruptedException if interrupted while awaiting results + * @throws JobException if the query job fails + */ + 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() && (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"); + } + return createArrowQueryResultFromTable( + destinationTable, job.getJobId(), "completed query"); + } + + 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"); + } + + return createArrowQueryResultFromTable( + destinationTable, completedJob.getJobId(), "fallback query"); + } + } finally { + if (querySpan != null) { + querySpan.end(); + } + } + } + + /** + * Creates an {@link ArrowQueryResult} backed by a BigQuery Storage Read API session on the given + * destination table. + * + * @param destinationTable the destination table containing query results + * @param jobId the ID of the BigQuery query job + * @param contextMessage context describing why the ReadSession is being created (for error + * messages) + * @return a new {@link ArrowQueryResult} instance + * @throws BigQueryException if ReadSession creation fails + */ + private ArrowQueryResult createArrowQueryResultFromTable( + TableId destinationTable, JobId jobId, String contextMessage) { + 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 = 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 " + contextMessage, e); + } + + return ArrowQueryResultImpl.fromReadSession(readSession, jobId, client); + } + @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/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); From e61ade5f60f6ffce8ec99629191c953271b71ff4 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 16:35:44 -0400 Subject: [PATCH 02/12] fix(bigquery): address bot review comments on queryArrow and settings --- .../google/cloud/bigquery/BigQueryImpl.java | 33 +++++++++---------- 1 file changed, 15 insertions(+), 18 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 fd999540f4c1..53dfb40f6ee4 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 @@ -26,6 +26,7 @@ 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.gax.rpc.HeaderProvider; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; import com.google.api.services.bigquery.model.ProjectList; @@ -323,8 +324,9 @@ private static void configureReadSettings( } else { settingsBuilder.setCredentialsProvider(NoCredentialsProvider.create()); } - if (options.getMergedHeaderProvider(null) != null) { - settingsBuilder.setHeaderProvider(options.getMergedHeaderProvider(null)); + HeaderProvider headerProvider = options.getMergedHeaderProvider(null); + if (headerProvider != null) { + settingsBuilder.setHeaderProvider(headerProvider); } if (options.getUniverseDomain() != null) { settingsBuilder.setUniverseDomain(options.getUniverseDomain()); @@ -2379,15 +2381,16 @@ && getOptions().getOpenTelemetryTracer() != null) { } 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(); + if (arrowConfig.getQueryResultsFormat() != QueryResultsFormat.ARROW + || arrowConfig.getJobCreationMode() == null) { + QueryJobConfiguration.Builder builder = configuration.toBuilder(); + if (arrowConfig.getQueryResultsFormat() != QueryResultsFormat.ARROW) { + builder.setQueryResultsFormat(QueryResultsFormat.ARROW); + } + if (arrowConfig.getJobCreationMode() == null) { + builder.setJobCreationMode(QueryJobConfiguration.JobCreationMode.JOB_CREATION_OPTIONAL); + } + arrowConfig = builder.build(); } QueryRequestInfo requestInfo = @@ -2528,13 +2531,7 @@ && getOptions().getOpenTelemetryTracer() != null) { } 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; - } + Job completedJob = job.waitFor(); if (completedJob == null) { throw new BigQueryException(0, "Job no longer exists or could not be retrieved."); From f56b410e425d22c6a87184f26d1fa815cb228d50 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Mon, 14 Sep 2026 16:39:50 -0400 Subject: [PATCH 03/12] fix(bigquery): support IPv6 loopback for plaintext channel configuration --- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 53dfb40f6ee4..648153872a3e 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 @@ -351,7 +351,9 @@ private static void configureReadSettings( } int port = hostAndPort.getPortOrDefault(443); settingsBuilder.setEndpoint(endpointHost + ":" + port); - if (endpointHost.contains("localhost") || endpointHost.contains("127.0.0.1")) { + if (endpointHost.contains("localhost") + || endpointHost.contains("127.0.0.1") + || endpointHost.contains("::1")) { settingsBuilder.setTransportChannelProvider( BigQueryReadSettings.defaultGrpcTransportProviderBuilder() .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) From 741e9f1baf074ddceaa05d3c711de0509500bfd1 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 15:13:36 -0400 Subject: [PATCH 04/12] fix(bigquery): cache BigQueryReadClient by location and configure regional endpoint --- .../google/cloud/bigquery/BigQueryImpl.java | 67 ++++++++++++++++--- .../cloud/bigquery/BigQueryImplTest.java | 31 +++++++++ 2 files changed, 88 insertions(+), 10 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 648153872a3e..7661523510ec 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 @@ -281,7 +281,7 @@ public Page getNextPage() { } private final ReentrantLock readClientLock = new ReentrantLock(); - private transient BigQueryReadClient bqReadClient; + private transient Map bqReadClients; /** * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming @@ -292,18 +292,58 @@ public Page getNextPage() { * @throws BigQueryException if initializing the storage read client fails */ BigQueryReadClient getBigQueryReadClient() { + return getBigQueryReadClient(getOptions().getLocation()); + } + + /** + * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance for the specified + * location used for streaming Arrow query results, reusing credentials and channel configuration + * from this {@link BigQueryImpl}. + * + * @param location the regional location of the dataset/query + * @return the active BigQueryReadClient instance + * @throws BigQueryException if initializing the storage read client fails + */ + BigQueryReadClient getBigQueryReadClient(String location) { + String cacheKey = location != null ? location.toLowerCase() : "global"; readClientLock.lock(); try { - if (bqReadClient == null) { + if (bqReadClients == null) { + bqReadClients = Maps.newHashMap(); + } + BigQueryReadClient client = bqReadClients.get(cacheKey); + if (client == null && bqReadClients.containsKey("global")) { + client = bqReadClients.get("global"); + } + if (client == null) { BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder(); - configureReadSettings(settingsBuilder, getOptions()); + configureReadSettings(settingsBuilder, getOptions(), location); try { - bqReadClient = BigQueryReadClient.create(settingsBuilder.build()); + client = BigQueryReadClient.create(settingsBuilder.build()); + bqReadClients.put(cacheKey, client); } catch (IOException e) { - throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + throw new BigQueryException( + 0, "Failed to initialize BigQueryReadClient for location " + location, e); } } - return bqReadClient; + return client; + } finally { + readClientLock.unlock(); + } + } + + void setBigQueryReadClient(BigQueryReadClient client) { + setBigQueryReadClient(null, client); + } + + void setBigQueryReadClient(String location, BigQueryReadClient client) { + String cacheKey = location != null ? location.toLowerCase() : "global"; + readClientLock.lock(); + try { + if (bqReadClients == null) { + bqReadClients = Maps.newHashMap(); + } + bqReadClients.put(cacheKey, client); } finally { readClientLock.unlock(); } @@ -315,9 +355,10 @@ BigQueryReadClient getBigQueryReadClient() { * * @param settingsBuilder the builder to configure * @param options the source BigQueryOptions + * @param location the regional location of the dataset/query */ private static void configureReadSettings( - BigQueryReadSettings.Builder settingsBuilder, BigQueryOptions options) { + BigQueryReadSettings.Builder settingsBuilder, BigQueryOptions options, String location) { if (options.getCredentials() != null) { settingsBuilder.setCredentialsProvider( FixedCredentialsProvider.create(options.getCredentials())); @@ -359,6 +400,10 @@ private static void configureReadSettings( .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) .build()); } + } else if (location != null + && !location.equalsIgnoreCase("us") + && !location.equalsIgnoreCase("eu")) { + settingsBuilder.setEndpoint(location.toLowerCase() + "-bigquerystorage.googleapis.com:443"); } } @@ -2494,10 +2539,11 @@ && getOptions().getOpenTelemetryTracer() != null) { } String streamName = null; + String jobLocation = null; if (actualJobId != null && actualJobId.getJob() != null) { String jobProject = actualJobId.getProject() != null ? actualJobId.getProject() : projectId; - String jobLocation = + jobLocation = actualJobId.getLocation() != null ? actualJobId.getLocation() : (content.getLocation() != null @@ -2513,7 +2559,7 @@ && getOptions().getOpenTelemetryTracer() != null) { BigQueryReadClient client = null; if (streamName != null) { - client = getBigQueryReadClient(); + client = getBigQueryReadClient(jobLocation); } JobCreationReason jobCreationReason = @@ -2591,7 +2637,8 @@ private ArrowQueryResult createArrowQueryResultFromTable( "projects/%s/datasets/%s/tables/%s", destProject, destinationTable.getDataset(), destinationTable.getTable()); - BigQueryReadClient client = getBigQueryReadClient(); + String location = jobId != null ? jobId.getLocation() : getOptions().getLocation(); + BigQueryReadClient client = getBigQueryReadClient(location); CreateReadSessionRequest request = CreateReadSessionRequest.newBuilder() 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..e043f7d5ccd0 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java @@ -37,6 +37,7 @@ import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import static org.mockito.Mockito.withSettings; import com.google.api.client.googleapis.json.GoogleJsonError; import com.google.api.client.googleapis.json.GoogleJsonResponseException; @@ -70,6 +71,7 @@ import com.google.cloud.bigquery.spi.BigQueryRpcFactory; import com.google.cloud.bigquery.spi.v2.BigQueryRpc; import com.google.cloud.bigquery.spi.v2.HttpBigQueryRpc; +import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; import com.google.common.base.Function; import com.google.common.base.Supplier; import com.google.common.collect.ImmutableList; @@ -2940,6 +2942,35 @@ void testQueryArrowDefaultsToJobCreationOptional() throws IOException, Interrupt assertEquals("ARROW", requestPb.getQueryResultsFormat()); } + @Test + void testGetBigQueryReadClientCachingByLocation() { + BigQueryReadClient mockClientUs = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + BigQueryReadClient mockClientEu = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + + bigquery = options.getService(); + BigQueryImpl bigQueryImpl = (BigQueryImpl) bigquery; + bigQueryImpl.setBigQueryReadClient("us-east1", mockClientUs); + bigQueryImpl.setBigQueryReadClient("europe-west1", mockClientEu); + + assertSame(mockClientUs, bigQueryImpl.getBigQueryReadClient("us-east1")); + assertSame(mockClientEu, bigQueryImpl.getBigQueryReadClient("europe-west1")); + } + + @Test + void testGetBigQueryReadClientFallbackToGlobal() { + BigQueryReadClient mockGlobalClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + + bigquery = options.getService(); + BigQueryImpl bigQueryImpl = (BigQueryImpl) bigquery; + bigQueryImpl.setBigQueryReadClient(mockGlobalClient); + + assertSame(mockGlobalClient, bigQueryImpl.getBigQueryReadClient()); + assertSame(mockGlobalClient, bigQueryImpl.getBigQueryReadClient("asia-northeast1")); + } + @Test void testGetQueryResults() throws IOException { JobId queryJob = JobId.of(JOB); From 4c8eeb6583889ae3bdac288b1579714d39b55af9 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 15:20:25 -0400 Subject: [PATCH 05/12] fix(bigquery): throw UnsupportedOperationException when query() called with QueryResultsFormat.ARROW --- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 2 +- .../test/java/com/google/cloud/bigquery/BigQueryImplTest.java | 4 ++-- 2 files changed, 3 insertions(+), 3 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 7661523510ec..5b1adcd9ec50 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 @@ -2320,7 +2320,7 @@ public Object queryWithTimeout( Job.checkNotDryRun(configuration, "query"); if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) { - throw new IllegalArgumentException( + throw new UnsupportedOperationException( "QueryResultsFormat.ARROW is not supported with query(). Use queryArrow() instead."); } 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 e043f7d5ccd0..84ad8c7371b0 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 @@ -2913,8 +2913,8 @@ void testQueryThrowsWhenArrowResultsFormat() { .setQueryResultsFormat(QueryResultsFormat.ARROW) .build(); bigquery = options.getService(); - IllegalArgumentException exception = - assertThrows(IllegalArgumentException.class, () -> bigquery.query(config)); + UnsupportedOperationException exception = + assertThrows(UnsupportedOperationException.class, () -> bigquery.query(config)); assertTrue(exception.getMessage().contains("Use queryArrow() instead")); } From b5ebf19658da207038d5186eccb1e2cd95e91b7c Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 15:25:17 -0400 Subject: [PATCH 06/12] chore(bigquery): inline and remove unused queryArrowWithTimeout internal helper --- .../google/cloud/bigquery/BigQueryImpl.java | 22 +------------------ 1 file changed, 1 insertion(+), 21 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 5b1adcd9ec50..7d8fefec8879 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 @@ -2396,23 +2396,6 @@ public ArrowQueryResult queryArrow(QueryJobConfiguration configuration, JobOptio public ArrowQueryResult queryArrow( QueryJobConfiguration configuration, JobId jobId, JobOption... options) throws InterruptedException, JobException { - return queryArrowWithTimeout(configuration, jobId, null, options); - } - - /** - * Executes a query in Arrow format with an optional execution timeout. - * - * @param configuration query job configuration - * @param jobId job identifier, or {@code null} - * @param timeoutMs query timeout in milliseconds, or {@code null} - * @param options query job options - * @return an {@link ArrowQueryResult} for streaming results - * @throws InterruptedException if interrupted while awaiting results - * @throws JobException if the query job fails - */ - 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; @@ -2421,7 +2404,7 @@ && getOptions().getOpenTelemetryTracer() != null) { querySpan = getOptions() .getOpenTelemetryTracer() - .spanBuilder("com.google.cloud.bigquery.BigQuery.queryArrowWithTimeout") + .spanBuilder("com.google.cloud.bigquery.BigQuery.queryArrow") .setAllAttributes(jobId != null ? jobId.getOtelAttributes() : Attributes.empty()) .setAllAttributes(otelAttributesFromOptions(options)) .startSpan(); @@ -2457,9 +2440,6 @@ && getOptions().getOpenTelemetryTracer() != null) { } 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 = From 41ae83e24b71e6bb45b27934e91cde8b2e7c15c1 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 15:56:28 -0400 Subject: [PATCH 07/12] chore(bigquery): add explanatory comments to queryArrow --- .../java/com/google/cloud/bigquery/BigQueryImpl.java | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 7d8fefec8879..992657d8b3e0 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 @@ -2410,6 +2410,7 @@ && getOptions().getOpenTelemetryTracer() != null) { .startSpan(); } try (Scope queryScope = querySpan != null ? querySpan.makeCurrent() : null) { + // 1. Ensure QueryResultsFormat is ARROW and default JobCreationMode is JOB_CREATION_OPTIONAL QueryJobConfiguration arrowConfig = configuration; if (arrowConfig.getQueryResultsFormat() != QueryResultsFormat.ARROW || arrowConfig.getJobCreationMode() == null) { @@ -2426,10 +2427,12 @@ && getOptions().getOpenTelemetryTracer() != null) { QueryRequestInfo requestInfo = new QueryRequestInfo(arrowConfig, getOptions().getDataFormatOptions()); + // 2. Check if fast-path query execution is supported (no destination table or custom job ID) boolean useFastPath = requestInfo.isFastQuerySupported() && (jobId == null || jobId.getJob() == null); if (useFastPath) { + // Fast Path: Execute query directly via the jobs.query REST RPC String projectId = jobId != null && jobId.getProject() != null ? jobId.getProject() @@ -2464,6 +2467,8 @@ && getOptions().getOpenTelemetryTracer() != null) { JobId actualJobId = results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : jobId; + // If the query didn't complete within the fast-query timeout, wait for the job and stream + // from table if (results.getJobComplete() != null && !results.getJobComplete()) { if (actualJobId == null) { throw new BigQueryException( @@ -2494,6 +2499,7 @@ && getOptions().getOpenTelemetryTracer() != null) { destinationTable, job.getJobId(), "completed query"); } + // Deserialize Arrow schema and record batch from fast-path query response org.apache.arrow.vector.types.pojo.Schema arrowSchema = null; if (results.getArrowSchema() != null) { try { @@ -2518,6 +2524,8 @@ && getOptions().getOpenTelemetryTracer() != null) { initialBatchBytes = results.getArrowRecordBatch().decodeSerializedRecordBatch(); } + // Construct default Storage Read API stream name for streaming subsequent pages (if job + // created) String streamName = null; String jobLocation = null; if (actualJobId != null && actualJobId.getJob() != null) { @@ -2557,7 +2565,8 @@ && getOptions().getOpenTelemetryTracer() != null) { streamName, client); } else { - // Fallback path: jobs.insert + BigQuery Storage Read API + // Fallback Path: Submit query job via jobs.insert and stream destination table via Storage + // Read API Job job = create(JobInfo.of(jobId, arrowConfig), options); Job completedJob = job.waitFor(); From e1e593dc98ada9225cec33ed2d8660e17b24c611 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 16:56:04 -0400 Subject: [PATCH 08/12] fix(bigquery): isolate user-provided global read client and improve location fallback --- .../java/com/google/cloud/bigquery/BigQueryImpl.java | 11 +++++++++-- .../com/google/cloud/bigquery/BigQueryImplTest.java | 4 ++++ 2 files changed, 13 insertions(+), 2 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index 992657d8b3e0..79eb5f9954c5 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 @@ -282,6 +282,7 @@ public Page getNextPage() { private final ReentrantLock readClientLock = new ReentrantLock(); private transient Map bqReadClients; + private transient boolean isGlobalClientUserProvided; /** * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming @@ -312,7 +313,7 @@ BigQueryReadClient getBigQueryReadClient(String location) { bqReadClients = Maps.newHashMap(); } BigQueryReadClient client = bqReadClients.get(cacheKey); - if (client == null && bqReadClients.containsKey("global")) { + if (client == null && isGlobalClientUserProvided && bqReadClients.containsKey("global")) { client = bqReadClients.get("global"); } if (client == null) { @@ -344,6 +345,9 @@ void setBigQueryReadClient(String location, BigQueryReadClient client) { bqReadClients = Maps.newHashMap(); } bqReadClients.put(cacheKey, client); + if ("global".equals(cacheKey)) { + isGlobalClientUserProvided = true; + } } finally { readClientLock.unlock(); } @@ -2626,7 +2630,10 @@ private ArrowQueryResult createArrowQueryResultFromTable( "projects/%s/datasets/%s/tables/%s", destProject, destinationTable.getDataset(), destinationTable.getTable()); - String location = jobId != null ? jobId.getLocation() : getOptions().getLocation(); + String location = + (jobId != null && jobId.getLocation() != null) + ? jobId.getLocation() + : getOptions().getLocation(); BigQueryReadClient client = getBigQueryReadClient(location); CreateReadSessionRequest request = 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 84ad8c7371b0..9b40d0dde64c 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 @@ -2962,11 +2962,15 @@ void testGetBigQueryReadClientCachingByLocation() { void testGetBigQueryReadClientFallbackToGlobal() { BigQueryReadClient mockGlobalClient = mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); + BigQueryReadClient mockRegionalClient = + mock(BigQueryReadClient.class, withSettings().withoutAnnotations()); bigquery = options.getService(); BigQueryImpl bigQueryImpl = (BigQueryImpl) bigquery; bigQueryImpl.setBigQueryReadClient(mockGlobalClient); + bigQueryImpl.setBigQueryReadClient("us-east1", mockRegionalClient); + assertSame(mockRegionalClient, bigQueryImpl.getBigQueryReadClient("us-east1")); assertSame(mockGlobalClient, bigQueryImpl.getBigQueryReadClient()); assertSame(mockGlobalClient, bigQueryImpl.getBigQueryReadClient("asia-northeast1")); } From 976c20c7a90ab57e9c3655c2d7a1ceed53b41627 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 17:38:11 -0400 Subject: [PATCH 09/12] fix(bigquery): address review feedback on queryArrow and client lifecycle --- .../com/google/cloud/bigquery/BigQuery.java | 5 ++- .../google/cloud/bigquery/BigQueryImpl.java | 40 ++++++++++++++++--- 2 files changed, 39 insertions(+), 6 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..00c99654f210 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,10 @@ * * @see Google Cloud BigQuery */ -public interface BigQuery extends Service { +public interface BigQuery extends Service, AutoCloseable { + + @Override + default 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 79eb5f9954c5..d788f446937a 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 @@ -280,10 +280,34 @@ public Page getNextPage() { } } - private final ReentrantLock readClientLock = new ReentrantLock(); + private transient ReentrantLock readClientLock = new ReentrantLock(); private transient Map bqReadClients; private transient boolean isGlobalClientUserProvided; + private void readObject(java.io.ObjectInputStream in) throws IOException, ClassNotFoundException { + in.defaultReadObject(); + readClientLock = new ReentrantLock(); + } + + @Override + public void close() { + if (readClientLock != null) { + readClientLock.lock(); + try { + if (bqReadClients != null) { + for (BigQueryReadClient client : bqReadClients.values()) { + if (client != null) { + client.close(); + } + } + bqReadClients.clear(); + } + } finally { + readClientLock.unlock(); + } + } + } + /** * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming * Arrow query results, reusing credentials and channel configuration from this {@link @@ -2433,7 +2457,9 @@ && getOptions().getOpenTelemetryTracer() != null) { // 2. Check if fast-path query execution is supported (no destination table or custom job ID) boolean useFastPath = - requestInfo.isFastQuerySupported() && (jobId == null || jobId.getJob() == null); + requestInfo.isFastQuerySupported() + && (jobId == null || jobId.getJob() == null) + && (options == null || options.length == 0); if (useFastPath) { // Fast Path: Execute query directly via the jobs.query REST RPC @@ -2462,7 +2488,7 @@ && getOptions().getOpenTelemetryTracer() != null) { throw BigQueryException.translateAndThrow(e); } - if (results.getErrors() != null) { + if (results.getErrors() != null && !results.getErrors().isEmpty()) { List bigQueryErrors = Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION); throw new BigQueryException(bigQueryErrors); @@ -2510,7 +2536,7 @@ && getOptions().getOpenTelemetryTracer() != null) { arrowSchema = ArrowDeserializer.deserializeSchema( results.getArrowSchema().decodeSerializedSchema()); - } catch (IOException e) { + } catch (IOException | IllegalArgumentException e) { throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e); } } @@ -2525,7 +2551,11 @@ && getOptions().getOpenTelemetryTracer() != null) { byte[] initialBatchBytes = null; if (results.getArrowRecordBatch() != null && results.getArrowRecordBatch().getSerializedRecordBatch() != null) { - initialBatchBytes = results.getArrowRecordBatch().decodeSerializedRecordBatch(); + try { + initialBatchBytes = results.getArrowRecordBatch().decodeSerializedRecordBatch(); + } catch (IllegalArgumentException e) { + throw new BigQueryException(0, "Failed to decode Arrow record batch from response", e); + } } // Construct default Storage Read API stream name for streaming subsequent pages (if job From 21f59a2774b32148b5a7d2658042e21d0f288be7 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 20:38:39 -0400 Subject: [PATCH 10/12] fix(bigquery): improve emulator endpoint detection and exception-safe client shutdown --- .../java/com/google/cloud/bigquery/BigQueryImpl.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index d788f446937a..e9e449b499d0 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 @@ -297,7 +297,11 @@ public void close() { if (bqReadClients != null) { for (BigQueryReadClient client : bqReadClients.values()) { if (client != null) { - client.close(); + try { + client.close(); + } catch (Exception ignored) { + // Suppress exception to ensure all other clients are closed + } } } bqReadClients.clear(); @@ -420,7 +424,8 @@ private static void configureReadSettings( } int port = hostAndPort.getPortOrDefault(443); settingsBuilder.setEndpoint(endpointHost + ":" + port); - if (endpointHost.contains("localhost") + if (host.startsWith("http://") + || endpointHost.contains("localhost") || endpointHost.contains("127.0.0.1") || endpointHost.contains("::1")) { settingsBuilder.setTransportChannelProvider( From 92f3ee96a19d7938c7844924bbd43a8003c5d8a8 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 20:44:24 -0400 Subject: [PATCH 11/12] fix(bigquery): preserve default ADC resolution when options credentials are null --- .../src/main/java/com/google/cloud/bigquery/BigQueryImpl.java | 3 --- 1 file changed, 3 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 e9e449b499d0..20fa5db9eb5d 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,7 +24,6 @@ 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.gax.rpc.HeaderProvider; import com.google.api.services.bigquery.model.ErrorProto; @@ -394,8 +393,6 @@ private static void configureReadSettings( if (options.getCredentials() != null) { settingsBuilder.setCredentialsProvider( FixedCredentialsProvider.create(options.getCredentials())); - } else { - settingsBuilder.setCredentialsProvider(NoCredentialsProvider.create()); } HeaderProvider headerProvider = options.getMergedHeaderProvider(null); if (headerProvider != null) { From ba202a99e69b87d4abecaaca60fbd291c66026a1 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 16 Sep 2026 21:18:37 -0400 Subject: [PATCH 12/12] fix(bigquery): revert AutoCloseable and simplify client synchronization --- .../com/google/cloud/bigquery/BigQuery.java | 5 +- .../google/cloud/bigquery/BigQueryImpl.java | 90 ++++++------------- 2 files changed, 26 insertions(+), 69 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 00c99654f210..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,10 +40,7 @@ * * @see Google Cloud BigQuery */ -public interface BigQuery extends Service, AutoCloseable { - - @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 20fa5db9eb5d..be9c7c279ec7 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 @@ -74,7 +74,6 @@ 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,38 +278,9 @@ public Page getNextPage() { } } - private transient ReentrantLock readClientLock = new ReentrantLock(); private transient Map bqReadClients; private transient boolean isGlobalClientUserProvided; - private void readObject(java.io.ObjectInputStream in) throws IOException, ClassNotFoundException { - in.defaultReadObject(); - readClientLock = new ReentrantLock(); - } - - @Override - public void close() { - if (readClientLock != null) { - readClientLock.lock(); - try { - if (bqReadClients != null) { - for (BigQueryReadClient client : bqReadClients.values()) { - if (client != null) { - try { - client.close(); - } catch (Exception ignored) { - // Suppress exception to ensure all other clients are closed - } - } - } - bqReadClients.clear(); - } - } finally { - readClientLock.unlock(); - } - } - } - /** * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming * Arrow query results, reusing credentials and channel configuration from this {@link @@ -332,51 +302,41 @@ BigQueryReadClient getBigQueryReadClient() { * @return the active BigQueryReadClient instance * @throws BigQueryException if initializing the storage read client fails */ - BigQueryReadClient getBigQueryReadClient(String location) { + synchronized BigQueryReadClient getBigQueryReadClient(String location) { String cacheKey = location != null ? location.toLowerCase() : "global"; - readClientLock.lock(); - try { - if (bqReadClients == null) { - bqReadClients = Maps.newHashMap(); - } - BigQueryReadClient client = bqReadClients.get(cacheKey); - if (client == null && isGlobalClientUserProvided && bqReadClients.containsKey("global")) { - client = bqReadClients.get("global"); - } - if (client == null) { - BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder(); - configureReadSettings(settingsBuilder, getOptions(), location); - try { - client = BigQueryReadClient.create(settingsBuilder.build()); - bqReadClients.put(cacheKey, client); - } catch (IOException e) { - throw new BigQueryException( - 0, "Failed to initialize BigQueryReadClient for location " + location, e); - } + if (bqReadClients == null) { + bqReadClients = Maps.newHashMap(); + } + BigQueryReadClient client = bqReadClients.get(cacheKey); + if (client == null && isGlobalClientUserProvided && bqReadClients.containsKey("global")) { + client = bqReadClients.get("global"); + } + if (client == null) { + BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder(); + configureReadSettings(settingsBuilder, getOptions(), location); + try { + client = BigQueryReadClient.create(settingsBuilder.build()); + bqReadClients.put(cacheKey, client); + } catch (IOException e) { + throw new BigQueryException( + 0, "Failed to initialize BigQueryReadClient for location " + location, e); } - return client; - } finally { - readClientLock.unlock(); } + return client; } void setBigQueryReadClient(BigQueryReadClient client) { setBigQueryReadClient(null, client); } - void setBigQueryReadClient(String location, BigQueryReadClient client) { + synchronized void setBigQueryReadClient(String location, BigQueryReadClient client) { String cacheKey = location != null ? location.toLowerCase() : "global"; - readClientLock.lock(); - try { - if (bqReadClients == null) { - bqReadClients = Maps.newHashMap(); - } - bqReadClients.put(cacheKey, client); - if ("global".equals(cacheKey)) { - isGlobalClientUserProvided = true; - } - } finally { - readClientLock.unlock(); + if (bqReadClients == null) { + bqReadClients = Maps.newHashMap(); + } + bqReadClients.put(cacheKey, client); + if ("global".equals(cacheKey)) { + isGlobalClientUserProvided = true; } }