perf: reduce heap allocs in expression/logical-plan iteration

crates/polars-plan/src/logical_plan/aexpr/mod.rs

@@ -255,38 +255,38 @@ impl AExpr {
     }
 
     /// Push nodes at this level to a pre-allocated stack
-    pub(crate) fn nodes(&self, container: &mut Vec<Node>) {
+    pub(crate) fn nodes<C: PushNode>(&self, container: &mut C) {
         use AExpr::*;
 
         match self {
             Nth(_) | Column(_) | Literal(_) | Wildcard | Len => {},
-            Alias(e, _) => container.push(*e),
+            Alias(e, _) => container.push_node(*e),
             BinaryExpr { left, op: _, right } => {
                 // reverse order so that left is popped first
-                container.push(*right);
-                container.push(*left);
+                container.push_node(*right);
+                container.push_node(*left);
             },
-            Cast { expr, .. } => container.push(*expr),
-            Sort { expr, .. } => container.push(*expr),
+            Cast { expr, .. } => container.push_node(*expr),
+            Sort { expr, .. } => container.push_node(*expr),
             Gather { expr, idx, .. } => {
-                container.push(*idx);
+                container.push_node(*idx);
                 // latest, so that it is popped first
-                container.push(*expr);
+                container.push_node(*expr);
             },
             SortBy { expr, by, .. } => {
                 for node in by {
-                    container.push(*node)
+                    container.push_node(*node)
                 }
                 // latest, so that it is popped first
-                container.push(*expr);
+                container.push_node(*expr);
             },
             Filter { input, by } => {
-                container.push(*by);
+                container.push_node(*by);
                 // latest, so that it is popped first
-                container.push(*input);
+                container.push_node(*input);
             },
             Agg(agg_e) => match agg_e.get_input() {
-                NodeInputs::Single(node) => container.push(node),
+                NodeInputs::Single(node) => container.push_node(node),
                 NodeInputs::Many(nodes) => container.extend_from_slice(&nodes),
                 NodeInputs::Leaf => {},
             },
@@ -295,10 +295,10 @@ impl AExpr {
                 falsy,
                 predicate,
             } => {
-                container.push(*predicate);
-                container.push(*falsy);
+                container.push_node(*predicate);
+                container.push_node(*falsy);
                 // latest, so that it is popped first
-                container.push(*truthy);
+                container.push_node(*truthy);
             },
             AnonymousFunction { input, .. } | Function { input, .. } =>
             // we iterate in reverse order, so that the lhs is popped first and will be found
@@ -308,29 +308,29 @@ impl AExpr {
                     .iter()
                     .rev()
                     .copied()
-                    .for_each(|node| container.push(node))
+                    .for_each(|node| container.push_node(node))
             },
-            Explode(e) => container.push(*e),
+            Explode(e) => container.push_node(*e),
             Window {
                 function,
                 partition_by,
                 options: _,
             } => {
                 for e in partition_by.iter().rev() {
-                    container.push(*e);
+                    container.push_node(*e);
                 }
                 // latest so that it is popped first
-                container.push(*function);
+                container.push_node(*function);
             },
             Slice {
                 input,
                 offset,
                 length,
             } => {
-                container.push(*length);
-                container.push(*offset);
+                container.push_node(*length);
+                container.push_node(*offset);
                 // latest so that it is popped first
-                container.push(*input);
+                container.push_node(*input);
             },
         }
     }

crates/polars-plan/src/logical_plan/iterator.rs

@@ -177,7 +177,7 @@ impl<'a> IntoIterator for &'a Expr {
 }
 
 pub struct AExprIter<'a> {
-    stack: Vec<Node>,
+    stack: UnitVec<Node>,
     arena: Option<&'a Arena<AExpr>>,
 }
 
@@ -203,8 +203,7 @@ pub trait ArenaExprIter<'a> {
 
 impl<'a> ArenaExprIter<'a> for &'a Arena<AExpr> {
     fn iter(&self, root: Node) -> AExprIter<'a> {
-        let mut stack = Vec::with_capacity(4);
-        stack.push(root);
+        let stack = unitvec![root];
         AExprIter {
             stack,
             arena: Some(self),
@@ -213,7 +212,7 @@ impl<'a> ArenaExprIter<'a> for &'a Arena<AExpr> {
 }
 
 pub struct AlpIter<'a> {
-    stack: Vec<Node>,
+    stack: UnitVec<Node>,
     arena: &'a Arena<ALogicalPlan>,
 }
 
@@ -223,7 +222,7 @@ pub trait ArenaLpIter<'a> {
 
 impl<'a> ArenaLpIter<'a> for &'a Arena<ALogicalPlan> {
     fn iter(&self, root: Node) -> AlpIter<'a> {
-        let stack = vec![root];
+        let stack = unitvec![root];
         AlpIter { stack, arena: self }
     }
 }

crates/polars-plan/src/utils.rs

@@ -43,18 +43,28 @@ pub(crate) fn fmt_column_delimited<S: AsRef<str>>(
 
 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: &ALogicalPlan) -> bool {
@@ -64,16 +74,6 @@ pub(crate) fn is_scan(plan: &ALogicalPlan) -> bool {
     )
 }
 
-impl PushNode for &mut [Option<Node>] {
-    fn push_node(&mut self, value: Node) {
-        if self[0].is_some() {
-            self[1] = Some(value)
-        } else {
-            self[0] = Some(value)
-        }
-    }
-}
-
 /// A projection that only takes a column or a column + alias.
 #[cfg(feature = "meta")]
 pub(crate) fn aexpr_is_simple_projection(current_node: Node, arena: &Arena<AExpr>) -> bool {

crates/polars-utils/src/idx_vec.rs

@@ -42,7 +42,13 @@ impl<T> UnitVec<T> {
         }
     }
 
+    #[inline]
     pub fn new() -> Self {
+        // This is optimized away, all const.
+        assert!(
+            std::mem::size_of::<T>() <= std::mem::size_of::<*mut T>()
+                && std::mem::align_of::<T>() <= std::mem::align_of::<*mut T>()
+        );
         Self {
             len: 0,
             capacity: NonZeroUsize::new(1).unwrap(),
@@ -120,18 +126,22 @@ impl<T> UnitVec<T> {
         new
     }
 
+    #[inline]
     pub fn iter(&self) -> std::slice::Iter<'_, T> {
         self.as_slice().iter()
     }
 
+    #[inline]
     pub fn iter_mut(&mut self) -> std::slice::IterMut<'_, T> {
         self.as_mut_slice().iter_mut()
     }
 
+    #[inline]
     pub fn as_slice(&self) -> &[T] {
         self.as_ref()
     }
 
+    #[inline]
     pub fn as_mut_slice(&mut self) -> &mut [T] {
         self.as_mut()
     }

py-polars/tests/unit/test_projections.py

@@ -391,7 +391,7 @@ def test_rolling_key_projected_13617() -> None:
     assert out.to_dict(as_series=False) == {"value": [["a"], ["b"]]}
 
 
-def test_projection_drop_with_series_lit_() -> None:
+def test_projection_drop_with_series_lit_14382() -> None:
     df = pl.DataFrame({"b": [1, 6, 8, 7]})
     df2 = pl.DataFrame({"a": [1, 2, 4, 4], "b": [True, True, True, False]})