REBALANCE followed by GROUP BY or a window returns duplicate groups on skewed data #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:

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.