perf: Make LogicalPlan immutable (#15416)

pola-rs · Apr 1, 2024 · 758b55a · 758b55a
1 parent 82f717b
commit 758b55a
Show file tree

Hide file tree

Showing 5 changed files with 67 additions and 66 deletions.
diff --git a/crates/polars-lazy/src/frame/mod.rs b/crates/polars-lazy/src/frame/mod.rs
@@ -751,7 +751,7 @@ impl LazyFrame {
     ) -> PolarsResult<()> {
         self.opt_state.streaming = true;
         self.logical_plan = LogicalPlan::Sink {
-            input: Box::new(self.logical_plan),
+            input: Arc::new(self.logical_plan),
             payload: SinkType::Cloud {
                 uri: Arc::new(uri),
                 cloud_options,
@@ -806,7 +806,7 @@ impl LazyFrame {
     fn sink(mut self, payload: SinkType, msg_alternative: &str) -> Result<(), PolarsError> {
         self.opt_state.streaming = true;
         self.logical_plan = LogicalPlan::Sink {
-            input: Box::new(self.logical_plan),
+            input: Arc::new(self.logical_plan),
             payload,
         };
         let (mut state, mut physical_plan, is_streaming) = self.prepare_collect(true)?;
@@ -1846,7 +1846,7 @@ impl LazyGroupBy {
         let options = GroupbyOptions { slice: None };
 
         let lp = LogicalPlan::Aggregate {
-            input: Box::new(self.logical_plan),
+            input: Arc::new(self.logical_plan),
             keys: Arc::new(self.keys),
             aggs: vec![],
             schema,

diff --git a/crates/polars-lazy/src/physical_plan/executors/group_by_partitioned.rs b/crates/polars-lazy/src/physical_plan/executors/group_by_partitioned.rs
@@ -257,7 +257,7 @@ impl PartitionGroupByExec {
         }
         .into();
         let lp = LogicalPlan::Aggregate {
-            input: Box::new(original_df.lazy().logical_plan),
+            input: Arc::new(original_df.lazy().logical_plan),
             keys: Arc::new(std::mem::take(&mut self.keys)),
             aggs: std::mem::take(&mut self.aggs),
             schema: self.output_schema.clone(),

diff --git a/crates/polars-plan/src/logical_plan/builder.rs b/crates/polars-plan/src/logical_plan/builder.rs
@@ -68,7 +68,7 @@ macro_rules! raise_err {
                 let err = $err.wrap_msg(&format_err_outer);
 
                 LogicalPlan::Error {
-                    input: Box::new(input),
+                    input: Arc::new(input),
                     err: err.into(), // PolarsError -> ErrorState
                 }
             },
@@ -424,7 +424,7 @@ impl LogicalPlanBuilder {
     }
 
     pub fn cache(self) -> Self {
-        let input = Box::new(self.0);
+        let input = Arc::new(self.0);
         let id = input.as_ref() as *const LogicalPlan as usize;
         LogicalPlan::Cache {
             input,
@@ -461,7 +461,7 @@ impl LogicalPlanBuilder {
         } else {
             LogicalPlan::Projection {
                 expr: columns,
-                input: Box::new(self.0),
+                input: Arc::new(self.0),
                 schema: Arc::new(output_schema),
                 options: ProjectionOptions {
                     run_parallel: false,
@@ -486,7 +486,7 @@ impl LogicalPlanBuilder {
         } else {
             LogicalPlan::Projection {
                 expr: exprs,
-                input: Box::new(self.0),
+                input: Arc::new(self.0),
                 schema: Arc::new(schema),
                 options,
             }
@@ -576,7 +576,7 @@ impl LogicalPlanBuilder {
         }
 
         LogicalPlan::HStack {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             exprs,
             schema: Arc::new(new_schema),
             options,
@@ -586,7 +586,7 @@ impl LogicalPlanBuilder {
 
     pub fn add_err(self, err: PolarsError) -> Self {
         LogicalPlan::Error {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             err: err.into(),
         }
         .into()
@@ -608,7 +608,7 @@ impl LogicalPlanBuilder {
             }
         }
         LogicalPlan::ExtContext {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             contexts,
             schema: Arc::new(schema),
         }
@@ -692,7 +692,7 @@ impl LogicalPlanBuilder {
 
         LogicalPlan::Selection {
             predicate,
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
         }
         .into()
     }
@@ -777,7 +777,7 @@ impl LogicalPlanBuilder {
         };
 
         LogicalPlan::Aggregate {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             keys: Arc::new(keys),
             aggs,
             schema: Arc::new(schema),
@@ -814,7 +814,7 @@ impl LogicalPlanBuilder {
         let schema = try_delayed!(self.0.schema(), &self.0, into);
         let by_column = try_delayed!(rewrite_projections(by_column, &schema, &[]), &self.0, into);
         LogicalPlan::Sort {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             by_column,
             args: SortArguments {
                 descending,
@@ -846,7 +846,7 @@ impl LogicalPlanBuilder {
         try_delayed!(explode_schema(&mut schema, &columns), &self.0, into);
 
         LogicalPlan::MapFunction {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             function: FunctionNode::Explode {
                 columns,
                 schema: Arc::new(schema),
@@ -859,7 +859,7 @@ impl LogicalPlanBuilder {
         let schema = try_delayed!(self.0.schema(), &self.0, into);
         let schema = det_melt_schema(&args, &schema);
         LogicalPlan::MapFunction {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             function: FunctionNode::Melt { args, schema },
         }
         .into()
@@ -871,7 +871,7 @@ impl LogicalPlanBuilder {
         row_index_schema(schema_mut, name);
 
         LogicalPlan::MapFunction {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             function: FunctionNode::RowIndex {
                 name: ColumnName::from(name),
                 offset,
@@ -883,15 +883,15 @@ impl LogicalPlanBuilder {
 
     pub fn distinct(self, options: DistinctOptions) -> Self {
         LogicalPlan::Distinct {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             options,
         }
         .into()
     }
 
     pub fn slice(self, offset: i64, len: IdxSize) -> Self {
         LogicalPlan::Slice {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             offset,
             len,
         }
@@ -908,7 +908,7 @@ impl LogicalPlanBuilder {
         for e in left_on.iter().chain(right_on.iter()) {
             if has_expr(e, |e| matches!(e, Expr::Alias(_, _))) {
                 return LogicalPlan::Error {
-                    input: Box::new(self.0),
+                    input: Arc::new(self.0),
                     err: polars_err!(
                         ComputeError:
                         "'alias' is not allowed in a join key, use 'with_columns' first",
@@ -929,8 +929,8 @@ impl LogicalPlanBuilder {
         );
 
         LogicalPlan::Join {
-            input_left: Box::new(self.0),
-            input_right: Box::new(other),
+            input_left: Arc::new(self.0),
+            input_right: Arc::new(other),
             schema,
             left_on,
             right_on,
@@ -940,7 +940,7 @@ impl LogicalPlanBuilder {
     }
     pub fn map_private(self, function: FunctionNode) -> Self {
         LogicalPlan::MapFunction {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             function,
         }
         .into()
@@ -955,7 +955,7 @@ impl LogicalPlanBuilder {
         validate_output: bool,
     ) -> Self {
         LogicalPlan::MapFunction {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             function: FunctionNode::OpaquePython {
                 function,
                 schema,
@@ -981,7 +981,7 @@ impl LogicalPlanBuilder {
         let function = Arc::new(function);
 
         LogicalPlan::MapFunction {
-            input: Box::new(self.0),
+            input: Arc::new(self.0),
             function: FunctionNode::Opaque {
                 function,
                 schema,