diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorHolder.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorHolder.java index f8c0d6dd69b8..2b749256b0d9 100644 --- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorHolder.java +++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorHolder.java @@ -185,4 +185,34 @@ public int numValues() { return numRows; } } + + public static class VariantVectorHolder extends VectorHolder { + private final VectorHolder metadataHolder; + private final VectorHolder valueHolder; + private final int numRows; + + public VariantVectorHolder( + Types.NestedField icebergField, + int numRows, + VectorHolder metadataHolder, + VectorHolder valueHolder) { + super(icebergField); + this.numRows = numRows; + this.metadataHolder = metadataHolder; + this.valueHolder = valueHolder; + } + + @Override + public int numValues() { + return numRows; + } + + public VectorHolder metadataHolder() { + return metadataHolder; + } + + public VectorHolder valueHolder() { + return valueHolder; + } + } } diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java index 11c73aea4a44..dce8741438fc 100644 --- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java +++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedArrowReader.java @@ -965,4 +965,57 @@ public String toString() { @Override public void setBatchSize(int batchSize) {} } + + public static class VectorizedVariantReader extends VectorizedArrowReader { + private final VectorizedArrowReader metadataReader; + private final VectorizedArrowReader valueReader; + + public VectorizedVariantReader( + Types.NestedField icebergField, + VectorizedArrowReader metadataReader, + VectorizedArrowReader valueReader) { + super(icebergField); + this.metadataReader = metadataReader; + this.valueReader = valueReader; + } + + @Override + public VectorHolder read(VectorHolder reuse, int numValsToRead) { + VectorHolder reuseMetadata = null; + VectorHolder reuseValue = null; + if (reuse instanceof VectorHolder.VariantVectorHolder) { + VectorHolder.VariantVectorHolder variantReuse = (VectorHolder.VariantVectorHolder) reuse; + reuseMetadata = variantReuse.metadataHolder(); + reuseValue = variantReuse.valueHolder(); + } + VectorHolder metadataHolder = metadataReader.read(reuseMetadata, numValsToRead); + VectorHolder valueHolder = valueReader.read(reuseValue, numValsToRead); + return new VectorHolder.VariantVectorHolder( + icebergField(), numValsToRead, metadataHolder, valueHolder); + } + + @Override + public void setRowGroupInfo( + PageReadStore source, Map metadata) { + metadataReader.setRowGroupInfo(source, metadata); + valueReader.setRowGroupInfo(source, metadata); + } + + @Override + public void setBatchSize(int batchSize) { + metadataReader.setBatchSize(batchSize); + valueReader.setBatchSize(batchSize); + } + + @Override + public void close() { + metadataReader.close(); + valueReader.close(); + } + + @Override + public String toString() { + return "VectorizedVariantReader"; + } + } } diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedReaderBuilder.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedReaderBuilder.java index 3fbd797c26fb..59023fe7c877 100644 --- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedReaderBuilder.java +++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedReaderBuilder.java @@ -27,6 +27,7 @@ import org.apache.iceberg.Schema; import org.apache.iceberg.arrow.ArrowAllocation; import org.apache.iceberg.arrow.vectorized.VectorizedArrowReader.ConstantVectorReader; +import org.apache.iceberg.parquet.ParquetVariantVisitor; import org.apache.iceberg.parquet.TypeWithSchemaVisitor; import org.apache.iceberg.parquet.VectorizedReader; import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; @@ -154,14 +155,16 @@ public VectorizedReader struct( return null; } + @Override + public ParquetVariantVisitor> variantVisitor() { + return new VectorizedVariantVisitor( + currentPath(), parquetSchema, icebergSchema, rootAllocator, setArrowValidityVector); + } + @Override public VectorizedReader variant( Types.VariantType iVariant, GroupType variant, VectorizedReader result) { - if (iVariant != null) { - throw new UnsupportedOperationException( - "Vectorized reads are not supported yet for variant fields"); - } - return null; + return result; } @Override diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedVariantVisitor.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedVariantVisitor.java new file mode 100644 index 000000000000..11d079e17766 --- /dev/null +++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/VectorizedVariantVisitor.java @@ -0,0 +1,142 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.iceberg.arrow.vectorized; + +import java.util.List; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.iceberg.Schema; +import org.apache.iceberg.parquet.ParquetVariantVisitor; +import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.types.Types; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.io.InvalidRecordException; +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType; + +/** Creates Arrow readers for a variant column's metadata and value leaves. */ +class VectorizedVariantVisitor extends ParquetVariantVisitor> { + private final String[] variantGroupPath; + private final MessageType parquetSchema; + private final Schema icebergSchema; + private final BufferAllocator allocator; + private final boolean setArrowValidityVector; + + VectorizedVariantVisitor( + String[] variantGroupPath, + MessageType parquetSchema, + Schema icebergSchema, + BufferAllocator allocator, + boolean setArrowValidityVector) { + this.variantGroupPath = variantGroupPath; + this.parquetSchema = parquetSchema; + this.icebergSchema = icebergSchema; + this.allocator = allocator; + this.setArrowValidityVector = setArrowValidityVector; + } + + @Override + public VectorizedReader metadata(PrimitiveType metadata) { + ColumnDescriptor desc = resolveDescriptor(metadata); + if (desc == null) { + return null; + } + + Types.NestedField field = + Types.NestedField.required(-1, metadata.getName(), Types.BinaryType.get()); + return new VectorizedArrowReader(desc, field, allocator, setArrowValidityVector); + } + + @Override + public VectorizedReader serialized(PrimitiveType value) { + ColumnDescriptor desc = resolveDescriptor(value); + if (desc == null) { + return null; + } + + Types.NestedField field = + Types.NestedField.optional(-1, value.getName(), Types.BinaryType.get()); + return new VectorizedArrowReader(desc, field, allocator, setArrowValidityVector); + } + + @Override + public VectorizedReader variant( + GroupType variant, VectorizedReader metadataResult, VectorizedReader valueResult) { + if (metadataResult instanceof VectorizedArrowReader + && valueResult instanceof VectorizedArrowReader) { + Types.NestedField field = findVariantField(variant); + if (field != null) { + return new VectorizedArrowReader.VectorizedVariantReader( + field, (VectorizedArrowReader) metadataResult, (VectorizedArrowReader) valueResult); + } + } + + return null; + } + + @Override + public VectorizedReader primitive(PrimitiveType primitive) { + throw new UnsupportedOperationException("Unsupported variant: shredded typed_value primitive"); + } + + @Override + public VectorizedReader value( + GroupType value, VectorizedReader valueResult, VectorizedReader typedResult) { + if (typedResult != null) { + throw new UnsupportedOperationException( + "Unsupported variant: shredded typed_value primitive"); + } + return valueResult; + } + + @Override + public VectorizedReader object( + GroupType object, VectorizedReader valueResult, List> fieldResults) { + throw new UnsupportedOperationException( + "Unsupported variant: shredded typed_value object with " + fieldResults.size() + " fields"); + } + + @Override + public VectorizedReader array( + GroupType array, VectorizedReader valueResult, VectorizedReader elementResult) { + throw new UnsupportedOperationException("Unsupported variant: shredded typed_value array"); + } + + private ColumnDescriptor resolveDescriptor(PrimitiveType primitive) { + // Build full column path: variant group path + primitive name + // e.g., ["v1", "metadata"] or ["v2", "value"] + String[] path = new String[variantGroupPath.length + 1]; + System.arraycopy(variantGroupPath, 0, path, 0, variantGroupPath.length); + path[variantGroupPath.length] = primitive.getName(); + + try { + return parquetSchema.getColumnDescription(path); + } catch (InvalidRecordException e) { + return null; + } + } + + private Types.NestedField findVariantField(GroupType variant) { + if (variant.getId() != null) { + return icebergSchema.findField(variant.getId().intValue()); + } + + return null; + } +} diff --git a/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedReaderBuilder.java b/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedReaderBuilder.java index e3d76515bcc7..148d0c0ecb82 100644 --- a/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedReaderBuilder.java +++ b/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedReaderBuilder.java @@ -19,7 +19,6 @@ package org.apache.iceberg.arrow.vectorized; import static org.assertj.core.api.Assertions.assertThatNoException; -import static org.assertj.core.api.Assertions.assertThatThrownBy; import org.apache.iceberg.Schema; import org.apache.iceberg.parquet.TypeWithSchemaVisitor; @@ -38,7 +37,7 @@ public class TestVectorizedReaderBuilder { @Test - public void testVariantNotSupportedInVectorizedReads() { + public void testVariantSupportedInVectorizedReads() { Schema icebergSchema = new Schema( NestedField.required(1, "id", IntegerType.get()), @@ -50,10 +49,9 @@ public void testVariantNotSupportedInVectorizedReads() { new VectorizedReaderBuilder( icebergSchema, parquetSchema, false, ImmutableMap.of(), readers -> null); - assertThatThrownBy( - () -> TypeWithSchemaVisitor.visit(icebergSchema.asStruct(), parquetSchema, builder)) - .isInstanceOf(UnsupportedOperationException.class) - .hasMessageContaining("Vectorized reads are not supported yet for variant fields"); + assertThatNoException() + .isThrownBy( + () -> TypeWithSchemaVisitor.visit(icebergSchema.asStruct(), parquetSchema, builder)); } @Test diff --git a/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedVariantVisitor.java b/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedVariantVisitor.java new file mode 100644 index 000000000000..cb3963f757c1 --- /dev/null +++ b/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedVariantVisitor.java @@ -0,0 +1,153 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.iceberg.arrow.vectorized; + +import static org.assertj.core.api.Assertions.assertThatNoException; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import org.apache.iceberg.Schema; +import org.apache.iceberg.parquet.ParquetVariantVisitor; +import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.types.Types.NestedField; +import org.apache.iceberg.types.Types.VariantType; +import org.apache.iceberg.variants.Variant; +import org.apache.parquet.schema.GroupType; +import org.apache.parquet.schema.LogicalTypeAnnotation; +import org.apache.parquet.schema.MessageType; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import org.apache.parquet.schema.Type; +import org.apache.parquet.schema.Types; +import org.junit.jupiter.api.Test; + +public class TestVectorizedVariantVisitor { + + private static final Schema ICEBERG_SCHEMA = + new Schema(NestedField.optional(1, "data", VariantType.get())); + + @Test + public void testUnshreddedVariantDoesNotThrow() { + MessageType schema = variantSchema(unshreddedGroup()); + assertThatNoException().isThrownBy(() -> visit(schema)); + } + + @Test + public void testShreddedPrimitiveThrows() { + MessageType schema = variantSchema(shreddedPrimitiveGroup()); + assertThatThrownBy(() -> visit(schema)) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessage("Unsupported variant: shredded typed_value primitive"); + } + + @Test + public void testShreddedObjectThrows() { + GroupType variantGroup = shreddedObjectGroup(); + int fieldCount = variantGroup.getType("typed_value").asGroupType().getFieldCount(); + MessageType schema = variantSchema(variantGroup); + assertThatThrownBy(() -> visit(schema)) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessage( + "Unsupported variant: shredded typed_value object with " + fieldCount + " fields"); + } + + @Test + public void testShreddedArrayThrows() { + MessageType schema = variantSchema(shreddedArrayGroup()); + assertThatThrownBy(() -> visit(schema)) + .isInstanceOf(UnsupportedOperationException.class) + .hasMessage("Unsupported variant: shredded typed_value array"); + } + + private static VectorizedReader visit(MessageType schema) { + GroupType variantGroup = schema.getType("data").asGroupType(); + VectorizedVariantVisitor visitor = + new VectorizedVariantVisitor(new String[] {"data"}, schema, ICEBERG_SCHEMA, null, false); + return ParquetVariantVisitor.visit(variantGroup, visitor); + } + + private static MessageType variantSchema(GroupType variantGroup) { + return Types.buildMessage().addField(variantGroup).named("table"); + } + + private static GroupType unshreddedGroup() { + return Types.buildGroup(Type.Repetition.OPTIONAL) + .as(LogicalTypeAnnotation.variantType(Variant.VARIANT_SPEC_VERSION)) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.REQUIRED).named("metadata")) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.OPTIONAL).named("value")) + .id(1) + .named("data"); + } + + private static GroupType shreddedPrimitiveGroup() { + return Types.buildGroup(Type.Repetition.OPTIONAL) + .as(LogicalTypeAnnotation.variantType(Variant.VARIANT_SPEC_VERSION)) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.REQUIRED).named("metadata")) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.OPTIONAL).named("value")) + .addField( + Types.primitive(PrimitiveTypeName.INT32, Type.Repetition.OPTIONAL).named("typed_value")) + .id(1) + .named("data"); + } + + private static GroupType shreddedObjectGroup() { + GroupType typedField = + Types.buildGroup(Type.Repetition.OPTIONAL) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.OPTIONAL).named("value")) + .named("field_a"); + GroupType typedValue = + Types.buildGroup(Type.Repetition.OPTIONAL).addField(typedField).named("typed_value"); + return Types.buildGroup(Type.Repetition.OPTIONAL) + .as(LogicalTypeAnnotation.variantType(Variant.VARIANT_SPEC_VERSION)) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.REQUIRED).named("metadata")) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.OPTIONAL).named("value")) + .addField(typedValue) + .id(1) + .named("data"); + } + + private static GroupType shreddedArrayGroup() { + GroupType elementValue = + Types.buildGroup(Type.Repetition.OPTIONAL) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.OPTIONAL).named("value")) + .named("element"); + GroupType repeated = + Types.buildGroup(Type.Repetition.REPEATED).addField(elementValue).named("list"); + GroupType typedValue = + Types.buildGroup(Type.Repetition.OPTIONAL) + .as(LogicalTypeAnnotation.listType()) + .addField(repeated) + .named("typed_value"); + return Types.buildGroup(Type.Repetition.OPTIONAL) + .as(LogicalTypeAnnotation.variantType(Variant.VARIANT_SPEC_VERSION)) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.REQUIRED).named("metadata")) + .addField( + Types.primitive(PrimitiveTypeName.BINARY, Type.Repetition.OPTIONAL).named("value")) + .addField(typedValue) + .id(1) + .named("data"); + } +} diff --git a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java index 61616a9f233c..2a14c2262c6a 100644 --- a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java +++ b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java @@ -27,7 +27,9 @@ class ColumnVectorBuilder { public ColumnVector build(VectorHolder holder, int numRows) { - if (holder.isDummy()) { + if (holder instanceof VectorHolder.VariantVectorHolder) { + return new VariantColumnVector((VectorHolder.VariantVectorHolder) holder); + } else if (holder.isDummy()) { if (holder instanceof VectorHolder.DeletedVectorHolder) { return new DeletedColumnVector(Types.BooleanType.get()); } else if (holder instanceof ConstantVectorHolder) { diff --git a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java index edaaaeda2515..f11b84443fd8 100644 --- a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java +++ b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java @@ -20,6 +20,7 @@ import org.apache.spark.sql.types.Decimal; import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.types.VariantType; import org.apache.spark.sql.vectorized.ColumnVector; import org.apache.spark.sql.vectorized.ColumnarArray; import org.apache.spark.sql.vectorized.ColumnarMap; @@ -143,6 +144,12 @@ public ColumnVector getChild(int ordinal) { for (int index = 0; index < structType.length(); index++) { children[index] = new ColumnVectorWithFilter(delegate.getChild(index), rowIdMapping); } + } else if (dataType() instanceof VariantType) { + this.children = + new ColumnVectorWithFilter[] { + new ColumnVectorWithFilter(delegate.getChild(0), rowIdMapping), + new ColumnVectorWithFilter(delegate.getChild(1), rowIdMapping) + }; } else { throw new UnsupportedOperationException("Unsupported nested type: " + dataType()); } diff --git a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VariantColumnVector.java b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VariantColumnVector.java new file mode 100644 index 000000000000..2831b1f83cdb --- /dev/null +++ b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VariantColumnVector.java @@ -0,0 +1,141 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.iceberg.spark.data.vectorized; + +import org.apache.iceberg.arrow.vectorized.VectorHolder; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Decimal; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarArray; +import org.apache.spark.sql.vectorized.ColumnarMap; +import org.apache.spark.unsafe.types.UTF8String; + +public class VariantColumnVector extends ColumnVector { + private final ColumnVector valueChild; + private final ColumnVector metadataChild; + + VariantColumnVector(VectorHolder.VariantVectorHolder holder) { + super(DataTypes.VariantType); + this.valueChild = new IcebergArrowColumnVector(holder.valueHolder()); + this.metadataChild = new IcebergArrowColumnVector(holder.metadataHolder()); + } + + @Override + public void close() { + valueChild.close(); + metadataChild.close(); + } + + @Override + public void closeIfFreeable() { + // See SPARK-50235, SPARK-50463 + } + + @Override + public boolean hasNull() { + return valueChild.hasNull(); + } + + @Override + public int numNulls() { + return valueChild.numNulls(); + } + + @Override + public boolean isNullAt(int rowId) { + return valueChild.isNullAt(rowId); + } + + // getChild is what getVariant() calls: child(0) = value, child(1) = metadata + @Override + public ColumnVector getChild(int ordinal) { + if (ordinal == 0) { + return valueChild; + } else if (ordinal == 1) { + return metadataChild; + } + + throw new IllegalArgumentException( + "Variant column has only 2 children, got ordinal: " + ordinal); + } + + @Override + public boolean getBoolean(int rowId) { + throw unsupported(); + } + + @Override + public byte getByte(int rowId) { + throw unsupported(); + } + + @Override + public short getShort(int rowId) { + throw unsupported(); + } + + @Override + public int getInt(int rowId) { + throw unsupported(); + } + + @Override + public long getLong(int rowId) { + throw unsupported(); + } + + @Override + public float getFloat(int rowId) { + throw unsupported(); + } + + @Override + public double getDouble(int rowId) { + throw unsupported(); + } + + @Override + public Decimal getDecimal(int rowId, int precision, int scale) { + throw unsupported(); + } + + @Override + public UTF8String getUTF8String(int rowId) { + throw unsupported(); + } + + @Override + public byte[] getBinary(int rowId) { + throw unsupported(); + } + + @Override + public ColumnarArray getArray(int rowId) { + throw unsupported(); + } + + @Override + public ColumnarMap getMap(int rowId) { + throw unsupported(); + } + + private UnsupportedOperationException unsupported() { + return new UnsupportedOperationException("Variant column only supports getVariant()"); + } +} diff --git a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index 2109936c96b9..f90423900d62 100644 --- a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -18,17 +18,23 @@ */ package org.apache.iceberg.spark.source; +import java.nio.ByteBuffer; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.function.Supplier; import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.MetricsConfig; +import org.apache.iceberg.MetricsModes; +import org.apache.iceberg.MetricsUtil; import org.apache.iceberg.ScanTask; import org.apache.iceberg.ScanTaskGroup; import org.apache.iceberg.Schema; import org.apache.iceberg.SchemaParser; import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; import org.apache.iceberg.io.FileIO; import org.apache.iceberg.spark.ImmutableOrcBatchReadConf; import org.apache.iceberg.spark.ImmutableParquetBatchReadConf; @@ -37,6 +43,7 @@ import org.apache.iceberg.spark.SparkReadConf; import org.apache.iceberg.spark.SparkUtil; import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.PropertyUtil; import org.apache.spark.api.java.JavaSparkContext; import org.apache.spark.broadcast.Broadcast; import org.apache.spark.sql.connector.read.Batch; @@ -149,7 +156,7 @@ private OrcBatchReadConf orcBatchReadConf() { // conditions for using Parquet batch reads: // - Parquet vectorization is enabled - // - only primitives or metadata columns are projected + // - only primitives, unshredded variant, or metadata columns are projected // - all tasks are of FileScanTask type and read only Parquet files private boolean useParquetBatchReads() { return readConf.parquetVectorizationEnabled() @@ -164,7 +171,18 @@ private boolean supportsParquetBatchReads(ScanTask task) { } else if (task.isFileScanTask() && !task.isDataTask()) { FileScanTask fileScanTask = task.asFileScanTask(); - return fileScanTask.file().format() == FileFormat.PARQUET; + if (fileScanTask.file().format() != FileFormat.PARQUET) { + return false; + } + Map lowerBounds = fileScanTask.file().lowerBounds(); + if (lowerBounds != null) { + for (Types.NestedField field : expectedSchema.columns()) { + if (field.type().isVariantType() && lowerBounds.containsKey(field.fieldId())) { + return false; + } + } + } + return true; } else { return false; @@ -172,7 +190,27 @@ private boolean supportsParquetBatchReads(ScanTask task) { } private boolean supportsParquetBatchReads(Types.NestedField field) { - return field.type().isPrimitiveType() || MetadataColumns.isMetadataColumn(field.fieldId()); + if (field.type().isVariantType()) { + boolean shredVariants = + PropertyUtil.propertyAsBoolean( + table.properties(), + TableProperties.PARQUET_SHRED_VARIANTS, + TableProperties.PARQUET_SHRED_VARIANTS_DEFAULT); + if (shredVariants) { + return false; + } + + MetricsConfig metricsConfig = MetricsConfig.forTable(table); + MetricsModes.MetricsMode mode = + MetricsUtil.metricsMode(table.schema(), metricsConfig, field.fieldId()); + if (mode == MetricsModes.None.get() || mode == MetricsModes.Counts.get()) { + return false; + } + } + + return field.type().isPrimitiveType() + || field.type().isVariantType() + || MetadataColumns.isMetadataColumn(field.fieldId()); } // conditions for using ORC batch reads: diff --git a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java index 8b75906a628f..79c2456fffd8 100644 --- a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java +++ b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java @@ -481,6 +481,11 @@ private org.apache.iceberg.Scan buildBatchScan( if (withStats) { scan = scan.includeColumnStats(); + } else { + List variantColumns = variantColumnNames(expectedSchema); + if (!variantColumns.isEmpty()) { + scan = scan.includeColumnStats(variantColumns); + } } if (snapshotId != null) { @@ -515,6 +520,11 @@ private org.apache.iceberg.Scan buildIncrementalAppendScan( if (withStats) { scan = scan.includeColumnStats(); + } else { + List variantColumns = variantColumnNames(expectedSchema); + if (!variantColumns.isEmpty()) { + scan = scan.includeColumnStats(variantColumns); + } } if (endSnapshotId != null) { @@ -524,6 +534,13 @@ private org.apache.iceberg.Scan buildIncrementalAppendScan( return configureSplitPlanning(scan); } + private List variantColumnNames(Schema projection) { + return projection.columns().stream() + .filter(field -> field.type().isVariantType()) + .map(Types.NestedField::name) + .collect(Collectors.toList()); + } + @SuppressWarnings("CyclomaticComplexity") public Scan buildChangelogScan() { Preconditions.checkArgument( diff --git a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkVariantRead.java b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkVariantRead.java index 2d6e919a91ee..79e52d6e339c 100644 --- a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkVariantRead.java +++ b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkVariantRead.java @@ -19,13 +19,33 @@ package org.apache.iceberg.spark.sql; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assumptions.assumeThat; +import java.io.File; +import java.io.IOException; import java.util.List; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.Files; +import org.apache.iceberg.SnapshotChanges; +import org.apache.iceberg.Table; +import org.apache.iceberg.data.FileHelpers; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.relocated.com.google.common.collect.Iterables; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.spark.Spark3Util; import org.apache.iceberg.spark.SparkCatalog; import org.apache.iceberg.spark.TestBase; +import org.apache.iceberg.util.CharSequenceSet; +import org.apache.iceberg.util.Pair; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.catalyst.parser.ParseException; import org.apache.spark.types.variant.Variant; import org.apache.spark.unsafe.types.VariantVal; import org.junit.jupiter.api.AfterEach; @@ -76,7 +96,6 @@ public void cleanup() { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testVariantColumnProjection_singleVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); setVectorization(vectorized); Dataset df = spark.table(TABLE).select("id", "v1").orderBy("id"); assertThat(df.schema().fieldNames()).containsExactly("id", "v1"); @@ -106,7 +125,6 @@ public void testVariantColumnProjection_singleVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testVariantColumnProjectionNoVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); setVectorization(vectorized); Dataset df = spark.table(TABLE).select("id"); assertThat(df.schema().fieldNames()).containsExactly("id"); @@ -117,7 +135,6 @@ public void testVariantColumnProjectionNoVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testFilterOnVariantColumnOnWholeValue(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); setVectorization(vectorized); sql("INSERT INTO %s SELECT 3, NULL, NULL", TABLE); @@ -147,7 +164,6 @@ public void testFilterOnVariantColumnOnWholeValue(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testVariantNullValueProjection(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); setVectorization(vectorized); // insert a row with NULL variant values @@ -162,10 +178,65 @@ public void testVariantNullValueProjection(boolean vectorized) { assertThat(row.isNullAt(1)).isTrue(); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + public void testVariantReadAfterDelete(boolean vectorized) + throws IOException, NoSuchTableException, ParseException { + String deleteTable = CATALOG + ".default.var_delete"; + + sql("DROP TABLE IF EXISTS %s", deleteTable); + sql( + "CREATE TABLE %s (id BIGINT, v1 VARIANT) USING iceberg " + + "TBLPROPERTIES ('format-version'='3')", + deleteTable); + setVectorization(deleteTable, vectorized); + + spark + .sql( + "SELECT 1L AS id, parse_json('{\"a\":1}') AS v1 " + + "UNION ALL SELECT 2L, parse_json('{\"b\":2}')") + .coalesce(1) + .writeTo(deleteTable) + .append(); + + Table table = Spark3Util.loadIcebergTable(spark, deleteTable); + DataFile dataFile = + Iterables.getOnlyElement(SnapshotChanges.builderFor(table).build().addedDataFiles()); + + Pair deletes = + FileHelpers.writeDeleteFile( + table, + Files.localOutput(File.createTempFile("dv-", ".puffin")), + null, + Lists.newArrayList(Pair.of(dataFile.location(), 0L)), + 3); + + table + .newRowDelta() + .addDeletes(deletes.first()) + .validateDataFilesExist(deletes.second()) + .commit(); + + sql("REFRESH TABLE %s", deleteTable); + + Dataset df = spark.table(deleteTable).select("id", "v1").orderBy("id"); + List rows = df.collectAsList(); + + assertThat(rows).hasSize(1); + assertThat(rows.get(0).getLong(0)).isEqualTo(2L); + + Variant v1 = + new Variant( + ((VariantVal) rows.get(0).get(1)).getValue(), + ((VariantVal) rows.get(0).get(1)).getMetadata()); + assertThat(v1.getFieldByKey("b").getLong()).isEqualTo(2L); + + sql("DROP TABLE IF EXISTS %s", deleteTable); + } + @ParameterizedTest @ValueSource(booleans = {false, true}) public void testNestedStructVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); String structTable = CATALOG + ".default.var_struct"; sql("DROP TABLE IF EXISTS %s", structTable); @@ -200,7 +271,6 @@ public void testNestedStructVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testNestedArrayVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); String arrayTable = CATALOG + ".default.var_array"; sql("DROP TABLE IF EXISTS %s", arrayTable); @@ -249,7 +319,6 @@ public void testNestedArrayVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testNestedMapVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); String mapTable = CATALOG + ".default.var_map"; sql("DROP TABLE IF EXISTS %s", mapTable); @@ -305,8 +374,6 @@ public void testNestedMapVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testMergeIntoWithVariant(boolean vectorized) { - // Variant columns are not vectorized yet, but MERGE INTO should not crash regardless of the - // vectorization setting. The reader falls back to non-vectorized for variant columns. String mergeTable = CATALOG + ".default.var_merge"; sql("DROP TABLE IF EXISTS %s", mergeTable); sql( @@ -351,6 +418,97 @@ public void testMergeIntoWithVariant(boolean vectorized) { sql("DROP TABLE IF EXISTS %s", mergeTable); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + public void testReadShreddedAfterPropertyToggled(boolean vectorized) + throws IOException, NoSuchTableException, ParseException { + String toggleTable = CATALOG + ".default.var_toggle"; + sql("DROP TABLE IF EXISTS %s", toggleTable); + sql( + "CREATE TABLE %s (id BIGINT, v VARIANT) USING iceberg " + + "TBLPROPERTIES ('format-version'='3', 'write.parquet.shred-variants'='true')", + toggleTable); + + spark.conf().set("spark.sql.iceberg.shred-variants", "true"); + try { + sql( + "INSERT INTO %s VALUES " + + "(1, parse_json('{\"name\":\"alice\",\"age\":30}')), " + + "(2, parse_json('{\"name\":\"bob\",\"age\":25}'))", + toggleTable); + } finally { + spark.conf().unset("spark.sql.iceberg.shred-variants"); + } + + Table table = Spark3Util.loadIcebergTable(spark, toggleTable); + assertHasTypedValueSubtree(table); + + sql("ALTER TABLE %s SET TBLPROPERTIES ('write.parquet.shred-variants'='false')", toggleTable); + setVectorization(toggleTable, vectorized); + + List rows = spark.table(toggleTable).select("id", "v").orderBy("id").collectAsList(); + assertThat(rows).hasSize(2); + Variant v1 = + new Variant( + ((VariantVal) rows.get(0).get(1)).getValue(), + ((VariantVal) rows.get(0).get(1)).getMetadata()); + assertThat(v1.getFieldByKey("name").getString()).isEqualTo("alice"); + assertThat(v1.getFieldByKey("age").getLong()).isEqualTo(30L); + Variant v2 = + new Variant( + ((VariantVal) rows.get(1).get(1)).getValue(), + ((VariantVal) rows.get(1).get(1)).getMetadata()); + assertThat(v2.getFieldByKey("name").getString()).isEqualTo("bob"); + assertThat(v2.getFieldByKey("age").getLong()).isEqualTo(25L); + + sql("DROP TABLE IF EXISTS %s", toggleTable); + } + + @ParameterizedTest + @ValueSource(strings = {"none", "counts"}) + public void testReadShreddedWithMetricsDisabled(String metricsMode) + throws IOException, NoSuchTableException, ParseException { + String noStatsTable = CATALOG + ".default.var_no_stats"; + sql("DROP TABLE IF EXISTS %s", noStatsTable); + sql( + "CREATE TABLE %s (id BIGINT, v VARIANT) USING iceberg " + + "TBLPROPERTIES (" + + "'format-version'='3', " + + "'write.parquet.shred-variants'='true', " + + "'write.metadata.metrics.default'='%s')", + noStatsTable, metricsMode); + + spark.conf().set("spark.sql.iceberg.shred-variants", "true"); + try { + sql( + "INSERT INTO %s VALUES " + + "(1, parse_json('{\"name\":\"alice\",\"age\":30}')), " + + "(2, parse_json('{\"name\":\"bob\",\"age\":25}'))", + noStatsTable); + } finally { + spark.conf().unset("spark.sql.iceberg.shred-variants"); + } + + Table table = Spark3Util.loadIcebergTable(spark, noStatsTable); + assertHasTypedValueSubtree(table); + setVectorization(noStatsTable, true); + + List rows = spark.table(noStatsTable).select("id", "v").orderBy("id").collectAsList(); + assertThat(rows).hasSize(2); + Variant v1 = + new Variant( + ((VariantVal) rows.get(0).get(1)).getValue(), + ((VariantVal) rows.get(0).get(1)).getMetadata()); + assertThat(v1.getFieldByKey("name").getString()).isEqualTo("alice"); + Variant v2 = + new Variant( + ((VariantVal) rows.get(1).get(1)).getValue(), + ((VariantVal) rows.get(1).get(1)).getMetadata()); + assertThat(v2.getFieldByKey("name").getString()).isEqualTo("bob"); + + sql("DROP TABLE IF EXISTS %s", noStatsTable); + } + private void setVectorization(boolean on) { sql( "ALTER TABLE %s SET TBLPROPERTIES ('read.parquet.vectorization.enabled'='%s')", @@ -362,4 +520,31 @@ private void setVectorization(String table, boolean on) { "ALTER TABLE %s SET TBLPROPERTIES ('read.parquet.vectorization.enabled'='%s')", table, Boolean.toString(on)); } + + private static void assertHasTypedValueSubtree(Table table) throws IOException { + try (CloseableIterable tasks = table.newScan().planFiles()) { + assertThat(tasks).isNotEmpty(); + for (FileScanTask task : tasks) { + HadoopInputFile inputFile = + HadoopInputFile.fromPath(new Path(task.file().location()), new Configuration()); + try (ParquetFileReader reader = ParquetFileReader.open(inputFile)) { + assertThat(containsTypedValue(reader.getFileMetaData().getSchema())) + .as("Expected variant column to be shredded with a typed_value subtree") + .isTrue(); + } + } + } + } + + private static boolean containsTypedValue(org.apache.parquet.schema.Type type) { + if (type.isPrimitive()) { + return false; + } + for (org.apache.parquet.schema.Type child : type.asGroupType().getFields()) { + if (child.getName().equals("typed_value") || containsTypedValue(child)) { + return true; + } + } + return false; + } } diff --git a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/variant/TestVariantShredding.java b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/variant/TestVariantShredding.java index 23b2311fa13f..4b14f32fb13c 100644 --- a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/variant/TestVariantShredding.java +++ b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/variant/TestVariantShredding.java @@ -110,7 +110,12 @@ public static void startMetastoreAndSpark() { public void before() { super.before(); validationCatalog.createTable( - tableIdent, SCHEMA, null, Map.of(TableProperties.FORMAT_VERSION, "3")); + tableIdent, + SCHEMA, + null, + Map.of( + TableProperties.FORMAT_VERSION, "3", + TableProperties.PARQUET_SHRED_VARIANTS, "true")); } @AfterEach diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java index 61616a9f233c..2a14c2262c6a 100644 --- a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java @@ -27,7 +27,9 @@ class ColumnVectorBuilder { public ColumnVector build(VectorHolder holder, int numRows) { - if (holder.isDummy()) { + if (holder instanceof VectorHolder.VariantVectorHolder) { + return new VariantColumnVector((VectorHolder.VariantVectorHolder) holder); + } else if (holder.isDummy()) { if (holder instanceof VectorHolder.DeletedVectorHolder) { return new DeletedColumnVector(Types.BooleanType.get()); } else if (holder instanceof ConstantVectorHolder) { diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java index edaaaeda2515..f11b84443fd8 100644 --- a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorWithFilter.java @@ -20,6 +20,7 @@ import org.apache.spark.sql.types.Decimal; import org.apache.spark.sql.types.StructType; +import org.apache.spark.sql.types.VariantType; import org.apache.spark.sql.vectorized.ColumnVector; import org.apache.spark.sql.vectorized.ColumnarArray; import org.apache.spark.sql.vectorized.ColumnarMap; @@ -143,6 +144,12 @@ public ColumnVector getChild(int ordinal) { for (int index = 0; index < structType.length(); index++) { children[index] = new ColumnVectorWithFilter(delegate.getChild(index), rowIdMapping); } + } else if (dataType() instanceof VariantType) { + this.children = + new ColumnVectorWithFilter[] { + new ColumnVectorWithFilter(delegate.getChild(0), rowIdMapping), + new ColumnVectorWithFilter(delegate.getChild(1), rowIdMapping) + }; } else { throw new UnsupportedOperationException("Unsupported nested type: " + dataType()); } diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VariantColumnVector.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VariantColumnVector.java new file mode 100644 index 000000000000..2831b1f83cdb --- /dev/null +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VariantColumnVector.java @@ -0,0 +1,141 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.iceberg.spark.data.vectorized; + +import org.apache.iceberg.arrow.vectorized.VectorHolder; +import org.apache.spark.sql.types.DataTypes; +import org.apache.spark.sql.types.Decimal; +import org.apache.spark.sql.vectorized.ColumnVector; +import org.apache.spark.sql.vectorized.ColumnarArray; +import org.apache.spark.sql.vectorized.ColumnarMap; +import org.apache.spark.unsafe.types.UTF8String; + +public class VariantColumnVector extends ColumnVector { + private final ColumnVector valueChild; + private final ColumnVector metadataChild; + + VariantColumnVector(VectorHolder.VariantVectorHolder holder) { + super(DataTypes.VariantType); + this.valueChild = new IcebergArrowColumnVector(holder.valueHolder()); + this.metadataChild = new IcebergArrowColumnVector(holder.metadataHolder()); + } + + @Override + public void close() { + valueChild.close(); + metadataChild.close(); + } + + @Override + public void closeIfFreeable() { + // See SPARK-50235, SPARK-50463 + } + + @Override + public boolean hasNull() { + return valueChild.hasNull(); + } + + @Override + public int numNulls() { + return valueChild.numNulls(); + } + + @Override + public boolean isNullAt(int rowId) { + return valueChild.isNullAt(rowId); + } + + // getChild is what getVariant() calls: child(0) = value, child(1) = metadata + @Override + public ColumnVector getChild(int ordinal) { + if (ordinal == 0) { + return valueChild; + } else if (ordinal == 1) { + return metadataChild; + } + + throw new IllegalArgumentException( + "Variant column has only 2 children, got ordinal: " + ordinal); + } + + @Override + public boolean getBoolean(int rowId) { + throw unsupported(); + } + + @Override + public byte getByte(int rowId) { + throw unsupported(); + } + + @Override + public short getShort(int rowId) { + throw unsupported(); + } + + @Override + public int getInt(int rowId) { + throw unsupported(); + } + + @Override + public long getLong(int rowId) { + throw unsupported(); + } + + @Override + public float getFloat(int rowId) { + throw unsupported(); + } + + @Override + public double getDouble(int rowId) { + throw unsupported(); + } + + @Override + public Decimal getDecimal(int rowId, int precision, int scale) { + throw unsupported(); + } + + @Override + public UTF8String getUTF8String(int rowId) { + throw unsupported(); + } + + @Override + public byte[] getBinary(int rowId) { + throw unsupported(); + } + + @Override + public ColumnarArray getArray(int rowId) { + throw unsupported(); + } + + @Override + public ColumnarMap getMap(int rowId) { + throw unsupported(); + } + + private UnsupportedOperationException unsupported() { + return new UnsupportedOperationException("Variant column only supports getVariant()"); + } +} diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index ee5d62a3ed83..5cafa1c9fa04 100644 --- a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -18,17 +18,23 @@ */ package org.apache.iceberg.spark.source; +import java.nio.ByteBuffer; import java.util.List; +import java.util.Map; import java.util.Objects; import java.util.function.Supplier; import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.MetadataColumns; +import org.apache.iceberg.MetricsConfig; +import org.apache.iceberg.MetricsModes; +import org.apache.iceberg.MetricsUtil; import org.apache.iceberg.ScanTask; import org.apache.iceberg.ScanTaskGroup; import org.apache.iceberg.Schema; import org.apache.iceberg.SchemaParser; import org.apache.iceberg.Table; +import org.apache.iceberg.TableProperties; import org.apache.iceberg.io.FileIO; import org.apache.iceberg.spark.ImmutableOrcBatchReadConf; import org.apache.iceberg.spark.ImmutableParquetBatchReadConf; @@ -38,6 +44,7 @@ import org.apache.iceberg.spark.SparkUtil; import org.apache.iceberg.types.Type; import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.PropertyUtil; import org.apache.spark.api.java.JavaSparkContext; import org.apache.spark.broadcast.Broadcast; import org.apache.spark.sql.connector.read.Batch; @@ -147,8 +154,8 @@ private OrcBatchReadConf orcBatchReadConf() { // conditions for using Parquet batch reads: // - Parquet vectorization is enabled - // - only primitives or metadata columns are projected, excluding geometry and geography which - // are primitives with no Arrow vector yet + // - only primitives, unshredded variant, or metadata columns are projected, excluding geometry + // and geography which are primitives with no Arrow vector yet // - all tasks are of FileScanTask type and read only Parquet files private boolean useParquetBatchReads() { return readConf.parquetVectorizationEnabled() @@ -163,7 +170,18 @@ private boolean supportsParquetBatchReads(ScanTask task) { } else if (task.isFileScanTask() && !task.isDataTask()) { FileScanTask fileScanTask = task.asFileScanTask(); - return fileScanTask.file().format() == FileFormat.PARQUET; + if (fileScanTask.file().format() != FileFormat.PARQUET) { + return false; + } + Map lowerBounds = fileScanTask.file().lowerBounds(); + if (lowerBounds != null) { + for (Types.NestedField field : projection.columns()) { + if (field.type().isVariantType() && lowerBounds.containsKey(field.fieldId())) { + return false; + } + } + } + return true; } else { return false; @@ -182,7 +200,25 @@ private boolean supportsParquetBatchReads(Types.NestedField field) { return false; } - return type.isPrimitiveType(); + if (type.isVariantType()) { + boolean shredVariants = + PropertyUtil.propertyAsBoolean( + table.properties(), + TableProperties.PARQUET_SHRED_VARIANTS, + TableProperties.PARQUET_SHRED_VARIANTS_DEFAULT); + if (shredVariants) { + return false; + } + + MetricsConfig metricsConfig = MetricsConfig.forTable(table); + MetricsModes.MetricsMode mode = + MetricsUtil.metricsMode(table.schema(), metricsConfig, field.fieldId()); + if (mode == MetricsModes.None.get() || mode == MetricsModes.Counts.get()) { + return false; + } + } + + return type.isPrimitiveType() || type.isVariantType(); } // conditions for using ORC batch reads: diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java index 69b6314a7f2b..c952417a547d 100644 --- a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java @@ -21,6 +21,7 @@ import java.io.IOException; import java.util.List; import java.util.Objects; +import java.util.stream.Collectors; import org.apache.iceberg.BaseMetadataTable; import org.apache.iceberg.BaseTable; import org.apache.iceberg.BatchScan; @@ -47,6 +48,7 @@ import org.apache.iceberg.spark.SparkTableUtil; import org.apache.iceberg.spark.TimeTravel; import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.Types; import org.apache.iceberg.util.Pair; import org.apache.spark.sql.SparkSession; import org.apache.spark.sql.catalyst.InternalRow; @@ -327,6 +329,11 @@ private IncrementalAppendScan buildIcebergIncrementalAppendScan( if (withStats) { scan = scan.includeColumnStats(); + } else { + List variantColumns = variantColumnNames(projection); + if (!variantColumns.isEmpty()) { + scan = scan.includeColumnStats(variantColumns); + } } if (endSnapshotId != null) { @@ -365,11 +372,23 @@ private BatchScan buildIcebergBatchScan( if (withStats) { scan = scan.includeColumnStats(); + } else { + List variantColumns = variantColumnNames(projection); + if (!variantColumns.isEmpty()) { + scan = scan.includeColumnStats(variantColumns); + } } return configureSplitPlanning(scan); } + private List variantColumnNames(Schema projection) { + return projection.columns().stream() + .filter(field -> field.type().isVariantType()) + .map(Types.NestedField::name) + .collect(Collectors.toList()); + } + private BatchScan newIcebergBatchScan() { if (readConf().distributedPlanningEnabled()) { return new SparkDistributedDataScan(spark(), table(), readConf()); diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkVariantRead.java b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkVariantRead.java index 2d6e919a91ee..79e52d6e339c 100644 --- a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkVariantRead.java +++ b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/sql/TestSparkVariantRead.java @@ -19,13 +19,33 @@ package org.apache.iceberg.spark.sql; import static org.assertj.core.api.Assertions.assertThat; -import static org.assertj.core.api.Assumptions.assumeThat; +import java.io.File; +import java.io.IOException; import java.util.List; +import org.apache.hadoop.conf.Configuration; +import org.apache.hadoop.fs.Path; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.Files; +import org.apache.iceberg.SnapshotChanges; +import org.apache.iceberg.Table; +import org.apache.iceberg.data.FileHelpers; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.relocated.com.google.common.collect.Iterables; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.spark.Spark3Util; import org.apache.iceberg.spark.SparkCatalog; import org.apache.iceberg.spark.TestBase; +import org.apache.iceberg.util.CharSequenceSet; +import org.apache.iceberg.util.Pair; +import org.apache.parquet.hadoop.ParquetFileReader; +import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; +import org.apache.spark.sql.catalyst.analysis.NoSuchTableException; +import org.apache.spark.sql.catalyst.parser.ParseException; import org.apache.spark.types.variant.Variant; import org.apache.spark.unsafe.types.VariantVal; import org.junit.jupiter.api.AfterEach; @@ -76,7 +96,6 @@ public void cleanup() { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testVariantColumnProjection_singleVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); setVectorization(vectorized); Dataset df = spark.table(TABLE).select("id", "v1").orderBy("id"); assertThat(df.schema().fieldNames()).containsExactly("id", "v1"); @@ -106,7 +125,6 @@ public void testVariantColumnProjection_singleVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testVariantColumnProjectionNoVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); setVectorization(vectorized); Dataset df = spark.table(TABLE).select("id"); assertThat(df.schema().fieldNames()).containsExactly("id"); @@ -117,7 +135,6 @@ public void testVariantColumnProjectionNoVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testFilterOnVariantColumnOnWholeValue(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); setVectorization(vectorized); sql("INSERT INTO %s SELECT 3, NULL, NULL", TABLE); @@ -147,7 +164,6 @@ public void testFilterOnVariantColumnOnWholeValue(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testVariantNullValueProjection(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); setVectorization(vectorized); // insert a row with NULL variant values @@ -162,10 +178,65 @@ public void testVariantNullValueProjection(boolean vectorized) { assertThat(row.isNullAt(1)).isTrue(); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + public void testVariantReadAfterDelete(boolean vectorized) + throws IOException, NoSuchTableException, ParseException { + String deleteTable = CATALOG + ".default.var_delete"; + + sql("DROP TABLE IF EXISTS %s", deleteTable); + sql( + "CREATE TABLE %s (id BIGINT, v1 VARIANT) USING iceberg " + + "TBLPROPERTIES ('format-version'='3')", + deleteTable); + setVectorization(deleteTable, vectorized); + + spark + .sql( + "SELECT 1L AS id, parse_json('{\"a\":1}') AS v1 " + + "UNION ALL SELECT 2L, parse_json('{\"b\":2}')") + .coalesce(1) + .writeTo(deleteTable) + .append(); + + Table table = Spark3Util.loadIcebergTable(spark, deleteTable); + DataFile dataFile = + Iterables.getOnlyElement(SnapshotChanges.builderFor(table).build().addedDataFiles()); + + Pair deletes = + FileHelpers.writeDeleteFile( + table, + Files.localOutput(File.createTempFile("dv-", ".puffin")), + null, + Lists.newArrayList(Pair.of(dataFile.location(), 0L)), + 3); + + table + .newRowDelta() + .addDeletes(deletes.first()) + .validateDataFilesExist(deletes.second()) + .commit(); + + sql("REFRESH TABLE %s", deleteTable); + + Dataset df = spark.table(deleteTable).select("id", "v1").orderBy("id"); + List rows = df.collectAsList(); + + assertThat(rows).hasSize(1); + assertThat(rows.get(0).getLong(0)).isEqualTo(2L); + + Variant v1 = + new Variant( + ((VariantVal) rows.get(0).get(1)).getValue(), + ((VariantVal) rows.get(0).get(1)).getMetadata()); + assertThat(v1.getFieldByKey("b").getLong()).isEqualTo(2L); + + sql("DROP TABLE IF EXISTS %s", deleteTable); + } + @ParameterizedTest @ValueSource(booleans = {false, true}) public void testNestedStructVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); String structTable = CATALOG + ".default.var_struct"; sql("DROP TABLE IF EXISTS %s", structTable); @@ -200,7 +271,6 @@ public void testNestedStructVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testNestedArrayVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); String arrayTable = CATALOG + ".default.var_array"; sql("DROP TABLE IF EXISTS %s", arrayTable); @@ -249,7 +319,6 @@ public void testNestedArrayVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testNestedMapVariant(boolean vectorized) { - assumeThat(vectorized).as("Variant vectorized Parquet read is not implemented yet").isFalse(); String mapTable = CATALOG + ".default.var_map"; sql("DROP TABLE IF EXISTS %s", mapTable); @@ -305,8 +374,6 @@ public void testNestedMapVariant(boolean vectorized) { @ParameterizedTest @ValueSource(booleans = {false, true}) public void testMergeIntoWithVariant(boolean vectorized) { - // Variant columns are not vectorized yet, but MERGE INTO should not crash regardless of the - // vectorization setting. The reader falls back to non-vectorized for variant columns. String mergeTable = CATALOG + ".default.var_merge"; sql("DROP TABLE IF EXISTS %s", mergeTable); sql( @@ -351,6 +418,97 @@ public void testMergeIntoWithVariant(boolean vectorized) { sql("DROP TABLE IF EXISTS %s", mergeTable); } + @ParameterizedTest + @ValueSource(booleans = {false, true}) + public void testReadShreddedAfterPropertyToggled(boolean vectorized) + throws IOException, NoSuchTableException, ParseException { + String toggleTable = CATALOG + ".default.var_toggle"; + sql("DROP TABLE IF EXISTS %s", toggleTable); + sql( + "CREATE TABLE %s (id BIGINT, v VARIANT) USING iceberg " + + "TBLPROPERTIES ('format-version'='3', 'write.parquet.shred-variants'='true')", + toggleTable); + + spark.conf().set("spark.sql.iceberg.shred-variants", "true"); + try { + sql( + "INSERT INTO %s VALUES " + + "(1, parse_json('{\"name\":\"alice\",\"age\":30}')), " + + "(2, parse_json('{\"name\":\"bob\",\"age\":25}'))", + toggleTable); + } finally { + spark.conf().unset("spark.sql.iceberg.shred-variants"); + } + + Table table = Spark3Util.loadIcebergTable(spark, toggleTable); + assertHasTypedValueSubtree(table); + + sql("ALTER TABLE %s SET TBLPROPERTIES ('write.parquet.shred-variants'='false')", toggleTable); + setVectorization(toggleTable, vectorized); + + List rows = spark.table(toggleTable).select("id", "v").orderBy("id").collectAsList(); + assertThat(rows).hasSize(2); + Variant v1 = + new Variant( + ((VariantVal) rows.get(0).get(1)).getValue(), + ((VariantVal) rows.get(0).get(1)).getMetadata()); + assertThat(v1.getFieldByKey("name").getString()).isEqualTo("alice"); + assertThat(v1.getFieldByKey("age").getLong()).isEqualTo(30L); + Variant v2 = + new Variant( + ((VariantVal) rows.get(1).get(1)).getValue(), + ((VariantVal) rows.get(1).get(1)).getMetadata()); + assertThat(v2.getFieldByKey("name").getString()).isEqualTo("bob"); + assertThat(v2.getFieldByKey("age").getLong()).isEqualTo(25L); + + sql("DROP TABLE IF EXISTS %s", toggleTable); + } + + @ParameterizedTest + @ValueSource(strings = {"none", "counts"}) + public void testReadShreddedWithMetricsDisabled(String metricsMode) + throws IOException, NoSuchTableException, ParseException { + String noStatsTable = CATALOG + ".default.var_no_stats"; + sql("DROP TABLE IF EXISTS %s", noStatsTable); + sql( + "CREATE TABLE %s (id BIGINT, v VARIANT) USING iceberg " + + "TBLPROPERTIES (" + + "'format-version'='3', " + + "'write.parquet.shred-variants'='true', " + + "'write.metadata.metrics.default'='%s')", + noStatsTable, metricsMode); + + spark.conf().set("spark.sql.iceberg.shred-variants", "true"); + try { + sql( + "INSERT INTO %s VALUES " + + "(1, parse_json('{\"name\":\"alice\",\"age\":30}')), " + + "(2, parse_json('{\"name\":\"bob\",\"age\":25}'))", + noStatsTable); + } finally { + spark.conf().unset("spark.sql.iceberg.shred-variants"); + } + + Table table = Spark3Util.loadIcebergTable(spark, noStatsTable); + assertHasTypedValueSubtree(table); + setVectorization(noStatsTable, true); + + List rows = spark.table(noStatsTable).select("id", "v").orderBy("id").collectAsList(); + assertThat(rows).hasSize(2); + Variant v1 = + new Variant( + ((VariantVal) rows.get(0).get(1)).getValue(), + ((VariantVal) rows.get(0).get(1)).getMetadata()); + assertThat(v1.getFieldByKey("name").getString()).isEqualTo("alice"); + Variant v2 = + new Variant( + ((VariantVal) rows.get(1).get(1)).getValue(), + ((VariantVal) rows.get(1).get(1)).getMetadata()); + assertThat(v2.getFieldByKey("name").getString()).isEqualTo("bob"); + + sql("DROP TABLE IF EXISTS %s", noStatsTable); + } + private void setVectorization(boolean on) { sql( "ALTER TABLE %s SET TBLPROPERTIES ('read.parquet.vectorization.enabled'='%s')", @@ -362,4 +520,31 @@ private void setVectorization(String table, boolean on) { "ALTER TABLE %s SET TBLPROPERTIES ('read.parquet.vectorization.enabled'='%s')", table, Boolean.toString(on)); } + + private static void assertHasTypedValueSubtree(Table table) throws IOException { + try (CloseableIterable tasks = table.newScan().planFiles()) { + assertThat(tasks).isNotEmpty(); + for (FileScanTask task : tasks) { + HadoopInputFile inputFile = + HadoopInputFile.fromPath(new Path(task.file().location()), new Configuration()); + try (ParquetFileReader reader = ParquetFileReader.open(inputFile)) { + assertThat(containsTypedValue(reader.getFileMetaData().getSchema())) + .as("Expected variant column to be shredded with a typed_value subtree") + .isTrue(); + } + } + } + } + + private static boolean containsTypedValue(org.apache.parquet.schema.Type type) { + if (type.isPrimitive()) { + return false; + } + for (org.apache.parquet.schema.Type child : type.asGroupType().getFields()) { + if (child.getName().equals("typed_value") || containsTypedValue(child)) { + return true; + } + } + return false; + } } diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/variant/TestVariantShredding.java b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/variant/TestVariantShredding.java index 23b2311fa13f..4b14f32fb13c 100644 --- a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/variant/TestVariantShredding.java +++ b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/variant/TestVariantShredding.java @@ -110,7 +110,12 @@ public static void startMetastoreAndSpark() { public void before() { super.before(); validationCatalog.createTable( - tableIdent, SCHEMA, null, Map.of(TableProperties.FORMAT_VERSION, "3")); + tableIdent, + SCHEMA, + null, + Map.of( + TableProperties.FORMAT_VERSION, "3", + TableProperties.PARQUET_SHRED_VARIANTS, "true")); } @AfterEach