From ae530c268d6a7617f851308143a1a917bec6764c Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Sun, 27 Sep 2026 22:19:29 +0700 Subject: [PATCH 1/6] 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. --- .../spark/sql/comet/CometNativeScanExec.scala | 11 +- .../spark/sql/comet/CometScanUtils.scala | 21 +--- .../comet/CometIcebergNativeSuite.scala | 85 ++++++++++++++ .../apache/comet/exec/CometExecSuite.scala | 111 ++++++++++++++++++ 4 files changed, 206 insertions(+), 22 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 a4365c00750..5c5ccd3d3d7 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 @@ -289,14 +289,13 @@ 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. + // Canonicalize originalPlan without any of its DPP partition filters. originalPlan carries + // column selection and schema info needed for equals/hashCode, but its DPP filters are a stale + // copy that no DPP rule rewrites. The top-level partitionFilters below carry the scan's DPP + // identity. val canonOriginal = if (originalPlan != null) { val stripped = originalPlan.copy(partitionFilters = - CometScanUtils.filterUnusedDynamicPruningExpressions(originalPlan.partitionFilters)) + originalPlan.partitionFilters.filterNot(_.isInstanceOf[DynamicPruningExpression])) stripped.doCanonicalize() } else { null 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..163e3b7168c 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,6 @@ package org.apache.spark.sql.comet import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, Literal} -import org.apache.spark.sql.execution.{InSubqueryExec, SubqueryAdaptiveBroadcastExec} object CometScanUtils { @@ -29,20 +28,10 @@ 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)) } } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index 256ed221355..b3c89df0d4e 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -5244,6 +5244,91 @@ 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. + def dppScans(plan: SparkPlan): Seq[CometIcebergNativeScanExec] = + collectIcebergNativeScans(plan).filter(_.runtimeFilters.nonEmpty) + assert( + dppScans(cometPlan).size == 2, + s"Expected one DPP scan per branch:\n${cometPlan.treeString}") + val reusedDpp = collect(cometPlan) { + case r: ReusedExchangeExec if dppScans(r.child).nonEmpty => r + } + assert( + reusedDpp.isEmpty, + s"An exchange over a DPP scan was reused by the other branch:\n${cometPlan.treeString}") + + 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 cbe86fadf34..f8d50e9aae7 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -1568,6 +1568,117 @@ 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 + } + } + } + + // Runs a UNION ALL of one CTE joined to dimension group 'a' and to group 'b'. 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 checkDppBranchesKeepTheirRows(query: String): 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") { + if (isSpark35Plus) { + val (_, cometPlan) = checkSparkAnswerAndOperator(sql(query)) + def dppScans(plan: SparkPlan): Seq[CometNativeScanExec] = collect(plan) { + case s: CometNativeScanExec + if s.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]) => + s + } + assert( + dppScans(cometPlan).map(_.metrics("numPartitions").value).sorted == Seq(1L, 3L), + "Expected one DPP scan per branch, pruned to 1 and 3 partitions:\n" + + cometPlan.treeString) + val reusedDpp = collect(cometPlan) { + case r: ReusedExchangeExec if dppScans(r.child).nonEmpty => r + } + assert( + reusedDpp.isEmpty, + s"An exchange over a DPP scan was reused by the other branch:\n${cometPlan.treeString}") + } 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("""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) + } + } + + // Guards the scan-stage path, which passes without keeping the DPP placeholder in the canonical + // form: the dimension join sits in the scan's own stage, and AQE keys each scan stage after its + // DPP filter is converted. + test("AQE DPP: DPP join in the scan's own stage keeps both branches (#6264)") { + withDppBranchTables { + checkDppBranchesKeepTheirRows("""SELECT f.store_id, f.item, f.amount FROM dpp_branch_fact f + |JOIN dpp_branch_dim d ON f.store_id = d.store_id WHERE d.grp = 'a' + |UNION ALL + |SELECT f.store_id, f.item, f.amount FROM dpp_branch_fact f + |JOIN dpp_branch_dim d ON f.store_id = d.store_id WHERE d.grp = 'b'""".stripMargin) + } + } + // 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)") { From da0a37921c4acf7e3f02228530871aa4a87ccce1 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Mon, 5 Oct 2026 21:01:03 +0700 Subject: [PATCH 2/6] 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. --- .../source/contributor-guide/adding_a_new_operator.md | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index f5b8f00a5c1..5d6acb46b00 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -371,8 +371,15 @@ 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` strips every dynamic pruning filter from `originalPlan` before +canonicalizing it, because that 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 copy that second rule: 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. From 5cf465440025c0760b992a14f692f30ba9857424 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Tue, 13 Oct 2026 18:13:35 +0700 Subject: [PATCH 3/6] 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. --- .../sql/comet/CometDeltaNativeScanExec.scala | 4 +- .../delta/CometDeltaNativeScanSuite.scala | 53 ++++++++++++++++++- 2 files changed, 55 insertions(+), 2 deletions(-) diff --git a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala index c1737172dd2..38831df6a76 100644 --- a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala +++ b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala @@ -225,9 +225,11 @@ case class CometDeltaNativeScanExec( } override def doCanonicalize(): CometDeltaNativeScanExec = { + // originalPlan's DPP filters are a stale copy that no DPP rule rewrites, so drop them all. + // runtimeFilters below carry the scan's DPP identity. val canonOriginal = if (originalPlan != null) { val stripped = originalPlan.copy(partitionFilters = - CometScanUtils.filterUnusedDynamicPruningExpressions(originalPlan.partitionFilters)) + originalPlan.partitionFilters.filterNot(_.isInstanceOf[DynamicPruningExpression])) stripped.doCanonicalize() } else { null diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala index 3b82cb32654..be353b00206 100644 --- a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala @@ -27,10 +27,11 @@ 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 +import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.functions.{col, lit, to_json} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ByteType, LongType, StringType, StructField, StructType} @@ -1931,6 +1932,56 @@ 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}") + + def reused(plan: SparkPlan): Int = + collect(plan) { case r: ReusedExchangeExec => r }.size + assert(reused(sparkPlan) == 1, s"Spark should reuse the shuffle:\n$sparkPlan") + assert( + reused(cometPlan) == reused(sparkPlan), + s"expected ${reused(sparkPlan)} reused exchange, got ${reused(cometPlan)}:\n" + + cometPlan.treeString) + } + } + } + } + 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 From 91a3ae17f063a789f04c4e72246cf35cb0eb2ef3 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Tue, 13 Oct 2026 18:54:05 +0700 Subject: [PATCH 4/6] 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. --- .../delta/CometDeltaNativeScanSuite.scala | 64 +++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala index be353b00206..478b3fde780 100644 --- a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala @@ -1982,6 +1982,70 @@ class CometDeltaNativeScanSuite extends CometDeltaTestBase { } } + // 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) + + def dppScans(plan: SparkPlan): Seq[CometDeltaNativeScanExec] = collect(plan) { + case s: CometDeltaNativeScanExec + if s.runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression]) => + s + } + assert( + dppScans(cometPlan).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) + val reusedDpp = collect(cometPlan) { + case r: ReusedExchangeExec if dppScans(r.child).nonEmpty => r + } + assert( + reusedDpp.isEmpty, + "An exchange over a DPP scan was reused by the other branch:\n" + + cometPlan.treeString) + } + } + } + } + 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 From eb188dd6a9f2a3cdfcd376427313ff571ab7dfa0 Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Wed, 14 Oct 2026 10:44:37 +0700 Subject: [PATCH 5/6] 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. --- .../sql/comet/CometDeltaNativeScanExec.scala | 11 +-- .../delta/CometDeltaNativeScanSuite.scala | 13 +-- .../adding_a_new_operator.md | 14 +-- .../spark/sql/comet/CometNativeScanExec.scala | 13 +-- .../spark/sql/comet/CometScanUtils.scala | 18 ++++ .../comet/CometIcebergNativeSuite.scala | 14 ++- .../apache/comet/exec/CometExecSuite.scala | 86 +++++++++++-------- .../org/apache/spark/sql/CometTestBase.scala | 9 ++ 8 files changed, 98 insertions(+), 80 deletions(-) diff --git a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala index 38831df6a76..bcd1d82f906 100644 --- a/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala +++ b/contrib/delta-spark/src/main/scala/org/apache/spark/sql/comet/CometDeltaNativeScanExec.scala @@ -225,15 +225,6 @@ case class CometDeltaNativeScanExec( } override def doCanonicalize(): CometDeltaNativeScanExec = { - // originalPlan's DPP filters are a stale copy that no DPP rule rewrites, so drop them all. - // runtimeFilters below carry the scan's DPP identity. - val canonOriginal = if (originalPlan != null) { - val stripped = originalPlan.copy(partitionFilters = - originalPlan.partitionFilters.filterNot(_.isInstanceOf[DynamicPruningExpression])) - stripped.doCanonicalize() - } else { - null - } CometDeltaNativeScanExec( nativeOp, output.map(QueryPlan.normalizeExpressions(_, output)), @@ -243,7 +234,7 @@ case class CometDeltaNativeScanExec( output), QueryPlan.normalizePredicates(dataFilters, output), relation, - canonOriginal, + CometScanUtils.canonicalizeOriginalPlan(originalPlan), SerializedPlan(None), "") } diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala index 478b3fde780..327963d3561 100644 --- a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala @@ -2025,22 +2025,17 @@ class CometDeltaNativeScanSuite extends CometDeltaTestBase { |JOIN dpp_branch_dim d ON y.store_id = d.store_id WHERE d.grp = 'b'""".stripMargin val (_, cometPlan) = checkSparkAnswer(query) - def dppScans(plan: SparkPlan): Seq[CometDeltaNativeScanExec] = collect(plan) { + val dppScan: PartialFunction[SparkPlan, CometDeltaNativeScanExec] = { case s: CometDeltaNativeScanExec if s.runtimeFilters.exists(_.isInstanceOf[DynamicPruningExpression]) => s } assert( - dppScans(cometPlan).map(_.metrics("numPartitions").value).sorted == Seq(1L, 3L), + 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) - val reusedDpp = collect(cometPlan) { - case r: ReusedExchangeExec if dppScans(r.child).nonEmpty => r - } - assert( - reusedDpp.isEmpty, - "An exchange over a DPP scan was reused by the other branch:\n" + - cometPlan.treeString) + assertNoReusedExchangeOverDppScan(cometPlan)(dppScan) } } } diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 0ccc79302b8..caf45198336 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/docs/source/contributor-guide/adding_a_new_operator.md @@ -383,14 +383,16 @@ 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. Its `doCanonicalize` strips every dynamic pruning filter from `originalPlan` before -canonicalizing it, because that copy is stale and no DPP rule rewrites it. The scan's DPP identity -comes from its top-level `partitionFilters` instead, where +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 copy that second rule: 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 +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. 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 f55b8743c12..8c8456b074a 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 @@ -285,17 +285,6 @@ case class CometNativeScanExec( } override def doCanonicalize(): CometNativeScanExec = { - // Canonicalize originalPlan without any of its DPP partition filters. originalPlan carries - // column selection and schema info needed for equals/hashCode, but its DPP filters are a stale - // copy that no DPP rule rewrites. The top-level partitionFilters below carry the scan's DPP - // identity. - val canonOriginal = if (originalPlan != null) { - val stripped = originalPlan.copy(partitionFilters = - originalPlan.partitionFilters.filterNot(_.isInstanceOf[DynamicPruningExpression])) - stripped.doCanonicalize() - } else { - null - } CometNativeScanExec( nativeOp, relation, @@ -309,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 "" 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 163e3b7168c..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,6 +20,7 @@ package org.apache.spark.sql.comet import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Expression, Literal} +import org.apache.spark.sql.execution.FileSourceScanExec object CometScanUtils { @@ -34,4 +35,21 @@ object CometScanUtils { // 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 2ab25c468fd..58a266be65d 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala @@ -6170,17 +6170,13 @@ class CometIcebergNativeSuite // Also checked on 3.4: the Iceberg scan stays in Comet there and Spark's own DPP rule // plans its filter. - def dppScans(plan: SparkPlan): Seq[CometIcebergNativeScanExec] = - collectIcebergNativeScans(plan).filter(_.runtimeFilters.nonEmpty) - assert( - dppScans(cometPlan).size == 2, - s"Expected one DPP scan per branch:\n${cometPlan.treeString}") - val reusedDpp = collect(cometPlan) { - case r: ReusedExchangeExec if dppScans(r.child).nonEmpty => r + val dppScan: PartialFunction[SparkPlan, CometIcebergNativeScanExec] = { + case s: CometIcebergNativeScanExec if s.runtimeFilters.nonEmpty => s } assert( - reusedDpp.isEmpty, - s"An exchange over a DPP scan was reused by the other branch:\n${cometPlan.treeString}") + 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") } 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 f95c6aa8be0..a788e03d854 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -1804,32 +1804,45 @@ class CometExecSuite extends CometTestBase { } } - // Runs a UNION ALL of one CTE joined to dimension group 'a' and to group 'b'. 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 checkDppBranchesKeepTheirRows(query: String): Unit = { + // 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") { + 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)) - def dppScans(plan: SparkPlan): Seq[CometNativeScanExec] = collect(plan) { - case s: CometNativeScanExec - if s.partitionFilters.exists(_.isInstanceOf[DynamicPruningExpression]) => - s - } assert( - dppScans(cometPlan).map(_.metrics("numPartitions").value).sorted == Seq(1L, 3L), + 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) - val reusedDpp = collect(cometPlan) { - case r: ReusedExchangeExec if dppScans(r.child).nonEmpty => r - } - assert( - reusedDpp.isEmpty, - s"An exchange over a DPP scan was reused by the other branch:\n${cometPlan.treeString}") + assertNoReusedExchangeOverDppScan(cometPlan)(dppScan) } else { // On 3.4, AQE DPP falls back to Spark. checkSparkAnswer(sql(query)) @@ -1856,27 +1869,32 @@ class CometExecSuite extends CometTestBase { test("AQE DPP: shuffle over a DPP scan stage is not reused across branches (#6264)") { withDppBranchTables { - checkDppBranchesKeepTheirRows("""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) + checkDppBranchesKeepTheirRows(dppShuffleQuery("a", "b")) } } - // Guards the scan-stage path, which passes without keeping the DPP placeholder in the canonical - // form: the dimension join sits in the scan's own stage, and AQE keys each scan stage after its - // DPP filter is converted. - test("AQE DPP: DPP join in the scan's own stage keeps both branches (#6264)") { + // 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 { - checkDppBranchesKeepTheirRows("""SELECT f.store_id, f.item, f.amount FROM dpp_branch_fact f - |JOIN dpp_branch_dim d ON f.store_id = d.store_id WHERE d.grp = 'a' - |UNION ALL - |SELECT f.store_id, f.item, f.amount FROM dpp_branch_fact f - |JOIN dpp_branch_dim d ON f.store_id = d.store_id WHERE d.grp = 'b'""".stripMargin) + withDppBranchConf { + val df = sql(dppShuffleQuery("a", "a")) + val (sparkPlan, cometPlan) = + if (isSpark35Plus) checkSparkAnswerAndOperator(df) else checkSparkAnswer(df) + def reused(plan: SparkPlan): Int = + collect(plan) { case r: ReusedExchangeExec => r }.size + assert(reused(sparkPlan) > 0, s"Spark should reuse the shuffle:\n$sparkPlan") + assert( + reused(cometPlan) == reused(sparkPlan), + s"expected ${reused(sparkPlan)} reused exchanges, got ${reused(cometPlan)}:\n" + + cometPlan.treeString) + if (isSpark35Plus) { + assertExchangeReuseOver(cometPlan, "Expected the shuffle over a DPP scan to be reused")( + dppScan) + } + } } } 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 5ff3f7ed08f..7f538432e48 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -78,6 +78,15 @@ abstract class CometTestBase assert(reused.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 = { + val reused = collect(plan) { + case exchange: ReusedExchangeExec if collect(exchange.child)(dppScan).nonEmpty => exchange + } + assert(reused.isEmpty, s"An exchange over a DPP scan was reused by another branch:\n$plan") + } + /** * Sets each Spark-to-Arrow conversion that [[sparkConf]] turns on to `enabled`. All but the * `OneRowRelation` conversion are off by default. These are the conversions From 4ac089278eb59f2d7387adf59e5598c6e576762b Mon Sep 17 00:00:00 2001 From: Dustin Smith Date: Fri, 16 Oct 2026 00:41:55 +0700 Subject: [PATCH 6/6] 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. --- .../delta/CometDeltaNativeScanSuite.scala | 10 ++--- .../apache/comet/exec/CometExecSuite.scala | 8 +--- .../org/apache/spark/sql/CometTestBase.scala | 41 ++++++++++++++----- 3 files changed, 34 insertions(+), 25 deletions(-) diff --git a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala index 327963d3561..a4dec885391 100644 --- a/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala +++ b/contrib/delta-spark/src/test/scala/org/apache/comet/contrib/delta/CometDeltaNativeScanSuite.scala @@ -31,7 +31,6 @@ import org.apache.spark.sql.catalyst.expressions.{AttributeReference, DynamicPru 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 -import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.functions.{col, lit, to_json} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{ByteType, LongType, StringType, StructField, StructType} @@ -1970,13 +1969,10 @@ class CometDeltaNativeScanSuite extends CometDeltaTestBase { collect(cometPlan) { case s: CometDeltaNativeScanExec => s }.nonEmpty, s"expected a native Delta scan:\n${cometPlan.treeString}") - def reused(plan: SparkPlan): Int = - collect(plan) { case r: ReusedExchangeExec => r }.size - assert(reused(sparkPlan) == 1, s"Spark should reuse the shuffle:\n$sparkPlan") assert( - reused(cometPlan) == reused(sparkPlan), - s"expected ${reused(sparkPlan)} reused exchange, got ${reused(cometPlan)}:\n" + - cometPlan.treeString) + reusedExchanges(sparkPlan).size == 1, + s"Spark should reuse the shuffle once:\n$sparkPlan") + assertSameExchangeReuseAsSpark(sparkPlan, cometPlan) } } } 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 a788e03d854..7434d877831 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -1883,13 +1883,7 @@ class CometExecSuite extends CometTestBase { val df = sql(dppShuffleQuery("a", "a")) val (sparkPlan, cometPlan) = if (isSpark35Plus) checkSparkAnswerAndOperator(df) else checkSparkAnswer(df) - def reused(plan: SparkPlan): Int = - collect(plan) { case r: ReusedExchangeExec => r }.size - assert(reused(sparkPlan) > 0, s"Spark should reuse the shuffle:\n$sparkPlan") - assert( - reused(cometPlan) == reused(sparkPlan), - s"expected ${reused(sparkPlan)} reused exchanges, got ${reused(cometPlan)}:\n" + - cometPlan.treeString) + assertSameExchangeReuseAsSpark(sparkPlan, cometPlan) if (isSpark35Plus) { assertExchangeReuseOver(cometPlan, "Expected the shuffle over a DPP scan to be reused")( dppScan) 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 12fc7eed94c..2f7ba3896d4 100644 --- a/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala +++ b/spark/src/test/scala/org/apache/spark/sql/CometTestBase.scala @@ -70,21 +70,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 = { - val reused = collect(plan) { - case exchange: ReusedExchangeExec if collect(exchange.child)(dppScan).nonEmpty => exchange - } - assert(reused.isEmpty, s"An exchange over a DPP scan was reused by another branch:\n$plan") + 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) } /**