Skip to main content

datafusion_physical_expr/equivalence/properties/
mod.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
18mod dependency; // Submodule containing DependencyMap and Dependencies
19mod joins; // Submodule containing join_equivalence_properties
20mod union; // Submodule containing calculate_union
21
22pub use joins::*;
23pub use union::*;
24
25use std::fmt::{self, Display};
26use std::mem;
27use std::sync::Arc;
28
29use self::dependency::{
30    Dependencies, DependencyMap, construct_prefix_orderings,
31    generate_dependency_orderings, referred_dependencies,
32};
33use crate::equivalence::{
34    AcrossPartitions, EquivalenceGroup, OrderingEquivalenceClass, ProjectionMapping,
35};
36use crate::expressions::{Column, Literal, with_new_schema};
37use crate::{
38    ConstExpr, LexOrdering, LexRequirement, PhysicalExpr, PhysicalSortExpr,
39    PhysicalSortRequirement,
40};
41
42use arrow::datatypes::SchemaRef;
43use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
44use datafusion_common::{Constraint, Constraints, HashMap, Result, plan_err};
45use datafusion_expr::interval_arithmetic::Interval;
46use datafusion_expr::sort_properties::{ExprProperties, SortProperties};
47use datafusion_physical_expr_common::sort_expr::options_compatible;
48use datafusion_physical_expr_common::utils::ExprPropertiesNode;
49
50use indexmap::IndexSet;
51use itertools::Itertools;
52
53/// `EquivalenceProperties` stores information about the output of a plan node
54/// that can be used to optimize the plan. Currently, it keeps track of:
55/// - Sort expressions (orderings),
56/// - Equivalent expressions; i.e. expressions known to have the same value.
57/// - Constants expressions; i.e. expressions known to contain a single constant
58///   value.
59///
60/// Please see the [Using Ordering for Better Plans] blog for more details.
61///
62/// [Using Ordering for Better Plans]: https://datafusion.apache.org/blog/2025/03/11/ordering-analysis/
63///
64/// # Example equivalent sort expressions
65///
66/// Consider table below:
67///
68/// ```text
69/// ┌-------┐
70/// | a | b |
71/// |---|---|
72/// | 1 | 9 |
73/// | 2 | 8 |
74/// | 3 | 7 |
75/// | 5 | 5 |
76/// └---┴---┘
77/// ```
78///
79/// In this case, both `a ASC` and `b DESC` can describe the table ordering.
80/// `EquivalenceProperties` tracks these different valid sort expressions and
81/// treat `a ASC` and `b DESC` on an equal footing. For example, if the query
82/// specifies the output sorted by EITHER `a ASC` or `b DESC`, the sort can be
83/// avoided.
84///
85/// # Example equivalent expressions
86///
87/// Similarly, consider the table below:
88///
89/// ```text
90/// ┌-------┐
91/// | a | b |
92/// |---|---|
93/// | 1 | 1 |
94/// | 2 | 2 |
95/// | 3 | 3 |
96/// | 5 | 5 |
97/// └---┴---┘
98/// ```
99///
100/// In this case,  columns `a` and `b` always have the same value. With this
101/// information, Datafusion can optimize various operations. For example, if
102/// the partition requirement is `Hash(a)` and output partitioning is
103/// `Hash(b)`, then DataFusion avoids repartitioning the data as the existing
104/// partitioning satisfies the requirement.
105///
106/// # Code Example
107/// ```
108/// # use std::sync::Arc;
109/// # use arrow::datatypes::{Schema, Field, DataType, SchemaRef};
110/// # use datafusion_physical_expr::{ConstExpr, EquivalenceProperties};
111/// # use datafusion_physical_expr::expressions::col;
112/// use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
113/// # let schema: SchemaRef = Arc::new(Schema::new(vec![
114/// #   Field::new("a", DataType::Int32, false),
115/// #   Field::new("b", DataType::Int32, false),
116/// #   Field::new("c", DataType::Int32, false),
117/// # ]));
118/// # let col_a = col("a", &schema).unwrap();
119/// # let col_b = col("b", &schema).unwrap();
120/// # let col_c = col("c", &schema).unwrap();
121/// // This object represents data that is sorted by a ASC, c DESC
122/// // with a single constant value of b
123/// let mut eq_properties = EquivalenceProperties::new(schema);
124/// eq_properties.add_constants(vec![ConstExpr::from(col_b)]);
125/// eq_properties.add_ordering([
126///     PhysicalSortExpr::new_default(col_a).asc(),
127///     PhysicalSortExpr::new_default(col_c).desc(),
128/// ]);
129///
130/// assert_eq!(
131///     eq_properties.to_string(),
132///     "order: [[a@0 ASC, c@2 DESC]], eq: [{members: [b@1], constant: (heterogeneous)}]"
133/// );
134/// ```
135#[derive(Clone, Debug)]
136pub struct EquivalenceProperties {
137    /// Distinct equivalence classes (i.e. expressions with the same value).
138    eq_group: EquivalenceGroup,
139    /// Equivalent sort expressions (i.e. those define the same ordering).
140    oeq_class: OrderingEquivalenceClass,
141    /// Cache storing equivalent sort expressions in normal form (i.e. without
142    /// constants/duplicates and in standard form) and a map associating leading
143    /// terms with full sort expressions.
144    oeq_cache: OrderingEquivalenceCache,
145    /// Table constraints that factor in equivalence calculations.
146    constraints: Constraints,
147    /// Schema associated with this object.
148    schema: SchemaRef,
149}
150
151/// This object serves as a cache for storing equivalent sort expressions
152/// in normal form, and a map associating leading sort expressions with
153/// full lexicographical orderings. With this information, DataFusion can
154/// efficiently determine whether a given ordering is satisfied by the
155/// existing orderings, and discover new orderings based on the existing
156/// equivalence properties.
157#[derive(Clone, Debug, Default)]
158struct OrderingEquivalenceCache {
159    /// Equivalent sort expressions in normal form.
160    normal_cls: OrderingEquivalenceClass,
161    /// Map associating leading sort expressions with full lexicographical
162    /// orderings. Values are indices into `normal_cls`.
163    leading_map: HashMap<Arc<dyn PhysicalExpr>, Vec<usize>>,
164}
165
166impl OrderingEquivalenceCache {
167    /// Creates a new `OrderingEquivalenceCache` object with the given
168    /// equivalent orderings, which should be in normal form.
169    pub fn new(
170        orderings: impl IntoIterator<Item = impl IntoIterator<Item = PhysicalSortExpr>>,
171    ) -> Self {
172        let mut cache = Self {
173            normal_cls: OrderingEquivalenceClass::new(orderings),
174            leading_map: HashMap::new(),
175        };
176        cache.update_map();
177        cache
178    }
179
180    /// Updates/reconstructs the leading expression map according to the normal
181    /// ordering equivalence class within.
182    pub fn update_map(&mut self) {
183        self.leading_map.clear();
184        for (idx, ordering) in self.normal_cls.iter().enumerate() {
185            let expr = Arc::clone(&ordering.first().expr);
186            self.leading_map.entry(expr).or_default().push(idx);
187        }
188    }
189
190    /// Clears the cache, removing all orderings and leading expressions.
191    pub fn clear(&mut self) {
192        self.normal_cls.clear();
193        self.leading_map.clear();
194    }
195}
196
197impl EquivalenceProperties {
198    /// Helper used by the ordering equivalence rule when considering whether
199    /// an expression can replace an existing sort key without invalidating
200    /// the ordering.
201    ///
202    /// The substitution is only allowed when, treating the sort key as the
203    /// only ordered input, the expression reports the same ordering *and*
204    /// that it is a one-to-one, order-preserving function of it (see
205    /// [`ExprProperties::strictly_order_preserving`]). For example, a
206    /// widening `CAST` of the sort key qualifies, while a narrowing one does
207    /// not, as it could collapse distinct values and violate the existing
208    /// sort order.
209    fn substitute_order_preserving_ordering(
210        r_expr: Arc<dyn PhysicalExpr>,
211        sort_expr: &PhysicalSortExpr,
212        schema: &SchemaRef,
213    ) -> Option<PhysicalSortExpr> {
214        if r_expr.eq(&sort_expr.expr) {
215            // No point in substituting an expression with itself.
216            return None;
217        }
218        let dependencies = Dependencies::new(std::iter::once(sort_expr.clone()));
219        let properties = get_expr_properties(&r_expr, &dependencies, schema).ok()?;
220        (properties.strictly_order_preserving
221            && properties.sort_properties == SortProperties::Ordered(sort_expr.options))
222        .then(|| PhysicalSortExpr::new(r_expr, sort_expr.options))
223    }
224
225    /// Creates an empty `EquivalenceProperties` object.
226    pub fn new(schema: SchemaRef) -> Self {
227        Self {
228            eq_group: EquivalenceGroup::default(),
229            oeq_class: OrderingEquivalenceClass::default(),
230            oeq_cache: OrderingEquivalenceCache::default(),
231            constraints: Constraints::default(),
232            schema,
233        }
234    }
235
236    /// Adds constraints to the properties.
237    pub fn set_constraints(&mut self, constraints: Constraints) {
238        self.constraints = constraints;
239    }
240
241    /// Adds constraints to the properties.
242    pub fn with_constraints(mut self, constraints: Constraints) -> Self {
243        self.set_constraints(constraints);
244        self
245    }
246
247    /// Creates a new `EquivalenceProperties` object with the given orderings.
248    pub fn new_with_orderings(
249        schema: SchemaRef,
250        orderings: impl IntoIterator<Item = impl IntoIterator<Item = PhysicalSortExpr>>,
251    ) -> Self {
252        let eq_group = EquivalenceGroup::default();
253        let oeq_class = OrderingEquivalenceClass::new(orderings);
254        // Here, we can avoid performing a full normalization, and get by with
255        // only removing constants because the equivalence group is empty.
256        let normal_orderings = oeq_class.iter().cloned().map(|o| {
257            o.into_iter()
258                .filter(|sort_expr| eq_group.is_expr_constant(&sort_expr.expr).is_none())
259        });
260        Self {
261            oeq_cache: OrderingEquivalenceCache::new(normal_orderings),
262            oeq_class,
263            eq_group,
264            constraints: Constraints::default(),
265            schema,
266        }
267    }
268
269    /// Returns the associated schema.
270    pub fn schema(&self) -> &SchemaRef {
271        &self.schema
272    }
273
274    /// Returns a reference to the ordering equivalence class within.
275    pub fn oeq_class(&self) -> &OrderingEquivalenceClass {
276        &self.oeq_class
277    }
278
279    /// Returns a reference to the equivalence group within.
280    pub fn eq_group(&self) -> &EquivalenceGroup {
281        &self.eq_group
282    }
283
284    /// Returns a reference to the constraints within.
285    pub fn constraints(&self) -> &Constraints {
286        &self.constraints
287    }
288
289    /// Returns all the known constants expressions.
290    pub fn constants(&self) -> Vec<ConstExpr> {
291        self.eq_group
292            .iter()
293            .flat_map(|c| {
294                c.iter().filter_map(|expr| {
295                    c.constant
296                        .as_ref()
297                        .map(|across| ConstExpr::new(Arc::clone(expr), across.clone()))
298                })
299            })
300            .collect()
301    }
302
303    /// Returns the output ordering of the properties.
304    pub fn output_ordering(&self) -> Option<LexOrdering> {
305        let concat = self.oeq_class.iter().flat_map(|o| o.iter().cloned());
306        self.normalize_sort_exprs(concat)
307    }
308
309    /// Extends this `EquivalenceProperties` with the `other` object.
310    pub fn extend(mut self, other: Self) -> Result<Self> {
311        self.constraints.extend(other.constraints);
312        self.add_equivalence_group(other.eq_group)?;
313        self.add_orderings(other.oeq_class);
314        Ok(self)
315    }
316
317    /// Clears (empties) the ordering equivalence class within this object.
318    /// Call this method when existing orderings are invalidated.
319    pub fn clear_orderings(&mut self) {
320        self.oeq_class.clear();
321        self.oeq_cache.clear();
322    }
323
324    /// Removes constant expressions that may change across partitions.
325    /// This method should be used when merging data from different partitions.
326    pub fn clear_per_partition_constants(&mut self) {
327        if self.eq_group.clear_per_partition_constants() {
328            // Renormalize orderings if the equivalence group changes:
329            let normal_orderings = self
330                .oeq_class
331                .iter()
332                .cloned()
333                .map(|o| self.eq_group.normalize_sort_exprs(o));
334            self.oeq_cache = OrderingEquivalenceCache::new(normal_orderings);
335        }
336    }
337
338    /// Adds new orderings into the existing ordering equivalence class.
339    pub fn add_orderings(
340        &mut self,
341        orderings: impl IntoIterator<Item = impl IntoIterator<Item = PhysicalSortExpr>>,
342    ) {
343        let orderings: Vec<_> =
344            orderings.into_iter().filter_map(LexOrdering::new).collect();
345        let normal_orderings: Vec<_> = orderings
346            .iter()
347            .cloned()
348            .filter_map(|o| self.normalize_sort_exprs(o))
349            .collect();
350        if !normal_orderings.is_empty() {
351            self.oeq_class.extend(orderings);
352            // Normalize given orderings to update the cache:
353            self.oeq_cache.normal_cls.extend(normal_orderings);
354            // TODO: If no ordering is found to be redundant during extension, we
355            //       can use a shortcut algorithm to update the leading map.
356            self.oeq_cache.update_map();
357        }
358    }
359
360    /// Adds a single ordering to the existing ordering equivalence class.
361    pub fn add_ordering(&mut self, ordering: impl IntoIterator<Item = PhysicalSortExpr>) {
362        self.add_orderings(std::iter::once(ordering));
363    }
364
365    fn update_oeq_cache(&mut self) -> Result<()> {
366        // Renormalize orderings if the equivalence group changes:
367        let normal_cls = mem::take(&mut self.oeq_cache.normal_cls);
368        let normal_orderings = normal_cls
369            .into_iter()
370            .map(|o| self.eq_group.normalize_sort_exprs(o));
371        self.oeq_cache.normal_cls = OrderingEquivalenceClass::new(normal_orderings);
372        self.oeq_cache.update_map();
373        // Discover any new orderings based on the new equivalence classes:
374        let leading_exprs: Vec<_> = self.oeq_cache.leading_map.keys().cloned().collect();
375        for expr in leading_exprs {
376            self.discover_new_orderings(expr)?;
377        }
378        Ok(())
379    }
380
381    /// Incorporates the given equivalence group to into the existing
382    /// equivalence group within.
383    pub fn add_equivalence_group(
384        &mut self,
385        other_eq_group: EquivalenceGroup,
386    ) -> Result<()> {
387        if !other_eq_group.is_empty() {
388            self.eq_group.extend(other_eq_group);
389            self.update_oeq_cache()?;
390        }
391        Ok(())
392    }
393
394    /// Returns the ordering equivalence class within in normal form.
395    /// Normalization standardizes expressions according to the equivalence
396    /// group within, and removes constants/duplicates.
397    pub fn normalized_oeq_class(&self) -> OrderingEquivalenceClass {
398        self.oeq_class
399            .iter()
400            .cloned()
401            .filter_map(|ordering| self.normalize_sort_exprs(ordering))
402            .collect::<Vec<_>>()
403            .into()
404    }
405
406    /// Adds a new equality condition into the existing equivalence group.
407    /// If the given equality defines a new equivalence class, adds this new
408    /// equivalence class to the equivalence group.
409    pub fn add_equal_conditions(
410        &mut self,
411        left: Arc<dyn PhysicalExpr>,
412        right: Arc<dyn PhysicalExpr>,
413    ) -> Result<()> {
414        // Add equal expressions to the state:
415        if self.eq_group.add_equal_conditions(left, right) {
416            self.update_oeq_cache()?;
417        }
418        self.update_oeq_cache()?;
419        Ok(())
420    }
421
422    /// Track/register physical expressions with constant values.
423    pub fn add_constants(
424        &mut self,
425        constants: impl IntoIterator<Item = ConstExpr>,
426    ) -> Result<()> {
427        // Add the new constant to the equivalence group:
428        for constant in constants {
429            self.eq_group.add_constant(constant);
430        }
431        // Renormalize the orderings after adding new constants by removing
432        // the constants from existing orderings:
433        let normal_cls = mem::take(&mut self.oeq_cache.normal_cls);
434        let normal_orderings = normal_cls.into_iter().map(|ordering| {
435            ordering.into_iter().filter(|sort_expr| {
436                self.eq_group.is_expr_constant(&sort_expr.expr).is_none()
437            })
438        });
439        self.oeq_cache.normal_cls = OrderingEquivalenceClass::new(normal_orderings);
440        self.oeq_cache.update_map();
441        // Discover any new orderings based on the constants:
442        let leading_exprs: Vec<_> = self.oeq_cache.leading_map.keys().cloned().collect();
443        for expr in leading_exprs {
444            self.discover_new_orderings(expr)?;
445        }
446        Ok(())
447    }
448
449    /// Discover new valid orderings in light of a new equality. Accepts a single
450    /// argument (`expr`) which is used to determine the orderings to update.
451    /// When constants or equivalence classes change, there may be new orderings
452    /// that can be discovered with the new equivalence properties.
453    /// For a discussion, see: <https://github.com/apache/datafusion/issues/9812>
454    fn discover_new_orderings(
455        &mut self,
456        normal_expr: Arc<dyn PhysicalExpr>,
457    ) -> Result<()> {
458        let Some(ordering_idxs) = self.oeq_cache.leading_map.get(&normal_expr) else {
459            return Ok(());
460        };
461        let eq_class = self
462            .eq_group
463            .get_equivalence_class(&normal_expr)
464            .map_or_else(|| vec![normal_expr], |class| class.clone().into());
465
466        let mut new_orderings = vec![];
467        for idx in ordering_idxs {
468            let ordering = &self.oeq_cache.normal_cls[*idx];
469            let leading_ordering_options = ordering[0].options;
470
471            'exprs: for equivalent_expr in &eq_class {
472                let children = equivalent_expr.children();
473                if children.is_empty() {
474                    continue;
475                }
476                // Check if all children match the next expressions in the ordering:
477                let mut child_properties = vec![];
478                // Build properties for each child based on the next expression:
479                for (i, child) in children.into_iter().enumerate() {
480                    let Some(next) = ordering.get(i + 1) else {
481                        break 'exprs;
482                    };
483                    if !next.expr.eq(child) {
484                        break 'exprs;
485                    }
486                    let data_type = child.data_type(&self.schema)?;
487                    child_properties.push(ExprProperties {
488                        sort_properties: SortProperties::Ordered(next.options),
489                        range: Interval::make_unbounded(&data_type)?,
490                        preserves_lex_ordering: true,
491                        strictly_order_preserving: true,
492                    });
493                }
494                // Check if the expression is monotonic in all arguments:
495                let expr_properties =
496                    equivalent_expr.get_properties(&child_properties)?;
497                if expr_properties.preserves_lex_ordering
498                    && expr_properties.sort_properties
499                        == SortProperties::Ordered(leading_ordering_options)
500                {
501                    // Assume that `[c ASC, a ASC, b ASC]` is among existing
502                    // orderings. If equality `c = f(a, b)` is given, ordering
503                    // `[a ASC, b ASC]` implies the ordering `[c ASC]`. Thus,
504                    // ordering `[a ASC, b ASC]` is also a valid ordering.
505                    new_orderings.push(ordering[1..].to_vec());
506                    break;
507                }
508            }
509        }
510
511        if !new_orderings.is_empty() {
512            self.add_orderings(new_orderings);
513        }
514        Ok(())
515    }
516
517    /// Updates the ordering equivalence class within assuming that the table
518    /// is re-sorted according to the argument `ordering`, and returns whether
519    /// this operation resulted in any change. Note that equivalence classes
520    /// (and constants) do not change as they are unaffected by a re-sort. If
521    /// the given ordering is already satisfied, the function does nothing.
522    pub fn reorder(
523        &mut self,
524        ordering: impl IntoIterator<Item = PhysicalSortExpr>,
525    ) -> Result<bool> {
526        let (ordering, ordering_tee) = ordering.into_iter().tee();
527        // First, standardize the given ordering:
528        let Some(normal_ordering) = self.normalize_sort_exprs(ordering) else {
529            // If the ordering vanishes after normalization, it is satisfied:
530            return Ok(false);
531        };
532        if normal_ordering.len() != self.common_sort_prefix_length(&normal_ordering)? {
533            // If the ordering is unsatisfied, replace existing orderings:
534            self.clear_orderings();
535            self.add_ordering(ordering_tee);
536            return Ok(true);
537        }
538        Ok(false)
539    }
540
541    /// Normalizes the given sort expressions (i.e. `sort_exprs`) using the
542    /// equivalence group within. Returns a `LexOrdering` instance if the
543    /// expressions define a proper lexicographical ordering. For more details,
544    /// see [`EquivalenceGroup::normalize_sort_exprs`].
545    pub fn normalize_sort_exprs(
546        &self,
547        sort_exprs: impl IntoIterator<Item = PhysicalSortExpr>,
548    ) -> Option<LexOrdering> {
549        LexOrdering::new(self.eq_group.normalize_sort_exprs(sort_exprs))
550    }
551
552    /// Normalizes the given sort requirements (i.e. `sort_reqs`) using the
553    /// equivalence group within. Returns a `LexRequirement` instance if the
554    /// expressions define a proper lexicographical requirement. For more
555    /// details, see [`EquivalenceGroup::normalize_sort_exprs`].
556    pub fn normalize_sort_requirements(
557        &self,
558        sort_reqs: impl IntoIterator<Item = PhysicalSortRequirement>,
559    ) -> Option<LexRequirement> {
560        LexRequirement::new(self.eq_group.normalize_sort_requirements(sort_reqs))
561    }
562
563    /// Iteratively checks whether the given ordering is satisfied by any of
564    /// the existing orderings. See [`Self::ordering_satisfy_requirement`] for
565    /// more details and examples.
566    pub fn ordering_satisfy(
567        &self,
568        given: impl IntoIterator<Item = PhysicalSortExpr>,
569    ) -> Result<bool> {
570        // First, standardize the given ordering:
571        let Some(normal_ordering) = self.normalize_sort_exprs(given) else {
572            // If the ordering vanishes after normalization, it is satisfied:
573            return Ok(true);
574        };
575        Ok(normal_ordering.len() == self.common_sort_prefix_length(&normal_ordering)?)
576    }
577
578    /// Iteratively checks whether the given sort requirement is satisfied by
579    /// any of the existing orderings.
580    ///
581    /// ### Example Scenarios
582    ///
583    /// In these scenarios, assume that all expressions share the same sort
584    /// properties.
585    ///
586    /// #### Case 1: Sort Requirement `[a, c]`
587    ///
588    /// **Existing orderings:** `[[a, b, c], [a, d]]`, **constants:** `[]`
589    /// 1. The function first checks the leading requirement `a`, which is
590    ///    satisfied by `[a, b, c].first()`.
591    /// 2. `a` is added as a constant for the next iteration.
592    /// 3. Normal orderings become `[[b, c], [d]]`.
593    /// 4. The function fails for `c` in the second iteration, as neither
594    ///    `[b, c]` nor `[d]` satisfies `c`.
595    ///
596    /// #### Case 2: Sort Requirement `[a, d]`
597    ///
598    /// **Existing orderings:** `[[a, b, c], [a, d]]`, **constants:** `[]`
599    /// 1. The function first checks the leading requirement `a`, which is
600    ///    satisfied by `[a, b, c].first()`.
601    /// 2. `a` is added as a constant for the next iteration.
602    /// 3. Normal orderings become `[[b, c], [d]]`.
603    /// 4. The function returns `true` as `[d]` satisfies `d`.
604    pub fn ordering_satisfy_requirement(
605        &self,
606        given: impl IntoIterator<Item = PhysicalSortRequirement>,
607    ) -> Result<bool> {
608        // First, standardize the given requirement:
609        let Some(normal_reqs) = self.normalize_sort_requirements(given) else {
610            // If the requirement vanishes after normalization, it is satisfied:
611            return Ok(true);
612        };
613        // Then, check whether given requirement is satisfied by constraints:
614        if self.satisfied_by_constraints(&normal_reqs) {
615            return Ok(true);
616        }
617        let schema = self.schema();
618        let mut eq_properties = self.clone();
619        for element in normal_reqs {
620            // Check whether given requirement is satisfied:
621            let ExprProperties {
622                sort_properties, ..
623            } = eq_properties.get_expr_properties(Arc::clone(&element.expr));
624            let satisfy = match sort_properties {
625                SortProperties::Ordered(options) => element.options.is_none_or(|opts| {
626                    let nullable = element.expr.nullable(schema).unwrap_or(true);
627                    options_compatible(&options, &opts, nullable)
628                }),
629                // Singleton expressions satisfy any requirement.
630                SortProperties::Singleton => true,
631                SortProperties::Unordered => false,
632            };
633            if !satisfy {
634                return Ok(false);
635            }
636            // Treat satisfied keys (and the sub-expressions they pin down) as
637            // constants in subsequent iterations. See
638            // [`Self::add_satisfied_key_constants`] for the rationale.
639            eq_properties.add_satisfied_key_constants(element.expr)?;
640        }
641        Ok(true)
642    }
643
644    /// Registers a satisfied sort key as a constant for subsequent iterations
645    /// of the ordering satisfaction checks. We can do this because the "next"
646    /// key only matters in a lexicographical ordering when the keys to its
647    /// left have the same values (i.e. within a single tie group). Note that
648    /// these expressions are not properly "constants"; this is just an
649    /// implementation strategy confined to the satisfaction checks.
650    ///
651    /// For example, assume that the requirement is `[a ASC, (b + c) ASC]`,
652    /// and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`.
653    /// Once we deduce that `[a ASC]` is satisfied, we add column `a` as a
654    /// constant to the algorithm state. This enables us to deduce that
655    /// `(b + c) ASC` is satisfied, given `a` is constant.
656    ///
657    /// In addition to the key itself, this also registers any sub-expressions
658    /// whose values the key pins down: if an expression is strictly
659    /// order-preserving, equal outputs imply equal values of its ordered
660    /// children, so within a tie group of the key those children are constant
661    /// as well. For example, if data is sorted by `[a, b]`, the requirement
662    /// `[CAST(a AS BIGINT) ASC, b ASC]` is satisfied: `a` is constant within
663    /// each group of equal `CAST(a AS BIGINT)` values, and hence `b` is
664    /// sorted within each such group.
665    fn add_satisfied_key_constants(&mut self, expr: Arc<dyn PhysicalExpr>) -> Result<()> {
666        let mut stack = vec![expr];
667        while let Some(expr) = stack.pop() {
668            let properties = self.get_expr_properties(Arc::clone(&expr));
669            if properties.strictly_order_preserving {
670                for child in expr.children() {
671                    let child_properties = self.get_expr_properties(Arc::clone(child));
672                    if matches!(
673                        child_properties.sort_properties,
674                        SortProperties::Ordered(_)
675                    ) {
676                        stack.push(Arc::clone(child));
677                    }
678                }
679            }
680            self.add_constants(std::iter::once(ConstExpr::from(expr)))?;
681        }
682        Ok(())
683    }
684
685    /// Returns the number of consecutive sort expressions (starting from the
686    /// left) that are satisfied by the existing ordering.
687    fn common_sort_prefix_length(&self, normal_ordering: &LexOrdering) -> Result<usize> {
688        let full_length = normal_ordering.len();
689        // Check whether the given ordering is satisfied by constraints:
690        if self.satisfied_by_constraints_ordering(normal_ordering) {
691            // If constraints satisfy all sort expressions, return the full
692            // length:
693            return Ok(full_length);
694        }
695        let schema = self.schema();
696        let mut eq_properties = self.clone();
697        for (idx, element) in normal_ordering.into_iter().enumerate() {
698            // Check whether given ordering is satisfied:
699            let ExprProperties {
700                sort_properties, ..
701            } = eq_properties.get_expr_properties(Arc::clone(&element.expr));
702            let satisfy = match sort_properties {
703                SortProperties::Ordered(options) => options_compatible(
704                    &options,
705                    &element.options,
706                    element.expr.nullable(schema).unwrap_or(true),
707                ),
708                // Singleton expressions satisfy any ordering.
709                SortProperties::Singleton => true,
710                SortProperties::Unordered => false,
711            };
712            if !satisfy {
713                // As soon as one sort expression is unsatisfied, return how
714                // many we've satisfied so far:
715                return Ok(idx);
716            }
717            // Treat satisfied keys (and the sub-expressions they pin down) as
718            // constants in subsequent iterations. See
719            // [`Self::add_satisfied_key_constants`] for the rationale.
720            eq_properties.add_satisfied_key_constants(Arc::clone(&element.expr))?;
721        }
722        // All sort expressions are satisfied, return full length:
723        Ok(full_length)
724    }
725
726    /// Determines the longest normal prefix of `ordering` satisfied by the
727    /// existing ordering. Returns that prefix as a new `LexOrdering`, and a
728    /// boolean indicating whether all the sort expressions are satisfied.
729    pub fn extract_common_sort_prefix(
730        &self,
731        ordering: LexOrdering,
732    ) -> Result<(Vec<PhysicalSortExpr>, bool)> {
733        // First, standardize the given ordering:
734        let Some(normal_ordering) = self.normalize_sort_exprs(ordering) else {
735            // If the ordering vanishes after normalization, it is satisfied:
736            return Ok((vec![], true));
737        };
738        let prefix_len = self.common_sort_prefix_length(&normal_ordering)?;
739        let flag = prefix_len == normal_ordering.len();
740        let mut sort_exprs: Vec<_> = normal_ordering.into();
741        if !flag {
742            sort_exprs.truncate(prefix_len);
743        }
744        Ok((sort_exprs, flag))
745    }
746
747    /// Checks if the sort expressions are satisfied by any of the table
748    /// constraints (primary key or unique). Returns true if any constraint
749    /// fully satisfies the expressions (i.e. constraint indices form a valid
750    /// prefix of an existing ordering that matches the expressions). For
751    /// unique constraints, also verifies nullable columns.
752    fn satisfied_by_constraints_ordering(
753        &self,
754        normal_exprs: &[PhysicalSortExpr],
755    ) -> bool {
756        self.constraints.iter().any(|constraint| match constraint {
757            Constraint::PrimaryKey(indices) | Constraint::Unique(indices) => {
758                let check_null = matches!(constraint, Constraint::Unique(_));
759                let normalized_size = normal_exprs.len();
760                indices.len() <= normalized_size
761                    && self.oeq_class.iter().any(|ordering| {
762                        let length = ordering.len();
763                        if indices.len() > length || normalized_size < length {
764                            return false;
765                        }
766                        // Build a map of column positions in the ordering:
767                        let mut col_positions = HashMap::with_capacity(length);
768                        for (pos, req) in ordering.iter().enumerate() {
769                            if let Some(col) = req.expr.downcast_ref::<Column>() {
770                                let nullable = col.nullable(&self.schema).unwrap_or(true);
771                                col_positions.insert(col.index(), (pos, nullable));
772                            }
773                        }
774                        // Check if all constraint indices appear in valid positions:
775                        if !indices.iter().all(|idx| {
776                            col_positions.get(idx).is_some_and(|&(pos, nullable)| {
777                                // For unique constraints, verify column is not nullable if it's first/last:
778                                !check_null
779                                    || !nullable
780                                    || (pos != 0 && pos != length - 1)
781                            })
782                        }) {
783                            return false;
784                        }
785                        // Check if this ordering matches the prefix:
786                        normal_exprs.iter().zip(ordering).all(|(given, existing)| {
787                            existing.satisfy_expr(given, &self.schema)
788                        })
789                    })
790            }
791        })
792    }
793
794    /// Checks if the sort requirements are satisfied by any of the table
795    /// constraints (primary key or unique). Returns true if any constraint
796    /// fully satisfies the requirements (i.e. constraint indices form a valid
797    /// prefix of an existing ordering that matches the requirements). For
798    /// unique constraints, also verifies nullable columns.
799    fn satisfied_by_constraints(&self, normal_reqs: &[PhysicalSortRequirement]) -> bool {
800        self.constraints.iter().any(|constraint| match constraint {
801            Constraint::PrimaryKey(indices) | Constraint::Unique(indices) => {
802                let check_null = matches!(constraint, Constraint::Unique(_));
803                let normalized_size = normal_reqs.len();
804                indices.len() <= normalized_size
805                    && self.oeq_class.iter().any(|ordering| {
806                        let length = ordering.len();
807                        if indices.len() > length || normalized_size < length {
808                            return false;
809                        }
810                        // Build a map of column positions in the ordering:
811                        let mut col_positions = HashMap::with_capacity(length);
812                        for (pos, req) in ordering.iter().enumerate() {
813                            if let Some(col) = req.expr.downcast_ref::<Column>() {
814                                let nullable = col.nullable(&self.schema).unwrap_or(true);
815                                col_positions.insert(col.index(), (pos, nullable));
816                            }
817                        }
818                        // Check if all constraint indices appear in valid positions:
819                        if !indices.iter().all(|idx| {
820                            col_positions.get(idx).is_some_and(|&(pos, nullable)| {
821                                // For unique constraints, verify column is not nullable if it's first/last:
822                                !check_null
823                                    || !nullable
824                                    || (pos != 0 && pos != length - 1)
825                            })
826                        }) {
827                            return false;
828                        }
829                        // Check if this ordering matches the prefix:
830                        normal_reqs.iter().zip(ordering).all(|(given, existing)| {
831                            existing.satisfy(given, &self.schema)
832                        })
833                    })
834            }
835        })
836    }
837
838    /// Checks whether the `given` sort requirements are equal or more specific
839    /// than the `reference` sort requirements.
840    pub fn requirements_compatible(
841        &self,
842        given: LexRequirement,
843        reference: LexRequirement,
844    ) -> bool {
845        let Some(normal_given) = self.normalize_sort_requirements(given) else {
846            return true;
847        };
848        let Some(normal_reference) = self.normalize_sort_requirements(reference) else {
849            return true;
850        };
851
852        (normal_reference.len() <= normal_given.len())
853            && normal_reference
854                .into_iter()
855                .zip(normal_given)
856                .all(|(reference, given)| given.compatible(&reference))
857    }
858
859    /// Modify existing orderings by substituting sort expressions with appropriate
860    /// targets from the projection mapping. We substitute a sort expression when
861    /// its physical expression has a one-to-one functional relationship with a
862    /// target expression in the mapping.
863    ///
864    /// After substitution, we may generate more than one `LexOrdering` for each
865    /// existing equivalent ordering. For example, `[a ASC, b ASC]` will turn
866    /// into `[CAST(a) ASC, b ASC]` and `[a ASC, b ASC]` when applying projection
867    /// expressions `a, b, CAST(a)`.
868    ///
869    /// TODO: Handle all scenarios that allow substitution; e.g. when `x` is
870    ///       sorted, `atan(x + 1000)` should also be substituted. For now, we
871    ///       consider widening `CAST` expressions and single-child expressions
872    ///       that declare themselves one-to-one order-preserving via
873    ///       [`ExprProperties::strictly_order_preserving`].
874    fn substitute_oeq_class(
875        schema: &SchemaRef,
876        mapping: &ProjectionMapping,
877        oeq_class: OrderingEquivalenceClass,
878    ) -> OrderingEquivalenceClass {
879        let new_orderings = oeq_class.into_iter().flat_map(|order| {
880            // Modify/expand existing orderings by substituting sort
881            // expressions with appropriate targets from the mapping:
882            order
883                .into_iter()
884                .map(|sort_expr| {
885                    let original_sort_expr = sort_expr.clone();
886                    mapping
887                        .iter()
888                        .map(|(source, _target)| source)
889                        .filter(|source| expr_refers(source, &original_sort_expr.expr))
890                        .cloned()
891                        .filter_map(|r_expr| {
892                            Self::substitute_order_preserving_ordering(
893                                r_expr,
894                                &original_sort_expr,
895                                schema,
896                            )
897                        })
898                        .chain(std::iter::once(sort_expr))
899                        .collect::<Vec<_>>()
900                })
901                // Generate all valid orderings given substituted expressions:
902                .multi_cartesian_product()
903        });
904        OrderingEquivalenceClass::new(new_orderings)
905    }
906
907    /// Projects argument `expr` according to the projection described by
908    /// `mapping`, taking equivalences into account.
909    ///
910    /// For example, assume that columns `a` and `c` are always equal, and that
911    /// the projection described by `mapping` encodes the following:
912    ///
913    /// ```text
914    /// a -> a1
915    /// b -> b1
916    /// ```
917    ///
918    /// Then, this function projects `a + b` to `Some(a1 + b1)`, `c + b` to
919    /// `Some(a1 + b1)` and `d` to `None`, meaning that it is not projectable.
920    pub fn project_expr(
921        &self,
922        expr: &Arc<dyn PhysicalExpr>,
923        mapping: &ProjectionMapping,
924    ) -> Option<Arc<dyn PhysicalExpr>> {
925        self.eq_group.project_expr(mapping, expr)
926    }
927
928    /// Projects the given `expressions` according to the projection described
929    /// by `mapping`, taking equivalences into account. This function is similar
930    /// to [`Self::project_expr`], but projects multiple expressions at once
931    /// more efficiently than calling `project_expr` for each expression.
932    pub fn project_expressions<'a>(
933        &'a self,
934        expressions: impl IntoIterator<Item = &'a Arc<dyn PhysicalExpr>> + 'a,
935        mapping: &'a ProjectionMapping,
936    ) -> impl Iterator<Item = Option<Arc<dyn PhysicalExpr>>> + 'a {
937        self.eq_group.project_expressions(mapping, expressions)
938    }
939
940    /// Constructs a dependency map based on existing orderings referred to in
941    /// the projection.
942    ///
943    /// This function analyzes the orderings in the normalized order-equivalence
944    /// class and builds a dependency map. The dependency map captures relationships
945    /// between expressions within the orderings, helping to identify dependencies
946    /// and construct valid projected orderings during projection operations.
947    ///
948    /// # Parameters
949    ///
950    /// - `mapping`: A reference to the `ProjectionMapping` that defines the
951    ///   relationship between source and target expressions.
952    ///
953    /// # Returns
954    ///
955    /// A [`DependencyMap`] representing the dependency map, where each
956    /// \[`DependencyNode`\] contains dependencies for the key [`PhysicalSortExpr`].
957    ///
958    /// # Example
959    ///
960    /// Assume we have two equivalent orderings: `[a ASC, b ASC]` and `[a ASC, c ASC]`,
961    /// and the projection mapping is `[a -> a_new, b -> b_new, b + c -> b + c]`.
962    /// Then, the dependency map will be:
963    ///
964    /// ```text
965    /// a ASC: Node {Some(a_new ASC), HashSet{}}
966    /// b ASC: Node {Some(b_new ASC), HashSet{a ASC}}
967    /// c ASC: Node {None, HashSet{a ASC}}
968    /// ```
969    fn construct_dependency_map(
970        &self,
971        oeq_class: OrderingEquivalenceClass,
972        mapping: &ProjectionMapping,
973    ) -> DependencyMap {
974        let mut map = DependencyMap::default();
975        for ordering in oeq_class.into_iter() {
976            // Previous expression is a dependency. Note that there is no
977            // dependency for the leading expression.
978            if !self.insert_to_dependency_map(
979                mapping,
980                ordering[0].clone(),
981                None,
982                &mut map,
983            ) {
984                continue;
985            }
986            for (dependency, sort_expr) in ordering.into_iter().tuple_windows() {
987                if !self.insert_to_dependency_map(
988                    mapping,
989                    sort_expr,
990                    Some(dependency),
991                    &mut map,
992                ) {
993                    // If we can't project, stop constructing the dependency map
994                    // as remaining dependencies will be invalid post projection.
995                    break;
996                }
997            }
998        }
999        map
1000    }
1001
1002    /// Projects the sort expression according to the projection mapping and
1003    /// inserts it into the dependency map with the given dependency. Returns
1004    /// a boolean flag indicating whether the given expression is projectable.
1005    fn insert_to_dependency_map(
1006        &self,
1007        mapping: &ProjectionMapping,
1008        sort_expr: PhysicalSortExpr,
1009        dependency: Option<PhysicalSortExpr>,
1010        map: &mut DependencyMap,
1011    ) -> bool {
1012        let target_sort_expr = self
1013            .project_expr(&sort_expr.expr, mapping)
1014            .map(|expr| PhysicalSortExpr::new(expr, sort_expr.options));
1015        let projectable = target_sort_expr.is_some();
1016        if projectable
1017            || mapping
1018                .iter()
1019                .any(|(source, _)| expr_refers(source, &sort_expr.expr))
1020        {
1021            // Add sort expressions that can be projected or referred to
1022            // by any of the projection expressions to the dependency map:
1023            map.insert(sort_expr, target_sort_expr, dependency);
1024        }
1025        projectable
1026    }
1027
1028    /// Returns a new `ProjectionMapping` where source expressions are in normal
1029    /// form. Normalization ensures that source expressions are transformed into
1030    /// a consistent representation, which is beneficial for algorithms that rely
1031    /// on exact equalities, as it allows for more precise and reliable comparisons.
1032    ///
1033    /// # Parameters
1034    ///
1035    /// - `mapping`: A reference to the original `ProjectionMapping` to normalize.
1036    ///
1037    /// # Returns
1038    ///
1039    /// A new `ProjectionMapping` with source expressions in normal form.
1040    fn normalize_mapping(&self, mapping: &ProjectionMapping) -> ProjectionMapping {
1041        mapping
1042            .iter()
1043            .map(|(source, target)| {
1044                let normal_source = self.eq_group.normalize_expr(Arc::clone(source));
1045                (normal_source, target.clone())
1046            })
1047            .collect()
1048    }
1049
1050    /// Computes projected orderings based on a given projection mapping.
1051    ///
1052    /// This function takes a `ProjectionMapping` and computes the possible
1053    /// orderings for the projected expressions. It considers dependencies
1054    /// between expressions and generates valid orderings according to the
1055    /// specified sort properties.
1056    ///
1057    /// # Parameters
1058    ///
1059    /// - `mapping`: A reference to the `ProjectionMapping` that defines the
1060    ///   relationship between source and target expressions.
1061    /// - `oeq_class`: The `OrderingEquivalenceClass` containing the orderings
1062    ///   to project.
1063    ///
1064    /// # Returns
1065    ///
1066    /// A vector of all valid (but not in normal form) orderings after projection.
1067    fn projected_orderings(
1068        &self,
1069        mapping: &ProjectionMapping,
1070        mut oeq_class: OrderingEquivalenceClass,
1071    ) -> Vec<LexOrdering> {
1072        // Normalize source expressions in the mapping:
1073        let mapping = self.normalize_mapping(mapping);
1074        // Get dependency map for existing orderings:
1075        oeq_class = Self::substitute_oeq_class(&self.schema, &mapping, oeq_class);
1076        let dependency_map = self.construct_dependency_map(oeq_class, &mapping);
1077        let orderings = mapping.iter().flat_map(|(source, targets)| {
1078            referred_dependencies(&dependency_map, source)
1079                .into_iter()
1080                .filter_map(|deps| {
1081                    let ep = get_expr_properties(source, &deps, &self.schema);
1082                    let sort_properties = ep.map(|prop| prop.sort_properties);
1083                    if let Ok(SortProperties::Ordered(options)) = sort_properties {
1084                        Some((options, deps))
1085                    } else {
1086                        // Do not consider unordered cases.
1087                        None
1088                    }
1089                })
1090                .flat_map(|(options, relevant_deps)| {
1091                    // Generate dependent orderings (i.e. prefixes for targets):
1092                    let dependency_orderings =
1093                        generate_dependency_orderings(&relevant_deps, &dependency_map);
1094                    let sort_exprs = targets.iter().map(|(target, _)| {
1095                        PhysicalSortExpr::new(Arc::clone(target), options)
1096                    });
1097                    if dependency_orderings.is_empty() {
1098                        sort_exprs.map(|sort_expr| [sort_expr].into()).collect()
1099                    } else {
1100                        sort_exprs
1101                            .flat_map(|sort_expr| {
1102                                let mut result = dependency_orderings.clone();
1103                                for ordering in result.iter_mut() {
1104                                    ordering.push(sort_expr.clone());
1105                                }
1106                                result
1107                            })
1108                            .collect::<Vec<_>>()
1109                    }
1110                })
1111        });
1112
1113        // Add valid projected orderings. For example, if existing ordering is
1114        // `a + b` and projection is `[a -> a_new, b -> b_new]`, we need to
1115        // preserve `a_new + b_new` as ordered. Please note that `a_new` and
1116        // `b_new` themselves need not be ordered. Such dependencies cannot be
1117        // deduced via the pass above.
1118        let projected_orderings = dependency_map.iter().flat_map(|(sort_expr, node)| {
1119            let mut prefixes = construct_prefix_orderings(sort_expr, &dependency_map);
1120            if prefixes.is_empty() {
1121                // If prefix is empty, there is no dependency. Insert
1122                // empty ordering:
1123                if let Some(target) = &node.target {
1124                    prefixes.push([target.clone()].into());
1125                }
1126            } else {
1127                // Append current ordering on top its dependencies:
1128                for ordering in prefixes.iter_mut() {
1129                    if let Some(target) = &node.target {
1130                        ordering.push(target.clone());
1131                    }
1132                }
1133            }
1134            prefixes
1135        });
1136
1137        // Simplify each ordering by removing redundant sections:
1138        orderings.chain(projected_orderings).collect()
1139    }
1140
1141    /// Projects constraints according to the given projection mapping.
1142    ///
1143    /// This function takes a projection mapping and extracts column indices of
1144    /// target columns. It then projects the constraints to only include
1145    /// relationships between columns that exist in the projected output.
1146    ///
1147    /// # Parameters
1148    ///
1149    /// * `mapping` - A reference to the `ProjectionMapping` that defines the
1150    ///   projection operation.
1151    ///
1152    /// # Returns
1153    ///
1154    /// Returns an optional `Constraints` object containing only the constraints
1155    /// that are valid for the projected columns (if any exists).
1156    fn projected_constraints(&self, mapping: &ProjectionMapping) -> Option<Constraints> {
1157        let indices = mapping
1158            .iter()
1159            .flat_map(|(_, targets)| {
1160                targets.iter().flat_map(|(target, _)| {
1161                    target.downcast_ref::<Column>().map(|c| c.index())
1162                })
1163            })
1164            .collect::<Vec<_>>();
1165        self.constraints.project(&indices)
1166    }
1167
1168    /// Projects the equivalences within according to `mapping` and
1169    /// `output_schema`.
1170    pub fn project(&self, mapping: &ProjectionMapping, output_schema: SchemaRef) -> Self {
1171        let eq_group = self.eq_group.project(mapping);
1172        let orderings =
1173            self.projected_orderings(mapping, self.oeq_cache.normal_cls.clone());
1174        let normal_orderings = orderings
1175            .iter()
1176            .cloned()
1177            .map(|o| eq_group.normalize_sort_exprs(o));
1178        Self {
1179            oeq_cache: OrderingEquivalenceCache::new(normal_orderings),
1180            oeq_class: OrderingEquivalenceClass::new(orderings),
1181            constraints: self.projected_constraints(mapping).unwrap_or_default(),
1182            schema: output_schema,
1183            eq_group,
1184        }
1185    }
1186
1187    /// Returns the longest (potentially partial) permutation satisfying the
1188    /// existing ordering. For example, if we have the equivalent orderings
1189    /// `[a ASC, b ASC]` and `[c DESC]`, with `exprs` containing `[c, b, a, d]`,
1190    /// then this function returns `([a ASC, b ASC, c DESC], [2, 1, 0])`.
1191    /// This means that the specification `[a ASC, b ASC, c DESC]` is satisfied
1192    /// by the existing ordering, and `[a, b, c]` resides at indices: `2, 1, 0`
1193    /// inside the argument `exprs` (respectively). For the mathematical
1194    /// definition of "partial permutation", see:
1195    ///
1196    /// <https://en.wikipedia.org/wiki/Permutation#k-permutations_of_n>
1197    pub fn find_longest_permutation(
1198        &self,
1199        exprs: &[Arc<dyn PhysicalExpr>],
1200    ) -> Result<(Vec<PhysicalSortExpr>, Vec<usize>)> {
1201        let mut eq_properties = self.clone();
1202        let mut result = vec![];
1203        // The algorithm is as follows:
1204        // - Iterate over all the expressions and insert ordered expressions
1205        //   into the result.
1206        // - Treat inserted expressions as constants (i.e. add them as constants
1207        //   to the state).
1208        // - Continue the above procedure until no expression is inserted; i.e.
1209        //   the algorithm reaches a fixed point.
1210        // This algorithm should reach a fixed point in at most `exprs.len()`
1211        // iterations.
1212        let mut search_indices = (0..exprs.len()).collect::<IndexSet<_>>();
1213        for _ in 0..exprs.len() {
1214            // Get ordered expressions with their indices.
1215            let ordered_exprs = search_indices
1216                .iter()
1217                .filter_map(|&idx| {
1218                    let ExprProperties {
1219                        sort_properties, ..
1220                    } = eq_properties.get_expr_properties(Arc::clone(&exprs[idx]));
1221                    match sort_properties {
1222                        SortProperties::Ordered(options) => {
1223                            let expr = Arc::clone(&exprs[idx]);
1224                            Some((PhysicalSortExpr::new(expr, options), idx))
1225                        }
1226                        SortProperties::Singleton => {
1227                            // Assign default ordering to constant expressions:
1228                            let expr = Arc::clone(&exprs[idx]);
1229                            Some((PhysicalSortExpr::new_default(expr), idx))
1230                        }
1231                        SortProperties::Unordered => None,
1232                    }
1233                })
1234                .collect::<Vec<_>>();
1235            // We reached a fixed point, exit.
1236            if ordered_exprs.is_empty() {
1237                break;
1238            }
1239            // Remove indices that have an ordering from `search_indices`, and
1240            // treat ordered expressions as constants in subsequent iterations.
1241            // We can do this because the "next" key only matters in a lexicographical
1242            // ordering when the keys to its left have the same values.
1243            //
1244            // Note that these expressions are not properly "constants". This is just
1245            // an implementation strategy confined to this function.
1246            for (PhysicalSortExpr { expr, .. }, idx) in &ordered_exprs {
1247                let const_expr = ConstExpr::from(Arc::clone(expr));
1248                eq_properties.add_constants(std::iter::once(const_expr))?;
1249                search_indices.shift_remove(idx);
1250            }
1251            // Add new ordered section to the state.
1252            result.extend(ordered_exprs);
1253        }
1254        Ok(result.into_iter().unzip())
1255    }
1256
1257    /// This function determines whether the provided expression is constant
1258    /// based on the known constants. For example, if columns `a` and `b` are
1259    /// constant, then expressions `a`, `b` and `a + b` will all return `true`
1260    /// whereas expression `c` will return `false`.
1261    ///
1262    /// # Parameters
1263    ///
1264    /// - `expr`: A reference to a `Arc<dyn PhysicalExpr>` representing the
1265    ///   expression to be checked.
1266    ///
1267    /// # Returns
1268    ///
1269    /// Returns a `Some` value if the expression is constant according to
1270    /// equivalence group, and `None` otherwise. The `Some` variant contains
1271    /// an `AcrossPartitions` value indicating whether the expression is
1272    /// constant across partitions, and its actual value (if available).
1273    pub fn is_expr_constant(
1274        &self,
1275        expr: &Arc<dyn PhysicalExpr>,
1276    ) -> Option<AcrossPartitions> {
1277        self.eq_group.is_expr_constant(expr)
1278    }
1279
1280    /// Retrieves the properties for a given physical expression.
1281    ///
1282    /// This function constructs an [`ExprProperties`] object for the given
1283    /// expression, which encapsulates information about the expression's
1284    /// properties, including its [`SortProperties`] and [`Interval`].
1285    ///
1286    /// # Parameters
1287    ///
1288    /// - `expr`: An `Arc<dyn PhysicalExpr>` representing the physical expression
1289    ///   for which ordering information is sought.
1290    ///
1291    /// # Returns
1292    ///
1293    /// Returns an [`ExprProperties`] object containing the ordering and range
1294    /// information for the given expression.
1295    pub fn get_expr_properties(&self, expr: Arc<dyn PhysicalExpr>) -> ExprProperties {
1296        ExprPropertiesNode::new_unknown(expr)
1297            .transform_up(|expr| update_properties(expr, self))
1298            .data()
1299            .map(|node| node.data)
1300            .unwrap_or_else(|_| ExprProperties::new_unknown())
1301    }
1302
1303    /// Transforms this `EquivalenceProperties` by mapping columns in the
1304    /// original schema to columns in the new schema by index.
1305    pub fn with_new_schema(mut self, schema: SchemaRef) -> Result<Self> {
1306        // The new schema and the original schema is aligned when they have the
1307        // same number of columns, and fields at the same index have the same
1308        // type in both schemas.
1309        let schemas_aligned = (self.schema.fields.len() == schema.fields.len())
1310            && self
1311                .schema
1312                .fields
1313                .iter()
1314                .zip(schema.fields.iter())
1315                .all(|(lhs, rhs)| lhs.data_type().eq(rhs.data_type()));
1316        if !schemas_aligned {
1317            // Rewriting equivalence properties in terms of new schema is not
1318            // safe when schemas are not aligned:
1319            return plan_err!(
1320                "Schemas have to be aligned to rewrite equivalences:\n Old schema: {}\n New schema: {}",
1321                self.schema,
1322                schema
1323            );
1324        }
1325
1326        // Rewrite equivalence classes according to the new schema:
1327        let mut eq_classes = vec![];
1328        for mut eq_class in self.eq_group {
1329            // Rewrite the expressions in the equivalence class:
1330            eq_class.exprs = eq_class
1331                .exprs
1332                .into_iter()
1333                .map(|expr| with_new_schema(expr, &schema))
1334                .collect::<Result<_>>()?;
1335            // Rewrite the constant value (if available and known):
1336            let data_type = eq_class
1337                .canonical_expr()
1338                .map(|e| e.data_type(&schema))
1339                .transpose()?;
1340            if let (Some(data_type), Some(AcrossPartitions::Uniform(Some(value)))) =
1341                (data_type, &mut eq_class.constant)
1342            {
1343                match value.cast_to(&data_type) {
1344                    Ok(cast_value) => *value = cast_value,
1345                    Err(_) => {
1346                        // This is optimizer metadata. If a stale constant
1347                        // value cannot be represented after schema rewrite,
1348                        // drop the constant instead of failing planning.
1349                        eq_class.constant = None;
1350                    }
1351                }
1352            }
1353            if eq_class.is_trivial() {
1354                continue;
1355            }
1356            eq_classes.push(eq_class);
1357        }
1358        self.eq_group = eq_classes.into();
1359
1360        // Rewrite orderings according to new schema:
1361        self.oeq_class = self.oeq_class.with_new_schema(&schema)?;
1362        self.oeq_cache.normal_cls = self.oeq_cache.normal_cls.with_new_schema(&schema)?;
1363
1364        // Update the schema:
1365        self.schema = schema;
1366
1367        Ok(self)
1368    }
1369}
1370
1371impl From<EquivalenceProperties> for OrderingEquivalenceClass {
1372    fn from(eq_properties: EquivalenceProperties) -> Self {
1373        eq_properties.oeq_class
1374    }
1375}
1376
1377/// More readable display version of the `EquivalenceProperties`.
1378///
1379/// Format:
1380/// ```text
1381/// order: [[b@1 ASC NULLS LAST]], eq: [{members: [a@0], constant: (heterogeneous)}]
1382/// ```
1383impl Display for EquivalenceProperties {
1384    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1385        let empty_eq_group = self.eq_group.is_empty();
1386        let empty_oeq_class = self.oeq_class.is_empty();
1387        if empty_oeq_class && empty_eq_group {
1388            write!(f, "No properties")?;
1389        } else if !empty_oeq_class {
1390            write!(f, "order: {}", self.oeq_class)?;
1391            if !empty_eq_group {
1392                write!(f, ", eq: {}", self.eq_group)?;
1393            }
1394        } else {
1395            write!(f, "eq: {}", self.eq_group)?;
1396        }
1397        Ok(())
1398    }
1399}
1400
1401/// Calculates the properties of a given [`ExprPropertiesNode`].
1402///
1403/// Order information can be retrieved as:
1404/// - If it is a leaf node, we directly find the order of the node by looking
1405///   at the given sort expression and equivalence properties if it is a `Column`
1406///   leaf, or we mark it as unordered. In the case of a `Literal` leaf, we mark
1407///   it as singleton so that it can cooperate with all ordered columns.
1408/// - If it is an intermediate node, the children states matter. Each `PhysicalExpr`
1409///   and operator has its own rules on how to propagate the children orderings.
1410///   However, before we engage in recursion, we check whether this intermediate
1411///   node directly matches with the sort expression. If there is a match, the
1412///   sort expression emerges at that node immediately, discarding the recursive
1413///   result coming from its children.
1414///
1415/// Range information is calculated as:
1416/// - If it is a `Literal` node, we set the range as a point value. If it is a
1417///   `Column` node, we set the datatype of the range, but cannot give an interval
1418///   for the range, yet.
1419/// - If it is an intermediate node, the children states matter. Each `PhysicalExpr`
1420///   and operator has its own rules on how to propagate the children range.
1421fn update_properties(
1422    mut node: ExprPropertiesNode,
1423    eq_properties: &EquivalenceProperties,
1424) -> Result<Transformed<ExprPropertiesNode>> {
1425    // First, try to gather the information from the children:
1426    if !node.expr.children().is_empty() {
1427        // We have an intermediate (non-leaf) node, account for its children:
1428        let children_props = node.children.iter().map(|c| c.data.clone()).collect_vec();
1429        node.data = node.expr.get_properties(&children_props)?;
1430    } else if node.expr.is::<Literal>() {
1431        // We have a Literal, which is one of the two possible leaf node types:
1432        node.data = node.expr.get_properties(&[])?;
1433    } else if node.expr.is::<Column>() {
1434        // We have a Column, which is the other possible leaf node type:
1435        node.data.range =
1436            Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)?;
1437        // A column is the identity mapping of itself, which is trivially
1438        // strict:
1439        node.data.strictly_order_preserving = true;
1440    }
1441    // Now, check what we know about orderings:
1442    let normal_expr = eq_properties
1443        .eq_group
1444        .normalize_expr(Arc::clone(&node.expr));
1445    let oeq_class = &eq_properties.oeq_cache.normal_cls;
1446    if eq_properties.is_expr_constant(&normal_expr).is_some()
1447        || oeq_class.is_expr_partial_const(&normal_expr)
1448    {
1449        node.data.sort_properties = SortProperties::Singleton;
1450    } else if let Some(options) = oeq_class.get_options(&normal_expr) {
1451        node.data.sort_properties = SortProperties::Ordered(options);
1452    }
1453    Ok(Transformed::yes(node))
1454}
1455
1456/// This function examines whether a referring expression directly refers to a
1457/// given referred expression or if any of its children in the expression tree
1458/// refer to the specified expression.
1459///
1460/// # Parameters
1461///
1462/// - `referring_expr`: A reference to the referring expression (`Arc<dyn PhysicalExpr>`).
1463/// - `referred_expr`: A reference to the referred expression (`Arc<dyn PhysicalExpr>`)
1464///
1465/// # Returns
1466///
1467/// A boolean value indicating whether `referring_expr` refers (needs it to evaluate its result)
1468/// `referred_expr` or not.
1469fn expr_refers(
1470    referring_expr: &Arc<dyn PhysicalExpr>,
1471    referred_expr: &Arc<dyn PhysicalExpr>,
1472) -> bool {
1473    referring_expr.eq(referred_expr)
1474        || referring_expr
1475            .children()
1476            .iter()
1477            .any(|child| expr_refers(child, referred_expr))
1478}
1479
1480/// This function examines the given expression and its properties to determine
1481/// the ordering properties of the expression. The range knowledge is not utilized
1482/// yet in the scope of this function.
1483///
1484/// # Parameters
1485///
1486/// - `expr`: A reference to the source expression (`Arc<dyn PhysicalExpr>`) for
1487///   which ordering properties need to be determined.
1488/// - `dependencies`: A reference to `Dependencies`, containing sort expressions
1489///   referred to by `expr`.
1490/// - `schema``: A reference to the schema which the `expr` columns refer.
1491///
1492/// # Returns
1493///
1494/// A `SortProperties` indicating the ordering information of the given expression.
1495fn get_expr_properties(
1496    expr: &Arc<dyn PhysicalExpr>,
1497    dependencies: &Dependencies,
1498    schema: &SchemaRef,
1499) -> Result<ExprProperties> {
1500    if let Some(column_order) = dependencies.iter().find(|&order| expr.eq(&order.expr)) {
1501        // If exact match is found, return its ordering. This is a base case
1502        // of the recursion: the expression is treated as an atomic ordered
1503        // input from here on, so `strictly_order_preserving` states only that
1504        // it is a one-to-one mapping *of itself* (the identity), which holds
1505        // for any expression. It makes no claim about the expression being
1506        // one-to-one in its own inputs (e.g. `floor(x)` as a sort key), and
1507        // it does not need to: parent expressions are substituted for this
1508        // sort key, so their strictness only has to be relative to it.
1509        Ok(ExprProperties {
1510            sort_properties: SortProperties::Ordered(column_order.options),
1511            range: Interval::make_unbounded(&expr.data_type(schema)?)?,
1512            preserves_lex_ordering: false,
1513            strictly_order_preserving: true,
1514        })
1515    } else if expr.downcast_ref::<Column>().is_some() {
1516        Ok(ExprProperties {
1517            sort_properties: SortProperties::Unordered,
1518            range: Interval::make_unbounded(&expr.data_type(schema)?)?,
1519            preserves_lex_ordering: false,
1520            // A base case of the recursion: a column is the identity mapping
1521            // of itself, which is trivially one-to-one.
1522            strictly_order_preserving: true,
1523        })
1524    } else if let Some(literal) = expr.downcast_ref::<Literal>() {
1525        Ok(ExprProperties {
1526            sort_properties: SortProperties::Singleton,
1527            range: literal.value().into(),
1528            preserves_lex_ordering: true,
1529            // Vacuously true: a literal has no ordered inputs.
1530            strictly_order_preserving: true,
1531        })
1532    } else {
1533        // Find orderings of its children
1534        let child_states = expr
1535            .children()
1536            .iter()
1537            .map(|child| get_expr_properties(child, dependencies, schema))
1538            .collect::<Result<Vec<_>>>()?;
1539        // Calculate expression ordering using ordering of its children.
1540        expr.get_properties(&child_states)
1541    }
1542}