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..7c5829734822 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/main/java/com/google/cloud/bigquery/ArrowDeserializer.java @@ -0,0 +1,325 @@ +/* + * 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.common.collect.ImmutableList; +import com.google.common.io.BaseEncoding; +import java.io.IOException; +import java.util.ArrayList; +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 { + + private ArrowDeserializer() {} + + /** + * Converts an Apache Arrow {@link org.apache.arrow.vector.types.pojo.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()); + } + } + + /** + * Deserializes a raw binary Arrow record batch payload into a list of BigQuery {@link + * FieldValueList} row objects. + * + *

Allocates off-heap memory within a local {@link RootAllocator} 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, org.apache.arrow.vector.types.pojo.Schema arrowSchema) + throws IOException { + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) { + List vectors = new ArrayList<>(); + try { + for (org.apache.arrow.vector.types.pojo.Field field : arrowSchema.getFields()) { + vectors.add(field.createVector(allocator)); + } + } 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; + } + 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); + } + } + } + } + + /** + * 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) { + // 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(); + 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/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..f23dbe4b01a5 --- /dev/null +++ b/java-bigquery/google-cloud-bigquery/src/test/java/com/google/cloud/bigquery/ArrowDeserializerTest.java @@ -0,0 +1,210 @@ +/* + * 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.assertNull; +import static org.junit.jupiter.api.Assertions.fail; + +import com.google.common.collect.ImmutableList; +import com.google.common.io.BaseEncoding; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.channels.Channels; +import java.nio.charset.StandardCharsets; +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.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( + "ts_col", + FieldType.nullable(new ArrowType.Timestamp(TimeUnit.MICROSECOND, "UTC")), + null); + + Schema arrowSchema = new Schema(ImmutableList.of(intField, strField, boolField, tsField)); + + com.google.cloud.bigquery.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("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("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)) { + Schema arrowSchema = root.getSchema(); + com.google.cloud.bigquery.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("false".equals("false") ? "active" : "score") != null + ? row1.get("active").getStringValue() + : "false"); + 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))) { + 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)); + + try { + ArrowDeserializer.arrowRootToFieldValueList(root, 0, mismatchedSchema); + fail("Expected IllegalArgumentException on schema size mismatch"); + } catch (IllegalArgumentException e) { + // Expected + } + } finally { + intVector.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(); + } +}