feat: Enable DPP support with native_datafusion scan

Shekharrajak · Shekharrajak · commit 58f60e79737e · 2026-01-10T01:13:18.000+05:30
diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto
@@ -22,6 +22,7 @@ syntax = "proto3";
 package spark.spark_operator;
 
 import "expr.proto";
+import "literal.proto";
 import "partitioning.proto";
 import "types.proto";
 
@@ -108,6 +109,18 @@ message NativeScan {
   // the map.
   map<string, string> object_store_options = 13;
   bool encryption_enabled = 14;
+  repeated RuntimeFilterBound runtime_filter_bounds = 15;
+}
+
+// Runtime filter bound for row-group pruning in native scan.
+// Extracted from Spark's data filters containing range or IN predicates.
+message RuntimeFilterBound {
+  string column_name = 1;
+  int32 column_index = 2;
+  optional spark.spark_expression.Literal min_value = 3;
+  optional spark.spark_expression.Literal max_value = 4;
+  repeated spark.spark_expression.Literal in_values = 5;
+  string filter_type = 6;  // "minmax", "in", "bloom"
 }
 
 message IcebergScan {
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala
@@ -601,6 +601,25 @@ case class CometScanRule(session: SparkSession) extends Rule[SparkPlan] with Com
       partitionSchema: StructType,
       hadoopConf: Configuration): String = {
 
+    val cometExecEnabled = COMET_EXEC_ENABLED.get()
+
+    // Prefer native_datafusion for queries with Dynamic Partition Pruning (DPP)
+    // because it provides significant I/O reduction through partition filtering
+    val hasDPP = scanExec.partitionFilters.exists(isDynamicPruningFilter)
+    if (hasDPP && cometExecEnabled) {
+      val dppFallbackReasons = new ListBuffer[String]()
+      val typeChecker = CometScanTypeChecker(SCAN_NATIVE_DATAFUSION)
+      val schemaSupported =
+        typeChecker.isSchemaSupported(scanExec.requiredSchema, dppFallbackReasons)
+      val partitionSchemaSupported =
+        typeChecker.isSchemaSupported(partitionSchema, dppFallbackReasons)
+
+      if (schemaSupported && partitionSchemaSupported && dppFallbackReasons.isEmpty) {
+        logInfo(s"Auto scan mode selecting $SCAN_NATIVE_DATAFUSION for DPP query")
+        return SCAN_NATIVE_DATAFUSION
+      }
+    }
+
     val fallbackReasons = new ListBuffer[String]()
 
     // native_iceberg_compat only supports local filesystem and S3
@@ -621,7 +640,6 @@ case class CometScanRule(session: SparkSession) extends Rule[SparkPlan] with Com
     val partitionSchemaSupported =
       typeChecker.isSchemaSupported(partitionSchema, fallbackReasons)
 
-    val cometExecEnabled = COMET_EXEC_ENABLED.get()
     if (!cometExecEnabled) {
       fallbackReasons += s"$SCAN_NATIVE_ICEBERG_COMPAT requires ${COMET_EXEC_ENABLED.key}=true"
     }
diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometNativeScan.scala
@@ -58,10 +58,12 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with Logging {
       withInfo(scanExec, s"Full native scan disabled because ${COMET_EXEC_ENABLED.key} disabled")
     }
 
-    // Native DataFusion doesn't support subqueries/dynamic pruning
-    if (scanExec.partitionFilters.exists(isDynamicPruningFilter)) {
-      withInfo(scanExec, "Native DataFusion scan does not support subqueries/dynamic pruning")
-    }
+    // Dynamic partition pruning (DPP) is now supported!
+    // The dynamicallySelectedPartitions in CometScanExec evaluates DPP filters
+    // and returns the filtered file list. Native scan receives these pre-filtered
+    // files, so partition-level pruning works correctly.
+    // Note: DPP filters are excluded from dataFilters to avoid pushing subqueries
+    // to native execution (see supportedDataFilters in CometScanExec).
 
     if (SQLConf.get.ignoreCorruptFiles ||
       scanExec.relation.options
@@ -191,6 +193,10 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with Logging {
         }
       }
 
+      // Add runtime filter bounds if available
+      // These are pushed down from join operators to enable I/O reduction
+      addRuntimeFilterBounds(scan, nativeScanBuilder)
+
       Some(builder.setNativeScan(nativeScanBuilder).build())
 
     } else {
@@ -250,4 +256,70 @@ object CometNativeScan extends CometOperatorSerde[CometScanExec] with Logging {
   override def createExec(nativeOp: Operator, op: CometScanExec): CometNativeExec = {
     CometNativeScanExec(nativeOp, op.wrapped, op.session)
   }
+
+  /**
+   * Add runtime filter bounds to the native scan for row-group pruning. Runtime filters are
+   * extracted from data filters that contain range predicates (GreaterThanOrEqual,
+   * LessThanOrEqual) or IN predicates.
+   */
+  private def addRuntimeFilterBounds(
+      scan: CometScanExec,
+      nativeScanBuilder: OperatorOuterClass.NativeScan.Builder): Unit = {
+    import org.apache.spark.sql.catalyst.expressions._
+
+    // Extract runtime filter bounds from data filters
+    scan.supportedDataFilters.foreach {
+      case GreaterThanOrEqual(attr: AttributeReference, Literal(value, dataType)) =>
+        val boundBuilder = OperatorOuterClass.RuntimeFilterBound.newBuilder()
+        boundBuilder.setColumnName(attr.name)
+        boundBuilder.setColumnIndex(scan.output.indexWhere(_.name == attr.name))
+        boundBuilder.setFilterType("minmax")
+        exprToProto(Literal(value, dataType), scan.output).foreach { minProto =>
+          boundBuilder.setMinValue(minProto.getLiteral)
+        }
+        nativeScanBuilder.addRuntimeFilterBounds(boundBuilder.build())
+
+      case LessThanOrEqual(attr: AttributeReference, Literal(value, dataType)) =>
+        val boundBuilder = OperatorOuterClass.RuntimeFilterBound.newBuilder()
+        boundBuilder.setColumnName(attr.name)
+        boundBuilder.setColumnIndex(scan.output.indexWhere(_.name == attr.name))
+        boundBuilder.setFilterType("minmax")
+        exprToProto(Literal(value, dataType), scan.output).foreach { maxProto =>
+          boundBuilder.setMaxValue(maxProto.getLiteral)
+        }
+        nativeScanBuilder.addRuntimeFilterBounds(boundBuilder.build())
+
+      case And(
+            GreaterThanOrEqual(attr1: AttributeReference, Literal(minVal, minType)),
+            LessThanOrEqual(attr2: AttributeReference, Literal(maxVal, maxType)))
+          if attr1.name == attr2.name =>
+        // Combined range filter: column >= min AND column <= max
+        val boundBuilder = OperatorOuterClass.RuntimeFilterBound.newBuilder()
+        boundBuilder.setColumnName(attr1.name)
+        boundBuilder.setColumnIndex(scan.output.indexWhere(_.name == attr1.name))
+        boundBuilder.setFilterType("minmax")
+        exprToProto(Literal(minVal, minType), scan.output).foreach { minProto =>
+          boundBuilder.setMinValue(minProto.getLiteral)
+        }
+        exprToProto(Literal(maxVal, maxType), scan.output).foreach { maxProto =>
+          boundBuilder.setMaxValue(maxProto.getLiteral)
+        }
+        nativeScanBuilder.addRuntimeFilterBounds(boundBuilder.build())
+
+      case InSet(attr: AttributeReference, values) if values.size <= 10 =>
+        // Small IN filter - pass individual values
+        val boundBuilder = OperatorOuterClass.RuntimeFilterBound.newBuilder()
+        boundBuilder.setColumnName(attr.name)
+        boundBuilder.setColumnIndex(scan.output.indexWhere(_.name == attr.name))
+        boundBuilder.setFilterType("in")
+        values.foreach { value =>
+          exprToProto(Literal(value, attr.dataType), scan.output).foreach { valProto =>
+            boundBuilder.addInValues(valProto.getLiteral)
+          }
+        }
+        nativeScanBuilder.addRuntimeFilterBounds(boundBuilder.build())
+
+      case _ => // Other filters are handled by data_filters
+    }
+  }
 }
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometDPPSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometDPPSuite.scala
@@ -0,0 +1,178 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.exec
+
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.comet.CometNativeScanExec
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
+
+import org.apache.comet.CometConf
+
+/** Tests for Dynamic Partition Pruning (DPP) with native DataFusion scan. */
+class CometDPPSuite extends CometTestBase {
+
+  test("DPP with native_datafusion scan - basic join") {
+    withSQLConf(
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_NATIVE_SCAN_IMPL.key -> "native_datafusion") {
+      withTempPath { dir =>
+        spark
+          .range(10000)
+          .selectExpr("id", "id % 100 as dim_key", "rand() as value")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/fact")
+        spark
+          .range(100)
+          .selectExpr("id", "'name_' || id as name")
+          .where("id < 10")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/dim")
+
+        spark.read.parquet(s"${dir.getCanonicalPath}/fact").createOrReplaceTempView("fact")
+        spark.read.parquet(s"${dir.getCanonicalPath}/dim").createOrReplaceTempView("dim")
+
+        val df = spark.sql("SELECT f.*, d.name FROM fact f JOIN dim d ON f.dim_key = d.id")
+        val result = df.collect()
+
+        assert(result.forall(row => row.getLong(1) < 10))
+
+        val plan = df.queryExecution.executedPlan
+        val hasNativeScan = plan.collect { case _: CometNativeScanExec => true }.nonEmpty ||
+          plan
+            .collect { case a: AdaptiveSparkPlanExec =>
+              a.executedPlan.collect { case _: CometNativeScanExec => true }.nonEmpty
+            }
+            .exists(identity)
+        assert(hasNativeScan, "Expected CometNativeScanExec in plan")
+      }
+    }
+  }
+
+  test("DPP auto-selection for queries with dynamic pruning filters") {
+    withSQLConf(
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_NATIVE_SCAN_IMPL.key -> "auto") {
+      withTempPath { dir =>
+        spark
+          .range(1000)
+          .selectExpr("id", "id % 10 as dim_key")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/fact")
+        spark
+          .range(10)
+          .selectExpr("id", "'name_' || id as name")
+          .where("id < 5")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/dim")
+
+        spark.read.parquet(s"${dir.getCanonicalPath}/fact").createOrReplaceTempView("fact2")
+        spark.read.parquet(s"${dir.getCanonicalPath}/dim").createOrReplaceTempView("dim2")
+
+        val result = spark
+          .sql("SELECT f.*, d.name FROM fact2 f JOIN dim2 d ON f.dim_key = d.id")
+          .collect()
+        assert(result.forall(row => row.getLong(1) < 5))
+      }
+    }
+  }
+
+  test("DPP reduces output rows significantly") {
+    withSQLConf(
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_NATIVE_SCAN_IMPL.key -> "native_datafusion") {
+      withTempPath { dir =>
+        val factRows = 100000
+        val dimRows = 1000
+        val selectivity = 0.01
+
+        spark
+          .range(factRows)
+          .selectExpr("id", s"id % $dimRows as dim_key", "rand() as value")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/fact")
+        spark
+          .range(dimRows)
+          .selectExpr("id", "'name_' || id as name")
+          .where(s"id < ${(dimRows * selectivity).toInt}")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/dim")
+
+        spark.read.parquet(s"${dir.getCanonicalPath}/fact").createOrReplaceTempView("fact3")
+        spark.read.parquet(s"${dir.getCanonicalPath}/dim").createOrReplaceTempView("dim3")
+
+        val count =
+          spark.sql("SELECT f.*, d.name FROM fact3 f JOIN dim3 d ON f.dim_key = d.id").count()
+        val expectedMax = (factRows * selectivity * 2).toLong
+        assert(count <= expectedMax, s"Expected at most $expectedMax rows with DPP, got $count")
+      }
+    }
+  }
+
+  test("DPP with multiple join conditions") {
+    withSQLConf(
+      CometConf.COMET_ENABLED.key -> "true",
+      CometConf.COMET_EXEC_ENABLED.key -> "true",
+      CometConf.COMET_NATIVE_SCAN_IMPL.key -> "native_datafusion") {
+      withTempPath { dir =>
+        spark
+          .range(1000)
+          .selectExpr("id", "id % 10 as key1", "id % 5 as key2", "rand() as value")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/fact")
+        spark
+          .range(10)
+          .selectExpr("id as key1", "'dim1_' || id as name1")
+          .where("id < 3")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/dim1")
+        spark
+          .range(5)
+          .selectExpr("id as key2", "'dim2_' || id as name2")
+          .where("id < 2")
+          .write
+          .mode("overwrite")
+          .parquet(s"${dir.getCanonicalPath}/dim2")
+
+        spark.read.parquet(s"${dir.getCanonicalPath}/fact").createOrReplaceTempView("fact_multi")
+        spark.read.parquet(s"${dir.getCanonicalPath}/dim1").createOrReplaceTempView("dim1")
+        spark.read.parquet(s"${dir.getCanonicalPath}/dim2").createOrReplaceTempView("dim2")
+
+        val result = spark
+          .sql("""
+          SELECT f.*, d1.name1, d2.name2 FROM fact_multi f
+          JOIN dim1 d1 ON f.key1 = d1.key1 JOIN dim2 d2 ON f.key2 = d2.key2
+        """)
+          .collect()
+        assert(result.forall(row => row.getLong(1) < 3 && row.getLong(2) < 2))
+      }
+    }
+  }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDPPBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometDPPBenchmark.scala
