Theme

Blog · Systems ·

Optimise once, execute forever: smuggling DataFrames into a compiled Polars plan

Polars re-runs its whole query optimizer on every collect(). For one fixed expression scored against a stream of small batches, that is the entire runtime, so I disguised the input as a file that never existed and skipped the optimizer altogether.

  • Interactive
  • rust
  • python
  • polars
  • query-planning
  • performance

I had one Polars expression (a scoring function, a decision tree, the kind of thing you’d call a feature pipeline) and a stream of small DataFrames arriving one at a time that all needed the same expression evaluated against them. Nothing about the expression changed call to call. Only the data did.

LazyFrame.collect() doesn’t know that. Every call runs the whole Polars optimizer from scratch (predicate pushdown, projection pushdown, common subexpression elimination, type coercion, the works) over the entire plan, even though the plan is identical to the one it optimized a millisecond ago. The README states the general shape of the problem plainly: for a fixed expression applied to a stream of small batches, “this optimization cost dominates execution time and grows super-linearly with plan complexity.” The benchmark numbers back that up without needing to be dramatic about it: the README’s linear-chain table averages 503 ms per call across the 100–999 node range, and reports 1.55 s exactly at depth 1000, for the same thousand-row batch every time.

Time per 1000-row batch versus depth (nodes) for a linear when/then chain, linear x-axis, comparing plain LazyFrame, RTLF and Compiled RTLF.

LazyFrame’s time climbs roughly linearly with node count on a linear x-axis, which is what “grows super-linearly with plan complexity” looks like once you stop hiding it behind a log scale.

