From 7dadec5e59f549d4b74feb6236a75fc461c2c41d Mon Sep 17 00:00:00 2001 From: dustin Date: Fri, 16 Oct 2026 16:05:55 +0000 Subject: [PATCH] fix: keep unconverted DPP filters in a scan's canonical form (#6270) * fix: keep unconverted DPP filters in a scan's canonical form AQE canonicalizes a query stage from its exchange before the stage rules convert DPP placeholders. Comet dropped a DPP filter from the scan's canonical form while it was still a placeholder, so exchanges over scans with different DPP filters compared equal and AQE reused one for both, losing the other branch's rows. The shared helper now drops only DynamicPruningExpression(TrueLiteral), as FileSourceScanExec does. CometNativeScanExec.doCanonicalize drops every DPP filter from originalPlan, a stale copy no DPP rule rewrites, so an unused DPP scan still matches its twin. Closes #6264. Part of #6133. * docs: describe how a native scan keeps its DPP identity when canonicalized The plan identity section still said CometNativeScanExec canonicalizes originalPlan while removing unused dynamic pruning filters. It now strips every DPP filter from that stale copy, and the scan's DPP identity comes from the top-level partitionFilters, where only DynamicPruningExpression(TrueLiteral) is dropped. Say both, and point a new scan at the second rule with the reason it matters for AQE reuse. * fix: strip every DPP filter from the Delta scan's originalPlan when canonicalizing CometDeltaNativeScanExec kept a stale DPP placeholder in its originalPlan copy once filterUnusedDynamicPruningExpressions stopped dropping placeholders, so after AQE turned the live filter into TrueLiteral the scan no longer matched its twin and lost SPARK-32509 exchange reuse. Drop every DPP filter from originalPlan, as CometNativeScanExec does; runtimeFilters still carry the scan's DPP identity. Adds a Delta version of the SPARK-32509 reuse test. * test: check that DPP scan stages on a Delta table are not reused across branches Ports the #6264 UNION ALL test to the native Delta scan, so a Delta scan whose live DPP filter differs from its twin's keeps its own shuffle stage and its own rows. * refactor: share the scan originalPlan canonicalization and DPP reuse check CometNativeScanExec and CometDeltaNativeScanExec now call CometScanUtils.canonicalizeOriginalPlan instead of each stripping DPP filters from originalPlan inline, and the operator guide points new scans at it. The no-reuse-over-a-DPP-scan assertion lives once in CometTestBase and the Comet, Iceberg and Delta suites use it. The CometExecSuite test that passed on main is replaced by one where both branches use the same DPP filter: the shuffles must still be reused as in Spark, which fails if the canonical form keeps too much. * test: share the reused exchange helpers in CometTestBase assertExchangeReuseOver and assertNoReusedExchangeOverDppScan now call one reusedExchangesOver, and assertSameExchangeReuseAsSpark replaces the reuse count comparison written out in CometExecSuite and CometDeltaNativeScanSuite. --------- Co-authored-by: Andy Grove (cherry picked from commit e242944ec43faf76ad9e22217f22e57f7e801ef3) Adapted for branch-1.1: the Delta contrib (CometDeltaNativeScanExec and its suite) is not on this branch, so those two files are dropped, and so is the contributor guide change. CometScanUtils keeps only the FileSourceScanExec import, since #6753's FilePartition and NativeConfig code is not here. --- .../spark/sql/comet/CometNativeScanExec.scala | 14 +- .../spark/sql/comet/CometScanUtils.scala | 37 +++--- .../comet/CometIcebergNativeSuite.scala | 81 ++++++++++++ .../apache/comet/exec/CometExecSuite.scala | 123 ++++++++++++++++++ .../org/apache/spark/sql/CometTestBase.scala | 38 +++++- 5 files changed, 260 insertions(+), 33 deletions(-) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala index 74cdbdc6947..a761e762021 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala @@ -287,18 +287,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, @@ -312,7 +300,7 @@ case class CometNativeScanExec( QueryPlan.normalizePredicates(dataFilters, output), None, disableBucketedScan, - canonOriginal, + CometScanUtils.canonicalizeOriginalPlan(originalPlan), SerializedPlan(None), null, // Transient scan not needed for canonicalization "" diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala index 19a2b53af11..b8f48da922d 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometScanUtils.scala @@ -20,7 +20,7 @@ package org.apache.spark.sql.comet 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 object CometScanUtils { @@ -29,20 +29,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() } } } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 4d64149416e..8635243913d 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -5558,6 +5558,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. diff --git a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala index ec1ff6542d8..3b0b48ac5a0 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -1569,6 +1569,129 @@ class CometExecSuite extends CometTestBase { } } + // Tables for the #6264 tests. The fact table has 10 store partitions. Dimension group 'a' keeps + // store 1 and group 'b' keeps stores 2 to 4, so the two DPP scans prune to 1 and 3 partitions. + private def withDppBranchTables(f: => Unit): Unit = { + withTempDir { 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 + .partitionBy("store_id") + .parquet(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") + spark + .range(13) + .selectExpr("cast(id as int) as item", "cast(id * 10 as int) as w") + .write + .parquet(s"$path/other") + withTempView("dpp_branch_fact", "dpp_branch_dim", "dpp_branch_other") { + spark.read.parquet(s"$path/fact").createOrReplaceTempView("dpp_branch_fact") + spark.read.parquet(s"$path/dim").createOrReplaceTempView("dpp_branch_dim") + spark.read.parquet(s"$path/other").createOrReplaceTempView("dpp_branch_other") + f + } + } + } + + // Coalescing is off: a coalesced shuffle read carries per-partition sizes that would keep the two + // branches' parent exchanges apart and hide a wrong reuse. + private def withDppBranchConf(f: => Unit): Unit = { + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.DYNAMIC_PARTITION_PRUNING_REUSE_BROADCAST_ONLY.key -> "true", + SQLConf.COALESCE_PARTITIONS_ENABLED.key -> "false")(f) + } + + private val dppScan: PartialFunction[SparkPlan, CometNativeScanExec] = { + case s: CometNativeScanExec + if s.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]) => + s + } + + // A UNION ALL of one CTE joined to dimension groups grp1 and grp2, with a shuffle stage over each + // DPP scan. + private def dppShuffleQuery(grp1: String, grp2: String): String = + s"""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 = '$grp1' + |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 = '$grp2'""".stripMargin + + // Runs a UNION ALL of one CTE joined to dimension group 'a' and to group 'b'. + private def checkDppBranchesKeepTheirRows(query: String): Unit = { + withDppBranchConf { + if (isSpark35Plus) { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query)) + assert( + collect(cometPlan)(dppScan).map(_.metrics("numPartitions").value).sorted == + Seq(1L, 3L), + "Expected one DPP scan per branch, pruned to 1 and 3 partitions:\n" + + cometPlan.treeString) + assertNoReusedExchangeOverDppScan(cometPlan)(dppScan) + } else { + // On 3.4, AQE DPP falls back to Spark. + checkSparkAnswer(sql(query)) + } + } + } + + // TPC-DS q64 shape: a broadcast of (fact SMJ other), with the dimension join above it. The two + // broadcasts differ only in the DPP filter of the fact scan inside a child shuffle stage. + test("AQE DPP: broadcast over a join of DPP scans is not reused across branches (#6264)") { + withDppBranchTables { + checkDppBranchesKeepTheirRows("""WITH x AS ( + | SELECT /*+ MERGE(f, o) */ f.store_id, f.item, f.amount, o.w + | FROM dpp_branch_fact f JOIN dpp_branch_other o ON f.item = o.item) + |SELECT /*+ BROADCAST(x) */ o2.item, x.store_id, x.amount, x.w + |FROM dpp_branch_other o2 JOIN x ON o2.item = x.item + |JOIN dpp_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 dpp_branch_other o2 JOIN x ON o2.item = x.item + |JOIN dpp_branch_dim d ON x.store_id = d.store_id WHERE d.grp = 'b'""".stripMargin) + } + } + + test("AQE DPP: shuffle over a DPP scan stage is not reused across branches (#6264)") { + withDppBranchTables { + checkDppBranchesKeepTheirRows(dppShuffleQuery("a", "b")) + } + } + + // With the same dimension filter in both branches, AQE's stage cache lookup after the DPP + // placeholder is converted finds the first branch's stage, so the shuffles are reused as in + // Spark. This fails if the scan's canonical form keeps too much, such as originalPlan's DPP + // filters or unnormalized partition filters. + test("AQE DPP: shuffle over scans with the same DPP filter is reused across branches (#6264)") { + withDppBranchTables { + withDppBranchConf { + val df = sql(dppShuffleQuery("a", "a")) + val (sparkPlan, cometPlan) = + if (isSpark35Plus) checkSparkAnswerAndOperator(df) else checkSparkAnswer(df) + assertSameExchangeReuseAsSpark(sparkPlan, cometPlan) + if (isSpark35Plus) { + assertExchangeReuseOver(cometPlan, "Expected the shuffle over a DPP scan to be reused")( + dppScan) + } + } + } + } + // SPARK-34637: previously IgnoreComet(#4045). DPP side broadcast query stage // should be created before the main join's broadcast stage. test("AQE DPP: broadcast query stage creation order (SPARK-34637)") { diff --git a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala index cdd466f03d3..bc9dc7c9570 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -68,12 +68,40 @@ abstract class CometTestBase protected val shuffleManager: String = "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager" + /** The reused exchanges in `plan`. */ + protected def reusedExchanges(plan: SparkPlan): Seq[ReusedExchangeExec] = + collect(plan) { case r: ReusedExchangeExec => r } + + /** The reused exchanges in `plan` whose child contains a node that `pf` matches. */ + protected def reusedExchangesOver[T](plan: SparkPlan)( + pf: PartialFunction[SparkPlan, T]): Seq[ReusedExchangeExec] = + reusedExchanges(plan).filter(exchange => collect(exchange.child)(pf).nonEmpty) + + /** Asserts that some reused exchange in `plan` sits over a node that `pf` matches. */ protected def assertExchangeReuseOver[T](plan: SparkPlan, clue: String)( - pf: PartialFunction[SparkPlan, T]): Unit = { - val reused = collect(plan) { - case exchange: ReusedExchangeExec if collect(exchange.child)(pf).nonEmpty => exchange - } - assert(reused.nonEmpty, s"$clue:\n$plan") + pf: PartialFunction[SparkPlan, T]): Unit = + assert(reusedExchangesOver(plan)(pf).nonEmpty, s"$clue:\n$plan") + + /** Asserts that no reused exchange in `plan` sits over a scan that `dppScan` matches. */ + protected def assertNoReusedExchangeOverDppScan[T](plan: SparkPlan)( + dppScan: PartialFunction[SparkPlan, T]): Unit = + assert( + reusedExchangesOver(plan)(dppScan).isEmpty, + s"An exchange over a DPP scan was reused by another branch:\n$plan") + + /** + * Asserts that Spark reuses at least one exchange in `sparkPlan` and that Comet reuses as many + * in `cometPlan`. + */ + protected def assertSameExchangeReuseAsSpark( + sparkPlan: SparkPlan, + cometPlan: SparkPlan): Unit = { + val sparkReused = reusedExchanges(sparkPlan).size + val cometReused = reusedExchanges(cometPlan).size + assert(sparkReused > 0, s"Spark should reuse an exchange:\n$sparkPlan") + assert( + cometReused == sparkReused, + s"expected $sparkReused reused exchanges, got $cometReused:\n" + cometPlan.treeString) } protected def sparkConf: SparkConf = {