Tracking Issue: Refactor Expression Optimizer

[mhk197]

This is a tracking issue for refactoring expression optimization.

Motivation

(Logical) expression optimization currently operates on an unbound Expression and a root DType and is centered around scalar functions. For every scalar-function node, it repeatedly invokes three hooks in this order:

1. ScalarFnVTable::simplify_untyped;
2. ScalarFnVTable::simplify, using SimplifyCtx for type lookup; and
3. ScalarFnVTable::reduce, using ExpressionReduceNode.

It first optimizes the root, then its children, and then the root again. A separate find_between
pass runs afterward.

Moreover, expression constructors also call Expression::optimize before returning. At other points in the codebase, Expression::optimize_recursive is also called.

There are several problems with this design:

• Optimization policy enlarges both ScalarFnVTable and the private DynScalarFn object-safe
  interface.
• A function can provide only one implementation for each hook, so multiple independently named
  and ordered rules are awkward.
• Cross-function rules do not fit a function vtable, which is why find_between is a separate
  top-level pass.
• The optimizer operates on untyped expressions even though we have BoundExpressions now. This creates the need for a dtype cache (SimplifyCache) for performance.
• We want to optimize other types of expressions beyond ScalarFns, like incoming Higher Order Functions.
• We should optimize the entire tree only once.
• Adding generic constant folding, normalization, or common-subexpression elimination would add
  more special phases rather than extending one optimizer.

In general, it would be beneficial to have a rule-based expression optimizer that operates on a bound expression tree.

Design

BoundExpression Is the Optimizer IR

BoundExpression is already a type-checked tree whose nodes store their result dtype. It becomes
the only expression representation accepted and produced by the optimizer.

Expression
  -> bind against Scope
  -> BoundExpression
  -> optimize
  -> optimized BoundExpression

Rules read node.dtype() and child dtypes directly. Creating or rebuilding a scalar node goes
through BoundExpression::try_new, which recomputes and validates its return dtype.

This removes the need for SimplifyCtx, SimplifyCache, and simplify_untyped.

Rule Interface

The type-erased rule interface is intentionally small:

pub trait ExprRewriteRule: Debug + Send + Sync + 'static {
    fn name(&self) -> &'static str;

    fn rewrite(
        &self,
        node: &BoundExpression,
        ctx: &RewriteContext,
    ) -> VortexResult<Rewrite>;
}

pub enum Rewrite {
    Unchanged,
    Rewritten(BoundExpression),
}

RewriteContext gives a rule access to the VortexSession, optimizer configuration, and tracing
facilities. It does not expose arrays or an ExecutionCtx.

Most rules target a concrete scalar-function vtable and need typed access to its options. A typed
interface and erased adapter provide this without adding methods to DynScalarFn:

pub trait TypedExprRewriteRule<V: ScalarFnVTable>:
    Debug + Send + Sync + 'static
{
    fn name(&self) -> &'static str;

    fn rewrite(
        &self,
        options: &V::Options,
        node: &BoundExpression,
        ctx: &RewriteContext,
    ) -> VortexResult<Rewrite>;
}

The adapter checks node.as_opt::<V>(), obtains V::Options, and calls the typed rule. Rule source
files may remain next to their scalar functions; moving a rule outside the vtable does not require
moving it far from the function implementation.

Rule Registry

Expression rules live in a session-scoped ExpressionOptimizerSession, separate from
ScalarFnSession:

session
    .expression_optimizer()
    .register::<Binary>(BooleanLiteralIdentities);

session
    .expression_optimizer()
    .register::<Binary>(FindBetween);

The registry stores an ordered rule list keyed by root ScalarFnId. Rules registered for a root are
the only function-specific rules considered at that node. A short ordered fallback list supports
rules that apply to every scalar function, such as future generic constant folding.

Registration order is evaluation order. When a rule rewrites a node, evaluation restarts at the
first rule for the replacement root. The registry is snapshotted at the beginning of an optimization
so concurrent registrations cannot change a rewrite halfway through.

Plug-in initialization registers scalar functions and expression rules separately. A missing rule
may reduce performance but must never affect correctness. Replacing a scalar-function registration
does not make an old typed rule unsafe: its adapter declines nodes whose concrete vtable does not
match.

Built-in rules are installed by ExpressionOptimizerSession::default. ScalarFnSession remains
the registry for scalar-function serialization and deserialization only.

ExpressionOptimizerSession::snapshot returns an immutable ExpressionOptimizer containing the
ordered rule lists used for one optimization run. This keeps registry mutation out of the rewrite
driver and gives the public optimization API a small explicit dependency rather than the entire
session.

Rewrite Driver

The driver performs deterministic bottom-up rewriting with an explicit work stack rather than Rust
recursion:

1. Visit children in argument order.
2. Rebuild the parent if any child changed.
3. Try the rules registered for the parent root, followed by fallback rules.
4. If a rule fires, validate the replacement and enqueue the replacement subtree.
5. Revisit the parent after any replacement child reaches a fixpoint.
6. Stop when the root and every descendant are stable.

Reprocessing a replacement subtree is essential. A rule may introduce a node whose children were
not present during the first traversal.

The optimizer maintains a global rewrite budget. Exceeding it returns an error containing the most
recent rule names and expression roots. Returning a structurally identical replacement is also an
optimizer error rather than a change.

The initial driver uses first-match rule ordering. It does not enumerate alternatives or assign
costs. If Vortex later needs cost-based selection, the explicit rule and optimizer boundary provides
a place to add it without changing scalar functions.

Steps

• Bind expressions before optimizing them
• Add BoundExpressionOptimizer
• Public API stabilization

Unresolved questions

• None yet.

Implementation history