Skip to content

Commit

Permalink
refactor: move describe to IR instead of DSL (#16191)
Browse files Browse the repository at this point in the history
  • Loading branch information
coastalwhite committed May 14, 2024
1 parent 07d58ab commit 80442f1
Show file tree
Hide file tree
Showing 7 changed files with 658 additions and 17 deletions.
14 changes: 12 additions & 2 deletions crates/polars-lazy/src/frame/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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.
Expand Down
2 changes: 1 addition & 1 deletion crates/polars-plan/src/logical_plan/aexpr/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
Expand Down
Loading

0 comments on commit 80442f1

Please sign in to comment.