Table of Contents

• Introduction
• A Simple IR: One Big Enum
• A Simple Transformation: Join Commutativity
• A Heuristic Optimizer Framework
• A Cost-Based Optimizer Framework
• Apply the Rules Again: Transformations in the Memo Table
• A New Plan Representation

Introduction

Suppose we are going to write an SQL query optimizer. The first thing we need to decide is how to represent and store a plan in memory.

Suppose we have this query:

CREATE TABLE t1(x INT, y INT, z INT);
CREATE TABLE t2(y INT);
SELECT * FROM t1, t2 WHERE t1.y = t2.y AND t1.z = 3;

This gives us a query plan: ^{^1}

Filter #2 (t1.z) = 3
  Join #1 (t1.y) = #3 (t2.y)
    Scan t1
    Scan t2

^1: We assume that column references in plan predicates are indexes into the output columns. The join preserves the original column order when it combines the three-column left table (x, y, z) with the one-column right table (y), producing (t1.x, t1.y, t1.z, t2.y). Therefore, t1.y is #1 and t2.y is #3. We will revisit this convention later.

How do we store such query plans in Rust? Suppose we need to write a function like this:

fn get_plan() -> RelNode {
  return ...
}

How do we define RelNode, and what do we put at ... to construct the plan? That is the main question this post will answer.

A Simple IR: One Big Enum

[code for this section]

A natural first design is to use one structure for each kind of plan node. Each structure has named fields for the node’s children and properties. For example:

pub struct Scan {
    pub table: TableId,
}

pub struct Join {
    pub left: Arc<RelNode>,
    pub right: Arc<RelNode>,
    pub cond: Arc<RelNode>,
}

pub struct Filter {
    pub child: Arc<RelNode>,
    pub predicate: Arc<RelNode>,
}

We can then define one large enum representing every entity in the query language:

pub enum RelNode {
    Scan(Scan),
    Join(Join),
    Filter(Filter),
    Eq(EqPred),
    ColumnRef(ColumnRefPred),
    Const(ConstPred),
}

Helper functions can convert the types and wrap them in Arc automatically:

pub fn join(
    left: impl Into<Arc<RelNode>>,
    right: impl Into<Arc<RelNode>>,
    cond: impl Into<Arc<RelNode>>,
) -> RelNode {
    RelNode::Join(Join {
        left: left.into(),
        right: right.into(),
        cond: cond.into(),
    })
}

We can then construct the plan:

pub fn plan() -> RelNode {
    filter(
        join(
            scan(TableId(0)),
            scan(TableId(1)),
            eq_pred(column_ref_pred(1), column_ref_pred(3)),
        ),
        eq_pred(column_ref_pred(2), const_pred(3)),
    )
}

This works well for a first pass, but it has a problem that we will see shortly.

A Simple Transformation: Join Commutativity

[code for this section]

Once we have stored the plan in memory, the optimizer can begin its work. A query optimizer rewrites a plan into another plan that is potentially more efficient to execute. One example applies join commutativity: A INNER JOIN B should produce the same result as B INNER JOIN A. Let’s implement this transformation. ^{^2}

pub fn join_commute(node: Arc<RelNode>) -> Option<Arc<RelNode>> {
    if let RelNode::Join(ref a) = &*node {
        return Some(join(a.right.clone(), a.left.clone(), a.cond.clone()).into());
    }
    None
}

^2: Note that we are not considering the column order and the join condition. Otherwise, we must add a projection node to reorder the columns and rewrite the join condition.

If the join_commute transformation receives a join plan node, it returns an equivalent plan (in terms of execution results) with the join inputs swapped. Otherwise, it returns None, meaning that it cannot rewrite that plan node.

We can express such transformations in a more general form as rules. A rule takes a plan and rewrites it into another plan that produces the same result. The user provides the optimizer with a set of rules, and the optimizer decides when and where to fire them and whether the rewritten plan is better than the original.

Now we have a plan representation, some plan nodes, and a rule. We can start building a query optimizer!

A Heuristic Optimizer Framework

[code for this section]

We have defined a rule. A rule applies to a single plan node, but a plan is a tree: we cannot apply it only to the root. Instead, we must apply it recursively throughout the plan tree. In the original query, for example, the join node is a child of the filter node. We must start at the root, descend recursively, and apply the rule along the way.

A heuristic optimizer applies rules to plan nodes in a specified order. Users usually provide rules known to improve plan quality in typical cases, and the heuristic optimizer applies them whenever possible. For example, replacing a nested-loop join that has an equality condition with a hash join is usually a good choice, and pushing filters down can reduce computation. A cost-based optimizer instead uses a cost model, so users can also define transformations that do not always improve a plan.

