Polars Reorders Work to Avoid Running Out of Memory

TL;DR: Polars’ streaming engine does not execute a plan in plan order. It executes it in phases, and inside a phase it decides which node to run by looking at how much memory that node is holding onto. At most one memory-hungry node runs per phase, and among the candidates the scheduler prefers the one whose output can be consumed immediately, then the one that frees the most held data, then the one with the most unbounded work downstream.
If reordering is not enough, the out-of-core layer spills to disk once the process-wide allocation counter crosses a budget, and a global allocator aborts the process before the kernel’s OOM killer does. Where the engine can trade CPU for memory, it does: a select is split with a minimum cut so that wide intermediates are recomputed rather than stored.

I went looking for a “reorder the plan so we do not run out of memory” pass in the optimizer. There is not one, at least not under that name. What exists is much more interesting: the streaming engine treats execution order as a scheduling problem with memory as the objective function.

The problem with a fixed order

A streaming plan is a graph, not a list. Some nodes are cheap filters that pass data through; others are sort, group_by, join — nodes that must see all of their input before they can emit anything, and that therefore hold their input in memory while they work.

If you run such a graph in plan order, the classic failure is:

text

scan ──► filter ──► sort ──────────► join ──► sink
                    └─► group_by ─────────────┘

Run sort first and it buffers the entire dataset. Run group_by first and it builds a hash table over the entire dataset. Meanwhile the other branch could have been streaming through, keeping only a fraction of the data alive at any time. The peak memory of the query is not a property of any single node; it is a property of the order in which nodes are run.

So the engine needs to know, for each node, how much memory it is currently holding, and it needs to be able to change its mind about the order.

The graph

The streaming engine’s graph is two slot maps — nodes and pipes:

rust
// crates/polars-stream/src/graph.rs:23
/// Represents the compute graph.
///
/// The `nodes` perform computation and the `pipes` form the connections between nodes
/// that data is sent through.
#[derive(Default)]
pub struct Graph {
    pub nodes: SlotMap<GraphNodeKey, GraphNode>,
    pub pipes: SlotMap<LogicalPipeKey, LogicalPipe>,
}

rust
// crates/polars-stream/src/graph.rs:187
pub struct GraphNode {
    pub compute: Box<dyn ComputeNode>,
    pub inputs: Vec<LogicalPipeKey>,
    pub outputs: Vec<LogicalPipeKey>,
}

Every pipe has a state on each side, and there are exactly three states:

rust
// crates/polars-stream/src/graph.rs:212
#[derive(Copy, Clone, PartialEq, Eq, Debug, PartialOrd, Ord)]
pub enum PortState {
    Blocked,
    Ready,
    Done,
}

A node can only make progress when its input pipes are Ready (data available) and its output pipes are not Blocked (the downstream is not ready to take more). This is a classic dataflow readiness rule, and it is what makes re-ordering possible in the first place: a node’s eligibility depends on the state of the graph, not on its position in the plan.

Classifying how much memory a node holds

Each node reports its memory behaviour through one small enum:

rust
// crates/polars-stream/src/nodes/mod.rs:83
/// How a node's memory use relates to its input size, given its current state.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NodeMemoryUsage {
    /// Memory does not meaningfully grow with the input size, now or in a later
    /// state.
    Bounded,
    /// Memory may grow with the input size while data streams through the node.
    Unbounded,

    /// Stores nothing meaningful yet, but will be `Accumulating` in a later state.
    WillAccumulate,
    /// Stores (substantial parts of) its input while running. Each phase runs
    /// at most one such node, unless all such nodes remaining are sinks.
    Accumulating,
    /// Holds significant stored data that may be freed once this node runs to completion.
    HoldingUntilDone,
    /// Holds significant stored data that may be freed immediately while this node runs.
    Draining,
}

