Adding support for DML writes to DV enabled tables

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/DeleteCommandMetaBase.scala

@@ -20,8 +20,7 @@ import com.nvidia.spark.rapids.{DataFromReplacementRule, RapidsConf, RapidsMeta,
 import com.nvidia.spark.rapids.delta.RapidsDeltaUtils
 
 import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.delta.commands.{DeleteCommand, DeletionVectorUtils}
-import org.apache.spark.sql.delta.sources.DeltaSQLConf
+import org.apache.spark.sql.delta.commands.DeleteCommand
 
 abstract class DeleteCommandMetaBase(
     deleteCmd: DeleteCommand,
@@ -35,14 +34,6 @@ abstract class DeleteCommandMetaBase(
       willNotWorkOnGpu("Delta Lake output acceleration has been disabled. To enable set " +
         s"${RapidsConf.ENABLE_DELTA_WRITE} to true")
     }
-    val dvFeatureEnabled = DeletionVectorUtils.deletionVectorsWritable(
-      deleteCmd.deltaLog.unsafeVolatileSnapshot)
-    if (dvFeatureEnabled && deleteCmd.conf.getConf(
-        DeltaSQLConf.DELETE_USE_PERSISTENT_DELETION_VECTORS)) {
-      // https://github.com/NVIDIA/spark-rapids/issues/8554
-      willNotWorkOnGpu("Deletion vectors are not supported on GPU")
-    }
-
     RapidsDeltaUtils.tagForDeltaWrite(this, deleteCmd.target.schema, Some(deleteCmd.deltaLog),
       Map.empty, SparkSession.active)
   }

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/DeltaProviderBase.scala

@@ -29,7 +29,7 @@ import org.apache.spark.sql.delta.{DeltaLog, DeltaParquetFileFormat}
 import org.apache.spark.sql.delta.DeltaParquetFileFormat.IS_ROW_DELETED_COLUMN_NAME
 import org.apache.spark.sql.delta.catalog.DeltaCatalog
 import org.apache.spark.sql.delta.metric.IncrementMetric
-import org.apache.spark.sql.delta.rapids.DeltaRuntimeShim
+import org.apache.spark.sql.delta.rapids.{DeltaRuntimeShim, InputFileDictionaryId}
 import org.apache.spark.sql.delta.sources.DeltaSQLConf
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.datasources.{FileFormat, HadoopFsRelation, SaveIntoDataSourceCommand}
@@ -94,7 +94,8 @@ abstract class DeltaProviderBase extends DeltaIOProvider {
       "IncrementMetric",
       ExprChecks.unaryProject(TypeSig.all, TypeSig.all, TypeSig.all, TypeSig.all),
       (cpuInc, conf, p, r) => GpuIncrementMetricMeta(cpuInc, conf, p, r)
-    )
+    ),
+    InputFileDictionaryId.exprRule
   ).map(r => (r.getClassFor.asSubclass(classOf[Expression]), r)).toMap
 
   override def tagSupportForGpuFileSourceScan(meta: SparkPlanMeta[FileSourceScanExec]): Unit = {

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/GpuDeltaParquetFileFormatBase.scala

@@ -88,8 +88,14 @@ class GpuDeltaParquetFileFormatBase(
    * key to remove from the metadata, which does not exist in earlier versions.
    */
   override def prepareSchema(inputSchema: StructType): StructType = {
-    val schema = DeltaColumnMapping.createPhysicalSchema(
-      inputSchema, referenceSchema, columnMappingMode)
+    val internalColumnNames = Set(IS_ROW_DELETED_COLUMN_NAME, ROW_INDEX_COLUMN_NAME)
+    val dataSchema = StructType(inputSchema.fields.filterNot(
+      field => internalColumnNames.contains(field.name)))
+    val physicalDataFields = DeltaColumnMapping.createPhysicalSchema(
+      dataSchema, referenceSchema, columnMappingMode).fields.iterator
+    val schema = StructType(inputSchema.fields.map { field =>
+      if (internalColumnNames.contains(field.name)) field else physicalDataFields.next()
+    })
     if (columnMappingMode == NameMapping) {
       SchemaMergingUtils.transformColumns(schema) { (_, field, _) =>
         field.copy(metadata = new MetadataBuilder()
@@ -189,7 +195,6 @@ class GpuDeltaParquetFileFormatBase(
 
     // We don't have any additional columns to generate, just return the original reader as is.
     if (isRowDeletedColumn.isEmpty && rowIndexColumn.isEmpty) return dataReader
-    if (isRowDeletedColumn.isEmpty) return dataReader
 
     require(useMetadataRowIndex || !optimizationsEnabled,
       "Cannot generate row index related metadata with file splitting or predicate pushdown")
@@ -239,7 +244,11 @@ class GpuDeltaParquetFileFormatBase(
       // When it is true, combining small files is disabled. Since we don't currently support
       // combining small files with deletion vectors, we need to disable it when deletion vectors
       // exist (which is when tablePath is defined).
-      queryUsesInputFile = hasTablePath || fileScan.queryUsesInputFile)
+      // Explicit row indices must restart at zero for each input file. Treat these scans as
+      // input-file-sensitive so the multi-threaded reader does not combine files into a partition.
+      queryUsesInputFile = hasTablePath ||
+        fileScan.requiredSchema.fieldNames.contains(ROW_INDEX_COLUMN_NAME) ||
+        fileScan.queryUsesInputFile)
   }
 }
 
@@ -315,8 +324,12 @@ class DeltaMultiFileParquetPartitionReader(
 
   override def get(): ColumnarBatch = {
     val batch = reader.get()
-    if (isRowDeletedColumnOpt.isEmpty) {
+    if (isRowDeletedColumnOpt.isEmpty && rowIndexColumnOpt.isEmpty) {
       return batch
+    } else if (file == null && isRowDeletedColumnOpt.isEmpty && files.length == 1) {
+      file = files.head
+      rowIndex = 0
+      rowIndexFilterOpt = None
     } else if (file == null || !compareFile(file)) {
       file = filesMap(InputFileUtils.getCurInputFilePath())
       rowIndex = 0
@@ -450,6 +463,10 @@ object RapidsDeletionVectorUtils {
     batch: ColumnarBatch,
     indexVectorTuples: (Int, org.apache.spark.sql.vectorized.ColumnVector) *): ColumnarBatch = {
     val vectors = ArrayBuffer[org.apache.spark.sql.vectorized.ColumnVector]()
+    val appendedVectors = indexVectorTuples.filter(_._1 >= batch.numCols()).sortBy(_._1)
+    require(appendedVectors.zipWithIndex.forall { case ((index, _), offset) =>
+      index == batch.numCols() + offset
+    }, "Generated metadata columns must be contiguous after the physical batch columns")
     for (i <- 0 until batch.numCols()) {
       var replaced: Boolean = false
       for (indexVectorTuple <- indexVectorTuples) {
@@ -464,6 +481,7 @@ object RapidsDeletionVectorUtils {
         vectors += batch.column(i)
       }
     }
+    appendedVectors.foreach { case (_, vector) => vectors += vector }
     new ColumnarBatch(vectors.toArray, batch.numRows())
   }
 
@@ -536,9 +554,11 @@ object RapidsDeletionVectorUtils {
           indexVectorTuples += (rowIndexCol.index -> rowIndexGpuCol.incRefCount())
         }
         startTime = System.nanoTime()
-        val isRowDeletedVector = rowIndexFilterOpt.get.materializeIntoVector(rowIndexGpuCol)
-        metrics("isRowDeletedColumnGenTime") += System.nanoTime() - startTime
-        indexVectorTuples += (isRowDeletedColumnOpt.get.index -> isRowDeletedVector)
+        isRowDeletedColumnOpt.foreach { isRowDeletedColumn =>
+          val isRowDeletedVector = rowIndexFilterOpt.get.materializeIntoVector(rowIndexGpuCol)
+          metrics("isRowDeletedColumnGenTime") += System.nanoTime() - startTime
+          indexVectorTuples += (isRowDeletedColumn.index -> isRowDeletedVector)
+        }
         replaceVectors(batch, indexVectorTuples.toSeq: _*)
       } catch {
         case e: Throwable =>

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/GpuDeltaParquetFileFormatBase2.scala

@@ -1300,10 +1300,24 @@ case class DeltaParquetTableReader(
   override protected lazy val resources: Seq[AutoCloseable] =
     Seq(reader) ++ buffers ++ dvInfos.map(_.serializedBitmap)
 
+  private val rowIndexColumn = readDataSchema.fieldNames.indexOf(ROW_INDEX_COLUMN_NAME)
+
   override protected def postProcessChunk(chunk: Table): Table = {
-    // The cuDF reader prepends an extra index column in the output table.
-    // We need to drop it before returning as we don't use it.
-    RapidsDeletionVectors.dropFirstColumn(chunk)
+    // Keep the prepended cuDF physical index through schema evolution when Delta requests it.
+    if (rowIndexColumn >= 0) chunk else RapidsDeletionVectors.dropFirstColumn(chunk)
+  }
+
+  override protected def evolveSchemaAndClose(table: Table): Table = {
+    if (rowIndexColumn < 0) {
+      super.evolveSchemaAndClose(table)
+    } else {
+      withResource(table.getColumn(0).castTo(DType.INT64)) { physicalRowIndex =>
+        val dataTable = RapidsDeletionVectors.dropFirstColumn(table)
+        val evolvedTable = super.evolveSchemaAndClose(dataTable)
+        RapidsDeletionVectors.replaceColumnAndClose(
+          evolvedTable, rowIndexColumn, physicalRowIndex)
+      }
+    }
   }
 }
 
@@ -1365,24 +1379,34 @@ object MakeParquetTableWithDVProducer extends Logging {
           }
         }
       }
-      // The cuDF reader prepends an extra index column in the output table.
-      // We need to drop it before returning as we don't use it.
-      val tableWithoutIndex = RapidsDeletionVectors.dropFirstColumn(table)
-      closeOnExcept(tableWithoutIndex) { _ =>
-        GpuParquetScan.throwIfRebaseNeededInExceptionMode(tableWithoutIndex, dateRebaseMode,
-          timestampRebaseMode)
-        if (readDataSchema.length < tableWithoutIndex.getNumberOfColumns) {
-          throw new QueryExecutionException(s"Expected ${readDataSchema.length} columns " +
-            s"but read ${tableWithoutIndex.getNumberOfColumns} from ${splits.mkString("; ")}")
+      // Preserve cuDF physical row indexes only for Delta internal row-index scans.
+      val rowIndexColumn = readDataSchema.fieldNames.indexOf(ROW_INDEX_COLUMN_NAME)
+      val physicalRowIndex = if (rowIndexColumn >= 0) {
+        Some(table.getColumn(0).castTo(DType.INT64))
+      } else {
+        None
+      }
+      withResource(physicalRowIndex) { _ =>
+        val tableWithoutIndex = RapidsDeletionVectors.dropFirstColumn(table)
+        closeOnExcept(tableWithoutIndex) { _ =>
+          GpuParquetScan.throwIfRebaseNeededInExceptionMode(tableWithoutIndex, dateRebaseMode,
+            timestampRebaseMode)
+          if (readDataSchema.length < tableWithoutIndex.getNumberOfColumns) {
+            throw new QueryExecutionException(s"Expected ${readDataSchema.length} columns " +
+              s"but read ${tableWithoutIndex.getNumberOfColumns} from ${splits.mkString("; ")}")
+          }
         }
+        metrics(NUM_OUTPUT_BATCHES) += 1
+        val evolvedSchemaTable = ParquetSchemaUtils.evolveSchemaIfNeededAndClose(tableWithoutIndex,
+          clippedParquetSchema, readDataSchema, isSchemaCaseSensitive, useFieldId)
+        val tableWithRowIndex = physicalRowIndex.map { index =>
+          RapidsDeletionVectors.replaceColumnAndClose(evolvedSchemaTable, rowIndexColumn, index)
+        }.getOrElse(evolvedSchemaTable)
+        val outputTable = GpuParquetScan.rebaseDateTime(tableWithRowIndex, dateRebaseMode,
+          timestampRebaseMode)
+        GpuMetric.recordOutputBatchBytes(outputTable, metrics.get(GPU_OUTPUT_BATCH_BYTES))
+        new SingleGpuDataProducer(outputTable)
       }
-      metrics(NUM_OUTPUT_BATCHES) += 1
-      val evolvedSchemaTable = ParquetSchemaUtils.evolveSchemaIfNeededAndClose(tableWithoutIndex,
-        clippedParquetSchema, readDataSchema, isSchemaCaseSensitive, useFieldId)
-      val outputTable = GpuParquetScan.rebaseDateTime(evolvedSchemaTable, dateRebaseMode,
-        timestampRebaseMode)
-      GpuMetric.recordOutputBatchBytes(outputTable, metrics.get(GPU_OUTPUT_BATCH_BYTES))
-      new SingleGpuDataProducer(outputTable)
     }
   }
 }

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/MergeIntoCommandMetaBase.scala

@@ -21,8 +21,7 @@ import com.nvidia.spark.rapids.delta.RapidsDeltaUtils
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.SparkSession
-import org.apache.spark.sql.delta.commands.{DeletionVectorUtils, MergeIntoCommand}
-import org.apache.spark.sql.delta.sources.DeltaSQLConf
+import org.apache.spark.sql.delta.commands.MergeIntoCommand
 
 abstract class MergeIntoCommandMetaBase(
     mergeCmd: MergeIntoCommand,
@@ -43,15 +42,6 @@ abstract class MergeIntoCommandMetaBase(
       willNotWorkOnGpu("notMatchedBySourceClauses not supported on GPU")
     }
     val deltaLog = mergeCmd.targetFileIndex.deltaLog
-    val dvFeatureEnabled =
-      DeletionVectorUtils.deletionVectorsWritable(deltaLog.unsafeVolatileSnapshot)
-
-    if (dvFeatureEnabled && mergeCmd.conf.getConf(
-      DeltaSQLConf.MERGE_USE_PERSISTENT_DELETION_VECTORS)) {
-      // https://github.com/NVIDIA/spark-rapids/issues/8654
-      willNotWorkOnGpu("Deletion vectors are not supported on GPU")
-    }
-
     val targetSchema = mergeCmd.migratedSchema.getOrElse(mergeCmd.target.schema)
     RapidsDeltaUtils.tagForDeltaWrite(this, targetSchema, Some(deltaLog), Map.empty,
       SparkSession.active)

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/RapidsDeletionVectors.scala

@@ -237,6 +237,23 @@ object RapidsDeletionVectors extends Logging {
     }
   }
 
+  /**
+   * Replaces one column in a table and consumes the input table. The returned table owns
+   * references to the replacement and all unchanged columns.
+   */
+  def replaceColumnAndClose(
+      table: Table,
+      outputColumn: Int,
+      replacement: ColumnVector): Table = {
+    require(outputColumn >= 0 && outputColumn < table.getNumberOfColumns,
+      "Invalid replacement column position")
+    withResource(table) { input =>
+      val outputColumns = (0 until input.getNumberOfColumns).map(input.getColumn).toArray
+      outputColumns(outputColumn) = replacement
+      new Table(outputColumns: _*)
+    }
+  }
+
   def isIfNotContainedRowIndexFilter(filterTypeOpt: Option[RowIndexFilterType]): Boolean = {
     filterTypeOpt.contains(RowIndexFilterType.IF_NOT_CONTAINED)
   }

delta-lake/common/src/main/delta-33x-42x/scala/com/nvidia/spark/rapids/delta/common/UpdateCommandMetaBase.scala

@@ -19,8 +19,7 @@ package com.nvidia.spark.rapids.delta.common
 import com.nvidia.spark.rapids.{DataFromReplacementRule, RapidsConf, RapidsMeta, RunnableCommandMeta}
 import com.nvidia.spark.rapids.delta.RapidsDeltaUtils
 
-import org.apache.spark.sql.delta.commands.{DeletionVectorUtils, UpdateCommand}
-import org.apache.spark.sql.delta.sources.DeltaSQLConf
+import org.apache.spark.sql.delta.commands.UpdateCommand
 
 abstract class UpdateCommandMetaBase(
     updateCmd: UpdateCommand,
@@ -35,15 +34,6 @@ abstract class UpdateCommandMetaBase(
         s"${RapidsConf.ENABLE_DELTA_WRITE} to true")
     }
 
-    val dvFeatureEnabled = DeletionVectorUtils.deletionVectorsWritable(
-      updateCmd.tahoeFileIndex.deltaLog.unsafeVolatileSnapshot)
-
-    if (dvFeatureEnabled && updateCmd.conf.getConf(
-      DeltaSQLConf.DELETE_USE_PERSISTENT_DELETION_VECTORS)) {
-      // https://github.com/NVIDIA/spark-rapids/issues/8554
-      willNotWorkOnGpu("Deletion vectors are not supported on GPU")
-    }
-
     RapidsDeltaUtils.tagForDeltaWrite(this, updateCmd.target.schema,
       Some(updateCmd.tahoeFileIndex.deltaLog),
       Map.empty, updateCmd.tahoeFileIndex.spark)

delta-lake/common/src/main/delta-33x-42x/scala/org/apache/spark/sql/delta/rapids/DeltaCommandShims.scala

@@ -19,6 +19,8 @@ package org.apache.spark.sql.delta.rapids
 import org.apache.spark.sql.{Column, DataFrame, SparkSession}
 import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.delta.actions.FileAction
+import org.apache.spark.sql.delta.commands.TouchedFileWithDV
 
 /**
  * Trait to abstract version-specific Spark API differences between Delta 3.3.x and Spark 4.x
@@ -82,6 +84,14 @@ trait DeltaCommandShims {
    */
   def exprToColumn(expr: Expression): Column
 
+  /**
+   * Apply version-specific Delta statistics handling to new deletion-vector actions.
+   */
+  def processUnmodifiedData(
+      spark: OperationSparkSession,
+      touchedFiles: Seq[TouchedFileWithDV],
+      txn: GpuOptimisticTransactionBase): (Seq[FileAction], Map[String, Long])
+
   /**
    * Recache by plan with the correct SparkSession type.
    */