Generally, there are two application orders for each rule: top-down and bottom-up. The following functions are illustrative; s03_heuristics.rs contains the runnable bottom-up implementation using clone_with_children.

fn apply_rule_bottom_up(
    node: Arc<RelNode>,
    rule: &impl Fn(Arc<RelNode>) -> Option<Arc<RelNode>>,
) -> Arc<RelNode> {
    // Get optimized children
    let mut children = Vec::new();
    for child in node.children() {
        let child = apply_rule_bottom_up(child, rule);
        children.push(child);
    }
    let rel = create_new_plan_node(node, children);
    // Apply the rule after all children are processed
    rule(rel.clone()).unwrap_or(rel)
}

fn apply_rule_top_down(
    node: Arc<RelNode>,
    rule: &impl Fn(Arc<RelNode>) -> Option<Arc<RelNode>>,
) -> Arc<RelNode> {
    // Apply the rule before all children are processed
    let node = rule(node.clone()).unwrap_or(node);
    // Get optimized children
    let mut children = Vec::new();
    for child in node.children() {
        let child = apply_rule_top_down(child, rule);
        children.push(child);
    }
    create_new_plan_node(node, children)
}

The application order does not matter for join commutativity: if it is the only rule in the system, we will always get the same result. It is not a good heuristic-optimizer example because we cannot know which join order is better without a cost model. Instead, let’s use filter pushdown to illustrate rule application order. We can push a filter past a projection when we find a filter-projection pair.

As the figure shows, applying the filter-projection rule in top-down order pushes the filter to the scan node, while applying it in bottom-up order pushes it down by only one level.

Besides choosing where to start in the plan tree, a heuristic optimizer must decide which rule to invoke first and how many times to apply it.

With these two traversal orders, we can write a generic apply_rule_bottom_up. Notice the representation problem, however: the traversal must retrieve every node’s children and clone each node with a new list of children. We need either a trait on every plan node or methods on the RelNode enum to expose those operations.

impl Join {
    pub fn children(&self) -> Vec<Arc<RelNode>> {
        vec![self.left.clone(), self.right.clone(), self.cond.clone()]
    }

    pub fn clone_with_children(&self, children: Vec<Arc<RelNode>>) -> Self {
        Self {
            left: children[0].clone(),
            right: children[1].clone(),
            cond: children[2].clone(),
        }
    }
}


impl RelNode {
    pub fn children(&self) -> Vec<Arc<RelNode>> {
        match self {
            RelNode::Join(join) => join.children(),
            RelNode::Filter(filter) => filter.children(),
            // ...
        }
    }

    pub fn clone_with_children(&self, children: Vec<Arc<RelNode>>) -> Self {
        match self {
            RelNode::Join(join) => RelNode::Join(join.clone_with_children(children)),
            RelNode::Filter(filter) => RelNode::Filter(filter.clone_with_children(children)),
            // ...
        }
    }
}

Now, we can write the generic function for applying the rule bottom-up:

fn apply_rule_bottom_up(
    node: Arc<RelNode>,
    rule: &impl Fn(Arc<RelNode>) -> Option<Arc<RelNode>>,
) -> Arc<RelNode> {
    // Get optimized children
    let mut children = Vec::new();
    for child in node.children() { // <- defined by the trait
        let child = apply_rule_bottom_up(child, rule);
        children.push(child);
    }
    let rel = Arc::new(node.clone_with_children(children)); // <- defined by the trait
    // Apply the rule after all children are processed
    rule(rel.clone()).unwrap_or(rel)
}

Later posts will return to rule invocation order and repeated application.

A Cost-Based Optimizer Framework

[code for this section]

With the heuristic optimizer framework, we can rewrite a query plan with rules that we believe will produce a better one. Some rules, however, do not always improve the plan. Consider a set of rules that reorder the joins in this query:

SELECT * FROM t1, t2, t3 WHERE t1.x = t2.x AND t1.y = t3.y

The plan’s initial join order is (t1 join t2) join t3. All joins have equality conditions. If the join-reordering rules instead produce (t2 join t3) join t1, the join between t2 and t3 has no equality condition, so we must use a nested-loop join. The resulting plan is worse than the original.

To avoid generating a worse plan, the optimizer must know which plan is better. Users can define a cost model to help it decide. A simple cost model might say that joins with equality conditions are always better than those without them. A more complex, statistics-based model might use histograms and base-table cardinalities to choose a join order.

