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..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,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), "") } 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..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 @@ -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 diff --git a/docs/source/contributor-guide/adding_a_new_operator.md b/docs/source/contributor-guide/adding_a_new_operator.md index 85946f44fd9..caf45198336 100644 --- a/docs/source/contributor-guide/adding_a_new_operator.md +++ b/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. 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 296fc4e6ccf..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,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 "" 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 652163957af..5c88b66e2e4 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 @@ -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() } } diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala index b00f3353e46..58a266be65d 100644 --- a/spark/src/test/scala/org/apache/comet/CometIcebergNativeSuite.scala +++ b/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. 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 8ec76483027..4ae13863fb5 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala @@ -2401,6 +2401,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 7fc3e66c169..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,12 +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 = + 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) } /**