This is a per-node classification, not a per-node budget. There is no “this node may use 200 MB”; the classification is qualitative, and the scheduler uses it to decide ordering. The distinction between HoldingUntilDone and Draining is the interesting one: both hold a lot, but the first can only release when it finishes, while the second releases as it runs. A Draining node is therefore a priority — run it and you get your memory back sooner.

Nodes implement it themselves. A sort, for instance:

rust
// crates/polars-stream/src/nodes/sort/mod.rs:370
    fn memory_usage(&self) -> NodeMemoryUsage {
        match &self.state {
            SortState::BufferAndSample(_) => NodeMemoryUsage::Accumulating,
            SortState::Flush(_) => NodeMemoryUsage::Draining,
            SortState::Source(src) => src.memory_usage(),
            SortState::Done => NodeMemoryUsage::Bounded,

The phase loop

Execution is a loop over phases. Each phase picks a subgraph to run, runs it to completion, and then re-evaluates the graph state. The subgraph selection is where the reordering happens:

rust
// crates/polars-stream/src/execute.rs:116
/// Finds a part of the graph which we can run.
fn find_runnable_subgraph(graph: &mut Graph) -> (PlHashSet<GraphNodeKey>, Vec<LogicalPipeKey>) {
    // Find phase sinks, choose a subset with at most one accumulating node, and
    // return the subgraph needed to feed them.
    let phase_sinks = find_runnable_phase_sinks(graph);
    let (accumulating, mut to_run): (Vec<_>, Vec<_>) = phase_sinks
        .into_iter()
        .partition(|n| graph.nodes[*n].compute.memory_usage() == NodeMemoryUsage::Accumulating);

    // If all accumulating nodes left are sinks (InMemorySink), we're not gaining
    // anything by only running a subset.
    let only_accumulating_sinks_left = accumulating
        .iter()
        .all(|node_key| graph.nodes[*node_key].outputs.is_empty());

    if only_accumulating_sinks_left {
        to_run.extend(accumulating);
    } else {
        to_run.push(choose_accumulating_node(graph, &accumulating));
    }

    expand_ready_subgraph(graph, to_run)
}

Read the comment carefully: choose a subset with at most one accumulating node. An Accumulating node is one that stores its input while it runs — a sort, a group-by, a join build side. Running two of them at once means two half-full hash tables in memory at the same time. The engine refuses: at most one per phase, so the peak is one accumulator plus whatever the cheap nodes need.

The exception is worth noting: if every remaining accumulator is a sink (the final in-memory result), there is nothing to be gained by serializing them, so they all run together.

Choosing which accumulator to run

When there is more than one candidate accumulator, the tie is broken by a scoring function that is, in effect, a memory-aware topological sort:

rust
// crates/polars-stream/src/execution_order.rs:13
/// Chooses which of the given accumulating nodes to run, in order of preference:
/// 1. nodes whose outputs are ready to be consumed, so the result is not held;
/// 2. nodes that read from the most nodes holding data (see [`num_held_upstream`]), so that data
///    can be released;
/// 3. nodes with the most non-bounded nodes downstream (see [`num_non_bounded_downstream`]), so
///    a short branch does not hold its result while a long branch is processed;
/// 4. nodes with more outputs, so in-memory sinks (which have none) run last.
pub fn choose_accumulating_node(graph: &Graph, candidates: &[GraphNodeKey]) -> GraphNodeKey {
    let num_held = num_held_upstream(graph, candidates);
    let num_non_bounded = num_non_bounded_downstream(graph, candidates);
    *candidates
        .iter()
        .max_by_key(|node_key| {
            let node = &graph.nodes[**node_key];
            let num_outputs_ready_to_recv = node
                .outputs
                .iter()
                .filter(|o| graph.pipes[**o].recv_state == PortState::Ready)
                .count();
            (
                num_outputs_ready_to_recv,
                num_held[**node_key],
                num_non_bounded[**node_key],
                node.outputs.len(),
            )
        })
        .unwrap()
}

Four rules, in priority order. Rule 1 is the anti-footgun rule: if the downstream is ready to consume, run this node now and hand the result over immediately, rather than holding the finished result in memory. Rule 2 is the memory rule proper — count how many holding nodes feed into this one, and prefer the node that, when it finishes, lets the most data go. Rule 3 keeps long branches moving: if a short branch would finish and sit on its result while a long branch is still working, prefer to start the long branch first. Rule 4 is a tiebreak that pushes final sinks last.

The two counting helpers are small recursive graph walks:

rust
// crates/polars-stream/src/execution_order.rs:39
/// For each of the given nodes and the nodes feeding them through pipes ready on both sides, the
/// number of [`NodeMemoryUsage::HoldingUntilDone`] or [`NodeMemoryUsage::Draining`] nodes among
/// itself and its inputs through such pipes, recursively. A node reachable through several such
/// paths is counted once per path.
fn num_held_upstream(graph: &Graph, nodes: &[GraphNodeKey]) -> SecondaryMap<GraphNodeKey, usize> {

rust
// crates/polars-stream/src/execution_order.rs:45
/// For each of the given nodes and the nodes downstream of them, the largest number of nodes
/// that are not [`NodeMemoryUsage::Bounded`] on a path from that node (inclusive) to a node
/// without outputs.
fn num_non_bounded_downstream(
    graph: &Graph,
    nodes: &[GraphNodeKey],
) -> SecondaryMap<GraphNodeKey, usize> {

Note that both walk only ready pipes. A node that is holding data but is currently blocked cannot contribute to freeing memory right now, so it does not count.

Inside a phase: topological spawn

Once the phase’s subgraph is chosen, the nodes are spawned in topological order — but backwards, from the sinks:

rust
// crates/polars-stream/src/execute.rs:163
    // We do a topological sort of the graph: we want to spawn each node,
    // starting with the sinks and moving backwards. This order is important
    // for the initialization of physical pipes - the receive port must be
    // initialized first.

The reason is mechanical: a physical pipe’s receive port must exist before the sender can use it, and the pipe is created between two nodes. Spawning sinks-first means every pipe is half-initialized by the time its other end is spawned. The sort is a simple count of not-yet-ready outputs:

rust
// crates/polars-stream/src/execute.rs:176
    let mut ready = Vec::new();
    let mut num_send_ports_not_yet_ready = SecondaryMap::new();
    for node_key in nodes {
        let node = &graph.nodes[*node_key];
        let num_outputs_in_subgraph = node
            .outputs
            .iter()
            .filter(|o| physical_pipes.contains_key(**o))
            .count();
        num_send_ports_not_yet_ready.insert(*node_key, num_outputs_in_subgraph);
        if num_outputs_in_subgraph == 0 {
            ready.push(*node_key);
        }
    }

So there are two levels of ordering: a global one that decides which memory-hungry node goes next, and a local one that makes the spawn valid.

The last resort: spill to disk

Reordering reduces peak memory; it cannot eliminate it. A single sort of a dataset that does not fit in RAM still does not fit in RAM. For that, Polars has an out-of-core layer, polars-ooc, with a global memory budget.

The budget is a fraction of total memory, capped by an absolute value:

rust
// crates/polars-config/src/lib.rs:384
        let budget_frac = f64::from_bits(self.ooc_memory_budget_fraction.load(Ordering::Relaxed));
        let bytes = u64::min(
            self.ooc_memory_budget_bytes.load(Ordering::Relaxed),
            (total_memory() as f64 * budget_frac) as u64,
        );
        self.ooc_effective_memory_budget_bytes
            .store(bytes, Ordering::Relaxed);

The decision to spill is made against a process-wide allocation counter, not against any per-node accounting:

rust
// crates/polars-ooc/src/memory_manager.rs:64
    fn should_spill(&self) -> bool {
        let usage = crate::estimate_memory_usage();
        let likely_dealt_with = self.est_spill_in_progress.load(Ordering::Relaxed);
        let likely_incoming = self.est_prefetch_in_progress.load(Ordering::Relaxed);
        (usage + likely_incoming).saturating_sub(likely_dealt_with)
            > config().ooc_memory_budget_bytes()
    }

and it is checked at explicit points in the hot loops of the nodes that accumulate — the group-by, the window nodes, the hash-shuffle batches:

rust
// crates/polars-stream/src/nodes/group_by/mod.rs:305
            for_each_parallel(
                partitions.iter_mut(),
                state.task_metrics(),
                &|(p, table, _)| async move {
                    for (pre_agg, groups) in pre_aggs.in_partition(*p) {
                        table.add_pre_agg(pre_agg, groups)?;
                    }
                    memory_manager().spill().await;
                    Ok(())
                },
            )?;

spill() is cheap when there is no pressure — it is a single comparison against the budget — and expensive only when it actually has to move data to disk.

The allocator that refuses to be OOM-killed

The counter that drives all of this is maintained by a custom global allocator, which is also where the hard limit lives:

rust
// crates/polars-ooc/src/global_alloc.rs:58
#[cold]
#[inline(never)]
fn flush_local_drift() {
    LOCAL_ALLOC_DRIFT.with(|drift| {
        let drift_bytes = drift.get();
        let old = GLOBAL_ALLOC_SIZE.fetch_add(drift_bytes as u64, Ordering::AcqRel);
        drift.set(0);

        let total = old.wrapping_add(drift_bytes as u64);
        if total <= i64::MAX as u64 && total > polars_config::get_oomkill_threshold_bytes() {
            oomkill();
        }
    })
}

#[cold]
#[inline(never)]
fn oomkill() -> ! {
    // Can't allocate here, so write directly to stderr.
    const MSG: &[u8] = b"polars: memory usage exceeded POLARS_OOMKILL_THRESHOLD_MB or POLARS_OOMKILL_THRESHOLD_FRACTION, aborting\n";
    unsafe {
        libc::write(2, MSG.as_ptr().cast(), MSG.len() as _);
        libc::abort()
    }
}

Per-thread allocation deltas are accumulated in a thread-local and only flushed to the global counter when they exceed a drift threshold, so the common case is one relaxed atomic add per allocation and no contention. The flush is also where the abort happens: rather than letting the Linux OOM killer pick a victim (which is often the wrong process, and gives you no diagnostics), Polars aborts itself with a message, once the allocation total crosses POLARS_OOMKILL_THRESHOLD_MB or its fractional counterpart.

Recompute instead of materialize

The last technique is the most interesting one, because it trades CPU for memory at plan time, and does it optimally.

Consider a group-by whose input is select(a, b, a + b). The naive plan computes a + b, stores it alongside a and b, and the group-by node aggregates over three columns. But the group-by only needs a + b as a reduction input, and it can recompute it from a and b inside the node, over the rows it actually needs. Storing it would mean the intermediate payload of the group-by is wider — and that payload is exactly what gets spilled to disk when memory runs out.

So the streaming engine splits the select into a pre-select (what is materialized) and a post-select (what is recomputed downstream), choosing the split to minimize the width of the intermediate:

rust
// crates/polars-stream/src/physical_plan/split_select.rs:266
/// Splits `exprs` into a pre-select and a post-select.
///
/// Selecting the returned pre-select on `input_schema` and then evaluating the returned
/// post-select expressions on that result, in order and each seeing the outputs of those
/// before it, produces exactly what selecting `exprs` on `input_schema` would, and:
///
/// * every post-select expression is elementwise,
/// * every expression in `must_preselect` is an output of the pre-select, under its own
///   output name,
/// * the expected bytes per row of the intermediate schema is minimal, with ties going
///   to materializing input columns rather than computed values,
/// * a post-select expression refers to an earlier post-select expression it contains by
///   its output name instead of computing it again.

The “minimal” claim is not a heuristic. The split is posed as a minimum-cut problem:

rust
// crates/polars-stream/src/physical_plan/split_select.rs:331
    // Minimize the intermediate width as a minimum cut. Per node we have two booleans:
    //
    //   post(n)  the node's operation is applied in the post-select
    //   avail(n) the node's value is available to the post-select
    //
    // with `avail(n) && !post(n)` meaning `n` is materialized, at cost `weight(n)`. The
    // only rule tying nodes together is that a node computed in the post-select needs its
    // inputs available there. Nodes are otherwise free to disagree, because an input of a
    // materialized node is recomputed inside the pre-select rather than passed through
    // the intermediate.
    //
    // Assigning "source side" to `true`, each cost is of the form
    // `w * [x == true && y == false]`, which is exactly an edge `x -> y` of capacity `w`,
    // so the minimum cut is the minimum intermediate width.

The graph is built with one node per expression sub-term, weights estimated from the output dtypes (a String column is assumed to cost 128 bits plus a 16-byte payload; a List is 64 bits plus 4 times its inner width), and the minimum cut is computed with Dinic’s algorithm:

rust
// crates/polars-stream/src/physical_plan/split_select.rs:356
    let post = |idx: usize| 2 * idx;
    let avail = |idx: usize| 2 * idx + 1;
    let src = 2 * dag.len();
    let dst = 2 * dag.len() + 1;

    // The weights are scaled so that a tie in width goes to the cut that materializes the
    // fewest computed values.
    let scale = dag.len() as u64 + 1;
    let mut network = FlowNetwork::new(2 * dag.len() + 2);
    for (idx, dag_node) in dag.iter().enumerate() {
        let is_computed = !matches!(expr_arena.get(dag_node.node), AExpr::Column(_));
        let cost = dag_node.weight * scale + (dag_node.weight > 0 && is_computed) as u64;
        network.add_edge(avail(idx), post(idx), cost);
        if dag_node.splittable {
            for input in &dag_node.inputs {
                network.add_edge(post(idx), avail(*input), INF);
            }
        } else {
            network.add_edge(post(idx), dst, INF);
        }
    }

The scale factor is a tiebreak: among cuts of equal width, prefer the one that materializes the fewest computed values, i.e. the one that recomputes more and stores less.

And here is the caller, with the reasoning stated in a comment:

rust
// crates/polars-stream/src/physical_plan/lower_group_by.rs:928
    // Keep the pre-select narrow: whatever is cheaper to recompute inside the group-by
    // node than to store in its (spilled) payload comes back as a post-select expression,
    // which the node evaluates over the rows that need it.
    let (mut pre_select_exprs, post_select_agg_inputs) = split_pre_post_select_minsize_elementwise(
        &splittable_agg_inputs,
        &must_preselect,
        &input_schema,
        expr_arena,
    )?;

The bottom line

There is no single “OOM protection policy” object in this engine. Memory safety is spread across four cooperating mechanisms, and it is worth seeing them as a ladder:

Reorder. Nodes are classified by how much they hold, and the scheduler picks the next accumulator by how much memory it will release. At most one accumulator per phase.
Narrow. Where a plan materializes an intermediate, the engine computes the minimum-width version of it with a min-cut, preferring to recompute a derived column instead of storing it.
Spill. Past a global budget, the out-of-core layer writes cold data to disk, and the group-by node pre-aggregates hot keys so that what does get spilled is as small as possible.
Abort. A global allocator counts every byte and aborts with a message before the kernel OOM killer does.

The trade-offs are real. Reordering costs a graph walk per phase. The min-cut is a polynomial-time computation at planning time, on a DAG of expression sub-terms. Spilling costs I/O and a re-read. And all of it is global rather than per-node: the budget is a process-wide byte count, and the scheduler’s notion of “holding” is a qualitative classification, not a measurement. What you get in exchange is that a query’s peak memory is a scheduling decision rather than a property of the plan you happened to write.