Opensource internals

Polars Evaluates Expressions in an Arena

11 min read Jay Hubbert

A Polars expression becomes an ExprIR handle over an Arena<AExpr>, then a tree of Arc<dyn PhysicalExpr>; here is why the middle step exists.

TL;DR: A Polars expression exists in three shapes. The one you write is a tree of Arc<Expr>. The one the optimizer reasons about is an AExpr node living in a flat Arena<AExpr>, addressed by an 8-byte Node index. The one that touches data is an Arc<dyn PhysicalExpr> tree.

The middle shape is the interesting one: because the plan is a vector plus integers, cloning it is copying indices, comparing two subtrees is comparing integers, and nothing in it is individually freed until the whole arena drops.

Every dataframe engine has to answer the same question twice: how do I describe an expression, and how do I run it? The naive answer is to use one structure for both, and to run it by walking the tree.

Polars does not. It has three distinct representations, and the transitions between them are explicit, one-way, and cheap. This post is about the middle one: the arena.

The shape you write: a tree of Arc<Expr>

The user-facing expression is an enum whose children are Arc<Expr>:

rust 
// crates/polars-plan/src/dsl/expr/mod.rs:76
pub enum Expr {
    /// Values in a `eval` context.
    ///
    /// Equivalent of `pl.element()`.
    Element,
    Alias(Arc<Expr>, PlSmallStr),
    Column(PlSmallStr),
    Selector(Selector),
    Literal(LiteralValue),
    DataTypeFunction(DataTypeFunction),
    BinaryExpr {
        left: Arc<Expr>,
        op: Operator,
        right: Arc<Expr>,
    },
    Cast {
        expr: Arc<Expr>,
        dtype: DataTypeExpr,
        options: CastOptions,
    },

This is a perfectly good API shape: it is cheap to clone (cloning an Arc is a refcount bump), it is easy to pattern match on, and it is what the Python bindings, the R bindings, and the Expr builder methods all produce.

It is also a bad shape for a plan. A tree of Arc means:

  • every node lives in its own heap allocation, scattered across the heap;
  • “the same expression” is a deep structural comparison, not a pointer comparison;
  • you cannot hand a sub-expression to another subsystem without either cloning it or keeping the Arc alive, which keeps the whole tree alive.

Polars needs all three capabilities constantly. The optimizer compares sub-expressions to find common subplans. The projection pushdown needs to answer “does this expression mention column a?”. The CSE pass needs to deduplicate identical subtrees. So the DSL tree gets converted, once, into something flat.

The arena

The whole arena is 172 lines in crates/polars-utils/src/arena.rs, and the relevant parts are tiny:

rust 
// crates/polars-utils/src/arena.rs:26
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Ord, PartialOrd)]
#[repr(transparent)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct Node(pub usize);

...

#[derive(Debug, Clone)]
#[cfg_attr(feature = "ir_serde", derive(Serialize, Deserialize))]
pub struct Arena<T> {
    version: u32,
    items: Vec<T>,
}

and adding a node is a Vec::push plus an index:

rust 
// crates/polars-utils/src/arena.rs:57
    pub fn add(&mut self, val: T) -> Node {
        let idx = self.items.len();
        self.items.push(val);
        Node(idx)
    }

Reading a node is a bounds-checked index:

rust 
// crates/polars-utils/src/arena.rs:106
    #[inline]
    pub fn get(&self, idx: Node) -> &T {
        self.items.get(idx.0).unwrap()
    }

That is the entire data structure. A Node is a usize — a Copy, Eq, Hash, Ord index into a Vec. Everything else in the design follows from that.

The IR node: AExpr

The expression IR is a single enum whose children are Nodes rather than pointers:

