refactor: make dsl immutable and cheap to clone (#15394)

crates/polars-lazy/src/physical_plan/exotic.rs

@@ -4,19 +4,12 @@ use crate::physical_plan::planner::create_physical_expr;
 use crate::prelude::*;
 
 #[cfg(feature = "pivot")]
-pub(crate) fn prepare_eval_expr(mut expr: Expr) -> Expr {
-    expr.mutate().apply(|e| match e {
-        Expr::Column(name) => {
-            *name = Arc::from("");
-            true
-        },
-        Expr::Nth(_) => {
-            *e = Expr::Column(Arc::from(""));
-            true
-        },
-        _ => true,
-    });
-    expr
+pub(crate) fn prepare_eval_expr(expr: Expr) -> Expr {
+    expr.map_expr(|e| match e {
+        Expr::Column(_) => Expr::Column(Arc::from("")),
+        Expr::Nth(_) => Expr::Column(Arc::from("")),
+        e => e,
+    })
 }
 
 pub(crate) fn prepare_expression_for_context(

crates/polars-plan/src/dsl/arity.rs

@@ -139,17 +139,17 @@ pub fn when<E: Into<Expr>>(condition: E) -> When {
 
 pub fn ternary_expr(predicate: Expr, truthy: Expr, falsy: Expr) -> Expr {
     Expr::Ternary {
-        predicate: Box::new(predicate),
-        truthy: Box::new(truthy),
-        falsy: Box::new(falsy),
+        predicate: Arc::new(predicate),
+        truthy: Arc::new(truthy),
+        falsy: Arc::new(falsy),
     }
 }
 
 /// Compute `op(l, r)` (or equivalently `l op r`). `l` and `r` must have types compatible with the Operator.
 pub fn binary_expr(l: Expr, op: Operator, r: Expr) -> Expr {
     Expr::BinaryExpr {
-        left: Box::new(l),
+        left: Arc::new(l),
         op,
-        right: Box::new(r),
+        right: Arc::new(r),
     }
 }

crates/polars-plan/src/dsl/expr.rs

@@ -12,30 +12,30 @@ use crate::prelude::*;
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub enum AggExpr {
     Min {
-        input: Box<Expr>,
+        input: Arc<Expr>,
         propagate_nans: bool,
     },
     Max {
-        input: Box<Expr>,
+        input: Arc<Expr>,
         propagate_nans: bool,
     },
-    Median(Box<Expr>),
-    NUnique(Box<Expr>),
-    First(Box<Expr>),
-    Last(Box<Expr>),
-    Mean(Box<Expr>),
-    Implode(Box<Expr>),
+    Median(Arc<Expr>),
+    NUnique(Arc<Expr>),
+    First(Arc<Expr>),
+    Last(Arc<Expr>),
+    Mean(Arc<Expr>),
+    Implode(Arc<Expr>),
     // include_nulls
-    Count(Box<Expr>, bool),
+    Count(Arc<Expr>, bool),
     Quantile {
-        expr: Box<Expr>,
-        quantile: Box<Expr>,
+        expr: Arc<Expr>,
+        quantile: Arc<Expr>,
         interpol: QuantileInterpolOptions,
     },
-    Sum(Box<Expr>),
-    AggGroups(Box<Expr>),
-    Std(Box<Expr>, u8),
-    Var(Box<Expr>, u8),
+    Sum(Arc<Expr>),
+    AggGroups(Arc<Expr>),
+    Std(Arc<Expr>, u8),
+    Var(Arc<Expr>, u8),
 }
 
 impl AsRef<Expr> for AggExpr {
@@ -67,42 +67,42 @@ impl AsRef<Expr> for AggExpr {
 #[must_use]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
 pub enum Expr {
-    Alias(Box<Expr>, Arc<str>),
+    Alias(Arc<Expr>, Arc<str>),
     Column(Arc<str>),
     Columns(Vec<String>),
     DtypeColumn(Vec<DataType>),
     Literal(LiteralValue),
     BinaryExpr {
-        left: Box<Expr>,
+        left: Arc<Expr>,
         op: Operator,
-        right: Box<Expr>,
+        right: Arc<Expr>,
     },
     Cast {
-        expr: Box<Expr>,
+        expr: Arc<Expr>,
         data_type: DataType,
         strict: bool,
     },
     Sort {
-        expr: Box<Expr>,
+        expr: Arc<Expr>,
         options: SortOptions,
     },
     Gather {
-        expr: Box<Expr>,
-        idx: Box<Expr>,
+        expr: Arc<Expr>,
+        idx: Arc<Expr>,
         returns_scalar: bool,
     },
     SortBy {
-        expr: Box<Expr>,
+        expr: Arc<Expr>,
         by: Vec<Expr>,
         descending: Vec<bool>,
     },
     Agg(AggExpr),
     /// A ternary operation
     /// if true then "foo" else "bar"
     Ternary {
-        predicate: Box<Expr>,
-        truthy: Box<Expr>,
-        falsy: Box<Expr>,
+        predicate: Arc<Expr>,
+        truthy: Arc<Expr>,
+        falsy: Arc<Expr>,
     },
     Function {
         /// function arguments
@@ -111,37 +111,37 @@ pub enum Expr {
         function: FunctionExpr,
         options: FunctionOptions,
     },
-    Explode(Box<Expr>),
+    Explode(Arc<Expr>),
     Filter {
-        input: Box<Expr>,
-        by: Box<Expr>,
+        input: Arc<Expr>,
+        by: Arc<Expr>,
     },
     /// See postgres window functions
     Window {
         /// Also has the input. i.e. avg("foo")
-        function: Box<Expr>,
+        function: Arc<Expr>,
         partition_by: Vec<Expr>,
         options: WindowType,
     },
     Wildcard,
     Slice {
-        input: Box<Expr>,
+        input: Arc<Expr>,
         /// length is not yet known so we accept negative offsets
-        offset: Box<Expr>,
-        length: Box<Expr>,
+        offset: Arc<Expr>,
+        length: Arc<Expr>,
     },
     /// Can be used in a select statement to exclude a column from selection
-    Exclude(Box<Expr>, Vec<Excluded>),
+    Exclude(Arc<Expr>, Vec<Excluded>),
     /// Set root name as Alias
-    KeepName(Box<Expr>),
+    KeepName(Arc<Expr>),
     Len,
     /// Take the nth column in the `DataFrame`
     Nth(i64),
     // skipped fields must be last otherwise serde fails in pickle
     #[cfg_attr(feature = "serde", serde(skip))]
     RenameAlias {
         function: SpecialEq<Arc<dyn RenameAliasFn>>,
-        expr: Box<Expr>,
+        expr: Arc<Expr>,
     },
     AnonymousFunction {
         /// function arguments

crates/polars-plan/src/dsl/functions/syntactic_sugar.rs

@@ -57,7 +57,7 @@ pub fn is_not_null(expr: Expr) -> Expr {
 /// nominal type of the column.
 pub fn cast(expr: Expr, data_type: DataType) -> Expr {
     Expr::Cast {
-        expr: Box::new(expr),
+        expr: Arc::new(expr),
         data_type,
         strict: false,
     }

crates/polars-plan/src/dsl/meta.rs

@@ -41,22 +41,13 @@ impl MetaNameSpace {
     }
 
     /// Undo any renaming operation like `alias`, `keep_name`.
-    pub fn undo_aliases(mut self) -> Expr {
-        self.0.mutate().apply(|e| match e {
+    pub fn undo_aliases(self) -> Expr {
+        self.0.map_expr(|e| match e {
             Expr::Alias(input, _)
             | Expr::KeepName(input)
-            | Expr::RenameAlias { expr: input, .. } => {
-                // remove this node
-                *e = *input.clone();
-
-                // continue iteration
-                true
-            },
-            // continue iteration
-            _ => true,
-        });
-
-        self.0
+            | Expr::RenameAlias { expr: input, .. } => Arc::unwrap_or_clone(input),
+            e => e,
+        })
     }
 
     /// Indicate if this expression expands to multiple expressions.