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..4709f6c88896 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,16 @@
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.gax.rpc.ServerStream;
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 +50,13 @@
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.ReadRowsRequest;
+import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
+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 +66,24 @@
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 com.google.common.primitives.Longs;
+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.ArrayDeque;
import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Iterator;
import java.util.List;
import java.util.Map;
+import java.util.Queue;
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 +288,256 @@ public Page getNextPage() {
}
}
+ /**
+ * NextPageFetcher implementation for queries returning results in Arrow format. Reads subsequent
+ * pages from the job's default gRPC storage read stream.
+ *
+ * Note: Neither {@link Page} nor {@link TableResult} implements {@link AutoCloseable}. The
+ * underlying gRPC stream is automatically canceled and resources released when iteration reaches
+ * the end (or maximum results requested) or when an error occurs. Callers that do not iterate to
+ * completion rely on server-side stream timeouts and garbage collection to release stream
+ * resources.
+ */
+ static class ArrowQueryPageFetcher implements NextPageFetcher {
+ private static final long serialVersionUID = 1L;
+ private static final long DEFAULT_PAGE_SIZE = 10000L;
+
+ private final JobId jobId;
+ private final Schema schema;
+ private final byte[] arrowSchemaBytes;
+ private final BigQueryOptions serviceOptions;
+ private final long maxResults;
+ private final Map optionsMap;
+
+ private transient org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo;
+ private transient BigQueryReadClient bqReadClient;
+ private transient ServerStream stream;
+ private transient Iterator streamIterator;
+ private final Queue buffer = new ArrayDeque<>();
+ private long totalRowsReturned = 0L;
+ private boolean streamClosed = false;
+
+ ArrowQueryPageFetcher(
+ JobId jobId,
+ Schema schema,
+ byte[] arrowSchemaBytes,
+ org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo,
+ BigQueryOptions serviceOptions,
+ long initialRowOffset,
+ Long maxResults,
+ Map optionsMap) {
+ this.jobId = jobId;
+ this.schema = schema;
+ this.arrowSchemaBytes = arrowSchemaBytes;
+ this.arrowSchemaPojo = arrowSchemaPojo;
+ this.serviceOptions = serviceOptions;
+ this.totalRowsReturned = initialRowOffset;
+ this.maxResults = maxResults != null ? maxResults : Long.MAX_VALUE;
+ this.optionsMap = optionsMap;
+ }
+
+ @Override
+ public Page getNextPage() {
+ if (streamClosed || totalRowsReturned >= maxResults) {
+ closeClient();
+ return null;
+ }
+
+ Number optionPageSize =
+ optionsMap != null ? (Number) optionsMap.get(BigQueryRpc.Option.MAX_RESULTS) : null;
+ long pageSize =
+ optionPageSize != null && optionPageSize.longValue() > 0
+ ? optionPageSize.longValue()
+ : DEFAULT_PAGE_SIZE;
+ List rowBatch = new ArrayList<>((int) Math.min(pageSize, 10000L));
+
+ try {
+ String location =
+ jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation();
+ if (location == null) {
+ throw new BigQueryException(
+ 0, "Location must be specified to read Arrow rows from storage stream");
+ }
+
+ if (bqReadClient == null) {
+ BigQuery service = serviceOptions.getService();
+ if (service instanceof BigQueryImpl) {
+ BigQueryImpl impl = (BigQueryImpl) service;
+ bqReadClient = impl.getBigQueryReadClient();
+ } else {
+ throw new IllegalStateException(
+ "Arrow query result pagination requires an instance of BigQueryImpl to manage BigQueryReadClient lifecycle");
+ }
+ }
+
+ if (streamIterator == null) {
+ String streamName =
+ String.format(
+ "projects/%s/locations/%s/jobs/%s/streams/_default",
+ jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(),
+ location,
+ jobId.getJob());
+
+ ReadRowsRequest readRowsRequest =
+ ReadRowsRequest.newBuilder()
+ .setReadStream(streamName)
+ .setOffset(totalRowsReturned)
+ .build();
+
+ stream = bqReadClient.readRowsCallable().call(readRowsRequest);
+ streamIterator = stream.iterator();
+ }
+
+ if (arrowSchemaPojo == null && arrowSchemaBytes != null) {
+ arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes);
+ }
+
+ boolean hasMore =
+ ArrowDeserializer.loadArrowRows(
+ streamIterator,
+ arrowSchemaPojo,
+ schema,
+ rowBatch,
+ buffer,
+ pageSize,
+ totalRowsReturned,
+ maxResults);
+
+ if (rowBatch.isEmpty()) {
+ streamClosed = true;
+ closeClient();
+ return null;
+ }
+
+ totalRowsReturned += rowBatch.size();
+
+ String nextPageToken = null;
+ if (hasMore && totalRowsReturned < maxResults) {
+ nextPageToken = String.valueOf(totalRowsReturned);
+ } else {
+ streamClosed = true;
+ closeClient();
+ }
+
+ return new PageImpl<>(this, nextPageToken, rowBatch);
+
+ } catch (BigQueryException e) {
+ streamClosed = true;
+ closeClient();
+ throw e;
+ } catch (Exception e) {
+ streamClosed = true;
+ closeClient();
+ throw new BigQueryException(0, "Failed to read Arrow rows from storage stream", e);
+ }
+ }
+
+ private void closeClient() {
+ if (stream != null) {
+ try {
+ stream.cancel();
+ } catch (Exception e) {
+ // Ignore cancellation exceptions
+ }
+ }
+ bqReadClient = null;
+ streamIterator = null;
+ stream = null;
+ }
+ }
+
+ private final ReentrantLock readClientLock = new ReentrantLock();
+ private transient BigQueryReadClient bqReadClient;
+
+ @VisibleForTesting
+ void setBigQueryReadClient(BigQueryReadClient client) {
+ readClientLock.lock();
+ try {
+ this.bqReadClient = client;
+ } 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
+ * 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 =
@@ -1840,7 +2112,7 @@ public TableDataList call() throws IOException {
}
}
- private static Iterable transformTableData(
+ private static List transformTableData(
Iterable tableDataPb, final Schema schema, boolean useInt64Timestamps) {
return ImmutableList.copyOf(
Iterables.transform(
@@ -2079,8 +2351,21 @@ public com.google.api.services.bigquery.model.QueryResponse call()
long numRows;
Schema schema;
- if (results.getJobComplete() && results.getSchema() != null) {
- schema = Schema.fromPb(results.getSchema());
+ boolean isArrow = results.getArrowSchema() != null;
+ org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo = null;
+ byte[] arrowSchemaBytes = null;
+ if (results.getJobComplete() && (results.getSchema() != null || isArrow)) {
+ if (isArrow) {
+ arrowSchemaBytes = results.getArrowSchema().decodeSerializedSchema();
+ try {
+ arrowSchemaPojo = ArrowDeserializer.deserializeSchema(arrowSchemaBytes);
+ } catch (IOException e) {
+ throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e);
+ }
+ schema = ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchemaPojo);
+ } else {
+ schema = Schema.fromPb(results.getSchema());
+ }
if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) {
numRows = 0L;
} else if (results.getNumDmlAffectedRows() != null) {
@@ -2108,25 +2393,75 @@ public com.google.api.services.bigquery.model.QueryResponse call()
SessionInfo sessionInfo =
results.getSessionInfo() != null ? SessionInfo.fromPb(results.getSessionInfo()) : null;
- if (results.getPageToken() != null) {
+ Collection firstPageRows;
+ if (isArrow) {
+ if (results.getArrowRecordBatch() != null
+ && results.getArrowRecordBatch().getSerializedRecordBatch() != null) {
+ try {
+ firstPageRows =
+ ArrowDeserializer.deserializeRecordBatch(
+ results.getArrowRecordBatch().decodeSerializedRecordBatch(),
+ schema,
+ arrowSchemaPojo);
+ } catch (IOException e) {
+ throw new BigQueryException(0, "Failed to deserialize Arrow record batch", e);
+ }
+ } else {
+ firstPageRows = ImmutableList.of();
+ }
+ } else {
+ firstPageRows =
+ results.getRows() != null
+ ? transformTableData(
+ results.getRows(),
+ schema,
+ getOptions().getDataFormatOptions().useInt64Timestamp())
+ : ImmutableList.of();
+ }
+
+ boolean hasMorePages = results.getPageToken() != null;
+ long initialRowOffset = 0L;
+ if (hasMorePages && isArrow) {
+ Long parsedOffset = Longs.tryParse(results.getPageToken());
+ initialRowOffset = parsedOffset != null ? parsedOffset : firstPageRows.size();
+ if (content.getMaxResults() != null
+ && (initialRowOffset >= content.getMaxResults()
+ || firstPageRows.size() >= content.getMaxResults())) {
+ hasMorePages = false;
+ }
+ }
+
+ if (hasMorePages) {
JobId jobId = JobId.fromPb(results.getJobReference());
String cursor = results.getPageToken();
+
+ NextPageFetcher pageFetcher;
+ if (isArrow) {
+ pageFetcher =
+ new ArrowQueryPageFetcher(
+ jobId,
+ schema,
+ arrowSchemaBytes,
+ arrowSchemaPojo,
+ getOptions(),
+ initialRowOffset,
+ content.getMaxResults(),
+ optionMap(options));
+ } else {
+ pageFetcher = new QueryPageFetcher(jobId, schema, getOptions(), cursor, optionMap(options));
+ }
+
return TableResult.newBuilder()
.setSchema(schema)
.setTotalRows(numRows)
.setPageNoSchema(
new PageImpl<>(
// fetch next pages of results
- new QueryPageFetcher(jobId, schema, getOptions(), cursor, optionMap(options)),
- cursor,
- transformTableData(
- results.getRows(),
- schema,
- getOptions().getDataFormatOptions().useInt64Timestamp())))
+ pageFetcher, cursor, firstPageRows))
.setJobId(jobId)
.setQueryId(results.getQueryId())
.setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
- .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
+ .setRowsInPage((long) firstPageRows.size())
.setStatementType(statementType)
.setTotalBytesBilled(totalBytesBilled)
.setTotalBytesProcessed(totalBytesProcessed)
@@ -2143,16 +2478,13 @@ public com.google.api.services.bigquery.model.QueryResponse call()
new PageImpl<>(
new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)),
null,
- transformTableData(
- results.getRows(),
- schema,
- getOptions().getDataFormatOptions().useInt64Timestamp())))
+ firstPageRows))
// Return the JobID of the successful job
.setJobId(
results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : null)
.setQueryId(results.getQueryId())
.setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason()))
- .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L)
+ .setRowsInPage((long) firstPageRows.size())
.setStatementType(statementType)
.setTotalBytesBilled(totalBytesBilled)
.setTotalBytesProcessed(totalBytesProcessed)
@@ -2232,6 +2564,12 @@ && getOptions().getOpenTelemetryTracer() != null) {
return queryRpc(projectId, content, options);
}
+
+ if (configuration.getQueryResultsFormat() == QueryResultsFormat.ARROW) {
+ throw new IllegalArgumentException(
+ "Arrow results format is only supported for fast query path execution (e.g. no destination table, no custom clustering, etc.).");
+ }
+
return create(JobInfo.of(jobId, configuration), options);
} finally {
if (querySpan != null) {
@@ -2240,6 +2578,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/ArrowQueryPageFetcherTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryPageFetcherTest.java
new file mode 100644
index 000000000000..0b9b5cb1bd02
--- /dev/null
+++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryPageFetcherTest.java
@@ -0,0 +1,341 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.cloud.bigquery;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.mockito.Mockito.withSettings;
+
+import com.google.api.gax.paging.Page;
+import com.google.api.gax.rpc.ServerStream;
+import com.google.api.gax.rpc.ServerStreamingCallable;
+import com.google.cloud.bigquery.spi.v2.BigQueryRpc;
+import com.google.cloud.bigquery.storage.v1.BigQueryReadClient;
+import com.google.cloud.bigquery.storage.v1.ReadRowsRequest;
+import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.protobuf.ByteString;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+import java.io.ObjectInputStream;
+import java.io.ObjectOutputStream;
+import java.nio.channels.Channels;
+import java.util.Collections;
+import java.util.List;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.VectorUnloader;
+import org.apache.arrow.vector.ipc.WriteChannel;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.ipc.message.MessageSerializer;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+public class ArrowQueryPageFetcherTest {
+
+ private static final String PROJECT = "test-project";
+ private static final String LOCATION = "US";
+ private static final String JOB = "test-job";
+
+ private BufferAllocator allocator;
+ private org.apache.arrow.vector.types.pojo.Schema arrowSchema;
+ private Schema bqSchema;
+ private byte[] schemaBytes;
+
+ @BeforeEach
+ void setUp() throws IOException {
+ allocator = new RootAllocator(Long.MAX_VALUE);
+ arrowSchema =
+ new org.apache.arrow.vector.types.pojo.Schema(
+ ImmutableList.of(
+ org.apache.arrow.vector.types.pojo.Field.nullable(
+ "id", new ArrowType.Int(64, true))));
+ bqSchema = ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema);
+
+ try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+ MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema);
+ schemaBytes = out.toByteArray();
+ }
+ }
+
+ @AfterEach
+ void tearDown() {
+ allocator.close();
+ }
+
+ private byte[] createBatchBytes(List values) throws IOException {
+ BigIntVector idVector = new BigIntVector("id", allocator);
+ idVector.allocateNew(values.size());
+ for (int i = 0; i < values.size(); i++) {
+ idVector.set(i, values.get(i));
+ }
+ idVector.setValueCount(values.size());
+
+ try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(idVector))) {
+ VectorUnloader unloader = new VectorUnloader(root);
+ try (ArrowRecordBatch recordBatch = unloader.getRecordBatch();
+ ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+ WriteChannel channel = new WriteChannel(Channels.newChannel(out));
+ MessageSerializer.serialize(channel, recordBatch);
+ return out.toByteArray();
+ }
+ } finally {
+ idVector.close();
+ }
+ }
+
+ private BigQueryReadClient createMockReadClient(
+ ServerStreamingCallable mockCallable) {
+ BigQueryReadClient mockClient =
+ mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
+ when(mockClient.readRowsCallable()).thenReturn(mockCallable);
+ return mockClient;
+ }
+
+ @Test
+ void testGetNextPage_singlePage() throws IOException {
+ byte[] batchBytes = createBatchBytes(ImmutableList.of(10L, 20L));
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch =
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder()
+ .setSerializedRecordBatch(ByteString.copyFrom(batchBytes))
+ .build();
+ ReadRowsResponse response =
+ ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build();
+
+ @SuppressWarnings("unchecked")
+ ServerStreamingCallable mockCallable =
+ mock(ServerStreamingCallable.class, withSettings().withoutAnnotations());
+ @SuppressWarnings("unchecked")
+ ServerStream mockServerStream =
+ mock(ServerStream.class, withSettings().withoutAnnotations());
+ when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream);
+ when(mockServerStream.iterator()).thenReturn(ImmutableList.of(response).iterator());
+
+ BigQueryReadClient mockReadClient = createMockReadClient(mockCallable);
+
+ BigQueryOptions options =
+ BigQueryOptions.newBuilder().setProjectId(PROJECT).setLocation(LOCATION).build();
+ BigQuery service = options.getService();
+ ((BigQueryImpl) service).setBigQueryReadClient(mockReadClient);
+
+ JobId jobId = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
+ BigQueryImpl.ArrowQueryPageFetcher fetcher =
+ new BigQueryImpl.ArrowQueryPageFetcher(
+ jobId,
+ bqSchema,
+ schemaBytes,
+ arrowSchema,
+ options,
+ /* initialRowOffset= */ 0L,
+ /* maxResults= */ 10L,
+ Collections.emptyMap());
+
+ Page page = fetcher.getNextPage();
+ assertNotNull(page);
+ List rows = ImmutableList.copyOf(page.getValues());
+ assertEquals(2, rows.size());
+ assertEquals("10", rows.get(0).get(0).getStringValue());
+ assertEquals("20", rows.get(1).get(0).getStringValue());
+ assertNull(page.getNextPageToken());
+ assertFalse(page.hasNextPage());
+
+ verify(mockCallable).call(any(ReadRowsRequest.class));
+ }
+
+ @Test
+ void testGetNextPage_multiplePagesWithPageSizeOption() throws IOException {
+ byte[] batch1Bytes = createBatchBytes(ImmutableList.of(1L, 2L));
+ byte[] batch2Bytes = createBatchBytes(ImmutableList.of(3L, 4L));
+
+ ReadRowsResponse resp1 =
+ ReadRowsResponse.newBuilder()
+ .setArrowRecordBatch(
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder()
+ .setSerializedRecordBatch(ByteString.copyFrom(batch1Bytes))
+ .build())
+ .build();
+ ReadRowsResponse resp2 =
+ ReadRowsResponse.newBuilder()
+ .setArrowRecordBatch(
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder()
+ .setSerializedRecordBatch(ByteString.copyFrom(batch2Bytes))
+ .build())
+ .build();
+
+ @SuppressWarnings("unchecked")
+ ServerStreamingCallable mockCallable =
+ mock(ServerStreamingCallable.class, withSettings().withoutAnnotations());
+ @SuppressWarnings("unchecked")
+ ServerStream mockServerStream =
+ mock(ServerStream.class, withSettings().withoutAnnotations());
+ when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream);
+ when(mockServerStream.iterator()).thenReturn(ImmutableList.of(resp1, resp2).iterator());
+
+ BigQueryReadClient mockReadClient = createMockReadClient(mockCallable);
+
+ BigQueryOptions options =
+ BigQueryOptions.newBuilder().setProjectId(PROJECT).setLocation(LOCATION).build();
+ BigQuery service = options.getService();
+ ((BigQueryImpl) service).setBigQueryReadClient(mockReadClient);
+
+ JobId jobId = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
+ BigQueryImpl.ArrowQueryPageFetcher fetcher =
+ new BigQueryImpl.ArrowQueryPageFetcher(
+ jobId,
+ bqSchema,
+ schemaBytes,
+ arrowSchema,
+ options,
+ /* initialRowOffset= */ 0L,
+ /* maxResults= */ 10L,
+ ImmutableMap.of(BigQueryRpc.Option.MAX_RESULTS, 2L));
+
+ Page page1 = fetcher.getNextPage();
+ assertNotNull(page1);
+ List page1Rows = ImmutableList.copyOf(page1.getValues());
+ assertEquals(2, page1Rows.size());
+ assertEquals("1", page1Rows.get(0).get(0).getStringValue());
+ assertEquals("2", page1Rows.get(1).get(0).getStringValue());
+ assertTrue(page1.hasNextPage());
+ assertEquals("2", page1.getNextPageToken());
+
+ Page page2 = page1.getNextPage();
+ assertNotNull(page2);
+ List page2Rows = ImmutableList.copyOf(page2.getValues());
+ assertEquals(2, page2Rows.size());
+ assertEquals("3", page2Rows.get(0).get(0).getStringValue());
+ assertEquals("4", page2Rows.get(1).get(0).getStringValue());
+ assertFalse(page2.hasNextPage());
+ assertNull(page2.getNextPageToken());
+ assertNull(page2.getNextPage());
+ }
+
+ @Test
+ void testGetNextPage_respectsMaxResults() throws IOException {
+ byte[] batchBytes = createBatchBytes(ImmutableList.of(1L, 2L, 3L));
+ ReadRowsResponse resp =
+ ReadRowsResponse.newBuilder()
+ .setArrowRecordBatch(
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder()
+ .setSerializedRecordBatch(ByteString.copyFrom(batchBytes))
+ .build())
+ .build();
+
+ @SuppressWarnings("unchecked")
+ ServerStreamingCallable mockCallable =
+ mock(ServerStreamingCallable.class, withSettings().withoutAnnotations());
+ @SuppressWarnings("unchecked")
+ ServerStream mockServerStream =
+ mock(ServerStream.class, withSettings().withoutAnnotations());
+ when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream);
+ when(mockServerStream.iterator()).thenReturn(ImmutableList.of(resp).iterator());
+
+ BigQueryReadClient mockReadClient = createMockReadClient(mockCallable);
+
+ BigQueryOptions options =
+ BigQueryOptions.newBuilder().setProjectId(PROJECT).setLocation(LOCATION).build();
+ BigQuery service = options.getService();
+ ((BigQueryImpl) service).setBigQueryReadClient(mockReadClient);
+
+ JobId jobId = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
+ // initialRowOffset = 1, maxResults = 2 -> only 1 more row allowed
+ BigQueryImpl.ArrowQueryPageFetcher fetcher =
+ new BigQueryImpl.ArrowQueryPageFetcher(
+ jobId,
+ bqSchema,
+ schemaBytes,
+ arrowSchema,
+ options,
+ /* initialRowOffset= */ 1L,
+ /* maxResults= */ 2L,
+ Collections.emptyMap());
+
+ Page page = fetcher.getNextPage();
+ assertNotNull(page);
+ List rows = ImmutableList.copyOf(page.getValues());
+ assertEquals(1, rows.size());
+ assertEquals("1", rows.get(0).get(0).getStringValue());
+ assertFalse(page.hasNextPage());
+ assertNull(page.getNextPageToken());
+ assertNull(page.getNextPage());
+ }
+
+ @Test
+ void testGetNextPage_missingLocationThrowsException() {
+ BigQueryOptions options = BigQueryOptions.newBuilder().setProjectId(PROJECT).build();
+ JobId jobId = JobId.of(PROJECT, JOB); // No location set on JobId or options
+
+ BigQueryImpl.ArrowQueryPageFetcher fetcher =
+ new BigQueryImpl.ArrowQueryPageFetcher(
+ jobId,
+ bqSchema,
+ schemaBytes,
+ arrowSchema,
+ options,
+ /* initialRowOffset= */ 0L,
+ /* maxResults= */ 10L,
+ Collections.emptyMap());
+
+ BigQueryException thrown = assertThrows(BigQueryException.class, fetcher::getNextPage);
+ assertTrue(thrown.getMessage().contains("Location must be specified"));
+ }
+
+ @Test
+ void testSerialization() throws Exception {
+ BigQueryOptions options =
+ BigQueryOptions.newBuilder().setProjectId(PROJECT).setLocation(LOCATION).build();
+ JobId jobId = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
+
+ BigQueryImpl.ArrowQueryPageFetcher fetcher =
+ new BigQueryImpl.ArrowQueryPageFetcher(
+ jobId,
+ bqSchema,
+ schemaBytes,
+ arrowSchema,
+ options,
+ /* initialRowOffset= */ 0L,
+ /* maxResults= */ 10L,
+ Collections.emptyMap());
+
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ try (ObjectOutputStream oos = new ObjectOutputStream(baos)) {
+ oos.writeObject(fetcher);
+ }
+
+ BigQueryImpl.ArrowQueryPageFetcher deserializedFetcher;
+ try (ObjectInputStream ois =
+ new ObjectInputStream(new ByteArrayInputStream(baos.toByteArray()))) {
+ deserializedFetcher = (BigQueryImpl.ArrowQueryPageFetcher) ois.readObject();
+ }
+
+ assertNotNull(deserializedFetcher);
+ }
+}
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..0d57d4cb3546 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;
@@ -45,6 +46,8 @@
import com.google.api.gax.paging.Page;
import com.google.api.gax.retrying.ResultRetryAlgorithm;
import com.google.api.gax.retrying.TimedAttemptSettings;
+import com.google.api.gax.rpc.ServerStream;
+import com.google.api.gax.rpc.ServerStreamingCallable;
import com.google.api.services.bigquery.model.ErrorProto;
import com.google.api.services.bigquery.model.GetQueryResultsResponse;
import com.google.api.services.bigquery.model.JobConfigurationQuery;
@@ -70,6 +73,9 @@
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.cloud.bigquery.storage.v1.ReadRowsRequest;
+import com.google.cloud.bigquery.storage.v1.ReadRowsResponse;
import com.google.common.base.Function;
import com.google.common.base.Supplier;
import com.google.common.collect.ImmutableList;
@@ -77,13 +83,27 @@
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
+import com.google.common.io.BaseEncoding;
+import com.google.protobuf.ByteString;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.net.ConnectException;
import java.net.UnknownHostException;
+import java.nio.channels.Channels;
import java.util.Collections;
import java.util.List;
import java.util.Map;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.memory.RootAllocator;
+import org.apache.arrow.vector.BigIntVector;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.VectorUnloader;
+import org.apache.arrow.vector.ipc.WriteChannel;
+import org.apache.arrow.vector.ipc.message.ArrowRecordBatch;
+import org.apache.arrow.vector.ipc.message.MessageSerializer;
+import org.apache.arrow.vector.types.pojo.ArrowType;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
@@ -2904,6 +2924,349 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException
assertEquals((Long) 1000L, requestPb.getTimeoutMs());
}
+ @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 testQueryArrowResultsFormatUnsupportedConfiguration() {
+ QueryJobConfiguration config =
+ QueryJobConfiguration.newBuilder("SELECT 1")
+ .setQueryResultsFormat(QueryResultsFormat.ARROW)
+ .setDestinationTable(TableId.of("dataset", "table"))
+ .build();
+ bigquery = options.getService();
+ IllegalArgumentException exception =
+ assertThrows(IllegalArgumentException.class, () -> bigquery.query(config));
+ assertTrue(
+ exception
+ .getMessage()
+ .contains("Arrow results format is only supported for fast query path execution"));
+ }
+
+ @Test
+ void testQueryWithArrowFormatFastPath() throws IOException, InterruptedException {
+ org.apache.arrow.vector.types.pojo.Schema arrowSchema =
+ new org.apache.arrow.vector.types.pojo.Schema(
+ ImmutableList.of(
+ org.apache.arrow.vector.types.pojo.Field.nullable(
+ "id", new ArrowType.Int(64, true))));
+
+ byte[] schemaBytes;
+ try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+ MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema);
+ schemaBytes = out.toByteArray();
+ }
+
+ com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
+ new com.google.api.services.bigquery.model.QueryResponse()
+ .setQueryId("q-arrow-1")
+ .setJobComplete(true)
+ .setTotalRows(BigInteger.ONE)
+ .setArrowSchema(
+ new com.google.api.services.bigquery.model.ArrowSchema()
+ .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes)));
+
+ ArgumentCaptor requestPbCapture = ArgumentCaptor.forClass(QueryRequest.class);
+ when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture()))
+ .thenReturn(queryResponsePb);
+
+ QueryJobConfiguration config =
+ QueryJobConfiguration.newBuilder("SELECT 1 as id")
+ .setQueryResultsFormat(QueryResultsFormat.ARROW)
+ .build();
+ bigquery = options.getService();
+ TableResult result = bigquery.query(config);
+ assertNotNull(result);
+ assertEquals("q-arrow-1", result.getQueryId());
+ assertNotNull(result.getSchema());
+ assertEquals(1, result.getSchema().getFields().size());
+ assertEquals("id", result.getSchema().getFields().get(0).getName());
+
+ QueryRequest requestPb = requestPbCapture.getValue();
+ assertEquals("ARROW", requestPb.getQueryResultsFormat());
+ }
+
+ @Test
+ void testQueryWithArrowFormatMultiplePages() throws IOException, InterruptedException {
+ org.apache.arrow.vector.types.pojo.Schema arrowSchema =
+ new org.apache.arrow.vector.types.pojo.Schema(
+ ImmutableList.of(
+ org.apache.arrow.vector.types.pojo.Field.nullable(
+ "id", new ArrowType.Int(64, true))));
+
+ byte[] schemaBytes;
+ try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+ MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema);
+ schemaBytes = out.toByteArray();
+ }
+
+ // Prepare page 2 Arrow batch for streaming
+ byte[] page2BatchBytes;
+ try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+ BigIntVector idVector = new BigIntVector("id", allocator);
+ idVector.allocateNew(1);
+ idVector.set(0, 2L);
+ idVector.setValueCount(1);
+ try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(idVector))) {
+ VectorUnloader unloader = new VectorUnloader(root);
+ try (ArrowRecordBatch recordBatch = unloader.getRecordBatch();
+ ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+ WriteChannel channel = new WriteChannel(Channels.newChannel(out));
+ MessageSerializer.serialize(channel, recordBatch);
+ page2BatchBytes = out.toByteArray();
+ }
+ } finally {
+ idVector.close();
+ }
+ }
+
+ JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
+ com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
+ new com.google.api.services.bigquery.model.QueryResponse()
+ .setQueryId("q-arrow-multipage")
+ .setJobComplete(true)
+ .setJobReference(queryJob.toPb())
+ .setTotalRows(BigInteger.valueOf(2L))
+ .setPageToken("1")
+ .setArrowSchema(
+ new com.google.api.services.bigquery.model.ArrowSchema()
+ .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes)));
+
+ when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class)))
+ .thenReturn(queryResponsePb);
+
+ // Mock BigQueryReadClient for page 2
+ @SuppressWarnings("unchecked")
+ ServerStreamingCallable mockCallable =
+ mock(ServerStreamingCallable.class, withSettings().withoutAnnotations());
+ @SuppressWarnings("unchecked")
+ ServerStream mockServerStream =
+ mock(ServerStream.class, withSettings().withoutAnnotations());
+ when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream);
+
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch =
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder()
+ .setSerializedRecordBatch(ByteString.copyFrom(page2BatchBytes))
+ .build();
+ ReadRowsResponse streamResponse =
+ ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build();
+ when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator());
+
+ BigQueryReadClient mockReadClient =
+ mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
+ when(mockReadClient.readRowsCallable()).thenReturn(mockCallable);
+
+ bigquery = options.getService();
+ ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient);
+
+ QueryJobConfiguration config =
+ QueryJobConfiguration.newBuilder("SELECT id FROM test")
+ .setQueryResultsFormat(QueryResultsFormat.ARROW)
+ .build();
+ TableResult result = bigquery.query(config);
+ assertNotNull(result);
+ assertEquals("q-arrow-multipage", result.getQueryId());
+ assertTrue(result.hasNextPage());
+ assertEquals("1", result.getNextPageToken());
+
+ Page page2 = result.getNextPage();
+ assertNotNull(page2);
+ List page2Rows = ImmutableList.copyOf(page2.getValues());
+ assertEquals(1, page2Rows.size());
+ assertEquals("2", page2Rows.get(0).get(0).getStringValue());
+
+ verify(mockCallable).call(any(ReadRowsRequest.class));
+ }
+
+ @Test
+ void testQueryWithArrowFormatMultiplePagesWithMaxResults()
+ throws IOException, InterruptedException {
+ org.apache.arrow.vector.types.pojo.Schema arrowSchema =
+ new org.apache.arrow.vector.types.pojo.Schema(
+ ImmutableList.of(
+ org.apache.arrow.vector.types.pojo.Field.nullable(
+ "id", new ArrowType.Int(64, true))));
+
+ byte[] schemaBytes;
+ try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+ MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema);
+ schemaBytes = out.toByteArray();
+ }
+
+ // Prepare page 2 Arrow batch for streaming with 2 rows
+ byte[] page2BatchBytes;
+ try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
+ BigIntVector idVector = new BigIntVector("id", allocator);
+ idVector.allocateNew(2);
+ idVector.set(0, 2L);
+ idVector.set(1, 3L);
+ idVector.setValueCount(2);
+ try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(idVector))) {
+ VectorUnloader unloader = new VectorUnloader(root);
+ try (ArrowRecordBatch recordBatch = unloader.getRecordBatch();
+ ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+ WriteChannel channel = new WriteChannel(Channels.newChannel(out));
+ MessageSerializer.serialize(channel, recordBatch);
+ page2BatchBytes = out.toByteArray();
+ }
+ } finally {
+ idVector.close();
+ }
+ }
+
+ JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
+ com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
+ new com.google.api.services.bigquery.model.QueryResponse()
+ .setQueryId("q-arrow-multipage-maxresults")
+ .setJobComplete(true)
+ .setJobReference(queryJob.toPb())
+ .setTotalRows(BigInteger.valueOf(3L))
+ .setPageToken("1")
+ .setArrowSchema(
+ new com.google.api.services.bigquery.model.ArrowSchema()
+ .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes)));
+
+ when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class)))
+ .thenReturn(queryResponsePb);
+
+ // Mock BigQueryReadClient for page 2
+ @SuppressWarnings("unchecked")
+ ServerStreamingCallable mockCallable =
+ mock(ServerStreamingCallable.class, withSettings().withoutAnnotations());
+ @SuppressWarnings("unchecked")
+ ServerStream mockServerStream =
+ mock(ServerStream.class, withSettings().withoutAnnotations());
+ when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream);
+
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch =
+ com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder()
+ .setSerializedRecordBatch(ByteString.copyFrom(page2BatchBytes))
+ .build();
+ ReadRowsResponse streamResponse =
+ ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build();
+ when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator());
+
+ BigQueryReadClient mockReadClient =
+ mock(BigQueryReadClient.class, withSettings().withoutAnnotations());
+ when(mockReadClient.readRowsCallable()).thenReturn(mockCallable);
+
+ bigquery = options.getService();
+ ((BigQueryImpl) bigquery).setBigQueryReadClient(mockReadClient);
+
+ QueryJobConfiguration config =
+ QueryJobConfiguration.newBuilder("SELECT id FROM test")
+ .setQueryResultsFormat(QueryResultsFormat.ARROW)
+ .setMaxResults(2L)
+ .build();
+ TableResult result = bigquery.query(config);
+ assertNotNull(result);
+ assertTrue(result.hasNextPage());
+
+ Page page2 = result.getNextPage();
+ assertNotNull(page2);
+ List page2Rows = ImmutableList.copyOf(page2.getValues());
+ // Since maxResults is 2 and initialRowOffset is 1, page2 should only contain 1 row even though
+ // stream returned 2 rows
+ assertEquals(1, page2Rows.size());
+ assertEquals("2", page2Rows.get(0).get(0).getStringValue());
+ // Since totalRowsReturned == maxResults, hasNextPage must be false
+ assertFalse(page2.hasNextPage());
+ assertNull(page2.getNextPage());
+
+ // When maxResults is 1, initialRowOffset (1) already reaches maxResults, so hasNextPage is
+ // false immediately
+ QueryJobConfiguration configMax1 =
+ QueryJobConfiguration.newBuilder("SELECT id FROM test")
+ .setQueryResultsFormat(QueryResultsFormat.ARROW)
+ .setMaxResults(1L)
+ .build();
+ TableResult resultMax1 = bigquery.query(configMax1);
+ assertNotNull(resultMax1);
+ assertFalse(resultMax1.hasNextPage());
+ assertNull(resultMax1.getNextPage());
+ }
+
+ @Test
+ void testArrowQueryPageFetcherSerialization() throws Exception {
+ org.apache.arrow.vector.types.pojo.Schema arrowSchema =
+ new org.apache.arrow.vector.types.pojo.Schema(
+ ImmutableList.of(
+ org.apache.arrow.vector.types.pojo.Field.nullable(
+ "id", new ArrowType.Int(64, true))));
+
+ byte[] schemaBytes;
+ try (ByteArrayOutputStream out = new ByteArrayOutputStream()) {
+ MessageSerializer.serialize(new WriteChannel(Channels.newChannel(out)), arrowSchema);
+ schemaBytes = out.toByteArray();
+ }
+
+ JobId queryJob = JobId.of(PROJECT, JOB).toBuilder().setLocation(LOCATION).build();
+ com.google.api.services.bigquery.model.QueryResponse queryResponsePb =
+ new com.google.api.services.bigquery.model.QueryResponse()
+ .setQueryId("q-arrow-multipage-ser")
+ .setJobComplete(true)
+ .setJobReference(queryJob.toPb())
+ .setTotalRows(BigInteger.valueOf(2L))
+ .setPageToken("1")
+ .setArrowSchema(
+ new com.google.api.services.bigquery.model.ArrowSchema()
+ .setSerializedSchema(BaseEncoding.base64().encode(schemaBytes)));
+
+ when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), any(QueryRequest.class)))
+ .thenReturn(queryResponsePb);
+
+ bigquery = options.getService();
+
+ QueryJobConfiguration config =
+ QueryJobConfiguration.newBuilder("SELECT id FROM test")
+ .setQueryResultsFormat(QueryResultsFormat.ARROW)
+ .build();
+ TableResult result = bigquery.query(config);
+ assertNotNull(result);
+ assertTrue(result.hasNextPage());
+ assertEquals("1", result.getNextPageToken());
+
+ // Serialize and deserialize TableResult
+ ByteArrayOutputStream baos = new ByteArrayOutputStream();
+ try (java.io.ObjectOutputStream oos = new java.io.ObjectOutputStream(baos)) {
+ oos.writeObject(result);
+ }
+
+ TableResult deserializedResult;
+ try (java.io.ObjectInputStream ois =
+ new java.io.ObjectInputStream(new ByteArrayInputStream(baos.toByteArray()))) {
+ deserializedResult = (TableResult) ois.readObject();
+ }
+
+ assertNotNull(deserializedResult);
+ assertEquals(result.getSchema(), deserializedResult.getSchema());
+ assertEquals(result.getTotalRows(), deserializedResult.getTotalRows());
+ assertEquals(result.getQueryId(), deserializedResult.getQueryId());
+ assertEquals("1", deserializedResult.getNextPageToken());
+ assertTrue(deserializedResult.hasNextPage());
+ }
+
@Test
void testGetQueryResults() throws IOException {
JobId queryJob = JobId.of(JOB);
diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
index 9744eebc2d81..f90a72bd9795 100644
--- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
+++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java
@@ -45,6 +45,7 @@
import com.google.cloud.bigquery.Acl.DatasetAclEntity;
import com.google.cloud.bigquery.Acl.Expr;
import com.google.cloud.bigquery.Acl.User;
+import com.google.cloud.bigquery.ArrowQueryResult;
import com.google.cloud.bigquery.BigQuery;
import com.google.cloud.bigquery.BigQuery.DatasetField;
import com.google.cloud.bigquery.BigQuery.DatasetListOption;
@@ -118,6 +119,7 @@
import com.google.cloud.bigquery.QueryJobConfiguration.JobCreationMode;
import com.google.cloud.bigquery.QueryJobConfiguration.Priority;
import com.google.cloud.bigquery.QueryParameterValue;
+import com.google.cloud.bigquery.QueryResultsFormat;
import com.google.cloud.bigquery.Range;
import com.google.cloud.bigquery.RangePartitioning;
import com.google.cloud.bigquery.Routine;
@@ -203,6 +205,7 @@
import java.util.concurrent.TimeoutException;
import java.util.logging.Level;
import java.util.logging.Logger;
+import org.apache.arrow.vector.VectorSchemaRoot;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
@@ -7497,6 +7500,50 @@ void testQueryWithTimeout() throws InterruptedException {
assertTrue(millis < 1_000_000 * 2);
}
+ @Test
+ void testQueryResultsFormatArrow() throws InterruptedException {
+ RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
+ BigQuery bigQuery = bigqueryHelper.getOptions().getService();
+ String query = "SELECT 1 as id, 'hello' as name, TIMESTAMP('2026-08-10T12:00:00Z') as ts";
+ QueryJobConfiguration config =
+ QueryJobConfiguration.newBuilder(query)
+ .setQueryResultsFormat(QueryResultsFormat.ARROW)
+ .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL)
+ .build();
+ try (ArrowQueryResult result = bigQuery.queryArrow(config)) {
+ assertNotNull(result);
+ int batchCount = 0;
+ long totalRows = 0;
+ for (VectorSchemaRoot root : result) {
+ batchCount++;
+ totalRows += root.getRowCount();
+ assertEquals(1, root.getRowCount());
+ }
+ assertTrue(batchCount > 0);
+ assertEquals(1, totalRows);
+ }
+ }
+
+ @Test
+ void testQueryResultsFormatArrowMultiPage() throws InterruptedException {
+ RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();
+ BigQuery bigQuery = bigqueryHelper.getOptions().getService();
+ String query = "SELECT x FROM UNNEST(GENERATE_ARRAY(1, 15000)) AS x";
+ QueryJobConfiguration config =
+ QueryJobConfiguration.newBuilder(query)
+ .setQueryResultsFormat(QueryResultsFormat.ARROW)
+ .setJobCreationMode(JobCreationMode.JOB_CREATION_OPTIONAL)
+ .build();
+ try (ArrowQueryResult result = bigQuery.queryArrow(config)) {
+ assertNotNull(result);
+ long totalRows = 0;
+ for (VectorSchemaRoot root : result) {
+ totalRows += root.getRowCount();
+ }
+ assertEquals(15000, totalRows);
+ }
+ }
+
@Test
void testUniverseDomainWithInvalidUniverseDomain() {
RemoteBigQueryHelper bigqueryHelper = RemoteBigQueryHelper.create();