refactor(rust): Replace custom PushNode trait with Extend (pola-r…

…s#20107)
mcrumiller · Dec 2, 2024 · e805487 · e805487
1 parent 31b7bb9
commit e805487
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 78 deletions.
diff --git a/crates/polars-plan/src/plans/aexpr/traverse.rs b/crates/polars-plan/src/plans/aexpr/traverse.rs
@@ -2,85 +2,72 @@ use super::*;
 
 impl AExpr {
     /// Push nodes at this level to a pre-allocated stack.
-    pub(crate) fn nodes(&self, container: &mut impl PushNode) {
+    pub(crate) fn nodes<E>(&self, container: &mut E)
+    where
+        E: Extend<Node>,
+    {
         use AExpr::*;
 
         match self {
             Column(_) | Literal(_) | Len => {},
-            Alias(e, _) => container.push_node(*e),
+            Alias(e, _) => container.extend([*e]),
             BinaryExpr { left, op: _, right } => {
                 // reverse order so that left is popped first
-                container.push_node(*right);
-                container.push_node(*left);
+                container.extend([*right, *left]);
             },
-            Cast { expr, .. } => container.push_node(*expr),
-            Sort { expr, .. } => container.push_node(*expr),
+            Cast { expr, .. } => container.extend([*expr]),
+            Sort { expr, .. } => container.extend([*expr]),
             Gather { expr, idx, .. } => {
-                container.push_node(*idx);
-                // latest, so that it is popped first
-                container.push_node(*expr);
+                container.extend([*idx, *expr]);
             },
             SortBy { expr, by, .. } => {
-                for node in by {
-                    container.push_node(*node)
-                }
+                container.extend(by.iter().cloned());
                 // latest, so that it is popped first
-                container.push_node(*expr);
+                container.extend([*expr]);
             },
             Filter { input, by } => {
-                container.push_node(*by);
-                // latest, so that it is popped first
-                container.push_node(*input);
+                container.extend([*by, *input]);
             },
             Agg(agg_e) => match agg_e.get_input() {
-                NodeInputs::Single(node) => container.push_node(node),
-                NodeInputs::Many(nodes) => container.extend_from_slice(&nodes),
+                NodeInputs::Single(node) => container.extend([node]),
+                NodeInputs::Many(nodes) => container.extend(nodes),
                 NodeInputs::Leaf => {},
             },
             Ternary {
                 truthy,
                 falsy,
                 predicate,
             } => {
-                container.push_node(*predicate);
-                container.push_node(*falsy);
-                // latest, so that it is popped first
-                container.push_node(*truthy);
+                container.extend([*predicate, *falsy, *truthy]);
             },
             AnonymousFunction { input, .. } | Function { input, .. } =>
             // we iterate in reverse order, so that the lhs is popped first and will be found
             // as the root columns/ input columns by `_suffix` and `_keep_name` etc.
             {
-                input
-                    .iter()
-                    .rev()
-                    .for_each(|e| container.push_node(e.node()))
+                container.extend(input.iter().rev().map(|e| e.node()))
             },
-            Explode(e) => container.push_node(*e),
+            Explode(e) => container.extend([*e]),
             Window {
                 function,
                 partition_by,
                 order_by,
                 options: _,
             } => {
                 if let Some((n, _)) = order_by {
-                    container.push_node(*n);
-                }
-                for e in partition_by.iter().rev() {
-                    container.push_node(*e);
+                    container.extend([*n]);
                 }
+
+                container.extend(partition_by.iter().rev().cloned());
+
                 // latest so that it is popped first
-                container.push_node(*function);
+                container.extend([*function]);
             },
             Slice {
                 input,
                 offset,
                 length,
             } => {
-                container.push_node(*length);
-                container.push_node(*offset);
-                // latest so that it is popped first
-                container.push_node(*input);
+                container.extend([*length, *offset, *input]);
             },
         }
     }

diff --git a/crates/polars-plan/src/plans/aexpr/utils.rs b/crates/polars-plan/src/plans/aexpr/utils.rs
@@ -26,7 +26,7 @@ pub fn is_elementwise(stack: &mut UnitVec<Node>, ae: &AExpr, expr_arena: &Arena<
                 let rhs = rhs.node();
 
                 if matches!(expr_arena.get(rhs), AExpr::Literal { .. }) {
-                    stack.push_node(input[0].node());
+                    stack.extend([input[0].node()]);
                     return;
                 }
             };

diff --git a/crates/polars-plan/src/plans/ir/inputs.rs b/crates/polars-plan/src/plans/ir/inputs.rs
@@ -214,20 +214,16 @@ impl IR {
     /// or an in-memory DataFrame has none. A Union has multiple.
     pub fn copy_inputs<T>(&self, container: &mut T)
     where
-        T: PushNode,
+        T: Extend<Node>,
     {
         use IR::*;
         let input = match self {
             Union { inputs, .. } => {
-                for node in inputs {
-                    container.push_node(*node);
-                }
+                container.extend(inputs.iter().cloned());
                 return;
             },
             HConcat { inputs, .. } => {
-                for node in inputs {
-                    container.push_node(*node);
-                }
+                container.extend(inputs.iter().cloned());
                 return;
             },
             Slice { input, .. } => *input,
@@ -243,8 +239,7 @@ impl IR {
                 input_right,
                 ..
             } => {
-                container.push_node(*input_left);
-                container.push_node(*input_right);
+                container.extend([*input_left, *input_right]);
                 return;
             },
             HStack { input, .. } => *input,
@@ -254,9 +249,7 @@ impl IR {
             ExtContext {
                 input, contexts, ..
             } => {
-                for n in contexts {
-                    container.push_node(*n)
-                }
+                container.extend(contexts.iter().cloned());
                 *input
             },
             Scan { .. } => return,
@@ -265,7 +258,7 @@ impl IR {
             PythonScan { .. } => return,
             Invalid => unreachable!(),
         };
-        container.push_node(input)
+        container.extend([input])
     }
 
     pub fn get_inputs(&self) -> UnitVec<Node> {

diff --git a/crates/polars-plan/src/utils.rs b/crates/polars-plan/src/utils.rs
@@ -2,7 +2,6 @@ use std::fmt::Formatter;
 use std::iter::FlatMap;
 
 use polars_core::prelude::*;
-use polars_utils::idx_vec::UnitVec;
 
 use crate::constants::get_len_name;
 use crate::prelude::*;
@@ -40,33 +39,6 @@ pub(crate) fn fmt_column_delimited<S: AsRef<str>>(
     write!(f, "{container_end}")
 }
 
-// TODO: Remove this and use `Extend<Node>` instead.
-pub trait PushNode {
-    fn push_node(&mut self, value: Node);
-
-    fn extend_from_slice(&mut self, values: &[Node]);
-}
-
-impl PushNode for Vec<Node> {
-    fn push_node(&mut self, value: Node) {
-        self.push(value)
-    }
-
-    fn extend_from_slice(&mut self, values: &[Node]) {
-        Vec::extend_from_slice(self, values)
-    }
-}
-
-impl PushNode for UnitVec<Node> {
-    fn push_node(&mut self, value: Node) {
-        self.push(value)
-    }
-
-    fn extend_from_slice(&mut self, values: &[Node]) {
-        UnitVec::extend(self, values.iter().copied())
-    }
-}
-
 pub(crate) fn is_scan(plan: &IR) -> bool {
     matches!(plan, IR::Scan { .. } | IR::DataFrameScan { .. })
 }