The simplest way to implement a cost-based optimizer is to enumerate all possible plans by firing every rule until further applications produce no new plans. How many possible plans will we enumerate for this query? ^{^3}

(t1 join t2) join t3
(t2 join t1) join t3
t3 join (t1 join t2)
t3 join (t2 join t1)
(t1 join t3) join t2
(t3 join t1) join t2
t2 join (t1 join t3)
t2 join (t3 join t1)
t1 join (t2 join t3)
t1 join (t3 join t2)
(t2 join t3) join t1
(t3 join t2) join t1

^3: Let’s keep things simple for now. Filters are always part of the join conditions and never appear as standalone filter operators. We do not consider column order, so there are no projection nodes. We also ignore join implementations (hash join versus nested-loop join); otherwise, we would get even more plans.

The plan space contains 12 possible plans. We can run the cost model on each one and choose the best, but doing so repeats work. For example, we compute the cost of (t1 join t2) in more than one plan.

Moreover, (t1 join t2) produces the same result as (t2 join t1). If we know which order is cheaper, we can reuse that result while costing larger plans. We need a structure that stores the intermediate cost and winner for each equivalent set of plans.

This is where the most critical structure in a cost-based optimizer appears: the memo table. The memo table stores groups of equivalent plan subtrees.

• Each equivalence set forms a memo table group.
• Each group contains multiple memo table expressions. A memo expression contains only a plan node and links to its children’s groups. For example, storing (t1 join t2) creates three groups: group 1 contains scan t1, group 2 contains scan t2, and group 3 contains join group1 group2. Notice how the plan nodes are flattened into the memo table.
• Each memo table group can have one winner expression. We choose each group’s winner recursively from leaves to roots (that is, from scan nodes to join nodes).

Let’s see how to find the best plan in the memo table. The cost model prefers the smaller table on the left-hand side of a join. A join’s estimated output row count is min(left, right). We start at the leaves. Because each scan group has only one expression, that expression is necessarily the winner.

Next, we move to the two-way join nodes and find the winner in each group. We compute each expression’s cost from the costs and statistics of its child groups’ winners, then compare those costs to select the best expression.

Finally, in the top-level group 7, we find the winner among six expressions.

The best plan is (t1 join t2) join t3, which we can reconstruct by following the winner expressions through the memo table. By storing each group’s intermediate cost, statistics, and winner, the memo table avoids duplicate computation.

Returning to plan representation, how do we store memo table expressions? We need a new plan representation in which each node’s children are GroupIds.

#[derive(Copy, Debug, Clone, Hash, Eq, PartialEq)]
pub struct GroupId(usize);

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct MemoJoin {
    pub left: GroupId,
    pub right: GroupId,
    pub cond: GroupId,
}

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct MemoFilter {
    pub child: GroupId,
    pub predicate: GroupId,
}

When we receive the initial, unoptimized plan tree, we must memorize it in the memo table before applying transformations and deriving the best plan.

pub fn memorize_rel(memo: &mut Memo, rel: Arc<RelNode>) -> GroupId {
    let rel = match &*rel {
        RelNode::Scan(scan) => MemoRelNode::Scan(scan.clone()),
        RelNode::Join(join) => MemoRelNode::Join(MemoJoin {
            left: memorize_rel(memo, join.left.clone()),
            right: memorize_rel(memo, join.right.clone()),
            cond: memorize_rel(memo, join.cond.clone()),
        }),
        RelNode::Filter(filter) => MemoRelNode::Filter(MemoFilter {
            child: memorize_rel(memo, filter.child.clone()),
            predicate: memorize_rel(memo, filter.predicate.clone()),
        }),
        // ... more RelNodes, doesn't seem maintainable
    };
    memo.add_expr(rel)
}

We have covered how to populate the memo table with the initial plan tree and find a winner. But how do we apply transformations to produce multiple expressions in each group?

Apply the Rules Again: Transformations in the Memo Table

[code for this section]

Initially, the memo table contains the original plan and exactly one expression per group. We can apply transformations to each expression to create more equivalent expressions in the same group.

Join commutativity is straightforward: we find every node shaped like join A B and insert join B A into the same group.

Join associativity is more complex. We must match a pattern such as join (join A B) C, but each memo table group contains only one level of the plan tree. We must first find every expression shaped like join X C, look inside group X, and iterate over its join A B expressions. We then stitch the pieces together into a structure that matches the original pattern. That structure is called a binding.

