comphead commented on code in PR #7231:
URL: https://github.com/apache/datafusion-comet/pull/7231#discussion_r4379458281
##########
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:
I don't think these two cases reach the join changes. The `GROUP BY a.k`
puts both aggregates above the join in the same stage, and the join passes its
children's partitioning through, so the aggregate's
`ClusteredDistribution(a.k)` already rejects the split. The cases would pass
with only the aggregate change.
Also, `PREFER_SORTMERGEJOIN=false` alone does not plan a shuffled hash join
here. With `AUTO_BROADCASTJOIN_THRESHOLD=-1`, `canBuildLocalHashMapBySize`
compares `sizeInBytes` with `-1 * numShufflePartitions`, and the two sides of a
self-join are the same size, so `muchSmaller` is false too. Both iterations
plan a sort-merge join, and `CometHashJoinExec` is never reached.
Could the join cases drop the grouping, for example `SELECT a.k, a.v, b.v
FROM $rebalanced a JOIN $rebalanced b ON a.k = b.k`, force the hash join with
`"spark.sql.join.forceApplyShuffledHashJoin" -> "true"` as `CometJoinSuite`
does, and assert the expected Comet join class is in the plan?
`checkSparkAnswer` does not check which operators ran natively, so a fallback
would also pass. I have not run this.
--
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]