rust 
// crates/polars-plan/src/plans/aexpr/mod.rs:180
/// IR expression node that is allocated in an [`Arena`][polars_utils::arena::Arena].
#[derive(Clone, Debug, Default)]
#[cfg_attr(feature = "ir_serde", derive(serde::Serialize, serde::Deserialize))]
pub enum AExpr {
    /// Values in a `eval` context.
    ///
    /// Equivalent of `pl.element()`.
    Element,
    Explode {
        expr: Node,
        options: ExplodeOptions,
    },
    Column(PlSmallStr),
    ...
    BinaryExpr {
        left: Node,
        op: Operator,
        right: Node,
    },
    Cast {
        expr: Node,
        dtype: DataType,
        options: CastOptions,
    },

Note what is not here: no Box, no Arc, no dyn. The recursive edges of the expression graph are plain usizes. The graph is stored in the arena; the enum only carries the edges.

The handle: ExprIR

A plan does not hold AExprs directly. It holds ExprIRs, which pair a node index with the output name and a memoized dtype:

rust 
// crates/polars-plan/src/plans/expr_ir.rs:70
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct ExprIR {
    /// Output name of this expression.
    output_name: OutputName,
    /// Output dtype of this expression
    /// Reduced expression.
    /// This expression is pruned from `alias` and already expanded.
    node: Node,
    #[cfg_attr(feature = "serde", serde(skip))]
    output_dtype: OnceLock<DataType>,
}

ExprIR is Clone, and cloning it copies a Node (8 bytes), an OutputName, and a OnceLock<DataType>. It does not copy the expression subtree, because the subtree lives in the arena and is shared by reference. This is the “cheap cloning of expression trees” property, and it is not an optimization bolted on later — it falls out of the representation.

ExprIR also caches the dtype:

rust 
// crates/polars-plan/src/plans/expr_ir.rs:291
    pub fn dtype(&self, schema: &Schema, expr_arena: &Arena<AExpr>) -> PolarsResult<&DataType> {
        match self.output_dtype.get() {
            Some(dtype) => Ok(dtype),
            None => {
                let dtype = expr_arena
                    .get(self.node)
                    .to_dtype(&ToFieldContext::new(expr_arena, schema))?;
                let _ = self.output_dtype.set(dtype);
                Ok(self.output_dtype.get().unwrap())
            },
        }
    }

Since the arena is immutable once built, the dtype of a node is a property of the node, so every ExprIR referring to it agrees. The OnceLock turns a repeated schema query into a pointer read.

From Expr to ExprIR

The conversion is a recursive walk that bottoms out in a single arena.add:

rust 
// crates/polars-plan/src/plans/conversion/dsl_to_ir/expr_to_ir.rs:38
/// Converts expression to AExpr and adds it to the arena, which uses an arena (Vec) for allocation.
#[recursive]
pub(super) fn to_aexpr_impl(
    expr: Expr,
    ctx: &mut ExprToIRContext,
) -> PolarsResult<(Node, PlSmallStr)> {
rust 
// crates/polars-plan/src/plans/conversion/dsl_to_ir/expr_to_ir.rs:170
        Expr::BinaryExpr { left, op, right } => {
            let (l, output_name) = recurse_arc!(left)?;
            let (r, _) = recurse_arc!(right)?;
            (
                AExpr::BinaryExpr {
                    left: l,
                    op,
                    right: r,
                },
                output_name,
            )
        },
rust 
// crates/polars-plan/src/plans/conversion/dsl_to_ir/expr_to_ir.rs:681
    Ok((ctx.arena.add(v), output_name))
}

The recursion is bottom-up: children are converted first, so by the time the parent is built, its children already have indices. The parent is pushed onto the arena last. The resulting arena is in post-order, which is also the order the physical planner wants to walk.

The whole plan uses two arenas — one for plan nodes, one for expressions:

rust 
// crates/polars-plan/src/plans/conversion/dsl_to_ir/mod.rs:58
pub fn to_alp(
    lp: DslPlan,
    expr_arena: &mut Arena<AExpr>,
    lp_arena: &mut Arena<IR>,
    // Only `SIMPLIFY_EXPR`, `TYPE_COERCION`, `TYPE_CHECK`, and `PREDICATE_PUSHDOWN` are respected.
    opt_flags: &mut OptFlags,
) -> PolarsResult<Node> {

Lowering to physical expressions

Now the interesting part: the arena is walked to build the operators that actually touch data. The walk is a single match on the IR enum — this is where the plan-time dispatch happens:

rust 
// crates/polars-expr/src/planner.rs:132
#[recursive]
fn create_physical_expr_inner(
    expression: Node,
    expr_arena: &mut Arena<AExpr>,
    schema: &SchemaRef, // Schema of the input.
    state: &mut ExpressionConversionState,
) -> PolarsResult<Arc<dyn PhysicalExpr>> {
    use AExpr::*;

    let aexpr = expr_arena.get(expression);
    match aexpr.clone() {
rust 
// crates/polars-expr/src/planner.rs:301
        BinaryExpr { left, op, right } => {
            let output_field = expr_arena
                .get(expression)
                .to_field(&ToFieldContext::new(expr_arena, schema))?;
            let is_scalar = is_scalar_ae(expression, expr_arena);
            let lhs = create_physical_expr_inner(left, expr_arena, schema, state)?;
            let rhs = create_physical_expr_inner(right, expr_arena, schema, state)?;
            Ok(Arc::new(phys_expr::BinaryExpr::new(
                lhs,
                op,
                rhs,
                node_to_expr(expression, expr_arena),
                state.local.has_lit,
                state.allow_threading,
                is_scalar,
                output_field,
            )))
        },

So the dynamic dispatch that a naive engine would pay per row is paid per plan node, at planning time. The match is on a #[derive(Clone)] enum, so the compiler can, in principle, turn the whole arm dispatch into a jump table or a set of specialized branches. What survives into the hot loop is Arc<dyn PhysicalExpr> — a vtable call per operator, not per row.

Note the node_to_expr(expression, expr_arena) argument. Every physical expression carries a copy of the original DSL expression, purely so that error messages can say “this failed in col("a") + col("b")”. That is a deliberate, expensive-in-the-plan-but-free-at-runtime trade: the arena makes it possible to reconstruct the DSL tree lazily, only on the error path.

rust 
// crates/polars-plan/src/plans/conversion/ir_to_dsl.rs:7
/// converts a node from the AExpr arena to Expr
#[recursive]
pub fn node_to_expr(node: Node, expr_arena: &Arena<AExpr>) -> Expr {
    let expr = expr_arena.get(node).clone();

    match expr {
        AExpr::Element => Expr::Element,
        AExpr::Explode { expr, options } => Expr::Explode {
            input: Arc::new(node_to_expr(expr, expr_arena)),
            options,
        },

The dispatch table

Not every IR node becomes a bespoke physical operator. Most function expressions are lowered into a generic ApplyExpr that calls a closure. The mapping from IR function to closure lives in one big match in crates/polars-expr/src/dispatch/mod.rs:

rust 
// crates/polars-expr/src/dispatch/mod.rs:175
/// `input` and `expr_arena` let a function pick a variant based on its arguments, e.g. a
/// constant argument that can be prepared once.
pub fn function_expr_to_udf(
    func: IRFunctionExpr,
    input: &[ExprIR],
    expr_arena: &Arena<AExpr>,
) -> SpecialEq<Arc<dyn ColumnsUdf>> {
    use IRFunctionExpr as F;
    match func {
        // Namespaces
        #[cfg(feature = "dtype-array")]
        F::ArrayExpr(func) => array::function_expr_to_udf(func),
        F::BinaryExpr(func) => binary::function_expr_to_udf(func),
        #[cfg(feature = "dtype-categorical")]
        F::Categorical(func) => cat::function_expr_to_udf(func),
        ...
        F::Boolean(func) => boolean::function_expr_to_udf(func, input, expr_arena),

and the closures are built with a small macro that fixes the calling convention (first column is the root, the rest are arguments):

rust 
// crates/polars-expr/src/dispatch/mod.rs:58
/// `Fn(&Series, args)`
#[macro_export]
macro_rules! map {
    ($func:path) => {{
        let f = move |c: &mut [::polars_core::prelude::Column]| {
            let c = &c[0];
            $func(c)
        };

        SpecialEq::new(Arc::new(f))
    }};

That is the whole extension mechanism for user-defined and built-in functions: an IR enum variant, one arm in the dispatch table, one function taking &Column.

Where the dyn actually is

It is worth being precise about the dynamic dispatch, because the usual claim (“Polars avoids Box<dyn>”) is only half true and the half that is false matters.

The DSL layer is an enum with Arc<Expr> children — there is no Box<dyn Expr> anywhere in it. The IR layer is an enum with Node children, and it has no Box<dyn> either. The dyn appears exactly twice, and both times it is bounded by the number of plan nodes, never by the number of rows:

  • Arc<dyn PhysicalExpr>, the operator tree. One vtable call per operator per chunk of data, and the kind of operator is resolved at planning time by the match you saw above.
  • Arc<dyn ColumnsUdf>, the function layer. One vtable call per function application per chunk.

So the cost of dynamic dispatch in Polars is O(operators) per chunk, not O(rows). Compare that with the design where the expression tree itself is Box<dyn Expr> and is walked per row: there the cost is O(rows × operators). The arena is what makes it possible to have the first shape without paying the second.

The other thing the arena changes is where the type information lives. In a Box<dyn Expr> design, the type of a sub-expression is a vtable lookup at runtime. In the IR, the type of a node is a property of the enum variant, known at compile time; the only thing resolved at runtime is the data type, and even that is memoized in the ExprIR.

What the arena buys, concretely

Cheap equality and hashing. Two ExprIRs that refer to the same arena node are the same expression. The canonicalization map exploits this to assign structural equivalence classes to subtrees:

rust 
// crates/polars-plan/src/plans/aexpr/canonical.rs:23
slotmap::new_key_type! {
    /// Identifies an `AExpr` up to structural equality: two nodes get the same
    /// `CanonicalExprId` if and only if they represent the same expression.
    pub struct CanonicalExprId;
}

Cheap duplication. Duplicating a subtree is one clone of the item plus one push:

rust 
// crates/polars-utils/src/arena.rs:144
impl<T: Clone> Arena<T> {
    pub fn duplicate(&mut self, node: Node) -> Node {
        let item = self.items[node.0].clone();
        self.add(item)
    }
}

One allocation strategy. Every expression node in the whole plan lives in one Vec. No per-node allocator calls, no Box, no Arc refcount traffic while the plan is being built and rewritten. When the optimizer pushes expressions around — predicate pushdown rewriting a filter, projection pushdown splitting a select, CSE inserting a cache — it copies Nodes, not trees.

Node lifetime. The trade-off is real and worth stating plainly: an arena never frees an individual node. Arena has clear() and pop(), but there is no “free this node”. A rewritten plan leaks the expressions it replaced until the arena itself is dropped. For a query plan, whose lifetime is one collect() call, that is the right trade. It is the wrong trade for a long-lived structure that mutates one node at a time — which is exactly why Arena is not used for, say, the physical operator tree, where each operator has its own lifetime and is dropped independently.

Caching the arena across calls

Because the arena is the plan, and the plan is expensive to rebuild, a LazyFrame keeps it around:

rust 
// crates/polars-lazy/src/frame/cached_arenas.rs:9
pub(crate) struct CachedArena {
    lp_arena: Arena<IR>,
    expr_arena: Arena<AExpr>,
}

impl LazyFrame {
    pub fn set_cached_arena(&self, lp_arena: Arena<IR>, expr_arena: Arena<AExpr>) {
        let mut cached = self.cached_arena.lock().unwrap();
        *cached = Some(CachedArena {
            lp_arena,
            expr_arena,
        });
    }
rust 
// crates/polars-lazy/src/frame/cached_arenas.rs:116
    pub(super) fn get_arenas(&mut self) -> (Arena<IR>, Arena<AExpr>) {
        match self.cached_arena.lock().unwrap().as_mut() {
            Some(arenas) => (arenas.lp_arena.clone(), arenas.expr_arena.clone()),
            None => (Arena::with_capacity(16), Arena::with_capacity(16)),
        }
    }

collect_schema() on a LazyFrame converts the DSL plan to IR once, caches both arenas, and reuses them for every subsequent schema query or collect(). Since Arena is Clone and cloning a Vec<AExpr> is a memcpy of the node vector, “reuse the plan” is nearly free. Note the arena has a version counter, bumped on every clear(), so a stale cached plan can be detected rather than silently reused.

The bottom line

The arena is not a performance micro-optimization; it is a representation decision that makes a class of algorithms possible.

  • The DSL tree (Arc<Expr>) is the right shape for a user API: cheap to build, cheap to share.
  • The IR (Arena<AExpr> + Node) is the right shape for a plan: flat, index-addressed, comparable by integer, duplicable by memcpy.
  • The physical tree (Arc<dyn PhysicalExpr>) is the right shape for execution: one vtable call per operator per chunk, and the dispatch on the kind of operator happens at planning time.

The one thing to remember when reading the code: Node is not a pointer, it is an index. That single fact explains why the plan is cheap to copy, why the optimizer can rewrite it freely, why common subplan elimination works, and why every physical operator can afford to carry a copy of the original expression for error messages.

PolarsInternalsLazyPerformanceArrow