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 new file mode 100644 index 000000000000..3def579e2abf --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java @@ -0,0 +1,396 @@ +/* + * 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.cloud.bigquery.storage.v1.ReadRowsResponse; +import com.google.common.collect.ImmutableList; +import com.google.common.io.BaseEncoding; +import java.io.IOException; +import java.nio.channels.Channels; +import java.util.ArrayList; +import java.util.Iterator; +import java.util.List; +import java.util.Locale; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.TimeStampVector; +import org.apache.arrow.vector.VectorLoader; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.complex.ListVector; +import org.apache.arrow.vector.complex.StructVector; +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.ArrowType; +import org.apache.arrow.vector.util.ByteArrayReadableSeekableByteChannel; + +/** + * Internal helper utility for converting Apache Arrow schemas and record batches into BigQuery + * Veneer objects. + */ +final class ArrowDeserializer { + + /** Lazy initialization holder for the root {@link BufferAllocator}. */ + private static class AllocatorHolder { + private static final BufferAllocator ALLOCATOR = new RootAllocator(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 + * construction. + * + * @param arrowSchema the Apache Arrow schema definition + * @param allocator the buffer allocator to bind the vectors to + * @return a new VectorSchemaRoot containing allocated field vectors + */ + private static VectorSchemaRoot createVectorSchemaRoot( + org.apache.arrow.vector.types.pojo.Schema arrowSchema, BufferAllocator allocator) { + List vectors = ArrowPojoUtils.createVectors(arrowSchema, allocator); + try { + return new VectorSchemaRoot(vectors); + } catch (Throwable t) { + for (int i = vectors.size() - 1; i >= 0; i--) { + try { + vectors.get(i).close(); + } catch (Exception e) { + t.addSuppressed(e); + } + } + throw t; + } + } + + private ArrowDeserializer() {} + + /** + * Deserializes a raw binary Arrow schema payload into an Apache Arrow Schema object. + * + * @param schemaBytes the raw binary Arrow schema payload + * @return the deserialized Apache Arrow Schema object + * @throws IOException if deserialization of the Arrow schema fails + */ + static Object deserializeSchema(byte[] schemaBytes) throws IOException { + try (ByteArrayReadableSeekableByteChannel byteChannel = + new ByteArrayReadableSeekableByteChannel(schemaBytes); + ReadChannel readChannel = new ReadChannel(byteChannel)) { + return MessageSerializer.deserializeSchema(readChannel); + } + } + + /** + * Serializes an Apache Arrow Schema object to its JSON string representation. + * + * @param arrowSchema the Apache Arrow schema object + * @return the JSON string representation, or null if arrowSchema is null + */ + static String arrowSchemaToJson(Object arrowSchema) { + if (arrowSchema == null) { + return null; + } + return ((org.apache.arrow.vector.types.pojo.Schema) arrowSchema).toJson(); + } + + /** + * Deserializes an Apache Arrow Schema object from its JSON string representation. + * + * @param json the JSON string representation of the Arrow schema + * @return the deserialized Apache Arrow Schema object, or null if json is null + * @throws IllegalArgumentException if the JSON string cannot be parsed as an Arrow schema + */ + static Object jsonToArrowSchema(String json) { + if (json == null) { + return null; + } + try { + return org.apache.arrow.vector.types.pojo.Schema.fromJSON(json); + } catch (IOException e) { + throw new IllegalArgumentException("Invalid Arrow schema JSON", e); + } + } + + /** + * Resolves an Apache Arrow Schema from either an in-memory Schema POJO or a serialized JSON + * string. + * + * @param arrowSchema the Arrow schema POJO or JSON string representation + * @return the resolved Apache Arrow Schema, or null if schema cannot be resolved + * @throws IOException if parsing JSON fails + */ + private static org.apache.arrow.vector.types.pojo.Schema resolveArrowSchema(Object arrowSchema) + throws IOException { + if (arrowSchema instanceof org.apache.arrow.vector.types.pojo.Schema) { + return (org.apache.arrow.vector.types.pojo.Schema) arrowSchema; + } + if (arrowSchema instanceof String) { + return org.apache.arrow.vector.types.pojo.Schema.fromJSON((String) arrowSchema); + } + return null; + } + + /** + * Reads and decodes a batch of Arrow rows from the provided stream iterator into the row batch. + * + * @param iterator the stream iterator providing ReadRowsResponse messages + * @param arrowSchema the Arrow schema POJO or serialized JSON representation + * @param schema the BigQuery target Schema + * @param rowBatch the destination list for decoded rows + * @param pageSize the maximum number of rows to decode in this batch + * @param totalRowsReturned the running count of rows returned so far + * @param maxResults the maximum total rows allowed across all pages + * @return true if more rows are available in the stream and maxResults has not been reached + * @throws IOException if deserialization fails + */ + static boolean loadArrowRows( + Iterator iterator, + Object arrowSchema, + Schema schema, + List rowBatch, + long pageSize, + long totalRowsReturned, + long maxResults) + throws IOException { + org.apache.arrow.vector.types.pojo.Schema resolvedSchema = resolveArrowSchema(arrowSchema); + + if (resolvedSchema == null) { + return false; + } + + org.apache.arrow.vector.types.pojo.Schema arrowSchemaFinal = resolvedSchema; + + try (BufferAllocator childAllocator = + AllocatorHolder.ALLOCATOR.newChildAllocator("loadArrowRows", 0, Long.MAX_VALUE); + VectorSchemaRoot closedRoot = createVectorSchemaRoot(arrowSchemaFinal, childAllocator)) { + VectorLoader loader = new VectorLoader(closedRoot); + boolean hasMore = false; + while (rowBatch.size() < pageSize + && iterator.hasNext() + && (totalRowsReturned + rowBatch.size() < maxResults)) { + ReadRowsResponse response = iterator.next(); + if (response.hasArrowRecordBatch()) { + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch batch = + response.getArrowRecordBatch(); + try (ReadChannel readChannel = + new ReadChannel( + Channels.newChannel(batch.getSerializedRecordBatch().newInput())); + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch(readChannel, childAllocator)) { + loader.load(deserializedBatch); + int batchRowCount = closedRoot.getRowCount(); + int i = 0; + for (; i < batchRowCount; i++) { + if (rowBatch.size() >= pageSize + || totalRowsReturned + rowBatch.size() >= maxResults) { + break; + } + rowBatch.add(arrowRootToFieldValueList(closedRoot, i, schema)); + } + if (i < batchRowCount && (totalRowsReturned + rowBatch.size() < maxResults)) { + hasMore = true; + } + closedRoot.clear(); + } + } + } + if (!hasMore) { + hasMore = iterator.hasNext() && (totalRowsReturned + rowBatch.size() < maxResults); + } + return hasMore; + } + } + + /** + * 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(Object arrowSchema) { + return ArrowPojoUtils.arrowSchemaToBigQuerySchema( + (org.apache.arrow.vector.types.pojo.Schema) arrowSchema); + } + + /** + * Deserializes a raw binary Arrow record batch payload into a list of BigQuery {@link + * FieldValueList} row objects. + * + *

Allocates off-heap memory within a local child allocator scope and closes all Arrow vector + * resources before returning, guaranteeing that native memory is released. + * + * @param recordBatchBytes the raw binary Arrow record batch payload + * @param schema the target BigQuery Schema + * @param arrowSchema the Arrow schema describing the record batch structure + * @return an immutable list of FieldValueList row objects + * @throws IOException if deserialization of the Arrow record batch fails + */ + static List deserializeRecordBatch( + byte[] recordBatchBytes, Schema schema, Object arrowSchema) throws IOException { + org.apache.arrow.vector.types.pojo.Schema schemaPojo = + arrowSchema instanceof org.apache.arrow.vector.types.pojo.Schema + ? (org.apache.arrow.vector.types.pojo.Schema) arrowSchema + : (arrowSchema instanceof String + ? (org.apache.arrow.vector.types.pojo.Schema) + jsonToArrowSchema((String) arrowSchema) + : null); + if (schemaPojo == null) { + throw new IllegalArgumentException("Arrow schema must not be null"); + } + try (BufferAllocator childAllocator = + AllocatorHolder.ALLOCATOR.newChildAllocator( + "deserializeRecordBatch", 0, Long.MAX_VALUE); + VectorSchemaRoot closedRoot = createVectorSchemaRoot(schemaPojo, childAllocator); + ByteArrayReadableSeekableByteChannel byteChannel = + new ByteArrayReadableSeekableByteChannel(recordBatchBytes); + ReadChannel readChannel = new ReadChannel(byteChannel); + ArrowRecordBatch deserializedBatch = + MessageSerializer.deserializeRecordBatch(readChannel, childAllocator)) { + VectorLoader loader = new VectorLoader(closedRoot); + loader.load(deserializedBatch); + int rowCount = closedRoot.getRowCount(); + List rows = new ArrayList<>(rowCount); + for (int i = 0; i < rowCount; i++) { + rows.add(arrowRootToFieldValueList(closedRoot, i, schema)); + } + return ImmutableList.copyOf(rows); + } + } + + /** + * Extracts a single row at the specified index from a {@link VectorSchemaRoot} into a {@link + * FieldValueList}. + * + * @param root the VectorSchemaRoot containing column vectors + * @param rowIndex the 0-based row index to extract + * @param schema the BigQuery schema corresponding to the vectors + * @return the extracted FieldValueList row object + * @throws IllegalArgumentException if vector count does not match schema field count + */ + static FieldValueList arrowRootToFieldValueList( + VectorSchemaRoot root, int rowIndex, Schema schema) { + if (root.getFieldVectors().size() != schema.getFields().size()) { + throw new IllegalArgumentException( + String.format( + "Schema mismatch: Arrow vector count (%d) does not match BigQuery schema field count (%d)", + root.getFieldVectors().size(), schema.getFields().size())); + } + List fieldValues = new ArrayList<>(); + for (int colIndex = 0; colIndex < root.getFieldVectors().size(); colIndex++) { + FieldVector vector = root.getVector(colIndex); + Field bqField = schema.getFields().get(colIndex); + fieldValues.add(arrowVectorToFieldValue(vector, rowIndex, bqField)); + } + return FieldValueList.of(fieldValues, schema.getFields()); + } + + /** + * Converts a single cell value within a {@link FieldVector} to a BigQuery {@link FieldValue}. + * + *

Handles null values, repeated list vectors, nested struct vectors, and primitive type + * conversions. + * + * @param vector the Arrow column vector + * @param rowIndex the 0-based row index + * @param bqField the corresponding BigQuery Field definition + * @return the converted FieldValue object + */ + private static FieldValue arrowVectorToFieldValue( + FieldVector vector, int rowIndex, Field bqField) { + if (vector.isNull(rowIndex)) { + return FieldValue.of(FieldValue.Attribute.PRIMITIVE, null); + } + + // Handle repeated fields + if (bqField.getMode() == Field.Mode.REPEATED) { + ListVector listVector = (ListVector) vector; + FieldVector dataVector = (FieldVector) listVector.getDataVector(); + int start = listVector.getElementStartIndex(rowIndex); + int end = listVector.getElementEndIndex(rowIndex); + List elements = new ArrayList<>(end - start); + Field.Builder elementBuilder = Field.newBuilder(bqField.getName(), bqField.getType()); + if (bqField.getType() == LegacySQLTypeName.RECORD && bqField.getSubFields() != null) { + elementBuilder.setType(LegacySQLTypeName.RECORD, bqField.getSubFields()); + } + Field elementBqField = elementBuilder.setMode(Field.Mode.NULLABLE).build(); + for (int k = start; k < end; k++) { + elements.add(arrowVectorToFieldValue(dataVector, k, elementBqField)); + } + return FieldValue.of( + FieldValue.Attribute.REPEATED, FieldValueList.of(elements, bqField.getSubFields())); + } + + // Handle RECORD/STRUCT fields + if (bqField.getType() == LegacySQLTypeName.RECORD) { + StructVector structVector = (StructVector) vector; + if (structVector.size() != bqField.getSubFields().size()) { + throw new IllegalArgumentException( + String.format( + "Schema mismatch for field '%s': Arrow struct size (%d) does not match BigQuery subfields size (%d)", + bqField.getName(), structVector.size(), bqField.getSubFields().size())); + } + List elements = new ArrayList<>(structVector.size()); + for (int colIndex = 0; colIndex < structVector.size(); colIndex++) { + FieldVector childVector = (FieldVector) structVector.getChildByOrdinal(colIndex); + Field childBqField = bqField.getSubFields().get(colIndex); + elements.add(arrowVectorToFieldValue(childVector, rowIndex, childBqField)); + } + return FieldValue.of( + FieldValue.Attribute.RECORD, FieldValueList.of(elements, bqField.getSubFields())); + } + + // Handle primitive types + String stringVal; + if (bqField.getType() == LegacySQLTypeName.TIMESTAMP) { + TimeStampVector tsVector = (TimeStampVector) vector; + long rawVal = tsVector.get(rowIndex); + ArrowType.Timestamp tsType = (ArrowType.Timestamp) vector.getField().getType(); + long micros; + switch (tsType.getUnit()) { + case SECOND: + micros = rawVal * 1_000_000L; + break; + case MILLISECOND: + micros = rawVal * 1_000L; + break; + case MICROSECOND: + micros = rawVal; + break; + case NANOSECOND: + micros = rawVal / 1_000L; + break; + default: + micros = rawVal; + } + long seconds = micros / 1_000_000L; + long remainingMicros = Math.abs(micros % 1_000_000L); + if (micros < 0 && seconds == 0) { + stringVal = String.format(Locale.US, "-0.%06d", remainingMicros); + } else { + stringVal = String.format(Locale.US, "%d.%06d", seconds, remainingMicros); + } + } else { + Object value = vector.getObject(rowIndex); + if (value instanceof byte[]) { + stringVal = BaseEncoding.base64().encode((byte[]) value); + } else { + stringVal = String.valueOf(value); + } + } + + return FieldValue.of(FieldValue.Attribute.PRIMITIVE, stringVal); + } +} 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..8835ebcb682d --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowPojoUtils.java @@ -0,0 +1,169 @@ +/* + * 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.cloud.bigquery.Field.Mode; +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; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.Schema; + +/** + * Internal helper utility for converting Apache Arrow POJO Schema and Field definitions into + * BigQuery Veneer {@link Schema} and {@link com.google.cloud.bigquery.Field} models. + */ +final class ArrowPojoUtils { + + private ArrowPojoUtils() {} + + /** + * Converts an Apache Arrow {@link Schema} into a BigQuery Veneer {@link Schema}. + * + * @param arrowSchema the Apache Arrow schema definition + * @return the corresponding BigQuery Veneer Schema + */ + static com.google.cloud.bigquery.Schema arrowSchemaToBigQuerySchema(Schema arrowSchema) { + List fields = new ArrayList<>(); + for (Field arrowField : arrowSchema.getFields()) { + fields.add(arrowFieldToBigQueryField(arrowField)); + } + return com.google.cloud.bigquery.Schema.of(fields); + } + + /** + * Recursively converts an Apache Arrow {@link Field} into a BigQuery Veneer {@link + * com.google.cloud.bigquery.Field}. + * + *

Handles primitive types, repeated/list types, and nested struct/record types. + * + * @param arrowField the Apache Arrow field definition + * @return the corresponding BigQuery Veneer Field + * @throws IllegalArgumentException if an Arrow List field contains no child elements + */ + static com.google.cloud.bigquery.Field arrowFieldToBigQueryField(Field arrowField) { + String name = arrowField.getName(); + ArrowType type = arrowField.getType(); + com.google.cloud.bigquery.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); + } + Field innerField = arrowField.getChildren().get(0); + if (!innerField.getChildren().isEmpty()) { + List subFields = new ArrayList<>(); + for (Field childField : innerField.getChildren()) { + subFields.add(arrowFieldToBigQueryField(childField)); + } + builder = + com.google.cloud.bigquery.Field.newBuilder( + name, LegacySQLTypeName.RECORD, FieldList.of(subFields)); + } else { + LegacySQLTypeName innerType = arrowTypeToLegacySQLTypeName(innerField.getType()); + builder = com.google.cloud.bigquery.Field.newBuilder(name, innerType); + } + builder.setMode(Mode.REPEATED); + } else { + if (!arrowField.getChildren().isEmpty()) { + List subFields = new ArrayList<>(); + for (Field childField : arrowField.getChildren()) { + subFields.add(arrowFieldToBigQueryField(childField)); + } + builder = + com.google.cloud.bigquery.Field.newBuilder( + name, LegacySQLTypeName.RECORD, FieldList.of(subFields)); + } else { + LegacySQLTypeName bqType = arrowTypeToLegacySQLTypeName(type); + builder = com.google.cloud.bigquery.Field.newBuilder(name, bqType); + } + if (arrowField.isNullable()) { + builder.setMode(Mode.NULLABLE); + } else { + builder.setMode(Mode.REQUIRED); + } + } + return builder.build(); + } + + /** + * Instantiates a list of {@link FieldVector} instances corresponding to the fields in the + * provided Arrow schema using the specified allocator. + * + *

Guarantees exception-safe LIFO cleanup of already-allocated vectors if an allocation fails + * halfway through. + * + * @param arrowSchema the Apache Arrow schema definition + * @param allocator the buffer allocator to allocate vector memory from + * @return the list of allocated FieldVector instances + */ + static List createVectors(Schema arrowSchema, BufferAllocator allocator) { + List vectors = new ArrayList<>(); + try { + for (Field field : arrowSchema.getFields()) { + vectors.add(field.createVector(allocator)); + } + return vectors; + } catch (Throwable t) { + for (int i = vectors.size() - 1; i >= 0; i--) { + try { + vectors.get(i).close(); + } catch (Exception e) { + t.addSuppressed(e); + } + } + throw t; + } + } + + /** + * Maps an Apache {@link ArrowType} to its corresponding BigQuery {@link LegacySQLTypeName}. + * + * @param type the Apache Arrow type + * @return the matching 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()); + } + } +} 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 new file mode 100644 index 000000000000..585427d0ca1b --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java @@ -0,0 +1,366 @@ +/* + * 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.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assertions.fail; + +import com.google.cloud.bigquery.storage.v1.ReadRowsResponse; +import com.google.common.collect.ImmutableList; +import com.google.common.io.BaseEncoding; +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.ArrayList; +import java.util.Arrays; +import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.BitVector; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.Float8Vector; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.TimeStampMicroVector; +import org.apache.arrow.vector.VarBinaryVector; +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.TimeUnit; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.junit.jupiter.api.Test; + +public class ArrowDeserializerTest { + + @Test + public void testArrowSchemaToBigQuerySchema() { + 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); + + org.apache.arrow.vector.types.pojo.Schema arrowSchema = + new org.apache.arrow.vector.types.pojo.Schema( + ImmutableList.of(intField, strField, boolField, tsField)); + + 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(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(Field.Mode.REQUIRED, bqSchema.getFields().get(1).getMode()); + + assertEquals("bool_col", bqSchema.getFields().get(2).getName()); + assertEquals(LegacySQLTypeName.BOOLEAN, bqSchema.getFields().get(2).getType()); + + assertEquals("ts_col", bqSchema.getFields().get(3).getName()); + assertEquals(LegacySQLTypeName.TIMESTAMP, bqSchema.getFields().get(3).getType()); + } + + @Test + public void testDeserializeRecordBatchPrimitives() throws IOException { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + IntVector intVector = new IntVector("id", allocator); + intVector.allocateNew(2); + intVector.set(0, 101); + intVector.set(1, 102); + intVector.setValueCount(2); + + VarCharVector nameVector = new VarCharVector("name", allocator); + nameVector.allocateNew(2); + nameVector.set(0, "Alice".getBytes(StandardCharsets.UTF_8)); + nameVector.set(1, "Bob".getBytes(StandardCharsets.UTF_8)); + nameVector.setValueCount(2); + + Float8Vector scoreVector = new Float8Vector("score", allocator); + scoreVector.allocateNew(2); + scoreVector.set(0, 95.5); + scoreVector.setNull(1); + scoreVector.setValueCount(2); + + BitVector activeVector = new BitVector("active", allocator); + activeVector.allocateNew(2); + activeVector.set(0, 1); + activeVector.set(1, 0); + activeVector.setValueCount(2); + + VarBinaryVector bytesVector = new VarBinaryVector("data", allocator); + bytesVector.allocateNew(2); + bytesVector.set(0, "test_bytes".getBytes(StandardCharsets.UTF_8)); + bytesVector.setNull(1); + bytesVector.setValueCount(2); + + TimeStampMicroVector tsVector = new TimeStampMicroVector("ts", allocator); + tsVector.allocateNew(2); + // 1408452095220000 microsecond timestamp -> "1408452095.220000" + tsVector.set(0, 1408452095220000L); + tsVector.setNull(1); + tsVector.setValueCount(2); + + List vectors = + ImmutableList.of(intVector, nameVector, scoreVector, activeVector, bytesVector, tsVector); + + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + org.apache.arrow.vector.types.pojo.Schema arrowSchema = root.getSchema(); + Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); + + byte[] recordBatchBytes = serializeVectorSchemaRoot(root, allocator); + + List rows = + ArrowDeserializer.deserializeRecordBatch(recordBatchBytes, bqSchema, arrowSchema); + + assertEquals(2, rows.size()); + + // Row 0 + FieldValueList row0 = rows.get(0); + assertEquals("101", row0.get("id").getStringValue()); + assertEquals("Alice", row0.get("name").getStringValue()); + assertEquals("95.5", row0.get("score").getStringValue()); + assertEquals("true", row0.get("active").getStringValue()); + assertEquals( + BaseEncoding.base64().encode("test_bytes".getBytes(StandardCharsets.UTF_8)), + row0.get("data").getStringValue()); + assertEquals("1408452095.220000", row0.get("ts").getStringValue()); + + // Row 1 + FieldValueList row1 = rows.get(1); + assertEquals("102", row1.get("id").getStringValue()); + assertEquals("Bob", row1.get("name").getStringValue()); + assertNull(row1.get("score").getValue()); + assertEquals("false", row1.get("active").getStringValue()); + assertNull(row1.get("data").getValue()); + assertNull(row1.get("ts").getValue()); + } finally { + for (FieldVector vector : vectors) { + vector.close(); + } + } + } + } + + @Test + public void testSchemaMismatchThrowsException() { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + IntVector intVector = new IntVector("col1", allocator); + intVector.allocateNew(1); + intVector.set(0, 1); + intVector.setValueCount(1); + + try (VectorSchemaRoot root = new VectorSchemaRoot(ImmutableList.of(intVector))) { + Schema mismatchedSchema = + Schema.of( + Field.of("col1", LegacySQLTypeName.INTEGER), + Field.of("col2", LegacySQLTypeName.STRING)); + + try { + ArrowDeserializer.arrowRootToFieldValueList(root, 0, mismatchedSchema); + fail("Expected IllegalArgumentException on schema size mismatch"); + } catch (IllegalArgumentException e) { + // Expected + } + } finally { + intVector.close(); + } + } + } + + @Test + public void testLoadArrowRows_multiBatchStream() throws IOException { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + ReadRowsResponse r1 = + createReadRowsResponse(Arrays.asList(1, 2), Arrays.asList("item1", "item2"), allocator); + ReadRowsResponse r2 = + createReadRowsResponse(Arrays.asList(3, 4), Arrays.asList("item3", "item4"), allocator); + + org.apache.arrow.vector.types.pojo.Schema arrowSchema = createSimpleArrowSchema(); + Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); + + List rowBatch = new ArrayList<>(); + boolean hasMore = + ArrowDeserializer.loadArrowRows( + Arrays.asList(r1, r2).iterator(), arrowSchema, bqSchema, rowBatch, 10L, 0L, 10L); + + assertFalse(hasMore); + assertEquals(4, rowBatch.size()); + assertEquals("1", rowBatch.get(0).get("id").getStringValue()); + assertEquals("item1", rowBatch.get(0).get("name").getStringValue()); + assertEquals("4", rowBatch.get(3).get("id").getStringValue()); + assertEquals("item4", rowBatch.get(3).get("name").getStringValue()); + } + } + + @Test + public void testLoadArrowRows_respectsPageSize() throws IOException { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + ReadRowsResponse r1 = + createReadRowsResponse(Arrays.asList(1, 2), Arrays.asList("item1", "item2"), allocator); + ReadRowsResponse r2 = + createReadRowsResponse(Arrays.asList(3, 4), Arrays.asList("item3", "item4"), allocator); + + org.apache.arrow.vector.types.pojo.Schema arrowSchema = createSimpleArrowSchema(); + Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); + + List rowBatch = new ArrayList<>(); + boolean hasMore = + ArrowDeserializer.loadArrowRows( + Arrays.asList(r1, r2).iterator(), arrowSchema, bqSchema, rowBatch, 2L, 0L, 10L); + + assertTrue(hasMore); + assertEquals(2, rowBatch.size()); + assertEquals("1", rowBatch.get(0).get("id").getStringValue()); + assertEquals("2", rowBatch.get(1).get("id").getStringValue()); + } + } + + @Test + public void testLoadArrowRows_respectsMaxResults() throws IOException { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + ReadRowsResponse r1 = + createReadRowsResponse(Arrays.asList(1, 2), Arrays.asList("item1", "item2"), allocator); + ReadRowsResponse r2 = + createReadRowsResponse(Arrays.asList(3, 4), Arrays.asList("item3", "item4"), allocator); + + org.apache.arrow.vector.types.pojo.Schema arrowSchema = createSimpleArrowSchema(); + Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); + + List rowBatch = new ArrayList<>(); + boolean hasMore = + ArrowDeserializer.loadArrowRows( + Arrays.asList(r1, r2).iterator(), arrowSchema, bqSchema, rowBatch, 10L, 0L, 3L); + + assertFalse(hasMore); + assertEquals(3, rowBatch.size()); + assertEquals("1", rowBatch.get(0).get("id").getStringValue()); + assertEquals("2", rowBatch.get(1).get("id").getStringValue()); + assertEquals("3", rowBatch.get(2).get("id").getStringValue()); + } + } + + @Test + public void testLoadArrowRows_unconsumedBatchRowsSignalHasMore() throws IOException { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + ReadRowsResponse r1 = + createReadRowsResponse( + Arrays.asList(1, 2, 3, 4), + Arrays.asList("item1", "item2", "item3", "item4"), + allocator); + + org.apache.arrow.vector.types.pojo.Schema arrowSchema = createSimpleArrowSchema(); + Schema bqSchema = ArrowDeserializer.arrowSchemaToBigQuerySchema(arrowSchema); + + List rowBatch = new ArrayList<>(); + boolean hasMore = + ArrowDeserializer.loadArrowRows( + Arrays.asList(r1).iterator(), arrowSchema, bqSchema, rowBatch, 2L, 0L, 10L); + + assertTrue(hasMore); + assertEquals(2, rowBatch.size()); + assertEquals("1", rowBatch.get(0).get("id").getStringValue()); + assertEquals("2", rowBatch.get(1).get("id").getStringValue()); + } + } + + @Test + public void testLoadArrowRows_nullSchemaReturnsFalse() throws IOException { + List rowBatch = new ArrayList<>(); + boolean hasMore = + ArrowDeserializer.loadArrowRows( + Arrays.asList().iterator(), + null, + Schema.of(), + rowBatch, + 10L, + 0L, + 10L); + assertFalse(hasMore); + } + + private static org.apache.arrow.vector.types.pojo.Schema createSimpleArrowSchema() { + org.apache.arrow.vector.types.pojo.Field intField = + new org.apache.arrow.vector.types.pojo.Field( + "id", FieldType.nullable(new ArrowType.Int(32, true)), null); + org.apache.arrow.vector.types.pojo.Field strField = + new org.apache.arrow.vector.types.pojo.Field( + "name", FieldType.nullable(new ArrowType.Utf8()), null); + return new org.apache.arrow.vector.types.pojo.Schema(ImmutableList.of(intField, strField)); + } + + private ReadRowsResponse createReadRowsResponse( + List ids, List names, BufferAllocator allocator) throws IOException { + IntVector intVector = new IntVector("id", allocator); + intVector.allocateNew(ids.size()); + for (int i = 0; i < ids.size(); i++) { + intVector.set(i, ids.get(i)); + } + intVector.setValueCount(ids.size()); + + VarCharVector nameVector = new VarCharVector("name", allocator); + nameVector.allocateNew(names.size()); + for (int i = 0; i < names.size(); i++) { + nameVector.set(i, names.get(i).getBytes(StandardCharsets.UTF_8)); + } + nameVector.setValueCount(names.size()); + + List vectors = ImmutableList.of(intVector, nameVector); + try (VectorSchemaRoot root = new VectorSchemaRoot(vectors)) { + byte[] bytes = serializeVectorSchemaRoot(root, allocator); + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch protoBatch = + com.google.cloud.bigquery.storage.v1.ArrowRecordBatch.newBuilder() + .setSerializedRecordBatch(ByteString.copyFrom(bytes)) + .build(); + return ReadRowsResponse.newBuilder().setArrowRecordBatch(protoBatch).build(); + } finally { + for (FieldVector vector : vectors) { + vector.close(); + } + } + } + + private byte[] serializeVectorSchemaRoot(VectorSchemaRoot root, BufferAllocator allocator) + throws IOException { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + WriteChannel channel = new WriteChannel(Channels.newChannel(out)); + + VectorUnloader unloader = new VectorUnloader(root); + try (ArrowRecordBatch batch = unloader.getRecordBatch()) { + MessageSerializer.serialize(channel, batch); + } + return out.toByteArray(); + } +} diff --git a/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowPojoUtilsTest.java b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowPojoUtilsTest.java new file mode 100644 index 000000000000..b040a0e8fd7d --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowPojoUtilsTest.java @@ -0,0 +1,208 @@ +/* + * 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.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.google.cloud.bigquery.Field.Mode; +import com.google.common.collect.ImmutableList; +import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.FieldVector; +import org.apache.arrow.vector.types.DateUnit; +import org.apache.arrow.vector.types.FloatingPointPrecision; +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 ArrowPojoUtilsTest { + + @Test + public void testArrowSchemaToBigQuerySchema_Primitives() { + Field intField = new Field("int_col", FieldType.nullable(new ArrowType.Int(64, 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 bytesField = new Field("bytes_col", FieldType.nullable(new ArrowType.Binary()), null); + Field floatField = + new Field( + "float_col", + FieldType.nullable(new ArrowType.FloatingPoint(FloatingPointPrecision.DOUBLE)), + null); + Field decimalField = + new Field("num_col", FieldType.nullable(new ArrowType.Decimal(38, 9, 128)), null); + Field dateField = + new Field("date_col", FieldType.nullable(new ArrowType.Date(DateUnit.DAY)), null); + Field timeField = + new Field( + "time_col", FieldType.nullable(new ArrowType.Time(TimeUnit.MICROSECOND, 64)), null); + Field tsField = + new Field( + "ts_col", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")), + null); + + Schema arrowSchema = + new Schema( + ImmutableList.of( + intField, + strField, + boolField, + bytesField, + floatField, + decimalField, + dateField, + timeField, + tsField)); + + com.google.cloud.bigquery.Schema bqSchema = + ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema); + + assertEquals(9, bqSchema.getFields().size()); + assertEquals(LegacySQLTypeName.INTEGER, bqSchema.getFields().get(0).getType()); + assertEquals(Mode.NULLABLE, bqSchema.getFields().get(0).getMode()); + + assertEquals(LegacySQLTypeName.STRING, bqSchema.getFields().get(1).getType()); + assertEquals(Mode.REQUIRED, bqSchema.getFields().get(1).getMode()); + + assertEquals(LegacySQLTypeName.BOOLEAN, bqSchema.getFields().get(2).getType()); + assertEquals(LegacySQLTypeName.BYTES, bqSchema.getFields().get(3).getType()); + assertEquals(LegacySQLTypeName.FLOAT, bqSchema.getFields().get(4).getType()); + assertEquals(LegacySQLTypeName.NUMERIC, bqSchema.getFields().get(5).getType()); + assertEquals(LegacySQLTypeName.DATE, bqSchema.getFields().get(6).getType()); + assertEquals(LegacySQLTypeName.TIME, bqSchema.getFields().get(7).getType()); + assertEquals(LegacySQLTypeName.TIMESTAMP, bqSchema.getFields().get(8).getType()); + } + + @Test + public void testArrowSchemaToBigQuerySchema_NestedStruct() { + Field innerInt = new Field("id", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field innerStr = new Field("name", FieldType.nullable(new ArrowType.Utf8()), null); + Field structField = + new Field( + "person", + FieldType.nullable(new ArrowType.Struct()), + ImmutableList.of(innerInt, innerStr)); + + Schema arrowSchema = new Schema(ImmutableList.of(structField)); + com.google.cloud.bigquery.Schema bqSchema = + ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema); + + assertEquals(1, bqSchema.getFields().size()); + com.google.cloud.bigquery.Field personField = bqSchema.getFields().get(0); + assertEquals("person", personField.getName()); + assertEquals(LegacySQLTypeName.RECORD, personField.getType()); + assertEquals(2, personField.getSubFields().size()); + assertEquals("id", personField.getSubFields().get(0).getName()); + assertEquals(LegacySQLTypeName.INTEGER, personField.getSubFields().get(0).getType()); + assertEquals("name", personField.getSubFields().get(1).getName()); + assertEquals(LegacySQLTypeName.STRING, personField.getSubFields().get(1).getType()); + } + + @Test + public void testArrowSchemaToBigQuerySchema_ListPrimitives() { + Field itemField = new Field("item", FieldType.notNullable(new ArrowType.Utf8()), null); + Field listField = + new Field("tags", FieldType.nullable(new ArrowType.List()), ImmutableList.of(itemField)); + + Schema arrowSchema = new Schema(ImmutableList.of(listField)); + com.google.cloud.bigquery.Schema bqSchema = + ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema); + + assertEquals(1, bqSchema.getFields().size()); + com.google.cloud.bigquery.Field tagsField = bqSchema.getFields().get(0); + assertEquals("tags", tagsField.getName()); + assertEquals(LegacySQLTypeName.STRING, tagsField.getType()); + assertEquals(Mode.REPEATED, tagsField.getMode()); + } + + @Test + public void testArrowSchemaToBigQuerySchema_ListOfStruct() { + Field innerKey = new Field("key", FieldType.nullable(new ArrowType.Utf8()), null); + Field innerVal = new Field("value", FieldType.nullable(new ArrowType.Int(64, true)), null); + Field structField = + new Field( + "item", + FieldType.nullable(new ArrowType.Struct()), + ImmutableList.of(innerKey, innerVal)); + Field listField = + new Field( + "entries", FieldType.nullable(new ArrowType.List()), ImmutableList.of(structField)); + + Schema arrowSchema = new Schema(ImmutableList.of(listField)); + com.google.cloud.bigquery.Schema bqSchema = + ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema); + + assertEquals(1, bqSchema.getFields().size()); + com.google.cloud.bigquery.Field entriesField = bqSchema.getFields().get(0); + assertEquals("entries", entriesField.getName()); + assertEquals(LegacySQLTypeName.RECORD, entriesField.getType()); + assertEquals(Mode.REPEATED, entriesField.getMode()); + assertEquals(2, entriesField.getSubFields().size()); + assertEquals("key", entriesField.getSubFields().get(0).getName()); + assertEquals("value", entriesField.getSubFields().get(1).getName()); + } + + @Test + public void testArrowSchemaToBigQuerySchema_EmptyListThrowsException() { + Field emptyList = + new Field("empty_list", FieldType.nullable(new ArrowType.List()), ImmutableList.of()); + Schema arrowSchema = new Schema(ImmutableList.of(emptyList)); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema)); + assertTrue(thrown.getMessage().contains("must have at least one child field")); + } + + @Test + public void testArrowSchemaToBigQuerySchema_UnsupportedTypeThrowsException() { + Field unsupportedField = + new Field("unsupported", FieldType.nullable(new ArrowType.Null()), null); + Schema arrowSchema = new Schema(ImmutableList.of(unsupportedField)); + + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> ArrowPojoUtils.arrowSchemaToBigQuerySchema(arrowSchema)); + assertTrue(thrown.getMessage().contains("Unsupported Arrow type")); + } + + @Test + public void testCreateVectors_Success() { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + Field intField = new Field("int_col", FieldType.nullable(new ArrowType.Int(32, true)), null); + Field strField = new Field("str_col", FieldType.nullable(new ArrowType.Utf8()), null); + Schema arrowSchema = new Schema(ImmutableList.of(intField, strField)); + + List vectors = ArrowPojoUtils.createVectors(arrowSchema, allocator); + assertEquals(2, vectors.size()); + assertEquals("int_col", vectors.get(0).getName()); + assertEquals("str_col", vectors.get(1).getName()); + + for (FieldVector v : vectors) { + v.close(); + } + } + } +}