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..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
@@ -18,11 +18,15 @@
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.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 +49,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 +63,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 +280,88 @@ 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());
+ }
+ 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 (endpointHost.contains("localhost")
+ || endpointHost.contains("127.0.0.1")
+ || endpointHost.contains("::1")) {
+ settingsBuilder.setTransportChannelProvider(
+ BigQueryReadSettings.defaultGrpcTransportProviderBuilder()
+ .setChannelConfigurator(ManagedChannelBuilder::usePlaintext)
+ .build());
+ }
+ }
+ }
+
private final HttpBigQueryRpc bigQueryRpc;
private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG =
@@ -2178,6 +2274,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 +2341,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 {
+ 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.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());
+
+ 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 = 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());
+
+ 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);