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