comphead commented on code in PR #6270:
URL: https://github.com/apache/datafusion-comet/pull/6270#discussion_r4265558207
##########
spark/src/main/scala/org/apache/spark/sql/comet/CometNativeScanExec.scala:
##########
@@ -285,14 +285,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]))
Review Comment:
The same block (null check, `originalPlan.copy(partitionFilters =
...filterNot(_.isInstanceOf[DynamicPruningExpression]))`,
`stripped.doCanonicalize()`) is now in both
`CometNativeScanExec.doCanonicalize` and
`CometDeltaNativeScanExec.doCanonicalize`, with a comment explaining it in
each, and the docs tell new scans to copy it. Could it live next to
`filterUnusedDynamicPruningExpressions` in `CometScanUtils`, taking the
nullable `FileSourceScanExec` and returning the canonical copy? Then both scans
and the docs point at one place, and the two comments become one.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -1769,6 +1769,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) {
Review Comment:
The `dppScans` collector plus the assertion that no `ReusedExchangeExec`
sits over a DPP scan is written three times: here, in `CometIcebergNativeSuite`
(the `reusedDpp` block) and in `CometDeltaNativeScanSuite`. Only the scan class
and the way DPP is detected differ. All three suites extend `CometTestBase`,
which already mixes in `AdaptiveSparkPlanHelper`, so a single
`assertNoReusedExchangeOverDppScan(plan)(isDppScan: SparkPlan => Boolean)`
there would hold the assertion and its message once.
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -1769,6 +1769,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)") {
Review Comment:
By its own comment this test passes on main, so it does not pin this fix.
Meanwhile nothing in the PR pins the other half of the rationale, that scans
with identical DPP filters (the same dimension filter in both branches) still
share a stage after the placeholder is kept. `checkDppBranchesKeepTheirRows`
cannot cover it because it requires pruning to 1 and 3 partitions. Could this
test become that case, with `grp = 'a'` in both branches and the Comet
`ReusedExchangeExec` count compared with the Spark plan from
`checkSparkAnswer`, as the Delta SPARK-32509 test does? I have not checked
whether Spark reuses the stage there, but comparing against the Spark plan
gives the right expectation either way (not run).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
