From f3d1ebb9596cb631cab8dbf651b9f9e6ddefbf7f Mon Sep 17 00:00:00 2001 From: Neelesh Salian Date: Mon, 14 Sep 2026 20:54:33 -0700 Subject: [PATCH] Arrow, Parquet, Spark: Support vectorized Parquet reads for struct type --- .../arrow/vectorized/VectorHolder.java | 32 ++++ .../vectorized/VectorizedArrowReader.java | 137 +++++++++++++ .../vectorized/VectorizedReaderBuilder.java | 86 +++++++-- .../parquet/VectorizedColumnIterator.java | 5 + .../parquet/VectorizedPageIterator.java | 17 ++ ...ectorizedParquetDefinitionLevelReader.java | 45 +++++ .../TestVectorizedStructReader.java | 116 +++++++++++ .../iceberg/data/RandomGenericData.java | 7 + .../iceberg/parquet/MessageTypeToType.java | 4 + .../iceberg/parquet/ParquetSchemaUtil.java | 16 +- .../data/vectorized/ColumnVectorBuilder.java | 4 +- .../data/vectorized/StructColumnVector.java | 143 ++++++++++++++ .../iceberg/spark/source/SparkBatch.java | 16 +- .../parquet/TestParquetVectorizedReads.java | 180 ++++++++++++++++-- .../spark/source/TestReadProjection.java | 51 +++++ .../spark/source/TestSparkReaderDeletes.java | 66 +++++++ .../data/vectorized/ColumnVectorBuilder.java | 2 + .../data/vectorized/StructColumnVector.java | 143 ++++++++++++++ .../iceberg/spark/source/SparkBatch.java | 37 +++- .../parquet/TestParquetVectorizedReads.java | 180 ++++++++++++++++-- .../spark/source/TestReadProjection.java | 51 +++++ .../spark/source/TestSparkReaderDeletes.java | 66 +++++++ .../data/vectorized/ColumnVectorBuilder.java | 2 + .../data/vectorized/StructColumnVector.java | 143 ++++++++++++++ .../iceberg/spark/source/SparkBatch.java | 38 +++- .../parquet/TestParquetVectorizedReads.java | 179 +++++++++++++++-- .../spark/source/TestReadProjection.java | 51 +++++ .../spark/source/TestSparkReaderDeletes.java | 66 +++++++ .../data/vectorized/ColumnVectorBuilder.java | 2 + .../data/vectorized/StructColumnVector.java | 143 ++++++++++++++ .../iceberg/spark/source/SparkBatch.java | 38 +++- .../parquet/TestParquetVectorizedReads.java | 179 +++++++++++++++-- .../spark/source/TestReadProjection.java | 51 +++++ .../spark/source/TestSparkReaderDeletes.java | 66 +++++++ 34 files changed, 2259 insertions(+), 103 deletions(-) create mode 100644 arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedStructReader.java create mode 100644 spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java create mode 100644 spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java create mode 100644 spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java create mode 100644 spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java 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 2b749256b0d9..8deff9b84274 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 @@ -18,6 +18,7 @@ */ package org.apache.iceberg.arrow.vectorized; +import java.util.List; import org.apache.arrow.vector.FieldVector; import org.apache.iceberg.MetadataColumns; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; @@ -215,4 +216,35 @@ public VectorHolder valueHolder() { return valueHolder; } } + + public static class StructVectorHolder extends VectorHolder { + private final int numRows; + private final List childHolders; + private final NullabilityHolder structNulls; + + public StructVectorHolder( + Types.NestedField icebergField, + int numRows, + List childHolders, + NullabilityHolder structNulls) { + super(icebergField); + this.numRows = numRows; + this.childHolders = childHolders; + this.structNulls = structNulls; + } + + @Override + public int numValues() { + return numRows; + } + + public List childHolders() { + return childHolders; + } + + @Override + public NullabilityHolder nullabilityHolder() { + return structNulls; + } + } } 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 fa8ec0222464..c592087121aa 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 @@ -18,6 +18,7 @@ */ package org.apache.iceberg.arrow.vectorized; +import java.util.List; import java.util.Map; import java.util.Optional; import org.apache.arrow.memory.ArrowBuf; @@ -47,6 +48,7 @@ import org.apache.iceberg.parquet.ParquetUtil; import org.apache.iceberg.parquet.VectorizedReader; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.iceberg.types.Types; import org.apache.parquet.column.ColumnDescriptor; import org.apache.parquet.column.Dictionary; @@ -136,6 +138,14 @@ public void setBatchSize(int batchSize) { this.vectorizedColumnIterator.setBatchSize(batchSize); } + void registerStructPresence(NullabilityHolder structNulls, int structDefinitionLevel) { + vectorizedColumnIterator.addStructPresence(structNulls, structDefinitionLevel); + } + + protected VectorizedArrowReader fileBackedLeaf() { + return vectorizedColumnIterator != null ? this : null; + } + @Override public VectorHolder read(VectorHolder reuse, int numValsToRead) { boolean dictEncoded = vectorizedColumnIterator.producesDictionaryEncodedVector(); @@ -386,6 +396,11 @@ public void setRowGroupInfo(PageReadStore source, Map metadata) { + return columnDescriptor != null + && metadata.containsKey(ColumnPath.get(columnDescriptor.getPath())); + } + @Override public void close() { if (vec != null) { @@ -1054,6 +1069,12 @@ public VectorizedVariantReader( this.valueReader = valueReader; } + @Override + protected VectorizedArrowReader fileBackedLeaf() { + VectorizedArrowReader metadataLeaf = metadataReader.fileBackedLeaf(); + return metadataLeaf != null ? metadataLeaf : valueReader.fileBackedLeaf(); + } + @Override public VectorHolder read(VectorHolder reuse, int numValsToRead) { VectorHolder reuseMetadata = null; @@ -1093,4 +1114,120 @@ public String toString() { return "VectorizedVariantReader"; } } + + static class VectorizedStructReader extends VectorizedArrowReader { + private final List> childReaders; + private final int structDefinitionLevel; + private final VectorizedArrowReader presenceReader; + private NullabilityHolder structNulls; + private VectorHolder reusePresence; + private boolean presenceColumnInRowGroup; + + VectorizedStructReader( + Types.NestedField icebergField, + List> childReaders, + int structDefinitionLevel, + VectorizedArrowReader presenceReader) { + super(icebergField); + this.childReaders = childReaders; + this.structDefinitionLevel = structDefinitionLevel; + this.presenceReader = presenceReader; + } + + @Override + protected VectorizedArrowReader fileBackedLeaf() { + for (VectorizedReader child : childReaders) { + if (child instanceof VectorizedArrowReader) { + VectorizedArrowReader leaf = ((VectorizedArrowReader) child).fileBackedLeaf(); + if (leaf != null) { + return leaf; + } + } + } + + // expose the presence leaf so an ancestor reuses it instead of double-reading the column + return presenceReader; + } + + @Override + public VectorHolder read(VectorHolder reuse, int numValsToRead) { + if (structNulls != null) { + structNulls.reset(); + } + + // when presenceColumnInRowGroup is false the struct is always present (structNulls left all + // not-null) + if (presenceReader != null && presenceColumnInRowGroup) { + // populate structNulls as a side effect of a value-column batch read; the values are unused + this.reusePresence = presenceReader.read(reusePresence, numValsToRead); + } + + List reuseChildren = null; + if (reuse instanceof VectorHolder.StructVectorHolder) { + reuseChildren = ((VectorHolder.StructVectorHolder) reuse).childHolders(); + } + + List childHolders = Lists.newArrayListWithExpectedSize(childReaders.size()); + for (int idx = 0; idx < childReaders.size(); idx++) { + VectorHolder reuseChild = reuseChildren == null ? null : reuseChildren.get(idx); + VectorizedArrowReader child = (VectorizedArrowReader) childReaders.get(idx); + childHolders.add(child.read(reuseChild, numValsToRead)); + } + + return new VectorHolder.StructVectorHolder( + icebergField(), numValsToRead, childHolders, structNulls); + } + + @Override + public void setRowGroupInfo( + PageReadStore source, Map metadata) { + for (VectorizedReader child : childReaders) { + child.setRowGroupInfo(source, metadata); + } + + // a partition-constant presence column is absent from the row group, so the struct is present + this.presenceColumnInRowGroup = + presenceReader != null && presenceReader.isColumnInRowGroup(metadata); + if (presenceColumnInRowGroup) { + presenceReader.setRowGroupInfo(source, metadata); + } + } + + @Override + public void setBatchSize(int batchSize) { + int resolvedBatchSize = (batchSize == 0) ? DEFAULT_BATCH_SIZE : batchSize; + for (VectorizedReader child : childReaders) { + child.setBatchSize(resolvedBatchSize); + } + + if (presenceReader != null) { + presenceReader.setBatchSize(resolvedBatchSize); + } + + if (structDefinitionLevel > 0) { + // per-row presence from a file leaf under the struct (child or shared presence leaf) + VectorizedArrowReader presenceLeaf = fileBackedLeaf(); + if (presenceLeaf != null) { + this.structNulls = new NullabilityHolder(resolvedBatchSize); + presenceLeaf.registerStructPresence(structNulls, structDefinitionLevel); + } + } + } + + @Override + public void close() { + for (VectorizedReader child : childReaders) { + child.close(); + } + + if (presenceReader != null) { + presenceReader.close(); + } + } + + @Override + public String toString() { + return "VectorizedStructReader(" + childReaders.size() + ")"; + } + } } 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 59023fe7c877..75bdf9556669 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.ParquetSchemaUtil; import org.apache.iceberg.parquet.ParquetVariantVisitor; import org.apache.iceberg.parquet.TypeWithSchemaVisitor; import org.apache.iceberg.parquet.VectorizedReader; @@ -102,27 +103,32 @@ protected VectorizedReaderBuilder( @Override public VectorizedReader message( Types.StructType expected, MessageType message, List> fieldReaders) { - GroupType groupType = message.asGroupType(); - Map> readersById = Maps.newHashMap(); - List fields = groupType.getFields(); - - IntStream.range(0, fields.size()) - .filter(pos -> fields.get(pos).getId() != null) - .forEach(pos -> readersById.put(fields.get(pos).getId().intValue(), fieldReaders.get(pos))); - List icebergFields = expected != null ? expected.fields() : ImmutableList.of(); + return vectorizedReader( + reorderFields(icebergFields, message.asGroupType().getFields(), fieldReaders)); + } - List> reorderedFields = - Lists.newArrayListWithExpectedSize(icebergFields.size()); - - for (Types.NestedField field : icebergFields) { + private List> reorderFields( + List expectedFields, + List parquetFields, + List> fieldReaders) { + Map> readersById = Maps.newHashMap(); + IntStream.range(0, parquetFields.size()) + .filter(pos -> parquetFields.get(pos).getId() != null) + .forEach( + pos -> + readersById.put(parquetFields.get(pos).getId().intValue(), fieldReaders.get(pos))); + + List> reordered = Lists.newArrayListWithExpectedSize(expectedFields.size()); + for (Types.NestedField field : expectedFields) { VectorizedReader reader = VectorizedArrowReader.replaceWithMetadataReader( field, readersById.get(field.fieldId()), idToConstant, setArrowValidityVector); - reorderedFields.add(defaultReader(field, reader)); + reordered.add(defaultReader(field, reader)); } - return vectorizedReader(reorderedFields); + + return reordered; } private VectorizedReader defaultReader(Types.NestedField field, VectorizedReader reader) { @@ -148,11 +154,55 @@ protected VectorizedReader vectorizedReader(List> reorder @Override public VectorizedReader struct( Types.StructType expected, GroupType groupType, List> fieldReaders) { - if (expected != null) { - throw new UnsupportedOperationException( - "Vectorized reads are not supported yet for struct fields"); + if (expected == null) { + return null; + } + + // no field ID / no matching Iceberg field: fall back like primitive() + if (groupType.getId() == null) { + return null; } - return null; + + Types.NestedField structField = icebergSchema.findField(groupType.getId().intValue()); + if (structField == null) { + return null; + } + + List> reorderedFields = + reorderFields(expected.fields(), groupType.getFields(), fieldReaders); + + int structDefinitionLevel = parquetSchema.getMaxDefinitionLevel(currentPath()); + + VectorizedArrowReader presenceReader = null; + // must agree with ParquetSchemaUtil.PresenceColumnSelector, which retains this presence column + if (structDefinitionLevel > 0 && !hasFileBackedLeaf(reorderedFields)) { + ColumnDescriptor presence = + ParquetSchemaUtil.selectPresenceColumn(parquetSchema, currentPath()); + // list/map are not vectorized, so the presence leaf is always rep level 0 (iterator + // precondition) + if (presence != null && presence.getMaxRepetitionLevel() == 0) { + presenceReader = + new VectorizedArrowReader( + presence, + ParquetSchemaUtil.presenceField(presence), + rootAllocator, + setArrowValidityVector); + } + } + + return new VectorizedArrowReader.VectorizedStructReader( + structField, reorderedFields, structDefinitionLevel, presenceReader); + } + + private boolean hasFileBackedLeaf(List> readers) { + for (VectorizedReader child : readers) { + if (child instanceof VectorizedArrowReader + && ((VectorizedArrowReader) child).fileBackedLeaf() != null) { + return true; + } + } + + return false; } @Override diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedColumnIterator.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedColumnIterator.java index 8e52fd1a28c3..8a58a42bfcda 100644 --- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedColumnIterator.java +++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedColumnIterator.java @@ -49,6 +49,11 @@ public VectorizedColumnIterator( public void setBatchSize(int batchSize) { this.batchSize = batchSize; + vectorizedPageIterator.clearStructPresences(); + } + + public void addStructPresence(NullabilityHolder structNulls, int structDefinitionLevel) { + vectorizedPageIterator.addStructPresence(structNulls, structDefinitionLevel); } public Dictionary setRowGroupInfo(PageReader store, boolean allPagesDictEncoded) { diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedPageIterator.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedPageIterator.java index 4a06f64b5af9..ead46a91d44a 100644 --- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedPageIterator.java +++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedPageIterator.java @@ -19,11 +19,13 @@ package org.apache.iceberg.arrow.vectorized.parquet; import java.io.IOException; +import java.util.List; import org.apache.arrow.vector.FieldVector; import org.apache.arrow.vector.IntVector; import org.apache.iceberg.arrow.vectorized.NullabilityHolder; import org.apache.iceberg.parquet.BasePageIterator; import org.apache.iceberg.parquet.ParquetUtil; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.parquet.CorruptDeltaByteArrays; import org.apache.parquet.bytes.ByteBufferInputStream; import org.apache.parquet.bytes.BytesUtils; @@ -58,6 +60,19 @@ private enum DictionaryDecodeMode { private DictionaryDecodeMode dictionaryDecodeMode; + private final List structPresences = + Lists.newArrayList(); + + void addStructPresence(NullabilityHolder structNulls, int structDefinitionLevel) { + structPresences.add( + new VectorizedParquetDefinitionLevelReader.StructPresence( + structNulls, structDefinitionLevel)); + } + + void clearStructPresences() { + structPresences.clear(); + } + public void setAllPagesDictEncoded(boolean allDictEncoded) { this.allPagesDictEncoded = allDictEncoded; } @@ -148,6 +163,7 @@ protected void initDefinitionLevelsReader( this.vectorizedDefinitionLevelReader = new VectorizedParquetDefinitionLevelReader( bitWidth, desc.getMaxDefinitionLevel(), setArrowValidityVector); + this.vectorizedDefinitionLevelReader.setStructPresences(structPresences); this.vectorizedDefinitionLevelReader.initFromPage(triplesCount, in); } @@ -159,6 +175,7 @@ protected void initDefinitionLevelsReader(DataPageV2 dataPageV2, ColumnDescripto this.vectorizedDefinitionLevelReader = new VectorizedParquetDefinitionLevelReader( bitWidth, desc.getMaxDefinitionLevel(), false, setArrowValidityVector); + this.vectorizedDefinitionLevelReader.setStructPresences(structPresences); this.vectorizedDefinitionLevelReader.initFromPage( dataPageV2.getValueCount(), dataPageV2.getDefinitionLevels().toInputStream()); } diff --git a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedParquetDefinitionLevelReader.java b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedParquetDefinitionLevelReader.java index d093d4c97989..b4cb9f2e08c4 100644 --- a/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedParquetDefinitionLevelReader.java +++ b/arrow/src/main/java/org/apache/iceberg/arrow/vectorized/parquet/VectorizedParquetDefinitionLevelReader.java @@ -20,6 +20,7 @@ import java.nio.ByteBuffer; import java.nio.ByteOrder; +import java.util.List; import org.apache.arrow.memory.ArrowBuf; import org.apache.arrow.vector.BaseVariableWidthVector; import org.apache.arrow.vector.BitVector; @@ -29,6 +30,7 @@ import org.apache.arrow.vector.IntVector; import org.apache.iceberg.arrow.vectorized.NullabilityHolder; import org.apache.iceberg.parquet.ParquetUtil; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; import org.apache.parquet.column.Dictionary; public final class VectorizedParquetDefinitionLevelReader @@ -44,6 +46,48 @@ public VectorizedParquetDefinitionLevelReader( super(bitWidth, maxDefLevel, readLength, setArrowValidityVector); } + private List structPresences = ImmutableList.of(); + + void setStructPresences(List presences) { + this.structPresences = presences; + } + + static class StructPresence { + private final NullabilityHolder structNulls; + private final int structDefinitionLevel; + + StructPresence(NullabilityHolder structNulls, int structDefinitionLevel) { + this.structNulls = structNulls; + this.structDefinitionLevel = structDefinitionLevel; + } + } + + private void recordStructPresence(Mode currentMode, int startIndex, int numValues) { + if (structPresences.isEmpty()) { + return; + } + + for (StructPresence presence : structPresences) { + if (currentMode == Mode.RLE) { + if (currentValue >= presence.structDefinitionLevel) { + presence.structNulls.setNotNulls(startIndex, numValues); + } else { + presence.structNulls.setNulls(startIndex, numValues); + } + } else { + // PACKED: must read before the consumer advances packedValuesBufferIdx for this run + for (int offset = 0; offset < numValues; offset++) { + if (packedValuesBuffer[packedValuesBufferIdx + offset] + >= presence.structDefinitionLevel) { + presence.structNulls.setNotNull(startIndex + offset); + } else { + presence.structNulls.setNull(startIndex + offset); + } + } + } + } + } + @FunctionalInterface interface ReaderFunction { void apply(Mode mode, int idx, int numValues, byte[] byteArray, ArrowBuf validityBuffer); @@ -70,6 +114,7 @@ private void nextBatch( } ArrowBuf validityBuffer = vector.getValidityBuffer(); + recordStructPresence(mode, idx, numValues); consumer.apply(mode, idx, numValues, byteArray, validityBuffer); idx += numValues; left -= numValues; diff --git a/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedStructReader.java b/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedStructReader.java new file mode 100644 index 000000000000..d9377a6e2535 --- /dev/null +++ b/arrow/src/test/java/org/apache/iceberg/arrow/vectorized/TestVectorizedStructReader.java @@ -0,0 +1,116 @@ +/* + * 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.assertThat; + +import org.apache.iceberg.arrow.ArrowAllocation; +import org.apache.iceberg.parquet.VectorizedReader; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.apache.iceberg.types.Types.BinaryType; +import org.apache.iceberg.types.Types.IntegerType; +import org.apache.iceberg.types.Types.LongType; +import org.apache.iceberg.types.Types.NestedField; +import org.apache.iceberg.types.Types.StructType; +import org.apache.iceberg.types.Types.VariantType; +import org.apache.parquet.column.ColumnDescriptor; +import org.apache.parquet.schema.PrimitiveType.PrimitiveTypeName; +import org.apache.parquet.schema.Types; +import org.junit.jupiter.api.Test; + +class TestVectorizedStructReader { + + @Test + void structAssemblesChildHoldersInOrder() { + NestedField lat = NestedField.required(2, "lat", IntegerType.get()); + NestedField lng = NestedField.optional(3, "lng", IntegerType.get()); + NestedField structField = NestedField.required(1, "s", StructType.of(lat, lng)); + + VectorizedArrowReader latReader = new VectorizedArrowReader.ConstantVectorReader<>(lat, 7); + VectorizedArrowReader lngReader = new VectorizedArrowReader.ConstantVectorReader<>(lng, 9); + VectorizedArrowReader.VectorizedStructReader reader = + new VectorizedArrowReader.VectorizedStructReader( + structField, ImmutableList.>of(latReader, lngReader), 0, null); + + reader.setBatchSize(128); + VectorHolder.StructVectorHolder holder = (VectorHolder.StructVectorHolder) reader.read(null, 3); + + assertThat(holder.numValues()).isEqualTo(3); + assertThat(holder.childHolders()).hasSize(2); + assertThat(((VectorHolder.ConstantVectorHolder) holder.childHolders().get(0)).getConstant()) + .isEqualTo(7); + assertThat(((VectorHolder.ConstantVectorHolder) holder.childHolders().get(1)).getConstant()) + .isEqualTo(9); + } + + @Test + void variantChildFileBackedLeafResolvesThroughMetadataReader() { + NestedField metadataField = NestedField.optional(3, "metadata", BinaryType.get()); + NestedField variantField = NestedField.optional(2, "v", VariantType.get()); + ColumnDescriptor metadataColumn = + new ColumnDescriptor( + new String[] {"metadata"}, + Types.optional(PrimitiveTypeName.BINARY).named("metadata"), + 0, + 1); + VectorizedArrowReader metadataLeaf = + new VectorizedArrowReader( + metadataColumn, metadataField, ArrowAllocation.rootAllocator(), false); + + VectorizedArrowReader.VectorizedVariantReader fileBacked = + new VectorizedArrowReader.VectorizedVariantReader( + variantField, metadataLeaf, VectorizedArrowReader.nulls()); + assertThat(fileBacked.fileBackedLeaf()).isSameAs(metadataLeaf); + + VectorizedArrowReader.VectorizedVariantReader noLeaf = + new VectorizedArrowReader.VectorizedVariantReader( + variantField, VectorizedArrowReader.nulls(), VectorizedArrowReader.nulls()); + assertThat(noLeaf.fileBackedLeaf()).isNull(); + } + + @Test + void nestedStructReusesDescendantPresenceReader() { + NestedField added = NestedField.optional(4, "added", IntegerType.get()); + NestedField inner = NestedField.optional(2, "inner", StructType.of(added)); + NestedField outer = NestedField.optional(1, "outer", StructType.of(inner)); + ColumnDescriptor presenceColumn = + new ColumnDescriptor( + new String[] {"outer", "inner", "innerId"}, + Types.optional(PrimitiveTypeName.INT64).named("innerId"), + 0, + 2); + VectorizedArrowReader presenceLeaf = + new VectorizedArrowReader( + presenceColumn, + NestedField.optional(3, "innerId", LongType.get()), + ArrowAllocation.rootAllocator(), + false); + VectorizedArrowReader constantChild = + new VectorizedArrowReader.ConstantVectorReader<>(added, 7); + VectorizedArrowReader.VectorizedStructReader innerReader = + new VectorizedArrowReader.VectorizedStructReader( + inner, ImmutableList.>of(constantChild), 2, presenceLeaf); + VectorizedArrowReader.VectorizedStructReader outerReader = + new VectorizedArrowReader.VectorizedStructReader( + outer, ImmutableList.>of(innerReader), 1, null); + + // outer has no file-backed child, so it reuses inner's presence reader rather than build one + assertThat(outerReader.fileBackedLeaf()).isSameAs(presenceLeaf); + } +} diff --git a/data/src/test/java/org/apache/iceberg/data/RandomGenericData.java b/data/src/test/java/org/apache/iceberg/data/RandomGenericData.java index 583513d3f90d..afecf890c7f9 100644 --- a/data/src/test/java/org/apache/iceberg/data/RandomGenericData.java +++ b/data/src/test/java/org/apache/iceberg/data/RandomGenericData.java @@ -54,6 +54,13 @@ public static List generate(Schema schema, int numRecords, long seed) { generateIcebergGenerics(schema, numRecords, () -> new RandomRecordGenerator(seed))); } + public static List generate( + Schema schema, int numRecords, long seed, float nullPercentage) { + return Lists.newArrayList( + generateIcebergGenerics( + schema, numRecords, () -> new RandomRecordGenerator(seed, nullPercentage))); + } + public static Iterable generateFallbackRecords( Schema schema, int numRecords, long seed, long numDictRows) { return generateIcebergGenerics( diff --git a/parquet/src/main/java/org/apache/iceberg/parquet/MessageTypeToType.java b/parquet/src/main/java/org/apache/iceberg/parquet/MessageTypeToType.java index 2b01bf882e75..5d791e66f35c 100644 --- a/parquet/src/main/java/org/apache/iceberg/parquet/MessageTypeToType.java +++ b/parquet/src/main/java/org/apache/iceberg/parquet/MessageTypeToType.java @@ -149,6 +149,10 @@ public Type map(GroupType map, Type keyType, Type valueType) { @Override public Type primitive(PrimitiveType primitive) { + return convertPrimitive(primitive); + } + + static Type convertPrimitive(PrimitiveType primitive) { // first, use the logical type annotation, if present LogicalTypeAnnotation logicalType = primitive.getLogicalTypeAnnotation(); if (logicalType != null) { diff --git a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetSchemaUtil.java b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetSchemaUtil.java index a8daec60bda0..2fd29e6f3a1f 100644 --- a/parquet/src/main/java/org/apache/iceberg/parquet/ParquetSchemaUtil.java +++ b/parquet/src/main/java/org/apache/iceberg/parquet/ParquetSchemaUtil.java @@ -25,6 +25,7 @@ import java.util.function.Function; import org.apache.iceberg.Schema; import org.apache.iceberg.mapping.NameMapping; +import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import org.apache.iceberg.relocated.com.google.common.collect.Lists; import org.apache.iceberg.relocated.com.google.common.collect.Sets; import org.apache.iceberg.types.TypeUtil; @@ -184,10 +185,23 @@ public Void variant(Types.VariantType expected, GroupType variantGroup, Void res } /** Shallowest leaf under path; its definition level shows whether the struct is present. */ - static ColumnDescriptor selectPresenceColumn(MessageType fileSchema, String[] path) { + public static ColumnDescriptor selectPresenceColumn(MessageType fileSchema, String[] path) { return selectPresenceColumn(fileSchema, path, leafColumns(fileSchema, path)); } + /** Iceberg field for a non-projected presence leaf, so a reader can allocate a value vector. */ + public static Types.NestedField presenceField(ColumnDescriptor leaf) { + PrimitiveType primitive = leaf.getPrimitiveType(); + Preconditions.checkArgument( + primitive.getId() != null, + "Cannot build presence field for column without an ID: %s", + leaf); + return Types.NestedField.optional( + primitive.getId().intValue(), + primitive.getName(), + MessageTypeToType.convertPrimitive(primitive)); + } + private static ColumnDescriptor selectPresenceColumn( MessageType fileSchema, String[] path, List leaves) { if (fileSchema.getMaxDefinitionLevel(path) <= 0) { diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java index 61616a9f233c..b9b80d02cd25 100644 --- a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java +++ b/spark/v3.5/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.StructVectorHolder) { + return new StructColumnVector((VectorHolder.StructVectorHolder) 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/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java new file mode 100644 index 000000000000..59e9c9bf1de7 --- /dev/null +++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java @@ -0,0 +1,143 @@ +/* + * 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 java.util.List; +import org.apache.iceberg.arrow.vectorized.NullabilityHolder; +import org.apache.iceberg.arrow.vectorized.VectorHolder; +import org.apache.iceberg.spark.SparkSchemaUtil; +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 StructColumnVector extends ColumnVector { + private final ColumnVector[] children; + private final NullabilityHolder nullabilityHolder; + + StructColumnVector(VectorHolder.StructVectorHolder holder) { + super(SparkSchemaUtil.convert(holder.icebergType())); + List childHolders = holder.childHolders(); + int numRows = holder.numValues(); + ColumnVectorBuilder builder = new ColumnVectorBuilder(); + this.children = new ColumnVector[childHolders.size()]; + for (int idx = 0; idx < childHolders.size(); idx++) { + children[idx] = builder.build(childHolders.get(idx), numRows); + } + + this.nullabilityHolder = holder.nullabilityHolder(); + } + + @Override + public void close() { + for (ColumnVector child : children) { + child.close(); + } + } + + @Override + public void closeIfFreeable() { + // See SPARK-50235, SPARK-50463 + } + + @Override + public boolean hasNull() { + return nullabilityHolder != null && nullabilityHolder.hasNulls(); + } + + @Override + public int numNulls() { + return nullabilityHolder == null ? 0 : nullabilityHolder.numNulls(); + } + + @Override + public boolean isNullAt(int rowId) { + return nullabilityHolder != null && nullabilityHolder.isNullAt(rowId) == 1; + } + + @Override + public ColumnVector getChild(int ordinal) { + return children[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("Struct column only supports getChild()"); + } +} diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index 2109936c96b9..7d2cbcc585a8 100644 --- a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -172,7 +172,21 @@ private boolean supportsParquetBatchReads(ScanTask task) { } private boolean supportsParquetBatchReads(Types.NestedField field) { - return field.type().isPrimitiveType() || MetadataColumns.isMetadataColumn(field.fieldId()); + if (MetadataColumns.isMetadataColumn(field.fieldId())) { + return true; + } + + if (field.type().isStructType()) { + for (Types.NestedField child : field.type().asStructType().fields()) { + if (!supportsParquetBatchReads(child)) { + return false; + } + } + + return true; + } + + return field.type().isPrimitiveType(); } // conditions for using ORC batch reads: diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java index 1a5e9cd93f17..c5913380ba2c 100644 --- a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java +++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java @@ -38,6 +38,7 @@ import java.util.function.Consumer; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; +import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.Files; import org.apache.iceberg.Schema; import org.apache.iceberg.arrow.ArrowAllocation; @@ -45,6 +46,7 @@ import org.apache.iceberg.data.Record; import org.apache.iceberg.data.parquet.GenericParquetReaders; import org.apache.iceberg.data.parquet.GenericParquetWriter; +import org.apache.iceberg.expressions.Literal; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.FileAppender; import org.apache.iceberg.parquet.Parquet; @@ -55,7 +57,6 @@ import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.collect.Iterables; import org.apache.iceberg.relocated.com.google.common.collect.Lists; -import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.apache.iceberg.spark.data.AvroDataTestBase; import org.apache.iceberg.spark.data.GenericsHelpers; import org.apache.iceberg.spark.data.RandomData; @@ -64,11 +65,13 @@ import org.apache.iceberg.types.Type.PrimitiveType; import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; +import org.apache.parquet.column.Encoding; import org.apache.parquet.column.ParquetProperties; +import org.apache.parquet.hadoop.ParquetFileReader; import org.apache.parquet.hadoop.ParquetOutputFormat; -import org.apache.parquet.schema.GroupType; -import org.apache.parquet.schema.MessageType; -import org.apache.parquet.schema.Type; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.vectorized.ColumnarBatch; import org.junit.jupiter.api.Test; @@ -234,7 +237,7 @@ Iterable generateData( long seed, float nullPercentage, Function transform) { - Iterable data = RandomGenericData.generate(schema, numRecords, seed); + Iterable data = RandomGenericData.generate(schema, numRecords, seed, nullPercentage); return transform == IDENTITY ? data : Iterables.transform(data, transform); } @@ -320,17 +323,162 @@ void assertRecordsMatch( @Test @Override - public void testNestedStruct() { - assertThatThrownBy( - () -> - VectorizedSparkParquetReaders.buildReader( - TypeUtil.assignIncreasingFreshIds( - new Schema(required(1, "struct", SUPPORTED_PRIMITIVES))), - new MessageType( - "struct", new GroupType(Type.Repetition.OPTIONAL, "struct").withId(1)), - Maps.newHashMap())) - .isInstanceOf(UnsupportedOperationException.class) - .hasMessage("Vectorized reads are not supported yet for struct fields"); + public void testNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(required(1, "struct", SUPPORTED_PRIMITIVES)))); + } + + @Test + void nullableNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(optional(1, "struct", SUPPORTED_PRIMITIVES))), + getNumRows(), + 0L, + 0.5f, + true); + } + + @Test + void threeLevelNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds( + new Schema( + optional( + 1, + "outer", + Types.StructType.of( + optional( + 2, + "middle", + Types.StructType.of(required(3, "inner", SUPPORTED_PRIMITIVES))))))), + getNumRows(), + 0L, + 0.5f, + true); + } + + @Test + void dictionaryEncodedStructChild() throws IOException { + // Assert children are dictionary-encoded so this is not silently a plain-encoding test. + Schema schema = + TypeUtil.assignIncreasingFreshIds( + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.IntegerType.get()))))); + int numRows = 1000; + File dataFile = temp.resolve("dict-struct.parquet").toFile(); + Iterable data = + RandomGenericData.generateDictionaryEncodableRecords(schema, numRows, 0L, 0.5f); + try (FileAppender writer = getParquetWriter(schema, dataFile)) { + writer.addAll(data); + } + + try (ParquetFileReader reader = + ParquetFileReader.open( + HadoopInputFile.fromPath( + new org.apache.hadoop.fs.Path(dataFile.toString()), new Configuration()))) { + boolean childUsesDictionary = false; + for (BlockMetaData block : reader.getFooter().getBlocks()) { + for (ColumnChunkMetaData column : block.getColumns()) { + if (column.getPath().toDotString().endsWith("int_data")) { + childUsesDictionary = column.getEncodings().stream().anyMatch(Encoding::usesDictionary); + } + } + } + + assertThat(childUsesDictionary) + .as("struct child int_data should be dictionary-encoded") + .isTrue(); + } + + assertRecordsMatch(schema, numRows, data, dataFile, false, BATCH_SIZE); + } + + @Test + void nestedStructWithReorderedAndDefaultedFields() throws IOException { + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "data", Types.StringType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + optional(3, "data", Types.StringType.get()), + required(2, "id", Types.LongType.get()), + Types.NestedField.optional("added") + .withId(4) + .ofType(Types.IntegerType.get()) + .withInitialDefault(Literal.of(42)) + .build()))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void nestedStructProjectingOnlyAddedField() throws IOException { + // project only a field added after the file was written: the struct must read per-row null + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "data", Types.StringType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + Types.NestedField.optional(4, "added", Types.IntegerType.get())))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void nestedStructWithPromotedChild() throws IOException { + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.IntegerType.get()), + optional(4, "float_data", Types.FloatType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.LongType.get()), + optional(4, "float_data", Types.DoubleType.get())))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void allNullNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(optional(1, "struct", SUPPORTED_PRIMITIVES))), + getNumRows(), + 0L, + 1.0f, + true); } @Test diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java index 9cf8e435270c..849c5024055d 100644 --- a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java +++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java @@ -22,6 +22,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.within; +import static org.assertj.core.api.Assumptions.assumeThat; import java.io.IOException; import java.nio.file.Path; @@ -209,6 +210,56 @@ public void testRename() throws Exception { assertThat(cmp).as("Should contain the correct data/renamed value").isEqualTo(0); } + @TestTemplate + void requiredInnerStructUnderPresentParentReadsPresent() throws Exception { + assumeThat(format).isEqualTo(FileFormat.PARQUET); + Schema writeSchema = + new Schema( + Types.NestedField.optional( + 1, + "outer", + Types.StructType.of( + Types.NestedField.required(2, "real", Types.LongType.get()), + Types.NestedField.required( + 3, + "inner", + Types.StructType.of( + Types.NestedField.required(4, "orig", Types.LongType.get())))))); + + Types.StructType outerType = writeSchema.findType("outer").asStructType(); + Record record = GenericRecord.create(writeSchema); + Record outer = GenericRecord.create(outerType); + Record inner = GenericRecord.create(outerType.field("inner").type().asStructType()); + inner.setField("orig", 200L); + outer.setField("real", 100L); + outer.setField("inner", inner); + record.setField("outer", outer); + + // project outer's stored field plus only a field added to the required inner after the write + Schema readSchema = + new Schema( + Types.NestedField.optional( + 1, + "outer", + Types.StructType.of( + Types.NestedField.required(2, "real", Types.LongType.get()), + Types.NestedField.required( + 3, + "inner", + Types.StructType.of( + Types.NestedField.optional(5, "added", Types.IntegerType.get())))))); + + Record projected = writeAndRead("required_inner_added_field", writeSchema, readSchema, record); + Record projectedOuter = (Record) projected.getField("outer"); + assertThat(projectedOuter).as("present parent should project").isNotNull(); + assertThat((long) projectedOuter.getField("real")).isEqualTo(100L); + Record projectedInner = (Record) projectedOuter.getField("inner"); + assertThat(projectedInner) + .as("required inner under a present parent should be present, not null") + .isNotNull(); + assertThat(projectedInner.getField("added")).as("added field should be null").isNull(); + } + @TestTemplate public void testNestedStructProjection() throws Exception { Schema writeSchema = diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java index 2fb6933ce933..47006ab5dd1f 100644 --- a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java +++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java @@ -20,6 +20,7 @@ import static org.apache.hadoop.hive.conf.HiveConf.ConfVars.METASTOREURIS; import static org.apache.iceberg.spark.source.SparkSQLExecutionHelper.lastExecutedMetricValue; +import static org.apache.iceberg.types.Types.NestedField.optional; import static org.apache.iceberg.types.Types.NestedField.required; import static org.apache.spark.sql.types.DataTypes.IntegerType; import static org.assertj.core.api.Assertions.assertThat; @@ -379,6 +380,71 @@ public void testPosDeletesAllRowsInBatch() throws IOException { checkDeleteCount(4L); } + @TestTemplate + void posDeletesOnStructColumn() throws IOException { + Schema structSchema = + new Schema( + required(1, "id", Types.IntegerType.get()), + optional( + 2, + "struct_col", + Types.StructType.of( + required(3, "nested_id", Types.LongType.get()), + optional(4, "nested_data", Types.StringType.get())))); + + String structTableName = "test_struct_pos_deletes"; + dropTable(structTableName); + Table structTable = createTable(structTableName, structSchema, PartitionSpec.unpartitioned()); + + Types.StructType nestedType = structSchema.findField("struct_col").type().asStructType(); + GenericRecord nested = GenericRecord.create(nestedType); + GenericRecord row = GenericRecord.create(structSchema); + List structRecords = Lists.newArrayList(); + for (int idx = 0; idx < 7; idx++) { + structRecords.add( + row.copy( + "id", + idx, + "struct_col", + nested.copy("nested_id", (long) (idx * 10), "nested_data", "v" + idx))); + } + + DataFile structDataFile = + FileHelpers.writeDataFile( + structTable, + Files.localOutput(File.createTempFile("junit", null, temp.toFile())), + structRecords); + structTable.newAppend().appendFile(structDataFile).commit(); + + // delete positions 0 and 5 (ids 0 and 5), spanning the 4-row vectorized batch boundary + List> deletes = + Lists.newArrayList( + Pair.of(structDataFile.location(), 0L), Pair.of(structDataFile.location(), 5L)); + + Pair posDeletes = + FileHelpers.writeDeleteFile( + structTable, + Files.localOutput(File.createTempFile("junit", null, temp.toFile())), + TestHelpers.Row.of(), + deletes, + formatVersion); + + structTable + .newRowDelta() + .addDeletes(posDeletes.first()) + .validateDataFilesExist(posDeletes.second()) + .commit(); + + StructLikeSet expected = rowSetWithoutIds(structTable, structRecords, 0, 5); + StructLikeSet actual = rowSet(structTableName, structTable, "*"); + + assertThat(actual) + .as("Struct column rows should survive positional deletes") + .isEqualTo(expected); + + dropTable(structTableName); + } + @TestTemplate public void testPosDeletesWithDeletedColumn() throws IOException { // read.parquet.vectorization.batch-size is set to 4, so the 4 rows in the first batch are all 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 2a14c2262c6a..cf37a2a3a28b 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 @@ -29,6 +29,8 @@ class ColumnVectorBuilder { public ColumnVector build(VectorHolder holder, int numRows) { if (holder instanceof VectorHolder.VariantVectorHolder) { return new VariantColumnVector((VectorHolder.VariantVectorHolder) holder); + } else if (holder instanceof VectorHolder.StructVectorHolder) { + return new StructColumnVector((VectorHolder.StructVectorHolder) holder); } else if (holder.isDummy()) { if (holder instanceof VectorHolder.DeletedVectorHolder) { return new DeletedColumnVector(Types.BooleanType.get()); diff --git a/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java new file mode 100644 index 000000000000..59e9c9bf1de7 --- /dev/null +++ b/spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java @@ -0,0 +1,143 @@ +/* + * 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 java.util.List; +import org.apache.iceberg.arrow.vectorized.NullabilityHolder; +import org.apache.iceberg.arrow.vectorized.VectorHolder; +import org.apache.iceberg.spark.SparkSchemaUtil; +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 StructColumnVector extends ColumnVector { + private final ColumnVector[] children; + private final NullabilityHolder nullabilityHolder; + + StructColumnVector(VectorHolder.StructVectorHolder holder) { + super(SparkSchemaUtil.convert(holder.icebergType())); + List childHolders = holder.childHolders(); + int numRows = holder.numValues(); + ColumnVectorBuilder builder = new ColumnVectorBuilder(); + this.children = new ColumnVector[childHolders.size()]; + for (int idx = 0; idx < childHolders.size(); idx++) { + children[idx] = builder.build(childHolders.get(idx), numRows); + } + + this.nullabilityHolder = holder.nullabilityHolder(); + } + + @Override + public void close() { + for (ColumnVector child : children) { + child.close(); + } + } + + @Override + public void closeIfFreeable() { + // See SPARK-50235, SPARK-50463 + } + + @Override + public boolean hasNull() { + return nullabilityHolder != null && nullabilityHolder.hasNulls(); + } + + @Override + public int numNulls() { + return nullabilityHolder == null ? 0 : nullabilityHolder.numNulls(); + } + + @Override + public boolean isNullAt(int rowId) { + return nullabilityHolder != null && nullabilityHolder.isNullAt(rowId) == 1; + } + + @Override + public ColumnVector getChild(int ordinal) { + return children[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("Struct column only supports getChild()"); + } +} 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 f90423900d62..9813a3e49391 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 @@ -22,7 +22,9 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.function.Supplier; +import java.util.stream.Collectors; import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.MetadataColumns; @@ -42,6 +44,7 @@ import org.apache.iceberg.spark.ParquetBatchReadConf; import org.apache.iceberg.spark.SparkReadConf; import org.apache.iceberg.spark.SparkUtil; +import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; import org.apache.iceberg.util.PropertyUtil; import org.apache.spark.api.java.JavaSparkContext; @@ -65,6 +68,7 @@ class SparkBatch implements Batch { private final boolean executorCacheLocalityEnabled; private final int scanHashCode; private final boolean cacheDeleteFilesOnExecutors; + private final Set projectedVariantIds; SparkBatch( JavaSparkContext sparkContext, @@ -88,6 +92,11 @@ class SparkBatch implements Batch { this.executorCacheLocalityEnabled = readConf.executorCacheLocalityEnabled(); this.scanHashCode = scanHashCode; this.cacheDeleteFilesOnExecutors = readConf.cacheDeleteFilesOnExecutors(); + this.projectedVariantIds = + TypeUtil.indexById(expectedSchema.asStruct()).values().stream() + .filter(field -> field.type().isVariantType()) + .map(Types.NestedField::fieldId) + .collect(Collectors.toSet()); } @Override @@ -175,12 +184,8 @@ private boolean supportsParquetBatchReads(ScanTask task) { 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; - } - } + if (lowerBounds != null && hasVariantWithLowerBound(lowerBounds)) { + return false; } return true; @@ -189,6 +194,16 @@ private boolean supportsParquetBatchReads(ScanTask task) { } } + private boolean hasVariantWithLowerBound(Map lowerBounds) { + for (Integer variantId : projectedVariantIds) { + if (lowerBounds.containsKey(variantId)) { + return true; + } + } + + return false; + } + private boolean supportsParquetBatchReads(Types.NestedField field) { if (field.type().isVariantType()) { boolean shredVariants = @@ -208,6 +223,16 @@ private boolean supportsParquetBatchReads(Types.NestedField field) { } } + if (field.type().isStructType()) { + for (Types.NestedField child : field.type().asStructType().fields()) { + if (!supportsParquetBatchReads(child)) { + return false; + } + } + + return true; + } + return field.type().isPrimitiveType() || field.type().isVariantType() || MetadataColumns.isMetadataColumn(field.fieldId()); diff --git a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java index bf6ff7907f75..e45b650fd1ea 100644 --- a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java +++ b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java @@ -38,6 +38,7 @@ import java.util.function.Consumer; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; +import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.Files; import org.apache.iceberg.Schema; import org.apache.iceberg.arrow.ArrowAllocation; @@ -45,6 +46,7 @@ import org.apache.iceberg.data.Record; import org.apache.iceberg.data.parquet.GenericParquetReaders; import org.apache.iceberg.data.parquet.GenericParquetWriter; +import org.apache.iceberg.expressions.Literal; import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.io.FileAppender; import org.apache.iceberg.parquet.Parquet; @@ -55,7 +57,6 @@ import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.collect.Iterables; import org.apache.iceberg.relocated.com.google.common.collect.Lists; -import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.apache.iceberg.spark.data.AvroDataTestBase; import org.apache.iceberg.spark.data.GenericsHelpers; import org.apache.iceberg.spark.data.RandomData; @@ -64,11 +65,13 @@ import org.apache.iceberg.types.Type.PrimitiveType; import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; +import org.apache.parquet.column.Encoding; import org.apache.parquet.column.ParquetProperties; +import org.apache.parquet.hadoop.ParquetFileReader; import org.apache.parquet.hadoop.ParquetOutputFormat; -import org.apache.parquet.schema.GroupType; -import org.apache.parquet.schema.MessageType; -import org.apache.parquet.schema.Type; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.vectorized.ColumnarBatch; import org.junit.jupiter.api.Test; @@ -234,7 +237,7 @@ Iterable generateData( long seed, float nullPercentage, Function transform) { - Iterable data = RandomGenericData.generate(schema, numRecords, seed); + Iterable data = RandomGenericData.generate(schema, numRecords, seed, nullPercentage); return transform == IDENTITY ? data : Iterables.transform(data, transform); } @@ -320,17 +323,162 @@ void assertRecordsMatch( @Test @Override - public void testNestedStruct() { - assertThatThrownBy( - () -> - VectorizedSparkParquetReaders.buildReader( - TypeUtil.assignIncreasingFreshIds( - new Schema(required(1, "struct", SUPPORTED_PRIMITIVES))), - new MessageType( - "struct", new GroupType(Type.Repetition.OPTIONAL, "struct").withId(1)), - Maps.newHashMap())) - .isInstanceOf(UnsupportedOperationException.class) - .hasMessage("Vectorized reads are not supported yet for struct fields"); + public void testNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(required(1, "struct", SUPPORTED_PRIMITIVES)))); + } + + @Test + void nullableNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(optional(1, "struct", SUPPORTED_PRIMITIVES))), + getNumRows(), + 0L, + 0.5f, + true); + } + + @Test + void threeLevelNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds( + new Schema( + optional( + 1, + "outer", + Types.StructType.of( + optional( + 2, + "middle", + Types.StructType.of(required(3, "inner", SUPPORTED_PRIMITIVES))))))), + getNumRows(), + 0L, + 0.5f, + true); + } + + @Test + void dictionaryEncodedStructChild() throws IOException { + // Assert children are dictionary-encoded so this is not silently a plain-encoding test. + Schema schema = + TypeUtil.assignIncreasingFreshIds( + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.IntegerType.get()))))); + int numRows = 1000; + File dataFile = temp.resolve("dict-struct.parquet").toFile(); + Iterable data = + RandomGenericData.generateDictionaryEncodableRecords(schema, numRows, 0L, 0.5f); + try (FileAppender writer = getParquetWriter(schema, dataFile)) { + writer.addAll(data); + } + + try (ParquetFileReader reader = + ParquetFileReader.open( + HadoopInputFile.fromPath( + new org.apache.hadoop.fs.Path(dataFile.toString()), new Configuration()))) { + boolean childUsesDictionary = false; + for (BlockMetaData block : reader.getFooter().getBlocks()) { + for (ColumnChunkMetaData column : block.getColumns()) { + if (column.getPath().toDotString().endsWith("int_data")) { + childUsesDictionary = column.getEncodings().stream().anyMatch(Encoding::usesDictionary); + } + } + } + + assertThat(childUsesDictionary) + .as("struct child int_data should be dictionary-encoded") + .isTrue(); + } + + assertRecordsMatch(schema, numRows, data, dataFile, false, BATCH_SIZE); + } + + @Test + void nestedStructWithReorderedAndDefaultedFields() throws IOException { + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "data", Types.StringType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + optional(3, "data", Types.StringType.get()), + required(2, "id", Types.LongType.get()), + Types.NestedField.optional("added") + .withId(4) + .ofType(Types.IntegerType.get()) + .withInitialDefault(Literal.of(42)) + .build()))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void nestedStructProjectingOnlyAddedField() throws IOException { + // project only a field added after the file was written: the struct must read per-row null + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "data", Types.StringType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + Types.NestedField.optional(4, "added", Types.IntegerType.get())))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void nestedStructWithPromotedChild() throws IOException { + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.IntegerType.get()), + optional(4, "float_data", Types.FloatType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.LongType.get()), + optional(4, "float_data", Types.DoubleType.get())))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void allNullNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(optional(1, "struct", SUPPORTED_PRIMITIVES))), + getNumRows(), + 0L, + 1.0f, + true); } @Test diff --git a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java index 9cf8e435270c..849c5024055d 100644 --- a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java +++ b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java @@ -22,6 +22,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.within; +import static org.assertj.core.api.Assumptions.assumeThat; import java.io.IOException; import java.nio.file.Path; @@ -209,6 +210,56 @@ public void testRename() throws Exception { assertThat(cmp).as("Should contain the correct data/renamed value").isEqualTo(0); } + @TestTemplate + void requiredInnerStructUnderPresentParentReadsPresent() throws Exception { + assumeThat(format).isEqualTo(FileFormat.PARQUET); + Schema writeSchema = + new Schema( + Types.NestedField.optional( + 1, + "outer", + Types.StructType.of( + Types.NestedField.required(2, "real", Types.LongType.get()), + Types.NestedField.required( + 3, + "inner", + Types.StructType.of( + Types.NestedField.required(4, "orig", Types.LongType.get())))))); + + Types.StructType outerType = writeSchema.findType("outer").asStructType(); + Record record = GenericRecord.create(writeSchema); + Record outer = GenericRecord.create(outerType); + Record inner = GenericRecord.create(outerType.field("inner").type().asStructType()); + inner.setField("orig", 200L); + outer.setField("real", 100L); + outer.setField("inner", inner); + record.setField("outer", outer); + + // project outer's stored field plus only a field added to the required inner after the write + Schema readSchema = + new Schema( + Types.NestedField.optional( + 1, + "outer", + Types.StructType.of( + Types.NestedField.required(2, "real", Types.LongType.get()), + Types.NestedField.required( + 3, + "inner", + Types.StructType.of( + Types.NestedField.optional(5, "added", Types.IntegerType.get())))))); + + Record projected = writeAndRead("required_inner_added_field", writeSchema, readSchema, record); + Record projectedOuter = (Record) projected.getField("outer"); + assertThat(projectedOuter).as("present parent should project").isNotNull(); + assertThat((long) projectedOuter.getField("real")).isEqualTo(100L); + Record projectedInner = (Record) projectedOuter.getField("inner"); + assertThat(projectedInner) + .as("required inner under a present parent should be present, not null") + .isNotNull(); + assertThat(projectedInner.getField("added")).as("added field should be null").isNull(); + } + @TestTemplate public void testNestedStructProjection() throws Exception { Schema writeSchema = diff --git a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java index 2fb6933ce933..47006ab5dd1f 100644 --- a/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java +++ b/spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java @@ -20,6 +20,7 @@ import static org.apache.hadoop.hive.conf.HiveConf.ConfVars.METASTOREURIS; import static org.apache.iceberg.spark.source.SparkSQLExecutionHelper.lastExecutedMetricValue; +import static org.apache.iceberg.types.Types.NestedField.optional; import static org.apache.iceberg.types.Types.NestedField.required; import static org.apache.spark.sql.types.DataTypes.IntegerType; import static org.assertj.core.api.Assertions.assertThat; @@ -379,6 +380,71 @@ public void testPosDeletesAllRowsInBatch() throws IOException { checkDeleteCount(4L); } + @TestTemplate + void posDeletesOnStructColumn() throws IOException { + Schema structSchema = + new Schema( + required(1, "id", Types.IntegerType.get()), + optional( + 2, + "struct_col", + Types.StructType.of( + required(3, "nested_id", Types.LongType.get()), + optional(4, "nested_data", Types.StringType.get())))); + + String structTableName = "test_struct_pos_deletes"; + dropTable(structTableName); + Table structTable = createTable(structTableName, structSchema, PartitionSpec.unpartitioned()); + + Types.StructType nestedType = structSchema.findField("struct_col").type().asStructType(); + GenericRecord nested = GenericRecord.create(nestedType); + GenericRecord row = GenericRecord.create(structSchema); + List structRecords = Lists.newArrayList(); + for (int idx = 0; idx < 7; idx++) { + structRecords.add( + row.copy( + "id", + idx, + "struct_col", + nested.copy("nested_id", (long) (idx * 10), "nested_data", "v" + idx))); + } + + DataFile structDataFile = + FileHelpers.writeDataFile( + structTable, + Files.localOutput(File.createTempFile("junit", null, temp.toFile())), + structRecords); + structTable.newAppend().appendFile(structDataFile).commit(); + + // delete positions 0 and 5 (ids 0 and 5), spanning the 4-row vectorized batch boundary + List> deletes = + Lists.newArrayList( + Pair.of(structDataFile.location(), 0L), Pair.of(structDataFile.location(), 5L)); + + Pair posDeletes = + FileHelpers.writeDeleteFile( + structTable, + Files.localOutput(File.createTempFile("junit", null, temp.toFile())), + TestHelpers.Row.of(), + deletes, + formatVersion); + + structTable + .newRowDelta() + .addDeletes(posDeletes.first()) + .validateDataFilesExist(posDeletes.second()) + .commit(); + + StructLikeSet expected = rowSetWithoutIds(structTable, structRecords, 0, 5); + StructLikeSet actual = rowSet(structTableName, structTable, "*"); + + assertThat(actual) + .as("Struct column rows should survive positional deletes") + .isEqualTo(expected); + + dropTable(structTableName); + } + @TestTemplate public void testPosDeletesWithDeletedColumn() throws IOException { // read.parquet.vectorization.batch-size is set to 4, so the 4 rows in the first batch are all 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 2a14c2262c6a..cf37a2a3a28b 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 @@ -29,6 +29,8 @@ class ColumnVectorBuilder { public ColumnVector build(VectorHolder holder, int numRows) { if (holder instanceof VectorHolder.VariantVectorHolder) { return new VariantColumnVector((VectorHolder.VariantVectorHolder) holder); + } else if (holder instanceof VectorHolder.StructVectorHolder) { + return new StructColumnVector((VectorHolder.StructVectorHolder) holder); } else if (holder.isDummy()) { if (holder instanceof VectorHolder.DeletedVectorHolder) { return new DeletedColumnVector(Types.BooleanType.get()); diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java new file mode 100644 index 000000000000..59e9c9bf1de7 --- /dev/null +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java @@ -0,0 +1,143 @@ +/* + * 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 java.util.List; +import org.apache.iceberg.arrow.vectorized.NullabilityHolder; +import org.apache.iceberg.arrow.vectorized.VectorHolder; +import org.apache.iceberg.spark.SparkSchemaUtil; +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 StructColumnVector extends ColumnVector { + private final ColumnVector[] children; + private final NullabilityHolder nullabilityHolder; + + StructColumnVector(VectorHolder.StructVectorHolder holder) { + super(SparkSchemaUtil.convert(holder.icebergType())); + List childHolders = holder.childHolders(); + int numRows = holder.numValues(); + ColumnVectorBuilder builder = new ColumnVectorBuilder(); + this.children = new ColumnVector[childHolders.size()]; + for (int idx = 0; idx < childHolders.size(); idx++) { + children[idx] = builder.build(childHolders.get(idx), numRows); + } + + this.nullabilityHolder = holder.nullabilityHolder(); + } + + @Override + public void close() { + for (ColumnVector child : children) { + child.close(); + } + } + + @Override + public void closeIfFreeable() { + // See SPARK-50235, SPARK-50463 + } + + @Override + public boolean hasNull() { + return nullabilityHolder != null && nullabilityHolder.hasNulls(); + } + + @Override + public int numNulls() { + return nullabilityHolder == null ? 0 : nullabilityHolder.numNulls(); + } + + @Override + public boolean isNullAt(int rowId) { + return nullabilityHolder != null && nullabilityHolder.isNullAt(rowId) == 1; + } + + @Override + public ColumnVector getChild(int ordinal) { + return children[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("Struct column only supports getChild()"); + } +} 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 5cafa1c9fa04..8f42b938a3f3 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 @@ -22,7 +22,9 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.function.Supplier; +import java.util.stream.Collectors; import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.MetadataColumns; @@ -43,6 +45,7 @@ import org.apache.iceberg.spark.SparkReadConf; import org.apache.iceberg.spark.SparkUtil; import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; import org.apache.iceberg.util.PropertyUtil; import org.apache.spark.api.java.JavaSparkContext; @@ -65,6 +68,7 @@ class SparkBatch implements Batch { private final boolean executorCacheLocalityEnabled; private final int scanHashCode; private final boolean cacheDeleteFilesOnExecutors; + private final Set projectedVariantIds; SparkBatch( JavaSparkContext sparkContext, @@ -87,6 +91,11 @@ class SparkBatch implements Batch { this.executorCacheLocalityEnabled = readConf.executorCacheLocalityEnabled(); this.scanHashCode = scanHashCode; this.cacheDeleteFilesOnExecutors = readConf.cacheDeleteFilesOnExecutors(); + this.projectedVariantIds = + TypeUtil.indexById(projection.asStruct()).values().stream() + .filter(field -> field.type().isVariantType()) + .map(Types.NestedField::fieldId) + .collect(Collectors.toSet()); } @Override @@ -174,12 +183,8 @@ private boolean supportsParquetBatchReads(ScanTask task) { 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; - } - } + if (lowerBounds != null && hasVariantWithLowerBound(lowerBounds)) { + return false; } return true; @@ -188,6 +193,17 @@ private boolean supportsParquetBatchReads(ScanTask task) { } } + // a variant with stored bounds may be shredded (unsupported vectorized); check all nesting levels + private boolean hasVariantWithLowerBound(Map lowerBounds) { + for (Integer variantId : projectedVariantIds) { + if (lowerBounds.containsKey(variantId)) { + return true; + } + } + + return false; + } + private boolean supportsParquetBatchReads(Types.NestedField field) { if (MetadataColumns.isMetadataColumn(field.fieldId())) { return true; @@ -218,6 +234,16 @@ private boolean supportsParquetBatchReads(Types.NestedField field) { } } + if (type.isStructType()) { + for (Types.NestedField child : type.asStructType().fields()) { + if (!supportsParquetBatchReads(child)) { + return false; + } + } + + return true; + } + return type.isPrimitiveType() || type.isVariantType(); } diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java index 12047bd12847..1c59e7bff75c 100644 --- a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java +++ b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java @@ -39,6 +39,7 @@ import java.util.function.Consumer; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; +import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.Files; import org.apache.iceberg.Schema; import org.apache.iceberg.arrow.ArrowAllocation; @@ -57,7 +58,6 @@ import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.collect.Iterables; import org.apache.iceberg.relocated.com.google.common.collect.Lists; -import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.apache.iceberg.spark.data.AvroDataTestBase; import org.apache.iceberg.spark.data.GenericsHelpers; import org.apache.iceberg.spark.data.RandomData; @@ -66,11 +66,13 @@ import org.apache.iceberg.types.Type.PrimitiveType; import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; +import org.apache.parquet.column.Encoding; import org.apache.parquet.column.ParquetProperties; +import org.apache.parquet.hadoop.ParquetFileReader; import org.apache.parquet.hadoop.ParquetOutputFormat; -import org.apache.parquet.schema.GroupType; -import org.apache.parquet.schema.MessageType; -import org.apache.parquet.schema.Type; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.vectorized.ColumnarBatch; import org.junit.jupiter.api.Test; @@ -236,7 +238,7 @@ Iterable generateData( long seed, float nullPercentage, Function transform) { - Iterable data = RandomGenericData.generate(schema, numRecords, seed); + Iterable data = RandomGenericData.generate(schema, numRecords, seed, nullPercentage); return transform == IDENTITY ? data : Iterables.transform(data, transform); } @@ -332,17 +334,162 @@ void assertRecordsMatch( @Test @Override - public void testNestedStruct() { - assertThatThrownBy( - () -> - VectorizedSparkParquetReaders.buildReader( - TypeUtil.assignIncreasingFreshIds( - new Schema(required(1, "struct", SUPPORTED_PRIMITIVES))), - new MessageType( - "struct", new GroupType(Type.Repetition.OPTIONAL, "struct").withId(1)), - Maps.newHashMap())) - .isInstanceOf(UnsupportedOperationException.class) - .hasMessage("Vectorized reads are not supported yet for struct fields"); + public void testNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(required(1, "struct", SUPPORTED_PRIMITIVES)))); + } + + @Test + void nullableNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(optional(1, "struct", SUPPORTED_PRIMITIVES))), + getNumRows(), + 0L, + 0.5f, + true); + } + + @Test + void threeLevelNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds( + new Schema( + optional( + 1, + "outer", + Types.StructType.of( + optional( + 2, + "middle", + Types.StructType.of(required(3, "inner", SUPPORTED_PRIMITIVES))))))), + getNumRows(), + 0L, + 0.5f, + true); + } + + @Test + void dictionaryEncodedStructChild() throws IOException { + // Assert children are dictionary-encoded so this is not silently a plain-encoding test. + Schema schema = + TypeUtil.assignIncreasingFreshIds( + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.IntegerType.get()))))); + int numRows = 1000; + File dataFile = temp.resolve("dict-struct.parquet").toFile(); + Iterable data = + RandomGenericData.generateDictionaryEncodableRecords(schema, numRows, 0L, 0.5f); + try (FileAppender writer = getParquetWriter(schema, dataFile)) { + writer.addAll(data); + } + + try (ParquetFileReader reader = + ParquetFileReader.open( + HadoopInputFile.fromPath( + new org.apache.hadoop.fs.Path(dataFile.toString()), new Configuration()))) { + boolean childUsesDictionary = false; + for (BlockMetaData block : reader.getFooter().getBlocks()) { + for (ColumnChunkMetaData column : block.getColumns()) { + if (column.getPath().toDotString().endsWith("int_data")) { + childUsesDictionary = column.getEncodings().stream().anyMatch(Encoding::usesDictionary); + } + } + } + + assertThat(childUsesDictionary) + .as("struct child int_data should be dictionary-encoded") + .isTrue(); + } + + assertRecordsMatch(schema, numRows, data, dataFile, false, BATCH_SIZE); + } + + @Test + void nestedStructWithReorderedAndDefaultedFields() throws IOException { + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "data", Types.StringType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + optional(3, "data", Types.StringType.get()), + required(2, "id", Types.LongType.get()), + Types.NestedField.optional("added") + .withId(4) + .ofType(Types.IntegerType.get()) + .withInitialDefault(Literal.of(42)) + .build()))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void nestedStructProjectingOnlyAddedField() throws IOException { + // project only a field added after the file was written: the struct must read per-row null + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "data", Types.StringType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + Types.NestedField.optional(4, "added", Types.IntegerType.get())))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void nestedStructWithPromotedChild() throws IOException { + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.IntegerType.get()), + optional(4, "float_data", Types.FloatType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.LongType.get()), + optional(4, "float_data", Types.DoubleType.get())))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void allNullNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(optional(1, "struct", SUPPORTED_PRIMITIVES))), + getNumRows(), + 0L, + 1.0f, + true); } @Test diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java index 9cf8e435270c..096581b20c8e 100644 --- a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java +++ b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java @@ -22,6 +22,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.within; +import static org.assertj.core.api.Assumptions.assumeThat; import java.io.IOException; import java.nio.file.Path; @@ -284,6 +285,56 @@ public void testNestedStructProjection() throws Exception { .isCloseTo(-1.539054f, within(0.000001f)); } + @TestTemplate + void requiredInnerStructUnderPresentParentReadsPresent() throws Exception { + assumeThat(format).isEqualTo(FileFormat.PARQUET); + Schema writeSchema = + new Schema( + Types.NestedField.optional( + 1, + "outer", + Types.StructType.of( + Types.NestedField.required(2, "real", Types.LongType.get()), + Types.NestedField.required( + 3, + "inner", + Types.StructType.of( + Types.NestedField.required(4, "orig", Types.LongType.get())))))); + + Types.StructType outerType = writeSchema.findType("outer").asStructType(); + Record record = GenericRecord.create(writeSchema); + Record outer = GenericRecord.create(outerType); + Record inner = GenericRecord.create(outerType.field("inner").type().asStructType()); + inner.setField("orig", 200L); + outer.setField("real", 100L); + outer.setField("inner", inner); + record.setField("outer", outer); + + // project outer's stored field plus only a field added to the required inner after the write + Schema readSchema = + new Schema( + Types.NestedField.optional( + 1, + "outer", + Types.StructType.of( + Types.NestedField.required(2, "real", Types.LongType.get()), + Types.NestedField.required( + 3, + "inner", + Types.StructType.of( + Types.NestedField.optional(5, "added", Types.IntegerType.get())))))); + + Record projected = writeAndRead("required_inner_added_field", writeSchema, readSchema, record); + Record projectedOuter = (Record) projected.getField("outer"); + assertThat(projectedOuter).as("present parent should project").isNotNull(); + assertThat((long) projectedOuter.getField("real")).isEqualTo(100L); + Record projectedInner = (Record) projectedOuter.getField("inner"); + assertThat(projectedInner) + .as("required inner under a present parent should be present, not null") + .isNotNull(); + assertThat(projectedInner.getField("added")).as("added field should be null").isNull(); + } + @TestTemplate public void testMapProjection() throws IOException { Schema writeSchema = diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java index 0d619305716e..1b7c3ee2f18e 100644 --- a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java +++ b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java @@ -20,6 +20,7 @@ import static org.apache.hadoop.hive.conf.HiveConf.ConfVars.METASTOREURIS; import static org.apache.iceberg.spark.source.SparkSQLExecutionHelper.lastExecutedMetricValue; +import static org.apache.iceberg.types.Types.NestedField.optional; import static org.apache.iceberg.types.Types.NestedField.required; import static org.apache.spark.sql.types.DataTypes.IntegerType; import static org.assertj.core.api.Assertions.assertThat; @@ -381,6 +382,71 @@ public void testPosDeletesAllRowsInBatch() throws IOException { checkDeleteCount(4L); } + @TestTemplate + void posDeletesOnStructColumn() throws IOException { + Schema structSchema = + new Schema( + required(1, "id", Types.IntegerType.get()), + optional( + 2, + "struct_col", + Types.StructType.of( + required(3, "nested_id", Types.LongType.get()), + optional(4, "nested_data", Types.StringType.get())))); + + String structTableName = "test_struct_pos_deletes"; + dropTable(structTableName); + Table structTable = createTable(structTableName, structSchema, PartitionSpec.unpartitioned()); + + Types.StructType nestedType = structSchema.findField("struct_col").type().asStructType(); + GenericRecord nested = GenericRecord.create(nestedType); + GenericRecord row = GenericRecord.create(structSchema); + List structRecords = Lists.newArrayList(); + for (int idx = 0; idx < 7; idx++) { + structRecords.add( + row.copy( + "id", + idx, + "struct_col", + nested.copy("nested_id", (long) (idx * 10), "nested_data", "v" + idx))); + } + + DataFile structDataFile = + FileHelpers.writeDataFile( + structTable, + Files.localOutput(File.createTempFile("junit", null, temp.toFile())), + structRecords); + structTable.newAppend().appendFile(structDataFile).commit(); + + // delete positions 0 and 5 (ids 0 and 5), spanning the 4-row vectorized batch boundary + List> deletes = + Lists.newArrayList( + Pair.of(structDataFile.location(), 0L), Pair.of(structDataFile.location(), 5L)); + + Pair posDeletes = + FileHelpers.writeDeleteFile( + structTable, + Files.localOutput(File.createTempFile("junit", null, temp.toFile())), + TestHelpers.Row.of(), + deletes, + formatVersion); + + structTable + .newRowDelta() + .addDeletes(posDeletes.first()) + .validateDataFilesExist(posDeletes.second()) + .commit(); + + StructLikeSet expected = rowSetWithoutIds(structTable, structRecords, 0, 5); + StructLikeSet actual = rowSet(structTableName, structTable, "*"); + + assertThat(actual) + .as("Struct column rows should survive positional deletes") + .isEqualTo(expected); + + dropTable(structTableName); + } + @TestTemplate public void testPosDeletesWithDeletedColumn() throws IOException { // read.parquet.vectorization.batch-size is set to 4, so the 4 rows in the first batch are all diff --git a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java index 2a14c2262c6a..cf37a2a3a28b 100644 --- a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java +++ b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/ColumnVectorBuilder.java @@ -29,6 +29,8 @@ class ColumnVectorBuilder { public ColumnVector build(VectorHolder holder, int numRows) { if (holder instanceof VectorHolder.VariantVectorHolder) { return new VariantColumnVector((VectorHolder.VariantVectorHolder) holder); + } else if (holder instanceof VectorHolder.StructVectorHolder) { + return new StructColumnVector((VectorHolder.StructVectorHolder) holder); } else if (holder.isDummy()) { if (holder instanceof VectorHolder.DeletedVectorHolder) { return new DeletedColumnVector(Types.BooleanType.get()); diff --git a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java new file mode 100644 index 000000000000..59e9c9bf1de7 --- /dev/null +++ b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/StructColumnVector.java @@ -0,0 +1,143 @@ +/* + * 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 java.util.List; +import org.apache.iceberg.arrow.vectorized.NullabilityHolder; +import org.apache.iceberg.arrow.vectorized.VectorHolder; +import org.apache.iceberg.spark.SparkSchemaUtil; +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 StructColumnVector extends ColumnVector { + private final ColumnVector[] children; + private final NullabilityHolder nullabilityHolder; + + StructColumnVector(VectorHolder.StructVectorHolder holder) { + super(SparkSchemaUtil.convert(holder.icebergType())); + List childHolders = holder.childHolders(); + int numRows = holder.numValues(); + ColumnVectorBuilder builder = new ColumnVectorBuilder(); + this.children = new ColumnVector[childHolders.size()]; + for (int idx = 0; idx < childHolders.size(); idx++) { + children[idx] = builder.build(childHolders.get(idx), numRows); + } + + this.nullabilityHolder = holder.nullabilityHolder(); + } + + @Override + public void close() { + for (ColumnVector child : children) { + child.close(); + } + } + + @Override + public void closeIfFreeable() { + // See SPARK-50235, SPARK-50463 + } + + @Override + public boolean hasNull() { + return nullabilityHolder != null && nullabilityHolder.hasNulls(); + } + + @Override + public int numNulls() { + return nullabilityHolder == null ? 0 : nullabilityHolder.numNulls(); + } + + @Override + public boolean isNullAt(int rowId) { + return nullabilityHolder != null && nullabilityHolder.isNullAt(rowId) == 1; + } + + @Override + public ColumnVector getChild(int ordinal) { + return children[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("Struct column only supports getChild()"); + } +} diff --git a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java index 5cafa1c9fa04..8f42b938a3f3 100644 --- a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java +++ b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatch.java @@ -22,7 +22,9 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Set; import java.util.function.Supplier; +import java.util.stream.Collectors; import org.apache.iceberg.FileFormat; import org.apache.iceberg.FileScanTask; import org.apache.iceberg.MetadataColumns; @@ -43,6 +45,7 @@ import org.apache.iceberg.spark.SparkReadConf; import org.apache.iceberg.spark.SparkUtil; import org.apache.iceberg.types.Type; +import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; import org.apache.iceberg.util.PropertyUtil; import org.apache.spark.api.java.JavaSparkContext; @@ -65,6 +68,7 @@ class SparkBatch implements Batch { private final boolean executorCacheLocalityEnabled; private final int scanHashCode; private final boolean cacheDeleteFilesOnExecutors; + private final Set projectedVariantIds; SparkBatch( JavaSparkContext sparkContext, @@ -87,6 +91,11 @@ class SparkBatch implements Batch { this.executorCacheLocalityEnabled = readConf.executorCacheLocalityEnabled(); this.scanHashCode = scanHashCode; this.cacheDeleteFilesOnExecutors = readConf.cacheDeleteFilesOnExecutors(); + this.projectedVariantIds = + TypeUtil.indexById(projection.asStruct()).values().stream() + .filter(field -> field.type().isVariantType()) + .map(Types.NestedField::fieldId) + .collect(Collectors.toSet()); } @Override @@ -174,12 +183,8 @@ private boolean supportsParquetBatchReads(ScanTask task) { 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; - } - } + if (lowerBounds != null && hasVariantWithLowerBound(lowerBounds)) { + return false; } return true; @@ -188,6 +193,17 @@ private boolean supportsParquetBatchReads(ScanTask task) { } } + // a variant with stored bounds may be shredded (unsupported vectorized); check all nesting levels + private boolean hasVariantWithLowerBound(Map lowerBounds) { + for (Integer variantId : projectedVariantIds) { + if (lowerBounds.containsKey(variantId)) { + return true; + } + } + + return false; + } + private boolean supportsParquetBatchReads(Types.NestedField field) { if (MetadataColumns.isMetadataColumn(field.fieldId())) { return true; @@ -218,6 +234,16 @@ private boolean supportsParquetBatchReads(Types.NestedField field) { } } + if (type.isStructType()) { + for (Types.NestedField child : type.asStructType().fields()) { + if (!supportsParquetBatchReads(child)) { + return false; + } + } + + return true; + } + return type.isPrimitiveType() || type.isVariantType(); } diff --git a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java index 12047bd12847..1c59e7bff75c 100644 --- a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java +++ b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/data/vectorized/parquet/TestParquetVectorizedReads.java @@ -39,6 +39,7 @@ import java.util.function.Consumer; import java.util.stream.Stream; import org.apache.arrow.memory.BufferAllocator; +import org.apache.hadoop.conf.Configuration; import org.apache.iceberg.Files; import org.apache.iceberg.Schema; import org.apache.iceberg.arrow.ArrowAllocation; @@ -57,7 +58,6 @@ import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap; import org.apache.iceberg.relocated.com.google.common.collect.Iterables; import org.apache.iceberg.relocated.com.google.common.collect.Lists; -import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.apache.iceberg.spark.data.AvroDataTestBase; import org.apache.iceberg.spark.data.GenericsHelpers; import org.apache.iceberg.spark.data.RandomData; @@ -66,11 +66,13 @@ import org.apache.iceberg.types.Type.PrimitiveType; import org.apache.iceberg.types.TypeUtil; import org.apache.iceberg.types.Types; +import org.apache.parquet.column.Encoding; import org.apache.parquet.column.ParquetProperties; +import org.apache.parquet.hadoop.ParquetFileReader; import org.apache.parquet.hadoop.ParquetOutputFormat; -import org.apache.parquet.schema.GroupType; -import org.apache.parquet.schema.MessageType; -import org.apache.parquet.schema.Type; +import org.apache.parquet.hadoop.metadata.BlockMetaData; +import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.util.HadoopInputFile; import org.apache.spark.sql.catalyst.InternalRow; import org.apache.spark.sql.vectorized.ColumnarBatch; import org.junit.jupiter.api.Test; @@ -236,7 +238,7 @@ Iterable generateData( long seed, float nullPercentage, Function transform) { - Iterable data = RandomGenericData.generate(schema, numRecords, seed); + Iterable data = RandomGenericData.generate(schema, numRecords, seed, nullPercentage); return transform == IDENTITY ? data : Iterables.transform(data, transform); } @@ -332,17 +334,162 @@ void assertRecordsMatch( @Test @Override - public void testNestedStruct() { - assertThatThrownBy( - () -> - VectorizedSparkParquetReaders.buildReader( - TypeUtil.assignIncreasingFreshIds( - new Schema(required(1, "struct", SUPPORTED_PRIMITIVES))), - new MessageType( - "struct", new GroupType(Type.Repetition.OPTIONAL, "struct").withId(1)), - Maps.newHashMap())) - .isInstanceOf(UnsupportedOperationException.class) - .hasMessage("Vectorized reads are not supported yet for struct fields"); + public void testNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(required(1, "struct", SUPPORTED_PRIMITIVES)))); + } + + @Test + void nullableNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(optional(1, "struct", SUPPORTED_PRIMITIVES))), + getNumRows(), + 0L, + 0.5f, + true); + } + + @Test + void threeLevelNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds( + new Schema( + optional( + 1, + "outer", + Types.StructType.of( + optional( + 2, + "middle", + Types.StructType.of(required(3, "inner", SUPPORTED_PRIMITIVES))))))), + getNumRows(), + 0L, + 0.5f, + true); + } + + @Test + void dictionaryEncodedStructChild() throws IOException { + // Assert children are dictionary-encoded so this is not silently a plain-encoding test. + Schema schema = + TypeUtil.assignIncreasingFreshIds( + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.IntegerType.get()))))); + int numRows = 1000; + File dataFile = temp.resolve("dict-struct.parquet").toFile(); + Iterable data = + RandomGenericData.generateDictionaryEncodableRecords(schema, numRows, 0L, 0.5f); + try (FileAppender writer = getParquetWriter(schema, dataFile)) { + writer.addAll(data); + } + + try (ParquetFileReader reader = + ParquetFileReader.open( + HadoopInputFile.fromPath( + new org.apache.hadoop.fs.Path(dataFile.toString()), new Configuration()))) { + boolean childUsesDictionary = false; + for (BlockMetaData block : reader.getFooter().getBlocks()) { + for (ColumnChunkMetaData column : block.getColumns()) { + if (column.getPath().toDotString().endsWith("int_data")) { + childUsesDictionary = column.getEncodings().stream().anyMatch(Encoding::usesDictionary); + } + } + } + + assertThat(childUsesDictionary) + .as("struct child int_data should be dictionary-encoded") + .isTrue(); + } + + assertRecordsMatch(schema, numRows, data, dataFile, false, BATCH_SIZE); + } + + @Test + void nestedStructWithReorderedAndDefaultedFields() throws IOException { + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "data", Types.StringType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + optional(3, "data", Types.StringType.get()), + required(2, "id", Types.LongType.get()), + Types.NestedField.optional("added") + .withId(4) + .ofType(Types.IntegerType.get()) + .withInitialDefault(Literal.of(42)) + .build()))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void nestedStructProjectingOnlyAddedField() throws IOException { + // project only a field added after the file was written: the struct must read per-row null + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "data", Types.StringType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + Types.NestedField.optional(4, "added", Types.IntegerType.get())))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void nestedStructWithPromotedChild() throws IOException { + Schema writeSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.IntegerType.get()), + optional(4, "float_data", Types.FloatType.get())))); + Schema expectedSchema = + new Schema( + optional( + 1, + "struct", + Types.StructType.of( + required(2, "id", Types.LongType.get()), + optional(3, "int_data", Types.LongType.get()), + optional(4, "float_data", Types.DoubleType.get())))); + writeAndValidate( + writeSchema, expectedSchema, getNumRows(), 0L, 0.5f, true, BATCH_SIZE, IDENTITY); + } + + @Test + void allNullNestedStruct() throws IOException { + writeAndValidate( + TypeUtil.assignIncreasingFreshIds(new Schema(optional(1, "struct", SUPPORTED_PRIMITIVES))), + getNumRows(), + 0L, + 1.0f, + true); } @Test diff --git a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java index 9cf8e435270c..849c5024055d 100644 --- a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java +++ b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestReadProjection.java @@ -22,6 +22,7 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.assertj.core.api.Assertions.within; +import static org.assertj.core.api.Assumptions.assumeThat; import java.io.IOException; import java.nio.file.Path; @@ -209,6 +210,56 @@ public void testRename() throws Exception { assertThat(cmp).as("Should contain the correct data/renamed value").isEqualTo(0); } + @TestTemplate + void requiredInnerStructUnderPresentParentReadsPresent() throws Exception { + assumeThat(format).isEqualTo(FileFormat.PARQUET); + Schema writeSchema = + new Schema( + Types.NestedField.optional( + 1, + "outer", + Types.StructType.of( + Types.NestedField.required(2, "real", Types.LongType.get()), + Types.NestedField.required( + 3, + "inner", + Types.StructType.of( + Types.NestedField.required(4, "orig", Types.LongType.get())))))); + + Types.StructType outerType = writeSchema.findType("outer").asStructType(); + Record record = GenericRecord.create(writeSchema); + Record outer = GenericRecord.create(outerType); + Record inner = GenericRecord.create(outerType.field("inner").type().asStructType()); + inner.setField("orig", 200L); + outer.setField("real", 100L); + outer.setField("inner", inner); + record.setField("outer", outer); + + // project outer's stored field plus only a field added to the required inner after the write + Schema readSchema = + new Schema( + Types.NestedField.optional( + 1, + "outer", + Types.StructType.of( + Types.NestedField.required(2, "real", Types.LongType.get()), + Types.NestedField.required( + 3, + "inner", + Types.StructType.of( + Types.NestedField.optional(5, "added", Types.IntegerType.get())))))); + + Record projected = writeAndRead("required_inner_added_field", writeSchema, readSchema, record); + Record projectedOuter = (Record) projected.getField("outer"); + assertThat(projectedOuter).as("present parent should project").isNotNull(); + assertThat((long) projectedOuter.getField("real")).isEqualTo(100L); + Record projectedInner = (Record) projectedOuter.getField("inner"); + assertThat(projectedInner) + .as("required inner under a present parent should be present, not null") + .isNotNull(); + assertThat(projectedInner.getField("added")).as("added field should be null").isNull(); + } + @TestTemplate public void testNestedStructProjection() throws Exception { Schema writeSchema = diff --git a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java index 0d619305716e..1b7c3ee2f18e 100644 --- a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java +++ b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkReaderDeletes.java @@ -20,6 +20,7 @@ import static org.apache.hadoop.hive.conf.HiveConf.ConfVars.METASTOREURIS; import static org.apache.iceberg.spark.source.SparkSQLExecutionHelper.lastExecutedMetricValue; +import static org.apache.iceberg.types.Types.NestedField.optional; import static org.apache.iceberg.types.Types.NestedField.required; import static org.apache.spark.sql.types.DataTypes.IntegerType; import static org.assertj.core.api.Assertions.assertThat; @@ -381,6 +382,71 @@ public void testPosDeletesAllRowsInBatch() throws IOException { checkDeleteCount(4L); } + @TestTemplate + void posDeletesOnStructColumn() throws IOException { + Schema structSchema = + new Schema( + required(1, "id", Types.IntegerType.get()), + optional( + 2, + "struct_col", + Types.StructType.of( + required(3, "nested_id", Types.LongType.get()), + optional(4, "nested_data", Types.StringType.get())))); + + String structTableName = "test_struct_pos_deletes"; + dropTable(structTableName); + Table structTable = createTable(structTableName, structSchema, PartitionSpec.unpartitioned()); + + Types.StructType nestedType = structSchema.findField("struct_col").type().asStructType(); + GenericRecord nested = GenericRecord.create(nestedType); + GenericRecord row = GenericRecord.create(structSchema); + List structRecords = Lists.newArrayList(); + for (int idx = 0; idx < 7; idx++) { + structRecords.add( + row.copy( + "id", + idx, + "struct_col", + nested.copy("nested_id", (long) (idx * 10), "nested_data", "v" + idx))); + } + + DataFile structDataFile = + FileHelpers.writeDataFile( + structTable, + Files.localOutput(File.createTempFile("junit", null, temp.toFile())), + structRecords); + structTable.newAppend().appendFile(structDataFile).commit(); + + // delete positions 0 and 5 (ids 0 and 5), spanning the 4-row vectorized batch boundary + List> deletes = + Lists.newArrayList( + Pair.of(structDataFile.location(), 0L), Pair.of(structDataFile.location(), 5L)); + + Pair posDeletes = + FileHelpers.writeDeleteFile( + structTable, + Files.localOutput(File.createTempFile("junit", null, temp.toFile())), + TestHelpers.Row.of(), + deletes, + formatVersion); + + structTable + .newRowDelta() + .addDeletes(posDeletes.first()) + .validateDataFilesExist(posDeletes.second()) + .commit(); + + StructLikeSet expected = rowSetWithoutIds(structTable, structRecords, 0, 5); + StructLikeSet actual = rowSet(structTableName, structTable, "*"); + + assertThat(actual) + .as("Struct column rows should survive positional deletes") + .isEqualTo(expected); + + dropTable(structTableName); + } + @TestTemplate public void testPosDeletesWithDeletedColumn() throws IOException { // read.parquet.vectorization.batch-size is set to 4, so the 4 rows in the first batch are all