chore: Avoid duplicated collect()

apache · Dec 6, 2024 · e4be775 · e4be775
1 parent bf008c7
commit e4be775
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 13 deletions.
diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs
@@ -211,6 +211,14 @@ impl IntoIterator for EquivalenceClass {
     }
 }
 
+impl FromIterator<Arc<dyn PhysicalExpr>> for EquivalenceClass {
+    fn from_iter<T: IntoIterator<Item = Arc<dyn PhysicalExpr>>>(iter: T) -> Self {
+        Self {
+            exprs: iter.into_iter().collect(),
+        }
+    }
+}
+
 impl EquivalenceClass {
     /// Create a new empty equivalence class
     pub fn new_empty() -> Self {
@@ -279,13 +287,12 @@ impl EquivalenceClass {
     /// Return a new equivalence class that have the specified offset added to
     /// each expression (used when schemas are appended such as in joins)
     pub fn with_offset(&self, offset: usize) -> Self {
-        let new_exprs = self
+        let new_exprs_iter = self
             .exprs
             .iter()
             .cloned()
-            .map(|e| add_offset_to_expr(e, offset))
-            .collect();
-        Self::new(new_exprs)
+            .map(|e| add_offset_to_expr(e, offset));
+        Self::from_iter(new_exprs_iter)
     }
 }
 
@@ -689,12 +696,12 @@ mod tests {
         for (entries, expected) in test_cases {
             let entries = entries
                 .into_iter()
-                .map(|entry| entry.into_iter().map(lit).collect::<Vec<_>>())
+                .map(|entry| entry.into_iter().map(lit)).collect::<Vec<_>>();
                 .map(EquivalenceClass::new)
                 .collect::<Vec<_>>();
             let expected = expected
                 .into_iter()
-                .map(|entry| entry.into_iter().map(lit).collect::<Vec<_>>())
+                .map(|entry| entry.into_iter().map(lit)).collect::<Vec<_>>();
                 .map(EquivalenceClass::new)
                 .collect::<Vec<_>>();
             let mut eq_groups = EquivalenceGroup::new(entries.clone());

diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs b/datafusion/physical-expr/src/equivalence/ordering.rs
@@ -264,6 +264,7 @@ mod tests {
     use crate::utils::tests::TestScalarUDF;
     use crate::{ConstExpr, PhysicalExpr, PhysicalSortExpr};
 
+    use arrow::compute::kernels::cmp::eq;
     use arrow::datatypes::{DataType, Field, Schema};
     use arrow_schema::SortOptions;
     use datafusion_common::{DFSchema, Result};
@@ -575,10 +576,7 @@ mod tests {
             eq_properties.add_new_orderings(orderings);
             let eq_group = eq_group
                 .into_iter()
-                .map(|eq_class| {
-                    let eq_classes = eq_class.into_iter().cloned().collect::<Vec<_>>();
-                    EquivalenceClass::new(eq_classes)
-                })
+                .map(|eq_class| EquivalenceClass::new(eq_class.into_iter()))
                 .collect::<Vec<_>>();
             let eq_group = EquivalenceGroup::new(eq_group);
             eq_properties.add_equivalence_group(eq_group);

diff --git a/datafusion/physical-expr/src/equivalence/properties.rs b/datafusion/physical-expr/src/equivalence/properties.rs
@@ -1124,9 +1124,8 @@ impl EquivalenceProperties {
         for eq_class in self.eq_group.classes {
             let new_eq_exprs = eq_class
                 .into_iter()
-                .map(|expr| with_new_schema(expr, &schema))
-                .collect::<Result<_>>()?;
-            eq_classes.push(EquivalenceClass::new(new_eq_exprs));
+                .map(|expr| with_new_schema(expr, &schema).unwrap());
+            eq_classes.push(EquivalenceClass::from_iter(new_eq_exprs));
         }
 
         // Construct the resulting equivalence properties: