diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/LastUpdatedSeqColumnVector.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/LastUpdatedSeqColumnVector.java new file mode 100644 index 000000000000..da65c0871489 --- /dev/null +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/LastUpdatedSeqColumnVector.java @@ -0,0 +1,126 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.spark.data.vectorized; + +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.types.Types; +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 LastUpdatedSeqColumnVector extends ColumnVector { + + private final long fileLastUpdated; + private final ColumnVector storedSeqs; + + LastUpdatedSeqColumnVector(long fileLastUpdated, ColumnVector storedSeqs) { + super(SparkSchemaUtil.convert(Types.LongType.get())); + this.fileLastUpdated = fileLastUpdated; + this.storedSeqs = storedSeqs; + } + + @Override + public long getLong(int rowId) { + if (storedSeqs != null && !storedSeqs.isNullAt(rowId)) { + return storedSeqs.getLong(rowId); + } + + return fileLastUpdated; + } + + @Override + public void close() {} + + @Override + public boolean hasNull() { + return false; + } + + @Override + public int numNulls() { + return 0; + } + + @Override + public boolean isNullAt(int rowId) { + return false; + } + + @Override + public boolean getBoolean(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public byte getByte(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public short getShort(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public int getInt(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public float getFloat(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public double getDouble(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public ColumnarArray getArray(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public ColumnarMap getMap(int ordinal) { + throw new UnsupportedOperationException(); + } + + @Override + public Decimal getDecimal(int rowId, int precision, int scale) { + throw new UnsupportedOperationException(); + } + + @Override + public UTF8String getUTF8String(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public byte[] getBinary(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public ColumnVector getChild(int ordinal) { + throw new UnsupportedOperationException(); + } +} diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/RowIdColumnVector.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/RowIdColumnVector.java new file mode 100644 index 000000000000..01a0b31e1da6 --- /dev/null +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/RowIdColumnVector.java @@ -0,0 +1,128 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.spark.data.vectorized; + +import org.apache.iceberg.spark.SparkSchemaUtil; +import org.apache.iceberg.types.Types; +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 RowIdColumnVector extends ColumnVector { + + private final long firstRowId; + private final long batchOffsetInFile; + private final ColumnVector storedIds; + + RowIdColumnVector(long firstRowId, long batchOffsetInFile, ColumnVector storedIds) { + super(SparkSchemaUtil.convert(Types.LongType.get())); + this.firstRowId = firstRowId; + this.batchOffsetInFile = batchOffsetInFile; + this.storedIds = storedIds; + } + + @Override + public long getLong(int rowId) { + if (storedIds != null && !storedIds.isNullAt(rowId)) { + return storedIds.getLong(rowId); + } + + return firstRowId + batchOffsetInFile + rowId; + } + + @Override + public void close() {} + + @Override + public boolean hasNull() { + return false; + } + + @Override + public int numNulls() { + return 0; + } + + @Override + public boolean isNullAt(int rowId) { + return false; + } + + @Override + public boolean getBoolean(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public byte getByte(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public short getShort(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public int getInt(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public float getFloat(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public double getDouble(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public ColumnarArray getArray(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public ColumnarMap getMap(int ordinal) { + throw new UnsupportedOperationException(); + } + + @Override + public Decimal getDecimal(int rowId, int precision, int scale) { + throw new UnsupportedOperationException(); + } + + @Override + public UTF8String getUTF8String(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public byte[] getBinary(int rowId) { + throw new UnsupportedOperationException(); + } + + @Override + public ColumnVector getChild(int ordinal) { + throw new UnsupportedOperationException(); + } +} diff --git a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkOrcReaders.java b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkOrcReaders.java index 4f324239881e..4286573b18a9 100644 --- a/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkOrcReaders.java +++ b/spark/v4.1/spark/src/main/java/org/apache/iceberg/spark/data/vectorized/VectorizedSparkOrcReaders.java @@ -23,11 +23,14 @@ import java.util.stream.IntStream; import org.apache.iceberg.MetadataColumns; import org.apache.iceberg.Schema; +import org.apache.iceberg.orc.ORCSchemaUtil; import org.apache.iceberg.orc.OrcBatchReader; import org.apache.iceberg.orc.OrcSchemaWithTypeVisitor; import org.apache.iceberg.orc.OrcValueReader; import org.apache.iceberg.orc.OrcValueReaders; +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.Maps; import org.apache.iceberg.spark.SparkSchemaUtil; import org.apache.iceberg.spark.data.SparkOrcValueReaders; import org.apache.iceberg.types.Type; @@ -104,7 +107,7 @@ public Converter record( TypeDescription record, List names, List fields) { - return new StructConverter(iStruct, fields, idToConstant); + return new StructConverter(iStruct, record, fields, idToConstant); } @Override @@ -432,16 +435,85 @@ public ColumnarMap getMap(int rowId) { private static class StructConverter implements Converter { private final Types.StructType structType; + private final Map fieldIdToOrcIndex; private final List fieldConverters; private final Map idToConstant; private StructConverter( Types.StructType structType, + TypeDescription orcType, List fieldConverters, Map idToConstant) { this.structType = structType; + this.fieldIdToOrcIndex = buildFieldIdToOrcIndex(orcType); this.fieldConverters = fieldConverters; this.idToConstant = idToConstant; + + Preconditions.checkState( + fieldConverters.size() <= orcType.getChildren().size(), + "Invalid ORC schema: %s converters for %s columns", + fieldConverters.size(), + orcType.getChildren().size()); + } + + private ColumnVector storedVector( + Types.NestedField field, + StructColumnVector structVector, + int batchSize, + long batchOffsetInFile, + boolean isSelectedInUse, + int[] selected) { + Integer orcIndex = fieldIdToOrcIndex.get(field.fieldId()); + if (orcIndex == null) { + return null; + } + + return fieldConverters + .get(orcIndex) + .convert( + structVector.fields[orcIndex], + batchSize, + batchOffsetInFile, + isSelectedInUse, + selected); + } + + private ColumnVector rowIdVector( + Types.NestedField field, + StructColumnVector structVector, + int batchSize, + long batchOffsetInFile, + boolean isSelectedInUse, + int[] selected) { + Long firstRowId = (Long) idToConstant.get(field.fieldId()); + if (firstRowId == null) { + return new ConstantColumnVector(field.type(), batchSize, null); + } + + return new RowIdColumnVector( + firstRowId, + batchOffsetInFile, + storedVector( + field, structVector, batchSize, batchOffsetInFile, isSelectedInUse, selected)); + } + + private ColumnVector lastUpdatedSeqVector( + Types.NestedField field, + StructColumnVector structVector, + int batchSize, + long batchOffsetInFile, + boolean isSelectedInUse, + int[] selected) { + Long fileLastUpdated = (Long) idToConstant.get(field.fieldId()); + Long firstRowId = (Long) idToConstant.get(MetadataColumns.ROW_ID.fieldId()); + if (fileLastUpdated == null || firstRowId == null) { + return new ConstantColumnVector(field.type(), batchSize, null); + } + + return new LastUpdatedSeqColumnVector( + fileLastUpdated, + storedVector( + field, structVector, batchSize, batchOffsetInFile, isSelectedInUse, selected)); } @Override @@ -456,7 +528,15 @@ public ColumnVector convert( List fieldVectors = Lists.newArrayListWithExpectedSize(fields.size()); for (int pos = 0, vectorIndex = 0; pos < fields.size(); pos += 1) { Types.NestedField field = fields.get(pos); - if (idToConstant.containsKey(field.fieldId())) { + if (field.equals(MetadataColumns.ROW_ID)) { + fieldVectors.add( + rowIdVector( + field, structVector, batchSize, batchOffsetInFile, isSelectedInUse, selected)); + } else if (field.equals(MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER)) { + fieldVectors.add( + lastUpdatedSeqVector( + field, structVector, batchSize, batchOffsetInFile, isSelectedInUse, selected)); + } else if (idToConstant.containsKey(field.fieldId())) { fieldVectors.add( new ConstantColumnVector(field.type(), batchSize, idToConstant.get(field.fieldId()))); } else if (field.equals(MetadataColumns.ROW_POSITION)) { @@ -488,5 +568,15 @@ public ColumnVector getChild(int ordinal) { } }; } + + private static Map buildFieldIdToOrcIndex(TypeDescription orcType) { + List children = orcType.getChildren(); + Map mapping = Maps.newHashMap(); + for (int i = 0; i < children.size(); i++) { + mapping.put(ORCSchemaUtil.fieldId(children.get(i)), i); + } + + return mapping; + } } } diff --git a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestVectorizedOrcDataReader.java b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestVectorizedOrcDataReader.java index 891cb9481ce6..147c2ead9783 100644 --- a/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestVectorizedOrcDataReader.java +++ b/spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/data/TestVectorizedOrcDataReader.java @@ -26,7 +26,9 @@ import java.util.Arrays; import java.util.Collections; import java.util.Iterator; +import java.util.Map; import org.apache.iceberg.Files; +import org.apache.iceberg.MetadataColumns; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.Schema; import org.apache.iceberg.data.GenericRecord; @@ -115,6 +117,12 @@ private void validateAllRows(Iterator rows) { assertThat(rowCount).isEqualTo(5); } + private void assertRow(InternalRow row, long expectedId, long expectedRowId, long expectedSeq) { + assertThat(row.getLong(0)).isEqualTo(expectedId); + assertThat(row.getLong(1)).isEqualTo(expectedRowId); + assertThat(row.getLong(2)).isEqualTo(expectedSeq); + } + @Test public void testReader() throws IOException { try (CloseableIterable reader = @@ -164,4 +172,66 @@ public void testWithFilterWithSelected() throws IOException { assertThat(rows).isExhausted(); } } + + @Test + public void testRowLineage() throws IOException { + Schema dataSchema = new Schema(Types.NestedField.required(1, "id", Types.LongType.get())); + Schema writeSchema = MetadataColumns.schemaWithRowLineage(dataSchema); + + GenericRecord record = GenericRecord.create(writeSchema); + ImmutableList.Builder builder = ImmutableList.builder(); + for (int i = 0; i < 4; i++) { + builder.add( + i % 2 == 0 + ? record.copy( + ImmutableMap.of( + "id", + (long) i, + MetadataColumns.ROW_ID.name(), + 555L + i, + MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.name(), + 7L)) + : record.copy(ImmutableMap.of("id", (long) i))); + } + + OutputFile lineageFile = + Files.localOutput(File.createTempFile("lineage", ".orc", temp.toFile())); + try (DataWriter dataWriter = + ORC.writeData(lineageFile) + .schema(writeSchema) + .createWriterFunc(GenericOrcWriter::buildWriter) + .overwrite() + .withSpec(PartitionSpec.unpartitioned()) + .build()) { + for (Record rec : builder.build()) { + dataWriter.write(rec); + } + } + + Schema projection = + new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + MetadataColumns.ROW_ID, + MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER); + Map idToConstant = + ImmutableMap.of( + MetadataColumns.ROW_ID.fieldId(), 100L, + MetadataColumns.LAST_UPDATED_SEQUENCE_NUMBER.fieldId(), 5L); + + try (CloseableIterable reader = + ORC.read(lineageFile.toInputFile()) + .project(projection) + .createBatchedReaderFunc( + readOrcSchema -> + VectorizedSparkOrcReaders.buildReader(projection, readOrcSchema, idToConstant)) + .build()) { + Iterator rows = batchesToRows(reader.iterator()); + + assertRow(rows.next(), 0L, 555L, 7L); + assertRow(rows.next(), 1L, 101L, 5L); + assertRow(rows.next(), 2L, 557L, 7L); + assertRow(rows.next(), 3L, 103L, 5L); + assertThat(rows).isExhausted(); + } + } }