fix: keep unconverted DPP filters in a scan's canonical form

contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala

@@ -225,13 +225,6 @@ case class CometDeltaNativeScanExec(
   }
 
   override def doCanonicalize(): CometDeltaNativeScanExec = {
-    val canonOriginal = if (originalPlan != null) {
-      val stripped = originalPlan.copy(partitionFilters =
-        CometScanUtils.filterUnusedDynamicPruningExpressions(originalPlan.partitionFilters))
-      stripped.doCanonicalize()
-    } else {
-      null
-    }
     CometDeltaNativeScanExec(
       nativeOp,
       output.map(QueryPlan.normalizeExpressions(_, output)),
@@ -241,7 +234,7 @@ case class CometDeltaNativeScanExec(
         output),
       QueryPlan.normalizePredicates(dataFilters, output),
       relation,
-      canonOriginal,
+      CometScanUtils.canonicalizeOriginalPlan(originalPlan),
       SerializedPlan(None),
       "")
   }

contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala

@@ -27,7 +27,7 @@ import scala.collection.mutable.ListBuffer
 import org.apache.spark.CometListenerBusUtils
 import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
 import org.apache.spark.sql.{DataFrame, Row}
-import org.apache.spark.sql.catalyst.expressions.{AttributeReference, DynamicPruningExpression, Expression, NamedExpression, StructsToJson}
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, DynamicPruningExpression, DynamicPruningSubquery, Expression, NamedExpression, StructsToJson}
 import org.apache.spark.sql.comet.{CometDeltaNativeScanExec, CometSubqueryBroadcastExec}
 import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, QueryExecution, ReusedSubqueryExec, ScalarSubquery, SparkPlan, SubqueryBroadcastExec, SubqueryExec}
 import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec
@@ -1931,6 +1931,112 @@ class CometDeltaNativeScanSuite extends CometDeltaTestBase {
     }
   }
 
