initial

amogh-jahagirdar · amogh-jahagirdar · commit 53e1db6dc739 · 2025-04-28T22:49:06.000-06:00
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
@@ -167,6 +167,13 @@ public PositionVectorHolder(
     }
   }
 
+  public static class RowIdVectorHolder extends VectorHolder {
+    public RowIdVectorHolder(
+        FieldVector vector, Types.NestedField icebergField, NullabilityHolder nulls) {
+      super(vector, icebergField, nulls);
+    }
+  }
+
   public static class DeletedVectorHolder extends VectorHolder {
     private final int numRows;
 
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
@@ -461,6 +461,10 @@ public static VectorizedArrowReader positionsWithSetArrowValidityVector() {
     return new PositionVectorReader(true);
   }
 
+  public static VectorizedArrowReader rowIds(long firstRowId, VectorizedArrowReader idReader) {
+    return new RowIdVectorReader(firstRowId, idReader);
+  }
+
   private static final class NullVectorReader extends VectorizedArrowReader {
     private static final NullVectorReader INSTANCE = new NullVectorReader();
 
@@ -567,6 +571,76 @@ public void close() {
     }
   }
 
+  private static final class RowIdVectorReader extends VectorizedArrowReader {
+    private static final Field ROW_ID_ARROW_FIELD = ArrowSchemaUtil.convert(MetadataColumns.ROW_ID);
+
+    private final long firstRowId;
+    private final VectorizedArrowReader idReader;
+    private final VectorizedArrowReader posReader;
+    private NullabilityHolder nulls;
+
+    private RowIdVectorReader(long firstRowId, VectorizedArrowReader idReader) {
+      this.firstRowId = firstRowId;
+      this.idReader = idReader;
+      this.posReader = new PositionVectorReader(true);
+    }
+
+    @Override
+    public VectorHolder read(VectorHolder reuse, int numValsToRead) {
+      FieldVector positions = posReader.read(null, numValsToRead).vector();
+      FieldVector ids = idReader.read(null, numValsToRead).vector();
+      BigIntVector vec = newVector(numValsToRead);
+      ArrowBuf dataBuffer = vec.getDataBuffer();
+      for (int i = 0; i < numValsToRead; i += 1) {
+        if (ids.isNull(i)) {
+          long rowId = firstRowId + (Long) positions.getObject(i);
+          dataBuffer.setLong((long) i * Long.BYTES, rowId);
+        } else {
+          dataBuffer.setLong((long) i * Long.BYTES, (Long) ids.getObject(i));
+        }
+      }
+
+      vec.setValueCount(numValsToRead);
+
+      return new VectorHolder.RowIdVectorHolder(vec, MetadataColumns.ROW_POSITION, nulls);
+    }
+
+    @Override
+    public void setRowGroupInfo(
+        PageReadStore source, Map<ColumnPath, ColumnChunkMetaData> metadata) {
+      idReader.setRowGroupInfo(source, metadata);
+      posReader.setRowGroupInfo(source, metadata);
+    }
+
+    @Override
+    public void setBatchSize(int batchSize) {
+      if (nulls == null || nulls.size() < batchSize) {
+        this.nulls = newNullabilityHolder(batchSize);
+      }
+
+      idReader.setBatchSize(batchSize);
+      posReader.setBatchSize(batchSize);
+    }
+
+    @Override
+    public void close() {
+      // don't close vectors as they are not owned by readers
+    }
+
+    private static BigIntVector newVector(int valueCount) {
+      BigIntVector vector =
+          (BigIntVector) ROW_ID_ARROW_FIELD.createVector(ArrowAllocation.rootAllocator());
+      vector.allocateNew(valueCount);
+      return vector;
+    }
+
+    private static NullabilityHolder newNullabilityHolder(int size) {
+      NullabilityHolder nullabilityHolder = new NullabilityHolder(size);
+      nullabilityHolder.setNotNulls(0, size);
+      return nullabilityHolder;
+    }
+  }
+
   /**
    * A Dummy Vector Reader which doesn't actually read files, instead it returns a dummy
    * VectorHolder which indicates the constant value which should be used for this column.

Original file line number	Diff line number	Diff line change
`@@ -167,6 +167,13 @@ public PositionVectorHolder(`
`167`	`167`	`}`
`168`	`168`	`}`
`169`	`169`
	`170`	`+ public static class RowIdVectorHolder extends VectorHolder {`
	`171`	`+ public RowIdVectorHolder(`
	`172`	`+ FieldVector vector, Types.NestedField icebergField, NullabilityHolder nulls) {`
	`173`	`+ super(vector, icebergField, nulls);`
	`174`	`+ }`
	`175`	`+ }`
	`176`	`+`
`170`	`177`	`public static class DeletedVectorHolder extends VectorHolder {`
`171`	`178`	`private final int numRows;`
`172`	`179`