datafusion_physical_expr/equivalence/ordering.rs
1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements. See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership. The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License. You may obtain a copy of the License at
8//
9// http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied. See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::fmt::Display;
19use std::ops::Deref;
20use std::sync::Arc;
21use std::vec::IntoIter;
22
23use crate::expressions::with_new_schema;
24use crate::{LexOrdering, PhysicalExpr, add_offset_to_physical_sort_exprs};
25
26use arrow::compute::SortOptions;
27use arrow::datatypes::SchemaRef;
28use datafusion_common::{HashSet, Result};
29use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
30
31/// An `OrderingEquivalenceClass` keeps track of distinct alternative orderings
32/// than can describe a table. For example, consider the following table:
33///
34/// ```text
35/// ┌───┬───┬───┬───┐
36/// │ a │ b │ c │ d │
37/// ├───┼───┼───┼───┤
38/// │ 1 │ 4 │ 3 │ 1 │
39/// │ 2 │ 3 │ 3 │ 2 │
40/// │ 3 │ 1 │ 2 │ 2 │
41/// │ 3 │ 2 │ 1 │ 3 │
42/// └───┴───┴───┴───┘
43/// ```
44///
45/// Here, both `[a ASC, b ASC]` and `[c DESC, d ASC]` describe the table
46/// ordering. In this case, we say that these orderings are equivalent.
47///
48/// An `OrderingEquivalenceClass` is a set of such equivalent orderings, which
49/// is represented by a vector of `LexOrdering`s. The set does not store any
50/// redundant information by enforcing the invariant that no suffix of an
51/// ordering in the equivalence class is a prefix of another ordering in the
52/// equivalence class. The set can be empty, which means that there are no
53/// orderings that describe the table.
54#[derive(Clone, Debug, Default, Eq, PartialEq)]
55pub struct OrderingEquivalenceClass {
56 orderings: Vec<LexOrdering>,
57}
58
59impl OrderingEquivalenceClass {
60 /// Clears (empties) this ordering equivalence class.
61 pub fn clear(&mut self) {
62 self.orderings.clear();
63 }
64
65 /// Creates a new ordering equivalence class from the given orderings
66 /// and removes any redundant entries (if given).
67 pub fn new(
68 orderings: impl IntoIterator<Item = impl IntoIterator<Item = PhysicalSortExpr>>,
69 ) -> Self {
70 let mut result = Self {
71 orderings: orderings.into_iter().filter_map(LexOrdering::new).collect(),
72 };
73 result.remove_redundant_entries();
74 result
75 }
76
77 /// Extend this ordering equivalence class with the given orderings.
78 pub fn extend(&mut self, orderings: impl IntoIterator<Item = LexOrdering>) {
79 self.orderings.extend(orderings);
80 // Make sure that there are no redundant orderings:
81 self.remove_redundant_entries();
82 }
83
84 /// Adds new orderings into this ordering equivalence class.
85 pub fn add_orderings(
86 &mut self,
87 sort_exprs: impl IntoIterator<Item = impl IntoIterator<Item = PhysicalSortExpr>>,
88 ) {
89 self.orderings
90 .extend(sort_exprs.into_iter().filter_map(LexOrdering::new));
91 // Make sure that there are no redundant orderings:
92 self.remove_redundant_entries();
93 }
94
95 /// Removes redundant orderings from this ordering equivalence class.
96 ///
97 /// For instance, if we already have the ordering `[a ASC, b ASC, c DESC]`,
98 /// then there is no need to keep ordering `[a ASC, b ASC]` in the state.
99 fn remove_redundant_entries(&mut self) {
100 let mut work = true;
101 while work {
102 work = false;
103 let mut idx = 0;
104 'outer: while idx < self.orderings.len() {
105 let mut ordering_idx = idx + 1;
106 while ordering_idx < self.orderings.len() {
107 if let Some(remove) = self.resolve_overlap(idx, ordering_idx) {
108 work = true;
109 if remove {
110 self.orderings.swap_remove(idx);
111 continue 'outer;
112 }
113 }
114 if let Some(remove) = self.resolve_overlap(ordering_idx, idx) {
115 work = true;
116 if remove {
117 self.orderings.swap_remove(ordering_idx);
118 continue;
119 }
120 }
121 ordering_idx += 1;
122 }
123 idx += 1;
124 }
125 }
126 }
127
128 /// Trims `orderings[idx]` if some suffix of it overlaps with a prefix of
129 /// `orderings[pre_idx]`. If there is any overlap, returns a `Some(true)`
130 /// if any trimming took place, and `Some(false)` otherwise. If there is
131 /// no overlap, returns `None`.
132 ///
133 /// For example, if `orderings[idx]` is `[a ASC, b ASC, c DESC]` and
134 /// `orderings[pre_idx]` is `[b ASC, c DESC]`, then the function will trim
135 /// `orderings[idx]` to `[a ASC]`.
136 fn resolve_overlap(&mut self, idx: usize, pre_idx: usize) -> Option<bool> {
137 let length = self.orderings[idx].len();
138 let other_length = self.orderings[pre_idx].len();
139 for overlap in 1..=length.min(other_length) {
140 if self.orderings[idx][length - overlap..]
141 == self.orderings[pre_idx][..overlap]
142 {
143 return Some(!self.orderings[idx].truncate(length - overlap));
144 }
145 }
146 None
147 }
148
149 /// Returns the concatenation of all the orderings. This enables merge
150 /// operations to preserve all equivalent orderings simultaneously.
151 pub fn output_ordering(&self) -> Option<LexOrdering> {
152 self.orderings.iter().cloned().reduce(|mut cat, o| {
153 cat.extend(o);
154 cat
155 })
156 }
157
158 // Append orderings in `other` to all existing orderings in this ordering
159 // equivalence class.
160 pub fn join_suffix(mut self, other: &Self) -> Self {
161 let n_ordering = self.orderings.len();
162 // Replicate entries before cross product:
163 let n_cross = std::cmp::max(n_ordering, other.len() * n_ordering);
164 self.orderings = self.orderings.into_iter().cycle().take(n_cross).collect();
165 // Append sort expressions of `other` to the current orderings:
166 for (outer_idx, ordering) in other.iter().enumerate() {
167 let base = outer_idx * n_ordering;
168 // Use the cross product index:
169 for idx in base..(base + n_ordering) {
170 self.orderings[idx].extend(ordering.iter().cloned());
171 }
172 }
173 self
174 }
175
176 /// Adds `offset` value to the index of each expression inside this
177 /// ordering equivalence class.
178 pub fn add_offset(&mut self, offset: isize) -> Result<()> {
179 let orderings = std::mem::take(&mut self.orderings);
180 for ordering_result in orderings
181 .into_iter()
182 .map(|o| add_offset_to_physical_sort_exprs(o, offset))
183 {
184 self.orderings.extend(LexOrdering::new(ordering_result?));
185 }
186 Ok(())
187 }
188
189 /// Transforms this `OrderingEquivalenceClass` by mapping columns in the
190 /// original schema to columns in the new schema by index. The new schema
191 /// and the original schema needs to be aligned; i.e. they should have the
192 /// same number of columns, and fields at the same index have the same type
193 /// in both schemas.
194 pub fn with_new_schema(mut self, schema: &SchemaRef) -> Result<Self> {
195 self.orderings = self
196 .orderings
197 .into_iter()
198 .map(|ordering| {
199 ordering
200 .into_iter()
201 .map(|mut sort_expr| {
202 sort_expr.expr = with_new_schema(sort_expr.expr, schema)?;
203 Ok(sort_expr)
204 })
205 .collect::<Result<Vec<_>>>()
206 // The following `unwrap` is safe because the vector will always
207 // be non-empty.
208 .map(|v| LexOrdering::new(v).unwrap())
209 })
210 .collect::<Result<_>>()?;
211 Ok(self)
212 }
213
214 /// Gets sort options associated with this expression if it is a leading
215 /// ordering expression. Otherwise, returns `None`.
216 pub fn get_options(&self, expr: &Arc<dyn PhysicalExpr>) -> Option<SortOptions> {
217 for ordering in self.iter() {
218 let leading_ordering = &ordering[0];
219 if leading_ordering.expr.eq(expr) {
220 return Some(leading_ordering.options);
221 }
222 }
223 None
224 }
225
226 /// Checks whether the given expression is partially constant according to
227 /// this ordering equivalence class.
228 ///
229 /// This function determines whether `expr` appears in at least one combination
230 /// of `descending` and `nulls_first` options that indicate partial constantness
231 /// in a lexicographical ordering. Specifically, an expression is considered
232 /// a partial constant in this context if its `SortOptions` satisfies either
233 /// of the following conditions:
234 /// - It is `descending` with `nulls_first` and _also_ `ascending` with
235 /// `nulls_last`, OR
236 /// - It is `descending` with `nulls_last` and _also_ `ascending` with
237 /// `nulls_first`.
238 ///
239 /// The equivalence mechanism primarily uses `ConstExpr`s to represent globally
240 /// constant expressions. However, some expressions may only be partially
241 /// constant within a lexicographical ordering. This function helps identify
242 /// such cases. If an expression is constant within a prefix ordering, it is
243 /// added as a constant during `ordering_satisfy_requirement()` iterations
244 /// after the corresponding prefix requirement is satisfied.
245 ///
246 /// ### Future Improvements
247 ///
248 /// This function may become unnecessary if any of the following improvements
249 /// are implemented:
250 /// 1. `SortOptions` supports encoding constantness information.
251 /// 2. `EquivalenceProperties` gains `FunctionalDependency` awareness, eliminating
252 /// the need for `Constant` and `Constraints`.
253 pub fn is_expr_partial_const(&self, expr: &Arc<dyn PhysicalExpr>) -> bool {
254 let mut constantness_defining_pairs = [
255 HashSet::from([(false, false), (true, true)]),
256 HashSet::from([(false, true), (true, false)]),
257 ];
258
259 for ordering in self.iter() {
260 let leading_ordering = ordering.first();
261 if leading_ordering.expr.eq(expr) {
262 let opt = (
263 leading_ordering.options.descending,
264 leading_ordering.options.nulls_first,
265 );
266 constantness_defining_pairs[0].remove(&opt);
267 constantness_defining_pairs[1].remove(&opt);
268 }
269 }
270
271 constantness_defining_pairs
272 .iter()
273 .any(|pair| pair.is_empty())
274 }
275}
276
277impl Deref for OrderingEquivalenceClass {
278 type Target = [LexOrdering];
279
280 fn deref(&self) -> &Self::Target {
281 self.orderings.as_slice()
282 }
283}
284
285impl From<Vec<LexOrdering>> for OrderingEquivalenceClass {
286 fn from(orderings: Vec<LexOrdering>) -> Self {
287 let mut result = Self { orderings };
288 result.remove_redundant_entries();
289 result
290 }
291}
292
293/// Convert the `OrderingEquivalenceClass` into an iterator of `LexOrdering`s.
294impl IntoIterator for OrderingEquivalenceClass {
295 type Item = LexOrdering;
296 type IntoIter = IntoIter<Self::Item>;
297
298 fn into_iter(self) -> Self::IntoIter {
299 self.orderings.into_iter()
300 }
301}
302
303impl Display for OrderingEquivalenceClass {
304 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
305 write!(f, "[")?;
306 let mut iter = self.orderings.iter();
307 if let Some(ordering) = iter.next() {
308 write!(f, "[{ordering}]")?;
309 }
310 for ordering in iter {
311 write!(f, ", [{ordering}]")?;
312 }
313 write!(f, "]")
314 }
315}
316
317impl From<OrderingEquivalenceClass> for Vec<LexOrdering> {
318 fn from(oeq_class: OrderingEquivalenceClass) -> Self {
319 oeq_class.orderings
320 }
321}
322
323#[cfg(test)]
324mod tests {
325 use std::sync::Arc;
326
327 use crate::equivalence::tests::create_test_schema;
328 use crate::equivalence::{
329 EquivalenceClass, EquivalenceGroup, EquivalenceProperties,
330 OrderingEquivalenceClass, convert_to_orderings, convert_to_sort_exprs,
331 };
332 use crate::expressions::{BinaryExpr, CastExpr, Column, col};
333 use crate::utils::tests::TestScalarUDF;
334 use crate::{
335 AcrossPartitions, ConstExpr, PhysicalExpr, PhysicalExprRef, PhysicalSortExpr,
336 ScalarFunctionExpr,
337 };
338
339 use arrow::compute::SortOptions;
340 use arrow::datatypes::{DataType, Field, Schema};
341 use datafusion_common::Result;
342 use datafusion_common::config::ConfigOptions;
343 use datafusion_expr::{Operator, ScalarUDF};
344
345 #[test]
346 fn test_ordering_satisfy() -> Result<()> {
347 let input_schema = Arc::new(Schema::new(vec![
348 Field::new("a", DataType::Int64, true),
349 Field::new("b", DataType::Int64, true),
350 ]));
351 let crude = vec![PhysicalSortExpr {
352 expr: Arc::new(Column::new("a", 0)),
353 options: SortOptions::default(),
354 }];
355 let finer = vec![
356 PhysicalSortExpr {
357 expr: Arc::new(Column::new("a", 0)),
358 options: SortOptions::default(),
359 },
360 PhysicalSortExpr {
361 expr: Arc::new(Column::new("b", 1)),
362 options: SortOptions::default(),
363 },
364 ];
365 // finer ordering satisfies, crude ordering should return true
366 let eq_properties_finer = EquivalenceProperties::new_with_orderings(
367 Arc::clone(&input_schema),
368 [finer.clone()],
369 );
370 assert!(eq_properties_finer.ordering_satisfy(crude.clone())?);
371
372 // Crude ordering doesn't satisfy finer ordering. should return false
373 let eq_properties_crude =
374 EquivalenceProperties::new_with_orderings(Arc::clone(&input_schema), [crude]);
375 assert!(!eq_properties_crude.ordering_satisfy(finer)?);
376 Ok(())
377 }
378
379 #[test]
380 fn test_ordering_satisfy_strictly_order_preserving() -> Result<()> {
381 let schema = Arc::new(Schema::new(vec![
382 Field::new("a", DataType::Int32, true),
383 Field::new("b", DataType::Int64, true),
384 ]));
385 let col_a = col("a", &schema)?;
386 let col_b = col("b", &schema)?;
387 let asc = SortOptions::default();
388 let sort_a = PhysicalSortExpr::new(Arc::clone(&col_a), asc);
389 let sort_b = PhysicalSortExpr::new(Arc::clone(&col_b), asc);
390 let eq_properties = EquivalenceProperties::new_with_orderings(
391 Arc::clone(&schema),
392 [vec![sort_a.clone(), sort_b.clone()]],
393 );
394
395 assert!(eq_properties.ordering_satisfy(vec![sort_a.clone(), sort_b.clone()])?);
396 assert!(eq_properties.ordering_satisfy(vec![sort_a.clone()])?);
397
398 // A widening cast is strictly order-preserving: `a` is constant
399 // within each group of equal `CAST(a AS BIGINT)` values, so `b`
400 // remains sorted within those groups.
401 let widening = Arc::new(CastExpr::new(Arc::clone(&col_a), DataType::Int64, None))
402 as PhysicalExprRef;
403 let sort_widening = PhysicalSortExpr::new(widening, asc);
404 assert!(eq_properties.ordering_satisfy(vec![sort_widening, sort_b.clone()])?);
405
406 // A narrowing cast is only monotonic: it satisfies as a leading key,
407 // but it may collapse distinct `a` values, so `b` is not guaranteed
408 // to be sorted within its tie groups.
409 let narrowing = Arc::new(CastExpr::new(Arc::clone(&col_a), DataType::Int16, None))
410 as PhysicalExprRef;
411 let sort_narrowing = PhysicalSortExpr::new(narrowing, asc);
412 assert!(eq_properties.ordering_satisfy(vec![sort_narrowing.clone()])?);
413 assert!(!eq_properties.ordering_satisfy(vec![sort_narrowing, sort_b.clone()])?);
414
415 Ok(())
416 }
417
418 #[test]
419 fn test_ordering_satisfy_with_equivalence2() -> Result<()> {
420 let test_schema = create_test_schema()?;
421 let col_a = &col("a", &test_schema)?;
422 let col_b = &col("b", &test_schema)?;
423 let col_c = &col("c", &test_schema)?;
424 let col_d = &col("d", &test_schema)?;
425 let col_e = &col("e", &test_schema)?;
426 let col_f = &col("f", &test_schema)?;
427 let test_fun = Arc::new(ScalarUDF::new_from_impl(TestScalarUDF::new()));
428
429 let floor_a = Arc::new(ScalarFunctionExpr::try_new(
430 Arc::clone(&test_fun),
431 vec![Arc::clone(col_a)],
432 &test_schema,
433 Arc::new(ConfigOptions::default()),
434 )?) as PhysicalExprRef;
435 let floor_f = Arc::new(ScalarFunctionExpr::try_new(
436 Arc::clone(&test_fun),
437 vec![Arc::clone(col_f)],
438 &test_schema,
439 Arc::new(ConfigOptions::default()),
440 )?) as PhysicalExprRef;
441 let exp_a = Arc::new(ScalarFunctionExpr::try_new(
442 Arc::clone(&test_fun),
443 vec![Arc::clone(col_a)],
444 &test_schema,
445 Arc::new(ConfigOptions::default()),
446 )?) as PhysicalExprRef;
447
448 let a_plus_b = Arc::new(BinaryExpr::new(
449 Arc::clone(col_a),
450 Operator::Plus,
451 Arc::clone(col_b),
452 )) as Arc<dyn PhysicalExpr>;
453 let options = SortOptions {
454 descending: false,
455 nulls_first: false,
456 };
457
458 let test_cases = vec![
459 // ------------ TEST CASE 1 ------------
460 (
461 // orderings
462 vec![
463 // [a ASC, d ASC, b ASC]
464 vec![(col_a, options), (col_d, options), (col_b, options)],
465 // [c ASC]
466 vec![(col_c, options)],
467 ],
468 // equivalence classes
469 vec![vec![col_a, col_f]],
470 // constants
471 vec![col_e],
472 // requirement [a ASC, b ASC], requirement is not satisfied.
473 vec![(col_a, options), (col_b, options)],
474 // expected: requirement is not satisfied.
475 false,
476 ),
477 // ------------ TEST CASE 2 ------------
478 (
479 // orderings
480 vec![
481 // [a ASC, c ASC, b ASC]
482 vec![(col_a, options), (col_c, options), (col_b, options)],
483 // [d ASC]
484 vec![(col_d, options)],
485 ],
486 // equivalence classes
487 vec![vec![col_a, col_f]],
488 // constants
489 vec![col_e],
490 // requirement [floor(a) ASC],
491 vec![(&floor_a, options)],
492 // expected: requirement is satisfied.
493 true,
494 ),
495 // ------------ TEST CASE 2.1 ------------
496 (
497 // orderings
498 vec![
499 // [a ASC, c ASC, b ASC]
500 vec![(col_a, options), (col_c, options), (col_b, options)],
501 // [d ASC]
502 vec![(col_d, options)],
503 ],
504 // equivalence classes
505 vec![vec![col_a, col_f]],
506 // constants
507 vec![col_e],
508 // requirement [floor(f) ASC], (Please note that a=f)
509 vec![(&floor_f, options)],
510 // expected: requirement is satisfied.
511 true,
512 ),
513 // ------------ TEST CASE 3 ------------
514 (
515 // orderings
516 vec![
517 // [a ASC, c ASC, b ASC]
518 vec![(col_a, options), (col_c, options), (col_b, options)],
519 // [d ASC]
520 vec![(col_d, options)],
521 ],
522 // equivalence classes
523 vec![vec![col_a, col_f]],
524 // constants
525 vec![col_e],
526 // requirement [a ASC, c ASC, a+b ASC],
527 vec![(col_a, options), (col_c, options), (&a_plus_b, options)],
528 // expected: requirement is not satisfied because addition can wrap.
529 false,
530 ),
531 // ------------ TEST CASE 4 ------------
532 (
533 // orderings
534 vec![
535 // [a ASC, b ASC, c ASC, d ASC]
536 vec![
537 (col_a, options),
538 (col_b, options),
539 (col_c, options),
540 (col_d, options),
541 ],
542 ],
543 // equivalence classes
544 vec![vec![col_a, col_f]],
545 // constants
546 vec![col_e],
547 // requirement [floor(a) ASC, a+b ASC],
548 vec![(&floor_a, options), (&a_plus_b, options)],
549 // expected: requirement is satisfied.
550 false,
551 ),
552 // ------------ TEST CASE 5 ------------
553 (
554 // orderings
555 vec![
556 // [a ASC, b ASC, c ASC, d ASC]
557 vec![
558 (col_a, options),
559 (col_b, options),
560 (col_c, options),
561 (col_d, options),
562 ],
563 ],
564 // equivalence classes
565 vec![vec![col_a, col_f]],
566 // constants
567 vec![col_e],
568 // requirement [exp(a) ASC, a+b ASC],
569 vec![(&exp_a, options), (&a_plus_b, options)],
570 // expected: requirement is not satisfied.
571 // TODO: If we know that exp function is 1-to-1 function.
572 // we could have deduced that above requirement is satisfied.
573 false,
574 ),
575 // ------------ TEST CASE 6 ------------
576 (
577 // orderings
578 vec![
579 // [a ASC, d ASC, b ASC]
580 vec![(col_a, options), (col_d, options), (col_b, options)],
581 // [c ASC]
582 vec![(col_c, options)],
583 ],
584 // equivalence classes
585 vec![vec![col_a, col_f]],
586 // constants
587 vec![col_e],
588 // requirement [a ASC, d ASC, floor(a) ASC],
589 vec![(col_a, options), (col_d, options), (&floor_a, options)],
590 // expected: requirement is satisfied.
591 true,
592 ),
593 // ------------ TEST CASE 7 ------------
594 (
595 // orderings
596 vec![
597 // [a ASC, c ASC, b ASC]
598 vec![(col_a, options), (col_c, options), (col_b, options)],
599 // [d ASC]
600 vec![(col_d, options)],
601 ],
602 // equivalence classes
603 vec![vec![col_a, col_f]],
604 // constants
605 vec![col_e],
606 // requirement [a ASC, floor(a) ASC, a + b ASC],
607 vec![(col_a, options), (&floor_a, options), (&a_plus_b, options)],
608 // expected: requirement is not satisfied.
609 false,
610 ),
611 // ------------ TEST CASE 8 ------------
612 (
613 // orderings
614 vec![
615 // [a ASC, b ASC, c ASC]
616 vec![(col_a, options), (col_b, options), (col_c, options)],
617 // [d ASC]
618 vec![(col_d, options)],
619 ],
620 // equivalence classes
621 vec![vec![col_a, col_f]],
622 // constants
623 vec![col_e],
624 // requirement [a ASC, c ASC, floor(a) ASC, a + b ASC],
625 vec![
626 (col_a, options),
627 (col_c, options),
628 (&floor_a, options),
629 (&a_plus_b, options),
630 ],
631 // expected: requirement is not satisfied.
632 false,
633 ),
634 // ------------ TEST CASE 9 ------------
635 (
636 // orderings
637 vec![
638 // [a ASC, b ASC, c ASC, d ASC]
639 vec![
640 (col_a, options),
641 (col_b, options),
642 (col_c, options),
643 (col_d, options),
644 ],
645 ],
646 // equivalence classes
647 vec![vec![col_a, col_f]],
648 // constants
649 vec![col_e],
650 // requirement [a ASC, b ASC, c ASC, floor(a) ASC],
651 vec![
652 (col_a, options),
653 (col_b, options),
654 (col_c, options),
655 (&floor_a, options),
656 ],
657 // expected: requirement is satisfied.
658 true,
659 ),
660 // ------------ TEST CASE 10 ------------
661 (
662 // orderings
663 vec![
664 // [d ASC, b ASC]
665 vec![(col_d, options), (col_b, options)],
666 // [c ASC, a ASC]
667 vec![(col_c, options), (col_a, options)],
668 ],
669 // equivalence classes
670 vec![vec![col_a, col_f]],
671 // constants
672 vec![col_e],
673 // requirement [c ASC, d ASC, a + b ASC],
674 vec![(col_c, options), (col_d, options), (&a_plus_b, options)],
675 // expected: requirement is not satisfied because addition can wrap.
676 false,
677 ),
678 ];
679
680 for (orderings, eq_group, constants, reqs, expected) in test_cases {
681 let err_msg = format!(
682 "error in test orderings: {orderings:?}, eq_group: {eq_group:?}, constants: {constants:?}, reqs: {reqs:?}, expected: {expected:?}"
683 );
684 let mut eq_properties = EquivalenceProperties::new(Arc::clone(&test_schema));
685 let orderings = convert_to_orderings(&orderings);
686 eq_properties.add_orderings(orderings);
687 let classes = eq_group
688 .into_iter()
689 .map(|eq_class| EquivalenceClass::new(eq_class.into_iter().cloned()));
690 let eq_group = EquivalenceGroup::new(classes);
691 eq_properties.add_equivalence_group(eq_group)?;
692
693 let constants = constants.into_iter().map(|expr| {
694 ConstExpr::new(Arc::clone(expr), AcrossPartitions::Uniform(None))
695 });
696 eq_properties.add_constants(constants)?;
697
698 let reqs = convert_to_sort_exprs(&reqs);
699 assert_eq!(eq_properties.ordering_satisfy(reqs)?, expected, "{err_msg}");
700 }
701
702 Ok(())
703 }
704
705 #[test]
706 fn test_ordering_satisfy_different_lengths() -> Result<()> {
707 let test_schema = create_test_schema()?;
708 let col_a = &col("a", &test_schema)?;
709 let col_b = &col("b", &test_schema)?;
710 let col_c = &col("c", &test_schema)?;
711 let col_d = &col("d", &test_schema)?;
712 let col_e = &col("e", &test_schema)?;
713 let col_f = &col("f", &test_schema)?;
714 let options = SortOptions {
715 descending: false,
716 nulls_first: false,
717 };
718 // a=c (e.g they are aliases).
719 let mut eq_properties = EquivalenceProperties::new(test_schema);
720 eq_properties.add_equal_conditions(Arc::clone(col_a), Arc::clone(col_c))?;
721
722 let orderings = vec![
723 vec![(col_a, options)],
724 vec![(col_e, options)],
725 vec![(col_d, options), (col_f, options)],
726 ];
727 let orderings = convert_to_orderings(&orderings);
728
729 // Column [a ASC], [e ASC], [d ASC, f ASC] are all valid orderings for the schema.
730 eq_properties.add_orderings(orderings);
731
732 // First entry in the tuple is required ordering, second entry is the expected flag
733 // that indicates whether this required ordering is satisfied.
734 // ([a ASC], true) indicate a ASC requirement is already satisfied by existing orderings.
735 let test_cases = vec![
736 // [c ASC, a ASC, e ASC], expected represents this requirement is satisfied
737 (
738 vec![(col_c, options), (col_a, options), (col_e, options)],
739 true,
740 ),
741 (vec![(col_c, options), (col_b, options)], false),
742 (vec![(col_c, options), (col_d, options)], true),
743 (
744 vec![(col_d, options), (col_f, options), (col_b, options)],
745 false,
746 ),
747 (vec![(col_d, options), (col_f, options)], true),
748 ];
749
750 for (reqs, expected) in test_cases {
751 let err_msg =
752 format!("error in test reqs: {reqs:?}, expected: {expected:?}",);
753 let reqs = convert_to_sort_exprs(&reqs);
754 assert_eq!(eq_properties.ordering_satisfy(reqs)?, expected, "{err_msg}");
755 }
756
757 Ok(())
758 }
759
760 #[test]
761 fn test_remove_redundant_entries_oeq_class() -> Result<()> {
762 let schema = create_test_schema()?;
763 let col_a = &col("a", &schema)?;
764 let col_b = &col("b", &schema)?;
765 let col_c = &col("c", &schema)?;
766 let col_d = &col("d", &schema)?;
767 let col_e = &col("e", &schema)?;
768
769 let option_asc = SortOptions {
770 descending: false,
771 nulls_first: false,
772 };
773 let option_desc = SortOptions {
774 descending: true,
775 nulls_first: true,
776 };
777
778 // First entry in the tuple is the given orderings for the table
779 // Second entry is the simplest version of the given orderings that is functionally equivalent.
780 let test_cases = vec![
781 // ------- TEST CASE 1 ---------
782 (
783 // ORDERINGS GIVEN
784 vec![
785 // [a ASC, b ASC]
786 vec![(col_a, option_asc), (col_b, option_asc)],
787 ],
788 // EXPECTED orderings that is succinct.
789 vec![
790 // [a ASC, b ASC]
791 vec![(col_a, option_asc), (col_b, option_asc)],
792 ],
793 ),
794 // ------- TEST CASE 2 ---------
795 (
796 // ORDERINGS GIVEN
797 vec![
798 // [a ASC, b ASC]
799 vec![(col_a, option_asc), (col_b, option_asc)],
800 // [a ASC, b ASC, c ASC]
801 vec![
802 (col_a, option_asc),
803 (col_b, option_asc),
804 (col_c, option_asc),
805 ],
806 ],
807 // EXPECTED orderings that is succinct.
808 vec![
809 // [a ASC, b ASC, c ASC]
810 vec![
811 (col_a, option_asc),
812 (col_b, option_asc),
813 (col_c, option_asc),
814 ],
815 ],
816 ),
817 // ------- TEST CASE 3 ---------
818 (
819 // ORDERINGS GIVEN
820 vec![
821 // [a ASC, b DESC]
822 vec![(col_a, option_asc), (col_b, option_desc)],
823 // [a ASC]
824 vec![(col_a, option_asc)],
825 // [a ASC, c ASC]
826 vec![(col_a, option_asc), (col_c, option_asc)],
827 ],
828 // EXPECTED orderings that is succinct.
829 vec![
830 // [a ASC, b DESC]
831 vec![(col_a, option_asc), (col_b, option_desc)],
832 // [a ASC, c ASC]
833 vec![(col_a, option_asc), (col_c, option_asc)],
834 ],
835 ),
836 // ------- TEST CASE 4 ---------
837 (
838 // ORDERINGS GIVEN
839 vec![
840 // [a ASC, b ASC]
841 vec![(col_a, option_asc), (col_b, option_asc)],
842 // [a ASC, b ASC, c ASC]
843 vec![
844 (col_a, option_asc),
845 (col_b, option_asc),
846 (col_c, option_asc),
847 ],
848 // [a ASC]
849 vec![(col_a, option_asc)],
850 ],
851 // EXPECTED orderings that is succinct.
852 vec![
853 // [a ASC, b ASC, c ASC]
854 vec![
855 (col_a, option_asc),
856 (col_b, option_asc),
857 (col_c, option_asc),
858 ],
859 ],
860 ),
861 // ------- TEST CASE 5 ---------
862 // Empty ordering
863 (
864 vec![],
865 // No ordering in the state (empty ordering is ignored).
866 vec![],
867 ),
868 // ------- TEST CASE 6 ---------
869 (
870 // ORDERINGS GIVEN
871 vec![
872 // [a ASC, b ASC]
873 vec![(col_a, option_asc), (col_b, option_asc)],
874 // [b ASC]
875 vec![(col_b, option_asc)],
876 ],
877 // EXPECTED orderings that is succinct.
878 vec![
879 // [a ASC]
880 vec![(col_a, option_asc)],
881 // [b ASC]
882 vec![(col_b, option_asc)],
883 ],
884 ),
885 // ------- TEST CASE 7 ---------
886 // b, a
887 // c, a
888 // d, b, c
889 (
890 // ORDERINGS GIVEN
891 vec![
892 // [b ASC, a ASC]
893 vec![(col_b, option_asc), (col_a, option_asc)],
894 // [c ASC, a ASC]
895 vec![(col_c, option_asc), (col_a, option_asc)],
896 // [d ASC, b ASC, c ASC]
897 vec![
898 (col_d, option_asc),
899 (col_b, option_asc),
900 (col_c, option_asc),
901 ],
902 ],
903 // EXPECTED orderings that is succinct.
904 vec![
905 // [b ASC, a ASC]
906 vec![(col_b, option_asc), (col_a, option_asc)],
907 // [c ASC, a ASC]
908 vec![(col_c, option_asc), (col_a, option_asc)],
909 // [d ASC]
910 vec![(col_d, option_asc)],
911 ],
912 ),
913 // ------- TEST CASE 8 ---------
914 // b, e
915 // c, a
916 // d, b, e, c, a
917 (
918 // ORDERINGS GIVEN
919 vec![
920 // [b ASC, e ASC]
921 vec![(col_b, option_asc), (col_e, option_asc)],
922 // [c ASC, a ASC]
923 vec![(col_c, option_asc), (col_a, option_asc)],
924 // [d ASC, b ASC, e ASC, c ASC, a ASC]
925 vec![
926 (col_d, option_asc),
927 (col_b, option_asc),
928 (col_e, option_asc),
929 (col_c, option_asc),
930 (col_a, option_asc),
931 ],
932 ],
933 // EXPECTED orderings that is succinct.
934 vec![
935 // [b ASC, e ASC]
936 vec![(col_b, option_asc), (col_e, option_asc)],
937 // [c ASC, a ASC]
938 vec![(col_c, option_asc), (col_a, option_asc)],
939 // [d ASC]
940 vec![(col_d, option_asc)],
941 ],
942 ),
943 // ------- TEST CASE 9 ---------
944 // b
945 // a, b, c
946 // d, a, b
947 (
948 // ORDERINGS GIVEN
949 vec![
950 // [b ASC]
951 vec![(col_b, option_asc)],
952 // [a ASC, b ASC, c ASC]
953 vec![
954 (col_a, option_asc),
955 (col_b, option_asc),
956 (col_c, option_asc),
957 ],
958 // [d ASC, a ASC, b ASC]
959 vec![
960 (col_d, option_asc),
961 (col_a, option_asc),
962 (col_b, option_asc),
963 ],
964 ],
965 // EXPECTED orderings that is succinct.
966 vec![
967 // [b ASC]
968 vec![(col_b, option_asc)],
969 // [a ASC, b ASC, c ASC]
970 vec![
971 (col_a, option_asc),
972 (col_b, option_asc),
973 (col_c, option_asc),
974 ],
975 // [d ASC]
976 vec![(col_d, option_asc)],
977 ],
978 ),
979 ];
980 for (orderings, expected) in test_cases {
981 let orderings = convert_to_orderings(&orderings);
982 let expected = convert_to_orderings(&expected);
983 let actual = OrderingEquivalenceClass::from(orderings.clone());
984 let err_msg = format!(
985 "orderings: {orderings:?}, expected: {expected:?}, actual :{actual:?}"
986 );
987 assert_eq!(actual.len(), expected.len(), "{err_msg}");
988 for elem in actual {
989 assert!(expected.contains(&elem), "{}", err_msg);
990 }
991 }
992
993 Ok(())
994 }
995}