diff --git a/crates/polars-plan/src/plans/aexpr/traverse.rs b/crates/polars-plan/src/plans/aexpr/traverse.rs index 20e7a454169c..3850cbafca6e 100644 --- a/crates/polars-plan/src/plans/aexpr/traverse.rs +++ b/crates/polars-plan/src/plans/aexpr/traverse.rs @@ -2,39 +2,35 @@ 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(&self, container: &mut E) + where + E: Extend, + { 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 { @@ -42,21 +38,15 @@ impl AExpr { 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, @@ -64,23 +54,20 @@ impl AExpr { 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 index 92657acc5340..834aa4ff6a75 100644 --- 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, 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 index 2a7c14e300de..865d398959ff 100644 --- 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(&self, container: &mut T) where - T: PushNode, + T: Extend, { 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 { diff --git a/crates/polars-plan/src/utils.rs b/crates/polars-plan/src/utils.rs index fdf9c979738a..0d433e9e74d7 100644 --- 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>( write!(f, "{container_end}") } -// TODO: Remove this and use `Extend` instead. -pub trait PushNode { - fn push_node(&mut self, value: Node); - - fn extend_from_slice(&mut self, values: &[Node]); -} - -impl PushNode for Vec { - 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 { - 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 { .. }) }