[SPARK-58004][SQL][FOLLOWUP] Preserve AQE fan-in bound through local shuffle reads

[sunchao]

What changes were proposed in this pull request?

This is a follow-up to #57087.

OptimizeShuffleWithLocalRead now checks the mapper-oriented partition specs it proposes against spark.sql.adaptive.coalescePartitions.maxReducerPartitionsPerTask. If any PartialMapperPartitionSpec spans too many reducer partitions, or any CoalescedMapperPartitionSpec reads too many reducers, the rule preserves the existing bounded shuffle read. Safe local-read rewrites are unchanged.

The patch also adds an AQE regression test that exercises both mapper-oriented local-read spec forms and verifies that the bounded coalesced reads remain in the final plan.

Why are the changes needed?

SPARK-58004 added a hard limit on the number of reducer partitions that AQE can combine into one task. However, the local shuffle reader runs after shuffle partition coalescing and can replace the bounded CoalescedPartitionSpec values with mapper-oriented specs whose reducer spans exceed that limit.

Without this follow-up, enabling local shuffle reads can therefore bypass the new fan-in bound and recreate the excessive per-task fetch pressure that #57087 is intended to prevent.

Does this PR introduce any user-facing change?

Yes, compared with the current unreleased master branch. When a proposed local shuffle read would exceed spark.sql.adaptive.coalescePartitions.maxReducerPartitionsPerTask, Spark now keeps the existing coalesced shuffle read instead. Released Spark versions are not affected.

How was this patch tested?

Added and ran the focused regression test:

build/sbt 'sql/testOnly org.apache.spark.sql.execution.adaptive.AdaptiveQueryExecSuite -- -z "Preserve the reducer partition limit through local shuffle read optimization"'

Ran the adjacent positive-path local shuffle reader test:

build/sbt 'sql/testOnly org.apache.spark.sql.execution.adaptive.AdaptiveQueryExecSuite -- -z "Reuse the parallelism of coalesced shuffle in local shuffle read"'

Ran the coalescing-disabled finite-cap regression:

build/sbt 'sql/testOnly org.apache.spark.sql.execution.adaptive.AdaptiveQueryExecSuite -- -z "Reuse the default parallelism in local shuffle read"'

Ran style checks:

build/sbt sql/scalastyle sql/Test/scalastyle

Was this patch authored or co-authored using generative AI tooling?

Generated-by: OpenAI Codex (GPT-5)

[sunchao]

cc @dongjoon-hyun @viirya a follow-up of #57087

[dongjoon-hyun]

This is not a regression of SPARK-58004 itself because the fan-in bound is correctly preserved. However, the check looks too stricter than necessary, so users who set the config lose the local shuffle read optimization almost entirely.

The two spec types are not equivalent per unit of reducer span:

A CoalescedPartitionSpec spanning k reducers fetches k × numMappers blocks over the network — this is what SPARK-58004 bounds.

A PartialMapperPartitionSpec spanning k reducers reads only k blocks from a single mapper, typically locally.
Yet the PR compares both spans against the same limit. Since a bound-constrained coalesced read has spans ≈ max, the equivalent local-read spec spans ≈ max × numMappers, which always fails the check when numMappers ≥ 2. So the rewrite is reverted even when the local-read task would read the same number of blocks as the retained coalesced task — with less network traffic.

In short, setting the config effectively seems to disable local shuffle reads on those stages, forfeiting their benefit in cases where the fan-in pressure the config guards against would not actually recur.

Maybe, compare PartialMapperPartitionSpec spans against max × numMappers (equivalent per-task block count), or regenerate specs with higher parallelism instead of reverting.

[viirya]

Thanks @dongjoon-hyun for the detailed analysis. I agree with the concern, and I think it can be pushed one step further, which changes what the right fix looks like.

The local read rewrite is per-task block-count preserving by construction. It reuses the coalesced read's parallelism P, so each rewritten task covers numMappers * numReducers / P (mapper, reducer) cells — exactly the block count of the coalesced task it replaces. Only the shape changes: from "all mappers x narrow reducer range, fetched over the network" to "one mapper x wide reducer range, read locally". Since the overhead SPARK-58004 targets (per-task metadata and fetch fan-in) scales with block count, the rewrite never actually worsens it; only the literal reducer span grows, by a factor of ~numMappers.

A concrete example with numMappers = 10, numReducers = 1000, maxReducerPartitionsPerTask = 100:

	tasks	span/task	blocks/task	read
coalesced read (kept by this PR)	10	100	1000	remote
local read it rejects	10	1000	1000	local

The rejected plan is strictly better, which is why the check as written effectively disables local reads whenever the bound was binding.

This also means comparing PartialMapperPartitionSpec spans against max * numMappers would (modulo integer-division rounding) never trigger — it degenerates to no check at all. The same holds for the CoalescedMapperPartitionSpec branch, which is likewise block-neutral. So I think the real decision is what contract the config makes, and the two candidate fixes follow from it:

• If the config is a proxy for per-task fan-in overhead (which matches the SPARK-58004 motivation), the cleanest fix is to drop the guard entirely and instead document in the config doc that the limit constrains coalesced specs, and that a later local read rewrite may replace them with mapper-oriented specs whose literal reducer span exceeds the limit while keeping per-task block counts unchanged. Otherwise users who inspect the final plan will reasonably report the config as broken.
• If the literal span bound must hold for every task, then regenerating the specs with higher parallelism (your second suggestion) is the way to go: P' ~= numMappers * ceil(numReducers / max) gives tasks that each read at most max blocks from one mapper, locally. In the example above that is 100 tasks x 100 blocks, satisfying the bound while keeping the local read. The cost is a task count multiplied by numMappers, so it should probably fall back to keeping the coalesced read when P' becomes excessive (e.g. numReducers = 100k, max = 10, numMappers = 100 would yield 1M tasks), making the current PR's revert the last resort rather than the only behavior.

My preference is the first option, since I could not find a cost in the local read path that scales with the literal span rather than the block count — but either way the current all-or-nothing revert seems to be the least favorable choice for users who set this config.