jayzhan211 commented on code in PR #26459: URL: https://github.com/apache/datafusion/pull/26459#discussion_r4324940656
########## datafusion/physical-optimizer/src/ensure_requirements/mod.rs: ########## @@ -69,6 +70,17 @@ //! <https://github.com/apache/datafusion/issues/21973>, where the old //! two-rule pipeline could regress a parallel sort plan into a serial one //! on pass 2. +//! - **Requirements are provable**: Every ordering requirement in the +//! output is satisfied according to the children's +//! [`EquivalenceProperties`], the same check `SanityCheckPlan` applies. +//! `pushdown_sorts` removes a sort when it can map the requirement onto a +//! descendant that satisfies it, but the operators in between can report +//! weaker properties than that mapping assumed (for example, a hash join +//! that forgets a probe-side column is constant). Phase 3 then adds a sort +//! where the requirement cannot be proven, so the rule never emits a plan +//! that it cannot show to be valid. Review Comment: This no longer applies: following your comment on the root cause, the EnsureRequirements step is removed and the PR fixes the join instead. ########## datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs: ########## @@ -232,6 +235,71 @@ pub fn replace_with_partial_sort( Ok(plan) } +/// Adds a [`SortExec`] below every node whose ordering requirement is not +/// satisfied by its child's equivalence properties. +/// +/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a +/// descendant that satisfies it. That mapping is computed separately from the +/// equivalence properties of the operators in between, and the two can +/// disagree. The plan may then be correct, but its properties cannot show it, +/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan +/// produced by `EnsureRequirements` valid by the same check. +/// +/// A requirement that a new [`SortExec`] still cannot satisfy is left as is, +/// for `SanityCheckPlan` to report. +/// +/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts +pub fn add_sort_for_unproven_ordering( + plan: Arc<dyn ExecutionPlan>, +) -> Result<Transformed<Arc<dyn ExecutionPlan>>> { + let requirements = plan.required_input_ordering(); + if requirements.iter().all(Option::is_none) { + return Ok(Transformed::no(plan)); + } + let mut children: Vec<_> = plan.children().into_iter().cloned().collect(); + let mut changed = false; + for (child, requirement) in children.iter_mut().zip(requirements) { + let Some(requirement) = requirement else { + continue; + }; + let requirement = requirement.into_single(); + let eq_properties = child.equivalence_properties(); + if eq_properties.ordering_satisfy_requirement(requirement.clone())? { + continue; + } + let sort_reqs: Vec<_> = requirement.clone().into(); + let sort_exprs = sort_reqs + .into_iter() + .filter(|req| eq_properties.is_expr_constant(&req.expr).is_none()) + .map(Into::into); + let Some(ordering) = LexOrdering::new(sort_exprs) else { + continue; + }; + let sort: Arc<dyn ExecutionPlan> = Arc::new( + SortExec::new(ordering, Arc::clone(child)).with_preserve_partitioning( + child.output_partitioning().partition_count() > 1, + ), Review Comment: This no longer applies: following your comment on the root cause, the EnsureRequirements step is removed and the PR fixes the join instead. ########## datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs: ########## @@ -232,6 +235,71 @@ pub fn replace_with_partial_sort( Ok(plan) } +/// Adds a [`SortExec`] below every node whose ordering requirement is not +/// satisfied by its child's equivalence properties. +/// +/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a +/// descendant that satisfies it. That mapping is computed separately from the +/// equivalence properties of the operators in between, and the two can +/// disagree. The plan may then be correct, but its properties cannot show it, +/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan +/// produced by `EnsureRequirements` valid by the same check. +/// +/// A requirement that a new [`SortExec`] still cannot satisfy is left as is, +/// for `SanityCheckPlan` to report. +/// +/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts +pub fn add_sort_for_unproven_ordering( + plan: Arc<dyn ExecutionPlan>, +) -> Result<Transformed<Arc<dyn ExecutionPlan>>> { + let requirements = plan.required_input_ordering(); + if requirements.iter().all(Option::is_none) { + return Ok(Transformed::no(plan)); + } + let mut children: Vec<_> = plan.children().into_iter().cloned().collect(); + let mut changed = false; + for (child, requirement) in children.iter_mut().zip(requirements) { + let Some(requirement) = requirement else { + continue; + }; + let requirement = requirement.into_single(); + let eq_properties = child.equivalence_properties(); + if eq_properties.ordering_satisfy_requirement(requirement.clone())? { + continue; + } + let sort_reqs: Vec<_> = requirement.clone().into(); Review Comment: This no longer applies: following your comment on the root cause, the EnsureRequirements step is removed and the PR fixes the join instead. ########## datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs: ########## @@ -232,6 +235,71 @@ pub fn replace_with_partial_sort( Ok(plan) } +/// Adds a [`SortExec`] below every node whose ordering requirement is not +/// satisfied by its child's equivalence properties. +/// +/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a +/// descendant that satisfies it. That mapping is computed separately from the +/// equivalence properties of the operators in between, and the two can +/// disagree. The plan may then be correct, but its properties cannot show it, +/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan +/// produced by `EnsureRequirements` valid by the same check. +/// +/// A requirement that a new [`SortExec`] still cannot satisfy is left as is, +/// for `SanityCheckPlan` to report. +/// +/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts +pub fn add_sort_for_unproven_ordering( + plan: Arc<dyn ExecutionPlan>, +) -> Result<Transformed<Arc<dyn ExecutionPlan>>> { + let requirements = plan.required_input_ordering(); + if requirements.iter().all(Option::is_none) { + return Ok(Transformed::no(plan)); + } + let mut children: Vec<_> = plan.children().into_iter().cloned().collect(); + let mut changed = false; + for (child, requirement) in children.iter_mut().zip(requirements) { + let Some(requirement) = requirement else { + continue; + }; + let requirement = requirement.into_single(); + let eq_properties = child.equivalence_properties(); + if eq_properties.ordering_satisfy_requirement(requirement.clone())? { + continue; + } + let sort_reqs: Vec<_> = requirement.clone().into(); + let sort_exprs = sort_reqs + .into_iter() + .filter(|req| eq_properties.is_expr_constant(&req.expr).is_none()) + .map(Into::into); + let Some(ordering) = LexOrdering::new(sort_exprs) else { + continue; + }; + let sort: Arc<dyn ExecutionPlan> = Arc::new( + SortExec::new(ordering, Arc::clone(child)).with_preserve_partitioning( + child.output_partitioning().partition_count() > 1, + ), + ); + if !sort + .equivalence_properties() + .ordering_satisfy_requirement(requirement)? + { + continue; Review Comment: This no longer applies: following your comment on the root cause, the EnsureRequirements step is removed and the PR fixes the join instead. ########## datafusion/physical-optimizer/src/ensure_requirements/mod.rs: ########## @@ -267,10 +279,15 @@ impl PhysicalOptimizerRule for EnsureRequirements { assign_initial_requirements(&mut sort_pushdown); let adjusted = pushdown_sorts(sort_pushdown)?; - // 3d: Partial sort - adjusted + // 3d: Sort where an ordering requirement cannot be proven + use super::enforce_sorting::add_sort_for_unproven_ordering; + let plan = adjusted .plan Review Comment: This no longer applies: following your comment on the root cause, the EnsureRequirements step is removed and the PR fixes the join instead. ########## datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs: ########## @@ -232,6 +235,71 @@ pub fn replace_with_partial_sort( Ok(plan) } +/// Adds a [`SortExec`] below every node whose ordering requirement is not +/// satisfied by its child's equivalence properties. +/// +/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a +/// descendant that satisfies it. That mapping is computed separately from the +/// equivalence properties of the operators in between, and the two can +/// disagree. The plan may then be correct, but its properties cannot show it, +/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan +/// produced by `EnsureRequirements` valid by the same check. +/// +/// A requirement that a new [`SortExec`] still cannot satisfy is left as is, +/// for `SanityCheckPlan` to report. +/// +/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts +pub fn add_sort_for_unproven_ordering( + plan: Arc<dyn ExecutionPlan>, +) -> Result<Transformed<Arc<dyn ExecutionPlan>>> { + let requirements = plan.required_input_ordering(); + if requirements.iter().all(Option::is_none) { + return Ok(Transformed::no(plan)); + } + let mut children: Vec<_> = plan.children().into_iter().cloned().collect(); + let mut changed = false; + for (child, requirement) in children.iter_mut().zip(requirements) { + let Some(requirement) = requirement else { + continue; + }; + let requirement = requirement.into_single(); + let eq_properties = child.equivalence_properties(); + if eq_properties.ordering_satisfy_requirement(requirement.clone())? { + continue; + } + let sort_reqs: Vec<_> = requirement.clone().into(); + let sort_exprs = sort_reqs + .into_iter() + .filter(|req| eq_properties.is_expr_constant(&req.expr).is_none()) + .map(Into::into); + let Some(ordering) = LexOrdering::new(sort_exprs) else { + continue; + }; + let sort: Arc<dyn ExecutionPlan> = Arc::new( + SortExec::new(ordering, Arc::clone(child)).with_preserve_partitioning( + child.output_partitioning().partition_count() > 1, + ), + ); + if !sort + .equivalence_properties() + .ordering_satisfy_requirement(requirement)? + { + continue; + } + *child = sort; + changed = true; + } + if changed { + plan.replace_children( Review Comment: This no longer applies: following your comment on the root cause, the EnsureRequirements step is removed and the PR fixes the join instead. ########## datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/mod.rs: ########## @@ -232,6 +235,71 @@ pub fn replace_with_partial_sort( Ok(plan) } +/// Adds a [`SortExec`] below every node whose ordering requirement is not +/// satisfied by its child's equivalence properties. +/// +/// [`pushdown_sorts`] removes a sort when it can map the requirement onto a +/// descendant that satisfies it. That mapping is computed separately from the +/// equivalence properties of the operators in between, and the two can +/// disagree. The plan may then be correct, but its properties cannot show it, +/// and `SanityCheckPlan` rejects it. Adding the sort back keeps every plan +/// produced by `EnsureRequirements` valid by the same check. +/// +/// A requirement that a new [`SortExec`] still cannot satisfy is left as is, +/// for `SanityCheckPlan` to report. +/// +/// [`pushdown_sorts`]: sort_pushdown::pushdown_sorts +pub fn add_sort_for_unproven_ordering( + plan: Arc<dyn ExecutionPlan>, +) -> Result<Transformed<Arc<dyn ExecutionPlan>>> { + let requirements = plan.required_input_ordering(); + if requirements.iter().all(Option::is_none) { + return Ok(Transformed::no(plan)); + } + let mut children: Vec<_> = plan.children().into_iter().cloned().collect(); + let mut changed = false; + for (child, requirement) in children.iter_mut().zip(requirements) { + let Some(requirement) = requirement else { + continue; + }; + let requirement = requirement.into_single(); + let eq_properties = child.equivalence_properties(); + if eq_properties.ordering_satisfy_requirement(requirement.clone())? { + continue; + } + let sort_reqs: Vec<_> = requirement.clone().into(); + let sort_exprs = sort_reqs + .into_iter() + .filter(|req| eq_properties.is_expr_constant(&req.expr).is_none()) + .map(Into::into); + let Some(ordering) = LexOrdering::new(sort_exprs) else { + continue; + }; + let sort: Arc<dyn ExecutionPlan> = Arc::new( Review Comment: This no longer applies: following your comment on the root cause, the EnsureRequirements step is removed and the PR fixes the join instead. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
