refactor: move describe to IR instead of DSL (pola-rs#16191)

orlp · May 14, 2024 · 80442f1 · 80442f1
1 parent 07d58ab
commit 80442f1
Show file tree

Hide file tree

Showing 7 changed files with 658 additions and 17 deletions.
diff --git a/crates/polars-lazy/src/frame/mod.rs b/crates/polars-lazy/src/frame/mod.rs
@@ -215,7 +215,7 @@ impl LazyFrame {
         self.logical_plan.describe_tree_format()
     }
 
-    fn optimized_plan(&self) -> PolarsResult<DslPlan> {
+    fn optimized_plan_ir(&self) -> PolarsResult<IRPlan> {
         let mut expr_arena = Arena::with_capacity(64);
         let mut lp_arena = Arena::with_capacity(64);
         let lp_top = self.clone().optimize_with_scratch(
@@ -224,14 +224,24 @@ impl LazyFrame {
             &mut vec![],
             true,
         )?;
+
+        Ok(IRPlan::new(lp_top, lp_arena, expr_arena))
+    }
+
+    fn optimized_plan(&self) -> PolarsResult<DslPlan> {
+        let IRPlan {
+            lp_top,
+            mut lp_arena,
+            expr_arena,
+        } = self.optimized_plan_ir()?;
         Ok(node_to_lp(lp_top, &expr_arena, &mut lp_arena))
     }
 
     /// Return a String describing the optimized logical plan.
     ///
     /// Returns `Err` if optimizing the logical plan fails.
     pub fn describe_optimized_plan(&self) -> PolarsResult<String> {
-        Ok(self.optimized_plan()?.describe())
+        Ok(self.optimized_plan_ir()?.describe())
     }
 
     /// Return a String describing the optimized logical plan in tree format.

diff --git a/crates/polars-plan/src/logical_plan/aexpr/mod.rs b/crates/polars-plan/src/logical_plan/aexpr/mod.rs
@@ -121,7 +121,7 @@ impl From<AAggExpr> for GroupByMethod {
     }
 }
 
-// AExpr representation of Nodes which are allocated in an Arena
+/// IR expression node that is allocated in an [`Arena`][polars_utils::arena::Arena].
 #[derive(Clone, Debug, Default)]
 pub enum AExpr {
     Explode(Node),