Spark 4.1: Fix row lineage in vectorized ORC reads

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();
+  }
+}

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();
+  }
+}

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<String> names,
         List<Converter> 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<Integer, Integer> fieldIdToOrcIndex;
     private final List<Converter> fieldConverters;
     private final Map<Integer, ?> idToConstant;
 
     private StructConverter(
         Types.StructType structType,
+        TypeDescription orcType,
         List<Converter> fieldConverters,
         Map<Integer, ?> 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<ColumnVector> 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<Integer, Integer> buildFieldIdToOrcIndex(TypeDescription orcType) {
+      List<TypeDescription> children = orcType.getChildren();
+      Map<Integer, Integer> mapping = Maps.newHashMap();
+      for (int i = 0; i < children.size(); i++) {
+        mapping.put(ORCSchemaUtil.fieldId(children.get(i)), i);
+      }
+
+      return mapping;
+    }
   }
 }