Database Theory in Action: Search-Based Program Optimization

Program optimization in the compilers literature is typically viewed from the perspective of term rewriting. Given an input term $t$ and a set of rewrite rules $ℛ$ of the form ${\text{expr}}_1\to {\text{expr}}_2$, a system repeatedly chooses a rule from $ℛ$ to apply to the term $t$, hopefully moving it closer and closer to a “better” program. When $ℛ$ is convergent (i.e. terminating and confluent) this works great! No matter the order of rewrite application, this process finishes (termination) and arrives at the same term (confluence).

What happens when a term rewriting system does not enjoy properties like confluence? Consider trying to optimize the program $(a\times 2)/2$ to $a$. One promising transformation in general is $a\times 2\to a\ll 1$: replacing a multiplication with a cheaper bitshift. However, in this particular case, applying that locally-good transformation would obscure the globally optimal solution. In the compilers literature, this is called the phase-ordering problem: different orderings of analyses and transformations produce different results.

Equality saturation (EqSat) is technique that promises to mitigate these above challenges with conventional term rewriting. The key idea is that rewriting no longer destroys the old term, instead it monotonically adds an equality between the new and old term, so applying the “incorrect” rewrite at any given time is no longer as punishing. To represent this large space of programs, EqSat borrows the e-graph data structure [14] from the automated theorem proving community, where it sits in the core of SMT solvers like Z3 [4] and CVC [1] to reason about the equality of uninterpreted functions. The original EqSat work [20] demonstrated the e-graph’s promise as the basis for a search-based program optimization technique. Recently, EqSat has seen a flurry of new work advancing the technique [25, 28, 9, 29], developing new toolchains like egg, and applying it to new areas [17, 13, 27, 23, 22, 12].

Figure 1 shows an example e-graph and how EqSat applies rewrites to grow the space of programs that it represents. Using the definition from [25], an e-graph is a set of e-classes (equivalence classes) that contain equivalent e-nodes. An e-node is a function application that points to e-classes (not e-nodes) as children; this captures a notion of congruence: if $a=b$ then $f(a)=f(b)$. Rewrites over the e-graph are somewhat of a misnomer, they only add new e-nodes and e-classes to the e-graph. Figure 1b shows a potentially “problematic” rewrite that in the traditional setting would hinder the optimization of $(a\times 2)/2$ to $a$.

The typical equality saturation workflow is as follows:

1. 1.

   (Figure 1a) Insert the input program into the e-graph; say the input is stored in e-class $c$.

2. 2.

   (Figure 1b-d) Apply rewrites to the e-graph until a fixed point or a timeout is reached; note that a fixed point may not exist.

3. 3.

   Extract from $c$ the cheapest program according to some cost function.

Applying rewrites is typically the most expensive part of the EqSat workflow. This complexity comes from the fact that the e-graph represents many, many programs (in fact, potentially infinite programs in e-graphs with cycles) that could match the left-hand side of a rewrite. The e-matching (pattern matching modulo equality) problem is the task of, given a pattern $p$ and an e-class $c$, find all substitutions $\sigma$ such that $p\sigma \in c$ where $\sigma$ maps variables in $p$ to e-classes in the e-graph.

Notice the rule $x/x\to 1$ in Figure 1d is not always sound. If $x$ is instantiated with a term that may evaluate to $0$, the left-hand side $x/x$ is undefined and should not be rewritten to $1$. E-class analyses provides a mechanism for incorporating semantic information during EqSat: An e-class analysis annotates each e-class with an element from a join semi-lattice. As e-classes are merged through rewrites, semantic information about e-classes is propagated and aggregated by taking the join (least upper bound) of individual e-classes’ annotated values. For example, an e-classes analysis can track the possible ranges with interval arithmetic. If two e-classes with possible ranges $[-2,5]$ and $[1,7]$ are merged (discovered to be equivalent), it can be concluded that the merged e-class can only have possible range $[1,5]$. This semantic information can then be used to safeguard rules like $x/x\to 0$ with appropriate side conditions (e.g., range of $x$ should not contain $0$). There are, however, some limitations with e-class analyses: it only allows one e-class analysis per E-graph, and the analysis data are only propagated bottom up, but not top down, making certain analyses like type analysis infeasible.

