Support Iceberg v3 row lineage writes on GPU

iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkPositionDeltaWrite.scala

@@ -252,7 +252,7 @@ class GpuPositionDeltaWriterFactory(
 }
 
 
-trait GpuDeltaWriter extends DeltaWriter[ColumnarBatch] {
+trait GpuIcebergDeltaWriter extends DeltaWriter[ColumnarBatch] with GpuDeltaBatchWriter {
 
   def context: GpuWriteContext
 
@@ -470,7 +470,7 @@ class GpuBasePositionDeltaWriter(
  * Base trait for delta writers that handle both deletes and data writes.
  * This is the GPU equivalent of Java's DeleteAndDataDeltaWriter.
  */
-trait GpuDeleteAndDataDeltaWriter extends GpuDeltaWriter {
+trait GpuDeleteAndDataDeltaWriter extends GpuIcebergDeltaWriter {
   protected val table: Table
   protected val delegate: GpuBasePositionDeltaWriter
   protected val io: FileIO
@@ -492,6 +492,29 @@ trait GpuDeleteAndDataDeltaWriter extends GpuDeltaWriter {
 
   private var closed: Boolean = false
 
+  protected def insertData(row: ColumnarBatch): Unit
+
+  override def insert(row: ColumnarBatch): Unit = reinsert(null, row)
+
+  override def reinsert(metadata: ColumnarBatch, row: ColumnarBatch): Unit = {
+    val physicalRow = withResource(Seq(metadata, row)) { _ =>
+      GpuDataWriterWithRowLineage.appendLineage(
+        row, metadata, context.dataSparkType, context.metadataSparkType)
+    }
+    insertData(physicalRow)
+  }
+
+  override def insertAndReinsert(
+      metadata: ColumnarBatch,
+      row: ColumnarBatch,
+      reinsertMask: CudfColumnVector): Unit = {
+    val physicalRow = withResource(Seq(metadata, row, reinsertMask)) { _ =>
+      GpuDataWriterWithRowLineage.appendLineage(
+        row, metadata, context.dataSparkType, context.metadataSparkType, reinsertMask)
+    }
+    insertData(physicalRow)
+  }
+
   override def delete(metadata: ColumnarBatch, rowId: ColumnarBatch): Unit = {
     require(metadata != null, "Metadata batch must be non null")
 
@@ -569,7 +592,7 @@ class GpuDeleteOnlyDeltaWriter(
     table: Table,
     writerFactory: GpuSparkFileWriterFactory,
     deleteFileFactory: OutputFileFactory,
-    override val context: GpuWriteContext) extends GpuDeltaWriter {
+    override val context: GpuWriteContext) extends GpuIcebergDeltaWriter {
 
   private val io: FileIO = table.io()
   private val specs: mutable.Map[Integer, PartitionSpec] = table.specs().asScala
@@ -642,6 +665,19 @@ class GpuDeleteOnlyDeltaWriter(
     throw new UnsupportedOperationException("Delete-only writer does not support inserts")
   }
 
+  override def insertAndReinsert(
+      metadata: ColumnarBatch,
+      row: ColumnarBatch,
+      reinsertMask: CudfColumnVector): Unit = {
+    withResource(Seq(metadata, row, reinsertMask)) { _ =>
+      throw new UnsupportedOperationException("Delete-only writer does not support inserts")
+    }
+  }
+
+  override def reinsert(metadata: ColumnarBatch, row: ColumnarBatch): Unit = {
+    throw new UnsupportedOperationException("Delete-only writer does not support reinserts")
+  }
+
   override def commit(): WriterCommitMessage = {
     close()
     val result = delegate.result()
@@ -704,7 +740,7 @@ class GpuUnpartitionedDeltaWriter(
     delegate.writeDelete(batch, spec, partition)
   }
 
-  override def insert(row: ColumnarBatch): Unit = {
+  override protected def insertData(row: ColumnarBatch): Unit = {
     val spillBatch = closeOnExcept(row) { _ =>
       SpillableColumnarBatch(row, ACTIVE_ON_DECK_PRIORITY)
     }
@@ -756,7 +792,7 @@ class GpuPartitionedDeltaWriter(
     delegate.writeDelete(batch, spec, partition)
   }
 
-  override def insert(row: ColumnarBatch): Unit = {
+  override protected def insertData(row: ColumnarBatch): Unit = {
     // Partition the data and write each partition
     dataPartitioner.partition(row)
       .safeConsume { part =>

iceberg/common/src/main/scala/org/apache/iceberg/spark/source/GpuSparkWrite.scala

@@ -21,8 +21,9 @@ import java.util.Locale
 import scala.collection.JavaConverters._
 import scala.util.{Failure, Success}
 
+import ai.rapids.cudf.{ColumnVector => CudfColumnVector}
 import com.nvidia.spark.rapids._
-import com.nvidia.spark.rapids.Arm.closeOnExcept
+import com.nvidia.spark.rapids.Arm.{closeOnExcept, withResource}
 import com.nvidia.spark.rapids.RapidsPluginImplicits.AutoCloseableSeq
 import com.nvidia.spark.rapids.SpillPriorities.ACTIVE_ON_DECK_PRIORITY
 import com.nvidia.spark.rapids.fileio.iceberg.IcebergFileIO
@@ -47,8 +48,8 @@ import org.apache.spark.sql.execution.SparkPlan
 import org.apache.spark.sql.execution.datasources.v2.{AtomicCreateTableAsSelectExec, AtomicReplaceTableAsSelectExec}
 import org.apache.spark.sql.rapids.GpuWriteJobStatsTracker
 import org.apache.spark.sql.rapids.shims.SparkSessionUtils
-import org.apache.spark.sql.types.StructType
-import org.apache.spark.sql.vectorized.ColumnarBatch
+import org.apache.spark.sql.types.{LongType, StructType}
+import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector}
 import org.apache.spark.util.SerializableConfiguration
 
 
@@ -64,7 +65,7 @@ class GpuSparkWrite(cpu: Write) extends GpuWrite with RequiresDistributionAndOrd
     // Iceberg's SparkWrite returns different implementations based on write mode:
     // - BatchAppend for append operations
     // - DynamicOverwrite for dynamic partition overwrite
-    // - BatchRewrite for copy-on-write operations (DELETE)
+    // - CopyOnWriteOperation for row-level copy-on-write operations
     // Since these are private classes, we check the class name to determine which GPU version
     // to use
     val cpuBatch = cpu.toBatch
@@ -374,11 +375,18 @@ class GpuWriterFactory(val tableBroadcast: Broadcast[Table],
   val outputWriterFactory: ColumnarOutputWriterFactory,
   val statsTracker: GpuWriteJobStatsTracker,
   val hadoopConf: SerializableConfiguration
-) extends DataWriterFactory {
+) extends GpuDataWriterFactory {
 
   private lazy val fileIO: IcebergFileIO = new IcebergFileIO(tableBroadcast.value.io())
 
   override def createWriter(partitionId: Int, taskId: Long): DataWriter[InternalRow] = {
+    createWriter(partitionId, taskId, null)
+  }
+
+  override def createWriter(
+      partitionId: Int,
+      taskId: Long,
+      metadataSchema: StructType): DataWriter[InternalRow] = {
     val table = tableBroadcast.value
     val spec = table.specs().get(outputSpecId)
     val io = table.io()
@@ -402,23 +410,131 @@ class GpuWriterFactory(val tableBroadcast: Broadcast[Table],
       fileIO)
 
     if (spec.isUnpartitioned) {
-      new GpuUnpartitionedDataWriter(writerFactory, outputFileFactory, io, spec, targetFileSize)
+      new GpuUnpartitionedDataWriter(
+        writerFactory, outputFileFactory, io, spec, targetFileSize, metadataSchema)
         .asInstanceOf[DataWriter[InternalRow]]
     } else {
       new GpuPartitionedDataWriter(writerFactory, outputFileFactory, io, spec, writeSchema,
-        dsSchema, targetFileSize, useFanout)
+        dsSchema, targetFileSize, useFanout, metadataSchema)
         .asInstanceOf[DataWriter[InternalRow]]
     }
   }
 }
 
+trait GpuDataWriterWithRowLineage extends GpuDataWriter {
+  protected def dataSparkType: StructType
+  protected def metadataSchema: StructType
+
+  override def write(record: ColumnarBatch): Unit
+
+  override def write(
+      metadata: ColumnarBatch,
+      record: ColumnarBatch): Unit = {
+    write(GpuDataWriterWithRowLineage.appendLineage(
+      record, metadata, dataSparkType, metadataSchema))
+  }
+}
+
+object GpuDataWriterWithRowLineage {
+  val lineageColumnNames: Seq[String] = Seq("_row_id", "_last_updated_sequence_number")
+
+  /**
+   * Returns an owned physical row batch matching dataSparkType without consuming record, metadata,
+   * or reinsertMask. Records that already contain all physical columns keep their existing values.
+   * Otherwise, appends the missing _row_id and _last_updated_sequence_number columns at the end.
+   *
+   * Without metadata, both appended columns are null. With metadata, lineage columns are looked up
+   * by name in metadataSchema. If reinsertMask is absent, their values are copied for every row.
+   * With a mask, only REINSERT rows (true) copy metadata; INSERT rows (false) receive nulls even
+   * when their metadata contains values. Metadata nulls are preserved for Iceberg's lineage
+   * inheritance mechanism; this method does not assign row IDs or sequence numbers.
+   *
+   * For example, a batch containing a REINSERT followed by an INSERT:
+   * {{{
+   * record:
+   *   id  amount
+   *   1   200
+   *   2   300
+   *
+   * metadata:
+   *   _row_id  _last_updated_sequence_number
+   *   101      null
+   *   999      8
+   *
+   * reinsertMask: [true, false]
+   *
+   * result (columns in dataSparkType order):
+   *   id  amount  _row_id  _last_updated_sequence_number
+   *   1   200     101      null
+   *   2   300     null     null
+   * }}}
+   * The REINSERT preserves row ID 101; the INSERT ignores metadata values 999 and 8 so its lineage
+   * can be inherited. Input row order is unchanged.
+   */
+  def appendLineage(
+      record: ColumnarBatch,
+      metadata: ColumnarBatch,
+      dataSparkType: StructType,
+      metadataSchema: StructType,
+      reinsertMask: CudfColumnVector = null): ColumnarBatch = {
+    if (reinsertMask != null) {
+      require(reinsertMask.getRowCount == record.numRows(),
+        "Reinsert mask row count does not match record row count")
+    }
+    val missingColumnCount = dataSparkType.length - record.numCols()
+    if (missingColumnCount == 0) {
+      GpuColumnVector.combineColumns(record)
+    } else {
+      require(missingColumnCount == lineageColumnNames.length,
+        s"Expected ${lineageColumnNames.length} row lineage " +
+          s"columns but record is missing $missingColumnCount columns")
+      require(dataSparkType.takeRight(missingColumnCount).map(_.name).toSeq == lineageColumnNames,
+        "Expected row lineage columns at the end of the write schema")
+      if (metadata != null) {
+        require(metadata.numRows() == record.numRows(),
+          s"Metadata row count ${metadata.numRows()} does not match record row count " +
+            s"${record.numRows()}")
+      }
+
+      val lineageColumns = closeOnExcept(new Array[ColumnVector](missingColumnCount)) { columns =>
+        lineageColumnNames.zipWithIndex.foreach { case (name, index) =>
+          columns(index) = if (metadata == null) {
+            // Newly inserted rows inherit both lineage values from the Iceberg commit.
+            GpuColumnVector.fromNull(record.numRows(), LongType)
+          } else {
+            val column = metadata.column(metadataSchema.fieldIndex(name))
+              .asInstanceOf[GpuColumnVector]
+            if (reinsertMask == null) {
+              column.incRefCount()
+            } else {
+              // INSERT rows inherit lineage even when their metadata projection has values.
+              withResource(GpuScalar.from(null, LongType)) { nullValue =>
+                GpuColumnVector.from(
+                  reinsertMask.ifElse(column.getBase, nullValue), LongType)
+              }
+            }
+          }
+        }
+        columns
+      }
+
+      withResource(new ColumnarBatch(lineageColumns, record.numRows())) { lineage =>
+        GpuColumnVector.combineColumns(record, lineage)
+      }
+    }
+  }
+}
+
 class GpuUnpartitionedDataWriter(
   val fileWriterFactory: GpuSparkFileWriterFactory,
   val fileFactory: OutputFileFactory,
   val io: FileIO,
   val spec: PartitionSpec,
-  val targetFileSize: Long)
-  extends DataWriter[ColumnarBatch] {
+  val targetFileSize: Long,
+  override protected val metadataSchema: StructType)
+  extends GpuDataWriterWithRowLineage {
+  override protected def dataSparkType: StructType = fileWriterFactory.dataSparkType
+
   private val delegate = new GpuRollingDataWriter(
     fileWriterFactory,
     fileFactory,
@@ -460,10 +576,11 @@ class GpuPartitionedDataWriter(
   val io: FileIO,
   val spec: PartitionSpec,
   val dataSchema: Schema,
-  val dataSparkType: StructType,
+  override val dataSparkType: StructType,
   val targetFileSize: Long,
   val fanoutEnabled: Boolean,
-) extends DataWriter[ColumnarBatch] {
+  override protected val metadataSchema: StructType,
+) extends GpuDataWriterWithRowLineage {
 
   private val delegate: PartitioningWriter[SpillableColumnarBatch, DataWriteResult] =
     if (fanoutEnabled) {