datafusion_physical_optimizer/
utils.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::sync::Arc;
19
20use datafusion_common::Result;
21use datafusion_physical_expr::{Distribution, LexOrdering, LexRequirement};
22use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
23use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
24use datafusion_physical_plan::repartition::RepartitionExec;
25use datafusion_physical_plan::sorts::sort::SortExec;
26use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec;
27use datafusion_physical_plan::tree_node::PlanContext;
28use datafusion_physical_plan::union::UnionExec;
29use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec};
30use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
31
32/// This utility function adds a `SortExec` above an operator according to the
33/// given ordering requirements while preserving the original partitioning.
34///
35/// Note that this updates the plan in both the `PlanContext.children` and
36/// the `PlanContext.plan`'s children. Therefore its not required to sync
37/// the child plans with [`PlanContext::update_plan_from_children`].
38pub fn add_sort_above<T: Clone + Default>(
39    node: PlanContext<T>,
40    sort_requirements: LexRequirement,
41    fetch: Option<usize>,
42) -> PlanContext<T> {
43    let mut sort_reqs: Vec<_> = sort_requirements.into();
44    sort_reqs.retain(|sort_expr| {
45        node.plan
46            .equivalence_properties()
47            .is_expr_constant(&sort_expr.expr)
48            .is_none()
49    });
50    let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::<Vec<_>>();
51    let Some(ordering) = LexOrdering::new(sort_exprs) else {
52        return node;
53    };
54    let mut new_sort = SortExec::new(ordering, Arc::clone(&node.plan)).with_fetch(fetch);
55    if node.plan.output_partitioning().partition_count() > 1 {
56        new_sort = new_sort.with_preserve_partitioning(true);
57    }
58    PlanContext::new(Arc::new(new_sort), T::default(), vec![node])
59}
60
61/// Like [`add_sort_above`], but also inserts a [`SortPreservingMergeExec`] when
62/// the parent distribution requires a single partition and the input has
63/// multiple partitions. This prevents `SortExec(preserve_partitioning=true)`
64/// from violating `SinglePartition` requirements.
65pub fn add_sort_above_with_distribution<T: Clone + Default>(
66    node: PlanContext<T>,
67    sort_requirements: LexRequirement,
68    fetch: Option<usize>,
69    required_distribution: &Distribution,
70) -> PlanContext<T> {
71    let mut sort_reqs: Vec<_> = sort_requirements.into();
72    sort_reqs.retain(|sort_expr| {
73        node.plan
74            .equivalence_properties()
75            .is_expr_constant(&sort_expr.expr)
76            .is_none()
77    });
78    let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::<Vec<_>>();
79    let Some(ordering) = LexOrdering::new(sort_exprs) else {
80        return node;
81    };
82    let input_has_multiple_partitions =
83        node.plan.output_partitioning().partition_count() > 1;
84
85    let mut new_sort =
86        SortExec::new(ordering.clone(), Arc::clone(&node.plan)).with_fetch(fetch);
87    if input_has_multiple_partitions {
88        new_sort = new_sort.with_preserve_partitioning(true);
89    }
90
91    let sort_node = PlanContext::new(Arc::new(new_sort), T::default(), vec![node]);
92
93    // If the parent requires SinglePartition and the input has multiple partitions,
94    // wrap the partition-preserving sort in SortPreservingMergeExec.
95    if matches!(required_distribution, Distribution::SinglePartition)
96        && input_has_multiple_partitions
97    {
98        PlanContext::new(
99            Arc::new(
100                SortPreservingMergeExec::new(ordering, Arc::clone(&sort_node.plan))
101                    .with_fetch(fetch),
102            ),
103            T::default(),
104            vec![sort_node],
105        )
106    } else {
107        sort_node
108    }
109}
110
111/// This utility function adds a `SortExec` above an operator according to the
112/// given ordering requirements while preserving the original partitioning. If
113/// requirement is already satisfied no `SortExec` is added.
114pub fn add_sort_above_with_check<T: Clone + Default>(
115    node: PlanContext<T>,
116    sort_requirements: LexRequirement,
117    fetch: Option<usize>,
118) -> Result<PlanContext<T>> {
119    if !node
120        .plan
121        .equivalence_properties()
122        .ordering_satisfy_requirement(sort_requirements.clone())?
123    {
124        Ok(add_sort_above(node, sort_requirements, fetch))
125    } else {
126        Ok(node)
127    }
128}
129
130/// Checks whether the given operator is a [`SortExec`].
131pub fn is_sort(plan: &Arc<dyn ExecutionPlan>) -> bool {
132    plan.is::<SortExec>()
133}
134
135/// Checks whether the given operator is a window;
136/// i.e. either a [`WindowAggExec`] or a [`BoundedWindowAggExec`].
137pub fn is_window(plan: &Arc<dyn ExecutionPlan>) -> bool {
138    plan.is::<WindowAggExec>() || plan.is::<BoundedWindowAggExec>()
139}
140
141/// Checks whether the given operator is a [`UnionExec`].
142pub fn is_union(plan: &Arc<dyn ExecutionPlan>) -> bool {
143    plan.is::<UnionExec>()
144}
145
146/// Checks whether the given operator is a [`SortPreservingMergeExec`].
147pub fn is_sort_preserving_merge(plan: &Arc<dyn ExecutionPlan>) -> bool {
148    plan.is::<SortPreservingMergeExec>()
149}
150
151/// Checks whether the given operator is a [`CoalescePartitionsExec`].
152pub fn is_coalesce_partitions(plan: &Arc<dyn ExecutionPlan>) -> bool {
153    plan.is::<CoalescePartitionsExec>()
154}
155
156/// Checks whether the given operator is a [`RepartitionExec`].
157pub fn is_repartition(plan: &Arc<dyn ExecutionPlan>) -> bool {
158    plan.is::<RepartitionExec>()
159}
160
161/// Checks whether the given operator is a limit;
162/// i.e. either a [`LocalLimitExec`] or a [`GlobalLimitExec`].
163pub fn is_limit(plan: &Arc<dyn ExecutionPlan>) -> bool {
164    plan.is::<GlobalLimitExec>() || plan.is::<LocalLimitExec>()
165}
datafusion_physical_optimizer/utils.rs

datafusion_physical_optimizer/
utils.rs
