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..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
@@ -18,11 +18,14 @@
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.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;
@@ -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,11 +62,15 @@
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;
@@ -266,6 +278,125 @@ public Page getNextPage() {
}
}
+ private transient Map bqReadClients;
+ private transient boolean isGlobalClientUserProvided;
+
+ /**
+ * 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() {
+ 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
+ */
+ synchronized BigQueryReadClient getBigQueryReadClient(String location) {
+ String cacheKey = location != null ? location.toLowerCase() : "global";
+ 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;
+ }
+
+ void setBigQueryReadClient(BigQueryReadClient client) {
+ setBigQueryReadClient(null, client);
+ }
+
+ synchronized void setBigQueryReadClient(String location, BigQueryReadClient client) {
+ String cacheKey = location != null ? location.toLowerCase() : "global";
+ if (bqReadClients == null) {
+ bqReadClients = Maps.newHashMap();
+ }
+ bqReadClients.put(cacheKey, client);
+ if ("global".equals(cacheKey)) {
+ isGlobalClientUserProvided = true;
+ }
+ }
+
+ /**
+ * 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
+ * @param location the regional location of the dataset/query
+ */
+ private static void configureReadSettings(
+ BigQueryReadSettings.Builder settingsBuilder, BigQueryOptions options, String location) {
+ if (options.getCredentials() != null) {
+ settingsBuilder.setCredentialsProvider(
+ FixedCredentialsProvider.create(options.getCredentials()));
+ }
+ HeaderProvider headerProvider = options.getMergedHeaderProvider(null);
+ if (headerProvider != null) {
+ settingsBuilder.setHeaderProvider(headerProvider);
+ }
+ 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 (host.startsWith("http://")
+ || endpointHost.contains("localhost")
+ || endpointHost.contains("127.0.0.1")
+ || endpointHost.contains("::1")) {
+ settingsBuilder.setTransportChannelProvider(
+ BigQueryReadSettings.defaultGrpcTransportProviderBuilder()
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .build());
+ }
+ } else if (location != null
+ && !location.equalsIgnoreCase("us")
+ && !location.equalsIgnoreCase("eu")) {
+ settingsBuilder.setEndpoint(location.toLowerCase() + "-bigquerystorage.googleapis.com:443");
+ }
+ }
+
private final HttpBigQueryRpc bigQueryRpc;
private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG =
@@ -2178,6 +2309,11 @@ public Object queryWithTimeout(
throws InterruptedException, JobException {
Job.checkNotDryRun(configuration, "query");
+ if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) {
+ throw new UnsupportedOperationException(
+ "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 +2376,275 @@ && 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 {
+ 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.queryArrow")
+ .setAllAttributes(jobId != null ? jobId.getOtelAttributes() : Attributes.empty())
+ .setAllAttributes(otelAttributesFromOptions(options))
+ .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) {
+ 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 =
+ 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)
+ && (options == null || options.length == 0);
+
+ if (useFastPath) {
+ // Fast Path: Execute query directly via the jobs.query REST RPC
+ 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());
+ }
+ 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 && !results.getErrors().isEmpty()) {
+ 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 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(
+ 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");
+ }
+
+ // 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 {
+ arrowSchema =
+ ArrowDeserializer.deserializeSchema(
+ results.getArrowSchema().decodeSerializedSchema());
+ } catch (IOException | IllegalArgumentException 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) {
+ 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
+ // created)
+ String streamName = null;
+ String jobLocation = null;
+ if (actualJobId != null && actualJobId.getJob() != null) {
+ String jobProject =
+ actualJobId.getProject() != null ? actualJobId.getProject() : projectId;
+ 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(jobLocation);
+ }
+
+ 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: 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();
+
+ 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());
+
+ String location =
+ (jobId != null && jobId.getLocation() != null)
+ ? jobId.getLocation()
+ : getOptions().getLocation();
+ BigQueryReadClient client = getBigQueryReadClient(location);
+
+ 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..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
@@ -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;
@@ -2904,6 +2906,75 @@ 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();
+ UnsupportedOperationException exception =
+ assertThrows(UnsupportedOperationException.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 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());
+ 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"));
+ }
+
@Test
void testGetQueryResults() throws IOException {
JobId queryJob = JobId.of(JOB);