The problem this post is about (this is rtlf, published to PyPI as polars-rt, an experimental workaround for pola-rs/polars#25246) isn’t a numeric kernel. There’s no SIMD in it, no hand-rolled parallelism, no custom allocator. The entire optimisation is structural: don’t re-run the planner. The previous post covers what it took to ship this as a wheel anyone could pip install: the ABI hash, the 72-job CI matrix, the manifest-rewriting script. This one is about the 480 lines of Rust that make the trick work, and about the part of it that’s honestly still a little unsafe.

Attempt zero: just cache the LazyFrame?

The first thing that doesn’t work is the obvious thing. A LazyFrame is a DslPlan: Scan / Filter / Select / … nodes wrapping the data and the operations you asked for. The DataFrame you eventually collect is stitched into a Scan node when the LazyFrame is built, not passed in separately at collect() time. There’s no seam between “the plan” and “this call’s data” to cache one half of. The optimizer has to see something that looks like a real data source before it can do anything with the query at all.

So the seam had to be built.

The placeholder trick

read_placeholder() builds a LazyFrame that looks like a Parquet scan to the optimizer but is not backed by any file on disk:

// src/realtime.rs:119-165
pub fn read_placeholder(name: &str, schema: &Schema) -> LazyFrame {
    let schema_ref: SchemaRef = Arc::new(schema.clone());

    let sources = ScanSources::Paths(Buffer::from_iter([
        PlRefPath::new(PLACEHOLDER_TOKEN),
        PlRefPath::new(name),
    ]));

    let file_info = FileInfo {
        schema: schema_ref.clone(),
        reader_schema: None,
        row_estimation: (None, usize::MAX),
    };
    let ir_scan = IR::Scan {
        sources: sources.clone(),
        file_info,
        hive_parts: None,
        predicate: None,
        predicate_file_skip_applied: None,
        output_schema: None,
        scan_type: Box::new(FileScanIR::Parquet {
            options: polars_io::parquet::read::ParquetOptions::default(),
            metadata: None,
        }),
        unified_scan_args: Box::new(UnifiedScanArgs {
            glob: false,
            ..Default::default()
        }),
    };

    let scan = DslPlan::Scan {
        sources,
        scan_type: Box::new(FileScanDsl::Parquet {
            options: polars_io::parquet::read::ParquetOptions {
                schema: Some(schema_ref.clone()),
                ..Default::default()
            },
        }),
        unified_scan_args: Box::new(UnifiedScanArgs {
            glob: false,
            ..Default::default()
        }),
        cached_ir: Arc::new(Mutex::new(Some(ir_scan))),
    };

    LazyFrame::from(scan)
}

Two things are doing the actual work here. First, sources is ScanSources::Paths holding exactly two entries: _rtlf::placeholder (a constant marker, PLACEHOLDER_TOKEN) and the placeholder’s own name. Polars sees an ordinary two-file scan and optimizes around it exactly as it would around a real one: pushes predicates into it, includes it in projection pruning, all of it. Detecting one later is just checking the shape:

// src/executor.rs:88-98
pub fn placeholder_name_from_ir(ir: &IR) -> Option<String> {
    if let IR::Scan { sources, .. } = ir {
        let paths = sources.as_paths()?;
        if paths.len() != 2 || paths[0].as_str() != PLACEHOLDER_TOKEN {
            return None;
        }
        Some(paths[1].as_str().to_owned())
    } else {
        None
    }
}

Second, and this is the part that actually stops Polars from touching a filesystem: ParquetOptions { schema: Some(schema_ref), .. }. Normally a Parquet scan’s schema comes from reading the file’s own footer, which would mean opening _rtlf::placeholder and failing immediately. Setting the schema explicitly short-circuits that fetch entirely. The optimizer gets everything it needs to plan the query (column names, types, row-count estimate) without ever going near disk, because a file that doesn’t exist has nothing to be read.

Mode A: patch the arena, every call

RealtimeLazyFrame::new() runs the real optimizer exactly once, at construction, and then holds onto the two arenas to_alp_optimized() produced (the plan node arena and the expression arena) along with the arena index of every placeholder scan it found while walking the optimized tree:

// src/realtime.rs:38-58
impl RealtimeLazyFrame {
    pub fn new(lf: LazyFrame) -> PolarsResult<Self> {
        let ir_plan = lf.to_alp_optimized()?;
        let lp_top = ir_plan.lp_top;
        let lp_arena = ir_plan.lp_arena;
        let expr_arena = ir_plan.expr_arena;

        let mut placeholder_nodes = HashMap::new();
        for (node, ir) in lp_arena.iter(lp_top) {
            if let Some(name) = placeholder_name_from_ir(ir) {
                placeholder_nodes.insert(name, node);
            }
        }

        Ok(Self { lp_top, lp_arena, expr_arena, placeholder_nodes })
    }

Every subsequent collect() clones both arenas, walks straight to each recorded placeholder node, and swaps it for a real IR::DataFrameScan holding the DataFrame that call actually got:

// src/realtime.rs:71-113
pub fn collect(&self, mut inputs: HashMap<String, DataFrame>) -> PolarsResult<DataFrame> {
    for name in self.placeholder_nodes.keys() {
        if !inputs.contains_key(name) {
            polars_core::error::polars_bail!(
                ComputeError: "placeholder '{}' not provided; got: {:?}",
                name, inputs.keys().collect::<Vec<_>>()
            );
        }
    }

    let mut lp_arena = self.lp_arena.clone();
    let mut expr_arena = self.expr_arena.clone();

    for (name, &node) in &self.placeholder_nodes {
        // Extract schema and predicate before taking a mutable borrow.
        let (schema, predicate) = match lp_arena.get(node) {
            IR::Scan { file_info, predicate, .. } => {
                (file_info.schema.clone(), predicate.clone())
            },
            _ => unreachable!("placeholder node was not IR::Scan"),
        };
        let df = Arc::new(inputs.remove(name).expect("validated above"));
        let df_scan = IR::DataFrameScan { df, schema, output_schema: None };

        if let Some(pred) = predicate {
            // Predicate was pushed into the scan by the optimizer; restore it
            // as an explicit Filter so it actually executes.
            let scan_node = lp_arena.add(df_scan);
            *lp_arena.get_mut(node) = IR::Filter { input: scan_node, predicate: pred };
        } else {
            *lp_arena.get_mut(node) = df_scan;
        }
    }

    let mut physical_plan = create_physical_plan(
        self.lp_top, &mut lp_arena, &mut expr_arena,
        Some(polars_stream::build_streaming_query_executor),
    )?;
    physical_plan.execute(&mut ExecutionState::new())
}

That if let Some(pred) = predicate branch is worth pausing on, because skipping it is the kind of bug that doesn’t announce itself. If the query has a .filter(...) anywhere near the placeholder, predicate pushdown (one of the optimizer passes that already ran, once, at construction) will have folded that filter straight into the scan node’s predicate field, and deleted the standalone Filter node it came from. That’s normal, correct optimizer behaviour for a real file scan: a Parquet reader can skip row groups using a predicate, so pushing it into the scan is strictly better.

But IR::DataFrameScan cannot apply a predicate. It has no predicate field to put one in. It just is the DataFrame. If collect() swapped the placeholder for a bare DataFrameScan and dropped the predicate on the floor, the query would still run, still produce a DataFrame, and just quietly contain rows that should have been filtered out. No panic, no error, no strict=False xfail to catch it: a silently wrong answer. The fix is the four lines above it: wrap the new DataFrameScan in an explicit IR::Filter holding the predicate that got pushed down, so the row-level filtering that used to happen for free inside the scan now happens as its own step, and the result is correct instead of merely plausible-looking.

Mode B: compile once, inject by slot

Mode A eliminates the optimizer but still clones both arenas and calls create_physical_plan on every collect(). For a fixed plan, that’s still redundant work: the physical plan built from an unchanging IR is itself unchanging. CompiledRealtimeLazyFrame goes one step further: build the executor tree exactly once, and find a way to get fresh data into an already-built tree without touching it again.

The mechanism is a function-pointer hook Polars already exposes, StreamingExecutorBuilder, which create_physical_plan calls once per scan node while building the executor tree. Passing placeholder_builder as that hook means the compiler itself is what intercepts each placeholder, not a later patch pass:

// src/executor.rs:57-86
pub fn placeholder_builder(
    node: Node,
    lp_arena: &mut Arena<IR>,
    expr_arena: &mut polars_utils::arena::Arena<polars_plan::plans::AExpr>,
) -> PolarsResult<Box<dyn Executor>> {
    let ir = lp_arena.get(node);
    let name = placeholder_name_from_ir(ir);
    let projection: Option<SchemaRef> = if name.is_some() {
        if let IR::Scan { output_schema, .. } = ir { output_schema.clone() } else { None }
    } else {
        None
    };

    if let Some(name) = name {
        let slot: Slot = Arc::new(Mutex::new(None));
        PLACEHOLDER_REGISTRY.with(|r| {
            r.borrow_mut()
                .as_mut()
                .expect("placeholder_builder called outside of CompiledRealtimeLazyFrame::from_parts")
                .insert(name, slot.clone());
        });
        return Ok(Box::new(PlaceholderExec { slot, projection }));
    }
    polars_stream::build_streaming_query_executor(node, lp_arena, expr_arena)
}

Every placeholder scan becomes a PlaceholderExec wrapping a Slot (Arc<Mutex<Option<DataFrame>>>); every other node (joins, group-bys, real file scans) falls straight through to Polars’ own build_streaming_query_executor. At execute time, PlaceholderExec just takes whatever DataFrame is sitting in its slot:

// src/executor.rs:31-49
impl Executor for PlaceholderExec {
    fn execute(&mut self, _state: &mut ExecutionState) -> PolarsResult<DataFrame> {
        let df = self.slot
            .lock()
            .expect("placeholder slot poisoned")
            .take()
            .ok_or_else(|| {
                polars_core::error::polars_err!(
                    ComputeError: "placeholder slot empty — collect() was not called before execute()"
                )
            })?;

        if let Some(schema) = &self.projection {
            let cols: Vec<_> = schema.iter_names().cloned().collect();
            Ok(df.select(cols)?)
        } else {
            Ok(df)
        }
    }
}

That leaves one problem: placeholder_builder is a bare fn pointer, because that’s the type StreamingExecutorBuilder demands. A bare function pointer has no environment: it cannot close over a HashMap the way a closure could, so there’s no direct way for create_physical_plan’s caller to get the slots it just created back out. The fix is a thread-local acting as a side-channel purely for the duration of one compile call:

// src/executor.rs:19-22
thread_local! {
    pub static PLACEHOLDER_REGISTRY: RefCell<Option<HashMap<String, Slot>>> =
        const { RefCell::new(None) };
}

CompiledRealtimeLazyFrame::from_parts arms it right before compiling, and drains it right after:

// src/compiled.rs:62-79 (predicate re-lift above this, omitted: see below)
PLACEHOLDER_REGISTRY.with(|r| *r.borrow_mut() = Some(HashMap::new()));

let physical_plan = create_physical_plan(
    lp_top, lp_arena, expr_arena,
    Some(placeholder_builder as StreamingExecutorBuilder),
)?;

let placeholder_slots = PLACEHOLDER_REGISTRY
    .with(|r| r.borrow_mut().take())
    .unwrap_or_default();

Ok(Self { physical_plan: Mutex::new(physical_plan), placeholder_slots })

I don’t love this. A thread-local that only makes sense while a specific function further down the call stack is executing, written to by a callback with no idea it’s being watched, is exactly the kind of “why is this ugly” code that needs a comment explaining it’s not an accident, and it has one. But it’s not a workaround for bad design on rtlf’s side; it’s the direct consequence of StreamingExecutorBuilder being a fn pointer and not a Box<dyn FnMut(...)>. Given that constraint, a thread-local side-channel is close to the only way to get state out of a callback that can’t capture any.

The predicate-pushdown subtlety from Mode A applies here too, but it’s handled once, at compile time, rather than on every call: from_parts walks the arena for placeholder scans carrying a pushed-down predicate and re-lifts each one into an explicit IR::Filter before create_physical_plan ever runs (src/compiled.rs:37-60). The compiled executor tree is built with the filter already in the right place; collect() never has to think about it again:

// src/compiled.rs:81-99
pub fn collect(&self, mut inputs: HashMap<String, DataFrame>) -> PolarsResult<DataFrame> {
    for name in self.placeholder_slots.keys() {
        if !inputs.contains_key(name) {
            polars_core::error::polars_bail!(
                ComputeError: "placeholder '{}' not provided; got: {:?}",
                name, inputs.keys().collect::<Vec<_>>()
            );
        }
    }

    let mut plan = self.physical_plan.lock().expect("executor mutex poisoned");

    for (name, slot) in &self.placeholder_slots {
        *slot.lock().expect("placeholder slot poisoned") = inputs.remove(name);
    }

    plan.execute(&mut ExecutionState::new())
}

No clone. No arena walk. No create_physical_plan. Every call is a HashMap drain into a handful of mutexes and one execute() on a tree that was built exactly once.

The rejected design

That slot mechanism wasn’t the first thing tried, and the git history keeps the receipt. Commit a2fa2403, “Exec version using cache” (2026-05-24 07:08:38 +0200), routes DataFrames through ExecutionState’s built-in DataFrame cache instead of a hand-rolled slot, the sanctioned Polars mechanism for handing a physical plan pre-computed data:

// a2fa2403, src/core.rs (as it stood after this commit)
struct PlaceholderExec {
    id: UniqueId,
}

impl Executor for PlaceholderExec {
    fn execute(&mut self, state: &mut ExecutionState) -> PolarsResult<DataFrame> {
        Ok(state.get_df_cache(&self.id))
    }
}

pub fn collect(&self, inputs: HashMap<String, DataFrame>) -> PolarsResult<DataFrame> {
    // ...
    let mut state = ExecutionState::new();
    for (name, id) in &self.placeholders {
        // hit_count = 1: each placeholder is used exactly once per collect.
        state.set_df_cache(id, inputs[name].clone(), 1);
    }
    self.physical_plan.lock().expect("executor mutex poisoned").execute(&mut state)
}

It worked. Seventy-nine seconds later (literally: the next commit’s timestamp is 07:09:57, one minute and nineteen seconds after the first, 07:08:38) it was gone, replaced by the Arc<Mutex<Option<DataFrame>>> slot design this post has been describing all along. The commit message says why: 3d06d7c3, “change to slots for better perfomance zero df copy”. state.set_df_cache(id, inputs[name].clone(), 1) clones the DataFrame into the cache on every call; the slot version moves it: inputs.remove(name) straight into *slot.lock().unwrap() = ..., no clone, no Arc refcount bump on the frame itself. The ExecutionState cache is the API Polars actually publishes for this (indexed by UniqueId, with a hit-count so the cache knows when an entry is spent) and it was measurably slower than a plain mutex the sanctioned API doesn’t need to know about, for the mundane reason that “the sanctioned way” carries indirection (a cache lookup by ID, an allocation per entry) that a private field on a struct doesn’t have to pay for. The commit that keeps the slower, more “correct” version doesn’t exist in this repo: it’s the one-minute gap between two commit timestamps.

Where it breaks, and why that’s the interesting part

Reusing a compiled executor tree only works if every executor in it is safe to run twice. Most are. Some aren’t, and the reason is structural rather than a bug that could plausibly get fixed later: Executor::execute takes &mut self, and Polars’ streaming executors were written for a tree that gets built once and executed once, so several of them treat that first execute() as consuming something.

Union’s streaming implementation drains a one-shot channel: reading it a second time reads nothing. A cross join materialises its right-hand side once on the assumption it will be read exactly once after that. An asof join’s sort/merge step isn’t idempotent. A sort under maintain_order=True can carry state that doesn’t reset between runs. A real file scan (scan_parquet, scan_csv, scan_ipc) owns a streaming file cursor that isn’t rewound between calls, because normally there’d only be one call. None of this is a rtlf bug to fix; it’s what “an executor tree” means when nothing in its design anticipated being asked to run twice.

The README turns this into an actual table rather than a paragraph of caveats, and rtlf’s own test suite turns the table into a test:

# tests/test_basic.py:79-99
@pytest.mark.xfail(
    reason="CompiledRealtimeLazyFrame reuses the physical executor tree; polars' "
           "union/concat executor is stateful and cannot be re-executed. Use "
           "RealtimeLazyFrame for queries involving concat/union.",
    strict=False,
)
def test_two_placeholders_compiled():
    schema = pl.Schema({"v": pl.Int32})
    lf_a = placeholder("a", schema)
    lf_b = placeholder("b", schema)
    compiled = rtlf(pl.concat([lf_a, lf_b])).compile()

    r1 = compiled.collect({"a": pl.DataFrame({"v": [1]}), "b": pl.DataFrame({"v": [2]})})
    r2 = compiled.collect({"a": pl.DataFrame({"v": [9]}), "b": pl.DataFrame({"v": [8]})})
    assert sorted(r1["v"].to_list()) == [1, 2]
    assert sorted(r2["v"].to_list()) == [8, 9]

strict=False means this test is allowed to pass: on a future Polars version where Union’s executor happens not to be stateful, it would, and CI wouldn’t fail. Until then, it documents the exact failure mode instead of just asserting it: this specific query, compiled, breaks on the second call, not the first, because the first call is indistinguishable from the tree’s intended one-shot use. That’s a genuinely good use of xfail: not “this doesn’t work yet,” but “this can’t work by construction, here’s the proof, tell me the moment that stops being true.”

Which is the honest shape of the whole project: RealtimeLazyFrame re-derives the physical plan from a stored IR on every call, so it’s always safe: a fresh executor tree every time, just without the optimizer pass. CompiledRealtimeLazyFrame reuses the tree itself, which is where the real speedup lives, and that’s exactly the part of Polars that was never designed to be reused. The README’s caveat table isn’t a list of bugs to fix later. It’s the actual boundary of what “compile once” can mean for a system whose execution layer assumes it runs once.

Results

Time per 1000-row batch versus depth (nodes, log scale) for a linear when/then chain, comparing plain LazyFrame, RTLF and Compiled RTLF, out to depth 1000.

The same linear-chain benchmark on a log-x axis, out to the full depth-1000 range the budget-based early stop allowed. Compiled RTLF stays under 40 ms the whole way.

Time per 1000-row batch versus depth for a width-5 decision tree, log-x scale, out to depth 6.

The decision-tree benchmark: node count is exponential in depth, so the y-axis blows up by depth 6 for plain LazyFrame while Compiled RTLF barely moves.

There’s a small honest surprise buried in build_dtree_expr, too. “Width-5 tree” suggests node count grows as 5^depth, but the generator actually does this:

# benchmark.py:54-63
def build_dtree_expr(max_depth: int, width: int = 5):
    """Exponential tree from benchmark.py."""
    def _build(depth=0):
        if depth >= max_depth:
            return make_leaf()
        expr = pl
        for wi in range(width):
            expr = expr.when(pl.col(f"feat_{depth}") < wi).then(_build(depth + 1))
        return expr.otherwise(_build(depth + 1))
    return _build()

Five .when(...).then(...) branches inside the loop, plus one more .otherwise(...) after it: six recursive calls to _build(depth + 1) per level, not five. And unlike the linear chain’s build_linear_expr (which is @cache-decorated), build_dtree_expr isn’t memoised, so each of those six calls independently builds its own subtree rather than sharing one. Node count is 6^depth, not 5^depth: 55,987 individual _build() calls by depth 6, not the 15,625 a literal reading of “width-5” would suggest. The widget below gets this right: it’s ported straight from the function above, not from the name of it.

Plan Cost Explorer

Depth and shape drive four linked views: the Python benchmark.py would actually generate for that plan, the exact node count, an inline-SVG plan graph that fans out fast enough to watch, and a cursor on the replotted benchmark curves. A third control, batch size, feeds a small model of how a collect() call’s time splits between planning and execution, clearly marked as a model layered on the measured numbers, not a new measurement, since the README only ever measured batch size 1000.

InteractivePlan Cost Explorer

With JavaScript on, this becomes a shape and depth control (linear chain, or width-5 decision tree) driving four linked panels: the generated Python source, the exact node count, an SVG plan graph that visibly explodes past a couple of levels of the tree, and a cursor on the replotted README benchmark curves showing the speedup at that depth, interpolated between the README’s measured points where the reader’s chosen depth doesn’t land on one exactly, and labelled as such. A batch-size control estimates how that time would split between planning and execution at a batch size other than the 1000 rows the benchmark actually measured.

Watching one collect() call

The code above is easier to hold in your head as a diagram than as two collect() implementations read side by side, so here’s the same contrast as a step-through: the uncompiled path’s clone-and-patch against the compiled path’s slot fill, five stages, one step index driving both.

InteractiveArena vs. slot: stepping through one collect() call

With JavaScript on, Step (or Play, motion permitting) walks through both collect() implementations in lockstep: IR::Scan(placeholder) → clone arena → patch to IR::DataFrameScan (re-lifting the predicate where one was pushed down) → create_physical_plan → execute, next to the compiled path’s slot.lock() = df → execute, sitting idle through the stages the uncompiled path is still doing real work in.

The general shape of it

Nothing in rtlf is about DataFrames specifically. It’s about the fact that “figure out what to do” and “do it” are different costs, and a system that pays the first cost on every call is only correct, not fast, whenever the plan doesn’t change between calls. A query engine’s optimizer, a regex engine’s NFA-to-DFA compilation, a shader pipeline’s compile-then-dispatch split, all of them have exactly this shape, a planning phase whose cost is a function of the plan’s structure and an execution phase whose cost is a function of the data, and all of them get faster the same way: do the planning once, keep whatever it produced, and make sure what it produced can actually accept new data without being rebuilt.

That last clause is the part rtlf is honest about not having fully solved. Mode A always can: it rebuilds the physical plan every time, so there’s nothing left to go stale. Mode B can’t, for exactly the set of operations the README’s table names, because the executors underneath were never written with reuse in mind. The interesting engineering here wasn’t finding a workaround for that; it was finding the version of “compile once, execute forever” that’s honest about where the “forever” stops.