sunchao opened a new pull request, #57087: URL: https://github.com/apache/spark/pull/57087
### What changes were proposed in this pull request? This resolves [SPARK-58004](https://issues.apache.org/jira/browse/SPARK-58004). This PR adds `spark.sql.adaptive.coalescePartitions.maxReducerPartitionsPerTask`, a positive integer SQL configuration that bounds how many contiguous original reducer partitions a single `CoalescedPartitionSpec` may span. Its default is `Int.MaxValue`, so existing behavior is unchanged unless users opt in. The coalescing algorithm now enforces both the existing byte-size target and this reducer-partition span. Empty reducer partitions count toward the bound, and neither the in-loop backward merge nor the final small-tail merge may cross it. The limit is applied to both skew-aware and non-skew paths and consistently across all shuffle inputs in a coalesce group. Existing `PartialReducerPartitionSpec` entries remain intact. The existing six-argument `ShufflePartitionsUtil.coalescePartitions` method remains available and delegates to the new bounded overload, preserving source and binary compatibility for its current signature. A remote shuffle-block cap is not included because this AQE planning path does not have reliable block-count information. That may be considered separately if such information becomes available. ### Why are the changes needed? AQE shuffle partition coalescing currently uses post-shuffle bytes as its primary packing bound. With very high initial partition counts and sparse or tiny output, a task can remain below the advisory byte target while spanning many thousands of reducer partitions. This produces high per-task shuffle fan-in and can make AQE coalescing unsuitable for these workloads. The new optional hard bound lets users retain byte-based AQE coalescing while limiting this independent source of task overhead. ### Does this PR introduce _any_ user-facing change? Yes. Users may set, for example: ``` spark.sql.adaptive.coalescePartitions.maxReducerPartitionsPerTask=128 ``` AQE will then ensure that each generated `CoalescedPartitionSpec` spans at most 128 original reducer partitions. The default value is `2147483647`, which preserves existing behavior. ### How was this patch tested? Added focused utility tests for the unbounded default, many tiny partitions, a bound of one, long empty runs, the in-loop backward merge, the final small-tail merge, an oversized individual reducer, multiple shuffle inputs, skew partition specs, and invalid zero/negative configuration values. Added an end-to-end AQE test that verifies the configured bound in the executed shuffle-read plan. Ran: ``` ./build/sbt "sql/testOnly org.apache.spark.sql.execution.ShufflePartitionsUtilSuite" ./build/sbt "sql/testOnly org.apache.spark.sql.execution.CoalesceShufflePartitionsSuite" ./build/sbt "catalyst/scalastyle" "sql/scalastyle" "sql/Test/scalastyle" ``` All 19 `ShufflePartitionsUtilSuite` tests and all 20 `CoalesceShufflePartitionsSuite` tests passed. All three Scala style checks passed with zero findings. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: OpenAI Codex (GPT-5) -- 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]
