Polars Simplifies the Expression Tree Before It Runs Anything

TL;DR: Polars does not have a rule set of arithmetic identities like a + 0 -> a. What it has is a rewrite pass over the expression arena that runs during DSL-to-IR conversion, interleaved with type coercion, and that rewrites the tree to a fixed point: boolean identities (x AND true -> x), constant folding (1 + 2 -> 3), ternary folding (when(true).then(a).otherwise(b) -> a) and function-level rewrites (is_null().any() -> null_count() > 0). The cleverest part is a single bit of context: inside a filter predicate, null and false are the same value, so a rule may replace a column by a scalar there and nowhere else.

I went looking for a table of algebraic identities in crates/polars-plan/src/plans/optimizer/simplify_expr/ expecting to find a + 0 -> a, x * 1 -> x and 0 * x -> 0. They are not there, and the absence is instructive: Polars does not simplify arithmetic against identities. x + 0 is not x in Polars’ semantics in the general case (it can change the dtype of the result, and for a float column x + 0 is not bit-identical to x for -0.0), so the rule does not exist.

What does exist is a small, careful set of rewrites that are always valid, applied by a generic rewrite machinery. That machinery is the interesting part, because the same machinery has to be safe under three different contexts: rewriting an arbitrary expression, rewriting a filter predicate, and rewriting an expression whose dtype is not yet decided.

The rule abstraction

Every rewrite in Polars is a method on a type implementing OptimizationRule. The trait has two hooks: one for rewriting a plan node, one for rewriting an expression node:

rust
pub trait OptimizationRule {
    ///  Optimize (subplan) in LogicalPlan
    ///
    /// * `lp_arena` - LogicalPlan memory arena
    /// * `expr_arena` - Expression memory arena
    /// * `node` - node of the current LogicalPlan node
    fn optimize_plan(
        &mut self,
        _lp_arena: &mut Arena<IR>,
        _expr_arena: &mut Arena<AExpr>,
        _node: Node,
    ) -> PolarsResult<Option<IR>> {
        Ok(None)
    }
    fn optimize_expr(
        &mut self,
        expr_arena: &mut Arena<AExpr>,
        expr_node: Node,
        schema: &Schema,
        ctx: OptimizeExprContext,
    ) -> PolarsResult<Option<AExpr>> {
        _ = (expr_arena, expr_node, schema, ctx);
        Ok(None)
    }
}

crates/polars-plan/src/plans/optimizer/stack_opt.rs:204

The contract is deliberately minimal: return Some(new_expr) to replace the node, None to leave it alone. The rule does not recurse — the machinery does. That is what makes the rules composable: a rule only has to know how to rewrite one node shape, and the driver takes care of walking the tree and re-running the rule until nothing changes.

There are two drivers. The optimizer-phase one walks the whole plan and repeats until a fixed point:

rust
    pub fn optimize_loop(
        &self,
        rules: &mut [Box<dyn OptimizationRule>],
        expr_arena: &mut Arena<AExpr>,
        lp_arena: &mut Arena<IR>,
        lp_top: Node,
    ) -> PolarsResult<Node> {
        let mut changed = true;

        // Nodes of expressions and lp node from which the expressions are a member of.
        let mut plans = vec![];
        let mut exprs = vec![];
        let mut scratch = vec![];
        let mut children = vec![];
        let mut schema_stack: Vec<Arc<Schema>> = vec![];

        // Run loop until reaching fixed point.
        #[allow(clippy::field_reassign_with_default)]
        while changed {
            // Recurse into sub plans and expressions and apply rules.
            changed = false;

crates/polars-plan/src/plans/optimizer/stack_opt.rs:33

and the expression-level driver runs the rules in a tight loop over a single node:

rust
            if let Some(rule) = &mut self.simplify {
                while let Some(x) =
                    rule.optimize_expr(expr_arena, current_expr_node, schema, ctx)?
                {
                    expr_arena.replace(current_expr_node, x);
                }
            }
            if let Some(rule) = &mut self.coerce {
                while let Some(x) =
                    rule.optimize_expr(expr_arena, current_expr_node, schema, ctx)?
                {
                    expr_arena.replace(current_expr_node, x);
                }
            }

crates/polars-plan/src/plans/conversion/stack_opt.rs:193

That inner while let is what makes the rules rewrite rules rather than peel rules. A rule that turns x AND true into x may expose a new opportunity for the same rule one level up, and the loop keeps going until the tree stops changing. A rule that returned a node that the same rule would immediately rewrite again would loop forever, so the rules are written to be strictly decreasing in some well-founded measure — usually the number of nodes.

Where simplification actually happens

This is the part that surprised me: SimplifyExprRule is not in the list of rules the optimizer runs. The optimizer’s rule list contains FusedArithmetic, SimpleProjectionAndCollapse, DelayRechunk, FilterConstraintRule, SimplifyBooleanRule, FlattenMergeSortedRule and FlattenUnionRule — but not SimplifyExprRule.

SimplifyExprRule runs earlier, during conversion from the DSL to the IR, in the ConversionOptimizer:

rust
/// Applies expression simplification and type coercion during conversion to IR.
pub struct ConversionOptimizer {
    /// (expression, schema index, is filter predicate if in a filter)
    scratch: Vec<(Node, usize, bool)>,
    schemas: Vec<Schema>,

    simplify: Option<SimplifyExprRule>,
    coerce: Option<TypeCoercionRule>,
    check: Option<TypeCheckRule>,

crates/polars-plan/src/plans/conversion/stack_opt.rs:13

The order in the loop is simplify, then coerce — for the same node, over and over, until both are stable. That ordering is the whole design in one line: simplify the expression into its simplest form, then insert the casts that make it type-correct, then simplify again (a cast can create a new simplification opportunity), then coerce again.

The reason for doing it here rather than in the optimizer is that the DSL expressions are the ones users write, and the DSL is where the redundancy lives. By the time the plan is an IR, most of it is already gone.

The context bit that makes boolean rules safe

The single most important thing a simplification rule needs to know is whether the expression it is rewriting is a filter predicate, because in that position null and false are interchangeable and a column may be replaced by a scalar. Under Kleene logic, a filter keeps a row iff its predicate is true; a & b is true iff both are; a | b is true iff either is. So in a filter, x AND true really is x — including when x is null, because null AND true is null, and null and false both mean “drop the row”.

That is not true anywhere else. x AND true under a NOT, or as an argument to is_null, or inside an aggregation, has different null behaviour, and replacing a column with a scalar there would change the length of the result.

Polars encodes this in the context:

rust
#[derive(Default, Clone, Copy)]
pub struct OptimizeExprContext {
    pub in_pyarrow_scan: bool,
    pub in_io_plugin: bool,
    /// The expression is the predicate of a filter, or is reachable from it only through AND/OR.
    ///
    /// Such an expression only matters through whether it evaluates to `true`: a filter keeps a
    /// row iff its predicate is `true`, and under Kleene logic `a & b` is `true` iff both `a` and
    /// `b` are, and `a | b` iff either is. So rewrites may treat `null` and `false` as the same
    /// value here and may replace a column by a broadcasting scalar. Neither is valid for any
    /// other sub-expression of a filter, e.g. under `NOT`, `XOR`, `is_null` or an aggregation.
    pub is_filter_predicate: bool,
    pub has_inputs: bool,
}

crates/polars-plan/src/plans/optimizer/stack_opt.rs:157

and the flag is propagated down the tree only through AND and OR:

rust
/// Whether the inputs of `expr` are filter predicates if `expr` is one, see
/// [`OptimizeExprContext::is_filter_predicate`].
pub fn passes_on_filter_predicate(expr: &AExpr) -> bool {
    matches!(
        expr,
        AExpr::BinaryExpr {
            op: Operator::And | Operator::Or,
            ..
        }
    )
}

crates/polars-plan/src/plans/optimizer/stack_opt.rs:173

So a AND b AND c in a filter: all three are filter predicates. NOT(a AND b) in a filter: a and b are not filter predicates, because the NOT in between breaks the equivalence. This is the kind of detail that a rewrite engine gets wrong, and it is why the flag is computed by the driver rather than guessed by each rule.

The rules themselves

Boolean identities

The And and Or cases live in crates/polars-plan/src/plans/optimizer/simplify_expr/arity.rs:

rust
        O::And => {
            let left_ae = expr_arena.get(left);
            let right_ae = expr_arena.get(right);

            // true AND x => x
            if matches!(
                left_ae,
                AExpr::Literal(lv) if lv.bool() == Some(true)
            ) && is_filter_predicate
            {
                // Only in filter as we might change the name from "literal"
                // to whatever lhs columns is.
                return Some(right_ae.clone());
            }

            // x AND true => x
            if matches!(
                right_ae,
                AExpr::Literal(lv) if lv.bool() == Some(true)
            ) {
                return Some(left_ae.clone());
            }

            // x AND false -> false
            // TODO: we need an optimizer redesign to allow x & false to be optimized
            // in general as we can forget the length of a series otherwise.
            // In filter we allow it as the length is not important there.
            if (is_scalar_ae(left, expr_arena) | is_filter_predicate)
                && matches!(
                    right_ae,
                    AExpr::Literal(lv) if lv.bool() == Some(false)
                )
            {
                return Some(AExpr::Literal(Scalar::from(false).into()));
            }

crates/polars-plan/src/plans/optimizer/simplify_expr/arity.rs:44

Look at the asymmetry: x AND true -> x is always valid, because AND with true propagates the left operand’s value and length. true AND x -> x is only valid in a filter, because true AND x would be a scalar true and x a column, and outside a filter the length matters.

The TODO in the third rule is worth reading twice: x AND false -> false is not valid in general, because the result would be a scalar false where the caller expected a series of the same length as x. The rule is allowed where the length does not matter — a filter, or when the left side is already a scalar.

The same file also has two rules that are pure wins. A AND NOT(A) is always false, and two comparisons on the same operands that can never both hold (col > 5 AND col <= 5) are always false:

rust
                // `A AND NOT(A)` is always false (either operand may be the
                // `NOT`). Catches cases the earlier rewrites left alone, like
                // `NOT(col)` on a bool column or `NOT(is_in(col, values))`.
                if is_self_negation(left, left_ae, right_ae, expr_arena, maintain_errors)
                    || is_self_negation(right, right_ae, left_ae, expr_arena, maintain_errors)
                {
                    return Some(AExpr::Literal(Scalar::from(false).into()));
                }

                // Two comparisons on the same operands that can never both hold,
                // e.g. `(col > N) AND (col <= N)` or `(col > N) AND (col < N)`.

crates/polars-plan/src/plans/optimizer/simplify_expr/arity.rs:59

The contradiction check is done with a three-bit encoding of the comparison operators — a comparison is true in some subset of {<, ==, >} — and two comparisons contradict when their subsets are disjoint:

rust
// Bitset positions for the three outcomes of comparing two values.
const CMP_LT: u8 = 1; // x < y
const CMP_EQ: u8 = 2; // x == y
const CMP_GT: u8 = 4; // x > y

// The cases a comparison operator is true in, as a bitset over `<` / `==` / `>`.
// `None` for non-comparison operators.
fn comparison_cases(op: Operator) -> Option<u8> {
    use Operator::*;
    Some(match op {
        Lt => CMP_LT,
        LtEq => CMP_LT | CMP_EQ,
        Gt => CMP_GT,
        GtEq => CMP_GT | CMP_EQ,
        Eq => CMP_EQ,
        NotEq => CMP_LT | CMP_GT,
        EqValidity | NotEqValidity | And | Or | Xor | LogicalAnd | LogicalOr | Plus | Minus
        | Multiply | RustDivide | TrueDivide | FloorDivide | TruncDivide | Modulus => {
            return None;
        },
    })
}

// Whether `(x op1 y) AND (x op2 y)` can never both hold: the operators share no
// case (their true-sets over `<` / `==` / `>` are disjoint). `false` if either
// is not a comparison.
fn comparisons_contradict(op1: Operator, op2: Operator) -> bool {
    match (comparison_cases(op1), comparison_cases(op2)) {
        (Some(a), Some(b)) => a & b == 0,
        _ => false,
    }
}

crates/polars-plan/src/plans/optimizer/simplify_expr/arity.rs:225

Constant folding

Constant folding lives in SimplifyExprRule::optimize_expr and is table-driven by a macro that dispatches on the runtime type of the two literals:

rust
                // lit(left) + lit(right) => lit(left + right)
                use Operator::*;
                #[allow(clippy::manual_map)]
                let out = match op {
                    Plus => {
                        match eval_binary_same_type!(left_aexpr, right_aexpr, |l, r| l + r) {
                            Some(new) => Some(new),
                            None => {

crates/polars-plan/src/plans/optimizer/simplify_expr/mod.rs:870

The macro eval_binary_same_type! only folds when both sides have the same runtime type, and it returns the result as the same type — that is what the “same_type” in the name means, and it is why i64 + f64 does not fold to an f64 behind your back:

rust
                        (AnyValue::Int64($l), AnyValue::Int64($r)) => {
                            Some(AExpr::Literal(<Scalar as From<i64>>::from($ret).into()))
                        },

crates/polars-plan/src/plans/optimizer/simplify_expr/mod.rs:334

The comparisons fold the same way, and the float comparisons use total ordering so that NaN folds deterministically rather than randomly:

rust
                (AnyValue::Float64(l), AnyValue::Float64(r)) => Some(AExpr::Literal({ let x: bool = l.to_total_ord() $operand r.to_total_ord(); Scalar::from(x) }.into())),

crates/polars-plan/src/plans/optimizer/simplify_expr/mod.rs:395

Ternary folding

when(...).then(...).otherwise(...) with a constant predicate collapses to one of its branches, but only when the lengths work out:

rust
pub(super) fn simplify_ternary(
    predicate: Node,
    truthy: Node,
    falsy: Node,
    expr_arena: &mut Arena<AExpr>,
) -> Option<AExpr> {
    let predicate = expr_arena.get(predicate);

    if let AExpr::Literal(lv) = predicate
        && let Some(predicate) = lv.bool()
    {
        // Only replace if both are scalar or both are not scalar and are the same length,
        // the latter is tested by checking if they are elementwise.
        let t_is_scalar = is_scalar_ae(truthy, expr_arena);
        let f_is_scalar = is_scalar_ae(falsy, expr_arena);

        if t_is_scalar == f_is_scalar
            && (t_is_scalar
                || (is_elementwise_rec(truthy, expr_arena)
                    && is_elementwise_rec(falsy, expr_arena)))
        {
            let branch = if predicate { truthy } else { falsy };
            return Some(expr_arena.get(branch).clone());
        }
    }

    None
}

crates/polars-plan/src/plans/optimizer/simplify_expr/arity.rs:169

The t_is_scalar == f_is_scalar test is the length check: you cannot replace a ternary whose branches have different lengths with one of them, because the result would no longer be elementwise.

Function-level rewrites

crates/polars-plan/src/plans/optimizer/simplify_expr/simplify_functions.rs holds the rules that change the shape of a function call rather than folding a constant:

rust
        // is_null().any() -> null_count() > 0
        // is_not_null().any() -> ~is_empty(ignore_nulls=True)
        // CORRECTNESS: we can ignore 'ignore_nulls' since is_null/is_not_null never produces NULLS
        IRFunctionExpr::Boolean(IRBooleanFunction::Any { ignore_nulls: _ }) => {
            let input_node = expr_arena.get(input[0].node());
            match input_node {
                AExpr::Function {
                    input,
                    function: IRFunctionExpr::Boolean(IRBooleanFunction::IsNull),
                    options: _,
                } => Some(AExpr::BinaryExpr {
                    left: expr_arena.add(new_null_count(input)),
                    op: Operator::Gt,
                    right: expr_arena.add(AExpr::Literal(LiteralValue::new_idxsize(0))),
                }),

crates/polars-plan/src/plans/optimizer/simplify_expr/simplify_functions.rs:14

is_null().any() becomes null_count() > 0, which is a single pass over the column instead of a boolean mask followed by an any-reduction. And:

rust
        // sort().reverse() -> sort(reverse)
        // sort_by().reverse() -> sort_by(reverse)
        IRFunctionExpr::Reverse => {
            let input = expr_arena.get(input[0].node());
            match input {
                AExpr::Sort { expr, options } if !options.maintain_order => {
                    let mut options = *options;
                    options.descending = !options.descending;
                    options.nulls_last = !options.nulls_last;
                    Some(AExpr::Sort {
                        expr: *expr,
                        options,
                    })
                },

crates/polars-plan/src/plans/optimizer/simplify_expr/simplify_functions.rs:85

Reversing a sort is not a separate operator that has to materialise the sorted array and then walk it backwards; it is the same sort with the descending and nulls_last flags flipped. Note the !options.maintain_order guard: with maintain_order set, the sort is stable, and reversing a stable sort is not the same as an unstable sort with flipped flags.

There is also a de Morgan rewrite:

rust
                // not(a and b) => not(a) or not(b)
                AExpr::BinaryExpr {
                    left,
                    op: Operator::And | Operator::LogicalAnd,
                    right,
                } => {

crates/polars-plan/src/plans/optimizer/simplify_expr/simplify_functions.rs:170

which matters because the whole boolean-simplification machinery above operates on AND/OR chains: pushing a NOT down to the leaves makes more of those chains visible.

From the simplified tree to running code

The simplified IR expression is not what executes. Two more steps happen afterwards.

Type coercion. TypeCoercionRule runs in the same loop as the simplifier, immediately after it, and inserts the casts that make the expression type-correct. Because both rules loop until stable, a simplification that changes a dtype (say, is_null().any() becoming null_count() > 0) feeds back into the coercion, and the coercion’s inserted casts feed back into the simplifier.

Physical lowering. At execution-planning time the IR expression tree is turned into a tree of PhysicalExpr, which is what actually gets called per chunk:

rust
pub fn create_physical_expr(
    expr_ir: &ExprIR,
    expr_arena: &mut Arena<AExpr>,
    schema: &SchemaRef, // Schema of the input.
    state: &mut ExpressionConversionState,
) -> PolarsResult<Arc<dyn PhysicalExpr>> {
    let phys_expr = create_physical_expr_inner(expr_ir.node(), expr_arena, schema, state)?;

    if let Some(name) = expr_ir.get_alias() {
        Ok(Arc::new(AliasExpr::new(
            phys_expr,
            name.clone(),
            node_to_expr(expr_ir.node(), expr_arena),
        )))
    } else {
        Ok(phys_expr)
    }
}

crates/polars-expr/src/planner.rs:135

The simplifier never has to know anything about the physical layer. It rewrites the logical expression, and the physical lowering just walks whatever tree it is handed.

The bottom line

The trade-offs of this design:

The rules are small and always-valid, not clever. There is no x * 1 -> x, no 0 * x -> 0, no reassociation of a + b + c. Each rule that exists is one where the rewrite is provably value-preserving in every context it is enabled in, and the rules that need a context gate get one (is_filter_predicate, is_scalar_ae, is_elementwise_rec).
The context is the hard part, not the rules. A rewrite engine that does not know whether it is inside a filter, under a NOT, or inside an aggregation will produce wrong answers on nulls. Polars computes that information once, in the driver, and hands it to every rule.
The pass is a fixed point, not a single sweep. Rules are applied until the tree stops changing, which is what lets a rule expose work for another rule. The cost is that the pass can visit the same node many times; the mitigation is that rules are written to strictly decrease the size of the tree they produce.
Simplification runs before type coercion, in the same loop. This is why simplification lives in the conversion stage rather than the optimizer: at conversion time the expressions are still the ones the user wrote, before casts have been sprinkled over them.

The order of operations, end to end:

text

  DSL expression
        |
        |  ConversionOptimizer
        |    loop { simplify; coerce } until stable
        v
  IR expression (AExpr in an arena)
        |
        |  optimizer passes: predicate pushdown, projection pushdown, ...
        v
  optimized IR
        |
        |  create_physical_expr  (crates/polars-expr/src/planner.rs)
        v
  PhysicalExpr tree  ->  executed per chunk

The practical takeaway: when you look at a Polars plan and see an expression that looks dumber than what you wrote — a null_count() where you wrote is_null(), a sort with flipped flags where you wrote .reverse() — that is the simplifier, and it ran before the plan was even finished being built.