sunchao commented on code in PR #7231:
URL: https://github.com/apache/datafusion-comet/pull/7231#discussion_r4391668281
##########
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 completed five independent review passes on
`ef4199bd443a4968a8f82f59786d2a3f4fabc094` and agree that the join regression
should isolate and assert each intended Comet join. I found no additional
implementation defect. Spark's and Comet's defaults do not switch the second
iteration to shuffled hash join, and the aggregate requirement can mask either
missing join override.
One practical addition to the suggested test change: could we use a smaller
dedicated skewed fixture when removing the grouping? With the current data,
that ungrouped self-join produces 17,440,000 rows: `4000 * 4000 + 900 * 40 *
40`. `checkSparkAnswer` collects both results, so simply dropping `GROUP BY`
would make this regression unnecessarily expensive. A smaller fixture with
appropriately lowered skew thresholds can still exercise the split.
Please explicitly select and assert the native join class, and verify each
case fails with only that join's distribution delegation removed while the
aggregate fix stays enabled. I am leaving approval pending this existing
coverage issue. The current CI regression passes, but that does not establish
that it protects both join changes.
--
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]
