Skip to content

Commit

Permalink
Minor: remove confusing update_plan_from_children call from `Enforc…
Browse files Browse the repository at this point in the history
…eSorting` (#14650)

* Minor: remove unnecessary update_plan_from_children call

* update

* resolve comments
  • Loading branch information
xudong963 authored Feb 15, 2025
1 parent c2e004d commit 9681c3b
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 10 deletions.
23 changes: 15 additions & 8 deletions datafusion/physical-optimizer/src/enforce_sorting/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ impl EnforceSorting {
/// via its children.
pub type PlanWithCorrespondingSort = PlanContext<bool>;

fn update_sort_ctx_children(
fn update_sort_ctx_children_data(
mut node: PlanWithCorrespondingSort,
data: bool,
) -> Result<PlanWithCorrespondingSort> {
Expand Down Expand Up @@ -119,7 +119,7 @@ fn update_sort_ctx_children(
}

node.data = data;
node.update_plan_from_children()
Ok(node)
}

/// This object is used within the [`EnforceSorting`] rule to track the closest
Expand Down Expand Up @@ -322,7 +322,10 @@ pub fn parallelize_sorts(
pub fn ensure_sorting(
mut requirements: PlanWithCorrespondingSort,
) -> Result<Transformed<PlanWithCorrespondingSort>> {
requirements = update_sort_ctx_children(requirements, false)?;
// Before starting, making requirements' children's ExecutionPlan be same as the requirements' plan's children's ExecutionPlan.
// It should be guaranteed by previous code, but we need to make sure to avoid any potential missing.
requirements = requirements.update_plan_from_children()?;
requirements = update_sort_ctx_children_data(requirements, false)?;

// Perform naive analysis at the beginning -- remove already-satisfied sorts:
if requirements.children.is_empty() {
Expand Down Expand Up @@ -353,7 +356,8 @@ pub fn ensure_sorting(
child = update_child_to_remove_unnecessary_sort(idx, child, plan)?;
}
child = add_sort_above(child, required, None);
child = update_sort_ctx_children(child, true)?;
child = child.update_plan_from_children()?;
child = update_sort_ctx_children_data(child, true)?;
}
} else if physical_ordering.is_none()
|| !plan.maintains_input_order()[idx]
Expand Down Expand Up @@ -383,9 +387,10 @@ pub fn ensure_sorting(
Arc::new(LocalLimitExec::new(Arc::clone(&child_node.plan), fetch));
}
return Ok(Transformed::yes(child_node));
} else {
requirements = requirements.update_plan_from_children()?;
}

update_sort_ctx_children(requirements, false).map(Transformed::yes)
update_sort_ctx_children_data(requirements, false).map(Transformed::yes)
}

/// Analyzes a given [`SortExec`] (`plan`) to determine whether its input
Expand Down Expand Up @@ -609,8 +614,9 @@ fn remove_corresponding_sort_from_sub_plan(
}
})
.collect::<Result<_>>()?;
node = node.update_plan_from_children()?;
if any_connection || node.children.is_empty() {
node = update_sort_ctx_children(node, false)?;
node = update_sort_ctx_children_data(node, false)?;
}

// Replace with variants that do not preserve order.
Expand Down Expand Up @@ -643,7 +649,8 @@ fn remove_corresponding_sort_from_sub_plan(
Arc::new(CoalescePartitionsExec::new(plan)) as _
};
node = PlanWithCorrespondingSort::new(plan, false, vec![node]);
node = update_sort_ctx_children(node, false)?;
node = node.update_plan_from_children()?;
node = update_sort_ctx_children_data(node, false)?;
}
Ok(node)
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ use itertools::izip;
pub type OrderPreservationContext = PlanContext<bool>;

/// Updates order-preservation data for all children of the given node.
pub fn update_children(opc: &mut OrderPreservationContext) {
pub fn update_order_preservation_ctx_children_data(opc: &mut OrderPreservationContext) {
for PlanContext {
plan,
children,
Expand Down Expand Up @@ -244,7 +244,7 @@ pub fn replace_with_order_preserving_variants(
is_spm_better: bool,
config: &ConfigOptions,
) -> Result<Transformed<OrderPreservationContext>> {
update_children(&mut requirements);
update_order_preservation_ctx_children_data(&mut requirements);
if !(is_sort(&requirements.plan) && requirements.children[0].data) {
return Ok(Transformed::no(requirements));
}
Expand Down
2 changes: 2 additions & 0 deletions datafusion/physical-plan/src/tree_node.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ impl DynTreeNode for dyn ExecutionPlan {
/// A node object beneficial for writing optimizer rules, encapsulating an [`ExecutionPlan`] node with a payload.
/// Since there are two ways to access child plans—directly from the plan and through child nodes—it's recommended
/// to perform mutable operations via [`Self::update_plan_from_children`].
/// After update `children`, please do the sync updating for `plan`'s children.
/// Or after creating the `PlanContext`, if you can't guarantee they are consistent, call `update_plan_from_children` to sync.
#[derive(Debug)]
pub struct PlanContext<T: Sized> {
/// The execution plan associated with this context.
Expand Down

0 comments on commit 9681c3b

Please sign in to comment.