feat: AQE DPP for native Parquet scans with broadcast reuse (#4112)

Author: mbutrovich

dev/diffs/3.4.3.diff

@@ -417,7 +417,7 @@ index daef11ae4d6..9f3cc9181f2 100644
      assert(exchanges.size == 2)
    }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
-index f33432ddb6f..914afa6b01d 100644
+index f33432ddb6f..b375e285dde 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
 @@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen
@@ -477,35 +477,15 @@ index f33432ddb6f..914afa6b01d 100644
 
        assert(countSubqueryBroadcasts == 1)
        assert(countReusedSubqueryBroadcasts == 1)
-@@ -1215,7 +1231,8 @@ abstract class DynamicPartitionPruningSuiteBase
-   }
- 
-   test("SPARK-32509: Unused Dynamic Pruning filter shouldn't affect " +
--    "canonicalization and exchange reuse") {
-+    "canonicalization and exchange reuse",
-+    IgnoreComet("TODO: https://github.com/apache/datafusion-comet/issues/4045")) {
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
-         val df = sql(
-@@ -1423,7 +1440,8 @@ abstract class DynamicPartitionPruningSuiteBase
-     }
-   }
- 
--  test("SPARK-34637: DPP side broadcast query stage is created firstly") {
-+  test("SPARK-34637: DPP side broadcast query stage is created firstly",
-+    IgnoreComet("TODO: https://github.com/apache/datafusion-comet/issues/4045")) {
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       val df = sql(
-         """ WITH v as (
-@@ -1577,6 +1595,7 @@ abstract class DynamicPartitionPruningSuiteBase
+@@ -1577,6 +1593,7 @@ abstract class DynamicPartitionPruningSuiteBase
 
          val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) {
            case s: SubqueryBroadcastExec => s
 +          case s: CometSubqueryBroadcastExec => s
          }
          assert(subqueryBroadcastExecs.size === 1)
          subqueryBroadcastExecs.foreach { subqueryBroadcastExec =>
-@@ -1729,6 +1748,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
+@@ -1729,6 +1746,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
                case s: BatchScanExec =>
                  // we use f1 col for v2 tables due to schema pruning
                  s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1")))

dev/diffs/3.5.8.diff

@@ -398,7 +398,7 @@ index c4fb4fa943c..a04b23870a8 100644
      assert(exchanges.size == 2)
    }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
-index f33432ddb6f..914afa6b01d 100644
+index f33432ddb6f..b375e285dde 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
 @@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen
@@ -458,35 +458,15 @@ index f33432ddb6f..914afa6b01d 100644
 
        assert(countSubqueryBroadcasts == 1)
        assert(countReusedSubqueryBroadcasts == 1)
-@@ -1215,7 +1231,8 @@ abstract class DynamicPartitionPruningSuiteBase
-   }
- 
-   test("SPARK-32509: Unused Dynamic Pruning filter shouldn't affect " +
--    "canonicalization and exchange reuse") {
-+    "canonicalization and exchange reuse",
-+    IgnoreComet("TODO: https://github.com/apache/datafusion-comet/issues/4045")) {
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
-         val df = sql(
-@@ -1423,7 +1440,8 @@ abstract class DynamicPartitionPruningSuiteBase
-     }
-   }
- 
--  test("SPARK-34637: DPP side broadcast query stage is created firstly") {
-+  test("SPARK-34637: DPP side broadcast query stage is created firstly",
-+    IgnoreComet("TODO: https://github.com/apache/datafusion-comet/issues/4045")) {
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       val df = sql(
-         """ WITH v as (
-@@ -1577,6 +1595,7 @@ abstract class DynamicPartitionPruningSuiteBase
+@@ -1577,6 +1593,7 @@ abstract class DynamicPartitionPruningSuiteBase
 
          val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) {
            case s: SubqueryBroadcastExec => s
 +          case s: CometSubqueryBroadcastExec => s
          }
          assert(subqueryBroadcastExecs.size === 1)
          subqueryBroadcastExecs.foreach { subqueryBroadcastExec =>
-@@ -1729,6 +1748,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
+@@ -1729,6 +1746,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
                case s: BatchScanExec =>
                  // we use f1 col for v2 tables due to schema pruning
                  s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1")))

dev/diffs/4.0.2.diff

@@ -535,7 +535,7 @@ index 81713c777bc..b5f92ed9742 100644
      assert(exchanges.size == 2)
    }
 diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
-index 2c24cc7d570..8c214e7d05c 100644
+index 2c24cc7d570..12d897866da 100644
 --- a/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
 +++ b/sql/core/src/test/scala/org/apache/spark/sql/DynamicPartitionPruningSuite.scala
 @@ -22,6 +22,7 @@ import org.scalatest.GivenWhenThen
@@ -595,35 +595,15 @@ index 2c24cc7d570..8c214e7d05c 100644
 
        assert(countSubqueryBroadcasts == 1)
        assert(countReusedSubqueryBroadcasts == 1)
-@@ -1215,7 +1231,8 @@ abstract class DynamicPartitionPruningSuiteBase
-   }
- 
-   test("SPARK-32509: Unused Dynamic Pruning filter shouldn't affect " +
--    "canonicalization and exchange reuse") {
-+    "canonicalization and exchange reuse",
-+    IgnoreComet("TODO: https://github.com/apache/datafusion-comet/issues/4045")) {
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       withSQLConf(SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
-         val df = sql(
-@@ -1424,7 +1441,8 @@ abstract class DynamicPartitionPruningSuiteBase
-     }
-   }
- 
--  test("SPARK-34637: DPP side broadcast query stage is created firstly") {
-+  test("SPARK-34637: DPP side broadcast query stage is created firstly",
-+    IgnoreComet("TODO: https://github.com/apache/datafusion-comet/issues/4045")) {
-     withSQLConf(SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true") {
-       val df = sql(
-         """ WITH v as (
-@@ -1578,6 +1596,7 @@ abstract class DynamicPartitionPruningSuiteBase
+@@ -1578,6 +1594,7 @@ abstract class DynamicPartitionPruningSuiteBase
 
          val subqueryBroadcastExecs = collectWithSubqueries(df.queryExecution.executedPlan) {
            case s: SubqueryBroadcastExec => s
 +          case s: CometSubqueryBroadcastExec => s
          }
          assert(subqueryBroadcastExecs.size === 1)
          subqueryBroadcastExecs.foreach { subqueryBroadcastExec =>
-@@ -1730,6 +1749,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
+@@ -1730,6 +1747,10 @@ abstract class DynamicPartitionPruningV1Suite extends DynamicPartitionPruningDat
                case s: BatchScanExec =>
                  // we use f1 col for v2 tables due to schema pruning
                  s.output.exists(_.exists(_.argString(maxFields = 100).contains("f1")))

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -32,7 +32,7 @@ import org.apache.spark.sql.execution._
 import org.apache.spark.sql.internal.SQLConf
 
 import org.apache.comet.CometConf._
-import org.apache.comet.rules.{CometExecRule, CometReuseSubquery, CometScanRule, EliminateRedundantTransitions}
+import org.apache.comet.rules.{CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions}
 import org.apache.comet.shims.ShimCometSparkSessionExtensions
 
 /**
@@ -43,34 +43,44 @@ import org.apache.comet.shims.ShimCometSparkSessionExtensions
  *
  * Non-AQE (QueryExecution.preparations):
  * {{{
- *   1. PlanDynamicPruningFilters    -- Spark creates DPP filters
+ *   1. PlanDynamicPruningFilters    -- Spark creates non-AQE DPP (SubqueryBroadcastExec)
  *   2. PlanSubqueries               -- Spark creates SubqueryExec for scalar subqueries
  *   3. EnsureRequirements            -- Spark inserts shuffles/sorts
  *   4. ApplyColumnarRulesAndInsertTransitions:
- *      a. preColumnarTransitions:   CometScanRule, CometExecRule (replace Spark -> Comet nodes)
+ *      a. preColumnarTransitions:   CometScanRule, CometExecRule
+ *         - CometExecRule.convertSubqueryBroadcasts converts SubqueryBroadcastExec to
+ *           CometSubqueryBroadcastExec for exchange reuse with Comet broadcasts
  *      b. insertTransitions:        ColumnarToRow/RowToColumnar added
  *      c. postColumnarTransitions:  EliminateRedundantTransitions
  *   5. ReuseExchangeAndSubquery     -- Spark deduplicates subqueries (sees Comet nodes)
  * }}}
  *
- * AQE (AdaptiveSparkPlanExec):
+ * AQE (AdaptiveSparkPlanExec, Spark 3.5+):
  * {{{
  *   Initial plan:
- *     queryStagePreparationRules:   CometScanRule, CometExecRule (replace Spark -> Comet nodes)
+ *     PlanAdaptiveSubqueries:       creates SubqueryAdaptiveBroadcastExec (SAB) for AQE DPP
+ *     queryStagePreparationRules:   CometScanRule, CometExecRule
+ *       - CometExecRule.convertSubqueryBroadcasts wraps SABs in
+ *         CometSubqueryAdaptiveBroadcastExec to prevent Spark's
+ *         PlanAdaptiveDynamicPruningFilters from replacing DPP with Literal.TrueLiteral
  *
  *   Per stage (optimizeQueryStage + postStageCreationRules):
- *     1. queryStageOptimizerRules:  ReuseAdaptiveSubquery, CometReuseSubquery
+ *     1. queryStageOptimizerRules:
+ *        a. PlanAdaptiveDynamicPruningFilters (Spark) -- skips wrapped SABs
+ *        b. ReuseAdaptiveSubquery (Spark)
+ *        c. CometPlanAdaptiveDynamicPruningFilters   -- converts wrapped SABs to
+ *           CometSubqueryBroadcastExec with BroadcastQueryStageExec for broadcast reuse
+ *        d. CometReuseSubquery                       -- deduplicates converted subqueries
  *     2. postStageCreationRules -> ApplyColumnarRulesAndInsertTransitions:
  *        a. preColumnarTransitions: CometScanRule, CometExecRule (no-ops, already converted)
  *        b. insertTransitions
  *        c. postColumnarTransitions: EliminateRedundantTransitions
  * }}}
  *
- * CometReuseSubquery is needed in AQE because Spark's ReuseAdaptiveSubquery may run before
- * Comet's node replacements in the initial plan construction, and the replacements can disrupt
- * subquery reuse that was already applied. The shim-based registration
- * (injectQueryStageOptimizerRuleShim) handles API availability: Spark 3.5+ has
- * injectQueryStageOptimizerRule, Spark 3.4 does not (no-op).
+ * On Spark 3.4, injectQueryStageOptimizerRule is unavailable. CometExecRule does not wrap SABs,
+ * and CometPlanAdaptiveDynamicPruningFilters/CometReuseSubquery are not registered. AQE DPP scans
+ * fall back to Spark so that Spark's PlanAdaptiveDynamicPruningFilters handles them natively
+ * (with DPP).
  */
 class CometSparkSessionExtensions
     extends (SparkSessionExtensions => Unit)
@@ -79,8 +89,13 @@ class CometSparkSessionExtensions
   override def apply(extensions: SparkSessionExtensions): Unit = {
     extensions.injectColumnar { session => CometScanColumnar(session) }
     extensions.injectColumnar { session => CometExecColumnar(session) }
+    // Pre-3.5 only: tag AQE DPP regions so the conversion rules below leave them Spark-native.
+    // Registered before CometScanRule/CometExecRule so tags are in place when conversion runs.
+    // No-op on Spark 3.5+; see CometSpark34AqeDppFallbackRule's class docstring.
+    injectPreSpark35QueryStagePrepRuleShim(extensions, CometSpark34AqeDppFallbackRule)
     extensions.injectQueryStagePrepRule { session => CometScanRule(session) }
     extensions.injectQueryStagePrepRule { session => CometExecRule(session) }
+    injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters)
     injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery)
   }
 

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -96,6 +96,13 @@ object CometExecRule {
   val SKIP_COMET_SHUFFLE_TAG: org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit] =
     org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometShuffle")
 
+  /**
+   * Tag set on a `BroadcastExchangeExec` that should be left as a plain Spark broadcast rather
+   * than converted to `CometBroadcastExchangeExec`. Written by [[CometSpark34AqeDppFallbackRule]]
+   * on Spark < 3.5. See that rule's class docstring for the rationale.
+   */
+  val SKIP_COMET_BROADCAST_TAG: org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit] =
+    org.apache.spark.sql.catalyst.trees.TreeNodeTag[Unit]("comet.skipCometBroadcast")
 }
 
 /**
@@ -297,6 +304,11 @@ case class CometExecRule(session: SparkSession)
       // broadcast exchange is forced to be enabled by Comet config.
       case plan if plan.children.exists(_.isInstanceOf[BroadcastExchangeExec]) =>
         val newChildren = plan.children.map {
+          // Tagged by CometSpark34AqeDppFallbackRule on Spark < 3.5 to keep the build-side
+          // broadcast Spark-native so Spark's PlanAdaptiveDynamicPruningFilters can match it.
+          case b: BroadcastExchangeExec
+              if b.getTagValue(CometExecRule.SKIP_COMET_BROADCAST_TAG).isDefined =>
+            b
           case b: BroadcastExchangeExec if b.children.forall(_.isInstanceOf[CometNativeExec]) =>
             convertToComet(b, CometBroadcastExchangeExec).getOrElse(b)
           case other => other
@@ -381,18 +393,22 @@ case class CometExecRule(session: SparkSession)
   }
 
   /**
-   * Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in a node's expressions.
+   * Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in a node's expressions
+   * (non-AQE DPP), and wrap SubqueryAdaptiveBroadcastExec in CometSubqueryAdaptiveBroadcastExec
+   * (AQE DPP) to protect it from Spark's PlanAdaptiveDynamicPruningFilters.
    *
-   * When CometExecRule converts BroadcastExchangeExec to CometBroadcastExchangeExec on the join
-   * side, the DPP subquery still references the original BroadcastExchangeExec.
+   * Non-AQE DPP: When CometExecRule converts BroadcastExchangeExec to CometBroadcastExchangeExec
+   * on the join side, the DPP subquery still references the original BroadcastExchangeExec.
    * ReuseExchangeAndSubquery (which runs after Comet rules) can't match them because they have
    * different types. By replacing SubqueryBroadcastExec with CometSubqueryBroadcastExec (which
    * wraps a CometBroadcastExchangeExec), both sides have the same exchange type and reuse works.
    *
-   * The BroadcastExchangeExec in the subquery has a CometNativeColumnarToRowExec child (inserted
-   * by ApplyColumnarRulesAndInsertTransitions because BroadcastExchangeExec expects row input).
-   * We strip this transition and create CometBroadcastExchangeExec with the underlying Comet plan
-   * directly.
+   * AQE DPP: Spark's PlanAdaptiveDynamicPruningFilters (queryStageOptimizerRule) pattern-matches
+   * on SubqueryAdaptiveBroadcastExec. When it can't find BroadcastHashJoinExec (Comet replaced
+   * it), it replaces DPP with Literal.TrueLiteral. We wrap SABs in
+   * CometSubqueryAdaptiveBroadcastExec to prevent this. CometPlanAdaptiveDynamicPruningFilters (a
+   * later queryStageOptimizerRule) unwraps and converts them with access to the materialized
+   * BroadcastQueryStageExec.
    */
   private def convertSubqueryBroadcasts(plan: SparkPlan): SparkPlan = {
     plan.transformExpressionsUp { case inSub: InSubqueryExec =>
@@ -422,6 +438,32 @@ case class CometExecRule(session: SparkSession)
               }
             case _ => inSub
           }
+        case sab: SubqueryAdaptiveBroadcastExec if isSpark35Plus =>
+          // Wrap SABs to prevent Spark's PlanAdaptiveDynamicPruningFilters from
+          // converting them to Literal.TrueLiteral. Spark's rule pattern-matches for
+          // BroadcastHashJoinExec, which Comet replaced with CometBroadcastHashJoinExec.
+          // Without wrapping, DPP is disabled for both Comet native scans and non-Comet
+          // scans (e.g., V2 BatchScan). CometPlanAdaptiveDynamicPruningFilters
+          // (queryStageOptimizerRule, 3.5+) unwraps and converts them later.
+          //
+          // On Spark 3.4, injectQueryStageOptimizerRule is unavailable. The isSpark35Plus
+          // guard leaves SABs unwrapped; CometSpark34AqeDppFallbackRule then tags the
+          // matching BHJ's build broadcast so Spark's rule can match it natively.
+          assert(
+            sab.buildKeys.nonEmpty,
+            s"SubqueryAdaptiveBroadcastExec '${sab.name}' has empty buildKeys")
+          logInfo(
+            s"Wrapping SubqueryAdaptiveBroadcastExec '${sab.name}' in " +
+              "CometSubqueryAdaptiveBroadcastExec to preserve AQE DPP")
+          val indices = getSubqueryBroadcastIndices(sab)
+          val wrapped = CometSubqueryAdaptiveBroadcastExec(
+            sab.name,
+            indices,
+            sab.onlyInBroadcast,
+            sab.buildPlan,
+            sab.buildKeys,
+            sab.child)
+          inSub.withNewPlan(wrapped)
         case _ => inSub
       }
     }