Fix distribution enforcement for chained aggregates to avoid SanityCheckPlan failures

[kosiew]

Which issue does this PR close?

• Closes Sanity check failed when sort and aggregate on a multi-partitioned table #18989.

Rationale for this change

A panic could occur during SanityCheckPlan when optimizing / executing queries that chain multiple aggregations over a multi-partitioned input (often empty tables), e.g. Sort -> Aggregate(ts, region) -> Sort -> Aggregate(ts).

In these cases the first aggregation can produce a hash partitioning that is a superset of the second aggregation’s required distribution keys. The optimizer previously treated some of these situations as not needing a hash repartition (especially when the input was empty / stats suggested repartition wasn’t beneficial), which could leave the plan violating the downstream operator’s distribution requirement and trigger SanityCheckPlan failures.

This PR aligns the optimizer’s distribution enforcement decisions with the same partitioning satisfaction logic used by the sanity checker, ensuring that required repartitions are inserted whenever the downstream requirement is not satisfied.

What changes are included in this PR?

• Introduced DistributionSatisfactionResult and a shared distribution_satisfaction(...) helper in enforce_distribution to compute and carry:

  • the required Distribution
  • the child’s output_partitioning
  • the computed PartitioningSatisfaction (Exact / Subset / NotSatisfied)
  • a single, reusable requires_repartition(...) decision function.

• Updated add_hash_on_top and ensure_distribution to use the shared satisfaction result and decision logic, rather than ad-hoc checks.

• Simplified get_repartition_requirement_status by removing alignment logic based on hash_necessary flags and instead deferring to the satisfaction-based repartition decision.

• Updated SanityCheckPlan to reuse the same distribution_satisfaction(...) helper (so optimizer and sanity checker evaluate distribution satisfaction consistently).

• Re-exported PartitioningSatisfaction from datafusion_physical_expr for use in tests and downstream code.

• Added targeted regression tests:

  • Unit tests verifying PartitioningSatisfaction outcomes (Exact / Subset / NotSatisfied) match sanity checker behavior for hash partitioning.
  • A tokio regression test reproducing the chained-aggregate scenario from Sanity check failed when sort and aggregate on a multi-partitioned table #18989 and asserting the optimized plan either inserts a repartition between aggregates or preserves a single-partition stream via SortPreservingMergeExec, and executes without panic.

• Updated sqllogictest expected plans where distribution enforcement now inserts RepartitionExec and/or preserves partitioning through SortExec.

Are these changes tested?

Yes.

• Added new unit tests in datafusion/core/tests/physical_optimizer/enforce_distribution.rs to validate distribution satisfaction behavior and its alignment with SanityCheckPlan.
• Added an async regression test covering the multi-partition empty table + chained aggregations scenario from Sanity check failed when sort and aggregate on a multi-partitioned table #18989.
• Updated sqllogictest outputs to reflect the new physical plans.

Performance sanity checks:

• Ran the following benchmarks and verified no performance regressions:

  • cargo bench --bench topk_aggregate --profile=profiling -- --save-baseline before
  • cargo bench --bench distinct_query_sql --profile=profiling -- --save-baseline before
  • cargo bench --bench sort_preserving_merge --profile=profiling -- --save-baseline before

Are there any user-facing changes?

Potentially.

• Some queries may now include additional RepartitionExec operators (or preserve partitioning through sorts) to correctly satisfy downstream distribution requirements. This can change the displayed physical plan (EXPLAIN output) and may affect performance characteristics in edge cases, but fixes a correctness/stability issue (prevents optimizer-time / execution-time panics).
• No SQL syntax or API behavior changes are intended.

LLM-generated code disclosure

This PR includes LLM-generated code and comments. All LLM-generated content has been manually reviewed and tested.