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 = {