Allow AQE stage rewrites to require exact downstream consumer task counts

[sesteves]

Summary

Allow an AQE stage rewrite to require an exact number of downstream consumer tasks, rather than only providing an advisory task count.

This is needed when a rewrite changes partitioning based on an operator's estimated memory requirements. If AQE subsequently coalesces the consumer count below that requirement, each task receives more partitions than the rewrite planned for and the memory-safety estimate no longer holds.

Motivation

The existing AQE sampler is effective at sizing work from observed row and byte throughput, but that is not enough for operators whose retained memory depends on other properties, such as the number and width of distinct grouping keys and aggregate accumulator state.

An application-specific dynamic-stage handler can apply that policy. For example, a handler can estimate aggregate memory, rewrite the producer to use an appropriate number of hash partitions, and determine how many consumers are required. DataFusion Distributed still needs a way for the handler to communicate that the resulting consumer count is a correctness or resource-safety requirement rather than a preference.

This keeps operator-specific sizing policy outside DataFusion Distributed while letting its scheduler enforce the result.

Proposed API direction

Building on the dynamic-stage extension point in #570, a response could expose an API such as:

DynamicStageBuiltEventResponse::new(rewritten_plan)
    .with_exact_consumer_tasks(tasks)

The exact API shape is open for discussion. The important part is distinguishing an exact requirement from an advisory desired task count.

Semantics

• An exact consumer count takes precedence over advisory AQE task-count decisions.
• The scheduler must not silently coalesce below the exact count.
• The producer's hash partitioning and the downstream consumer task count must remain consistent.
• Conflicting exact requirements should be rejected rather than resolved arbitrarily.
• Global hard limits must still be respected; if the exact requirement cannot be satisfied, planning should fail explicitly.
• "Exact" describes the task count when the stage is consumed. It does not require every prepared task to execute when a query short-circuits. Skipped or cancelled tasks must still reach a terminal state and must not block metrics finalization, as discussed in Metrics finalization can hang on planned but never executed tasks with AQE or static planning #739.

Relationship to sampling

This issue does not propose embedding aggregate cardinality estimation in DataFusion Distributed. A caller could use catalog statistics or a future post-sampling statistics hook to estimate grouping-key cardinality, then use the exact consumer-count contract to apply its sizing decision safely.

Related work

• Add dynamic stage built events #570 adds the dynamic-stage extension point that this could build on.
• [Epic][part 2] Adaptive Query Execution, runtime plan optimization #522 tracks runtime query optimization through AQE.
• Leverage memory and network cost for dynamic task count decisions #520 discusses memory- and network-aware task sizing.
• feat: Reduce producer partition fan-out in NetworkShuffleExec #717 aligns shuffle partitioning between producers and consumers.
• Metrics finalization can hang on planned but never executed tasks with AQE or static planning #739 tracks AQE tasks that are prepared or sampled but never executed and can block metrics finalization.

I have a working prototype of this behavior and used it to complete a high-cardinality GROUP BY query that otherwise exhausted worker memory. Once the API direction is agreed, I can follow up with a PR.