Delta table scan should be optimized when deletion vectors don't exist [databricks]

delta-lake/delta-33x/src/main/scala/com/nvidia/spark/rapids/delta/delta33x/Delta33xProvider.scala

@@ -19,15 +19,16 @@ package com.nvidia.spark.rapids.delta.delta33x
 import com.nvidia.spark.rapids._
 import com.nvidia.spark.rapids.delta.common.DeltaProviderBase
 
+import org.apache.spark.internal.Logging
 import org.apache.spark.sql.connector.catalog.SupportsWrite
-import org.apache.spark.sql.delta.{DeltaDynamicPartitionOverwriteCommand, DeltaParquetFileFormat}
+import org.apache.spark.sql.delta.{DeltaConfigs, DeltaDynamicPartitionOverwriteCommand, DeltaParquetFileFormat}
 import org.apache.spark.sql.delta.catalog.DeltaTableV2
 import org.apache.spark.sql.delta.commands.{DeleteCommand, MergeIntoCommand, OptimizeTableCommand, UpdateCommand}
 import org.apache.spark.sql.execution.command.RunnableCommand
 import org.apache.spark.sql.execution.datasources.FileFormat
 import org.apache.spark.sql.execution.datasources.v2.AppendDataExecV1
 
-object Delta33xProvider extends DeltaProviderBase {
+object Delta33xProvider extends DeltaProviderBase with Logging {
 
   override def isSupportedWrite(write: Class[_ <: SupportsWrite]): Boolean = {
     write == classOf[DeltaTableV2] || write == classOf[GpuDeltaCatalog#GpuStagedDeltaTableV2]
@@ -72,10 +73,22 @@ object Delta33xProvider extends DeltaProviderBase {
     ).map(r => (r.getClassFor.asSubclass(classOf[RunnableCommand]), r)).toMap
   }
 
-  override protected def toGpuParquetFileFormat(fmt: DeltaParquetFileFormat): FileFormat =
+  override protected def toGpuParquetFileFormat(fmt: DeltaParquetFileFormat): FileFormat = {
+    val enableDVConfig = DeltaConfigs.ENABLE_DELETION_VECTORS_CREATION
+    val isDVEnabled = fmt.metadata.configuration.getOrElse(
+      enableDVConfig.key, enableDVConfig.defaultValue).toBoolean
+    val optimizationsEnabled = if (isDVEnabled) {
+      logWarning(s"Input Delta table has deletion vectors enabled. " +
+        s"Optimizations such as file splitting and predicate pushdown are currently not " +
+        s"supported for this table")
+      false
+    } else {
+      fmt.optimizationsEnabled
+    }
     GpuDelta33xParquetFileFormat(fmt.protocol, fmt.metadata, fmt.nullableRowTrackingFields,
-      false, // we don't support splits and predicate pushdown yet
+      optimizationsEnabled,
       fmt.tablePath, fmt.isCDCRead)
+  }
 
   override def convertToGpu(
       cpuExec: AppendDataExecV1,

integration_tests/src/main/python/delta_lake_test.py

@@ -81,6 +81,7 @@ def setup_tables(spark):
     assert_gpu_and_cpu_are_equal_collect(
         lambda spark: spark.sql("SELECT * FROM delta.`{}`".format(data_path)))
 
+
 @allow_non_gpu("FileSourceScanExec", "ColumnarToRowExec", *delta_meta_allow)
 @delta_lake
 @ignore_order(local=True)
@@ -101,6 +102,63 @@ def setup_tables(spark):
         lambda spark: spark.sql("SELECT * FROM delta.`{}`".format(data_path)),
         conf=conf)
 
+
+def do_test_scan_split(spark_tmp_path, enable_deletion_vectors, expected_num_partitions, scan_conf={}):
+    import os
+    import math
+
+    data_path = spark_tmp_path + "/DELTA_DATA"
+    num_rows = 2048
+    def setup_tables(spark):
+        setup_delta_dest_table(spark, data_path,
+                               dest_table_func=lambda spark: unary_op_df(spark, int_gen, length=num_rows, num_slices=1),
+                               use_cdf=False, enable_deletion_vectors=enable_deletion_vectors)
+    target_num_row_groups = 2
+    row_group_size = int(num_rows * 4 / target_num_row_groups) # num_rows * 4 bytes per int / target_num_row_groups
+    conf = {"parquet.block.size": str(row_group_size)}
+    with_cpu_session(setup_tables, conf)
+    # Verify that we have 1 file with 2 row groups
+    def verify_files_and_row_groups():
+        # list files in data_path
+        files = [f for f in os.listdir(data_path) if f.endswith(".parquet")]
+        assert len(files) == 1, "Expected 1 parquet file in the delta table"
+        parquet_file = f"{data_path}/{files[0]}"
+
+        import pyarrow.parquet as pq
+        metadata = pq.read_metadata(parquet_file)
+        assert metadata.num_row_groups == target_num_row_groups, f"Expected {target_num_row_groups} row groups in the parquet"
+        return parquet_file
+    data_file = verify_files_and_row_groups()
+    file_size = os.path.getsize(data_file)
+
+    conf = copy_and_update(scan_conf, {"spark.sql.files.maxPartitionBytes": str(math.ceil(file_size/2.0))})
+
+    def get_num_partitions(spark):
+        df = spark.sql("SELECT * from delta.`{}`".format(data_path))
+        return df.rdd.getNumPartitions()
+    num_partitions = with_gpu_session(get_num_partitions, conf=conf)
+    assert num_partitions == expected_num_partitions, f"Expected {expected_num_partitions} partitions for split read"
+
+
+@allow_non_gpu(*delta_meta_allow)
+@delta_lake
+def test_delta_scan_split(spark_tmp_path):
+    do_test_scan_split(spark_tmp_path, enable_deletion_vectors=False, expected_num_partitions=2)
+
+
+@allow_non_gpu(*delta_meta_allow)
+@delta_lake
+@pytest.mark.skipif(not supports_delta_lake_deletion_vectors(),
+                    reason="Delta Lake deletion vector support is required")
+def test_delta_scan_split_with_deletion_vector_enabled(spark_tmp_path):
+    if is_databricks_runtime():
+        # We support only perfile reader type on Databricks with deletion vectors
+        conf = {"spark.rapids.sql.format.parquet.reader.type": "PERFILE"}
+    else:
+        conf = {}
+    do_test_scan_split(spark_tmp_path, enable_deletion_vectors=True, expected_num_partitions=1, scan_conf=conf)
+
+
 # ID mapping is supported starting in Delta Lake 2.2, but currently cannot distinguish
 # Delta Lake 2.1 from 2.2 in tests. https://github.com/NVIDIA/spark-rapids/issues/9276
 column_mappings = ["name"]

integration_tests/src/main/python/spark_session.py

@@ -83,7 +83,12 @@ def reset_spark_session_conf():
 def _check_for_proper_return_values(something):
     """We don't want to return an DataFrame or Dataset from a with_spark_session. You will not get what you expect"""
     if (isinstance(something, DataFrame)):
-        raise RuntimeError("You should never return a DataFrame from a with_*_session, you will not get the results that you expect")
+        error_msg = """
+        You are trying to return a DataFrame from a with_*_session.
+        This is not allowed because the dataframe is processed lazily,
+        leading to discrepancies between the actual behavior and what you might expect.
+        """
+        raise RuntimeError(error_msg)
 
 @contextmanager
 def pyspark_compatibility_fixes():