Fix for delta skip row exception

Signed-off-by: Niranjan Artal <nartal@nvidia.com>

Author: nartal1

delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/DeltaSpark400DB173Provider.scala

@@ -40,7 +40,16 @@ import com.nvidia.spark.rapids.delta.shims.DeltaLogShim
 import com.nvidia.spark.rapids.shims.ShimPredicateHelper
 
 import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, EqualTo, Expression, If, IsNotNull, Literal, Not}
+import org.apache.spark.sql.catalyst.expressions.{
+  And,
+  AttributeReference,
+  EqualTo,
+  Expression,
+  If,
+  IsNotNull,
+  Literal,
+  Not
+}
 import org.apache.spark.sql.catalyst.plans.logical.TableSpec
 import org.apache.spark.sql.connector.write.V1Write
 import org.apache.spark.sql.execution.{FileSourceScanExec, FilterExec, ProjectExec, SparkPlan}
@@ -321,8 +330,20 @@ private object DB173DVPredicatePushdown extends ShimPredicateHelper {
 
     def pruneDeletionVectorSkipRowColumn(plan: SparkPlan): SparkPlan = {
       plan.transformUp {
+        case project @ ProjectExec(projectList, _) =>
+          project.copy(projectList = projectList.filterNot(isDeletionVectorSkipRowColumnRef))
         case project @ GpuProjectExec(projectList, _, _) =>
           project.copy(projectList = projectList.filterNot(isDeletionVectorSkipRowColumnRef))
+        case fsse: FileSourceScanExec =>
+          fsse.copy(
+            output = fsse.output.filterNot(attr => isDeletionVectorSkipRowColumn(attr.name)),
+            requiredSchema = StructType(fsse.requiredSchema.filterNot(field =>
+              isDeletionVectorSkipRowColumn(field.name))),
+            // AQE expects expressions in dataFilters to exist in the output of the scan.
+            // It will not reuse the stage of the scan otherwise. Since we are removing
+            // the deletion-vector skip-row column from scan's output, remove the
+            // corresponding filter from dataFilters as well.
+            dataFilters = fsse.dataFilters.filterNot(isDVCondition))
         case fsse: GpuFileSourceScanExec =>
           fsse.copy(
             originalOutput = fsse.originalOutput.filterNot(attr =>
@@ -341,37 +362,53 @@ private object DB173DVPredicatePushdown extends ShimPredicateHelper {
       // Only native GPU DV scans can replace the skip-row filter. DB DML bitmap-writing
       // plans may still need that filter even when the plugin is enabled.
       plan.exists {
+        case fsse: FileSourceScanExec =>
+          fsse.relation.fileFormat.isInstanceOf[GpuDeltaParquetFileFormatNativeDV]
         case fsse: GpuFileSourceScanExec =>
           fsse.relation.fileFormat.isInstanceOf[GpuDeltaParquetFileFormatNativeDV]
         case _ => false
       }
     }
 
+    def rewriteFilter(
+        condition: Expression,
+        child: SparkPlan,
+        combinePredicates: (Expression, Expression) => Expression,
+        copyFilter: (Expression, SparkPlan) => SparkPlan): Option[SparkPlan] = {
+      val conjuncts = splitConjunctivePredicates(condition)
+      val (dvPredicates, otherPredicates) = conjuncts.partition { predicate =>
+        predicate.references.size == 1 &&
+          predicate.references.exists(ref => isDeletionVectorSkipRowColumn(ref.name)) &&
+          isDVCondition(predicate)
+      }
+      val otherPredicatesReadingSkipRow = otherPredicates.exists { predicate =>
+        predicate.references.exists(ref => isDeletionVectorSkipRowColumn(ref.name))
+      }
+      if (dvPredicates.nonEmpty &&
+          !otherPredicatesReadingSkipRow &&
+          hasNativeDeletionVectorGpuScan(child)) {
+        val newChild = pruneDeletionVectorSkipRowColumn(child)
+        Some(if (otherPredicates.isEmpty) {
+          newChild
+        } else {
+          copyFilter(otherPredicates.reduce(combinePredicates), newChild)
+        })
+      } else {
+        None
+      }
+    }
+
     plan.transformUp {
       case filter @ GpuFilterExec(condition, child)
           if condition.references.exists(ref => isDeletionVectorSkipRowColumn(ref.name)) =>
-        val conjuncts = splitConjunctivePredicates(condition)
-        val (dvPredicates, otherPredicates) = conjuncts.partition { predicate =>
-          predicate.references.size == 1 &&
-            predicate.references.exists(ref => isDeletionVectorSkipRowColumn(ref.name)) &&
-            isDVCondition(predicate)
-        }
-        val otherPredicatesReadingSkipRow = otherPredicates.exists { predicate =>
-          predicate.references.exists(ref => isDeletionVectorSkipRowColumn(ref.name))
-        }
-        if (dvPredicates.nonEmpty &&
-            !otherPredicatesReadingSkipRow &&
-            hasNativeDeletionVectorGpuScan(child)) {
-          val newChild = pruneDeletionVectorSkipRowColumn(child)
-          if (otherPredicates.isEmpty) {
-            newChild
-          } else {
-            filter.copy(condition = otherPredicates.reduce(GpuAnd),
-              child = newChild)(filter.coalesceAfter)
-          }
-        } else {
-          filter
-        }
+        rewriteFilter(condition, child, GpuAnd(_, _),
+          (newCondition, newChild) => filter.copy(condition = newCondition,
+            child = newChild)(filter.coalesceAfter)).getOrElse(filter)
+      case filter @ FilterExec(condition, child)
+          if condition.references.exists(ref => isDeletionVectorSkipRowColumn(ref.name)) =>
+        rewriteFilter(condition, child, And(_, _),
+          (newCondition, newChild) => filter.copy(condition = newCondition, child = newChild))
+          .getOrElse(filter)
     }
   }
 

delta-lake/delta-spark400db173/src/main/scala/com/nvidia/spark/rapids/delta/GpuDeltaParquetFileFormatNativeDV.scala

@@ -26,6 +26,7 @@ import com.databricks.sql.io.{RowIndexFilterProvider, RowIndexFilterType}
 import com.databricks.sql.transaction.tahoe.{
   DeltaColumnMapping,
   DeltaColumnMappingMode,
+  DeltaParquetFileFormat,
   IdMapping,
   NameMapping,
   NoMapping
@@ -1436,14 +1437,74 @@ case class DeltaParquetTableReader(
   override protected lazy val resources: Seq[AutoCloseable] =
     Seq(reader) ++ buffers ++ dvInfos.map(_.serializedBitmap)
 
+  private lazy val deletionVectorSkipRowIndexes =
+    MakeParquetTableWithDVProducer.deletionVectorSkipRowIndexes(readDataSchema)
+
   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)
   }
+
+  override def next: Table = {
+    MakeParquetTableWithDVProducer.materializeDeletionVectorSkipRowColumnsAsFalseIfNeeded(
+      super.next, deletionVectorSkipRowIndexes)
+  }
 }
 
 object MakeParquetTableWithDVProducer extends Logging {
+  private def isDeletionVectorSkipRowColumn(name: String): Boolean =
+    name == DeltaParquetFileFormat.IS_ROW_DELETED_COLUMN_NAME ||
+      name == GpuDeltaParquetFileFormat.EDGE_COMPUTED_COLUMN_SKIP_ROW
+
+  private[delta] def deletionVectorSkipRowIndexes(readDataSchema: StructType): Array[Int] =
+    readDataSchema.fields.zipWithIndex.collect {
+      case (field, index) if isDeletionVectorSkipRowColumn(field.name) => index
+    }
+
+  // Returns the input table unchanged when the planner has pruned skip-row columns.
+  // If replacement is needed, this closes the input table and returns a new one.
+  private[delta] def materializeDeletionVectorSkipRowColumnsAsFalseIfNeeded(
+      table: Table,
+      skipRowIndexes: Array[Int]): Table = {
+    if (skipRowIndexes.nonEmpty) {
+      withResource(table) { tableToClose =>
+        materializeDeletionVectorSkipRowColumnsAsFalse(tableToClose, skipRowIndexes)
+      }
+    } else {
+      table
+    }
+  }
+
+  private[delta] def materializeDeletionVectorSkipRowColumnsAsFalse(
+      table: Table,
+      skipRowIndexes: Array[Int]): Table = {
+    require(skipRowIndexes.forall(_ < table.getNumberOfColumns),
+      s"Expected skip-row indexes ${skipRowIndexes.mkString(",")} within " +
+        s"${table.getNumberOfColumns} output columns")
+    val numRows = Math.toIntExact(table.getRowCount)
+    withResource(Scalar.fromBool(false)) { falseScalar =>
+      val columns = new Array[ColumnVector](table.getNumberOfColumns)
+      val replacementColumns = new ArrayBuffer[ColumnVector](skipRowIndexes.length)
+      try {
+        var i = 0
+        while (i < table.getNumberOfColumns) {
+          columns(i) = if (skipRowIndexes.contains(i)) {
+            val replacement = ColumnVector.fromScalar(falseScalar, numRows)
+            replacementColumns += replacement
+            replacement
+          } else {
+            table.getColumn(i)
+          }
+          i += 1
+        }
+        new Table(columns: _*)
+      } finally {
+        replacementColumns.safeClose()
+      }
+    }
+  }
+
   def apply(
       useChunkedReader: Boolean,
       maxChunkedReaderMemoryUsageSizeBytes: Long,
@@ -1479,6 +1540,7 @@ object MakeParquetTableWithDVProducer extends Logging {
         isSchemaCaseSensitive, useFieldId, readDataSchema, clippedParquetSchema,
         splits, debugDumpPrefix, debugDumpAlways, deletionVectorInfos)
     } else {
+      val skipRowIndexes = deletionVectorSkipRowIndexes(readDataSchema)
       val table = withResource(buffers) { _ =>
         withResource(deletionVectorInfos.map(_.serializedBitmap)) { _ =>
           try {
@@ -1517,7 +1579,8 @@ object MakeParquetTableWithDVProducer extends Logging {
         clippedParquetSchema, readDataSchema, isSchemaCaseSensitive, useFieldId)
       val outputTable = GpuParquetScan.rebaseDateTime(evolvedSchemaTable, dateRebaseMode,
         timestampRebaseMode)
-      new SingleGpuDataProducer(outputTable)
+      new SingleGpuDataProducer(
+        materializeDeletionVectorSkipRowColumnsAsFalseIfNeeded(outputTable, skipRowIndexes))
     }
   }
 }

integration_tests/src/main/python/delta_lake_test.py

@@ -888,6 +888,54 @@ def read_table(spark):
         assert_gpu_and_cpu_are_equal_collect(read_table, conf=conf)
 
 
+@allow_non_gpu("FilterExec", *delta_meta_allow)
+@delta_lake
+@ignore_order(local=True)
+@pytest.mark.skipif(not is_databricks173_or_later(),
+                    reason="This regression is specific to DBR 17.3 native DV reads")
+@pytest.mark.skipif(is_before_spark_353(),
+                    reason="Spark-RAPIDS supports scan with deletion vectors starting in Spark 3.5.3")
+def test_delta_dv_cpu_filter_after_native_scan(spark_tmp_path):
+    data_path = spark_tmp_path + "/DELTA_DATA"
+    conf = {
+        "spark.rapids.sql.delta.deletionVectors.predicatePushdown.enabled": "true",
+        "spark.databricks.delta.delete.deletionVectors.persistent": "true",
+        "spark.databricks.delta.deletionVectors.useMetadataRowIndex": "true",
+        "spark.rapids.sql.expression.In": "false",
+        "spark.rapids.sql.expression.InSet": "false",
+        "spark.rapids.sql.format.parquet.reader.type": "MULTITHREADED",
+        "spark.rapids.sql.reader.chunked": "true"
+    }
+
+    col_a_gen = IntegerGen(min_val=0, max_val=100, nullable=False, special_cases=[])
+    col_b_gen = IntegerGen(min_val=0, max_val=5, nullable=False, special_cases=[0, 1, 2, 3])
+
+    def create_delta(spark):
+        two_col_df(spark, col_a_gen, col_b_gen, length=4000).coalesce(1).write.format("delta") \
+            .option("delta.enableDeletionVectors", "true") \
+            .partitionBy("a").save(data_path)
+
+        count = spark.sql(f"DELETE FROM delta.`{data_path}` WHERE b = 0").collect()[0][0]
+        assert count > 100, "Expected enough rows to be deleted to create deletion vectors"
+
+    def read_table(spark):
+        df = spark.sql(f"SELECT a, b FROM delta.`{data_path}` WHERE b IN (1, 2, 3)")
+        is_gpu = str(spark.conf.get("spark.rapids.sql.enabled", "false")).lower() == "true"
+        if is_gpu:
+            _assert_db173_gpu_delta_scan_if_enabled(spark, df)
+            plan = df._jdf.queryExecution().executedPlan()
+            explain_str = str(plan)
+            callback = spark._sc._jvm.org.apache.spark.sql.rapids.ExecutionPlanCaptureCallback
+            assert callback.contains(plan, "GpuFileGpuScan"), explain_str
+            assert callback.contains(plan, "org.apache.spark.sql.execution.FilterExec"), \
+                explain_str
+            assert "_databricks_internal_edge_computed_column_skip_row" not in explain_str
+        return df
+
+    with_cpu_session(create_delta, conf=conf)
+    assert_gpu_and_cpu_are_equal_collect(read_table, conf=conf)
+
+
 @allow_non_gpu(*delta_meta_allow)
 @delta_lake
 @ignore_order(local=True)