Both e-matching and database queries aim to find homomorphisms: the former over the e-graph data structure and the latter over a database instance. As the e-matching problem grows more complex – patterns with repeated variable use, searching multiple correlated patterns simultaneously, the need for incremental search, etc. – the techniques employed become more complex as well [3]. Instead of building on existing e-matching techniques, recent work [29] encoded e-matching as conjunctive queries over a relational database instance, then used a Worst Case Optimal Join (WCOJ) [15] to perform e-matching. The work followed the standard flattening of terms into tables. For example, the term $y=f(f(a))$ can be encoded in a binary relation, call it $R_f(\text{child},\text{id})$, containing the tuples $(a,x)$, $(x,y)$, where $x,y$ are the identifiers of the nodes $f(a)$ and $f(f(a))$ respectively. To represent $z=g(f(f(a)),b)$, add a second relation $R_g({\text{child}}_1,{\text{child}}_2,\text{id})$, containing the tuple $(y,b,z)$. When two e-nodes are in the same e-class, they simply share the same id. This encoding naturally supports complex patterns, since a query optimizer can exploit the query structure for asymptotic speedup over the traditional e-matching algorithm [29].

egglog is a recent system that unifies EqSat and Datalog. There is a simple intuition here: both EqSat and Datalog perform non-destructive, fixpoint reasoning over some (term or relational) databases. The resulting system is beneficial to both EqSat and Datalog. For EqSat, this brings more expressive program analyses over e-graphs, powerful query optimization, and incrementalization via semi-naïve evaluation. Compared to Datalog, the union-find-based equational reasoning from EqSat makes tasks like unification-based pointer analysis [18] require asympototically less space. Figure 2 shows some egglog code in both a Datalog and EqSat style.

At the core of egglog is the concept of functional dependency (FD). egglog is based on a “functional” database design – functions are just relations with functional dependencies. Let $f$ be a binary function table; under the hood it is really just a ternary relation $R_f(c_1,c_2,c)$ with FD $(c_1,c_2)\to c$. What if the FDs are violated? For example, suppose $f(a,b)=c$ and later we learned also that $f(a,b)=c^{\prime }$. In egglog there are only two cases. If $c$ and $c^{\prime }$ are e-classes, then we can unify $c$ and $c^{\prime }$ – from now on, $c$ and $c^{\prime }$ should be viewed as the same e-class and will not be distinguishable. If $c$ and $c^{\prime }$ are lattice values, then we can take the join of $c$ and $c^{\prime }$ and claim $f(a,b)=c\vee c^{\prime }$, since the knowledge $c\vee c^{\prime }$ subsumes both $c$ and $c^{\prime }$. This functional dependency repair is implemented via a congruence-closure-style algorithm [5].

The dominant query optimization architecture in production database systems is derived from the Volcano [7] and Cascades [8] systems. Yongwen Xu’s Master thesis [26] describes an enhanced version of Cascades, and the optd project¹¹1The source code of optd is available at https://github.com/cmu-db/optd?tab=readme-ov-file. started at Carnegie Mellon University provides a modern implementation. These query optimizers can be characterized as rewrite-based, memoized search for all queries equivalent to the input, selecting the cheapest in a top-down manner. The rewrite-based aspect makes the system extensible; if a new logical or physical operator is introduced, one can just add new rewrites to enable optimization of queries with the new construct. The memoization of the search is critical for performance, as many transformations of (sub)queries will lead to programs that have already been explored. These characteristics are all present in EqSat: the memo table in Cascades corresponds to the e-graph data structure, which employs the classic union-find data structure [19] to efficiently store and maintain the equivalence classes. Unlike Cascades, EqSat computes the congruence closure [5] of the equivalence relation to further compress the e-graph. These similarities have already spawned interest in using EqSat for portions of query optimization [24] or implementing the entire process (e.g., in CubeSQL and RisingLightDB). The Volcano/Cascades system also critically employs top-down search to enable a branch-and-bound technique, while EqSat/Datalog are bottom-up. Yet both support top-down search, as we explain next.

Magic set optimization [11], demand transformation [21], and more recent works around ADTs and functional programming in Datalog [6, 16] all aim to provide some degree of top-down control to Datalog. Previous works employing Datalog for query optimization [10, 2] have also found it necessary to implement top-down search in a manual manner. EqSat and egglog offer an additional perspective on this problem. EqSat is bottom-up, but it provides an additional tool that can be used to “place demand” on certain tuples/terms. In EqSat, the term $f(a)$ can be considered on its own (in its own e-class) without knowledge of its output. This is similar to labeled nulls; the e-class can always be unified with the result at a later time.