From c05c58541e9212985618ab1eaecf068fbbb88891 Mon Sep 17 00:00:00 2001 From: Jin Seop Kim Date: Wed, 29 Jul 2026 10:33:48 -0400 Subject: [PATCH] feat(bigquery): integrate Arrow query response processing and stream pagination --- .../google/cloud/bigquery/BigQueryImpl.java | 253 ++++++++++++++++-- .../cloud/bigquery/QueryRequestInfo.java | 10 + 2 files changed, 245 insertions(+), 18 deletions(-) diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/BigQueryImpl.java index a5e91886d030..ae92f56c597f 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 @@ -22,7 +22,9 @@ import com.google.api.core.BetaApi; import com.google.api.core.InternalApi; +import com.google.api.gax.core.FixedCredentialsProvider; import com.google.api.gax.paging.Page; +import com.google.api.gax.rpc.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; @@ -43,6 +45,10 @@ 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.ReadRowsRequest; +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; import com.google.common.annotations.VisibleForTesting; import com.google.common.base.Function; import com.google.common.base.Strings; @@ -57,6 +63,7 @@ import io.opentelemetry.context.Scope; import java.io.IOException; import java.util.ArrayList; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.concurrent.Callable; @@ -264,6 +271,143 @@ public Page getNextPage() { } } + private static class ArrowQueryPageFetcher implements NextPageFetcher { + private static final long serialVersionUID = 1L; + + private final JobId jobId; + private final Schema schema; + private final org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo; + private final BigQueryOptions serviceOptions; + private final long maxResults; + + private transient BigQueryReadClient bqReadClient; + private transient ServerStream stream; + private transient Iterator streamIterator; + private long totalRowsReturned = 0L; + private boolean streamClosed = false; + + ArrowQueryPageFetcher( + JobId jobId, + Schema schema, + org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo, + BigQueryOptions serviceOptions, + long initialRowOffset, + Long maxResults) { + this.jobId = jobId; + this.schema = schema; + this.arrowSchemaPojo = arrowSchemaPojo; + this.serviceOptions = serviceOptions; + this.totalRowsReturned = initialRowOffset; + this.maxResults = maxResults != null ? maxResults : Long.MAX_VALUE; + } + + @Override + public Page getNextPage() { + if (streamClosed || totalRowsReturned >= maxResults) { + closeClient(); + return null; + } + + long pageSize = 100000L; + List rowBatch = new ArrayList<>(); + + try { + if (bqReadClient == null) { + BigQueryReadSettings settings = + BigQueryReadSettings.newBuilder() + .setCredentialsProvider( + FixedCredentialsProvider.create(serviceOptions.getCredentials())) + .build(); + bqReadClient = BigQueryReadClient.create(settings); + } + + if (streamIterator == null) { + String streamName = + String.format( + "projects/%s/locations/%s/jobs/%s/streams/_default", + jobId.getProject() != null ? jobId.getProject() : serviceOptions.getProjectId(), + jobId.getLocation() != null ? jobId.getLocation() : serviceOptions.getLocation(), + jobId.getJob()); + + ReadRowsRequest readRowsRequest = + ReadRowsRequest.newBuilder() + .setReadStream(streamName) + .setOffset(totalRowsReturned) + .build(); + + stream = bqReadClient.readRowsCallable().call(readRowsRequest); + streamIterator = stream.iterator(); + } + + try (org.apache.arrow.memory.BufferAllocator allocator = + new org.apache.arrow.memory.RootAllocator(Long.MAX_VALUE)) { + List vectors = new ArrayList<>(); + for (org.apache.arrow.vector.types.pojo.Field field : arrowSchemaPojo.getFields()) { + vectors.add((org.apache.arrow.vector.FieldVector) field.createVector(allocator)); + } + try (org.apache.arrow.vector.VectorSchemaRoot root = + new org.apache.arrow.vector.VectorSchemaRoot(vectors)) { + org.apache.arrow.vector.VectorLoader loader = + new org.apache.arrow.vector.VectorLoader(root); + + while (rowBatch.size() < pageSize && streamIterator.hasNext()) { + ReadRowsResponse response = streamIterator.next(); + if (response.hasArrowRecordBatch()) { + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = + response.getArrowRecordBatch(); + org.apache.arrow.vector.ipc.message.ArrowRecordBatch deserializedBatch = + org.apache.arrow.vector.ipc.message.MessageSerializer.deserializeRecordBatch( + new org.apache.arrow.vector.ipc.ReadChannel( + new org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel( + batch.getSerializedRecordBatch().toByteArray())), + allocator); + loader.load(deserializedBatch); + deserializedBatch.close(); + int batchRowCount = root.getRowCount(); + for (int i = 0; i < batchRowCount; i++) { + rowBatch.add(ArrowDeserializer.arrowRootToFieldValueList(root, i, schema)); + } + root.clear(); + } + } + } + } + + if (rowBatch.isEmpty()) { + streamClosed = true; + closeClient(); + return null; + } + + totalRowsReturned += rowBatch.size(); + + String nextPageToken = null; + if (streamIterator.hasNext() && totalRowsReturned < maxResults) { + nextPageToken = String.valueOf(totalRowsReturned); + } else { + streamClosed = true; + closeClient(); + } + + return new PageImpl<>(this, nextPageToken, rowBatch); + + } catch (Exception e) { + streamClosed = true; + closeClient(); + throw new BigQueryException(0, "Failed to read Arrow rows from storage stream", e); + } + } + + private void closeClient() { + if (bqReadClient != null) { + bqReadClient.close(); + bqReadClient = null; + } + streamIterator = null; + stream = null; + } + } + private final HttpBigQueryRpc bigQueryRpc; private static final BigQueryRetryConfig EMPTY_RETRY_CONFIG = @@ -2073,8 +2217,28 @@ 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 = false; + org.apache.arrow.vector.types.pojo.Schema arrowSchemaPojo = null; + + if (results.getJobComplete()) { + if (results.getSchema() != null) { + schema = Schema.fromPb(results.getSchema()); + } else if (results.getArrowSchema() != null) { + isArrow = true; + try { + arrowSchemaPojo = + org.apache.arrow.vector.ipc.message.MessageSerializer.deserializeSchema( + new org.apache.arrow.vector.ipc.ReadChannel( + new org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel( + results.getArrowSchema().decodeSerializedSchema()))); + schema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchemaPojo); + } catch (IOException e) { + throw new BigQueryException(0, "Failed to deserialize Arrow schema from response", e); + } + } else { + schema = null; + } + if (results.getNumDmlAffectedRows() == null && results.getTotalRows() == null) { numRows = 0L; } else if (results.getNumDmlAffectedRows() != null) { @@ -2094,42 +2258,91 @@ public com.google.api.services.bigquery.model.QueryResponse call() if (results.getPageToken() != null) { JobId jobId = JobId.fromPb(results.getJobReference()); String cursor = results.getPageToken(); + + Iterable firstPageRows; + NextPageFetcher pageFetcher; + + if (isArrow) { + if (results.getArrowRecordBatch() != 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(); + } + long initialRowOffset = + firstPageRows instanceof List ? ((List) firstPageRows).size() : 0L; + pageFetcher = + new ArrowQueryPageFetcher( + jobId, + schema, + arrowSchemaPojo, + getOptions(), + initialRowOffset, + null); // Or use maxResults from configuration if available + } else { + firstPageRows = + transformTableData( + results.getRows(), schema, getOptions().getDataFormatOptions().useInt64Timestamp()); + 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()))) + .setPageNoSchema(new PageImpl<>(pageFetcher, cursor, firstPageRows)) .setJobId(jobId) .setQueryId(results.getQueryId()) .setJobCreationReason(JobCreationReason.fromPb(results.getJobCreationReason())) - .setRowsInPage(results.getRows() != null ? (long) results.getRows().size() : 0L) + .setRowsInPage( + firstPageRows instanceof List ? (long) ((List) firstPageRows).size() : 0L) .build(); } // only 1 page of result + Iterable firstPageRows; + if (isArrow) { + if (results.getArrowRecordBatch() != 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 = + transformTableData( + results.getRows(), schema, getOptions().getDataFormatOptions().useInt64Timestamp()); + } + return TableResult.newBuilder() .setSchema(schema) .setTotalRows(numRows) .setPageNoSchema( new PageImpl<>( - new TableDataPageFetcher(null, schema, getOptions(), null, optionMap(options)), + isArrow + ? null + : 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(firstPageRows instanceof List ? (long) ((List) firstPageRows).size() : 0L) .build(); } @@ -2203,6 +2416,10 @@ && 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) { 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..ec09d140a28d 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) { @@ -66,6 +68,8 @@ final class QueryRequestInfo { this.formatOptions = dataFormatOptions.toPb(); 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; }