You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "zhuqi-lucas (via GitHub)" <gi...@apache.org> on 2026/05/01 15:29:14 UTC
[PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
zhuqi-lucas opened a new pull request, #21976:
URL: https://github.com/apache/datafusion/pull/21976
## Summary
This PR adds a new `EnsureRequirements` optimizer rule that merges `EnforceDistribution` and `EnforceSorting` into a single coordinated rule, and fixes `pushdown_sorts` to be distribution-aware, making the composition idempotent.
**Epic**: #21973
## Problem
`EnforceDistribution` and `EnforceSorting` run as separate rules, but sorting and distribution are coupled through `SortExec.preserve_partitioning`. `pushdown_sorts` (inside `EnforceSorting`) removes sorts and reinserts them via `add_sort_above()`, which has no distribution context. This creates `SortExec(preserve_partitioning=true)` without `SortPreservingMergeExec`, violating `SinglePartition` requirements from operators like `GlobalLimitExec`.
This caused production 502 errors for specific API users whose entitlements triggered 32-partition `StorageExec` with `in_list_partitioning`.
DataFusion is the only major query engine that separates distribution and sorting enforcement — Spark (`EnsureRequirements`), Presto/Trino (`AddExchanges`) handle both in a single rule.
## Changes
### 1. `EnsureRequirements` rule (`ensure_requirements/mod.rs` + `new_tests.rs`)
- Composes `EnforceDistribution::optimize()` + `EnforceSorting::optimize()` in a single rule
- 41 comprehensive tests including:
- Multi-partition sort + limit (2, 4, 8, 16, 32, 64 partitions)
- Union with mixed partition counts
- Projection, Filter, Repartition, Coalesce over multi-partition
- Hash joins, aggregates, window functions
- Fetch/skip preservation (#14150 regression)
- PR #53/#54 regression scenarios
- **Idempotency tests**: 1-64 partition sweep, triple-optimize convergence, all bug-triggering topologies
### 2. Distribution-aware `pushdown_sorts` (`sort_pushdown.rs`)
- Add `distribution_requirement: Distribution` field to `ParentRequirements`
- Initialize from `plan.required_input_distribution()` in `assign_initial_requirements`
- Propagate through recursion with `stronger_distribution()` helper
- Reset to `UnspecifiedDistribution` below partition-merging nodes (SPM, single-partition outputs)
- Switch both `add_sort_above` call sites to `add_sort_above_with_distribution`
### 3. `add_sort_above_with_distribution` (`utils.rs`)
- New utility: inserts `SortPreservingMergeExec` when parent requires `SinglePartition` and input has multiple partitions
## Testing
| Test Suite | Result |
|-----------|--------|
| EnsureRequirements (new) | **41 passed** |
| enforce_sorting (existing) | **124 passed, 0 failed** |
| enforce_distribution (existing) | **66 passed, 0 failed** |
| SLT (464 files) | **1 pre-existing failure only** |
| **Total** | **695 tests, 0 regressions** |
Key idempotency verifications:
- All partition counts 1-64 ✓
- PR #53 scenario (OutputRequirementExec + multi-partition) ✓
- PR #54 scenario (ProjectionExec + multi-partition) ✓
- SPM → Sort → multi-partition ✓
- Sort → Aggregate pattern (#18989) ✓
- Triple optimization convergence ✓
## Next Steps (future PRs)
- [ ] Replace default optimizer chain with `EnsureRequirements` (currently opt-in)
- [ ] Gradually eliminate `pushdown_sorts` by implementing its optimizations in the bottom-up `ensure_sorting` pass
- [ ] Single-pass architecture (like Spark's `EnsureRequirements`)
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4384563140
cc @alamb @adriangb @xudong963 @Dandandan @2010YOUY01
We've hit several production bugs caused by EnforceDistribution and EnforceSorting running as separate rules — they're coupled through SortExec.preserve_partitioning, which makes the composition non-idempotent.
Other engines (Spark's EnsureRequirements, Trino's AddExchanges) handle both in a single rule, which sidesteps this entire bug class.
I initially attempted a true single-pass rewrite, but the interactions between distribution and sorting requirements turned out to be deeper than anticipated.
This PR takes a more conservative approach — composing the existing rules into one with the necessary correctness fixes. A clean single-pass design is left as a follow-up once the foundation is in place.
If a default change feels too risky for a first step, I'm happy to gate it behind a config flag (e.g. `datafusion.optimizer.use_ensure_requirements`) so users can opt in and we can validate equivalence in CI before flipping the default.
Would appreciate any thoughts on the overall direction.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521086953
run benchmarks
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521535635
Ok, let's do it!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4421083625
> Thanks @alamb. Reading again I see two separate asks:
>
> * The inline on `new_tests.rs` (move it to `datafusion/physical-optimizer/tests/ensure_requirements.rs`) is something I can do inside this PR right away — will push that next.
> * The bigger one (move the existing `core/tests/physical_optimizer/{enforce_distribution.rs, enforce_sorting.rs, ...}` integration tests, ~9.8k lines, to live alongside the optimizer crate) — I'll do that as a separate prep PR after 54 ships, then rebase this on top so the porting diff is clean.
>
> Will ping you when the inline fix lands and again when the prep PR is up.
Thanks @zhuqi-lucas
FYI I think we can reorganize the tests before 54 as that has a low chance of introducing bugs (so doesn't need as long a "bake" time)
In terms of this PR, my plan for reviewing it is basically
1. Review what, if any tests have changed (will spend most effort here)
2. Run sql planning performance benchmarks to ensure no regressions
3. Review the code for anything that stands out
So TLDR is I plan to spend most of my time doing verification rather than code review
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "xudong963 (via GitHub)" <gi...@apache.org>.
xudong963 commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4418139727
> What I suggest we do is:
>
> 1. Prepare the tests
> First, make a separate PR to move all existing tests for enforce sorting and enforce distribution somewhere outside the existing modules. By having the tests in a separate module, it will be much easier to understand what, if anything, changed when you replace EnforceSorting/EnforceDistribtuion with EnforceRequirements
+1 for this
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4418242035
> > What I suggest we do is:
> >
> > 1. Prepare the tests
> > First, make a separate PR to move all existing tests for enforce sorting and enforce distribution somewhere outside the existing modules. By having the tests in a separate module, it will be much easier to understand what, if anything, changed when you replace EnforceSorting/EnforceDistribtuion with EnforceRequirements
>
> +1 for this
Thanks @xudong963 , will go this way.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521351053
I don't really expect any performance difference, so rerunning the benchmarks to see if I can reproduce anything
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4525048305
The GitHub is flaky for merging PR now.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4505952961
> The `release-54` branch has been made so I think we can merge this PR when ready
Thank you @alamb , i addressed comments now, please let me know if i am missing anything!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279072061
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
Review Comment:
Thanks @alamb, this is a good point, and addressed in the latest PR — tests moved core/tests/physical_optimizer/ensure_requirements.rs, Tier 1 (4 tests) deleted, Tier 2 folded into one combined test, Tier 3 trimmed from 10x sweep to assert_idempotent. Kept the pairs you flagged (#14150, projection vs OutputRequirement).
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521347657
run benchmark clickbench_partitioned
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3214871632
##########
datafusion/physical-optimizer/src/optimizer.rs:
##########
@@ -170,18 +169,12 @@ impl PhysicalOptimizer {
// those are handled by the later `FilterPushdown` rule.
// See `FilterPushdownPhase` for more details.
Arc::new(FilterPushdown::new()),
- // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution
- // requirements. Please make sure that the whole plan tree is determined before this rule.
- // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at
- // least one of the operators in the plan benefits from increased parallelism.
- Arc::new(EnforceDistribution::new()),
- // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule
+ // EnsureRequirements: merged EnforceDistribution + EnforceSorting into a
+ // single idempotent rule with distribution-aware pushdown_sorts.
+ // See https://github.com/apache/datafusion/issues/21973
Review Comment:
Thanks @2010YOUY01 for review and good suggestion!
Will apply the suggestion in the next push (per alamb's plan I'll land a test-extraction PR first and rebase this on top).
On "distribution-aware pushdown_sorts": the old pushdown_sorts didn't propagate the parent's required_input_distribution, so a SortExec could get pushed below something requiring SinglePartition (e.g. under GlobalLimitExec) — giving an invalid plan because the sort then runs on multiple partitions. The fix propagates the requirement and inserts a SortPreservingMergeExec where needed. Will link to #21973 in the comment block when I update.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4524539644
> Looks like there is a logical conflict -- I just merged up from main and the compile now fails
Thanks again @alamb for review, i fixed the conflicts and merging it now!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279097357
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + `GlobalLimitExec` must produce a valid plan.
+ /// Regression for the `SanityCheckPlan` failure that motivated this PR:
+ /// `pushdown_sorts` setting `preserve_partitioning=true` on multi-partition
+ /// input without inserting `SortPreservingMergeExec` violated the
+ /// `SinglePartition` requirement from `GlobalLimitExec`.
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
Review Comment:
I tried idempotency to use insta, but it seems hard for this type test to use it.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279192370
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + `GlobalLimitExec` must produce a valid plan.
+ /// Regression for the `SanityCheckPlan` failure that motivated this PR:
+ /// `pushdown_sorts` setting `preserve_partitioning=true` on multi-partition
+ /// input without inserting `SortPreservingMergeExec` violated the
+ /// `SinglePartition` requirement from `GlobalLimitExec`.
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Idempotency: multi-partition sort + limit
+ #[test]
+ fn test_idempotent_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_idempotent(limit);
+ }
+
+ /// Idempotency: union with mixed partitions
+ #[test]
+ fn test_idempotent_union_mixed_partitions() {
+ let live = Arc::new(MockMultiPartitionExec::new(8));
+ let hist = Arc::new(MockMultiPartitionExec::new(1));
+ let union = UnionExec::try_new(vec![live as _, hist as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(5)));
+
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Projection + multi-partition tests (pushdown_sorts trigger path)
+ // ========================================================================
+
+ /// ProjectionExec over multi-partition + sort DESC + limit.
+ /// This is the topology where pushdown_sorts pushes sort through projection
+ /// onto the multi-partition source. The optimizer must still produce a valid plan.
+ #[test]
+ fn test_projection_over_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection =
+ Arc::new(ProjectionExec::try_new(proj_exprs, source as _).unwrap());
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, projection));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(Arc::clone(&limit));
+ assert!(
+ result.is_ok(),
+ "SanityCheckPlan failed for projection over multi-partition: {:?}",
+ result.err()
+ );
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Single partition tests (no unnecessary operators)
+ // ========================================================================
+
+ /// Single partition source + sort + limit should NOT add SortPreservingMergeExec.
+ #[test]
+ fn test_single_partition_no_unnecessary_spm() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ let optimized = optimize_and_sanity_check(limit).unwrap();
+ let plan_str = datafusion_physical_plan::displayable(optimized.as_ref())
+ .indent(true)
+ .to_string();
+
+ // Single partition should not have SortPreservingMergeExec
+ assert!(
+ !plan_str.contains("SortPreservingMergeExec"),
+ "Unnecessary SortPreservingMergeExec for single partition:\n{plan_str}"
+ );
Review Comment:
Done in e8e0b9f15 — added an assert_ensure_requirements_plan! macro and converted the 11 plan-shape tests in this file to inline insta snapshots. Property tests (idempotency, partition-count sweeps, fetch preservation) kept assert! since insta doesn't fit invariant-style checks. Side benefit: the snapshots make the actual plan shape visible at the test site — test_repartition_sort_limit_idempotent for instance now shows that the optimizer eliminates the RoundRobin repartition over a 1-partition source, which the old is_ok() check hid.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521115210
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521086953)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4521086953-289-pvkgk 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing ensure-requirements (504fb08b720a182fd36ec80ca2bf958a81193225) to 0da8961 (merge-base) [diff](https://github.com/apache/datafusion/compare/0da89616e43299546ed8405cfa900d0f97cb2515..504fb08b720a182fd36ec80ca2bf958a81193225) using: tpcds
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4436562586
Thanks for spelling that out @alamb — that's much cleaner than what I had. Closing #22117 and switching to the in-place approach here. Plan:
1. Delete the disconnected `EnforceDistribution` / `EnforceSorting` modules (they're not registered in the pipeline anymore — just dead code at this point).
2. Update the existing tests in `core/tests/physical_optimizer/{enforce_distribution,enforce_sorting,enforce_sorting_monotonicity,replace_with_order_preserving_variants}.rs` to call `EnsureRequirements::new()`. Keeps the diff focused on the rule swap, exactly as you said.
3. Drop `ensure_requirements/new_tests.rs` and the inline tests in `ensure_requirements/mod.rs` that duplicate the existing coverage.
One question on #3: a handful of cases in `new_tests.rs` cover behavior that's genuinely new — idempotent re-application of pushdown_sorts, a couple of multi-rule integration cases that didn't exist before. Worth keeping those (either where they are or folded into the relevant existing files), or do you want everything dropped and the new behavior covered later in a follow-up?
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521155526
🤖 Criterion benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521087701)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4521087701-291-mkqlr 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing ensure-requirements (504fb08b720a182fd36ec80ca2bf958a81193225) to 0da8961 (merge-base) [diff](https://github.com/apache/datafusion/compare/0da89616e43299546ed8405cfa900d0f97cb2515..504fb08b720a182fd36ec80ca2bf958a81193225)
BENCH_NAME=sql_planner
BENCH_COMMAND=cargo bench --features=parquet --bench sql_planner
BENCH_FILTER=
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521380069
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521347657)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4521347657-292-jkjxc 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing ensure-requirements (504fb08b720a182fd36ec80ca2bf958a81193225) to 0da8961 (merge-base) [diff](https://github.com/apache/datafusion/compare/0da89616e43299546ed8405cfa900d0f97cb2515..504fb08b720a182fd36ec80ca2bf958a81193225) using: clickbench_partitioned
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4502874251
The `release-54` branch has been made so I think we can merge this PR when ready
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279111431
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + `GlobalLimitExec` must produce a valid plan.
+ /// Regression for the `SanityCheckPlan` failure that motivated this PR:
+ /// `pushdown_sorts` setting `preserve_partitioning=true` on multi-partition
+ /// input without inserting `SortPreservingMergeExec` violated the
+ /// `SinglePartition` requirement from `GlobalLimitExec`.
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Idempotency: multi-partition sort + limit
+ #[test]
+ fn test_idempotent_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_idempotent(limit);
+ }
+
+ /// Idempotency: union with mixed partitions
+ #[test]
+ fn test_idempotent_union_mixed_partitions() {
+ let live = Arc::new(MockMultiPartitionExec::new(8));
+ let hist = Arc::new(MockMultiPartitionExec::new(1));
+ let union = UnionExec::try_new(vec![live as _, hist as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(5)));
+
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Projection + multi-partition tests (pushdown_sorts trigger path)
+ // ========================================================================
+
+ /// ProjectionExec over multi-partition + sort DESC + limit.
+ /// This is the topology where pushdown_sorts pushes sort through projection
+ /// onto the multi-partition source. The optimizer must still produce a valid plan.
+ #[test]
+ fn test_projection_over_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection =
+ Arc::new(ProjectionExec::try_new(proj_exprs, source as _).unwrap());
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, projection));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(Arc::clone(&limit));
+ assert!(
Review Comment:
Good catch — switched to unwrap() / expect().
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4432401050
Follow on comment
- https://github.com/apache/datafusion/pull/22117#issuecomment-4432397769
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3243911160
##########
datafusion/physical-optimizer/src/enforce_distribution.rs:
##########
@@ -970,7 +813,10 @@ fn preserving_order_enables_streaming(
///
/// Updated node with an execution plan, where the desired single distribution
/// requirement is satisfied.
-fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
+fn add_merge_on_top(
Review Comment:
It seems like as a follow on some of this code coudl be consolidated into enforce_requirements
##########
datafusion/core/src/optimizer_rule_reference.md:
##########
@@ -67,27 +67,26 @@ Rule order matters. The default pipeline may change between releases.
The same rule name may appear more than once when the default pipeline runs it
in multiple phases.
-| order | rule | phase | summary |
-| ----- | ------------------------------ | ----------------------- | ------------------------------------------------------------------------------------------------------------ |
-| 1 | `OutputRequirements` | add phase | Adds helper nodes so output requirements survive later physical rewrites. |
-| 2 | `aggregate_statistics` | - | Uses exact source statistics to answer some aggregates without scanning data. |
-| 3 | `join_selection` | - | Chooses join implementation, build side, and partition mode from statistics and stream properties. |
-| 4 | `LimitedDistinctAggregation` | - | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed. |
-| 5 | `FilterPushdown` | pre-optimization phase | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. |
-| 6 | `EnforceDistribution` | - | Adds repartitioning only where needed to satisfy physical distribution requirements. |
-| 7 | `CombinePartialFinalAggregate` | - | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant. |
-| 8 | `EnforceSorting` | - | Adds or removes local sorts to satisfy required input orderings. |
-| 9 | `OptimizeAggregateOrder` | - | Updates aggregate expressions to use the best ordering once sort requirements are known. |
-| 10 | `WindowTopN` | - | Replaces eligible row-number window and filter patterns with per-partition TopK execution. |
-| 11 | `ProjectionPushdown` | early pass | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure. |
-| 12 | `OutputRequirements` | remove phase | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done. |
-| 13 | `LimitAggregation` | - | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets. |
-| 14 | `LimitPushPastWindows` | - | Pushes fetch limits through bounded window operators when doing so keeps the result correct. |
-| 15 | `HashJoinBuffering` | - | Adds buffering on the probe side of hash joins so probing can start before build completion. |
-| 16 | `LimitPushdown` | - | Moves physical limits into child operators or fetch-enabled variants to cut data early. |
-| 17 | `TopKRepartition` | - | Pushes TopK below hash repartition when the partition key is a prefix of the sort key. |
-| 18 | `ProjectionPushdown` | late pass | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities. |
-| 19 | `PushdownSort` | - | Pushes sort requirements into data sources that can already return sorted output. |
-| 20 | `EnsureCooperative` | - | Wraps non-cooperative plan parts so long-running tasks yield fairly. |
-| 21 | `FilterPushdown(Post)` | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving. |
-| 22 | `SanityCheckPlan` | - | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. |
+| order | rule | phase | summary |
+| ----- | ------------------------------ | ----------------------- | -------------------------------------------------------------------------------------------------------------------------------- |
+| 1 | `OutputRequirements` | add phase | Adds helper nodes so output requirements survive later physical rewrites. |
+| 2 | `aggregate_statistics` | - | Uses exact source statistics to answer some aggregates without scanning data. |
+| 3 | `join_selection` | - | Chooses join implementation, build side, and partition mode from statistics and stream properties. |
+| 4 | `LimitedDistinctAggregation` | - | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed. |
+| 5 | `FilterPushdown` | pre-optimization phase | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. |
+| 6 | `EnsureRequirements` | - | Enforces both distribution and sorting requirements in a single idempotent rule (replaces EnforceDistribution + EnforceSorting). |
Review Comment:
we probably don't have to mention it "replaces" as when someone reads this the code will no longer have refrences to EnforceDistribution + EnforceSorting
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
Review Comment:
Given this rule calls transform_up / transform_down several times I am not sure calling it a "single pass" is quite accurate. That being said, it seems like a very nice formulation and maybe as a follow on PR (s) we can consolidate the code more / combine some of the passes.
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
Review Comment:
There are over 1600 lines of tests in this module and I am finding it very hard to review. I wonder:
1. Why add tests in this module (rather than the existing pattern of `core/tests/physical_optimizer/` (where you can use actual ExecutionPlans rather than mocking everything out)
2. Are all these tests necessary? Since this seems mostly like an internal refactor and the exsiting code is covered by the existing tests, I wonder if there was a big testing gap you found? Or maybe the tests were intermediate work product that may not be necessary to commit to the repo?
I asked claude what was redundant and it said some non trival number are (I didn't double check but its answers are here):
<details>
Tier 1 — Clearly redundant (delete)
These use the identical plan shape SortExec(DESC) → MockMultiPartitionExec(N) → GlobalLimit(fetch) and the identical assertion as another test that strictly subsumes them:
1. test_idempotent_various_partition_counts (mod.rs:691, sweeps {2,4,8,16,32,64}) — strictly subsumed by test_idempotent_all_partition_counts_1_to_64 (mod.rs:1188, sweeps 1..=64). Same topology, same
assertion. Delete.
2. test_idempotent_multi_partition_sort_limit (mod.rs:477, N=16) — subsumed by test_idempotent_all_partition_counts_1_to_64 at N=16. Delete.
3. test_multi_partition_sort_limit_sanity_check (mod.rs:432, N=32) — subsumed by test_various_partition_counts_all_pass_sanity_check (mod.rs:663, sweeps {2,4,8,16,32,64}) at N=32. Same topology, same
optimize_and_sanity_check assertion. Delete.
4. test_triple_optimize_stable (mod.rs:1208, N=32, three passes) — same topology again. assert_idempotent proves f(f(x)) == f(x); for a deterministic optimizer that implies f^N(x) == f(x) for all N, so adding
a third pass is mathematically redundant. The partition sweep already covers N=32. Delete.
Tier 2 — Substantial overlap (consider combining)
5. test_fetch_preserved_across_passes (mod.rs:629) — checks fetch=5 appears after 1 pass on Sort → multi(4) → GlobalLimit(5). test_issue_14150_fetch_survives_multiple_passes (mod.rs:1259) checks the same
property after 3 passes on a richer plan (Repartition → Sort(fetch=5) → GlobalLimit(5)) and is more representative of the actual bug. The single-pass test adds little once the multi-pass one exists.
6. test_no_extra_spm_when_already_optimal (mod.rs:1006) and test_idempotent_spm_sort_multi_partition (mod.rs:1114) — identical input plan (SPM → Sort(preserve=true) → multi(10) → GlobalLimit(21)). First
asserts SPM count ≤ 1; second asserts 2-pass equality. Since the input already has exactly one SPM, idempotency on this input implies no extra SPM is added. Either fold the SPM-count assertion into the
idempotency test (so one test covers both), or drop one.
Tier 3 — Mathematically over-strong but topology unique (trim, don't delete)
7. test_idempotent_10x_complex (mod.rs:1763) — the union + projection + sort topology is unique among the new tests, so this should stay. But 10 passes is unnecessary — a 2-pass assert_idempotent call on the
same plan gives the same guarantee. Consider just reducing it to use assert_idempotent.
Tests that look similar but are NOT redundant
- test_issue_14150_fetch_survives_multiple_passes vs. test_issue_14150_fetch_survives_with_input_spm — both target #14150 but with deliberately different input shapes (the latter is the "sharper reproduce"
with an SPM(fetch) already in the input — the historical bug shape). Keep both.
- test_projection_over_multi_partition_sort_limit vs. test_idempotent_projection_over_multi_partition_with_single_partition_requirement — first uses GlobalLimit, second uses
OutputRequirementExec(SinglePartition). Different requirement source, different code path. Keep both.
- test_output_requirement_single_partition_over_multi_partition_source vs. test_idempotent_output_requirement_single_partition — same topology but one checks SanityCheckPlan, the other checks idempotency.
Mild overlap; could be merged into a single test asserting both, but each captures a distinct property.
</details>
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + `GlobalLimitExec` must produce a valid plan.
+ /// Regression for the `SanityCheckPlan` failure that motivated this PR:
+ /// `pushdown_sorts` setting `preserve_partitioning=true` on multi-partition
+ /// input without inserting `SortPreservingMergeExec` violated the
+ /// `SinglePartition` requirement from `GlobalLimitExec`.
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Idempotency: multi-partition sort + limit
+ #[test]
+ fn test_idempotent_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_idempotent(limit);
+ }
+
+ /// Idempotency: union with mixed partitions
+ #[test]
+ fn test_idempotent_union_mixed_partitions() {
+ let live = Arc::new(MockMultiPartitionExec::new(8));
+ let hist = Arc::new(MockMultiPartitionExec::new(1));
+ let union = UnionExec::try_new(vec![live as _, hist as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(5)));
+
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Projection + multi-partition tests (pushdown_sorts trigger path)
+ // ========================================================================
+
+ /// ProjectionExec over multi-partition + sort DESC + limit.
+ /// This is the topology where pushdown_sorts pushes sort through projection
+ /// onto the multi-partition source. The optimizer must still produce a valid plan.
+ #[test]
+ fn test_projection_over_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection =
+ Arc::new(ProjectionExec::try_new(proj_exprs, source as _).unwrap());
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, projection));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(Arc::clone(&limit));
+ assert!(
Review Comment:
I think this could be more concise using `unwrap()` (or `expect(...)` if youw ant the messages
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -1672,17 +1669,16 @@ fn multi_smj_joins() -> Result<()> {
// TODO(wiedld): show different test result if enforce distribution first.
assert_plan!(plan_sort, @r"
SortMergeJoinExec: join_type=..., on=[(a@0, c@2)]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[a@0 ASC], preserve_partitioning=[false]
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC
Review Comment:
this seems like a better plan to me (rather than sorting in a single partition is is sorting in parallel across all partitions)
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1804 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod new_tests;
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::enforce_distribution::EnforceDistribution;
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + GlobalLimitExec must produce valid plan.
+ /// Regression test for UXX0/HRC 502s (April 2026).
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Idempotency: multi-partition sort + limit
+ #[test]
+ fn test_idempotent_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_idempotent(limit);
+ }
+
+ /// Idempotency: union with mixed partitions
+ #[test]
+ fn test_idempotent_union_mixed_partitions() {
+ let live = Arc::new(MockMultiPartitionExec::new(8));
+ let hist = Arc::new(MockMultiPartitionExec::new(1));
+ let union = UnionExec::try_new(vec![live as _, hist as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(5)));
+
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Projection + multi-partition tests (pushdown_sorts trigger path)
+ // ========================================================================
+
+ /// ProjectionExec over multi-partition + sort DESC + limit.
+ /// This is the topology where pushdown_sorts pushes sort through projection
+ /// onto the multi-partition source. The optimizer must still produce a valid plan.
+ #[test]
+ fn test_projection_over_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection =
+ Arc::new(ProjectionExec::try_new(proj_exprs, source as _).unwrap());
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, projection));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(Arc::clone(&limit));
+ assert!(
+ result.is_ok(),
+ "SanityCheckPlan failed for projection over multi-partition: {:?}",
+ result.err()
+ );
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Single partition tests (no unnecessary operators)
+ // ========================================================================
+
+ /// Single partition source + sort + limit should NOT add SortPreservingMergeExec.
+ #[test]
+ fn test_single_partition_no_unnecessary_spm() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ let optimized = optimize_and_sanity_check(limit).unwrap();
+ let plan_str = datafusion_physical_plan::displayable(optimized.as_ref())
+ .indent(true)
+ .to_string();
+
+ // Single partition should not have SortPreservingMergeExec
+ assert!(
+ !plan_str.contains("SortPreservingMergeExec"),
+ "Unnecessary SortPreservingMergeExec for single partition:\n{plan_str}"
+ );
+ }
+
+ /// Source already has correct ordering → should not add SortExec.
+ #[test]
+ fn test_sort_already_satisfied_no_extra_sort() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+
+ // Sort ASC matches MockMultiPartitionExec's output ordering (a ASC)
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ let optimized = optimize_and_sanity_check(limit).unwrap();
+ let plan_str = datafusion_physical_plan::displayable(optimized.as_ref())
+ .indent(true)
+ .to_string();
+
+ // Sort should be eliminated since source already satisfies ordering
+ // The plan should just be limit + source (or limit + local limit + source)
+ assert!(
+ !plan_str.contains("SortExec: expr=[a@0 ASC"),
+ "Unnecessary SortExec when ordering already satisfied:\n{plan_str}"
+ );
+ }
+
+ // ========================================================================
Review Comment:
(see comment above about reundancy)
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
Review Comment:
as a follow on PR it might be nuce to move the enforce_sorting and enforce_distribution to submodules of ensure_requirements (just move the code / helpers around, as they are (now) internal implementation details)
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -695,17 +695,14 @@ impl TestConfig {
// TODO: End state payloads will be checked here.
}
- for run in optimizers_to_run {
- optimized = match run {
- Run::Distribution => {
- let optimizer = EnforceDistribution::new();
- optimizer.optimize(optimized, &self.config)?
- }
- Run::Sorting => {
- let optimizer = EnforceSorting::new();
- optimizer.optimize(optimized, &self.config)?
- }
- };
+ // With `EnsureRequirements`, distribution and sorting enforcement are
+ // composed into a single idempotent pass, so the historical sequence
+ // of `Run::Distribution` / `Run::Sorting` collapses to repeated calls
+ // of the same rule. The sequences are preserved so existing test
+ // assertions (which encode legacy run orders) remain stable.
Review Comment:
It would be a nice follow on PR to conslidate these tests (so it only ran EnsureREquirements once) (or maybe ran it once, and then ran it again and ensured it was idempotent)
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
Review Comment:
Is `EnsureRequirementsNoPushdown` used anywhere? I don't see it
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -1621,18 +1618,18 @@ fn multi_smj_joins() -> Result<()> {
_ => {
assert_plan!(plan_distrib, @r"
SortMergeJoinExec: join_type=..., on=[(a@0, c@2)]
- SortExec: expr=[a@0 ASC], preserve_partitioning=[true]
- RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10
+ RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC
Review Comment:
It looks like the repartitioning was pulled up after the join (the plans look very similar to me)
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + `GlobalLimitExec` must produce a valid plan.
+ /// Regression for the `SanityCheckPlan` failure that motivated this PR:
+ /// `pushdown_sorts` setting `preserve_partitioning=true` on multi-partition
+ /// input without inserting `SortPreservingMergeExec` violated the
+ /// `SinglePartition` requirement from `GlobalLimitExec`.
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
Review Comment:
These tets now also juse `assert` when most of the existing tests use `insta` snapshots to make updating easier (which may not be a thing anymore now that AI coding tools are around), but consistency would be nice
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -1875,25 +1870,21 @@ fn smj_join_key_ordering() -> Result<()> {
let plan_sort = test_config.to_plan(join, &SORT_DISTRIB_DISTRIB);
assert_plan!(plan_sort, @r"
SortMergeJoinExec: join_type=Inner, on=[(b3@1, b2@1), (a3@0, a2@0)]
- RepartitionExec: partitioning=Hash([b3@1, a3@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- ProjectionExec: expr=[a1@0 as a3, b1@1 as b3]
- ProjectionExec: expr=[a1@1 as a1, b1@0 as b1]
- AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[]
- RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10
- AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
- RepartitionExec: partitioning=Hash([b2@1, a2@0], 10), input_partitions=1, maintains_sort_order=true
- SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[false]
- CoalescePartitionsExec
- ProjectionExec: expr=[a@1 as a2, b@0 as b2]
- AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[]
- RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10
- AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]
- RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1
- DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet
+ SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[true]
Review Comment:
same here -- this plan is better
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + `GlobalLimitExec` must produce a valid plan.
+ /// Regression for the `SanityCheckPlan` failure that motivated this PR:
+ /// `pushdown_sorts` setting `preserve_partitioning=true` on multi-partition
+ /// input without inserting `SortPreservingMergeExec` violated the
+ /// `SinglePartition` requirement from `GlobalLimitExec`.
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Idempotency: multi-partition sort + limit
+ #[test]
+ fn test_idempotent_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_idempotent(limit);
+ }
+
+ /// Idempotency: union with mixed partitions
+ #[test]
+ fn test_idempotent_union_mixed_partitions() {
+ let live = Arc::new(MockMultiPartitionExec::new(8));
+ let hist = Arc::new(MockMultiPartitionExec::new(1));
+ let union = UnionExec::try_new(vec![live as _, hist as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(5)));
+
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Projection + multi-partition tests (pushdown_sorts trigger path)
+ // ========================================================================
+
+ /// ProjectionExec over multi-partition + sort DESC + limit.
+ /// This is the topology where pushdown_sorts pushes sort through projection
+ /// onto the multi-partition source. The optimizer must still produce a valid plan.
+ #[test]
+ fn test_projection_over_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection =
+ Arc::new(ProjectionExec::try_new(proj_exprs, source as _).unwrap());
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, projection));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(Arc::clone(&limit));
+ assert!(
+ result.is_ok(),
+ "SanityCheckPlan failed for projection over multi-partition: {:?}",
+ result.err()
+ );
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Single partition tests (no unnecessary operators)
+ // ========================================================================
+
+ /// Single partition source + sort + limit should NOT add SortPreservingMergeExec.
+ #[test]
+ fn test_single_partition_no_unnecessary_spm() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ let optimized = optimize_and_sanity_check(limit).unwrap();
+ let plan_str = datafusion_physical_plan::displayable(optimized.as_ref())
+ .indent(true)
+ .to_string();
+
+ // Single partition should not have SortPreservingMergeExec
+ assert!(
+ !plan_str.contains("SortPreservingMergeExec"),
+ "Unnecessary SortPreservingMergeExec for single partition:\n{plan_str}"
+ );
Review Comment:
This pattern is often combined into a single macro or function using insta_assert in other tests to reduce the test reptition and make it easier tp update
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4415328481
Thanks @alamb. Plan sounds right — I'll do the test-extraction PR first (move EnforceDistribution / EnforceSorting tests into something like `datafusion/physical-optimizer/tests/ensure_requirements.rs`, no logic changes), then rebase this on top so the diff makes the porting and behavior changes obvious.
I'll wait until 54 is shipped before pushing the rebased version. Will ping you once the test-extraction PR is up.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4438274721
Quick update — pushed all three steps. Branch is now `ensure-requirements` (commits `d5bd105e8` and `47cd38363` on top of the original).
**Steps 1+2 (`d5bd105e8`)** — removed the disconnected `EnforceDistribution` / `EnforceSorting` rule structs, and retargeted the existing tests in `core/tests/physical_optimizer/`:
- `enforce_distribution.rs` — `Run::Distribution` / `Run::Sorting` branches both call `EnsureRequirements::new()`. Legacy run sequences kept verbatim so existing assertions retain their shape; the merged rule is idempotent so the previously-different orderings now converge.
- `enforce_sorting.rs` — the historical `[Dist, Sort]` vs `[Sort, Dist, Sort]` idempotency comparison is rewritten as "running EnsureRequirements N times == running it once".
- 78 snapshots refreshed; consistent pattern is `SortExec + CoalescePartitionsExec` (blocking) → `SortPreservingMergeExec` (streaming), since EnsureRequirements now applies `parallelize_sorts` + `replace_with_order_preserving_variants` on inputs the single-rule path didn't reach. Improvements rather than regressions.
**Step 3 (`47cd38363`)** — went ahead and made the call: dropped `ensure_requirements/new_tests.rs` entirely.
For the "which new-behavior cases to keep" question I raised earlier: turns out every test in both `new_tests.rs` and `ensure_requirements/mod.rs` inline tests uses either `optimize_and_sanity_check` (production-502 regression) or `assert_idempotent` (merged-rule core property). None are pure smoke tests. The real redundancy was organizational — two files with duplicated helpers and substantial plan-shape overlap (hash join, aggregate, window, projection, fetch, filter all covered in both).
So I deleted `new_tests.rs` and kept the `mod.rs` inline tests as-is. All the named bug regressions (#14150, PR53, PR54) live there, plus a comprehensive sweep of plan shapes with idempotency + SanityCheckPlan-validity assertions. If you'd prefer a more aggressive cull (keep only `#14150` / PR53 / PR54 named regressions, drop the rest), happy to do that as a follow-up commit — just say the word.
Test results after the cleanup:
- `core_integration physical_optimizer::` 454 passed
- `datafusion-physical-optimizer` (lib) 59 passed (down from 79 with `new_tests.rs` dropped)
- `clippy --all-targets -- -D warnings` clean
- `cargo fmt --check` clean
CI is running on the push now.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3278975177
##########
datafusion/core/src/optimizer_rule_reference.md:
##########
@@ -67,27 +67,26 @@ Rule order matters. The default pipeline may change between releases.
The same rule name may appear more than once when the default pipeline runs it
in multiple phases.
-| order | rule | phase | summary |
-| ----- | ------------------------------ | ----------------------- | ------------------------------------------------------------------------------------------------------------ |
-| 1 | `OutputRequirements` | add phase | Adds helper nodes so output requirements survive later physical rewrites. |
-| 2 | `aggregate_statistics` | - | Uses exact source statistics to answer some aggregates without scanning data. |
-| 3 | `join_selection` | - | Chooses join implementation, build side, and partition mode from statistics and stream properties. |
-| 4 | `LimitedDistinctAggregation` | - | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed. |
-| 5 | `FilterPushdown` | pre-optimization phase | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. |
-| 6 | `EnforceDistribution` | - | Adds repartitioning only where needed to satisfy physical distribution requirements. |
-| 7 | `CombinePartialFinalAggregate` | - | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant. |
-| 8 | `EnforceSorting` | - | Adds or removes local sorts to satisfy required input orderings. |
-| 9 | `OptimizeAggregateOrder` | - | Updates aggregate expressions to use the best ordering once sort requirements are known. |
-| 10 | `WindowTopN` | - | Replaces eligible row-number window and filter patterns with per-partition TopK execution. |
-| 11 | `ProjectionPushdown` | early pass | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure. |
-| 12 | `OutputRequirements` | remove phase | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done. |
-| 13 | `LimitAggregation` | - | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets. |
-| 14 | `LimitPushPastWindows` | - | Pushes fetch limits through bounded window operators when doing so keeps the result correct. |
-| 15 | `HashJoinBuffering` | - | Adds buffering on the probe side of hash joins so probing can start before build completion. |
-| 16 | `LimitPushdown` | - | Moves physical limits into child operators or fetch-enabled variants to cut data early. |
-| 17 | `TopKRepartition` | - | Pushes TopK below hash repartition when the partition key is a prefix of the sort key. |
-| 18 | `ProjectionPushdown` | late pass | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities. |
-| 19 | `PushdownSort` | - | Pushes sort requirements into data sources that can already return sorted output. |
-| 20 | `EnsureCooperative` | - | Wraps non-cooperative plan parts so long-running tasks yield fairly. |
-| 21 | `FilterPushdown(Post)` | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving. |
-| 22 | `SanityCheckPlan` | - | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. |
+| order | rule | phase | summary |
+| ----- | ------------------------------ | ----------------------- | -------------------------------------------------------------------------------------------------------------------------------- |
+| 1 | `OutputRequirements` | add phase | Adds helper nodes so output requirements survive later physical rewrites. |
+| 2 | `aggregate_statistics` | - | Uses exact source statistics to answer some aggregates without scanning data. |
+| 3 | `join_selection` | - | Chooses join implementation, build side, and partition mode from statistics and stream properties. |
+| 4 | `LimitedDistinctAggregation` | - | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed. |
+| 5 | `FilterPushdown` | pre-optimization phase | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. |
+| 6 | `EnsureRequirements` | - | Enforces both distribution and sorting requirements in a single idempotent rule (replaces EnforceDistribution + EnforceSorting). |
Review Comment:
Good point — once those rules are gone, the "replaces" parenthetical only dangles a reference readers can't grep. Dropped it; now just "Enforces both distribution and sorting requirements in a single idempotent rule."
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279276073
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -695,17 +695,14 @@ impl TestConfig {
// TODO: End state payloads will be checked here.
}
- for run in optimizers_to_run {
- optimized = match run {
- Run::Distribution => {
- let optimizer = EnforceDistribution::new();
- optimizer.optimize(optimized, &self.config)?
- }
- Run::Sorting => {
- let optimizer = EnforceSorting::new();
- optimizer.optimize(optimized, &self.config)?
- }
- };
+ // With `EnsureRequirements`, distribution and sorting enforcement are
+ // composed into a single idempotent pass, so the historical sequence
+ // of `Run::Distribution` / `Run::Sorting` collapses to repeated calls
+ // of the same rule. The sequences are preserved so existing test
+ // assertions (which encode legacy run orders) remain stable.
Review Comment:
Folded into this PR — f6f8f5a2e. The macro now runs EnsureRequirements twice (snapshot pass 1, verify pass 2 produces the same plan, sanity-check both), and the 6 tests that previously called both the macro and
assert_idempotent are down to one call site. Cuts optimizer runs from 3 → 2 per test and means "correct plan + idempotent + sanity" is a single assertion.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521120509
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521086953)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4521086953-290-xnwcg 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing ensure-requirements (504fb08b720a182fd36ec80ca2bf958a81193225) to 0da8961 (merge-base) [diff](https://github.com/apache/datafusion/compare/0da89616e43299546ed8405cfa900d0f97cb2515..504fb08b720a182fd36ec80ca2bf958a81193225) using: tpch
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521276189
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521086953)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
Comparing HEAD and ensure-requirements
--------------------
Benchmark tpch_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ ensure-requirements ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 39.57 / 40.87 ±1.17 / 42.82 ms │ 38.15 / 38.80 ±1.07 / 40.93 ms │ +1.05x faster │
│ QQuery 2 │ 20.56 / 21.09 ±0.58 / 22.16 ms │ 19.64 / 19.85 ±0.29 / 20.41 ms │ +1.06x faster │
│ QQuery 3 │ 35.16 / 37.06 ±1.06 / 38.42 ms │ 33.17 / 35.34 ±1.18 / 36.76 ms │ no change │
│ QQuery 4 │ 18.25 / 18.87 ±0.66 / 19.74 ms │ 17.26 / 17.31 ±0.03 / 17.35 ms │ +1.09x faster │
│ QQuery 5 │ 44.14 / 44.60 ±0.29 / 44.93 ms │ 39.77 / 41.79 ±1.01 / 42.34 ms │ +1.07x faster │
│ QQuery 6 │ 17.02 / 17.39 ±0.42 / 18.21 ms │ 16.07 / 16.27 ±0.11 / 16.36 ms │ +1.07x faster │
│ QQuery 7 │ 49.46 / 51.18 ±1.63 / 53.44 ms │ 46.46 / 48.36 ±2.60 / 53.51 ms │ +1.06x faster │
│ QQuery 8 │ 46.49 / 46.63 ±0.14 / 46.80 ms │ 44.23 / 44.67 ±0.25 / 44.95 ms │ no change │
│ QQuery 9 │ 51.90 / 52.97 ±0.91 / 54.45 ms │ 50.10 / 50.72 ±0.76 / 52.17 ms │ no change │
│ QQuery 10 │ 65.96 / 66.38 ±0.43 / 67.07 ms │ 63.32 / 63.54 ±0.19 / 63.82 ms │ no change │
│ QQuery 11 │ 14.33 / 15.12 ±0.98 / 16.75 ms │ 13.14 / 13.46 ±0.34 / 14.02 ms │ +1.12x faster │
│ QQuery 12 │ 24.48 / 25.44 ±0.53 / 25.94 ms │ 23.97 / 24.27 ±0.28 / 24.78 ms │ no change │
│ QQuery 13 │ 34.71 / 37.50 ±1.90 / 40.39 ms │ 33.77 / 35.87 ±1.69 / 38.65 ms │ no change │
│ QQuery 14 │ 26.03 / 26.29 ±0.19 / 26.59 ms │ 25.45 / 25.57 ±0.14 / 25.83 ms │ no change │
│ QQuery 15 │ 31.87 / 32.70 ±0.76 / 34.11 ms │ 31.32 / 31.76 ±0.51 / 32.74 ms │ no change │
│ QQuery 16 │ 15.19 / 15.39 ±0.19 / 15.66 ms │ 14.61 / 14.87 ±0.25 / 15.19 ms │ no change │
│ QQuery 17 │ 78.57 / 83.05 ±6.91 / 96.78 ms │ 73.51 / 74.12 ±0.58 / 75.02 ms │ +1.12x faster │
│ QQuery 18 │ 64.00 / 65.56 ±1.12 / 66.73 ms │ 61.61 / 63.16 ±1.04 / 64.65 ms │ no change │
│ QQuery 19 │ 34.99 / 35.67 ±1.02 / 37.65 ms │ 33.77 / 33.99 ±0.12 / 34.13 ms │ no change │
│ QQuery 20 │ 37.59 / 38.45 ±0.52 / 39.03 ms │ 37.01 / 37.91 ±0.68 / 38.99 ms │ no change │
│ QQuery 21 │ 58.17 / 60.03 ±1.14 / 61.25 ms │ 55.84 / 57.27 ±1.58 / 60.35 ms │ no change │
│ QQuery 22 │ 24.27 / 24.68 ±0.59 / 25.85 ms │ 23.00 / 24.30 ±1.97 / 28.22 ms │ no change │
└───────────┴────────────────────────────────┴────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│ Total Time (HEAD) │ 856.91ms │
│ Total Time (ensure-requirements) │ 813.20ms │
│ Average Time (HEAD) │ 38.95ms │
│ Average Time (ensure-requirements) │ 36.96ms │
│ Queries Faster │ 8 │
│ Queries Slower │ 0 │
│ Queries with No Change │ 14 │
│ Queries with Failure │ 0 │
└────────────────────────────────────┴──────────┘
```
</p>
</details>
<details><summary>Resource Usage</summary>
**tpch — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 5.0s |
| Peak memory | 5.5 GiB |
| Avg memory | 5.0 GiB |
| CPU user | 30.9s |
| CPU sys | 2.6s |
| Peak spill | 0 B |
**tpch — branch**
| Metric | Value |
|--------|-------|
| Wall time | 5.0s |
| Peak memory | 5.5 GiB |
| Avg memory | 5.0 GiB |
| CPU user | 29.6s |
| CPU sys | 2.2s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521536570
Thank you @zhuqi-lucas -- this is pretty epic
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "github-actions[bot] (via GitHub)" <gi...@apache.org>.
github-actions[bot] commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4438532869
<!-- semver-check-comment -->
Thank you for opening this pull request!
Reviewer note: [cargo-semver-checks](https://github.com/obi1kenobi/cargo-semver-checks) reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch).
<details>
<summary>Details</summary>
```
Cloning apache/main
Building datafusion v53.1.0 (current)
Built [ 84.629s] (current)
Parsing datafusion v53.1.0 (current)
Parsed [ 0.034s] (current)
Building datafusion v53.1.0 (baseline)
Built [ 82.120s] (baseline)
Parsing datafusion v53.1.0 (baseline)
Parsed [ 0.035s] (baseline)
Checking datafusion v53.1.0 -> v53.1.0 (no change; assume patch)
Checked [ 0.594s] 222 checks: 222 pass, 30 skip
Summary no semver update required
Finished [ 169.001s] datafusion
Building datafusion-physical-optimizer v53.1.0 (current)
Built [ 38.500s] (current)
Parsing datafusion-physical-optimizer v53.1.0 (current)
Parsed [ 0.023s] (current)
Building datafusion-physical-optimizer v53.1.0 (baseline)
Built [ 39.477s] (baseline)
Parsing datafusion-physical-optimizer v53.1.0 (baseline)
Parsed [ 0.023s] (baseline)
Checking datafusion-physical-optimizer v53.1.0 -> v53.1.0 (no change; assume patch)
Checked [ 0.104s] 222 checks: 221 pass, 1 fail, 0 warn, 30 skip
--- failure struct_missing: pub struct removed or renamed ---
Description:
A publicly-visible struct cannot be imported by its prior path. A `pub use` may have been removed, or the struct itself may have been renamed or removed entirely.
ref: https://doc.rust-lang.org/cargo/reference/semver.html#item-remove
impl: https://github.com/obi1kenobi/cargo-semver-checks/tree/v0.47.0/src/lints/struct_missing.ron
Failed in:
struct datafusion_physical_optimizer::enforce_sorting::EnforceSorting, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/da1bc44dbfe2ea5ef369389ef86956738e6352d0/datafusion/physical-optimizer/src/enforce_sorting/mod.rs:79
struct datafusion_physical_optimizer::enforce_distribution::EnforceDistribution, previously in file /home/runner/work/datafusion/datafusion/target/semver-checks/git-apache_main/da1bc44dbfe2ea5ef369389ef86956738e6352d0/datafusion/physical-optimizer/src/enforce_distribution.rs:184
Summary semver requires new major version: 1 major and 0 minor checks failed
Finished [ 79.346s] datafusion-physical-optimizer
Building datafusion-sqllogictest v53.1.0 (current)
Built [ 140.279s] (current)
Parsing datafusion-sqllogictest v53.1.0 (current)
Parsed [ 0.022s] (current)
Building datafusion-sqllogictest v53.1.0 (baseline)
Built [ 143.933s] (baseline)
Parsing datafusion-sqllogictest v53.1.0 (baseline)
Parsed [ 0.024s] (baseline)
Checking datafusion-sqllogictest v53.1.0 -> v53.1.0 (no change; assume patch)
Checked [ 0.087s] 222 checks: 222 pass, 30 skip
Summary no semver update required
Finished [ 287.195s] datafusion-sqllogictest
```
</details>
<!-- semver-check-comment -->
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279033484
##########
datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs:
##########
@@ -970,7 +813,10 @@ fn preserving_order_enables_streaming(
///
/// Updated node with an execution plan, where the desired single distribution
/// requirement is satisfied.
-fn add_merge_on_top(input: DistributionContext) -> DistributionContext {
+fn add_merge_on_top(
Review Comment:
Moved enforce_distribution.rs and enforce_sorting/ into ensure_requirements/ as submodules.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521287192
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521086953)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
Comparing HEAD and ensure-requirements
--------------------
Benchmark tpcds_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ ensure-requirements ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 1 │ 5.91 / 6.47 ±0.94 / 8.35 ms │ 5.67 / 6.28 ±0.91 / 8.08 ms │ no change │
│ QQuery 2 │ 82.93 / 83.30 ±0.30 / 83.72 ms │ 81.52 / 81.76 ±0.17 / 81.97 ms │ no change │
│ QQuery 3 │ 29.81 / 30.08 ±0.19 / 30.37 ms │ 29.72 / 30.01 ±0.21 / 30.30 ms │ no change │
│ QQuery 4 │ 527.52 / 536.18 ±6.10 / 546.63 ms │ 535.69 / 544.40 ±5.10 / 549.71 ms │ no change │
│ QQuery 5 │ 51.83 / 52.51 ±0.50 / 53.16 ms │ 51.26 / 51.99 ±0.77 / 53.01 ms │ no change │
│ QQuery 6 │ 36.39 / 36.81 ±0.29 / 37.25 ms │ 36.51 / 37.06 ±0.36 / 37.56 ms │ no change │
│ QQuery 7 │ 108.81 / 110.30 ±1.23 / 112.11 ms │ 109.21 / 111.15 ±2.09 / 115.17 ms │ no change │
│ QQuery 8 │ 37.28 / 37.56 ±0.29 / 37.98 ms │ 37.29 / 37.43 ±0.18 / 37.77 ms │ no change │
│ QQuery 9 │ 54.47 / 56.08 ±1.22 / 57.77 ms │ 53.14 / 55.95 ±1.84 / 58.12 ms │ no change │
│ QQuery 10 │ 81.16 / 81.91 ±0.98 / 83.85 ms │ 80.90 / 82.15 ±1.42 / 84.85 ms │ no change │
│ QQuery 11 │ 334.27 / 366.29 ±50.31 / 465.86 ms │ 341.72 / 344.92 ±3.52 / 351.49 ms │ +1.06x faster │
│ QQuery 12 │ 29.64 / 30.24 ±0.42 / 30.83 ms │ 29.27 / 29.70 ±0.31 / 30.10 ms │ no change │
│ QQuery 13 │ 127.57 / 128.02 ±0.34 / 128.48 ms │ 128.00 / 128.22 ±0.21 / 128.48 ms │ no change │
│ QQuery 14 │ 509.61 / 517.35 ±7.15 / 530.84 ms │ 513.45 / 515.84 ±2.17 / 519.83 ms │ no change │
│ QQuery 15 │ 62.59 / 64.21 ±2.20 / 68.47 ms │ 64.38 / 64.89 ±0.34 / 65.36 ms │ no change │
│ QQuery 16 │ 6.54 / 6.69 ±0.22 / 7.13 ms │ 6.47 / 6.67 ±0.27 / 7.19 ms │ no change │
│ QQuery 17 │ 81.57 / 82.27 ±0.54 / 82.86 ms │ 81.96 / 82.62 ±0.39 / 83.17 ms │ no change │
│ QQuery 18 │ 153.10 / 154.14 ±0.66 / 155.13 ms │ 153.70 / 154.18 ±0.44 / 154.95 ms │ no change │
│ QQuery 19 │ 41.27 / 41.52 ±0.21 / 41.85 ms │ 41.86 / 42.31 ±0.34 / 42.90 ms │ no change │
│ QQuery 20 │ 36.18 / 39.97 ±2.64 / 42.88 ms │ 36.38 / 36.86 ±0.34 / 37.32 ms │ +1.08x faster │
│ QQuery 21 │ 19.36 / 19.74 ±0.28 / 20.19 ms │ 17.79 / 17.85 ±0.03 / 17.87 ms │ +1.11x faster │
│ QQuery 22 │ 63.65 / 73.34 ±5.13 / 77.57 ms │ 63.05 / 63.69 ±0.53 / 64.62 ms │ +1.15x faster │
│ QQuery 23 │ 492.49 / 498.24 ±4.08 / 503.27 ms │ 496.82 / 499.29 ±1.52 / 501.14 ms │ no change │
│ QQuery 24 │ 240.80 / 244.67 ±3.89 / 251.39 ms │ 238.40 / 242.46 ±3.25 / 248.28 ms │ no change │
│ QQuery 25 │ 114.23 / 115.46 ±1.33 / 118.02 ms │ 115.71 / 117.02 ±1.47 / 119.85 ms │ no change │
│ QQuery 26 │ 71.42 / 72.71 ±1.69 / 75.92 ms │ 72.16 / 73.10 ±0.85 / 74.69 ms │ no change │
│ QQuery 27 │ 6.77 / 6.92 ±0.19 / 7.28 ms │ 6.94 / 7.83 ±1.55 / 10.91 ms │ 1.13x slower │
│ QQuery 28 │ 63.17 / 63.82 ±0.43 / 64.49 ms │ 58.41 / 61.77 ±1.70 / 63.08 ms │ no change │
│ QQuery 29 │ 99.99 / 100.45 ±0.45 / 101.28 ms │ 100.36 / 102.19 ±1.10 / 103.64 ms │ no change │
│ QQuery 30 │ 31.11 / 32.30 ±1.81 / 35.90 ms │ 31.21 / 31.54 ±0.32 / 32.14 ms │ no change │
│ QQuery 31 │ 114.15 / 114.88 ±0.63 / 115.90 ms │ 115.02 / 115.58 ±0.45 / 116.38 ms │ no change │
│ QQuery 32 │ 21.10 / 21.71 ±0.78 / 23.14 ms │ 21.10 / 21.27 ±0.09 / 21.36 ms │ no change │
│ QQuery 33 │ 39.15 / 40.88 ±2.34 / 45.52 ms │ 39.58 / 40.57 ±1.06 / 42.46 ms │ no change │
│ QQuery 34 │ 9.90 / 10.69 ±0.59 / 11.38 ms │ 9.95 / 10.20 ±0.21 / 10.54 ms │ no change │
│ QQuery 35 │ 81.93 / 82.41 ±0.41 / 82.96 ms │ 82.41 / 83.57 ±1.13 / 85.51 ms │ no change │
│ QQuery 36 │ 5.99 / 6.15 ±0.18 / 6.50 ms │ 6.11 / 6.21 ±0.10 / 6.39 ms │ no change │
│ QQuery 37 │ 6.96 / 7.04 ±0.07 / 7.13 ms │ 7.10 / 7.17 ±0.08 / 7.33 ms │ no change │
│ QQuery 38 │ 69.82 / 70.61 ±0.75 / 71.92 ms │ 71.11 / 71.71 ±0.53 / 72.68 ms │ no change │
│ QQuery 39 │ 100.70 / 102.79 ±2.18 / 106.96 ms │ 102.36 / 103.25 ±0.49 / 103.73 ms │ no change │
│ QQuery 40 │ 23.66 / 24.15 ±0.40 / 24.85 ms │ 23.68 / 24.01 ±0.27 / 24.46 ms │ no change │
│ QQuery 41 │ 11.59 / 11.84 ±0.21 / 12.07 ms │ 11.91 / 12.13 ±0.14 / 12.31 ms │ no change │
│ QQuery 42 │ 24.35 / 24.94 ±0.53 / 25.89 ms │ 24.58 / 24.91 ±0.39 / 25.67 ms │ no change │
│ QQuery 43 │ 4.85 / 4.97 ±0.16 / 5.29 ms │ 4.80 / 4.91 ±0.20 / 5.31 ms │ no change │
│ QQuery 44 │ 10.90 / 11.15 ±0.17 / 11.36 ms │ 10.93 / 11.07 ±0.10 / 11.20 ms │ no change │
│ QQuery 45 │ 41.69 / 42.76 ±0.90 / 44.39 ms │ 42.56 / 43.56 ±0.73 / 44.81 ms │ no change │
│ QQuery 46 │ 13.29 / 13.78 ±0.64 / 15.03 ms │ 13.48 / 13.75 ±0.22 / 14.03 ms │ no change │
│ QQuery 47 │ 237.07 / 239.47 ±1.90 / 242.33 ms │ 241.35 / 246.97 ±3.33 / 251.12 ms │ no change │
│ QQuery 48 │ 103.81 / 104.98 ±1.25 / 107.32 ms │ 104.65 / 105.05 ±0.33 / 105.40 ms │ no change │
│ QQuery 49 │ 79.98 / 80.92 ±0.51 / 81.43 ms │ 81.26 / 82.70 ±0.93 / 83.76 ms │ no change │
│ QQuery 50 │ 60.38 / 61.01 ±0.49 / 61.86 ms │ 60.93 / 61.41 ±0.30 / 61.80 ms │ no change │
│ QQuery 51 │ 92.69 / 94.94 ±1.75 / 97.84 ms │ 93.06 / 95.90 ±2.24 / 99.51 ms │ no change │
│ QQuery 52 │ 25.07 / 25.49 ±0.34 / 25.98 ms │ 25.03 / 25.44 ±0.28 / 25.88 ms │ no change │
│ QQuery 53 │ 30.33 / 30.59 ±0.24 / 31.00 ms │ 30.28 / 30.67 ±0.28 / 31.12 ms │ no change │
│ QQuery 54 │ 55.52 / 55.92 ±0.32 / 56.32 ms │ 55.31 / 55.88 ±0.46 / 56.72 ms │ no change │
│ QQuery 55 │ 24.28 / 24.34 ±0.05 / 24.44 ms │ 24.29 / 24.63 ±0.37 / 25.33 ms │ no change │
│ QQuery 56 │ 39.14 / 39.74 ±0.45 / 40.48 ms │ 39.91 / 40.58 ±0.86 / 42.21 ms │ no change │
│ QQuery 57 │ 177.98 / 179.28 ±0.74 / 180.29 ms │ 177.17 / 180.01 ±1.53 / 181.42 ms │ no change │
│ QQuery 58 │ 119.50 / 120.10 ±0.41 / 120.67 ms │ 120.61 / 121.65 ±0.96 / 123.06 ms │ no change │
│ QQuery 59 │ 118.84 / 120.43 ±2.02 / 124.29 ms │ 117.84 / 118.83 ±0.55 / 119.51 ms │ no change │
│ QQuery 60 │ 40.48 / 41.19 ±0.75 / 42.60 ms │ 39.71 / 40.35 ±0.55 / 41.37 ms │ no change │
│ QQuery 61 │ 13.32 / 13.64 ±0.27 / 14.03 ms │ 12.97 / 13.10 ±0.10 / 13.24 ms │ no change │
│ QQuery 62 │ 47.29 / 47.49 ±0.21 / 47.80 ms │ 47.03 / 47.33 ±0.22 / 47.60 ms │ no change │
│ QQuery 63 │ 30.34 / 30.60 ±0.20 / 30.95 ms │ 30.53 / 30.74 ±0.20 / 31.02 ms │ no change │
│ QQuery 64 │ 465.05 / 470.51 ±5.35 / 478.58 ms │ 468.58 / 477.80 ±7.29 / 487.79 ms │ no change │
│ QQuery 65 │ 143.51 / 147.61 ±2.78 / 151.45 ms │ 150.16 / 152.60 ±1.59 / 154.67 ms │ no change │
│ QQuery 66 │ 80.50 / 81.74 ±1.24 / 83.57 ms │ 80.13 / 82.83 ±3.30 / 88.63 ms │ no change │
│ QQuery 67 │ 253.99 / 258.41 ±4.28 / 266.55 ms │ 260.23 / 266.53 ±4.43 / 272.96 ms │ no change │
│ QQuery 68 │ 13.57 / 13.73 ±0.19 / 14.09 ms │ 13.31 / 13.50 ±0.17 / 13.79 ms │ no change │
│ QQuery 69 │ 77.30 / 77.93 ±0.69 / 79.25 ms │ 77.35 / 78.52 ±1.66 / 81.82 ms │ no change │
│ QQuery 70 │ 105.87 / 108.42 ±4.43 / 117.27 ms │ 107.08 / 110.80 ±5.38 / 121.49 ms │ no change │
│ QQuery 71 │ 36.21 / 36.79 ±0.41 / 37.32 ms │ 36.39 / 36.65 ±0.21 / 37.03 ms │ no change │
│ QQuery 72 │ 2181.20 / 2290.08 ±70.56 / 2400.63 ms │ 2218.69 / 2284.55 ±54.23 / 2373.67 ms │ no change │
│ QQuery 73 │ 9.46 / 10.63 ±2.09 / 14.80 ms │ 9.65 / 11.77 ±2.85 / 17.40 ms │ 1.11x slower │
│ QQuery 74 │ 191.17 / 192.62 ±1.08 / 194.37 ms │ 195.52 / 199.28 ±2.29 / 201.66 ms │ no change │
│ QQuery 75 │ 149.10 / 150.45 ±0.93 / 151.94 ms │ 149.41 / 150.39 ±0.74 / 151.34 ms │ no change │
│ QQuery 76 │ 36.26 / 36.60 ±0.26 / 36.95 ms │ 35.98 / 36.60 ±0.37 / 37.07 ms │ no change │
│ QQuery 77 │ 61.20 / 62.90 ±2.94 / 68.77 ms │ 61.16 / 62.57 ±1.68 / 65.83 ms │ no change │
│ QQuery 78 │ 192.99 / 195.97 ±2.19 / 199.32 ms │ 194.12 / 196.93 ±1.97 / 199.44 ms │ no change │
│ QQuery 79 │ 68.32 / 68.77 ±0.32 / 69.13 ms │ 68.33 / 68.83 ±0.50 / 69.56 ms │ no change │
│ QQuery 80 │ 101.71 / 103.83 ±2.72 / 109.08 ms │ 102.23 / 105.41 ±2.35 / 109.10 ms │ no change │
│ QQuery 81 │ 24.75 / 25.04 ±0.25 / 25.35 ms │ 25.27 / 25.56 ±0.25 / 25.95 ms │ no change │
│ QQuery 82 │ 17.17 / 17.31 ±0.10 / 17.45 ms │ 16.99 / 17.22 ±0.14 / 17.43 ms │ no change │
│ QQuery 83 │ 38.08 / 38.45 ±0.40 / 39.19 ms │ 38.03 / 38.23 ±0.15 / 38.49 ms │ no change │
│ QQuery 84 │ 43.95 / 44.18 ±0.14 / 44.33 ms │ 44.39 / 45.30 ±1.56 / 48.41 ms │ no change │
│ QQuery 85 │ 137.90 / 140.95 ±2.18 / 144.14 ms │ 138.54 / 140.02 ±2.00 / 143.97 ms │ no change │
│ QQuery 86 │ 25.99 / 26.36 ±0.41 / 27.10 ms │ 25.76 / 26.15 ±0.28 / 26.54 ms │ no change │
│ QQuery 87 │ 70.83 / 71.30 ±0.31 / 71.67 ms │ 71.35 / 72.35 ±0.72 / 73.30 ms │ no change │
│ QQuery 88 │ 63.87 / 64.91 ±0.94 / 66.61 ms │ 64.05 / 65.78 ±2.47 / 70.64 ms │ no change │
│ QQuery 89 │ 36.74 / 37.23 ±0.35 / 37.81 ms │ 36.82 / 37.24 ±0.25 / 37.55 ms │ no change │
│ QQuery 90 │ 17.74 / 18.00 ±0.16 / 18.14 ms │ 17.54 / 17.68 ±0.13 / 17.86 ms │ no change │
│ QQuery 91 │ 52.64 / 53.34 ±1.00 / 55.32 ms │ 52.48 / 53.61 ±1.27 / 56.09 ms │ no change │
│ QQuery 92 │ 30.46 / 30.69 ±0.19 / 30.93 ms │ 30.33 / 31.07 ±0.75 / 32.26 ms │ no change │
│ QQuery 93 │ 51.18 / 51.85 ±0.50 / 52.59 ms │ 51.64 / 52.40 ±0.52 / 53.09 ms │ no change │
│ QQuery 94 │ 39.19 / 39.50 ±0.28 / 39.84 ms │ 38.90 / 39.52 ±0.39 / 40.10 ms │ no change │
│ QQuery 95 │ 86.62 / 87.11 ±0.50 / 88.03 ms │ 86.75 / 88.06 ±1.11 / 89.26 ms │ no change │
│ QQuery 96 │ 24.87 / 25.11 ±0.19 / 25.44 ms │ 24.91 / 25.01 ±0.08 / 25.13 ms │ no change │
│ QQuery 97 │ 46.63 / 47.28 ±0.38 / 47.69 ms │ 46.96 / 47.86 ±0.73 / 49.11 ms │ no change │
│ QQuery 98 │ 43.40 / 43.87 ±0.45 / 44.66 ms │ 43.96 / 44.21 ±0.28 / 44.59 ms │ no change │
│ QQuery 99 │ 70.54 / 71.22 ±0.65 / 72.21 ms │ 71.00 / 72.14 ±1.55 / 75.10 ms │ no change │
└───────────┴───────────────────────────────────────┴───────────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ Total Time (HEAD) │ 10831.08ms │
│ Total Time (ensure-requirements) │ 10851.20ms │
│ Average Time (HEAD) │ 109.40ms │
│ Average Time (ensure-requirements) │ 109.61ms │
│ Queries Faster │ 4 │
│ Queries Slower │ 2 │
│ Queries with No Change │ 93 │
│ Queries with Failure │ 0 │
└────────────────────────────────────┴────────────┘
```
</p>
</details>
<details><summary>Resource Usage</summary>
**tpcds — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 55.0s |
| Peak memory | 6.9 GiB |
| Avg memory | 6.3 GiB |
| CPU user | 242.7s |
| CPU sys | 5.9s |
| Peak spill | 0 B |
**tpcds — branch**
| Metric | Value |
|--------|-------|
| Wall time | 55.0s |
| Peak memory | 6.9 GiB |
| Avg memory | 6.2 GiB |
| CPU user | 242.4s |
| CPU sys | 6.0s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521509204
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521347657)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
Comparing HEAD and ensure-requirements
--------------------
Benchmark clickbench_partitioned.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ ensure-requirements ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 0 │ 1.14 / 4.68 ±6.94 / 18.56 ms │ 1.14 / 4.66 ±6.89 / 18.45 ms │ no change │
│ QQuery 1 │ 12.12 / 12.74 ±0.33 / 13.07 ms │ 12.48 / 12.63 ±0.15 / 12.83 ms │ no change │
│ QQuery 2 │ 35.25 / 35.68 ±0.36 / 36.15 ms │ 35.28 / 35.77 ±0.43 / 36.52 ms │ no change │
│ QQuery 3 │ 30.63 / 31.35 ±0.50 / 32.05 ms │ 30.45 / 30.72 ±0.16 / 30.90 ms │ no change │
│ QQuery 4 │ 224.67 / 227.46 ±2.04 / 230.68 ms │ 226.44 / 227.78 ±1.45 / 230.10 ms │ no change │
│ QQuery 5 │ 270.17 / 273.24 ±2.10 / 276.07 ms │ 269.53 / 271.76 ±2.55 / 276.53 ms │ no change │
│ QQuery 6 │ 1.17 / 1.32 ±0.23 / 1.77 ms │ 1.17 / 1.32 ±0.23 / 1.77 ms │ no change │
│ QQuery 7 │ 13.66 / 13.75 ±0.06 / 13.84 ms │ 13.59 / 13.73 ±0.15 / 14.01 ms │ no change │
│ QQuery 8 │ 316.83 / 320.20 ±2.56 / 323.94 ms │ 320.21 / 322.73 ±2.25 / 326.72 ms │ no change │
│ QQuery 9 │ 449.96 / 460.01 ±6.03 / 468.39 ms │ 443.32 / 453.74 ±5.34 / 458.42 ms │ no change │
│ QQuery 10 │ 69.24 / 69.64 ±0.33 / 70.11 ms │ 68.92 / 69.47 ±0.46 / 70.10 ms │ no change │
│ QQuery 11 │ 81.41 / 83.43 ±1.28 / 85.03 ms │ 80.17 / 82.84 ±3.48 / 89.67 ms │ no change │
│ QQuery 12 │ 260.92 / 267.97 ±4.36 / 273.33 ms │ 261.32 / 265.72 ±2.76 / 268.86 ms │ no change │
│ QQuery 13 │ 366.49 / 375.67 ±7.83 / 387.16 ms │ 361.91 / 372.53 ±11.02 / 391.74 ms │ no change │
│ QQuery 14 │ 276.09 / 280.42 ±4.41 / 288.78 ms │ 276.79 / 279.94 ±4.17 / 288.09 ms │ no change │
│ QQuery 15 │ 262.24 / 270.80 ±10.28 / 289.11 ms │ 261.89 / 269.14 ±6.60 / 279.01 ms │ no change │
│ QQuery 16 │ 602.68 / 609.42 ±4.69 / 615.23 ms │ 622.12 / 626.76 ±2.48 / 629.62 ms │ no change │
│ QQuery 17 │ 608.21 / 613.94 ±6.74 / 626.84 ms │ 619.63 / 628.38 ±8.19 / 638.27 ms │ no change │
│ QQuery 18 │ 1236.15 / 1246.89 ±9.69 / 1262.57 ms │ 1274.94 / 1286.31 ±7.96 / 1295.07 ms │ no change │
│ QQuery 19 │ 27.83 / 31.57 ±4.13 / 37.14 ms │ 27.74 / 27.83 ±0.05 / 27.88 ms │ +1.13x faster │
│ QQuery 20 │ 514.41 / 527.15 ±14.32 / 552.90 ms │ 518.71 / 522.13 ±2.27 / 524.92 ms │ no change │
│ QQuery 21 │ 594.45 / 598.44 ±3.03 / 603.44 ms │ 599.09 / 603.76 ±4.11 / 608.57 ms │ no change │
│ QQuery 22 │ 1052.55 / 1062.35 ±7.62 / 1071.23 ms │ 1064.84 / 1076.13 ±7.68 / 1087.48 ms │ no change │
│ QQuery 23 │ 3148.40 / 3174.59 ±19.51 / 3205.32 ms │ 3106.66 / 3156.07 ±26.08 / 3177.81 ms │ no change │
│ QQuery 24 │ 41.12 / 42.17 ±1.04 / 43.53 ms │ 41.00 / 45.96 ±4.02 / 49.97 ms │ 1.09x slower │
│ QQuery 25 │ 109.73 / 112.70 ±3.13 / 118.33 ms │ 110.14 / 111.73 ±1.25 / 113.50 ms │ no change │
│ QQuery 26 │ 41.86 / 42.49 ±0.54 / 43.40 ms │ 41.85 / 43.86 ±2.75 / 49.26 ms │ no change │
│ QQuery 27 │ 665.50 / 667.89 ±2.98 / 673.78 ms │ 667.53 / 671.91 ±5.45 / 682.45 ms │ no change │
│ QQuery 28 │ 3017.89 / 3036.69 ±10.54 / 3049.17 ms │ 3003.34 / 3028.72 ±21.99 / 3062.70 ms │ no change │
│ QQuery 29 │ 39.84 / 46.32 ±7.91 / 58.61 ms │ 39.41 / 41.84 ±3.88 / 49.47 ms │ +1.11x faster │
│ QQuery 30 │ 296.97 / 302.06 ±3.35 / 306.45 ms │ 296.29 / 300.34 ±2.25 / 303.09 ms │ no change │
│ QQuery 31 │ 282.05 / 290.71 ±6.65 / 299.25 ms │ 281.68 / 286.78 ±4.02 / 290.67 ms │ no change │
│ QQuery 32 │ 900.27 / 915.18 ±10.98 / 926.40 ms │ 925.33 / 951.63 ±20.09 / 976.70 ms │ no change │
│ QQuery 33 │ 1409.22 / 1426.91 ±11.24 / 1442.72 ms │ 1421.65 / 1440.27 ±16.07 / 1465.57 ms │ no change │
│ QQuery 34 │ 1450.12 / 1456.82 ±6.95 / 1468.55 ms │ 1413.69 / 1441.84 ±17.52 / 1460.83 ms │ no change │
│ QQuery 35 │ 275.67 / 288.24 ±12.15 / 305.53 ms │ 277.12 / 315.98 ±72.29 / 460.49 ms │ 1.10x slower │
│ QQuery 36 │ 64.83 / 68.82 ±2.72 / 73.16 ms │ 64.23 / 71.35 ±6.89 / 83.89 ms │ no change │
│ QQuery 37 │ 35.67 / 38.80 ±3.33 / 44.38 ms │ 35.10 / 36.73 ±2.59 / 41.85 ms │ +1.06x faster │
│ QQuery 38 │ 40.58 / 45.44 ±3.08 / 48.37 ms │ 39.95 / 43.23 ±1.94 / 45.55 ms │ no change │
│ QQuery 39 │ 142.75 / 150.63 ±4.16 / 154.57 ms │ 135.18 / 152.22 ±14.05 / 177.76 ms │ no change │
│ QQuery 40 │ 13.99 / 15.62 ±2.12 / 19.65 ms │ 13.67 / 16.33 ±4.82 / 25.98 ms │ no change │
│ QQuery 41 │ 13.29 / 13.57 ±0.20 / 13.87 ms │ 13.30 / 13.60 ±0.26 / 14.09 ms │ no change │
│ QQuery 42 │ 12.88 / 14.10 ±1.97 / 18.02 ms │ 12.82 / 14.37 ±2.87 / 20.10 ms │ no change │
└───────────┴───────────────────────────────────────┴───────────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ Total Time (HEAD) │ 19567.89ms │
│ Total Time (ensure-requirements) │ 19670.55ms │
│ Average Time (HEAD) │ 455.07ms │
│ Average Time (ensure-requirements) │ 457.45ms │
│ Queries Faster │ 3 │
│ Queries Slower │ 2 │
│ Queries with No Change │ 38 │
│ Queries with Failure │ 0 │
└────────────────────────────────────┴────────────┘
```
</p>
</details>
<details><summary>Resource Usage</summary>
**clickbench_partitioned — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 100.0s |
| Peak memory | 31.1 GiB |
| Avg memory | 23.4 GiB |
| CPU user | 1021.0s |
| CPU sys | 62.0s |
| Peak spill | 0 B |
**clickbench_partitioned — branch**
| Metric | Value |
|--------|-------|
| Wall time | 100.0s |
| Peak memory | 29.7 GiB |
| Avg memory | 23.3 GiB |
| CPU user | 1022.8s |
| CPU sys | 64.0s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521119108
🤖 Benchmark running (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521086953)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB) | `Linux bench-c4521086953-288-jqxvw 6.12.68+ #1 SMP Wed Apr 1 02:23:28 UTC 2026 aarch64 GNU/Linux`
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
Comparing ensure-requirements (504fb08b720a182fd36ec80ca2bf958a81193225) to 0da8961 (merge-base) [diff](https://github.com/apache/datafusion/compare/0da89616e43299546ed8405cfa900d0f97cb2515..504fb08b720a182fd36ec80ca2bf958a81193225) using: clickbench_partitioned
Results will be posted here when complete
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521087701
run benchmark sql_planner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279097357
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + `GlobalLimitExec` must produce a valid plan.
+ /// Regression for the `SanityCheckPlan` failure that motivated this PR:
+ /// `pushdown_sorts` setting `preserve_partitioning=true` on multi-partition
+ /// input without inserting `SortPreservingMergeExec` violated the
+ /// `SinglePartition` requirement from `GlobalLimitExec`.
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
Review Comment:
Tried using insta for the idempotency tests but it doesn't fit well for this property-based style. Will convert the plan-shape ones in a follow-up.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas merged PR #21976:
URL: https://github.com/apache/datafusion/pull/21976
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4548440317
EPIC!
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4468997723
> Thank you @zhuqi-lucas
>
> cc @ozankabak @berkaysynnada @mustafasrepo as I believe you have spent quite a lot of time in / on this code as well
>
> I went through the code and tests (except the new ones in enforce_distribution) carefully and I think this is a very nice code refactor and I think the PR could be merged as is. It is especially impressive that all the existing slt tests do not change.
>
> I think there are several improvements we should make, though given how large this PR is already is is pretty unwieldy (and it is so large it will be hard to keep re-reviewing) I recommend we get it in as soon as possible (or break it into some smaller pieces) and then iterate / cleanup as follow on PRs
>
> I suggest
>
> 1. We wait until we have cut the DataFusion 54 release to merge this one in (as I think given its size and importance, it has potential for unintended consequences)
> 2. Remove dead `EnsureRequirementsNoPushdown`
> 3. Re-review the tests added in ensure_distribution and remove anything that is redundant and port the rest to datafusion/core/tests/physical_optimizer/ensure_distribution.rs so the tests stay together (also maybe leave a comment in ensure_distribution.rs to point at the test locations)
> 4. Plan a few follow on PRs to clean up the tests (eg the tests that currently seem to test only EnsureDistribution and EnsureSorting individually)
Thank you @alamb for review, sure, i will address remaining comments and wait until we have cut the DataFusion 54 release .
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3192516930
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1804 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
Review Comment:
`TransformedResult` is imported but not used in this module. If CI runs `clippy -D warnings`, this will fail the build; please remove the unused import (or use it if intended).
##########
datafusion/physical-optimizer/src/optimizer.rs:
##########
@@ -170,18 +169,12 @@ impl PhysicalOptimizer {
// those are handled by the later `FilterPushdown` rule.
// See `FilterPushdownPhase` for more details.
Arc::new(FilterPushdown::new()),
- // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution
- // requirements. Please make sure that the whole plan tree is determined before this rule.
- // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at
- // least one of the operators in the plan benefits from increased parallelism.
- Arc::new(EnforceDistribution::new()),
- // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule
+ // EnsureRequirements: merged EnforceDistribution + EnforceSorting into a
+ // single idempotent rule with distribution-aware pushdown_sorts.
+ // See https://github.com/apache/datafusion/issues/21973
+ Arc::new(EnsureRequirements::new()),
+ // The CombinePartialFinalAggregate rule should be applied after distribution enforcement
Arc::new(CombinePartialFinalAggregate::new()),
Review Comment:
With the default optimizer chain change, EnsureRequirements now enforces sorting before `CombinePartialFinalAggregate` runs. `CombinePartialFinalAggregate` only matches directly-adjacent Partial→Final aggregates, but `ensure_sorting` can insert `SortExec` between them (e.g. when the Final aggregate has `required_input_ordering`), preventing the combine optimization. Consider moving `CombinePartialFinalAggregate` before the sorting phase (e.g. keep it between distribution and sorting inside EnsureRequirements, or otherwise preserve the previous relative ordering of this combine rule vs sort enforcement).
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1804 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
Review Comment:
Module docs claim EnsureRequirements performs a single bottom-up pass and has "No separate `pushdown_sorts`", but the implementation still runs `pushdown_sorts` in Phase 3c (and also runs multiple sequential passes: distribution enforcement, sorting enforcement, parallelize_sorts, order-preserving variants, pushdown_sorts, partial sort). Please update the module-level documentation to match the actual algorithm, or adjust the implementation to align with the documented design.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3194201853
##########
datafusion/physical-optimizer/src/optimizer.rs:
##########
@@ -170,18 +169,12 @@ impl PhysicalOptimizer {
// those are handled by the later `FilterPushdown` rule.
// See `FilterPushdownPhase` for more details.
Arc::new(FilterPushdown::new()),
- // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution
- // requirements. Please make sure that the whole plan tree is determined before this rule.
- // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at
- // least one of the operators in the plan benefits from increased parallelism.
- Arc::new(EnforceDistribution::new()),
- // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule
+ // EnsureRequirements: merged EnforceDistribution + EnforceSorting into a
+ // single idempotent rule with distribution-aware pushdown_sorts.
+ // See https://github.com/apache/datafusion/issues/21973
Review Comment:
```suggestion
// Ensures each input plan satisfies the distribution and ordering requirements
// declared by `ExecutionPlan::required_input_distribution` and
// `ExecutionPlan::required_input_ordering`.
// If the requirements are already satisfied, this rule leaves the plan
// unchanged. For example, it does not add sorting when the input is a file
// scan whose existing order already satisfies the required ordering.
// Otherwise, this rule inserts the necessary repartitioning and sorting
// operators.
// This used to be implemented as two separate rules: `EnforceDistribution`
// and `EnforceSorting`. It is now a single idempotent rule with
// distribution-aware `pushdown_sorts`.
// See https://github.com/apache/datafusion/issues/21973.
```
Added more comments since this is the entry point.
I also have a question: What is this 'distribution-aware pushdown_sorts'? We cloud link some reference to it.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3194201853
##########
datafusion/physical-optimizer/src/optimizer.rs:
##########
@@ -170,18 +169,12 @@ impl PhysicalOptimizer {
// those are handled by the later `FilterPushdown` rule.
// See `FilterPushdownPhase` for more details.
Arc::new(FilterPushdown::new()),
- // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution
- // requirements. Please make sure that the whole plan tree is determined before this rule.
- // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at
- // least one of the operators in the plan benefits from increased parallelism.
- Arc::new(EnforceDistribution::new()),
- // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule
+ // EnsureRequirements: merged EnforceDistribution + EnforceSorting into a
+ // single idempotent rule with distribution-aware pushdown_sorts.
+ // See https://github.com/apache/datafusion/issues/21973
Review Comment:
```suggestion
// Ensures each input plan satisfies the distribution and ordering requirements
// declared by `ExecutionPlan::required_input_distribution` and
// `ExecutionPlan::required_input_ordering`.
// If the requirements are already satisfied, this rule leaves the plan
// unchanged. For example, it does not add sorting when the input is a file
// scan whose existing order already satisfies the required ordering.
// Otherwise, this rule inserts the necessary repartitioning and sorting
// operators.
// This used to be implemented as two separate rules: `EnforceDistribution`
// and `EnforceSorting`. It is now a single idempotent rule with
// distribution-aware `pushdown_sorts`.
// See https://github.com/apache/datafusion/issues/21973.
```
Question: What is this 'distribution-aware pushdown_sorts'? We cloud link some reference to it.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279522574
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
Review Comment:
Fair point — tightened the module doc in 97728adb7 to drop the "single pass" framing and break down what each phase actually does. Phase 2 is the one combined distribution+sorting bottom-up; the join-key reorder, parallelize_sorts, replace_with_order_preserving_variants, pushdown_sorts, replace_with_partial_sort are still separate traversals. Agree consolidating is a good follow-up, but i tried something before, it's pretty hard.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3279222247
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1814 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
Review Comment:
Good catch, removed it now.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521320937
🤖 Benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521086953)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
Comparing HEAD and ensure-requirements
--------------------
Benchmark clickbench_partitioned.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┓
┃ Query ┃ HEAD ┃ ensure-requirements ┃ Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━┩
│ QQuery 0 │ 1.13 / 4.64 ±6.94 / 18.52 ms │ 1.14 / 4.55 ±6.77 / 18.09 ms │ no change │
│ QQuery 1 │ 12.38 / 12.81 ±0.23 / 13.02 ms │ 12.86 / 13.23 ±0.51 / 14.24 ms │ no change │
│ QQuery 2 │ 36.69 / 37.11 ±0.29 / 37.46 ms │ 35.43 / 35.68 ±0.25 / 36.15 ms │ no change │
│ QQuery 3 │ 31.60 / 32.13 ±0.60 / 33.29 ms │ 31.08 / 31.31 ±0.22 / 31.62 ms │ no change │
│ QQuery 4 │ 221.31 / 224.45 ±3.37 / 230.94 ms │ 225.05 / 228.46 ±4.53 / 237.39 ms │ no change │
│ QQuery 5 │ 267.34 / 270.93 ±2.38 / 274.67 ms │ 268.70 / 271.08 ±1.52 / 273.03 ms │ no change │
│ QQuery 6 │ 1.16 / 1.30 ±0.23 / 1.74 ms │ 1.18 / 1.32 ±0.21 / 1.74 ms │ no change │
│ QQuery 7 │ 13.44 / 13.62 ±0.13 / 13.81 ms │ 13.74 / 14.08 ±0.30 / 14.62 ms │ no change │
│ QQuery 8 │ 316.24 / 319.61 ±1.98 / 321.75 ms │ 318.89 / 322.66 ±4.55 / 331.32 ms │ no change │
│ QQuery 9 │ 446.84 / 455.58 ±7.10 / 467.49 ms │ 446.18 / 496.23 ±27.94 / 519.85 ms │ 1.09x slower │
│ QQuery 10 │ 69.12 / 70.83 ±1.89 / 74.03 ms │ 69.28 / 70.91 ±1.25 / 72.42 ms │ no change │
│ QQuery 11 │ 81.40 / 82.19 ±0.50 / 82.96 ms │ 80.36 / 84.76 ±5.40 / 95.37 ms │ no change │
│ QQuery 12 │ 263.08 / 267.19 ±4.36 / 272.73 ms │ 263.33 / 271.38 ±5.92 / 281.76 ms │ no change │
│ QQuery 13 │ 370.32 / 375.43 ±3.18 / 379.47 ms │ 362.19 / 366.08 ±3.31 / 371.95 ms │ no change │
│ QQuery 14 │ 277.22 / 280.83 ±3.71 / 287.77 ms │ 279.77 / 285.84 ±6.70 / 297.85 ms │ no change │
│ QQuery 15 │ 271.30 / 286.09 ±11.39 / 302.40 ms │ 266.41 / 270.67 ±3.54 / 274.47 ms │ +1.06x faster │
│ QQuery 16 │ 632.86 / 637.91 ±4.36 / 645.86 ms │ 611.53 / 620.40 ±9.41 / 635.67 ms │ no change │
│ QQuery 17 │ 638.46 / 645.84 ±5.62 / 652.66 ms │ 617.47 / 622.14 ±4.89 / 630.17 ms │ no change │
│ QQuery 18 │ 1301.33 / 1321.07 ±13.12 / 1341.84 ms │ 1257.02 / 1275.05 ±16.83 / 1304.71 ms │ no change │
│ QQuery 19 │ 29.00 / 30.83 ±3.00 / 36.79 ms │ 29.35 / 30.48 ±0.97 / 32.09 ms │ no change │
│ QQuery 20 │ 522.79 / 528.84 ±5.81 / 539.12 ms │ 527.02 / 531.39 ±2.76 / 535.08 ms │ no change │
│ QQuery 21 │ 616.10 / 619.06 ±2.62 / 622.86 ms │ 608.87 / 617.52 ±7.71 / 627.12 ms │ no change │
│ QQuery 22 │ 1059.01 / 1081.68 ±11.78 / 1091.47 ms │ 1079.29 / 1099.10 ±12.80 / 1113.17 ms │ no change │
│ QQuery 23 │ 3151.75 / 3245.95 ±87.44 / 3357.81 ms │ 3195.98 / 3270.24 ±55.58 / 3328.52 ms │ no change │
│ QQuery 24 │ 42.05 / 43.91 ±2.65 / 49.14 ms │ 43.52 / 45.31 ±1.42 / 47.63 ms │ no change │
│ QQuery 25 │ 111.24 / 113.24 ±2.83 / 118.86 ms │ 111.48 / 114.88 ±3.47 / 119.23 ms │ no change │
│ QQuery 26 │ 42.44 / 43.20 ±0.96 / 45.09 ms │ 42.51 / 43.70 ±1.82 / 47.29 ms │ no change │
│ QQuery 27 │ 666.31 / 678.79 ±10.68 / 698.11 ms │ 674.44 / 684.68 ±7.53 / 695.30 ms │ no change │
│ QQuery 28 │ 3005.05 / 3026.40 ±10.88 / 3035.02 ms │ 3021.91 / 3060.89 ±30.83 / 3096.86 ms │ no change │
│ QQuery 29 │ 40.66 / 46.06 ±10.11 / 66.27 ms │ 39.66 / 42.57 ±4.89 / 52.30 ms │ +1.08x faster │
│ QQuery 30 │ 304.82 / 324.65 ±20.11 / 362.41 ms │ 300.79 / 304.10 ±2.72 / 306.81 ms │ +1.07x faster │
│ QQuery 31 │ 289.84 / 297.03 ±4.66 / 303.04 ms │ 282.42 / 295.54 ±10.68 / 313.43 ms │ no change │
│ QQuery 32 │ 964.70 / 985.64 ±16.38 / 1009.78 ms │ 914.35 / 930.08 ±14.20 / 956.64 ms │ +1.06x faster │
│ QQuery 33 │ 1439.97 / 1527.46 ±44.58 / 1564.29 ms │ 1414.08 / 1468.94 ±51.46 / 1540.11 ms │ no change │
│ QQuery 34 │ 1443.14 / 1488.88 ±37.07 / 1541.03 ms │ 1505.85 / 1553.51 ±30.27 / 1591.77 ms │ no change │
│ QQuery 35 │ 321.61 / 354.65 ±38.07 / 429.07 ms │ 278.63 / 297.21 ±19.94 / 332.85 ms │ +1.19x faster │
│ QQuery 36 │ 65.38 / 69.56 ±2.39 / 72.50 ms │ 67.85 / 72.19 ±3.03 / 76.49 ms │ no change │
│ QQuery 37 │ 35.15 / 38.34 ±4.14 / 46.47 ms │ 37.55 / 43.44 ±7.10 / 53.02 ms │ 1.13x slower │
│ QQuery 38 │ 40.91 / 44.69 ±5.33 / 55.07 ms │ 44.83 / 50.03 ±3.57 / 54.49 ms │ 1.12x slower │
│ QQuery 39 │ 139.86 / 150.22 ±9.02 / 165.53 ms │ 146.75 / 162.10 ±8.24 / 171.62 ms │ 1.08x slower │
│ QQuery 40 │ 14.02 / 16.05 ±3.83 / 23.70 ms │ 14.85 / 18.61 ±5.83 / 30.20 ms │ 1.16x slower │
│ QQuery 41 │ 13.54 / 18.47 ±7.51 / 33.19 ms │ 14.46 / 15.74 ±1.39 / 18.41 ms │ +1.17x faster │
│ QQuery 42 │ 13.71 / 14.02 ±0.18 / 14.25 ms │ 13.92 / 14.04 ±0.11 / 14.22 ms │ no change │
└───────────┴───────────────────────────────────────┴───────────────────────────────────────┴───────────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━┓
┃ Benchmark Summary ┃ ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━┩
│ Total Time (HEAD) │ 20127.15ms │
│ Total Time (ensure-requirements) │ 20052.11ms │
│ Average Time (HEAD) │ 468.07ms │
│ Average Time (ensure-requirements) │ 466.33ms │
│ Queries Faster │ 6 │
│ Queries Slower │ 5 │
│ Queries with No Change │ 32 │
│ Queries with Failure │ 0 │
└────────────────────────────────────┴────────────┘
```
</p>
</details>
<details><summary>Resource Usage</summary>
**clickbench_partitioned — base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 105.0s |
| Peak memory | 30.7 GiB |
| Avg memory | 23.2 GiB |
| CPU user | 1045.8s |
| CPU sys | 68.7s |
| Peak spill | 0 B |
**clickbench_partitioned — branch**
| Metric | Value |
|--------|-------|
| Wall time | 105.0s |
| Peak memory | 31.1 GiB |
| Avg memory | 23.0 GiB |
| CPU user | 1038.7s |
| CPU sys | 69.4s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "adriangbot (via GitHub)" <gi...@apache.org>.
adriangbot commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521733540
🤖 Criterion benchmark completed (GKE) | [trigger](https://github.com/apache/datafusion/pull/21976#issuecomment-4521087701)
**Instance:** `c4a-highmem-16` (12 vCPU / 65 GiB)
<details><summary>CPU Details (lscpu)</summary>
```
Architecture: aarch64
CPU op-mode(s): 64-bit
Byte Order: Little Endian
CPU(s): 16
On-line CPU(s) list: 0-15
Vendor ID: ARM
Model name: Neoverse-V2
Model: 1
Thread(s) per core: 1
Core(s) per cluster: 16
Socket(s): -
Cluster(s): 1
Stepping: r0p1
BogoMIPS: 2000.00
Flags: fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache: 1 MiB (16 instances)
L1i cache: 1 MiB (16 instances)
L2 cache: 32 MiB (16 instances)
L3 cache: 80 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0-15
Vulnerability Gather data sampling: Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Not affected
Vulnerability Mds: Not affected
Vulnerability Meltdown: Not affected
Vulnerability Mmio stale data: Not affected
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Not affected
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1: Mitigation; __user pointer sanitization
Vulnerability Spectre v2: Mitigation; CSV2, BHB
Vulnerability Srbds: Not affected
Vulnerability Tsa: Not affected
Vulnerability Tsx async abort: Not affected
Vulnerability Vmscape: Not affected
```
</details>
<details><summary>Details</summary>
<p>
```
group ensure-requirements main
----- ------------------- ----
logical_aggregate_with_join 1.03 429.4±1.57µs ? ?/sec 1.00 416.9±1.00µs ? ?/sec
logical_correlated_subquery_exists 1.02 254.9±0.61µs ? ?/sec 1.00 250.1±0.73µs ? ?/sec
logical_correlated_subquery_in 1.02 255.7±1.54µs ? ?/sec 1.00 250.5±0.69µs ? ?/sec
logical_distinct_many_columns 1.01 320.1±0.65µs ? ?/sec 1.00 318.0±0.85µs ? ?/sec
logical_join_4_with_agg_and_filter 1.00 238.8±0.87µs ? ?/sec 1.00 237.9±0.77µs ? ?/sec
logical_join_8_with_agg_sort_limit 1.00 404.8±1.41µs ? ?/sec 1.02 412.2±1.58µs ? ?/sec
logical_join_chain_16 1.00 660.1±2.81µs ? ?/sec 1.01 667.1±5.02µs ? ?/sec
logical_join_chain_4 1.00 114.1±0.30µs ? ?/sec 1.01 115.1±0.38µs ? ?/sec
logical_join_chain_8 1.00 241.5±0.76µs ? ?/sec 1.00 242.5±0.72µs ? ?/sec
logical_multiple_subqueries 1.00 457.4±2.00µs ? ?/sec 1.00 457.1±1.36µs ? ?/sec
logical_nested_cte_4_levels 1.01 239.4±1.08µs ? ?/sec 1.00 237.8±1.26µs ? ?/sec
logical_plan_struct_join_agg_sort 1.00 171.6±0.69µs ? ?/sec 1.00 172.1±0.84µs ? ?/sec
logical_plan_tpcds_all 1.01 74.6±0.71ms ? ?/sec 1.00 74.0±0.12ms ? ?/sec
logical_plan_tpch_all 1.00 5.8±0.06ms ? ?/sec 1.01 5.8±0.02ms ? ?/sec
logical_scalar_subquery 1.02 274.7±1.43µs ? ?/sec 1.00 270.3±1.22µs ? ?/sec
logical_select_all_from_1000 1.04 8.3±0.03ms ? ?/sec 1.00 8.0±0.02ms ? ?/sec
logical_select_one_from_700 1.01 279.4±1.59µs ? ?/sec 1.00 277.3±1.87µs ? ?/sec
logical_trivial_join_high_numbered_columns 1.03 262.4±0.66µs ? ?/sec 1.00 255.7±0.69µs ? ?/sec
logical_trivial_join_low_numbered_columns 1.03 249.9±1.31µs ? ?/sec 1.00 243.8±0.64µs ? ?/sec
logical_union_4_branches 1.02 363.1±1.29µs ? ?/sec 1.00 355.5±1.15µs ? ?/sec
logical_union_8_branches 1.02 690.1±2.11µs ? ?/sec 1.00 679.1±2.05µs ? ?/sec
logical_wide_aggregate_100_exprs 1.00 3.7±0.01ms ? ?/sec 1.01 3.7±0.02ms ? ?/sec
logical_wide_case_50_exprs 1.00 1671.5±5.76µs ? ?/sec 1.00 1665.7±8.28µs ? ?/sec
logical_wide_filter_200_predicates 1.01 1302.0±7.45µs ? ?/sec 1.00 1293.2±6.87µs ? ?/sec
logical_wide_filter_50_predicates 1.00 373.4±1.77µs ? ?/sec 1.00 371.9±1.76µs ? ?/sec
optimizer_correlated_exists 1.02 256.4±1.13µs ? ?/sec 1.00 252.5±1.07µs ? ?/sec
optimizer_join_4_with_agg_filter 1.00 432.4±1.08µs ? ?/sec 1.01 436.1±1.46µs ? ?/sec
optimizer_join_chain_4 1.00 178.2±0.46µs ? ?/sec 1.01 179.1±0.59µs ? ?/sec
optimizer_join_chain_8 1.00 557.0±0.83µs ? ?/sec 1.00 557.4±0.83µs ? ?/sec
optimizer_select_all_from_1000 1.01 4.6±0.02ms ? ?/sec 1.00 4.6±0.02ms ? ?/sec
optimizer_select_one_from_700 1.02 264.2±0.61µs ? ?/sec 1.00 259.7±0.62µs ? ?/sec
optimizer_tpcds_all 1.01 259.6±0.42ms ? ?/sec 1.00 258.1±0.33ms ? ?/sec
optimizer_tpch_all 1.00 14.7±0.04ms ? ?/sec 1.00 14.7±0.03ms ? ?/sec
optimizer_wide_aggregate_100 1.00 2.1±0.00ms ? ?/sec 1.01 2.1±0.00ms ? ?/sec
optimizer_wide_filter_200 1.02 3.9±0.01ms ? ?/sec 1.00 3.9±0.01ms ? ?/sec
physical_intersection 1.02 558.3±2.71µs ? ?/sec 1.00 545.4±1.99µs ? ?/sec
physical_join_consider_sort 1.02 997.2±11.57µs ? ?/sec 1.00 977.4±2.18µs ? ?/sec
physical_join_distinct 1.02 243.4±0.54µs ? ?/sec 1.00 238.1±0.74µs ? ?/sec
physical_many_self_joins 1.01 7.8±0.03ms ? ?/sec 1.00 7.7±0.02ms ? ?/sec
physical_plan_clickbench_all 1.00 120.6±0.48ms ? ?/sec 1.03 123.8±0.38ms ? ?/sec
physical_plan_clickbench_q1 1.00 1284.0±4.14µs ? ?/sec 1.02 1304.6±7.88µs ? ?/sec
physical_plan_clickbench_q10 1.00 1893.7±7.08µs ? ?/sec 1.02 1928.8±10.13µs ? ?/sec
physical_plan_clickbench_q11 1.00 2.1±0.01ms ? ?/sec 1.01 2.1±0.02ms ? ?/sec
physical_plan_clickbench_q12 1.00 2.2±0.01ms ? ?/sec 1.01 2.2±0.01ms ? ?/sec
physical_plan_clickbench_q13 1.00 1934.8±5.58µs ? ?/sec 1.01 1954.1±8.21µs ? ?/sec
physical_plan_clickbench_q14 1.00 2.1±0.01ms ? ?/sec 1.01 2.1±0.02ms ? ?/sec
physical_plan_clickbench_q15 1.00 2.0±0.02ms ? ?/sec 1.01 2.0±0.01ms ? ?/sec
physical_plan_clickbench_q16 1.00 1659.3±14.88µs ? ?/sec 1.01 1673.5±9.58µs ? ?/sec
physical_plan_clickbench_q17 1.00 1687.8±4.88µs ? ?/sec 1.02 1718.2±21.44µs ? ?/sec
physical_plan_clickbench_q18 1.00 1559.5±5.03µs ? ?/sec 1.01 1576.7±9.82µs ? ?/sec
physical_plan_clickbench_q19 1.00 1922.6±6.60µs ? ?/sec 1.00 1927.7±9.35µs ? ?/sec
physical_plan_clickbench_q2 1.00 1708.8±6.65µs ? ?/sec 1.01 1722.5±109.74µs ? ?/sec
physical_plan_clickbench_q20 1.00 1498.4±4.89µs ? ?/sec 1.00 1492.6±6.17µs ? ?/sec
physical_plan_clickbench_q21 1.00 1716.0±5.25µs ? ?/sec 1.00 1709.1±9.17µs ? ?/sec
physical_plan_clickbench_q22 1.01 2.1±0.01ms ? ?/sec 1.00 2.1±0.01ms ? ?/sec
physical_plan_clickbench_q23 1.01 2.2±0.01ms ? ?/sec 1.00 2.2±0.01ms ? ?/sec
physical_plan_clickbench_q24 1.00 5.5±0.02ms ? ?/sec 1.00 5.5±0.03ms ? ?/sec
physical_plan_clickbench_q25 1.00 1838.6±5.71µs ? ?/sec 1.02 1872.5±5.31µs ? ?/sec
physical_plan_clickbench_q26 1.00 1678.8±5.33µs ? ?/sec 1.02 1704.2±20.80µs ? ?/sec
physical_plan_clickbench_q27 1.02 1897.5±5.93µs ? ?/sec 1.00 1866.8±6.65µs ? ?/sec
physical_plan_clickbench_q28 1.01 2.3±0.01ms ? ?/sec 1.00 2.3±0.01ms ? ?/sec
physical_plan_clickbench_q29 1.01 2.5±0.01ms ? ?/sec 1.00 2.4±0.01ms ? ?/sec
physical_plan_clickbench_q3 1.00 1564.4±5.51µs ? ?/sec 1.00 1571.0±27.17µs ? ?/sec
physical_plan_clickbench_q30 1.00 14.8±0.06ms ? ?/sec 1.00 14.8±0.06ms ? ?/sec
physical_plan_clickbench_q31 1.01 2.4±0.01ms ? ?/sec 1.00 2.3±0.01ms ? ?/sec
physical_plan_clickbench_q32 1.01 2.4±0.00ms ? ?/sec 1.00 2.3±0.01ms ? ?/sec
physical_plan_clickbench_q33 1.01 1918.1±5.66µs ? ?/sec 1.00 1893.5±5.79µs ? ?/sec
physical_plan_clickbench_q34 1.00 1662.0±5.00µs ? ?/sec 1.01 1670.7±7.47µs ? ?/sec
physical_plan_clickbench_q35 1.00 1725.7±6.80µs ? ?/sec 1.00 1732.0±6.65µs ? ?/sec
physical_plan_clickbench_q36 1.00 2.0±0.01ms ? ?/sec 1.01 2.0±0.01ms ? ?/sec
physical_plan_clickbench_q37 1.00 2.4±0.01ms ? ?/sec 1.01 2.4±0.01ms ? ?/sec
physical_plan_clickbench_q38 1.00 2.4±0.02ms ? ?/sec 1.00 2.4±0.01ms ? ?/sec
physical_plan_clickbench_q39 1.00 2.5±0.01ms ? ?/sec 1.00 2.5±0.01ms ? ?/sec
physical_plan_clickbench_q4 1.00 1387.7±6.08µs ? ?/sec 1.00 1390.5±5.67µs ? ?/sec
physical_plan_clickbench_q40 1.00 3.2±0.01ms ? ?/sec 1.00 3.2±0.02ms ? ?/sec
physical_plan_clickbench_q41 1.00 2.7±0.01ms ? ?/sec 1.00 2.7±0.01ms ? ?/sec
physical_plan_clickbench_q42 1.01 3.0±0.01ms ? ?/sec 1.00 2.9±0.02ms ? ?/sec
physical_plan_clickbench_q43 1.00 3.1±0.02ms ? ?/sec 1.00 3.1±0.02ms ? ?/sec
physical_plan_clickbench_q44 1.02 1501.9±12.49µs ? ?/sec 1.00 1471.6±9.76µs ? ?/sec
physical_plan_clickbench_q45 1.01 1487.1±5.28µs ? ?/sec 1.00 1474.7±6.20µs ? ?/sec
physical_plan_clickbench_q46 1.00 1767.7±6.34µs ? ?/sec 1.00 1773.1±5.57µs ? ?/sec
physical_plan_clickbench_q47 1.00 2.5±0.01ms ? ?/sec 1.00 2.5±0.01ms ? ?/sec
physical_plan_clickbench_q48 1.02 2.8±0.01ms ? ?/sec 1.00 2.7±0.02ms ? ?/sec
physical_plan_clickbench_q49 1.03 2.8±0.01ms ? ?/sec 1.00 2.7±0.01ms ? ?/sec
physical_plan_clickbench_q5 1.01 1527.6±5.63µs ? ?/sec 1.00 1519.7±5.64µs ? ?/sec
physical_plan_clickbench_q50 1.02 2.7±0.01ms ? ?/sec 1.00 2.6±0.01ms ? ?/sec
physical_plan_clickbench_q51 1.00 1852.4±6.15µs ? ?/sec 1.01 1871.5±7.79µs ? ?/sec
physical_plan_clickbench_q6 1.00 1512.7±6.69µs ? ?/sec 1.02 1538.4±4.72µs ? ?/sec
physical_plan_clickbench_q7 1.00 1338.8±5.46µs ? ?/sec 1.02 1369.1±5.96µs ? ?/sec
physical_plan_clickbench_q8 1.00 1842.0±8.18µs ? ?/sec 1.02 1875.3±9.03µs ? ?/sec
physical_plan_clickbench_q9 1.00 1796.3±5.06µs ? ?/sec 1.02 1834.5±7.84µs ? ?/sec
physical_plan_struct_join_agg_sort 1.03 1260.9±4.40µs ? ?/sec 1.00 1226.7±3.37µs ? ?/sec
physical_plan_tpcds_all 1.00 660.0±1.06ms ? ?/sec 1.00 657.2±1.06ms ? ?/sec
physical_plan_tpch_all 1.00 43.1±0.09ms ? ?/sec 1.00 43.2±0.12ms ? ?/sec
physical_plan_tpch_q1 1.00 1420.1±2.53µs ? ?/sec 1.00 1424.4±4.96µs ? ?/sec
physical_plan_tpch_q10 1.00 2.7±0.01ms ? ?/sec 1.00 2.7±0.00ms ? ?/sec
physical_plan_tpch_q11 1.00 2.1±0.00ms ? ?/sec 1.00 2.1±0.00ms ? ?/sec
physical_plan_tpch_q12 1.00 1132.8±2.42µs ? ?/sec 1.01 1147.5±3.24µs ? ?/sec
physical_plan_tpch_q13 1.00 874.3±2.15µs ? ?/sec 1.02 887.7±1.33µs ? ?/sec
physical_plan_tpch_q14 1.00 1306.5±4.08µs ? ?/sec 1.02 1338.1±3.54µs ? ?/sec
physical_plan_tpch_q16 1.00 1457.4±4.77µs ? ?/sec 1.02 1480.8±4.28µs ? ?/sec
physical_plan_tpch_q17 1.00 1606.9±3.94µs ? ?/sec 1.02 1638.6±11.83µs ? ?/sec
physical_plan_tpch_q18 1.00 1740.5±3.26µs ? ?/sec 1.01 1758.6±14.00µs ? ?/sec
physical_plan_tpch_q19 1.00 1655.9±2.54µs ? ?/sec 1.01 1669.5±2.83µs ? ?/sec
physical_plan_tpch_q2 1.00 3.9±0.01ms ? ?/sec 1.00 3.9±0.01ms ? ?/sec
physical_plan_tpch_q20 1.00 2.1±0.00ms ? ?/sec 1.00 2.1±0.00ms ? ?/sec
physical_plan_tpch_q21 1.01 2.8±0.01ms ? ?/sec 1.00 2.8±0.01ms ? ?/sec
physical_plan_tpch_q22 1.00 1487.7±2.99µs ? ?/sec 1.01 1501.0±4.06µs ? ?/sec
physical_plan_tpch_q3 1.00 1799.4±3.73µs ? ?/sec 1.00 1800.0±4.14µs ? ?/sec
physical_plan_tpch_q4 1.00 1088.7±3.81µs ? ?/sec 1.01 1099.2±5.39µs ? ?/sec
physical_plan_tpch_q5 1.00 2.4±0.00ms ? ?/sec 1.00 2.4±0.01ms ? ?/sec
physical_plan_tpch_q6 1.00 587.4±2.00µs ? ?/sec 1.01 591.6±1.86µs ? ?/sec
physical_plan_tpch_q7 1.00 2.9±0.01ms ? ?/sec 1.00 2.9±0.01ms ? ?/sec
physical_plan_tpch_q8 1.00 3.9±0.01ms ? ?/sec 1.00 3.9±0.01ms ? ?/sec
physical_plan_tpch_q9 1.00 2.7±0.00ms ? ?/sec 1.00 2.7±0.01ms ? ?/sec
physical_select_aggregates_from_200 1.01 13.5±0.04ms ? ?/sec 1.00 13.4±0.04ms ? ?/sec
physical_select_all_from_1000 1.03 17.4±0.05ms ? ?/sec 1.00 16.9±0.05ms ? ?/sec
physical_select_one_from_700 1.01 735.4±2.33µs ? ?/sec 1.00 726.7±1.96µs ? ?/sec
physical_sorted_union_order_by_10_int64 1.00 4.2±0.01ms ? ?/sec 1.00 4.2±0.01ms ? ?/sec
physical_sorted_union_order_by_10_uint64 1.01 8.3±0.03ms ? ?/sec 1.00 8.3±0.02ms ? ?/sec
physical_sorted_union_order_by_50_int64 1.00 104.8±0.36ms ? ?/sec 1.00 104.4±0.24ms ? ?/sec
physical_sorted_union_order_by_50_uint64 1.00 357.7±0.78ms ? ?/sec 1.00 357.1±0.76ms ? ?/sec
physical_theta_join_consider_sort 1.02 1033.2±3.04µs ? ?/sec 1.00 1008.6±3.18µs ? ?/sec
physical_unnest_to_join 1.04 648.1±4.97µs ? ?/sec 1.00 623.6±2.11µs ? ?/sec
physical_window_function_partition_by_12_on_values 1.01 697.4±1.42µs ? ?/sec 1.00 690.2±1.76µs ? ?/sec
physical_window_function_partition_by_30_on_values 1.00 1386.5±2.83µs ? ?/sec 1.00 1387.6±3.25µs ? ?/sec
physical_window_function_partition_by_4_on_values 1.00 414.8±1.28µs ? ?/sec 1.00 413.1±1.19µs ? ?/sec
physical_window_function_partition_by_7_on_values 1.01 516.6±1.86µs ? ?/sec 1.00 514.0±1.42µs ? ?/sec
physical_window_function_partition_by_8_on_values 1.01 559.7±1.70µs ? ?/sec 1.00 553.3±1.47µs ? ?/sec
with_param_values_many_columns 1.01 436.7±3.03µs ? ?/sec 1.00 434.5±2.44µs ? ?/sec
```
</p>
</details>
<details><summary>Resource Usage</summary>
**base (merge-base)**
| Metric | Value |
|--------|-------|
| Wall time | 1555.3s |
| Peak memory | 19.8 GiB |
| Avg memory | 19.7 GiB |
| CPU user | 1870.6s |
| CPU sys | 1.9s |
| Peak spill | 0 B |
**branch**
| Metric | Value |
|--------|-------|
| Wall time | 1545.3s |
| Peak memory | 19.8 GiB |
| Avg memory | 19.8 GiB |
| CPU user | 1862.0s |
| CPU sys | 1.6s |
| Peak spill | 0 B |
</details>
---
[File an issue](https://github.com/adriangb/datafusion-benchmarking/issues) against this benchmark runner
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4521724833
Looks like there is a logical conflict -- I just merged up from main and the compile now fails
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3214790498
##########
datafusion/physical-optimizer/src/ensure_requirements/new_tests.rs:
##########
@@ -0,0 +1,782 @@
+// Licensed to the Apache Software Foundation (ASF) under one
Review Comment:
Why is this a separate module? If we want tests in a separate module, I think putting it in an integration test might be easier to find (datafusion/physical-optimizer/tests/ensure_requirements.rs perhaps?)
##########
datafusion/physical-optimizer/src/ensure_requirements/mod.rs:
##########
@@ -0,0 +1,1804 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! [`EnsureRequirements`] optimizer rule that enforces both distribution and
+//! sorting requirements in a **single bottom-up pass**.
+//!
+//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting`
+//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements`
+//! and Presto/Trino's `AddExchanges`.
+//!
+//! # Motivation
+//!
+//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`)
+//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts`
+//! can break distribution invariants established by `EnforceDistribution`,
+//! because `SortExec.preserve_partitioning` couples sorting and distribution
+//! decisions. See <https://github.com/apache/datafusion/issues/21973> for details.
+//!
+//! # Architecture
+//!
+//! ```text
+//! EnsureRequirements::optimize(plan)
+//! │
+//! ├─ Phase 1 (optional): reorder_join_keys (top-down)
+//! │ └─ Same as existing adjust_input_keys_ordering
+//! │
+//! └─ Phase 2: ensure_requirements (single bottom-up pass)
+//! └─ For each node (bottom-up), for each child:
+//! Step 1: Ensure distribution requirement
+//! └─ Add RepartitionExec / CoalescePartitionsExec / SortPreservingMergeExec
+//! Step 2: Ensure ordering requirement (distribution-aware)
+//! └─ Add SortExec with correct preserve_partitioning + SPM if needed
+//! ```
+//!
+//! # Key Properties
+//!
+//! - **Idempotent**: Running the rule twice produces the same plan.
+//! - **Distribution before sorting**: For each child, distribution is resolved
+//! before ordering, so sorting decisions always have full distribution context.
+//! - **No separate `pushdown_sorts`**: Sort pushdown is implicit — the bottom-up
+//! pass only adds `SortExec` where the child doesn't already satisfy the
+//! ordering requirement, naturally placing sorts at the deepest valid position.
+
+use std::sync::Arc;
+
+use crate::PhysicalOptimizerRule;
+
+use datafusion_common::Result;
+use datafusion_common::config::ConfigOptions;
+use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_physical_plan::ExecutionPlan;
+
+// Internal functions used directly instead of calling EnforceDistribution/EnforceSorting
+// as opaque boxes. This gives us control over the pass ordering and enables
+// future merging into a true single-pass architecture.
+
+// For the no-pushdown variant (Phase 3)
+use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+};
+use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, ensure_sorting,
+ parallelize_sorts, replace_with_partial_sort,
+};
+
+/// Optimizer rule that enforces both distribution and sorting requirements.
+///
+/// This rule combines the functionality of `EnforceDistribution` and
+/// `EnforceSorting` into a coordinated sequence where distribution is
+/// always settled before sorting for each operator, preventing the
+/// non-idempotent interactions between the two separate rules.
+///
+/// See [module level documentation](self) for more details.
+#[derive(Default, Debug)]
+pub struct EnsureRequirements {}
+
+impl EnsureRequirements {
+ /// Create a new `EnsureRequirements` optimizer rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirements {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Phase 1: Join key reordering (top-down, from EnforceDistribution)
+ use crate::enforce_distribution::{
+ PlanWithKeyRequirements, adjust_input_keys_ordering,
+ };
+ let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down_join_key_reordering {
+ let ctx = PlanWithKeyRequirements::new_default(plan);
+ ctx.transform_down(adjust_input_keys_ordering).data()?.plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+
+ // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass)
+ // For each node: distribution first, then sorting.
+ use crate::enforce_distribution::{DistributionContext, ensure_distribution};
+ use crate::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting};
+
+ // Step 2a: Distribution enforcement (bottom-up)
+ let dist_ctx = DistributionContext::new_default(plan);
+ let dist_ctx = dist_ctx
+ .transform_up(|ctx| ensure_distribution(ctx, config))
+ .data()?;
+
+ // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan
+ let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan);
+ let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data;
+
+ // Phase 3: Optimization passes
+ // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort)
+ use crate::enforce_sorting::{
+ PlanWithCorrespondingCoalescePartitions, parallelize_sorts,
+ replace_with_partial_sort,
+ };
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ sort_ctx.plan
+ };
+
+ // 3b: Order-preserving variants
+ use crate::enforce_sorting::replace_with_order_preserving_variants::{
+ OrderPreservationContext, replace_with_order_preserving_variants,
+ };
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // 3c: Sort pushdown (distribution-aware)
+ use crate::enforce_sorting::sort_pushdown::{
+ SortPushDown, assign_initial_requirements, pushdown_sorts,
+ };
+ let mut sort_pushdown = SortPushDown::new_default(plan);
+ assign_initial_requirements(&mut sort_pushdown);
+ let adjusted = pushdown_sorts(sort_pushdown)?;
+
+ // 3d: Partial sort
+ adjusted
+ .plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirements"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+/// Phase 3 variant: no `pushdown_sorts`, sort placement handled entirely
+/// by bottom-up passes. Currently experimental — some plan shapes differ
+/// from the `pushdown_sorts` variant (less optimal but still correct).
+#[derive(Default, Debug)]
+pub struct EnsureRequirementsNoPushdown {}
+
+impl EnsureRequirementsNoPushdown {
+ /// Create a new rule.
+ pub fn new() -> Self {
+ Self {}
+ }
+}
+
+impl PhysicalOptimizerRule for EnsureRequirementsNoPushdown {
+ fn optimize(
+ &self,
+ plan: Arc<dyn ExecutionPlan>,
+ config: &ConfigOptions,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ // Step 1: Distribution enforcement
+ use crate::enforce_distribution::{
+ DistributionContext as DistCtx, PlanWithKeyRequirements as KeyReqs,
+ adjust_input_keys_ordering as adj_keys, ensure_distribution as ensure_dist,
+ };
+ let top_down = config.optimizer.top_down_join_key_reordering;
+ let plan = if top_down {
+ KeyReqs::new_default(plan)
+ .transform_down(adj_keys)
+ .data()?
+ .plan
+ } else {
+ use crate::enforce_distribution::reorder_join_keys_to_inputs;
+ plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?)))
+ .data()?
+ };
+ let dist_ctx = DistCtx::new_default(plan);
+ let plan = dist_ctx
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 2: ensure_sorting (bottom-up, NO pushdown_sorts)
+ let plan_requirements = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted = plan_requirements.transform_up(ensure_sorting)?.data;
+
+ // Step 3: parallelize_sorts (optional)
+ let plan = if config.optimizer.repartition_sorts {
+ let ctx = PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan);
+ ctx.transform_up(parallelize_sorts).data()?.plan
+ } else {
+ adjusted.plan
+ };
+
+ // Step 4: order-preserving variants
+ let ctx = OrderPreservationContext::new_default(plan);
+ let plan = ctx
+ .transform_up(|c| {
+ replace_with_order_preserving_variants(c, false, true, config)
+ })
+ .data()?
+ .plan;
+
+ // Step 5: partial sort
+ let plan = plan
+ .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?)))
+ .data()?;
+
+ // NO pushdown_sorts — sort placement is purely bottom-up.
+ // Step 6: Final distribution enforcement
+ let dist_ctx2 = DistCtx::new_default(plan);
+ let plan = dist_ctx2
+ .transform_up(|ctx| ensure_dist(ctx, config))
+ .data()?
+ .plan;
+
+ // Step 7: Fix any sorting violations the final distribution pass introduced.
+ let sort_ctx2 = PlanWithCorrespondingSort::new_default(plan);
+ let adjusted2 = sort_ctx2.transform_up(ensure_sorting)?.data;
+
+ Ok(adjusted2.plan)
+ }
+
+ fn name(&self) -> &str {
+ "EnsureRequirementsNoPushdown"
+ }
+
+ fn schema_check(&self) -> bool {
+ true
+ }
+}
+
+#[cfg(test)]
+mod new_tests;
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+
+ use std::sync::Arc;
+
+ use arrow::compute::SortOptions;
+ use arrow::datatypes::{DataType, Field, Schema};
+ use datafusion_common::Result;
+ use datafusion_common::tree_node::TreeNodeRecursion;
+ use datafusion_physical_expr::expressions::Column;
+ use datafusion_physical_expr::{
+ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr,
+ };
+ use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType};
+ use datafusion_physical_plan::limit::GlobalLimitExec;
+ use datafusion_physical_plan::sorts::sort::SortExec;
+ use datafusion_physical_plan::union::UnionExec;
+ use datafusion_physical_plan::{
+ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties,
+ Partitioning, PlanProperties, SendableRecordBatchStream,
+ };
+
+ use crate::enforce_distribution::EnforceDistribution;
+ use crate::output_requirements::OutputRequirementExec;
+ use crate::sanity_checker::SanityCheckPlan;
+
+ use datafusion_common::{JoinType, NullEquality};
+ use datafusion_physical_expr::Distribution;
+ use datafusion_physical_expr_common::sort_expr::OrderingRequirements;
+ use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
+ use datafusion_physical_plan::joins::{
+ HashJoinExec, PartitionMode, SortMergeJoinExec,
+ };
+ use datafusion_physical_plan::projection::ProjectionExec;
+ use datafusion_physical_plan::repartition::RepartitionExec;
+ use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
+
+ /// Mock ExecutionPlan with configurable partition count and output ordering.
+ #[derive(Debug)]
+ struct MockMultiPartitionExec {
+ properties: Arc<PlanProperties>,
+ }
+
+ impl MockMultiPartitionExec {
+ fn new(partition_count: usize) -> Self {
+ let schema = Arc::new(Schema::new(vec![
+ Field::new("a", DataType::Int64, false),
+ Field::new("b", DataType::Int64, false),
+ ]));
+ let mut eq = EquivalenceProperties::new(Arc::clone(&schema));
+ if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )]) {
+ eq.add_orderings(vec![ordering.into_iter().collect::<Vec<_>>()]);
+ }
+ let properties = PlanProperties::new(
+ eq,
+ Partitioning::UnknownPartitioning(partition_count),
+ EmissionType::Incremental,
+ Boundedness::Bounded,
+ );
+ Self {
+ properties: Arc::new(properties),
+ }
+ }
+ }
+
+ impl DisplayAs for MockMultiPartitionExec {
+ fn fmt_as(
+ &self,
+ _t: DisplayFormatType,
+ f: &mut std::fmt::Formatter,
+ ) -> std::fmt::Result {
+ write!(f, "MockMultiPartitionExec")
+ }
+ }
+
+ impl ExecutionPlan for MockMultiPartitionExec {
+ fn name(&self) -> &str {
+ "MockMultiPartitionExec"
+ }
+ fn properties(&self) -> &Arc<PlanProperties> {
+ &self.properties
+ }
+ fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> {
+ vec![]
+ }
+ fn with_new_children(
+ self: Arc<Self>,
+ _children: Vec<Arc<dyn ExecutionPlan>>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ Ok(self)
+ }
+ fn apply_expressions(
+ &self,
+ _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result<TreeNodeRecursion>,
+ ) -> Result<TreeNodeRecursion> {
+ Ok(TreeNodeRecursion::Continue)
+ }
+ fn execute(
+ &self,
+ _partition: usize,
+ _context: Arc<datafusion_execution::TaskContext>,
+ ) -> Result<SendableRecordBatchStream> {
+ unimplemented!()
+ }
+ }
+
+ /// Helper: run EnsureRequirements and verify SanityCheckPlan passes
+ fn optimize_and_sanity_check(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Arc<dyn ExecutionPlan>> {
+ let config = ConfigOptions::default();
+ let optimized = EnsureRequirements::new().optimize(plan, &config)?;
+ // SanityCheckPlan must pass
+ SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?;
+ Ok(optimized)
+ }
+
+ /// Helper: verify idempotency — running twice produces the same plan
+ fn assert_idempotent(plan: Arc<dyn ExecutionPlan>) {
+ let config = ConfigOptions::default();
+ let p1 = EnsureRequirements::new()
+ .optimize(plan, &config)
+ .expect("first optimize failed");
+ let p2 = EnsureRequirements::new()
+ .optimize(Arc::clone(&p1), &config)
+ .expect("second optimize failed");
+
+ let s1 = datafusion_physical_plan::displayable(p1.as_ref())
+ .indent(true)
+ .to_string();
+ let s2 = datafusion_physical_plan::displayable(p2.as_ref())
+ .indent(true)
+ .to_string();
+ assert_eq!(
+ s1, s2,
+ "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}"
+ );
+
+ // Both must pass SanityCheckPlan
+ SanityCheckPlan::new()
+ .optimize(p1, &config)
+ .expect("SanityCheckPlan failed on first pass");
+ SanityCheckPlan::new()
+ .optimize(p2, &config)
+ .expect("SanityCheckPlan failed on second pass");
+ }
+
+ /// Multi-partition sort + GlobalLimitExec must produce valid plan.
+ /// Regression test for UXX0/HRC 502s (April 2026).
+ #[test]
+ fn test_multi_partition_sort_limit_sanity_check() {
+ let source = Arc::new(MockMultiPartitionExec::new(32));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Union with mixed partition counts + sort + limit.
+ #[test]
+ fn test_union_mixed_partitions_sort_limit() {
+ let live = Arc::new(MockMultiPartitionExec::new(32));
+ let historical = Arc::new(MockMultiPartitionExec::new(1));
+
+ let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(limit);
+ assert!(result.is_ok(), "SanityCheckPlan failed: {:?}", result.err());
+ }
+
+ /// Idempotency: multi-partition sort + limit
+ #[test]
+ fn test_idempotent_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ assert_idempotent(limit);
+ }
+
+ /// Idempotency: union with mixed partitions
+ #[test]
+ fn test_idempotent_union_mixed_partitions() {
+ let live = Arc::new(MockMultiPartitionExec::new(8));
+ let hist = Arc::new(MockMultiPartitionExec::new(1));
+ let union = UnionExec::try_new(vec![live as _, hist as _]).unwrap();
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, union));
+ let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(5)));
+
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Projection + multi-partition tests (pushdown_sorts trigger path)
+ // ========================================================================
+
+ /// ProjectionExec over multi-partition + sort DESC + limit.
+ /// This is the topology where pushdown_sorts pushes sort through projection
+ /// onto the multi-partition source. The optimizer must still produce a valid plan.
+ #[test]
+ fn test_projection_over_multi_partition_sort_limit() {
+ let source = Arc::new(MockMultiPartitionExec::new(16));
+ // Identity projection
+ let proj_exprs: Vec<(Arc<dyn PhysicalExpr>, String)> = vec![
+ (Arc::new(Column::new("a", 0)), "a".to_string()),
+ (Arc::new(Column::new("b", 1)), "b".to_string()),
+ ];
+ let projection =
+ Arc::new(ProjectionExec::try_new(proj_exprs, source as _).unwrap());
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, projection));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(21)));
+
+ let result = optimize_and_sanity_check(Arc::clone(&limit));
+ assert!(
+ result.is_ok(),
+ "SanityCheckPlan failed for projection over multi-partition: {:?}",
+ result.err()
+ );
+ assert_idempotent(limit);
+ }
+
+ // ========================================================================
+ // Single partition tests (no unnecessary operators)
+ // ========================================================================
+
+ /// Single partition source + sort + limit should NOT add SortPreservingMergeExec.
+ #[test]
+ fn test_single_partition_no_unnecessary_spm() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: true,
+ nulls_first: true,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ let optimized = optimize_and_sanity_check(limit).unwrap();
+ let plan_str = datafusion_physical_plan::displayable(optimized.as_ref())
+ .indent(true)
+ .to_string();
+
+ // Single partition should not have SortPreservingMergeExec
+ assert!(
+ !plan_str.contains("SortPreservingMergeExec"),
+ "Unnecessary SortPreservingMergeExec for single partition:\n{plan_str}"
+ );
+ }
+
+ /// Source already has correct ordering → should not add SortExec.
+ #[test]
+ fn test_sort_already_satisfied_no_extra_sort() {
+ let source = Arc::new(MockMultiPartitionExec::new(1));
+
+ // Sort ASC matches MockMultiPartitionExec's output ordering (a ASC)
+ let sort_expr = LexOrdering::new(vec![PhysicalSortExpr::new(
+ Arc::new(Column::new("a", 0)),
+ SortOptions {
+ descending: false,
+ nulls_first: false,
+ },
+ )])
+ .unwrap();
+
+ let sort = Arc::new(SortExec::new(sort_expr, source));
+ let limit: Arc<dyn ExecutionPlan> =
+ Arc::new(GlobalLimitExec::new(sort, 0, Some(10)));
+
+ let optimized = optimize_and_sanity_check(limit).unwrap();
+ let plan_str = datafusion_physical_plan::displayable(optimized.as_ref())
+ .indent(true)
+ .to_string();
+
+ // Sort should be eliminated since source already satisfies ordering
+ // The plan should just be limit + source (or limit + local limit + source)
+ assert!(
+ !plan_str.contains("SortExec: expr=[a@0 ASC"),
+ "Unnecessary SortExec when ordering already satisfied:\n{plan_str}"
+ );
+ }
+
+ // ========================================================================
Review Comment:
many of these tests look similar/ the same as what is in enforce distribution, etc but it is hard to tell what changed
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on code in PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#discussion_r3214871632
##########
datafusion/physical-optimizer/src/optimizer.rs:
##########
@@ -170,18 +169,12 @@ impl PhysicalOptimizer {
// those are handled by the later `FilterPushdown` rule.
// See `FilterPushdownPhase` for more details.
Arc::new(FilterPushdown::new()),
- // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution
- // requirements. Please make sure that the whole plan tree is determined before this rule.
- // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at
- // least one of the operators in the plan benefits from increased parallelism.
- Arc::new(EnforceDistribution::new()),
- // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule
+ // EnsureRequirements: merged EnforceDistribution + EnforceSorting into a
+ // single idempotent rule with distribution-aware pushdown_sorts.
+ // See https://github.com/apache/datafusion/issues/21973
Review Comment:
Will apply the suggestion in the next push (per alamb's plan I'll land a test-extraction PR first and rebase this on top).
On "distribution-aware pushdown_sorts": the old pushdown_sorts didn't propagate the parent's required_input_distribution, so a SortExec could get pushed below something requiring SinglePartition (e.g. under GlobalLimitExec) — giving an invalid plan because the sort then runs on multiple partitions. The fix propagates the requirement and inserts a SortPreservingMergeExec where needed. Will link to #21973 in the comment block when I update.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [datafusion]
Posted by "zhuqi-lucas (via GitHub)" <gi...@apache.org>.
zhuqi-lucas commented on PR #21976:
URL: https://github.com/apache/datafusion/pull/21976#issuecomment-4427088981
Test-extraction PR is up: #22117. Moves the four `enforce_distribution` / `enforce_sorting` / `enforce_sorting_monotonicity` / `replace_with_order_preserving_variants` integration tests to `datafusion/physical-optimizer/tests/`. Git detects all four as 94–99% renames, so the actual review surface is just the `Cargo.toml` dev-deps and a copy of `test_utils.rs` (kept in `tests/common/` so cargo doesn't treat it as its own test binary).
Once #22117 merges I'll rebase this PR on top — the resulting diff should be the function-name swap from `EnforceDistribution::new()` / `EnforceSorting::new()` to `EnsureRequirements::new()` on those moved files plus the genuinely new tests, which is the comparison you wanted to be able to do.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org