dwsmith1983 commented on code in PR #6270:
URL: https://github.com/apache/datafusion-comet/pull/6270#discussion_r4266633364
##########
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:
Added `assertNoReusedExchangeOverDppScan(plan)(dppScan)` to `CometTestBase`,
next to `assertExchangeReuseOver`, which already takes a partial function for
the node to look under. The assertion and its message live there once, and each
suite passes its own DPP scan match: `partitionFilters` for
`CometNativeScanExec` and `runtimeFilters` for the Iceberg and Delta scans. The
Delta suite reaches it through `CometDeltaTestBase`, which extends
`CometTestBase`.
--
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]
