diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs b/datafusion/physical-expr/src/equivalence/ordering.rs index 24e2fc7dbaf5..22696f1fa527 100644 --- a/datafusion/physical-expr/src/equivalence/ordering.rs +++ b/datafusion/physical-expr/src/equivalence/ordering.rs @@ -39,7 +39,7 @@ use arrow_schema::SortOptions; /// ordering. In this case, we say that these orderings are equivalent. #[derive(Debug, Clone, Eq, PartialEq, Hash, Default)] pub struct OrderingEquivalenceClass { - pub orderings: Vec, + orderings: Vec, } impl OrderingEquivalenceClass { @@ -53,13 +53,20 @@ impl OrderingEquivalenceClass { self.orderings.clear(); } - /// Creates new ordering equivalence class from the given orderings. + /// Creates new ordering equivalence class from the given orderings + /// + /// Any redundant entries are removed, as described on [`Self::remove_redundant_entries`]. pub fn new(orderings: Vec) -> Self { let mut result = Self { orderings }; result.remove_redundant_entries(); result } + /// Converts this OrderingEquivalenceClass to a vector of orderings. + pub fn into_inner(self) -> Vec { + self.orderings + } + /// Checks whether `ordering` is a member of this equivalence class. pub fn contains(&self, ordering: &LexOrdering) -> bool { self.orderings.contains(ordering) @@ -67,10 +74,12 @@ impl OrderingEquivalenceClass { /// Adds `ordering` to this equivalence class. #[allow(dead_code)] + #[deprecated( + since = "45.0.0", + note = "use OrderingEquivalenceClass::add_new_ordering instead" + )] fn push(&mut self, ordering: LexOrdering) { - self.orderings.push(ordering); - // Make sure that there are no redundant orderings: - self.remove_redundant_entries(); + self.add_new_ordering(ordering) } /// Checks whether this ordering equivalence class is empty. @@ -79,6 +88,9 @@ impl OrderingEquivalenceClass { } /// Returns an iterator over the equivalent orderings in this class. + /// + /// Note this class also implements [`IntoIterator`] to return an iterator + /// over owned [`LexOrdering`]s. pub fn iter(&self) -> impl Iterator { self.orderings.iter() } @@ -95,7 +107,7 @@ impl OrderingEquivalenceClass { self.remove_redundant_entries(); } - /// Adds new orderings into this ordering equivalence class. + /// Adds new orderings into this ordering equivalence class pub fn add_new_orderings( &mut self, orderings: impl IntoIterator, @@ -110,9 +122,10 @@ impl OrderingEquivalenceClass { self.add_new_orderings([ordering]); } - /// Removes redundant orderings from this equivalence class. For instance, - /// if we already have the ordering `[a ASC, b ASC, c DESC]`, then there is - /// no need to keep ordering `[a ASC, b ASC]` in the state. + /// Removes redundant orderings from this equivalence class. + /// + /// For instance, if we already have the ordering `[a ASC, b ASC, c DESC]`, + /// then there is no need to keep ordering `[a ASC, b ASC]` in the state. fn remove_redundant_entries(&mut self) { let mut work = true; while work { @@ -198,6 +211,7 @@ impl OrderingEquivalenceClass { } } +/// Convert the `OrderingEquivalenceClass` into an iterator of LexOrderings impl IntoIterator for OrderingEquivalenceClass { type Item = LexOrdering; type IntoIter = IntoIter; @@ -293,13 +307,15 @@ mod tests { // finer ordering satisfies, crude ordering should return true let mut eq_properties_finer = EquivalenceProperties::new(Arc::clone(&input_schema)); - eq_properties_finer.oeq_class.push(finer.clone()); + eq_properties_finer + .oeq_class + .add_new_ordering(finer.clone()); assert!(eq_properties_finer.ordering_satisfy(crude.as_ref())); // Crude ordering doesn't satisfy finer ordering. should return false let mut eq_properties_crude = EquivalenceProperties::new(Arc::clone(&input_schema)); - eq_properties_crude.oeq_class.push(crude); + eq_properties_crude.oeq_class.add_new_ordering(crude); assert!(!eq_properties_crude.ordering_satisfy(finer.as_ref())); Ok(()) } diff --git a/datafusion/physical-expr/src/equivalence/projection.rs b/datafusion/physical-expr/src/equivalence/projection.rs index 25a05a2a5918..681484fd6bff 100644 --- a/datafusion/physical-expr/src/equivalence/projection.rs +++ b/datafusion/physical-expr/src/equivalence/projection.rs @@ -636,7 +636,7 @@ mod tests { let err_msg = format!( "test_idx: {:?}, actual: {:?}, expected: {:?}, projection_mapping: {:?}", - idx, orderings.orderings, expected, projection_mapping + idx, orderings, expected, projection_mapping ); assert_eq!(orderings.len(), expected.len(), "{}", err_msg); @@ -825,7 +825,7 @@ mod tests { let err_msg = format!( "test idx: {:?}, actual: {:?}, expected: {:?}, projection_mapping: {:?}", - idx, orderings.orderings, expected, projection_mapping + idx, orderings, expected, projection_mapping ); assert_eq!(orderings.len(), expected.len(), "{}", err_msg); @@ -971,7 +971,7 @@ mod tests { let err_msg = format!( "actual: {:?}, expected: {:?}, projection_mapping: {:?}", - orderings.orderings, expected, projection_mapping + orderings, expected, projection_mapping ); assert_eq!(orderings.len(), expected.len(), "{}", err_msg); diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs index c3d458103285..5c0c1e18b2df 100755 --- a/datafusion/physical-expr/src/equivalence/properties.rs +++ b/datafusion/physical-expr/src/equivalence/properties.rs @@ -431,8 +431,8 @@ impl EquivalenceProperties { let mut new_orderings = vec![filtered_exprs.clone()]; // Preserve valid suffixes from existing orderings - let orderings = mem::take(&mut self.oeq_class.orderings); - for existing in orderings { + let oeq_class = mem::take(&mut self.oeq_class); + for existing in oeq_class { if self.is_prefix_of(&filtered_exprs, &existing) { let mut extended = filtered_exprs.clone(); extended.extend(existing.into_iter().skip(filtered_exprs.len())); @@ -711,8 +711,8 @@ impl EquivalenceProperties { /// Since it would cause bug in dependency constructions, we should substitute the input order in order to get correct /// dependency map, happen in issue 8838: pub fn substitute_oeq_class(&mut self, mapping: &ProjectionMapping) -> Result<()> { - let orderings = &self.oeq_class.orderings; - let new_order = orderings + let new_order = self + .oeq_class .iter() .map(|order| self.substitute_ordering_component(mapping, order)) .collect::>>()?; @@ -1220,7 +1220,7 @@ impl EquivalenceProperties { // Rewrite orderings according to new schema: let mut new_orderings = vec![]; - for ordering in self.oeq_class.orderings { + for ordering in self.oeq_class { let new_ordering = ordering .inner .into_iter() @@ -2009,16 +2009,8 @@ fn calculate_union_binary( // Next, calculate valid orderings for the union by searching for prefixes // in both sides. let mut orderings = UnionEquivalentOrderingBuilder::new(); - orderings.add_satisfied_orderings( - lhs.normalized_oeq_class().orderings, - lhs.constants(), - &rhs, - ); - orderings.add_satisfied_orderings( - rhs.normalized_oeq_class().orderings, - rhs.constants(), - &lhs, - ); + orderings.add_satisfied_orderings(lhs.normalized_oeq_class(), lhs.constants(), &rhs); + orderings.add_satisfied_orderings(rhs.normalized_oeq_class(), rhs.constants(), &lhs); let orderings = orderings.build(); let mut eq_properties = @@ -2157,7 +2149,7 @@ impl UnionEquivalentOrderingBuilder { // for each equivalent ordering in properties, try and augment // `ordering` it with the constants to match - for existing_ordering in &properties.oeq_class.orderings { + for existing_ordering in properties.oeq_class.iter() { if let Some(augmented_ordering) = self.augment_ordering( ordering, constants, @@ -2438,17 +2430,12 @@ mod tests { Some(JoinSide::Left), &[], ); - let orderings = &join_eq.oeq_class.orderings; - let err_msg = format!("expected: {:?}, actual:{:?}", expected, orderings); - assert_eq!( - join_eq.oeq_class.orderings.len(), - expected.len(), - "{}", - err_msg - ); - for ordering in orderings { + let err_msg = + format!("expected: {:?}, actual:{:?}", expected, &join_eq.oeq_class); + assert_eq!(join_eq.oeq_class.len(), expected.len(), "{}", err_msg); + for ordering in join_eq.oeq_class { assert!( - expected.contains(ordering), + expected.contains(&ordering), "{}, ordering: {:?}", err_msg, ordering @@ -3767,8 +3754,8 @@ mod tests { // Check whether orderings are same. let lhs_orderings = lhs.oeq_class(); - let rhs_orderings = &rhs.oeq_class.orderings; - for rhs_ordering in rhs_orderings { + let rhs_orderings = rhs.oeq_class(); + for rhs_ordering in rhs_orderings.iter() { assert!( lhs_orderings.contains(rhs_ordering), "{err_msg}\nlhs: {lhs}\nrhs: {rhs}" @@ -3844,7 +3831,7 @@ mod tests { // Add equality condition c = concat(a, b) eq_properties.add_equal_conditions(&col_c, &a_concat_b)?; - let orderings = eq_properties.oeq_class().orderings.clone(); + let orderings = eq_properties.oeq_class(); let expected_ordering1 = LexOrdering::from(vec![ @@ -3895,7 +3882,7 @@ mod tests { // Add equality condition c = a * b eq_properties.add_equal_conditions(&col_c, &a_times_b)?; - let orderings = eq_properties.oeq_class().orderings.clone(); + let orderings = eq_properties.oeq_class(); // The ordering should remain unchanged since multiplication is not lex-monotonic assert_eq!(orderings.len(), 1); @@ -3935,7 +3922,7 @@ mod tests { // Add equality condition c = concat(a, b) eq_properties.add_equal_conditions(&col_c, &a_concat_b)?; - let orderings = eq_properties.oeq_class().orderings.clone(); + let orderings = eq_properties.oeq_class(); let expected_ordering1 = LexOrdering::from(vec![PhysicalSortExpr::new_default( Arc::clone(&a_concat_b), @@ -3979,8 +3966,9 @@ mod tests { // Should only contain b since a is constant assert_eq!(result.oeq_class().len(), 1); - assert_eq!(result.oeq_class().orderings[0].len(), 1); - assert!(result.oeq_class().orderings[0][0].expr.eq(&col_b)); + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering.len(), 1); + assert!(ordering[0].expr.eq(&col_b)); Ok(()) } @@ -4026,13 +4014,14 @@ mod tests { // Should only contain [a ASC, b DESC, c ASC] assert_eq!(result.oeq_class().len(), 1); - assert_eq!(result.oeq_class().orderings[0].len(), 3); - assert!(result.oeq_class().orderings[0][0].expr.eq(&col_a)); - assert!(result.oeq_class().orderings[0][0].options.eq(&asc)); - assert!(result.oeq_class().orderings[0][1].expr.eq(&col_b)); - assert!(result.oeq_class().orderings[0][1].options.eq(&desc)); - assert!(result.oeq_class().orderings[0][2].expr.eq(&col_c)); - assert!(result.oeq_class().orderings[0][2].options.eq(&asc)); + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering.len(), 3); + assert!(ordering[0].expr.eq(&col_a)); + assert!(ordering[0].options.eq(&asc)); + assert!(ordering[1].expr.eq(&col_b)); + assert!(ordering[1].options.eq(&desc)); + assert!(ordering[2].expr.eq(&col_c)); + assert!(ordering[2].options.eq(&asc)); Ok(()) } @@ -4075,11 +4064,12 @@ mod tests { assert_eq!(result.oeq_class().len(), 1); // Verify orderings - assert_eq!(result.oeq_class().orderings[0].len(), 2); - assert!(result.oeq_class().orderings[0][0].expr.eq(&col_b)); - assert!(result.oeq_class().orderings[0][0].options.eq(&asc)); - assert!(result.oeq_class().orderings[0][1].expr.eq(&col_c)); - assert!(result.oeq_class().orderings[0][1].options.eq(&asc)); + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering.len(), 2); + assert!(ordering[0].expr.eq(&col_b)); + assert!(ordering[0].options.eq(&asc)); + assert!(ordering[1].expr.eq(&col_c)); + assert!(ordering[1].options.eq(&asc)); Ok(()) } @@ -4120,7 +4110,8 @@ mod tests { // Should only contain the new ordering since options don't match assert_eq!(result.oeq_class().len(), 1); - assert_eq!(result.oeq_class().orderings[0], new_order); + let ordering = result.oeq_class().iter().next().unwrap(); + assert_eq!(ordering, &new_order); Ok(()) } @@ -4178,7 +4169,7 @@ mod tests { // Should preserve the original [d ASC, a ASC] ordering assert_eq!(result.oeq_class().len(), 1); - let ordering = &result.oeq_class().orderings[0]; + let ordering = result.oeq_class().iter().next().unwrap(); assert_eq!(ordering.len(), 2); // First expression should be either b or d (they're equivalent) diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index 521008ce9b02..ee8f9e1cbfea 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -261,7 +261,7 @@ impl MemoryExec { sort_information = base_eqp .project(&projection_mapping, self.schema()) .oeq_class - .orderings; + .into_inner(); } self.sort_information = sort_information; diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index 6e768a3d87bc..cedfcee6fabd 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -843,9 +843,9 @@ mod tests { ) { // Check whether orderings are same. let lhs_orderings = lhs.oeq_class(); - let rhs_orderings = &rhs.oeq_class.orderings; + let rhs_orderings = &rhs.oeq_class; assert_eq!(lhs_orderings.len(), rhs_orderings.len(), "{}", err_msg); - for rhs_ordering in rhs_orderings { + for rhs_ordering in rhs_orderings.iter() { assert!(lhs_orderings.contains(rhs_ordering), "{}", err_msg); } }