refactor: Don't take ownership of IRplan in new streaming engine (#21551

)
pola-rs · Mar 1, 2025 · 1ad0ec2 · 1ad0ec2
1 parent 6497981
commit 1ad0ec2
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/crates/polars-lazy/src/frame/mod.rs b/crates/polars-lazy/src/frame/mod.rs
@@ -871,7 +871,11 @@ impl LazyFrame {
 
             let _hold = StringCacheHolder::hold();
             let f = || {
-                polars_stream::run_query(stream_lp_top, alp_plan.lp_arena, &mut alp_plan.expr_arena)
+                polars_stream::run_query(
+                    stream_lp_top,
+                    &mut alp_plan.lp_arena,
+                    &mut alp_plan.expr_arena,
+                )
             };
             match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) {
                 Ok(v) => return Some(v),

diff --git a/crates/polars-stream/src/skeleton.rs b/crates/polars-stream/src/skeleton.rs
@@ -12,7 +12,7 @@ use slotmap::{SecondaryMap, SlotMap};
 
 pub fn run_query(
     node: Node,
-    mut ir_arena: Arena<IR>,
+    ir_arena: &mut Arena<IR>,
     expr_arena: &mut Arena<AExpr>,
 ) -> PolarsResult<Option<DataFrame>> {
     if let Ok(visual_path) = std::env::var("POLARS_VISUALIZE_IR") {
@@ -25,8 +25,7 @@ pub fn run_query(
         std::fs::write(visual_path, visualization).unwrap();
     }
     let mut phys_sm = SlotMap::with_capacity_and_key(ir_arena.len());
-    let root =
-        crate::physical_plan::build_physical_plan(node, &mut ir_arena, expr_arena, &mut phys_sm)?;
+    let root = crate::physical_plan::build_physical_plan(node, ir_arena, expr_arena, &mut phys_sm)?;
     if let Ok(visual_path) = std::env::var("POLARS_VISUALIZE_PHYSICAL_PLAN") {
         let visualization = crate::physical_plan::visualize_plan(root, &phys_sm, expr_arena);
         std::fs::write(visual_path, visualization).unwrap();