Bindings differ from the original plan tree because group IDs appear as leaves. We therefore need a third plan representation for rule-match bindings.

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct BindJoin {
    pub left: Arc<BindRelNode>,
    pub right: Arc<BindRelNode>,
    pub cond: Arc<BindRelNode>,
}

#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub struct BindFilter {
    pub child: Arc<BindRelNode>,
    pub predicate: Arc<BindRelNode>,
}


#[derive(Debug, Clone, Hash, Eq, PartialEq)]
pub enum BindRelNode {
    Scan(BindScan),
    Join(BindJoin),
    // ... Other RelNodes
    Group(GroupId),
}

Once the optimizer produces a binding from the memo table, users can define transformations that turn it into an equivalent plan:

 fn join_assoc_memo(node: Arc<BindRelNode>) -> Option<Arc<BindRelNode>> {
    if let BindRelNode::Join(ref a) = &*node {
        if let BindRelNode::Join(b) = &*a.left {
            return Some(Arc::new(BindRelNode::Join(BindJoin {
                left: b.left.clone(),
                right: Arc::new(BindRelNode::Join(BindJoin {
                    left: b.right.clone(),
                    right: a.right.clone(),
                    cond: a.cond.clone(),
                })),
                cond: b.cond.clone(),
            })));
        }
    }
    None
}

To find every binding that matches the join-associativity rule, we need a loop that expands the left child of the top-level match:

pub fn apply_join_assoc_rules_on_node(memo: &mut Memo, group: GroupId, node: MemoRelNode) {
    if let MemoRelNode::Join(node1) = node {
        // Expand the left group of the join
        for expr in memo.get_all_exprs_in_group(node1.left) {
            if let MemoRelNode::Join(node2) = expr {
                let binding = BindJoin {
                    left: Arc::new(BindRelNode::Join(BindJoin {
                        left: Arc::new(BindRelNode::Group(node2.left)),
                        right: Arc::new(BindRelNode::Group(node2.right)),
                        cond: Arc::new(BindRelNode::Group(node2.cond)),
                    })),
                    right: Arc::new(BindRelNode::Group(node1.right)),
                    cond: Arc::new(BindRelNode::Group(node1.cond)),
                };
                let applied = join_assoc_memo(Arc::new(BindRelNode::Join(binding))).unwrap();
                add_binding_to_memo(memo, group, applied);
            }
        }
    }
}

We need yet another function to add bindings back to the memo table:

pub fn add_binding_to_memo(memo: &mut Memo, group: GroupId, node: Arc<BindRelNode>) -> GroupId {
    fn add_binding_to_memo_inner(memo: &mut Memo, node: Arc<BindRelNode>) -> GroupId {
        let node = match &*node {
            BindRelNode::Scan(scan) => MemoRelNode::Scan(scan.clone()),
            BindRelNode::Join(join) => {
                let left = add_binding_to_memo_inner(memo, join.left.clone());
                let right = add_binding_to_memo_inner(memo, join.right.clone());
                let cond = add_binding_to_memo_inner(memo, join.cond.clone());
                MemoRelNode::Join(MemoJoin { left, right, cond })
            }
            BindRelNode::Filter(filter) => {
                let child = add_binding_to_memo_inner(memo, filter.child.clone());
                let predicate = add_binding_to_memo_inner(memo, filter.predicate.clone());
                MemoRelNode::Filter(MemoFilter { child, predicate })
            }
            // ... other RelNodes
            BindRelNode::Group(group) => return *group,
        };
        memo.add_expr(node.clone())
    }
    let new_group = add_binding_to_memo_inner(memo, node);
    if group != new_group {
        memo.merge_group(group, new_group)
    } else {
        group
    }
}

We are accumulating chunks of unmaintainable code. Every time a user adds a new plan node, they must modify many places:

1. We need alternative representations of the new plan node for the memo table (MemoXXX) and rule-matching bindings (BindXXX). The user therefore has to define the same plan node in three different structures.
2. The user must implement children and clone_with_children for the new plan node so that the optimizer can access and modify its children.
3. The user must add the new plan node to memorize_rel, add_binding_to_memo, and many other functions that convert between plan-node representations.

From a user’s point of view, these optimizer internals should not matter: how the optimizer stores memo entries, applies rules, and so on. Ideally, users could tell the optimizer, “Here are the plan nodes, transformation rules, and cost model; please give me the best plan.” The current representation and its surrounding design do not make that goal easy to achieve.

Is it possible to design a new plan representation that makes users’ lives easier?