+  // SPARK-32509: with no broadcast join the DPP filter on one side of the self-join degrades to
+  // TrueLiteral, so both fact scans must canonicalize alike and share one shuffle.
+  test("unused DPP filter does not prevent exchange reuse (SPARK-32509)") {
+    withTempPath { dir =>
+      val path = dir.getAbsolutePath
+      spark
+        .range(1, 26)
+        .selectExpr(
+          "cast(id as int) as store_id",
+          "cast(case when id = 15 then 70 else id % 5 * 10 end as int) as units_sold")
+        .write
+        .format("delta")
+        .partitionBy("store_id")
+        .save(path)
+
+      withTempView("dpp_reuse_fact") {
+        spark.read.format("delta").load(path).createOrReplaceTempView("dpp_reuse_fact")
+        withSQLConf(
+          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+          SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+          SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+          val df = sql("""WITH v AS (
+              |  SELECT store_id FROM dpp_reuse_fact WHERE units_sold = 70
+              |)
+              |SELECT * FROM v v1 JOIN v v2 ON v1.store_id = v2.store_id""".stripMargin)
+          assert(
+            df.queryExecution.optimizedPlan.exists(_.expressions.exists(_.exists {
+              case _: DynamicPruningSubquery => true
+              case _ => false
+            })),
+            s"expected a DPP filter in the optimized plan:\n${df.queryExecution.optimizedPlan}")
+
+          val (sparkPlan, cometPlan) = checkSparkAnswer(df)
+          checkAnswer(df, Row(15, 15) :: Nil)
+          assert(
+            collect(cometPlan) { case s: CometDeltaNativeScanExec => s }.nonEmpty,
+            s"expected a native Delta scan:\n${cometPlan.treeString}")
+
+          assert(
+            reusedExchanges(sparkPlan).size == 1,
+            s"Spark should reuse the shuffle once:\n$sparkPlan")
+          assertSameExchangeReuseAsSpark(sparkPlan, cometPlan)
+        }
+      }
+    }
+  }
+
+  // Each UNION ALL branch joins the same aggregate over the Delta fact table to a different
+  // dimension group, so the two DPP scans prune to 1 and 3 store partitions. The shuffles above
+  // them differ only in their DPP filter and must not share a query stage. Coalescing is off: a
+  // coalesced read carries per-partition sizes that would keep the parent exchanges apart anyway.
+  test("shuffle over a DPP scan stage is not reused across UNION ALL branches (#6264)") {
+    withTempPath { dir =>
+      val path = dir.getAbsolutePath
+      spark
+        .range(1000)
+        .selectExpr(
+          "cast(id % 10 as int) as store_id",
+          "cast(id % 13 as int) as item",
+          "cast(id as int) as amount")
+        .write
+        .format("delta")
+        .partitionBy("store_id")
+        .save(s"$path/fact")
+      spark
+        .range(10)
+        .selectExpr(
+          "cast(id as int) as store_id",
+          "case when id = 1 then 'a' when id in (2, 3, 4) then 'b' else 'c' end as grp")
+        .write
+        .parquet(s"$path/dim")
+
+      withTempView("dpp_branch_fact", "dpp_branch_dim") {
+        spark.read.format("delta").load(s"$path/fact").createOrReplaceTempView("dpp_branch_fact")
+        spark.read.parquet(s"$path/dim").createOrReplaceTempView("dpp_branch_dim")
+        withSQLConf(
+          SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+          SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+          SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false") {
+          val query = """WITH x AS (
+              |  SELECT store_id, item, sum(amount) AS s FROM dpp_branch_fact
+              |  GROUP BY store_id, item),
+              |y AS (SELECT store_id, s % 7 AS b, count(*) AS c FROM x GROUP BY store_id, s % 7)
+              |SELECT y.store_id, y.b, y.c FROM y
+              |JOIN dpp_branch_dim d ON y.store_id = d.store_id WHERE d.grp = 'a'
+              |UNION ALL
+              |SELECT y.store_id, y.b, y.c FROM y
+              |JOIN dpp_branch_dim d ON y.store_id = d.store_id WHERE d.grp = 'b'""".stripMargin
+          val (_, cometPlan) = checkSparkAnswer(query)
+
+          val dppScan: PartialFunction[SparkPlan, CometDeltaNativeScanExec] = {
+            case s: CometDeltaNativeScanExec
+                if s.runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression]) =>
+              s
+          }
+          assert(
+            collect(cometPlan)(dppScan).map(_.metrics("numPartitions").value).sorted ==
+              Seq(1L, 3L),
+            "Expected one native Delta DPP scan per branch, pruned to 1 and 3 partitions:\n" +
+              cometPlan.treeString)
+          assertNoReusedExchangeOverDppScan(cometPlan)(dppScan)
+        }
+      }
+    }
+  }
+
   test("scalar subquery in a partition filter does not force partitioning during AQE checks") {
     // Crash shape: a scalar subquery used directly as
     // a partition filter, e.g. `p = (SELECT max(p) FROM dim ...)`, references only the

docs/source/contributor-guide/adding_a_new_operator.md

@@ -383,8 +383,17 @@ The examples above follow the implementations in
   rather than serialization state; this method does not define equality.
 
 Some operators use a different convention. `CometNativeScanExec` compares `originalPlan` to retain
-scan identity, and its `doCanonicalize` canonicalizes that plan while removing unused dynamic
-pruning filters. `CometBroadcastExchangeExec` also compares `originalPlan` before canonicalization,
+scan identity. Its `doCanonicalize` passes `originalPlan` to
+`CometScanUtils.canonicalizeOriginalPlan`, which strips every dynamic pruning filter from that copy
+before canonicalizing it, because the copy is stale and no DPP rule rewrites it. The scan's DPP
+identity comes from its top-level `partitionFilters` instead, where
+`CometScanUtils.filterUnusedDynamicPruningExpressions` drops only
+`DynamicPruningExpression(TrueLiteral)`, as Spark's `FileSourceScanExec` does. A new scan with DPP
+support should call both helpers (the second when it keeps a `FileSourceScanExec` as
+`originalPlan`): a filter that still holds the adaptive broadcast placeholder must
+stay in the canonical form, because AQE canonicalizes a query stage as its exchange was before that
+placeholder was converted, and dropping it lets AQE reuse one stage for scans with different
+filters. `CometBroadcastExchangeExec` also compares `originalPlan` before canonicalization,
 but its `doCanonicalize` clears that reference and retains the canonicalized child. Follow each
 operator's equality and canonicalization together rather than copying an exclusion in isolation.
 

spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala

@@ -285,18 +285,6 @@ case class CometNativeScanExec(
   }
 
   override def doCanonicalize(): CometNativeScanExec = {
-    // Canonicalize originalPlan but strip its DPP partition filters.
-    // originalPlan carries column selection and schema info needed for
-    // equals/hashCode. But its partitionFilters may have stale DPP
-    // expressions (e.g., SABs not yet converted to TrueLiteral) that
-    // would prevent exchange reuse between otherwise-identical scans.
-    val canonOriginal = if (originalPlan != null) {
-      val stripped = originalPlan.copy(partitionFilters =
-        CometScanUtils.filterUnusedDynamicPruningExpressions(originalPlan.partitionFilters))
-      stripped.doCanonicalize()
-    } else {
-      null
-    }
     CometNativeScanExec(
       nativeOp,
       relation,
@@ -310,7 +298,7 @@ case class CometNativeScanExec(
       QueryPlan.normalizePredicates(dataFilters, output),
       None,
       disableBucketedScan,
-      canonOriginal,
+      CometScanUtils.canonicalizeOriginalPlan(originalPlan),
       SerializedPlan(None),
       null, // Transient scan not needed for canonicalization
       ""

spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala

@@ -25,7 +25,7 @@ import java.util.Locale
 import scala.collection.mutable
 
 import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, Literal}
-import org.apache.spark.sql.execution.{InSubqueryExec, SubqueryAdaptiveBroadcastExec}
+import org.apache.spark.sql.execution.FileSourceScanExec
 import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionedFile}
 
 import org.apache.comet.objectstore.NativeConfig
@@ -37,20 +37,27 @@ object CometScanUtils {
    * DynamicPruningExpression(Literal.TrueLiteral) during Physical Planning
    */
   def filterUnusedDynamicPruningExpressions(predicates: Seq[Expression]): Seq[Expression] = {
-    // Strip DPP expressions for canonicalization. Matches Spark's
-    // FileSourceScanExec.filterUnusedDynamicPruningExpressions (TrueLiteral).
-    // Also strips unconverted SAB wrappers because AQE stageCache canonicalizes
-    // before our queryStageOptimizerRule converts them, so they would prevent
-    // exchange reuse between otherwise-identical scans.
-    predicates.filterNot {
-      case DynamicPruningExpression(Literal.TrueLiteral) => true
-      case DynamicPruningExpression(
-            InSubqueryExec(_, _: CometSubqueryAdaptiveBroadcastExec, _, _, _, _)) =>
-        true
-      case DynamicPruningExpression(
-            InSubqueryExec(_, _: SubqueryAdaptiveBroadcastExec, _, _, _, _)) =>
-        true
-      case _ => false
+    // Matches Spark's FileSourceScanExec.filterUnusedDynamicPruningExpressions. A DPP filter that
+    // still holds the adaptive broadcast placeholder must stay: AQE keys a query stage by its
+    // exchange as it was before the stage optimizer rules converted that placeholder, so dropping
+    // it makes scans with different DPP filters equal and lets AQE reuse one stage for both.
+    predicates.filterNot(_ == DynamicPruningExpression(Literal.TrueLiteral))
+  }
+
+  /**
+   * Canonicalizes a Comet scan's `originalPlan` (which may be null) without any of its DPP
+   * partition filters. That copy carries the column selection and schema needed for equals and
+   * hashCode, but its DPP filters are stale because no DPP rule rewrites them. The scan's own DPP
+   * filters, passed through [[filterUnusedDynamicPruningExpressions]], carry its DPP identity.
+   */
+  def canonicalizeOriginalPlan(originalPlan: FileSourceScanExec): FileSourceScanExec = {
+    if (originalPlan == null) {
+      null
+    } else {
+      originalPlan
+        .copy(partitionFilters =
+          originalPlan.partitionFilters.filterNot(_.isInstanceOf[DynamicPruningExpression]))
+        .doCanonicalize()
     }
   }
 

spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala

@@ -6102,6 +6102,87 @@ class CometIcebergNativeSuite
     }
   }
 
+  // TPC-DS q64 shape (#6264): a broadcast of (fact SMJ other) per UNION ALL branch, with the
+  // dimension join above it. The two broadcasts differ only in the DPP filter of the fact scan
+  // inside a child shuffle stage, so they must not be reused for each other.
+  test("AQE DPP - broadcast over a join of DPP scans is not reused across branches (#6264)") {
+    assume(icebergAvailable, "Iceberg not available")
+    withTempIcebergDir { warehouseDir =>
+      val dimDir = new File(warehouseDir, "dim_parquet")
+      val otherDir = new File(warehouseDir, "other_parquet")
+      withSQLConf(
+        "spark.sql.catalog.aqe_cat" -> "org.apache.iceberg.spark.SparkCatalog",
+        "spark.sql.catalog.aqe_cat.type" -> "hadoop",
+        "spark.sql.catalog.aqe_cat.warehouse" -> warehouseDir.getAbsolutePath,
+        SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+        SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true",
+        // A coalesced shuffle read carries per-partition sizes that keep the two broadcasts apart.
+        SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false",
+        CometConf.COMET_ENABLED.key -> "true",
+        CometConf.COMET_EXEC_ENABLED.key -> "true",
+        CometConf.COMET_ICEBERG_NATIVE_ENABLED.key -> "true") {
+
+        spark.sql("""
+          CREATE TABLE aqe_cat.db.branch_fact (
+            store_id INT, item INT, amount INT
+          ) USING iceberg PARTITIONED BY (store_id)
+        """)
+        spark
+          .range(1000)
+          .selectExpr(
+            "cast(id % 10 as int) as store_id",
+            "cast(id % 13 as int) as item",
+            "cast(id as int) as amount")
+          .write
+          .format("iceberg")
+          .mode("append")
+          .saveAsTable("aqe_cat.db.branch_fact")
+        // Group 'a' keeps store 1 and group 'b' keeps stores 2 to 4.
+        spark
+          .range(10)
+          .selectExpr(
+            "cast(id as int) as store_id",
+            "case when id = 1 then 'a' when id in (2, 3, 4) then 'b' else 'c' end as grp")
+          .write
+          .parquet(dimDir.getAbsolutePath)
+        spark.read.parquet(dimDir.getAbsolutePath).createOrReplaceTempView("branch_dim")
+        spark
+          .range(13)
+          .selectExpr("cast(id as int) as item", "cast(id * 10 as int) as w")
+          .write
+          .parquet(otherDir.getAbsolutePath)
+        spark.read.parquet(otherDir.getAbsolutePath).createOrReplaceTempView("branch_other")
+
+        val query =
+          """WITH x AS (
+            |  SELECT /*+ MERGE(f, o) */ f.store_id, f.item, f.amount, o.w
+            |  FROM aqe_cat.db.branch_fact f JOIN branch_other o ON f.item = o.item)
+            |SELECT /*+ BROADCAST(x) */ o2.item, x.store_id, x.amount, x.w
+            |FROM branch_other o2 JOIN x ON o2.item = x.item
+            |JOIN branch_dim d ON x.store_id = d.store_id WHERE d.grp = 'a'
+            |UNION ALL
+            |SELECT /*+ BROADCAST(x) */ o2.item, x.store_id, x.amount, x.w
+            |FROM branch_other o2 JOIN x ON o2.item = x.item
+            |JOIN branch_dim d ON x.store_id = d.store_id WHERE d.grp = 'b'""".stripMargin
+        val (_, cometPlan) = checkSparkAnswer(query)
+        assertIcebergNativeScanPresent(cometPlan)
+        assertNoLeftoverCSAB(cometPlan)
+
+        // Also checked on 3.4: the Iceberg scan stays in Comet there and Spark's own DPP rule
+        // plans its filter.
+        val dppScan: PartialFunction[SparkPlan, CometIcebergNativeScanExec] = {
+          case s: CometIcebergNativeScanExec if s.runtimeFilters.nonEmpty => s
+        }
+        assert(
+          collect(cometPlan)(dppScan).size == 2,
+          s"Expected one DPP scan per branch:\n${cometPlan.treeString}")
+        assertNoReusedExchangeOverDppScan(cometPlan)(dppScan)
+
+        spark.sql("DROP TABLE aqe_cat.db.branch_fact")
+      }
+    }
+  }
+
   // Multi-key BHJ where SAB build keys must keep their position so the index lookup
   // selects the right DPP value. A bug in convertSAB key handling would either pick
   // the wrong column or produce SubqueryExec instead of broadcast reuse.