andygrove opened a new issue, #7208:
URL: https://github.com/apache/datafusion-comet/issues/7208
### Describe the bug
With AQE on (the default), a `REBALANCE(k)` hint followed by `GROUP BY k`
returns duplicate groups when the data is skewed. Spark returns one row per
key. Comet returns several rows for the skewed keys, each with a partial sum
and count. A `PARTITION BY k` window above the rebalance has the same problem,
so `row_number()` restarts inside a key.
Spark's `OptimizeSkewInRebalancePartitions` splits the skewed shuffle
partition (`AQEShuffleRead coalesced and skewed`). In plain Spark that is safe,
because `ValidateRequirements` rejects the split when an operator above needs
`ClusteredDistribution(k)`. Comet's aggregate and window operators declare no
`requiredChildDistribution`, so the check passes. The native `Partial` and
`Final` aggregates then run directly over the split partitions, and each piece
of a key gets its own output row.
- It reproduces with native shuffle, with JVM shuffle, and with
`spark.comet.shuffle.directRead.enabled=false`.
- It goes away with `spark.comet.exec.aggregate.enabled=false` or
`spark.sql.adaptive.optimizeSkewsInRebalancePartitions.enabled=false`.
The 1.1.0 release has the bug too. With the 1.1.0 jar on Spark 3.5.7, the
query below returns 2000 rows; branch-1.2 on Spark 4.1 returns 1357. No release
so far has declared a required distribution on these operators.
### Steps to reproduce
Default configs except the two below, which make the skew split happen on a
small input:
```scala
spark.conf.set("spark.sql.shuffle.partitions", "200")
spark.conf.set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "2k")
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("/tmp/rebal_t1")
spark.read.parquet("/tmp/rebal_t1").createOrReplaceTempView("t1")
spark.sql("""SELECT k, sum(v), count(*)
FROM (SELECT /*+ REBALANCE(k) */ * FROM t1) GROUP BY k""").collect().length
```
- Spark: 901 rows. Key 0 is one row, `(79980000, 4000)`.
- Comet (branch-1.2, Spark 4.1): 1357 rows. Key 0 is four rows: `(17495000,
1000)`, `(19995000, 1000)`, `(34995000, 1000)` and `(7495000, 1000)`.
- Comet 1.1.0 (Spark 3.5.7): 2000 rows.
Calling `.count()` on the DataFrame instead hides the bug, because the extra
aggregate changes the plan.
The plan:
```
CometHashAggregate [k], [Final]
+- CometHashAggregate [k], [Partial]
+- AQEShuffleRead coalesced and skewed
+- ShuffleQueryStage 0
+- CometExchange hashpartitioning(k, 200),
REBALANCE_PARTITIONS_BY_COL, CometNativeShuffle
+- CometNativeScan parquet [k, v]
```
The window case: `SELECT k, row_number() OVER (PARTITION BY k ORDER BY v) rn
FROM (SELECT /*+ REBALANCE(k) */ * FROM t1)` has 1357 rows with `rn = 1` in
Comet and 901 in Spark. The plan is `CometWindowExec` over `CometSort` over the
same `AQEShuffleRead coalesced and skewed`.
### Expected behavior
The same rows as Spark. Comet's operators that need their input clustered
(aggregates with grouping keys, windows with a partition spec, joins) declare
the same `requiredChildDistribution` as the Spark operators they replace, so
AQE does not split partitions under them.
### Additional context
Found in the branch-1.2 release audit (#7204). It is not a 1.2 regression.
Searches for rebalance, skew and `requiredChildDistribution` found no existing
issue.
--
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]