A New Plan Representation

[code for this section]

The plan representation we saw earlier is widely used in query optimizers, especially those written in object-oriented languages such as Java (for example, Apache Calcite). Using the same representation in Rust is much harder because Rust does not offer the same kind of polymorphism as Java. To adapt it, we need a separate structure for each representation of a plan node (for example, RelNode, BindRelNode, and MemoRelNode).

Let’s revisit these problems from the optimizer framework’s perspective.

1. The optimizer framework must understand what is inside the user’s RelNode enum. For example, memorize_rel must match on the user-provided RelNode enum to transform a node into the corresponding memo table expression. When a rule matches, the framework must also convert that expression into the corresponding binding.
2. The optimizer framework needs an easy way to access and manipulate children, which is why every plan node needs children and clone_with_children.
3. The framework has multiple representations of the same plan node and must convert between them. For example, memorize_rel converts the initial plan tree into a memo table representation, while apply_join_assoc_rules_on_node converts that representation into a binding.

Can we have a plan representation that is friendly for the optimizer framework to manipulate and pleasant for the user?

I propose a new plan representation. The next snippets show the stage-6 representation from s06_new_repr.rs, rather than the tagged RelNode enum used in the opening stages.

#[derive(Clone)]
pub enum RelNodeType {
    Scan,
    Filter,
    Join,
    Eq,
    ColumnRef,
    Const,
}

pub enum RelAttrType {
    TableId(TableId),
    ColumnRef(usize),
    Const(i64),
    None,
}

pub struct RelNode {
    pub typ: RelNodeType,
    pub children: Vec<Arc<RelNode>>,
    pub data: Arc<RelAttrType>,
}

To construct a plan node,

pub fn join(
    left: impl Into<Arc<RelNode>>,
    right: impl Into<Arc<RelNode>>,
    cond: impl Into<Arc<RelNode>>,
) -> RelNode {
    RelNode {
        typ: RelNodeType::Join,
        children: vec![left.into(), right.into(), cond.into()],
        data: Arc::new(RelAttrType::None),
    }
}

pub fn scan(table: TableId) -> RelNode {
    RelNode {
        typ: RelNodeType::Scan,
        children: vec![],
        data: Arc::new(RelAttrType::TableId(table)),
    }
}

We split the original plan node into three parts: its type, children, and data. From the framework’s perspective, users define RelNodeType and RelAttrType, while the framework defines RelNode. We can go one step further and make the representation generic:

pub struct RelNode<T: RelNodeType, D: RelAttrType> {
    pub typ: T,
    pub children: Vec<Arc<RelNode<T, D>>>,
    pub data: Arc<D>,
}

The optimizer framework does not need to know what T: RelNodeType is or how many plan-node variants there are. It only compares, hashes, and clones T. The user can choose an enum, a Box<dyn SomeTrait>, or another suitable type. This representation also lets the optimizer access and manipulate children directly. To convert a plan node into another representation, the optimizer can preserve T and D while rewriting children into the expected form.

We can then define every plan representation used in this post:

pub struct MemoRelNode {
    pub typ: RelNodeType,
    pub children: Vec<GroupId>,
    pub data: RelAttrType,
}

pub enum BindRelNode {
    RelNode {
        typ: RelNodeType,
        children: Vec<Arc<BindRelNode>>,
        data: Arc<RelAttrType>,
    },
    Group(GroupId),
}

// We will talk about this in the future
pub enum RelNodeMatcher {
    Match {
        typ: RelNodeType,
        children: Vec<RelNodeMatcher>,
    },
    Any
}

The conversion details are for a later post. For now, the key point is that this representation avoids matching on every plan-node variant. When users add a new plan node, they only need to define its RelNodeType and RelAttrType; the optimizer framework handles the rest.

Users still need to implement helper functions that interpret the structure of RelNode.

pub struct Join(Arc<RelNode>);

impl Join {
    pub fn left(&self) -> Arc<RelNode> {
        self.0.children[0].clone()
    }

    pub fn right(&self) -> Arc<RelNode> {
        self.0.children[1].clone()
    }

    pub fn cond(&self) -> Arc<RelNode> {
        self.0.children[2].clone()
    }
}

We can make these interpretations generic across RelNode, BindRelNode, and MemoRelNode, allowing a single impl to interpret all three. Later posts in this series will explore this plan representation and the Cascades/Volcano query optimization framework in more detail.

Do you think the proposed plan representation is a better alternative than the Apache Calcite representation? Continue the discussion in the corresponding GitHub Discussion.