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 7c5829734822..e131c157d6d1 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 @@ -45,105 +45,13 @@ final class ArrowDeserializer { private ArrowDeserializer() {} /** - * Converts an Apache Arrow {@link org.apache.arrow.vector.types.pojo.Schema} to a BigQuery Veneer - * {@link Schema}. + * Converts an Apache Arrow Schema to a BigQuery Veneer {@link Schema}. * * @param arrowSchema the Apache Arrow schema to convert * @return the corresponding BigQuery Veneer Schema */ - static Schema arrowSchemaToBigQuerySchema(org.apache.arrow.vector.types.pojo.Schema arrowSchema) { - List fields = new ArrayList<>(); - for (org.apache.arrow.vector.types.pojo.Field arrowField : arrowSchema.getFields()) { - fields.add(arrowFieldToBigQueryField(arrowField)); - } - return Schema.of(fields); - } - - /** - * Recursively converts an Apache Arrow {@link org.apache.arrow.vector.types.pojo.Field} to a - * BigQuery Veneer {@link Field}. - * - * @param arrowField the Arrow field to convert - * @return the corresponding BigQuery Veneer Field - */ - private static Field arrowFieldToBigQueryField( - org.apache.arrow.vector.types.pojo.Field arrowField) { - String name = arrowField.getName(); - ArrowType type = arrowField.getType(); - Field.Builder builder; - - if (type instanceof ArrowType.List) { - if (arrowField.getChildren().isEmpty()) { - throw new IllegalArgumentException( - "Arrow List field must have at least one child field: " + name); - } - org.apache.arrow.vector.types.pojo.Field innerField = arrowField.getChildren().get(0); - LegacySQLTypeName innerType = arrowTypeToLegacySQLTypeName(innerField.getType()); - builder = Field.newBuilder(name, innerType); - builder.setMode(Field.Mode.REPEATED); - if (!innerField.getChildren().isEmpty()) { - List subFields = new ArrayList<>(); - for (org.apache.arrow.vector.types.pojo.Field childField : innerField.getChildren()) { - subFields.add(arrowFieldToBigQueryField(childField)); - } - builder.setType(LegacySQLTypeName.RECORD, FieldList.of(subFields)); - } - } else { - LegacySQLTypeName bqType = arrowTypeToLegacySQLTypeName(type); - builder = Field.newBuilder(name, bqType); - if (arrowField.isNullable()) { - builder.setMode(Field.Mode.NULLABLE); - } else { - builder.setMode(Field.Mode.REQUIRED); - } - if (!arrowField.getChildren().isEmpty()) { - List subFields = new ArrayList<>(); - for (org.apache.arrow.vector.types.pojo.Field childField : innerFieldChildren(arrowField)) { - subFields.add(arrowFieldToBigQueryField(childField)); - } - builder.setType(LegacySQLTypeName.RECORD, FieldList.of(subFields)); - } - } - return builder.build(); - } - - private static List innerFieldChildren( - org.apache.arrow.vector.types.pojo.Field arrowField) { - return arrowField.getChildren(); - } - - /** - * Maps an Apache Arrow data type {@link ArrowType} to a BigQuery {@link LegacySQLTypeName}. - * - * @param type the Arrow data type to map - * @return the corresponding BigQuery LegacySQLTypeName - * @throws IllegalArgumentException if the Arrow type is unsupported - */ - private static LegacySQLTypeName arrowTypeToLegacySQLTypeName(ArrowType type) { - switch (type.getTypeID()) { - case Int: - return LegacySQLTypeName.INTEGER; - case FloatingPoint: - return LegacySQLTypeName.FLOAT; - case Utf8: - return LegacySQLTypeName.STRING; - case Bool: - return LegacySQLTypeName.BOOLEAN; - case Binary: - return LegacySQLTypeName.BYTES; - case Decimal: - return LegacySQLTypeName.NUMERIC; - case Timestamp: - return LegacySQLTypeName.TIMESTAMP; - case Date: - return LegacySQLTypeName.DATE; - case Time: - return LegacySQLTypeName.TIME; - case Struct: - return LegacySQLTypeName.RECORD; - default: - throw new IllegalArgumentException("Unsupported Arrow type: " + type.getTypeID()); - } + static Schema arrowSchemaToBigQuerySchema(Object arrowSchema) { + return ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema); } /** @@ -160,13 +68,24 @@ private static LegacySQLTypeName arrowTypeToLegacySQLTypeName(ArrowType type) { * @throws IOException if deserialization of the Arrow record batch fails */ static List deserializeRecordBatch( - byte[] recordBatchBytes, Schema schema, org.apache.arrow.vector.types.pojo.Schema arrowSchema) - throws IOException { + byte[] recordBatchBytes, Schema schema, Object arrowSchema) throws IOException { try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { - List vectors = new ArrayList<>(); + List vectors = ArrowPojoUtils.createVectors(arrowSchema, allocator); try { - for (org.apache.arrow.vector.types.pojo.Field field : arrowSchema.getFields()) { - vectors.add(field.createVector(allocator)); + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + VectorLoader loader = new VectorLoader(root); + try (ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch( + new ReadChannel(new ByteArrayReadableSeekableByteChannel(recordBatchBytes)), + allocator)) { + loader.load(deserializedBatch); + int rowCount = root.getRowCount(); + List rows = new ArrayList<>(rowCount); + for (int i = 0; i < rowCount; i++) { + rows.add(arrowRootToFieldValueList(root, i, schema)); + } + return ImmutableList.copyOf(rows); + } } } catch (Throwable t) { for (int i = vectors.size() - 1; i >= 0; i--) { @@ -178,21 +97,6 @@ static List deserializeRecordBatch( } throw t; } - try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { - VectorLoader loader = new VectorLoader(root); - try (ArrowRecordBatch deserializedBatch = - MessageSerializer.deserializeRecordBatch( - new ReadChannel(new ByteArrayReadableSeekableByteChannel(recordBatchBytes)), - allocator)) { - loader.load(deserializedBatch); - int rowCount = root.getRowCount(); - List rows = new ArrayList<>(rowCount); - for (int i = 0; i < rowCount; i++) { - rows.add(arrowRootToFieldValueList(root, i, schema)); - } - return ImmutableList.copyOf(rows); - } - } } } @@ -281,9 +185,6 @@ private static FieldValue arrowVectorToFieldValue( // Handle primitive types String stringVal; if (bqField.getType() == LegacySQLTypeName.TIMESTAMP) { - // Arrow timestamps are long values representing epoch seconds/millis/micros/nanos. - // Standard BigQuery JSON returns timestamps as string of epoch seconds with micro precision - // (e.g. "1408452095.220000"). TimeStampVector tsVector = (TimeStampVector) vector; long rawVal = tsVector.get(rowIndex); ArrowType.Timestamp tsType = (ArrowType.Timestamp) vector.getField().getType(); diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowPojoUtils.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowPojoUtils.java new file mode 100644 index 000000000000..f1bfa3e501a0 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowPojoUtils.java @@ -0,0 +1,118 @@ +/* + * 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 java.util.ArrayList; +import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.pojo.ArrowType; + +/** Internal helper for Apache Arrow Schema/Field conversions. */ +final class ArrowPojoUtils { + + private ArrowPojoUtils() {} + + static Schema arrowSchemaToBigQuerySchema(Object arrowSchemaObj) { + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + (org.apache.arrow.vector.types.pojo.Schema) arrowSchemaObj; + List fields = new ArrayList<>(); + for (org.apache.arrow.vector.types.pojo.Field arrowField : arrowSchema.getFields()) { + fields.add(arrowFieldToBigQueryField(arrowField)); + } + return Schema.of(fields); + } + + static Field arrowFieldToBigQueryField(Object arrowFieldObj) { + org.apache.arrow.vector.types.pojo.Field arrowField = + (org.apache.arrow.vector.types.pojo.Field) arrowFieldObj; + String name = arrowField.getName(); + ArrowType type = arrowField.getType(); + Field.Builder builder; + + if (type instanceof ArrowType.List) { + if (arrowField.getChildren().isEmpty()) { + throw new IllegalArgumentException( + "Arrow List field must have at least one child field: " + name); + } + org.apache.arrow.vector.types.pojo.Field innerField = arrowField.getChildren().get(0); + LegacySQLTypeName innerType = arrowTypeToLegacySQLTypeName(innerField.getType()); + builder = Field.newBuilder(name, innerType); + builder.setMode(Field.Mode.REPEATED); + if (!innerField.getChildren().isEmpty()) { + List subFields = new ArrayList<>(); + for (org.apache.arrow.vector.types.pojo.Field childField : innerField.getChildren()) { + subFields.add(arrowFieldToBigQueryField(childField)); + } + builder.setType(LegacySQLTypeName.RECORD, FieldList.of(subFields)); + } + } else { + LegacySQLTypeName bqType = arrowTypeToLegacySQLTypeName(type); + builder = Field.newBuilder(name, bqType); + if (arrowField.isNullable()) { + builder.setMode(Field.Mode.NULLABLE); + } else { + builder.setMode(Field.Mode.REQUIRED); + } + if (!arrowField.getChildren().isEmpty()) { + List subFields = new ArrayList<>(); + for (org.apache.arrow.vector.types.pojo.Field childField : arrowField.getChildren()) { + subFields.add(arrowFieldToBigQueryField(childField)); + } + builder.setType(LegacySQLTypeName.RECORD, FieldList.of(subFields)); + } + } + return builder.build(); + } + + static List createVectors(Object arrowSchemaObj, BufferAllocator allocator) { + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + (org.apache.arrow.vector.types.pojo.Schema) arrowSchemaObj; + List vectors = new ArrayList<>(); + for (org.apache.arrow.vector.types.pojo.Field field : arrowSchema.getFields()) { + vectors.add(field.createVector(allocator)); + } + return vectors; + } + + private static LegacySQLTypeName arrowTypeToLegacySQLTypeName(ArrowType type) { + switch (type.getTypeID()) { + case Int: + return LegacySQLTypeName.INTEGER; + case FloatingPoint: + return LegacySQLTypeName.FLOAT; + case Utf8: + return LegacySQLTypeName.STRING; + case Bool: + return LegacySQLTypeName.BOOLEAN; + case Binary: + return LegacySQLTypeName.BYTES; + case Decimal: + return LegacySQLTypeName.NUMERIC; + case Timestamp: + return LegacySQLTypeName.TIMESTAMP; + case Date: + return LegacySQLTypeName.DATE; + case Time: + return LegacySQLTypeName.TIME; + case Struct: + return LegacySQLTypeName.RECORD; + default: + throw new IllegalArgumentException("Unsupported Arrow type: " + type.getTypeID()); + } + } +} diff --git a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowSerializationOptions.java b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowSerializationOptions.java index 1ad88f8de84b..1c1c6295360f 100644 --- a/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowSerializationOptions.java +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowSerializationOptions.java @@ -77,8 +77,7 @@ com.google.api.services.bigquery.model.ArrowSerializationOptions toPb() { return ArrowSerializationOptionsConverter.toPb(this); } - static ArrowSerializationOptions fromPb( - com.google.api.services.bigquery.model.ArrowSerializationOptions optionsPb) { + static ArrowSerializationOptions fromPb(Object optionsPb) { return ArrowSerializationOptionsConverter.fromPb(optionsPb); } 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..d40724363f6c 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,11 +63,21 @@ 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; import java.util.regex.Matcher; import java.util.regex.Pattern; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +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.util.ByteArrayReadableSeekableByteChannel; import org.checkerframework.checker.nullness.qual.NonNull; final class BigQueryImpl extends BaseService implements BigQuery { @@ -264,6 +280,137 @@ 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 Object 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, + Object 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 (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + List vectors = ArrowPojoUtils.createVectors(arrowSchemaPojo, allocator); + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + VectorLoader loader = new 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(); + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch( + new ReadChannel( + new 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 = @@ -2077,8 +2224,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; + Object arrowSchemaPojo = null; + + if (results.getJobComplete()) { + if (results.getSchema() != null) { + schema = Schema.fromPb(results.getSchema()); + } else if (results.getArrowSchema() != null) { + isArrow = true; + try { + arrowSchemaPojo = + MessageSerializer.deserializeSchema( + new ReadChannel( + new 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) { @@ -2098,42 +2265,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(); } @@ -2207,6 +2423,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..ea5df0ff2cf2 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,23 +65,13 @@ 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(); } - /** - * Determines if the query can be executed via the "fast query" path (jobs.query API) instead of - * the "slow path" (jobs.insert API followed by jobs.getQueryResults). - * - *

The fast query path is preferred because it completes in a single RPC, significantly - * reducing end-to-end latency for small queries. - * - *

However, the jobs.query API does not support all configuration options available in - * jobs.insert (e.g., destination table, clustering, time partitioning). This method checks the - * QueryJobConfiguration for any unsupported options. If any are present, we must fall back to the - * jobs.insert path. - */ boolean isFastQuerySupported() { return config.getClustering() == null && config.getCreateDisposition() == null @@ -142,6 +134,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 +159,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/ArrowDeserializerTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java index f23dbe4b01a5..b1524d49d603 100644 --- a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java @@ -43,39 +43,42 @@ import org.apache.arrow.vector.ipc.message.MessageSerializer; import org.apache.arrow.vector.types.TimeUnit; 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.Test; public class ArrowDeserializerTest { @Test public void testArrowSchemaToBigQuerySchema() { - Field intField = new Field("int_col", FieldType.nullable(new ArrowType.Int(32, true)), null); - Field strField = new Field("str_col", FieldType.notNullable(new ArrowType.Utf8()), null); - Field boolField = new Field("bool_col", FieldType.nullable(new ArrowType.Bool()), null); - Field tsField = - new Field( + org.apache.arrow.vector.types.pojo.Field intField = + new org.apache.arrow.vector.types.pojo.Field( + "int_col", FieldType.nullable(new ArrowType.Int(32, true)), null); + org.apache.arrow.vector.types.pojo.Field strField = + new org.apache.arrow.vector.types.pojo.Field( + "str_col", FieldType.notNullable(new ArrowType.Utf8()), null); + org.apache.arrow.vector.types.pojo.Field boolField = + new org.apache.arrow.vector.types.pojo.Field( + "bool_col", FieldType.nullable(new ArrowType.Bool()), null); + org.apache.arrow.vector.types.pojo.Field tsField = + new org.apache.arrow.vector.types.pojo.Field( "ts_col", FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")), null); - Schema arrowSchema = new Schema(ImmutableList.of(intField, strField, boolField, tsField)); + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + new org.apache.arrow.vector.types.pojo.Schema( + ImmutableList.of(intField, strField, boolField, tsField)); - com.google.cloud.bigquery.Schema bqSchema = - ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); + Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); assertEquals(4, bqSchema.getFields().size()); assertEquals("int_col", bqSchema.getFields().get(0).getName()); assertEquals(LegacySQLTypeName.INTEGER, bqSchema.getFields().get(0).getType()); - assertEquals( - com.google.cloud.bigquery.Field.Mode.NULLABLE, bqSchema.getFields().get(0).getMode()); + assertEquals(Field.Mode.NULLABLE, bqSchema.getFields().get(0).getMode()); assertEquals("str_col", bqSchema.getFields().get(1).getName()); assertEquals(LegacySQLTypeName.STRING, bqSchema.getFields().get(1).getType()); - assertEquals( - com.google.cloud.bigquery.Field.Mode.REQUIRED, bqSchema.getFields().get(1).getMode()); + assertEquals(Field.Mode.REQUIRED, bqSchema.getFields().get(1).getMode()); assertEquals("bool_col", bqSchema.getFields().get(2).getName()); assertEquals(LegacySQLTypeName.BOOLEAN, bqSchema.getFields().get(2).getType()); @@ -128,9 +131,8 @@ public void testDeserializeRecordBatchPrimitives() throws IOException { ImmutableList.of(intVector, nameVector, scoreVector, activeVector, bytesVector, tsVector); try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { - Schema arrowSchema = root.getSchema(); - com.google.cloud.bigquery.Schema bqSchema = - ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); + Object arrowSchema = root.getSchema(); + Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); byte[] recordBatchBytes = serializeVectorSchemaRoot(root, allocator); @@ -155,11 +157,7 @@ public void testDeserializeRecordBatchPrimitives() throws IOException { assertEquals("102", row1.get("id").getStringValue()); assertEquals("Bob", row1.get("name").getStringValue()); assertNull(row1.get("score").getValue()); - assertEquals( - "false", - row1.get("false".equals("false") ? "active" : "score") != null - ? row1.get("active").getStringValue() - : "false"); + assertEquals("false", row1.get("active").getStringValue()); assertNull(row1.get("data").getValue()); assertNull(row1.get("ts").getValue()); } finally { @@ -179,10 +177,10 @@ public void testSchemaMismatchThrowsException() { intVector.setValueCount(1); try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(intVector))) { - com.google.cloud.bigquery.Schema mismatchedSchema = - com.google.cloud.bigquery.Schema.of( - com.google.cloud.bigquery.Field.of("col1", LegacySQLTypeName.INTEGER), - com.google.cloud.bigquery.Field.of("col2", LegacySQLTypeName.STRING)); + Schema mismatchedSchema = + Schema.of( + Field.of("col1", LegacySQLTypeName.INTEGER), + Field.of("col2", LegacySQLTypeName.STRING)); try { ArrowDeserializer.arrowRootToFieldValueList(root, 0, mismatchedSchema);