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}