andygrove commented on code in PR #7231:
URL: https://github.com/apache/datafusion-comet/pull/7231#discussion_r4395244637
##########
spark/src/test/scala/org/apache/comet/exec/CometExecSuite.scala:
##########
@@ -3994,6 +3994,51 @@ class CometExecSuite extends CometTestBase {
}
}
+ // https://github.com/apache/datafusion-comet/issues/7208
+ test("AQE does not split a skewed rebalance partition under operators that
need clustering") {
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "200",
+ SQLConf.ADVISORY_PARTITION_SIZE_IN_BYTES.key -> "2k",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") {
+ withTempPath { dir =>
+ spark
+ .range(0, 40000, 1, 8)
+ .selectExpr("CASE WHEN id % 10 = 0 THEN 0L ELSE id % 1000 END AS k",
"id AS v")
+ .write
+ .parquet(dir.getCanonicalPath)
+
+ withParquetTable(dir.getCanonicalPath, "rebalance_t1") {
+ def checkNotSplit(query: String): Unit = {
+ val (_, cometPlan) = checkSparkAnswer(query)
+ val skewedReads = collect(cometPlan) {
+ case r: AQEShuffleReadExec if r.hasSkewedPartition => r
+ }
+ assert(skewedReads.isEmpty, cometPlan)
+ }
+
+ val rebalanced = "(SELECT /*+ REBALANCE(k) */ * FROM rebalance_t1)"
+ val window =
+ s"SELECT k, row_number() OVER (PARTITION BY k ORDER BY v) rn FROM
$rebalanced"
+ checkNotSplit(s"SELECT k, sum(v), count(*) FROM $rebalanced GROUP BY
k")
+ // operators between the shuffle and the aggregate pass its
partitioning through
+ checkNotSplit(
+ s"SELECT k, sum(w) FROM (SELECT k, v + 1 AS w FROM $rebalanced
WHERE v > 5) GROUP BY k")
+ checkNotSplit(window)
+ checkNotSplit(s"SELECT * FROM ($window) WHERE rn <= 2")
+ Seq("true", "false").foreach { preferSmj =>
+ withSQLConf(SQLConf.PREFER_SORTMERGEJOIN.key -> preferSmj) {
Review Comment:
Thanks both, you were right: the aggregate above the join rejected the split
by itself, and `PREFER_SORTMERGEJOIN=false` never planned a shuffled hash join
here. Fixed in 0107cbb6d9.
The join cases now drop the `GROUP BY` and join the skewed side to a
one-row-per-key side, so the output is about 40k rows instead of 17M and no
separate fixture was needed. The first case uses the default sort-merge join
and asserts `CometSortMergeJoinExec` is in the plan. The second sets
`spark.sql.join.forceApplyShuffledHashJoin` with sort-merge preferred off, and
asserts `CometHashJoinExec`.
I checked each join on its own by removing
`CometDelegatesRequiredChildDistribution` from one join at a time, with the
aggregate change still in place. Both mutations fail the test, with
`PartitioningCollection requires all of its partitionings have the same
numPartitions` thrown from that join's `outputPartitioning`, which is the split
reaching the join. With the real implementation the test passes.
--
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]
