diff --git a/java-bigquery/google-cloud-bigquery/pom.xml b/java-bigquery/google-cloud-bigquery/pom.xml index 6716595c1995..12e2e97ef0c1 100644 --- a/java-bigquery/google-cloud-bigquery/pom.xml +++ b/java-bigquery/google-cloud-bigquery/pom.xml @@ -120,6 +120,15 @@ arrow-memory-netty + + com.google.api + gax-grpc + + + io.grpc + grpc-api + + com.google.errorprone error_prone_annotations diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java index 3def579e2abf..a608b0f445bd 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java @@ -50,6 +50,16 @@ private static class AllocatorHolder { private static final BufferAllocator ALLOCATOR = new RootAllocator(Long.MAX_VALUE); } + /** + * Creates a new child buffer allocator with the given name. + * + * @param name the child allocator name + * @return a new child buffer allocator + */ + static BufferAllocator createChildAllocator(String name) { + return AllocatorHolder.ALLOCATOR.newChildAllocator(name, 0, Long.MAX_VALUE); + } + /** * Instantiates a new {@link VectorSchemaRoot} for the given Arrow schema using vectors allocated * from the provided child allocator, ensuring LIFO cleanup if an error occurs during diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResult.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResult.java new file mode 100644 index 000000000000..13c7f2e11bb8 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResult.java @@ -0,0 +1,60 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +import com.google.api.core.BetaApi; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.Schema; + +/** + * [Beta] A query result container providing zero-copy access to Apache Arrow {@link + * VectorSchemaRoot} batches. + * + *

Implementations manage direct off-heap native memory buffers. Callers must invoke {@link + * #close()} (idiomatically via a {@code try-with-resources} block) to ensure native allocations and + * underlying gRPC streaming channels are deterministically released. + */ +@BetaApi +public interface ArrowQueryResult extends AutoCloseable, Iterable { + + /** Returns the Apache Arrow schema of the result vectors. */ + Schema getArrowSchema(); + + /** + * Returns the job ID associated with the query execution, or {@code null} if no job was created + * (e.g. when optional job creation was used). + */ + JobId getJobId(); + + /** Returns the query ID associated with the query execution, or {@code null} if unavailable. */ + String getQueryId(); + + /** + * Returns the reason a job was created when optional job creation was requested, or {@code null} + * if no job was created or if the query ran via the fallback path. + */ + JobCreationReason getJobCreationReason(); + + /** Returns the total number of rows across all batches if known, or {@code -1} if unknown. */ + long getTotalRows(); + + /** + * Releases underlying direct off-heap memory allocations and closes any active stream channels. + */ + @Override + void close(); +} diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java new file mode 100644 index 000000000000..6637f1483fbf --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowQueryResultImpl.java @@ -0,0 +1,329 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +import com.google.api.gax.rpc.ServerStream; +import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; +import com.google.cloud.bigquery.storage.v1.ReadRowsRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; +import com.google.cloud.bigquery.storage.v1.ReadSession; +import java.io.IOException; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.VectorLoader; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.ipc.ReadChannel; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.ipc.message.MessageSerializer; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; + +/** + * Implementation of {@link ArrowQueryResult} that provides zero-copy streaming of Apache Arrow + * {@link VectorSchemaRoot} batches across initial REST response and subsequent gRPC stream. + */ +class ArrowQueryResultImpl implements ArrowQueryResult { + + private final Schema arrowSchema; + private final JobId jobId; + private final String queryId; + private final JobCreationReason jobCreationReason; + private final long totalRows; + private final byte[] initialRecordBatchBytes; + private final String streamName; + private final BigQueryReadClient readClient; + + private final BufferAllocator allocator; + private final VectorSchemaRoot root; + private final VectorLoader loader; + + private final Object lock = new Object(); + private boolean closed = false; + private boolean iteratorCreated = false; + private ServerStream serverStream; + + ArrowQueryResultImpl( + Object arrowSchema, + JobId jobId, + long totalRows, + byte[] initialRecordBatchBytes, + String streamName, + BigQueryReadClient readClient) { + this( + arrowSchema, + jobId, + /* queryId= */ null, + /* jobCreationReason= */ null, + totalRows, + initialRecordBatchBytes, + streamName, + readClient); + } + + ArrowQueryResultImpl( + Object arrowSchema, + JobId jobId, + String queryId, + JobCreationReason jobCreationReason, + long totalRows, + byte[] initialRecordBatchBytes, + String streamName, + BigQueryReadClient readClient) { + if (arrowSchema instanceof Schema) { + this.arrowSchema = (Schema) arrowSchema; + } else { + this.arrowSchema = null; + } + this.jobId = jobId; + this.queryId = queryId; + this.jobCreationReason = jobCreationReason; + this.totalRows = totalRows; + this.initialRecordBatchBytes = initialRecordBatchBytes; + this.streamName = streamName; + this.readClient = readClient; + + if (this.arrowSchema != null) { + this.allocator = ArrowDeserializer.createChildAllocator("ArrowQueryResult"); + List vectors = ArrowPojoUtils.createVectors(this.arrowSchema, this.allocator); + this.root = new VectorSchemaRoot(vectors); + this.loader = new VectorLoader(this.root); + } else { + this.allocator = null; + this.root = null; + this.loader = null; + } + } + + static ArrowQueryResultImpl fromReadSession( + ReadSession readSession, JobId jobId, BigQueryReadClient readClient) { + Schema pojoSchema = null; + if (readSession.hasArrowSchema()) { + try { + pojoSchema = + (Schema) + ArrowDeserializer.deserializeSchema( + readSession.getArrowSchema().getSerializedSchema().toByteArray()); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow schema from ReadSession", e); + } + } + String streamName = + readSession.getStreamsCount() > 0 ? readSession.getStreams(0).getName() : null; + return new ArrowQueryResultImpl( + pojoSchema, + jobId, + /* queryId= */ null, + /* jobCreationReason= */ null, + /* totalRows= */ -1L, + /* initialRecordBatchBytes= */ null, + streamName, + readClient); + } + + @Override + public Schema getArrowSchema() { + return arrowSchema; + } + + @Override + public JobId getJobId() { + return jobId; + } + + @Override + public String getQueryId() { + return queryId; + } + + @Override + public JobCreationReason getJobCreationReason() { + return jobCreationReason; + } + + @Override + public long getTotalRows() { + return totalRows; + } + + @Override + public Iterator iterator() { + synchronized (lock) { + checkNotClosed(); + if (iteratorCreated) { + throw new IllegalStateException("ArrowQueryResult can only be iterated once"); + } + iteratorCreated = true; + return new VectorBatchIterator(); + } + } + + @Override + public void close() { + synchronized (lock) { + if (closed) { + return; + } + closed = true; + Throwable firstException = null; + + if (serverStream != null) { + try { + serverStream.cancel(); + } catch (Throwable t) { + firstException = t; + } + } + if (root != null) { + try { + root.close(); + } catch (Throwable t) { + if (firstException == null) { + firstException = t; + } else { + firstException.addSuppressed(t); + } + } + } + if (allocator != null) { + try { + allocator.close(); + } catch (Throwable t) { + if (firstException == null) { + firstException = t; + } else { + firstException.addSuppressed(t); + } + } + } + if (firstException instanceof RuntimeException) { + throw (RuntimeException) firstException; + } else if (firstException != null) { + throw new RuntimeException("Failed to close Arrow resources", firstException); + } + } + } + + private void checkNotClosed() { + if (closed) { + throw new IllegalStateException("ArrowQueryResult has already been closed"); + } + } + + private final class VectorBatchIterator implements Iterator { + private boolean yieldedInitialBatch = false; + private Iterator streamIterator = null; + private boolean streamInitialized = false; + private long totalRowsYielded = 0; + + @Override + public boolean hasNext() { + synchronized (lock) { + if (closed) { + return false; + } + if (!yieldedInitialBatch + && initialRecordBatchBytes != null + && initialRecordBatchBytes.length > 0) { + return true; + } + ensureStreamInitialized(); + if (streamIterator == null) { + return false; + } + return streamIterator.hasNext(); + } + } + + @Override + public VectorSchemaRoot next() { + synchronized (lock) { + checkNotClosed(); + + // 1. Yield initial batch from REST response if present + if (!yieldedInitialBatch + && initialRecordBatchBytes != null + && initialRecordBatchBytes.length > 0) { + yieldedInitialBatch = true; + try { + loadBatchBytes(initialRecordBatchBytes); + totalRowsYielded += root.getRowCount(); + return root; + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load initial Arrow record batch", e); + } + } + yieldedInitialBatch = true; + + // 2. Stream subsequent batches from gRPC + ensureStreamInitialized(); + if (streamIterator == null || !streamIterator.hasNext()) { + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } + + while (streamIterator.hasNext()) { + ReadRowsResponse response = streamIterator.next(); + if (response.hasArrowRecordBatch()) { + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = + response.getArrowRecordBatch(); + try { + loadBatchBytes(batch.getSerializedRecordBatch().toByteArray()); + totalRowsYielded += root.getRowCount(); + return root; + } catch (IOException e) { + throw new BigQueryException(0, "Failed to load streaming Arrow record batch", e); + } + } + } + throw new NoSuchElementException("No more Arrow batches available in query stream."); + } + } + + private void ensureStreamInitialized() { + if (streamInitialized) { + return; + } + streamInitialized = true; + if (totalRows >= 0 && totalRowsYielded >= totalRows && yieldedInitialBatch) { + return; + } + if (streamName != null && readClient != null) { + ReadRowsRequest request = + ReadRowsRequest.newBuilder() + .setReadStream(streamName) + .setOffset(totalRowsYielded) + .build(); + serverStream = readClient.readRowsCallable().call(request); + streamIterator = serverStream.iterator(); + } + } + + private void loadBatchBytes(byte[] bytes) throws IOException { + try (ByteArrayReadableSeekableByteChannel byteChannel = + new ByteArrayReadableSeekableByteChannel(bytes); + ReadChannel readChannel = new ReadChannel(byteChannel); + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch(readChannel, allocator)) { + if (deserializedBatch != null) { + loader.load(deserializedBatch); + } + } + } + } +} diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQuery.java index f92cb96c377a..b6895f407852 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 @@ -1636,6 +1636,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 2ad09c33d7cb..962756e5a04f 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java @@ -18,10 +18,12 @@ import static com.google.cloud.bigquery.PolicyHelper.convertFromApiPolicy; import static com.google.cloud.bigquery.PolicyHelper.convertToApiPolicy; import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkNotNull; import static java.net.HttpURLConnection.HTTP_NOT_FOUND; import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; +import com.google.api.gax.core.FixedCredentialsProvider; import com.google.api.gax.paging.Page; import com.google.api.services.bigquery.model.ErrorProto; import com.google.api.services.bigquery.model.GetQueryResultsResponse; @@ -43,6 +45,11 @@ import com.google.cloud.bigquery.InsertAllRequest.RowToInsert; 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; @@ -52,14 +59,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; @@ -264,6 +276,76 @@ public Page getNextPage() { } } + private final ReentrantLock readClientLock = new ReentrantLock(); + private transient BigQueryReadClient bqReadClient; + + /** + * Lazily creates or retrieves the shared {@link BigQueryReadClient} instance used for streaming + * Arrow query results, reusing credentials and channel configuration from this {@link + * BigQueryImpl}. + * + * @return the active BigQueryReadClient instance + * @throws IOException if initializing the storage read client fails + */ + BigQueryReadClient getBigQueryReadClient() throws IOException { + readClientLock.lock(); + try { + if (bqReadClient == null) { + BigQueryReadSettings.Builder settingsBuilder = BigQueryReadSettings.newBuilder(); + configureReadSettings(settingsBuilder, getOptions()); + bqReadClient = BigQueryReadClient.create(settingsBuilder.build()); + } + return bqReadClient; + } finally { + readClientLock.unlock(); + } + } + + /** + * Configures a {@link BigQueryReadSettings.Builder} with credentials, universe domain, custom + * endpoint, and transport settings mapped from the given {@link BigQueryOptions}. + * + * @param settingsBuilder the builder to configure + * @param options the source BigQueryOptions + */ + private static void configureReadSettings( + BigQueryReadSettings.Builder settingsBuilder, BigQueryOptions options) { + if (options.getCredentials() != null) { + settingsBuilder.setCredentialsProvider( + FixedCredentialsProvider.create(options.getCredentials())); + } + if (options.getUniverseDomain() != null) { + settingsBuilder.setUniverseDomain(options.getUniverseDomain()); + } + if (options.getHost() != null) { + String host = options.getHost(); + String target = host; + if (target.contains("://")) { + target = URI.create(target).getAuthority(); + } + HostAndPort hostAndPort = HostAndPort.fromString(target); + String endpointHost = hostAndPort.getHost(); + if (endpointHost.contains("bigquery.googleapis.com")) { + endpointHost = + endpointHost.replace("bigquery.googleapis.com", "bigquerystorage.googleapis.com"); + } else if (endpointHost.contains("bigquery.private.googleapis.com")) { + endpointHost = + endpointHost.replace( + "bigquery.private.googleapis.com", "bigquerystorage.private.googleapis.com"); + } else if (endpointHost.startsWith("bigquery.")) { + endpointHost = endpointHost.replaceFirst("^bigquery\\.", "bigquerystorage."); + } + int port = hostAndPort.getPortOrDefault(443); + settingsBuilder.setEndpoint(endpointHost + ":" + port); + if (endpointHost.contains("localhost") || endpointHost.contains("127.0.0.1")) { + settingsBuilder.setTransportChannelProvider( + BigQueryReadSettings.defaultGrpcTransportProviderBuilder() + .setChannelConfigurator(ManagedChannelBuilder::usePlaintext) + .build()); + } + } + } + private final HttpBigQueryRpc bigQueryRpc; private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG = @@ -2153,6 +2235,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 = @@ -2207,6 +2294,7 @@ && getOptions().getOpenTelemetryTracer() != null) { return queryRpc(projectId, content, options); } + return create(JobInfo.of(jobId, configuration), options); } finally { if (querySpan != null) { @@ -2215,6 +2303,229 @@ && getOptions().getOpenTelemetryTracer() != null) { } } + @Override + public ArrowQueryResult queryArrow(QueryJobConfiguration configuration, JobOption... options) + throws InterruptedException, JobException { + return queryArrow(configuration, (JobId) null, options); + } + + @Override + public ArrowQueryResult queryArrow( + QueryJobConfiguration configuration, JobId jobId, JobOption... options) + throws InterruptedException, JobException { + return queryArrowWithTimeout(configuration, jobId, null, options); + } + + private ArrowQueryResult queryArrowWithTimeout( + QueryJobConfiguration configuration, JobId jobId, Long timeoutMs, JobOption... options) + throws InterruptedException, JobException { + checkNotNull(configuration, "configuration cannot be null"); + Span querySpan = null; + if (getOptions().isOpenTelemetryTracingEnabled() + && getOptions().getOpenTelemetryTracer() != null) { + querySpan = + getOptions() + .getOpenTelemetryTracer() + .spanBuilder("com.google.cloud.bigquery.BigQuery.queryArrowWithTimeout") + .setAllAttributes(jobId != null ? jobId.getOtelAttributes() : Attributes.empty()) + .setAllAttributes(otelAttributesFromOptions(options)) + .startSpan(); + } + try (Scope queryScope = querySpan != null ? querySpan.makeCurrent() : null) { + QueryJobConfiguration arrowConfig = configuration; + if (arrowConfig.getQueryResultsFormat() != QueryResultsFormat.ARROW) { + arrowConfig = + configuration.toBuilder().setQueryResultsFormat(QueryResultsFormat.ARROW).build(); + } + if (arrowConfig.getJobCreationMode() == null) { + arrowConfig = + arrowConfig.toBuilder() + .setJobCreationMode(QueryJobConfiguration.JobCreationMode.JOB_CREATION_OPTIONAL) + .build(); + } + + QueryRequestInfo requestInfo = + new QueryRequestInfo(arrowConfig, getOptions().getDataFormatOptions()); + + boolean useFastPath = + requestInfo.isFastQuerySupported() + && arrowConfig.getDestinationTable() == null + && (jobId == null || jobId.getJob() == null); + + if (useFastPath) { + String projectId = + jobId != null && jobId.getProject() != null + ? jobId.getProject() + : getOptions().getProjectId(); + QueryRequest content = requestInfo.toPb(); + if (jobId != null && jobId.getLocation() != null) { + content.setLocation(jobId.getLocation()); + } else if (getOptions().getLocation() != null) { + content.setLocation(getOptions().getLocation()); + } + if (timeoutMs != null) { + content.setTimeoutMs(timeoutMs); + } + + Map optionsMap = optionMap(options); + com.google.api.services.bigquery.model.QueryResponse results; + try { + results = + BigQueryRetryHelper.runWithRetries( + new Callable() { + @Override + public com.google.api.services.bigquery.model.QueryResponse call() + throws IOException { + return bigQueryRpc.queryRpcSkipExceptionTranslation(projectId, content); + } + }, + getOptions().getRetrySettings(), + getOptions().getResultRetryAlgorithm(), + getOptions().getClock(), + DEFAULT_RETRY_CONFIG, + getOptions().isOpenTelemetryTracingEnabled(), + getOptions().getOpenTelemetryTracer()); + } catch (BigQueryRetryHelper.BigQueryRetryHelperException e) { + throw BigQueryException.translateAndThrow(e); + } + + if (results.getErrors() != null) { + List bigQueryErrors = + Lists.transform(results.getErrors(), BigQueryError.FROM_PB_FUNCTION); + throw new BigQueryException(bigQueryErrors); + } + + JobId actualJobId = + results.getJobReference() != null ? JobId.fromPb(results.getJobReference()) : jobId; + + Object arrowSchema = null; + if (results.getArrowSchema() != null) { + try { + arrowSchema = + ArrowDeserializer.deserializeSchema( + results.getArrowSchema().decodeSerializedSchema()); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e); + } + } + + long numRows = -1L; + if (results.getNumDmlAffectedRows() != null) { + numRows = results.getNumDmlAffectedRows(); + } else if (results.getTotalRows() != null) { + numRows = results.getTotalRows().longValue(); + } + + byte[] initialBatchBytes = null; + if (results.getArrowRecordBatch() != null + && results.getArrowRecordBatch().getSerializedRecordBatch() != null) { + initialBatchBytes = results.getArrowRecordBatch().decodeSerializedRecordBatch(); + } + + String streamName = null; + if (actualJobId != null && actualJobId.getJob() != null) { + String jobProject = + actualJobId.getProject() != null ? actualJobId.getProject() : projectId; + String jobLocation = + actualJobId.getLocation() != null + ? actualJobId.getLocation() + : (content.getLocation() != null + ? content.getLocation() + : getOptions().getLocation()); + if (jobLocation != null) { + streamName = + String.format( + "projects/%s/locations/%s/jobs/%s/streams/_default", + jobProject, jobLocation, actualJobId.getJob()); + } + } + + BigQueryReadClient client; + try { + client = getBigQueryReadClient(); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + } + + JobCreationReason jobCreationReason = + results.getJobCreationReason() != null + ? JobCreationReason.fromPb(results.getJobCreationReason()) + : null; + + return new ArrowQueryResultImpl( + arrowSchema, + actualJobId, + results.getQueryId(), + jobCreationReason, + numRows, + initialBatchBytes, + streamName, + client); + } else { + // Fallback path: jobs.insert + BigQuery Storage Read API + Job job = create(JobInfo.of(jobId, arrowConfig), options); + Job completedJob; + try { + completedJob = job.waitFor(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw e; + } + + if (completedJob.getStatus().getError() != null) { + throw new BigQueryException( + Collections.singletonList(completedJob.getStatus().getError())); + } + + TableId destinationTable = null; + if (completedJob.getConfiguration() instanceof QueryJobConfiguration) { + destinationTable = + ((QueryJobConfiguration) completedJob.getConfiguration()).getDestinationTable(); + } + if (destinationTable == null) { + destinationTable = arrowConfig.getDestinationTable(); + } + if (destinationTable == null) { + throw new BigQueryException(0, "Unable to resolve destination table for fallback query"); + } + + String destProject = + destinationTable.getProject() != null + ? destinationTable.getProject() + : (jobId != null && jobId.getProject() != null + ? jobId.getProject() + : getOptions().getProjectId()); + String parent = String.format("projects/%s", destProject); + String srcTable = + String.format( + "projects/%s/datasets/%s/tables/%s", + destProject, destinationTable.getDataset(), destinationTable.getTable()); + + BigQueryReadClient client; + try { + client = getBigQueryReadClient(); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to initialize BigQueryReadClient", e); + } + + CreateReadSessionRequest request = + CreateReadSessionRequest.newBuilder() + .setParent(parent) + .setReadSession( + ReadSession.newBuilder().setTable(srcTable).setDataFormat(DataFormat.ARROW)) + .setMaxStreamCount(1) + .build(); + ReadSession readSession = client.createReadSession(request); + + return ArrowQueryResultImpl.fromReadSession(readSession, completedJob.getJobId(), client); + } + } finally { + if (querySpan != null) { + querySpan.end(); + } + } + } + @Override public QueryResponse getQueryResults(JobId jobId, QueryResultsOption... options) { Map optionsMap = optionMap(options); diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java index c224bed5cc58..14d2c65fe78a 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/QueryRequestInfo.java @@ -46,6 +46,8 @@ final class QueryRequestInfo { private final DataFormatOptions formatOptions; private final String reservation; private final Long jobTimeoutMs; + private final QueryResultsFormat queryResultsFormat; + private final ArrowSerializationOptions arrowSerializationOptions; QueryRequestInfo( QueryJobConfiguration config, com.google.cloud.bigquery.DataFormatOptions dataFormatOptions) { @@ -63,9 +65,11 @@ final class QueryRequestInfo { this.useLegacySql = config.useLegacySql(); this.useQueryCache = config.useQueryCache(); this.jobCreationMode = config.getJobCreationMode(); - this.formatOptions = dataFormatOptions.toPb(); + this.formatOptions = dataFormatOptions != null ? dataFormatOptions.toPb() : null; this.reservation = config.getReservation(); this.jobTimeoutMs = config.getJobTimeoutMs(); + this.queryResultsFormat = config.getQueryResultsFormat(); + this.arrowSerializationOptions = config.getArrowSerializationOptions(); } /** @@ -142,6 +146,12 @@ QueryRequest toPb() { if (jobTimeoutMs != null) { request.setJobTimeoutMs(jobTimeoutMs); } + if (queryResultsFormat != null) { + request.setQueryResultsFormat(queryResultsFormat.toString()); + } + if (arrowSerializationOptions != null) { + request.setArrowSerializationOptions(arrowSerializationOptions.toPb()); + } return request; } @@ -161,7 +171,7 @@ public String toString() { .add("useQueryCache", useQueryCache) .add("useLegacySql", useLegacySql) .add("jobCreationMode", jobCreationMode) - .add("formatOptions", formatOptions.getUseInt64Timestamp()) + .add("formatOptions", formatOptions != null ? formatOptions.getUseInt64Timestamp() : null) .add("reservation", reservation) .add("jobTimeoutMs", jobTimeoutMs) .toString(); diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryResultTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryResultTest.java new file mode 100644 index 000000000000..e62155663478 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowQueryResultTest.java @@ -0,0 +1,354 @@ +/* + * Copyright 2026 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.google.cloud.bigquery; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.api.gax.rpc.ServerStream; +import com.google.api.gax.rpc.ServerStreamingCallable; +import com.google.cloud.bigquery.storage.v1.ArrowSchema; +import com.google.cloud.bigquery.storage.v1.BigQueryReadClient; +import com.google.cloud.bigquery.storage.v1.BigQueryReadSettings; +import com.google.cloud.bigquery.storage.v1.ReadRowsRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; +import com.google.cloud.bigquery.storage.v1.ReadSession; +import com.google.cloud.bigquery.storage.v1.ReadStream; +import com.google.cloud.bigquery.storage.v1.stub.EnhancedBigQueryReadStub; +import com.google.common.collect.ImmutableList; +import com.google.protobuf.ByteString; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.channels.Channels; +import java.nio.charset.StandardCharsets; +import java.util.Iterator; +import java.util.List; +import java.util.NoSuchElementException; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BigIntVector; +import org.apache.arrow.vector.VarCharVector; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.VectorUnloader; +import org.apache.arrow.vector.ipc.WriteChannel; +import org.apache.arrow.vector.ipc.message.ArrowRecordBatch; +import org.apache.arrow.vector.ipc.message.MessageSerializer; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +public class ArrowQueryResultTest { + + private BufferAllocator testAllocator; + + @BeforeEach + void setUp() { + testAllocator = new RootAllocator(Long.MAX_VALUE); + } + + @AfterEach + void tearDown() { + testAllocator.close(); + } + + private Schema createTestArrowSchema() { + Field idField = + new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), ImmutableList.of()); + Field nameField = + new Field("name", FieldType.nullable(new ArrowType.Utf8()), ImmutableList.of()); + return new Schema(ImmutableList.of(idField, nameField)); + } + + private byte[] createTestBatchBytes(List ids, List names) throws IOException { + BigIntVector idVector = new BigIntVector("id", testAllocator); + idVector.allocateNew(ids.size()); + for (int i = 0; i < ids.size(); i++) { + if (ids.get(i) != null) { + idVector.set(i, ids.get(i)); + } else { + idVector.setNull(i); + } + } + idVector.setValueCount(ids.size()); + + VarCharVector nameVector = new VarCharVector("name", testAllocator); + nameVector.allocateNew(names.size()); + for (int i = 0; i < names.size(); i++) { + if (names.get(i) != null) { + nameVector.set(i, names.get(i).getBytes(StandardCharsets.UTF_8)); + } else { + nameVector.setNull(i); + } + } + nameVector.setValueCount(names.size()); + + VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(idVector, nameVector)); + VectorUnloader unloader = new VectorUnloader(root); + ArrowRecordBatch recordBatch = unloader.getRecordBatch(); + + ByteArrayOutputStream out = new ByteArrayOutputStream(); + WriteChannel channel = new WriteChannel(Channels.newChannel(out)); + MessageSerializer.serialize(channel, recordBatch); + + recordBatch.close(); + root.close(); + return out.toByteArray(); + } + + private BigQueryReadClient createMockReadClient( + ServerStreamingCallable mockCallable) { + BigQueryReadClient mockClient = mock(BigQueryReadClient.class); + EnhancedBigQueryReadStub mockStub = mock(EnhancedBigQueryReadStub.class); + BigQueryReadSettings mockSettings = mock(BigQueryReadSettings.class); + try { + java.lang.reflect.Field settingsField = BigQueryReadClient.class.getDeclaredField("settings"); + settingsField.setAccessible(true); + settingsField.set(mockClient, mockSettings); + + java.lang.reflect.Field stubField = BigQueryReadClient.class.getDeclaredField("stub"); + stubField.setAccessible(true); + stubField.set(mockClient, mockStub); + } catch (ReflectiveOperationException e) { + throw new RuntimeException(e); + } + when(mockStub.readRowsCallable()).thenReturn(mockCallable); + return mockClient; + } + + @Test + void testSingleBatchIteration() throws IOException { + Schema arrowSchema = createTestArrowSchema(); + byte[] batchBytes = + createTestBatchBytes(ImmutableList.of(1L, 2L), ImmutableList.of("Alice", "Bob")); + JobId jobId = JobId.of("test-project", "job_123"); + + try (ArrowQueryResult result = + new ArrowQueryResultImpl(arrowSchema, jobId, 2L, batchBytes, null, null)) { + assertEquals(arrowSchema, result.getArrowSchema()); + assertEquals(jobId, result.getJobId()); + assertEquals(2L, result.getTotalRows()); + + Iterator it = result.iterator(); + assertTrue(it.hasNext()); + VectorSchemaRoot root = it.next(); + assertEquals(2, root.getRowCount()); + BigIntVector idVec = (BigIntVector) root.getVector("id"); + assertEquals(1L, idVec.get(0)); + assertEquals(2L, idVec.get(1)); + + assertFalse(it.hasNext()); + assertThrows(NoSuchElementException.class, it::next); + } + } + + @Test + void testIteratorCannotBeCreatedTwice() throws IOException { + Schema arrowSchema = createTestArrowSchema(); + byte[] batchBytes = createTestBatchBytes(ImmutableList.of(1L), ImmutableList.of("Alice")); + + try (ArrowQueryResult result = + new ArrowQueryResultImpl(arrowSchema, JobId.of("j1"), 1L, batchBytes, null, null)) { + Iterator it1 = result.iterator(); + assertNotNull(it1); + assertThrows(IllegalStateException.class, result::iterator); + } + } + + @Test + void testCloseIsIdempotentAndReleasesResources() throws IOException { + Schema arrowSchema = createTestArrowSchema(); + byte[] batchBytes = createTestBatchBytes(ImmutableList.of(1L), ImmutableList.of("Alice")); + + ArrowQueryResult result = + new ArrowQueryResultImpl(arrowSchema, JobId.of("j1"), 1L, batchBytes, null, null); + result.close(); + result.close(); + assertThrows(IllegalStateException.class, result::iterator); + } + + @Test + void testMultiBatchStreaming() throws IOException { + Schema arrowSchema = createTestArrowSchema(); + byte[] initialBatchBytes = + createTestBatchBytes(ImmutableList.of(1L, 2L), ImmutableList.of("A", "B")); + byte[] streamingBatchBytes = + createTestBatchBytes(ImmutableList.of(3L, 4L, 5L), ImmutableList.of("C", "D", "E")); + + @SuppressWarnings("unchecked") + ServerStreamingCallable mockCallable = + mock(ServerStreamingCallable.class); + + @SuppressWarnings("unchecked") + ServerStream mockServerStream = mock(ServerStream.class); + when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream); + + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(streamingBatchBytes)) + .build(); + ReadRowsResponse streamResponse = + ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build(); + + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(streamResponse).iterator()); + + BigQueryReadClient mockClient = createMockReadClient(mockCallable); + + String streamName = "projects/p/locations/l/jobs/j/streams/_default"; + try (ArrowQueryResult result = + new ArrowQueryResultImpl( + arrowSchema, JobId.of("j"), 5L, initialBatchBytes, streamName, mockClient)) { + Iterator it = result.iterator(); + + // Batch 1 (initial REST response) + assertTrue(it.hasNext()); + VectorSchemaRoot root1 = it.next(); + assertEquals(2, root1.getRowCount()); + + // Batch 2 (streaming gRPC response) + assertTrue(it.hasNext()); + VectorSchemaRoot root2 = it.next(); + assertEquals(3, root2.getRowCount()); + + assertFalse(it.hasNext()); + assertThrows(NoSuchElementException.class, it::next); + } + } + + @Test + void testQueryIdAndJobCreationReason() throws IOException { + Schema schema = createTestArrowSchema(); + byte[] batch1 = createTestBatchBytes(ImmutableList.of(1L), ImmutableList.of("Alice")); + JobId jobId = JobId.of("p", "j"); + String queryId = "query-12345"; + JobCreationReason reason = + JobCreationReason.fromPb( + new com.google.api.services.bigquery.model.JobCreationReason().setCode("REQUESTED")); + + ArrowQueryResultImpl result = + new ArrowQueryResultImpl(schema, jobId, queryId, reason, 1L, batch1, null, null); + + assertEquals(queryId, result.getQueryId()); + assertNotNull(result.getJobCreationReason()); + assertEquals(JobCreationReason.Code.REQUESTED, result.getJobCreationReason().getCode()); + assertEquals(jobId, result.getJobId()); + + try (ArrowQueryResult res = result) { + Iterator it = result.iterator(); + assertTrue(it.hasNext()); + VectorSchemaRoot root = it.next(); + assertEquals(1, root.getRowCount()); + assertFalse(it.hasNext()); + } + } + + @Test + void testStatelessSingleBatchIterationWithoutJobId() throws IOException { + Schema schema = createTestArrowSchema(); + byte[] batch1 = + createTestBatchBytes(ImmutableList.of(1L, 2L), ImmutableList.of("Alice", "Bob")); + + ArrowQueryResultImpl result = + new ArrowQueryResultImpl( + schema, + /* jobId= */ null, + /* queryId= */ "stateless-q-1", + /* jobCreationReason= */ null, + 2L, + batch1, + /* streamName= */ null, + /* readClient= */ null); + + assertNull(result.getJobId()); + assertEquals("stateless-q-1", result.getQueryId()); + assertNull(result.getJobCreationReason()); + + try (ArrowQueryResult res = result) { + Iterator it = result.iterator(); + assertTrue(it.hasNext()); + VectorSchemaRoot root = it.next(); + assertEquals(2, root.getRowCount()); + assertFalse(it.hasNext()); + } + } + + @Test + void testFromReadSessionFallback() throws IOException { + Schema schema = createTestArrowSchema(); + byte[] batch1 = createTestBatchBytes(ImmutableList.of(10L), ImmutableList.of("FallbackUser")); + + ByteArrayOutputStream schemaOut = new ByteArrayOutputStream(); + WriteChannel schemaChannel = new WriteChannel(Channels.newChannel(schemaOut)); + MessageSerializer.serialize(schemaChannel, schema); + + ArrowSchema arrowSchemaPb = + ArrowSchema.newBuilder() + .setSerializedSchema(ByteString.copyFrom(schemaOut.toByteArray())) + .build(); + ReadStream streamPb = + ReadStream.newBuilder().setName("projects/p/locations/l/sessions/s/streams/str1").build(); + ReadSession readSession = + ReadSession.newBuilder().setArrowSchema(arrowSchemaPb).addStreams(streamPb).build(); + + ReadRowsResponse response = + ReadRowsResponse.newBuilder() + .setArrowRecordBatch( + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(batch1)) + .build()) + .build(); + + @SuppressWarnings("unchecked") + ServerStream mockServerStream = mock(ServerStream.class); + when(mockServerStream.iterator()).thenReturn(ImmutableList.of(response).iterator()); + + @SuppressWarnings("unchecked") + ServerStreamingCallable mockCallable = + mock(ServerStreamingCallable.class); + when(mockCallable.call(any(ReadRowsRequest.class))).thenReturn(mockServerStream); + + BigQueryReadClient mockReadClient = createMockReadClient(mockCallable); + + JobId jobId = JobId.of("p", "fallback-job"); + ArrowQueryResultImpl result = + ArrowQueryResultImpl.fromReadSession(readSession, jobId, mockReadClient); + + assertEquals(jobId, result.getJobId()); + assertNull(result.getQueryId()); + assertNull(result.getJobCreationReason()); + + try (ArrowQueryResult res = result) { + Iterator it = result.iterator(); + assertTrue(it.hasNext()); + VectorSchemaRoot root = it.next(); + assertEquals(1, root.getRowCount()); + BigIntVector idVector = (BigIntVector) root.getVector("id"); + assertEquals(10L, idVector.get(0)); + assertFalse(it.hasNext()); + } + } +} diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/BigQueryImplTest.java index 9f2320ab3c35..d0838a17f99a 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 @@ -2887,6 +2887,42 @@ void testQueryWithTimeoutSetsTimeout() throws InterruptedException, IOException assertEquals((Long) 1000L, requestPb.getTimeoutMs()); } + @Test + void testQueryThrowsWhenArrowResultsFormat() { + QueryJobConfiguration config = + QueryJobConfiguration.newBuilder("SELECT 1") + .setQueryResultsFormat(QueryResultsFormat.ARROW) + .build(); + bigquery = options.getService(); + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, () -> bigquery.query(config)); + assertTrue(exception.getMessage().contains("Use queryArrow() instead")); + } + + @Test + void testQueryArrowDefaultsToJobCreationOptional() throws IOException, InterruptedException { + QueryJobConfiguration config = QueryJobConfiguration.newBuilder("SELECT 1").build(); + com.google.api.services.bigquery.model.QueryResponse queryResponsePb = + new com.google.api.services.bigquery.model.QueryResponse() + .setQueryId("q-optional-1") + .setJobComplete(true) + .setTotalRows(java.math.BigInteger.ZERO); + + ArgumentCaptor requestPbCapture = ArgumentCaptor.forClass(QueryRequest.class); + when(bigqueryRpcMock.queryRpcSkipExceptionTranslation(eq(PROJECT), requestPbCapture.capture())) + .thenReturn(queryResponsePb); + + bigquery = options.getService(); + ArrowQueryResult result = bigquery.queryArrow(config); + assertNotNull(result); + assertEquals("q-optional-1", result.getQueryId()); + assertNull(result.getJobId()); + + QueryRequest requestPb = requestPbCapture.getValue(); + assertEquals("JOB_CREATION_OPTIONAL", requestPb.getJobCreationMode()); + assertEquals("ARROW", requestPb.getQueryResultsFormat()); + } + @Test void testGetQueryResults() throws IOException { JobId queryJob = JobId.of(JOB); diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/it/ITBigQueryTest.java index 3c9613d20758..a87ab517f5fa 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; @@ -119,6 +120,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; @@ -210,6 +212,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; @@ -7561,6 +7564,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();