This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new b0ec2d61b3 Refactor tests for union sorting properties, add tests for
unions and constants (#12702)
b0ec2d61b3 is described below
commit b0ec2d61b3a1a5710d0964979e61adc73cf827f6
Author: Andrew Lamb <[email protected]>
AuthorDate: Wed Oct 2 10:18:25 2024 -0400
Refactor tests for union sorting properties, add tests for unions and
constants (#12702)
* Refactor tests for union sorting properties
* update doc test
* Undo import reordering
* remove unecessary static lifetimes
---
.../physical-expr/src/equivalence/ordering.rs | 18 +-
.../physical-expr/src/equivalence/properties.rs | 815 ++++++++++++---------
2 files changed, 467 insertions(+), 366 deletions(-)
diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs
b/datafusion/physical-expr/src/equivalence/ordering.rs
index 65423033d5..bb3e9218bc 100644
--- a/datafusion/physical-expr/src/equivalence/ordering.rs
+++ b/datafusion/physical-expr/src/equivalence/ordering.rs
@@ -18,6 +18,7 @@
use std::fmt::Display;
use std::hash::Hash;
use std::sync::Arc;
+use std::vec::IntoIter;
use crate::equivalence::add_offset_to_expr;
use crate::{LexOrdering, PhysicalExpr, PhysicalSortExpr};
@@ -36,7 +37,7 @@ use arrow_schema::SortOptions;
///
/// Here, both `vec![a ASC, b ASC]` and `vec![c DESC, d ASC]` describe the
table
/// ordering. In this case, we say that these orderings are equivalent.
-#[derive(Debug, Clone, Eq, PartialEq, Hash)]
+#[derive(Debug, Clone, Eq, PartialEq, Hash, Default)]
pub struct OrderingEquivalenceClass {
pub orderings: Vec<LexOrdering>,
}
@@ -44,7 +45,7 @@ pub struct OrderingEquivalenceClass {
impl OrderingEquivalenceClass {
/// Creates new empty ordering equivalence class.
pub fn empty() -> Self {
- Self { orderings: vec![] }
+ Default::default()
}
/// Clears (empties) this ordering equivalence class.
@@ -197,6 +198,15 @@ impl OrderingEquivalenceClass {
}
}
+impl IntoIterator for OrderingEquivalenceClass {
+ type Item = LexOrdering;
+ type IntoIter = IntoIter<LexOrdering>;
+
+ fn into_iter(self) -> Self::IntoIter {
+ self.orderings.into_iter()
+ }
+}
+
/// This function constructs a duplicate-free `LexOrdering` by filtering out
/// duplicate entries that have same physical expression inside. For example,
/// `vec![a ASC, a DESC]` collapses to `vec![a ASC]`.
@@ -229,10 +239,10 @@ impl Display for OrderingEquivalenceClass {
write!(f, "[")?;
let mut iter = self.orderings.iter();
if let Some(ordering) = iter.next() {
- write!(f, "{}", PhysicalSortExpr::format_list(ordering))?;
+ write!(f, "[{}]", PhysicalSortExpr::format_list(ordering))?;
}
for ordering in iter {
- write!(f, "{}", PhysicalSortExpr::format_list(ordering))?;
+ write!(f, ", [{}]", PhysicalSortExpr::format_list(ordering))?;
}
write!(f, "]")?;
Ok(())
diff --git a/datafusion/physical-expr/src/equivalence/properties.rs
b/datafusion/physical-expr/src/equivalence/properties.rs
index dc59a1eb83..8137b4f9da 100644
--- a/datafusion/physical-expr/src/equivalence/properties.rs
+++ b/datafusion/physical-expr/src/equivalence/properties.rs
@@ -118,7 +118,7 @@ use itertools::Itertools;
/// PhysicalSortExpr::new_default(col_c).desc(),
/// ]);
///
-/// assert_eq!(eq_properties.to_string(), "order: [a@0 ASC,c@2 DESC], const:
[b@1]")
+/// assert_eq!(eq_properties.to_string(), "order: [[a@0 ASC,c@2 DESC]], const:
[b@1]")
/// ```
#[derive(Debug, Clone)]
pub struct EquivalenceProperties {
@@ -2708,379 +2708,428 @@ mod tests {
))
}
- #[tokio::test]
- async fn test_union_equivalence_properties_multi_children() -> Result<()> {
- let schema = create_test_schema()?;
+ #[test]
+ fn test_union_equivalence_properties_multi_children_1() {
+ let schema = create_test_schema().unwrap();
let schema2 = append_fields(&schema, "1");
let schema3 = append_fields(&schema, "2");
- let test_cases = vec![
- // --------- TEST CASE 1 ----------
- (
- vec![
- // Children 1
- (
- // Orderings
- vec![vec!["a", "b", "c"]],
- Arc::clone(&schema),
- ),
- // Children 2
- (
- // Orderings
- vec![vec!["a1", "b1", "c1"]],
- Arc::clone(&schema2),
- ),
- // Children 3
- (
- // Orderings
- vec![vec!["a2", "b2"]],
- Arc::clone(&schema3),
- ),
- ],
- // Expected
- vec![vec!["a", "b"]],
- ),
- // --------- TEST CASE 2 ----------
- (
- vec![
- // Children 1
- (
- // Orderings
- vec![vec!["a", "b", "c"]],
- Arc::clone(&schema),
- ),
- // Children 2
- (
- // Orderings
- vec![vec!["a1", "b1", "c1"]],
- Arc::clone(&schema2),
- ),
- // Children 3
- (
- // Orderings
- vec![vec!["a2", "b2", "c2"]],
- Arc::clone(&schema3),
- ),
- ],
- // Expected
- vec![vec!["a", "b", "c"]],
- ),
- // --------- TEST CASE 3 ----------
- (
- vec![
- // Children 1
- (
- // Orderings
- vec![vec!["a", "b"]],
- Arc::clone(&schema),
- ),
- // Children 2
- (
- // Orderings
- vec![vec!["a1", "b1", "c1"]],
- Arc::clone(&schema2),
- ),
- // Children 3
- (
- // Orderings
- vec![vec!["a2", "b2", "c2"]],
- Arc::clone(&schema3),
- ),
- ],
- // Expected
- vec![vec!["a", "b"]],
- ),
- // --------- TEST CASE 4 ----------
- (
- vec![
- // Children 1
- (
- // Orderings
- vec![vec!["a", "b"]],
- Arc::clone(&schema),
- ),
- // Children 2
- (
- // Orderings
- vec![vec!["a1", "b1"]],
- Arc::clone(&schema2),
- ),
- // Children 3
- (
- // Orderings
- vec![vec!["b2", "c2"]],
- Arc::clone(&schema3),
- ),
- ],
- // Expected
- vec![],
- ),
- // --------- TEST CASE 5 ----------
- (
- vec![
- // Children 1
- (
- // Orderings
- vec![vec!["a", "b"], vec!["c"]],
- Arc::clone(&schema),
- ),
- // Children 2
- (
- // Orderings
- vec![vec!["a1", "b1"], vec!["c1"]],
- Arc::clone(&schema2),
- ),
- ],
- // Expected
- vec![vec!["a", "b"], vec!["c"]],
- ),
- ];
- for (children, expected) in test_cases {
- let children_eqs = children
- .iter()
- .map(|(orderings, schema)| {
- let orderings = orderings
- .iter()
- .map(|ordering| {
- ordering
- .iter()
- .map(|name| PhysicalSortExpr {
- expr: col(name, schema).unwrap(),
- options: SortOptions::default(),
- })
- .collect::<Vec<_>>()
- })
- .collect::<Vec<_>>();
- EquivalenceProperties::new_with_orderings(
- Arc::clone(schema),
- &orderings,
- )
- })
- .collect::<Vec<_>>();
- let actual = calculate_union(children_eqs, Arc::clone(&schema))?;
+ UnionEquivalenceTest::new(&schema)
+ // Children 1
+ .with_child_sort(vec![vec!["a", "b", "c"]], &schema)
+ // Children 2
+ .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)
+ // Children 3
+ .with_child_sort(vec![vec!["a2", "b2"]], &schema3)
+ .with_expected_sort(vec![vec!["a", "b"]])
+ .run()
+ }
- let expected_ordering = expected
- .into_iter()
- .map(|ordering| {
- ordering
- .into_iter()
- .map(|name| PhysicalSortExpr {
- expr: col(name, &schema).unwrap(),
- options: SortOptions::default(),
- })
- .collect::<Vec<_>>()
- })
- .collect::<Vec<_>>();
- let expected = EquivalenceProperties::new_with_orderings(
- Arc::clone(&schema),
- &expected_ordering,
- );
- assert_eq_properties_same(
- &actual,
- &expected,
- format!("expected: {:?}, actual: {:?}", expected, actual),
- );
- }
- Ok(())
+ #[test]
+ fn test_union_equivalence_properties_multi_children_2() {
+ let schema = create_test_schema().unwrap();
+ let schema2 = append_fields(&schema, "1");
+ let schema3 = append_fields(&schema, "2");
+ UnionEquivalenceTest::new(&schema)
+ // Children 1
+ .with_child_sort(vec![vec!["a", "b", "c"]], &schema)
+ // Children 2
+ .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)
+ // Children 3
+ .with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3)
+ .with_expected_sort(vec![vec!["a", "b", "c"]])
+ .run()
}
- #[tokio::test]
- async fn test_union_equivalence_properties_binary() -> Result<()> {
- let schema = create_test_schema()?;
+ #[test]
+ fn test_union_equivalence_properties_multi_children_3() {
+ let schema = create_test_schema().unwrap();
let schema2 = append_fields(&schema, "1");
- let col_a = &col("a", &schema)?;
- let col_b = &col("b", &schema)?;
- let col_c = &col("c", &schema)?;
- let col_a1 = &col("a1", &schema2)?;
- let col_b1 = &col("b1", &schema2)?;
- let options = SortOptions::default();
- let options_desc = !SortOptions::default();
- let test_cases = [
- //-----------TEST CASE 1----------//
- (
- (
- // First child orderings
- vec![
- // [a ASC]
- (vec![(col_a, options)]),
- ],
- // First child constants
- vec![col_b, col_c],
- Arc::clone(&schema),
- ),
- (
- // Second child orderings
- vec![
- // [b ASC]
- (vec![(col_b, options)]),
- ],
- // Second child constants
- vec![col_a, col_c],
- Arc::clone(&schema),
- ),
- (
- // Union expected orderings
- vec![
- // [a ASC]
- vec![(col_a, options)],
- // [b ASC]
- vec![(col_b, options)],
- ],
- // Union
- vec![col_c],
- ),
- ),
- //-----------TEST CASE 2----------//
+ let schema3 = append_fields(&schema, "2");
+ UnionEquivalenceTest::new(&schema)
+ // Children 1
+ .with_child_sort(vec![vec!["a", "b"]], &schema)
+ // Children 2
+ .with_child_sort(vec![vec!["a1", "b1", "c1"]], &schema2)
+ // Children 3
+ .with_child_sort(vec![vec!["a2", "b2", "c2"]], &schema3)
+ .with_expected_sort(vec![vec!["a", "b"]])
+ .run()
+ }
+
+ #[test]
+ fn test_union_equivalence_properties_multi_children_4() {
+ let schema = create_test_schema().unwrap();
+ let schema2 = append_fields(&schema, "1");
+ let schema3 = append_fields(&schema, "2");
+ UnionEquivalenceTest::new(&schema)
+ // Children 1
+ .with_child_sort(vec![vec!["a", "b"]], &schema)
+ // Children 2
+ .with_child_sort(vec![vec!["a1", "b1"]], &schema2)
+ // Children 3
+ .with_child_sort(vec![vec!["b2", "c2"]], &schema3)
+ .with_expected_sort(vec![])
+ .run()
+ }
+
+ #[test]
+ fn test_union_equivalence_properties_multi_children_5() {
+ let schema = create_test_schema().unwrap();
+ let schema2 = append_fields(&schema, "1");
+ UnionEquivalenceTest::new(&schema)
+ // Children 1
+ .with_child_sort(vec![vec!["a", "b"], vec!["c"]], &schema)
+ // Children 2
+ .with_child_sort(vec![vec!["a1", "b1"], vec!["c1"]], &schema2)
+ .with_expected_sort(vec![vec!["a", "b"], vec!["c"]])
+ .run()
+ }
+
+ #[test]
+ fn test_union_equivalence_properties_constants_1() {
+ let schema = create_test_schema().unwrap();
+ UnionEquivalenceTest::new(&schema)
+ .with_child_sort_and_const_exprs(
+ // First child: [a ASC], const [b, c]
+ vec![vec!["a"]],
+ vec!["b", "c"],
+ &schema,
+ )
+ .with_child_sort_and_const_exprs(
+ // Second child: [b ASC], const [a, c]
+ vec![vec!["b"]],
+ vec!["a", "c"],
+ &schema,
+ )
+ .with_expected_sort_and_const_exprs(
+ // Union expected orderings: [[a ASC], [b ASC]], const [c]
+ vec![vec!["a"], vec!["b"]],
+ vec!["c"],
+ )
+ .run()
+ }
+
+ #[test]
+ fn test_union_equivalence_properties_constants_2() {
+ let schema = create_test_schema().unwrap();
+ UnionEquivalenceTest::new(&schema)
// Meet ordering between [a ASC], [a ASC, b ASC] should be [a ASC]
- (
- (
- // First child orderings
- vec![
- // [a ASC]
- vec![(col_a, options)],
- ],
- // No constant
- vec![],
- Arc::clone(&schema),
- ),
- (
- // Second child orderings
- vec![
- // [a ASC, b ASC]
- vec![(col_a, options), (col_b, options)],
- ],
- // No constant
- vec![],
- Arc::clone(&schema),
- ),
- (
- // Union orderings
- vec![
- // [a ASC]
- vec![(col_a, options)],
- ],
- // No constant
- vec![],
- ),
- ),
- //-----------TEST CASE 3----------//
+ .with_child_sort_and_const_exprs(
+ // First child: [a ASC], const []
+ vec![vec!["a"]],
+ vec![],
+ &schema,
+ )
+ .with_child_sort_and_const_exprs(
+ // Second child: [a ASC, b ASC], const []
+ vec![vec!["a", "b"]],
+ vec![],
+ &schema,
+ )
+ .with_expected_sort_and_const_exprs(
+ // Union orderings: [a ASC], const []
+ vec![vec!["a"]],
+ vec![],
+ )
+ .run()
+ }
+
+ #[test]
+ fn test_union_equivalence_properties_constants_3() {
+ let schema = create_test_schema().unwrap();
+ UnionEquivalenceTest::new(&schema)
// Meet ordering between [a ASC], [a DESC] should be []
- (
- (
- // First child orderings
- vec![
- // [a ASC]
- vec![(col_a, options)],
- ],
- // No constant
- vec![],
- Arc::clone(&schema),
- ),
- (
- // Second child orderings
- vec![
- // [a DESC]
- vec![(col_a, options_desc)],
- ],
- // No constant
- vec![],
- Arc::clone(&schema),
- ),
- (
- // Union doesn't have any ordering
- vec![],
- // No constant
- vec![],
- ),
- ),
- //-----------TEST CASE 4----------//
- // Meet ordering between [a ASC], [a1 ASC, b1 ASC] should be [a
ASC]
- // Where a, and a1 ath the same index for their corresponding
schemas.
- (
- (
- // First child orderings
- vec![
- // [a ASC]
- vec![(col_a, options)],
- ],
- // No constant
- vec![],
- Arc::clone(&schema),
- ),
- (
- // Second child orderings
- vec![
- // [a1 ASC, b1 ASC]
- vec![(col_a1, options), (col_b1, options)],
- ],
- // No constant
- vec![],
- Arc::clone(&schema2),
- ),
- (
- // Union orderings
- vec![
- // [a ASC]
- vec![(col_a, options)],
- ],
- // No constant
- vec![],
- ),
- ),
- ];
+ .with_child_sort_and_const_exprs(
+ // First child: [a ASC], const []
+ vec![vec!["a"]],
+ vec![],
+ &schema,
+ )
+ .with_child_sort_and_const_exprs(
+ // Second child orderings: [a DESC], const []
+ vec![vec!["a DESC"]],
+ vec![],
+ &schema,
+ )
+ .with_expected_sort_and_const_exprs(
+ // Union doesn't have any ordering or constant
+ vec![],
+ vec![],
+ )
+ .run()
+ }
- for (
- test_idx,
- (
- (first_child_orderings, first_child_constants, first_schema),
- (second_child_orderings, second_child_constants,
second_schema),
- (union_orderings, union_constants),
- ),
- ) in test_cases.iter().enumerate()
- {
- let first_orderings = first_child_orderings
- .iter()
- .map(|ordering| convert_to_sort_exprs(ordering))
- .collect::<Vec<_>>();
- let first_constants = first_child_constants
- .iter()
- .map(|expr| ConstExpr::new(Arc::clone(expr)))
- .collect::<Vec<_>>();
- let mut lhs = EquivalenceProperties::new(Arc::clone(first_schema));
- lhs = lhs.with_constants(first_constants);
- lhs.add_new_orderings(first_orderings);
+ #[test]
+ fn test_union_equivalence_properties_constants_4() {
+ let schema = create_test_schema().unwrap();
+ let schema2 = append_fields(&schema, "1");
+ UnionEquivalenceTest::new(&schema)
+ .with_child_sort_and_const_exprs(
+ // First child orderings: [a ASC], const []
+ vec![vec!["a"]],
+ vec![],
+ &schema,
+ )
+ .with_child_sort_and_const_exprs(
+ // Second child orderings: [a1 ASC, b1 ASC], const []
+ vec![vec!["a1", "b1"]],
+ vec![],
+ &schema2,
+ )
+ .with_expected_sort_and_const_exprs(
+ // Union orderings:
+ // should be [a ASC]
+ //
+ // Where a, and a1 ath the same index for their corresponding
+ // schemas.
+ vec![vec!["a"]],
+ vec![],
+ )
+ .run()
+ }
- let second_orderings = second_child_orderings
- .iter()
- .map(|ordering| convert_to_sort_exprs(ordering))
- .collect::<Vec<_>>();
- let second_constants = second_child_constants
- .iter()
- .map(|expr| ConstExpr::new(Arc::clone(expr)))
- .collect::<Vec<_>>();
- let mut rhs =
EquivalenceProperties::new(Arc::clone(second_schema));
- rhs = rhs.with_constants(second_constants);
- rhs.add_new_orderings(second_orderings);
+ #[test]
+ #[ignore]
+ // ignored due to https://github.com/apache/datafusion/issues/12446
+ fn test_union_equivalence_properties_constants() {
+ let schema = create_test_schema().unwrap();
+ UnionEquivalenceTest::new(&schema)
+ .with_child_sort_and_const_exprs(
+ // First child orderings: [a ASC, c ASC], const [b]
+ vec![vec!["a", "c"]],
+ vec!["b"],
+ &schema,
+ )
+ .with_child_sort_and_const_exprs(
+ // Second child orderings: [b ASC, c ASC], const [a]
+ vec![vec!["b", "c"]],
+ vec!["a"],
+ &schema,
+ )
+ .with_expected_sort_and_const_exprs(
+ // Union orderings: [
+ // [a ASC, b ASC, c ASC],
+ // [b ASC, a ASC, c ASC]
+ // ], const []
+ vec![vec!["a", "b", "c"], vec!["b", "a", "c"]],
+ vec![],
+ )
+ .run()
+ }
+
+ #[test]
+ #[ignore]
+ // ignored due to https://github.com/apache/datafusion/issues/12446
+ fn test_union_equivalence_properties_constants_desc() {
+ let schema = create_test_schema().unwrap();
+ UnionEquivalenceTest::new(&schema)
+ .with_child_sort_and_const_exprs(
+ // NB `b DESC` in the second child
+ // First child orderings: [a ASC, c ASC], const [b]
+ vec![vec!["a", "c"]],
+ vec!["b"],
+ &schema,
+ )
+ .with_child_sort_and_const_exprs(
+ // Second child orderings: [b ASC, c ASC], const [a]
+ vec![vec!["b DESC", "c"]],
+ vec!["a"],
+ &schema,
+ )
+ .with_expected_sort_and_const_exprs(
+ // Union orderings: [
+ // [a ASC, b ASC, c ASC],
+ // [b ASC, a ASC, c ASC]
+ // ], const []
+ vec![vec!["a", "b DESC", "c"], vec!["b DESC", "a", "c"]],
+ vec![],
+ )
+ .run()
+ }
+
+ #[test]
+ #[ignore]
+ // ignored due to https://github.com/apache/datafusion/issues/12446
+ fn test_union_equivalence_properties_constants_middle() {
+ let schema = create_test_schema().unwrap();
+ UnionEquivalenceTest::new(&schema)
+ .with_child_sort_and_const_exprs(
+ // First child: [a ASC, b ASC, d ASC], const [c]
+ vec![vec!["a", "b", "d"]],
+ vec!["c"],
+ &schema,
+ )
+ .with_child_sort_and_const_exprs(
+ // Second child: [a ASC, c ASC, d ASC], const [b]
+ vec![vec!["a", "c", "d"]],
+ vec!["b"],
+ &schema,
+ )
+ .with_expected_sort_and_const_exprs(
+ // Union orderings:
+ // [a, b, d] (c constant)
+ // [a, c, d] (b constant)
+ vec![vec!["a", "c", "b", "d"], vec!["a", "b", "c", "d"]],
+ vec![],
+ )
+ .run()
+ }
+
+ #[test]
+ #[ignore]
+ // ignored due to https://github.com/apache/datafusion/issues/12446
+ fn test_union_equivalence_properties_constants_middle_desc() {
+ let schema = create_test_schema().unwrap();
+ UnionEquivalenceTest::new(&schema)
+ .with_child_sort_and_const_exprs(
+ // NB `b DESC` in the first child
+ //
+ // First child: [a ASC, b DESC, d ASC], const [c]
+ vec![vec!["a", "b DESC", "d"]],
+ vec!["c"],
+ &schema,
+ )
+ .with_child_sort_and_const_exprs(
+ // Second child: [a ASC, c ASC, d ASC], const [b]
+ vec![vec!["a", "c", "d"]],
+ vec!["b"],
+ &schema,
+ )
+ .with_expected_sort_and_const_exprs(
+ // Union orderings:
+ // [a, b, d] (c constant)
+ // [a, c, d] (b constant)
+ vec![vec!["a", "c", "b DESC", "d"], vec!["a", "b DESC", "c",
"d"]],
+ vec![],
+ )
+ .run()
+ }
+
+ // TODO tests with multiple constants
- let union_expected_orderings = union_orderings
+ #[derive(Debug)]
+ struct UnionEquivalenceTest {
+ /// The schema of the output of the Union
+ output_schema: SchemaRef,
+ /// The equivalence properties of each child to the union
+ child_properties: Vec<EquivalenceProperties>,
+ /// The expected output properties of the union. Must be set before
+ /// running `build`
+ expected_properties: Option<EquivalenceProperties>,
+ }
+
+ impl UnionEquivalenceTest {
+ fn new(output_schema: &SchemaRef) -> Self {
+ Self {
+ output_schema: Arc::clone(output_schema),
+ child_properties: vec![],
+ expected_properties: None,
+ }
+ }
+
+ /// Add a union input with the specified orderings
+ ///
+ /// See [`Self::make_props`] for the format of the strings in
`orderings`
+ fn with_child_sort(
+ mut self,
+ orderings: Vec<Vec<&str>>,
+ schema: &SchemaRef,
+ ) -> Self {
+ let properties = self.make_props(orderings, vec![], schema);
+ self.child_properties.push(properties);
+ self
+ }
+
+ /// Add a union input with the specified orderings and constant
+ /// equivalences
+ ///
+ /// See [`Self::make_props`] for the format of the strings in
+ /// `orderings` and `constants`
+ fn with_child_sort_and_const_exprs(
+ mut self,
+ orderings: Vec<Vec<&str>>,
+ constants: Vec<&str>,
+ schema: &SchemaRef,
+ ) -> Self {
+ let properties = self.make_props(orderings, constants, schema);
+ self.child_properties.push(properties);
+ self
+ }
+
+ /// Set the expected output sort order for the union of the children
+ ///
+ /// See [`Self::make_props`] for the format of the strings in
`orderings`
+ fn with_expected_sort(mut self, orderings: Vec<Vec<&str>>) -> Self {
+ let properties = self.make_props(orderings, vec![],
&self.output_schema);
+ self.expected_properties = Some(properties);
+ self
+ }
+
+ /// Set the expected output sort order and constant expressions for the
+ /// union of the children
+ ///
+ /// See [`Self::make_props`] for the format of the strings in
+ /// `orderings` and `constants`.
+ fn with_expected_sort_and_const_exprs(
+ mut self,
+ orderings: Vec<Vec<&str>>,
+ constants: Vec<&str>,
+ ) -> Self {
+ let properties = self.make_props(orderings, constants,
&self.output_schema);
+ self.expected_properties = Some(properties);
+ self
+ }
+
+ /// compute the union's output equivalence properties from the child
+ /// properties, and compare them to the expected properties
+ fn run(self) {
+ let Self {
+ output_schema,
+ child_properties,
+ expected_properties,
+ } = self;
+ let expected_properties =
+ expected_properties.expect("expected_properties not set");
+ let actual_properties =
+ calculate_union(child_properties, Arc::clone(&output_schema))
+ .expect("failed to calculate union equivalence
properties");
+ assert_eq_properties_same(
+ &actual_properties,
+ &expected_properties,
+ format!(
+ "expected: {expected_properties:?}\nactual:
{actual_properties:?}"
+ ),
+ );
+ }
+
+ /// Make equivalence properties for the specified columns named in
orderings and constants
+ ///
+ /// orderings: strings formatted like `"a"` or `"a DESC"`. See
[`parse_sort_expr`]
+ /// constants: strings formatted like `"a"`.
+ fn make_props(
+ &self,
+ orderings: Vec<Vec<&str>>,
+ constants: Vec<&str>,
+ schema: &SchemaRef,
+ ) -> EquivalenceProperties {
+ let orderings = orderings
.iter()
- .map(|ordering| convert_to_sort_exprs(ordering))
+ .map(|ordering| {
+ ordering
+ .iter()
+ .map(|name| parse_sort_expr(name, schema))
+ .collect::<Vec<_>>()
+ })
.collect::<Vec<_>>();
- let union_constants = union_constants
+
+ let constants = constants
.iter()
- .map(|expr| ConstExpr::new(Arc::clone(expr)))
+ .map(|col_name| ConstExpr::new(col(col_name, schema).unwrap()))
.collect::<Vec<_>>();
- let mut union_expected_eq =
EquivalenceProperties::new(Arc::clone(&schema));
- union_expected_eq =
union_expected_eq.with_constants(union_constants);
- union_expected_eq.add_new_orderings(union_expected_orderings);
- let actual_union_eq = calculate_union_binary(lhs, rhs)?;
- let err_msg = format!(
- "Error in test id: {:?}, test case: {:?}",
- test_idx, test_cases[test_idx]
- );
- assert_eq_properties_same(&actual_union_eq, &union_expected_eq,
err_msg);
+ EquivalenceProperties::new_with_orderings(Arc::clone(schema),
&orderings)
+ .with_constants(constants)
}
- Ok(())
}
fn assert_eq_properties_same(
@@ -3091,21 +3140,63 @@ mod tests {
// Check whether constants are same
let lhs_constants = lhs.constants();
let rhs_constants = rhs.constants();
- assert_eq!(lhs_constants.len(), rhs_constants.len(), "{}", err_msg);
for rhs_constant in rhs_constants {
assert!(
const_exprs_contains(lhs_constants, rhs_constant.expr()),
- "{}",
- err_msg
+ "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
);
}
+ assert_eq!(
+ lhs_constants.len(),
+ rhs_constants.len(),
+ "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
+ );
// Check whether orderings are same.
let lhs_orderings = lhs.oeq_class();
let rhs_orderings = &rhs.oeq_class.orderings;
- assert_eq!(lhs_orderings.len(), rhs_orderings.len(), "{}", err_msg);
for rhs_ordering in rhs_orderings {
- assert!(lhs_orderings.contains(rhs_ordering), "{}", err_msg);
+ assert!(
+ lhs_orderings.contains(rhs_ordering),
+ "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
+ );
}
+ assert_eq!(
+ lhs_orderings.len(),
+ rhs_orderings.len(),
+ "{err_msg}\nlhs: {lhs}\nrhs: {rhs}"
+ );
+ }
+
+ /// Converts a string to a physical sort expression
+ ///
+ /// # Example
+ /// * `"a"` -> (`"a"`, `SortOptions::default()`)
+ /// * `"a ASC"` -> (`"a"`, `SortOptions { descending: false, nulls_first:
false }`)
+ fn parse_sort_expr(name: &str, schema: &SchemaRef) -> PhysicalSortExpr {
+ let mut parts = name.split_whitespace();
+ let name = parts.next().expect("empty sort expression");
+ let mut sort_expr = PhysicalSortExpr::new(
+ col(name, schema).expect("invalid column name"),
+ SortOptions::default(),
+ );
+
+ if let Some(options) = parts.next() {
+ sort_expr = match options {
+ "ASC" => sort_expr.asc(),
+ "DESC" => sort_expr.desc(),
+ _ => panic!(
+ "unknown sort options. Expected 'ASC' or 'DESC', got {}",
+ options
+ ),
+ }
+ }
+
+ assert!(
+ parts.next().is_none(),
+ "unexpected tokens in column name. Expected 'name' / 'name ASC' /
'name DESC' but got '{name}'"
+ );
+
+ sort_expr
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]