marginparsep has been altered.
topmargin has been altered.
marginparwidth has been altered.
marginparpush has been altered.
The page layout violates the ICML style. Please do not change the page layout,
or include packages like geometry, savetrees, or fullpage, which change it for
you. We’re not able to reliably undo arbitrary changes to the style. Please
remove the offending package(s), or layout-changing commands and try again.

Learned Graph Rewriting with Equality Saturation: A New Paradigm in Relational
Query Rewrite and Beyond

Anonymous Authors1

###### Abstract

Query rewrite systems perform graph substitutions using rewrite rules to
generate optimal SQL query plans. Rewriting logical and physical relational
query plans is proven to be an NP-hard sequential decision-making problem with
a search space exponential in the number of rewrite rules. In this paper, we
address the query rewrite problem by interleaving Equality Saturation and
Graph Reinforcement Learning (RL). The proposed system, Aurora, rewrites
relational queries by guiding Equality Saturation, a method from compiler
literature to perform non-destructive graph rewriting, with a novel RL agent
that embeds both the spatial structure of the query graph as well as the
temporal dimension associated with the sequential construction of query plans.
Our results show Graph Reinforcement Learning for non-destructive graph
rewriting yields SQL plans orders of magnitude faster than existing equality
saturation solvers, while also achieving competitive results against
mainstream query optimisers.

††footnotetext:  1Anonymous Institution, Anonymous City, Anonymous Region,
Anonymous Country. Correspondence to: Anonymous Author
<EMAIL_ADDRESS>
Preprint.

##  1 Introduction

Databases (DBs) are intricate systems for storing and managing data. To
manipulate the data, DBs allow users to submit queries, declarative statements
through which the user specifies what data to fetch or manipulate. The DB
system manages how the query is executed through query plans. Rewriting query
plans is a fundamental problem in DB literature and a core driver for
performance Finance & Gardarin (1991). DBs employ query rewrite algorithms and
rewrite rules (e.g. relational algebra rules) to transform the original
database transaction into an equivalent and more efficient one. If executed
effectively, the rewrite process can speed up the original transaction by
orders of magnitude Zhou et al. (2021). The rewrite system takes as input a
database query and projects it into a logical plan. Subsequently, the plan is
rewired based on pre-defined rules (e.g. remove redundant operators) with two
overarching goals in mind: 1) to produce a plan with faster execution time
than the original, and 2) to ensure the output of the rewritten plan is
equivalent to the original.

Traditional query optimisers apply the rewrite rules in a predefined order,
with each rewrite operating on the output of its predecessor Begoli et al.
(2018). For example, PostgreSQL PostgreSQL (1996) rewrites queries in a top-
down fashion. Graph substitutions are performed sequentially starting from the
query plan root node, based on pattern matching. The inherent limitation of
this approach is the potential to fall into a local optimum when the rewrite
rules have inter-dependencies. Finding the optimal order to apply query
rewrite rules is generally NP-hard Finance & Gardarin (1991); Pirahesh et al.
(1992). An alternative approach to mainstream query optimisers is to apply all
rules (i.e. graph substitutions) at once and subsequently pick the best query
plan from the conglomerate. This paradigm is known as equality saturation, and
it has been successfully applied in a variety of optimisation problems Wang et
al. (2020); Yang et al. (2021). However, a naive application of equality
saturation to rewrite SQL queries in the general case is set for failure.
While equality saturation as proposed by Tate et al. (2009) employs equality
graphs (e-graphs) to store a potentially infinite number of expressions (e.g.
query plans) within finite memory, the algorithm’s computational requirements
remain impractical for local and cloud DB deployments with limited resources.
The limitations of equality saturation for query rewrite are not solely tied
to infeasible memory requirements but also hamper the algorithm’s execution
time which is reflected in the end-to-end latency of SQL queries.

These limitations underscore the need for an alternative approach capable of
constraining the non-destructive graph rewriting algorithm within reasonable
memory and time limits. On their own, e-graphs help solve one of the pressing
challenges in the query rewrite literature, which looks at how to represent a
large number of query plans Zhou et al. (2021). Through Aurora, we aim to
solve the remaining challenge of identifying optimal query plans given a set
of rewrite rules in relational DBs. We introduce a new Graph Reinforcement
Learning agent that is able to navigate the complex search space of e-graphs
by using a mix of Graph Neural Networks (GNNs) and Recurrent Neural Networks
(RNNs). We show that by guiding the construction of query e-graphs with RL,
Aurora outperforms existing state-of-the-art equality saturation solvers,
returning order-of-magnitude faster query plans within a service-level
agreement (SLA) on memory requirements. Furthermore, the RL-based equality
saturation paradigm we use to rewrite relational query plans can be adapted to
other programming languages than SQL with microscopic changes.

Concretely, our contributions are:

  1. 1.

A pioneering approach to the query rewrite problem using Deep Reinforcement
Learning (DRL) within the context of equality saturation.

  2. 2.

The design and implementation of Aurora, a database query optimizer that
leverages RL-based equality saturation. The methodology encompasses: a) a
novel spatio-temporal DRL algorithm for solving combinatorial problems over
graphs; and b) an integer linear programming (ILP) formula for extracting
query plans from e-graphs.

  3. 3.

An extensive experimental evaluation of Aurora against the state-of-the-art
equality saturation solvers. We show that Aurora reduces the cost of query
plans by orders of magnitude compared to the baselines, and it’s competitive
against traditional DB optimisers.

##  2 Background

This section introduces equality graphs, as a method to represent a
conglomerate of equivalent expressions, the equality saturation algorithm as
initially proposed by Tate et al. (2009), and reinforcement learning (RL) as a
method to optimise sequential decision-making.

###  2.1 Equality Graphs

Equality graphs (e-graphs) are data structures that efficiently represent
congruence relations Nelson & Oppen (1980); Nieuwenhuis & Oliveras (2005).
E-graphs, initially developed for automated theorem provers (ATPs) Willsey et
al. (2021), extend the union-find method proposed by Tarjan (1975) to
efficiently encode expressions in equivalence classes. A core property of the
e-graph is the ability to close the equivalence relation under congruence
(e.g. if $a\equiv b\rightarrow f(a)\equiv f(b)$).

E-graphs have become ubiquitous in rewrite-driven compiler optimisations, as
part of a technique coined equality saturation Tate et al. (2009); Wang et al.
(2020; 2022); Yang et al. (2021). Equality saturation receives an input
expression (e.g. $a+b$) and stores it in an e-graph. Subsequently, the e-graph
is expanded iteratively by applying rewrite rules over the information it
holds. An important remark is that growing the e-graph is a purely additive
process - the rewrites solely add information instead of changing the original
input expression or the expressions that follow from expanding the e-graph.
Upon reaching saturation, the e-graph is deemed to represent all possible
equivalent expressions that can be derived from the rewrite rules and the
input expression.

###  2.2 Equality Saturation

Figure 1: E-graph construction. The e-graph undergoes rewrite $x*2\rightarrow
x<<1$. Figure adapted from Willsey et al. (2021).

Figure 1 shows the construction of an e-graph for mathematical optimisation.
The rewrite rule $x*2\rightarrow x<<1$ only adds information to the original
expression $(x*2)/2$. If there exists a rule such that $(x*2)/2\rightarrow x$,
this rule can be subsequently fired over the e-graph to produce an optimal
plan in terms of the number of operations. In a classic term rewrite system,
the left-hand side expression $(x*2)/2$ would have been obfuscated by
$(x<<1)/2$, thus leading to an ordering problem. An e-graph consists of
e-classes (blue boxes) and e-nodes (circles). An e-class contains one or more
e-nodes. Picking any e-node from an e-class is ensured to produce a correct
and equivalent expression in reference to the original.

In equality saturation, all available rewrite rules are applied sequentially
until the e-graph no longer changes. Once saturation is reached, the optimal
form of an expression can be extracted from the e-graph via a greedy search or
a more sophisticated algorithm such as integer linear programming (ILP) Yang
et al. (2021). In Section 3.2, we discuss why equality saturation cannot be
trivially applied to complex programming languages and problems such as query
rewriting.

###  2.3 Reinforcement Learning

Reinforcement Learning (RL) encompasses a class of techniques for optimising
sequential decision-making processes. One formalisation of RL is Markov
Decision Processes (MDPs) Sutton & Barto (2018), which are instruments for
characterising decision-making problems over a discrete time dimension
Puterman (2014). The decisions are taken by an agent that acts over a Markov
chain Norris (1998), an environment where the agent’s future state depends
only on its previous state and the action taken Frydenberg (1990).

Formally, an MDP is a 4-tuple $(S,A,T,R)$ which describes the state space $S$,
the action space $A$, the dynamics of the environment $T$, and a reward signal
$R$, respectively. To learn a policy for the agent to follow, Artificial
Neural Networks (ANNs) are typically used to generalise across large state-
action (S-A) spaces, which is known as Deep Reinforcement Learning (DRL).

Aurora aims to learn a policy over the space of e-graphs (S) constructed from
input query plans, where the action space A is defined by the set of available
rewrite rules (e.g. relational algebra rules) and the reward is the
expressiveness of the e-graph (the latency of the extracted query plan)
derived within an SLA constraint such as the memory required to store the
e-graph. This formalisation serves as the core intuition behind Aurora.

##  3 Problem Overview

###  3.1 Challenges in Query Rewrite

Figure 2: The original SQL plan (left) is optimised in a top-down fashion by
creating a temporary view over the subquery (right).

Query rewriting is an NP-hard problem Pirahesh et al. (1992); Finance &
Gardarin (1991) due to the exponential search space. Database engines project
SQL queries into logical plans, which are subsequently optimised based on
hand-crafted rules (e.g. perform filtering foremost in the plan). The
application of rules (i.e. substitutions) commonly follows an ordering scheme
(e.g. arbitrary, heuristic) Zhou et al. (2021). In turn, this can lead to
suboptimal query plans, as rewrite rules may have inter-dependencies.
Concretely, substituting parts of the original plan may render more prosperous
future rewrites inapplicable.

To build our intuition, we refer to a study case provided by the authors of
Zhou et al. (2021), who investigate the inner workings of PostgreSQL.
PostgreSQL facilitates database transactions in a top-down manner. The engine
traverses the logical query plan starting from the root and performs local
rewrites iteratively. Figure 2 renders an example SQL query for which a top-
down optimiser finds a sub-optimal plan. The original plan tree is
representative of SQL templates such as “select _ from $table_{1}$ with
condition over (select _ from $table_{1}$, $table_{2}$ where _)”. Once the
optimiser reaches the $Subquery$ node in the plan, it naively creates an
inline view. Consequently, further optimisations are disabled over the view.
By disregarding the characteristics of the leaf nodes, a top-down application
of rewrite rules will not take advantage of the correlation between the
subquery and the outer query (note the same relation - $Table_{1}$ \- is
scanned twice). Ideally, the subquery is pulled up as a join to prevent a
redundant scan operation. In practice, this simple optimisation reduces the
execution time by a factor of $600$ as per the original authors Zhou et al.
(2021).

###  3.2 Challenges in Applying Equality Saturation

Equality Saturation (ES) starts from the promise to address the ordering
problem in classic rewrite-based systems. Different from term rewriting
systems, ES employs e-graphs to perform non-destructive graph rewriting. We
investigate the theoretical and practical limitations of equality saturation.

One of the concerns raised in literature is the prominence of the algorithm to
run into “explosions”. We coin an explosion a soar in the number of nodes in
the e-graph in the aftermath of a suboptimal rewrite rule application. As a
concrete example, Yang et al. (2021) investigate this dimension in the context
of tensor graph optimisation, showing that rewrite rules can quickly grow the
e-graph to the upper limit of $50k$ (thousands) nodes they set. The
consequences of a large e-graph are twofold: 1) the data structure may not fit
into memory in complex domains where explosions are frequent, and 2) it
hampers the time to extract the optimal expression from the e-graph. The
latter point is a pressing concern in systems that provide real-time feedback,
such as database query engines. In this scenario, the overhead associated with
growing the e-graph and extracting the optimal expression is included in the
end-to-end latency of a query execution. In this case, constraining equality
saturation with RL becomes mandatory to achieve competitive latency.

Finally, Willsey et al. Willsey et al. (2021) render a scenario in which
equality saturation is characterised by non-termination, where a purely
destructive approach would terminate. In this context, we denote destructive
any rewrite system which performs substitutions. The condition of non-
termination is linked to an unbounded e-graph expansion, one that never
saturates. Identifying an instance of this for SQL or relational algebra
remains a question for future research, but one that’s tackled inherently by
Aurora.

Figure 3: A 6-step schematic workflow of Aurora

##  4 Solution Overview

Aurora aims to jointly address the challenges underpinning query rewrite and
the limitations associated with existing equality saturation solvers. The
usage of e-graphs is motivated by the large number of equivalent plans that
can be derived from relational rules and beyond. To perform query rewrite with
e-graphs within reasonable planning latency and memory bounds, we introduce RL
into the equality saturation paradigm. We show this leads to query plans order
of magnitude faster than the ones derived by existing ES solvers. Figure 3
renders Aurora’s workflow. The system’s actions are clustered into five
overarching steps, with an optional sixth step that shows how the framework
can be used as a SaaS for mainstream databases (e.g. PostgreSQL).

  * •

Step 1: Aurora’s controller receives a SQL query.

  * •

Step 2: The SQL query is parsed into an extended relational algebra (RA)
expression, with the caveat that the logical plan can include physical
operators (e.g. hash-join, merge-join). The motivation behind the design is to
jointly rewrite the query and to set flags for physical execution.

  * •

Step 3: The derived RA expression is used as input for a new e-graph. The
expression is split into e-classes and e-node operators (e.g. project,
filter).

  * •

Step 4: The RL agent guides the expansion of the e-graph. To achieve this, we
apply pattern-based rewrite rules (e.g. $join(a,b)\rightarrow join(b,a)$) to
generate new e-nodes and e-classes. The logic is encapsulated in an E-Graph
Runner which uses e-graphs good (egg) Willsey et al. (2021), the state-of-the-
art equality saturation implementation.

  * •

Step 5: The e-graph and statistics from the database catalogue (e.g.
cardinality, selectivity) are sent to an extractor module. The module extracts
the best query from the e-graph using integer linear programming (ILP) and a
cost heuristic we explain in Section 6.

  * •

Step 6: The extracted relational plan can be converted back into a SQL
dialect, for instance, to be used in mainstream databases such as PostgreSQL.
This final step is optional in Aurora. The reason is that the framework acts
as an end-to-end SQL optimiser with access to direct execution in Risinglight
Risinglight (2022), a Rust-based DBMS for analytical queries.

##  5 Reinforcement Learning for Equality Saturation

Having established existing equality saturation solvers cannot be trivially
used to generate good query plans due to challenges such as memory
requirements and latency considerations (see Section 7.3 for experimental
results), we introduce RL as a method to execute equality saturation
efficiently. Intuitively, the goal is to encode in the e-graph as many
competitive equivalent query plans to the original as possible within a
limited time and memory budget.

###  5.1 Equality Saturation as Markov Decision Process

This section formalises equality saturation for query rewrite through the lens
of a Markov Decision Process (MDP).

####  5.1.1 State

Example: The SQL query “select * from (select * from t where c2=True) where
c1=True” is projected into relational algebra and rewritten by the query
planner as $\pi(\sigma_{c1}(\sigma_{c2}(t)))$. An example is provided in
Figure 4.

Figure 4: Example e-graph for RA expression
$\pi(\sigma_{c1}(\sigma_{c2}(t)))$. By using the rewrite rule
“$\sigma(c_{1},\sigma(c_{2},t))\rightarrow\sigma(c_{1}\land c_{2},t)$”, the
query e-graph adds the equivalent plan $\pi(\sigma_{c1\land c2}(t)))$. The
circles represent e-node operators and the blue boxes represent e-classes
Figure 5: Example e-graph encoding. The encoding scheme covers both the nodes
as well as the edges in the e-graph

We take a slice from the rewritten e-graph rendered in Figure 4, which shows
the top-most equivalence class (e-class $E1$), and two ‘filter’ operators
(Figure 5).

Our proposed node encoding over the query e-graph is a 5-tuple (Table 1). The
node type distinguishes between e-nodes and e-classes. An e-class is a pointer
to a set of e-nodes (e.g. filter, join). The operator type is the core
component of the encoding scheme. We note the rewrite actions over the e-graph
are solely facilitated by the presence of certain operator patterns (e.g. a
filter operation before a join operation). We propose encoding statistics from
the target database management system. If the node is an equivalence class,
then the feature vector stores the cardinality (number of rows) and the width
(the number of columns/data bytes of a tuple) returned by its e-nodes. It is
important to make a distinction here between e-classes and e-nodes. The
e-nodes are executable operators (e.g. filter) and they can only be described
through operator costs. To estimate the cost of an operator, we merely need to
know the size of the output of each of its children. We remark that e-classes
can be described through cardinality and width, which are oblivious of the
e-node picked. Finally, we add edge attributes to distinguish between e-node
to e-class edges and vice versa.

Table 1: The node encoding of the query equality graph Feature |  Details
---|---
node type |  One-hot encoding over node types in e-graphs. A node in the e-graph can either be an equivalence class or an e-node.
operator type |  One-hot encoding over the operators. The number of operators spanned by the extended RA language is 64.
cardinality |  The number of records in the e-class
width |  The number of columns in the e-class
cost |  The estimated execution cost

####  5.1.2 Action & Transition Function

The action space is the set of query plan rewrite rules. Aurora uses 37 rules
that encompass RA, mathematical, and boolean rewrite rules, by expanding
Risinglight Risinglight (2022).

Aligned with previous work Singh (2022), we use a Reset action. The Reset
action extracts the optimal plan from the current e-graph and wraps it into a
new e-graph. The aim is two-fold: 1) to induce a heuristic-like effect during
e-graph construction, and 2) to reduce the size of the e-graph while
maintaining opportunities for future optimisation.

####  5.1.3 Reward Signal

The reward system encourages latency improvements by benchmarking the final
query plan extracted from an e-graph in an optimisation episode. Positive
rewards are computed by measuring the cost improvement between the previous
and newly-extracted plans from the e-graph. In terms of negative rewards, the
agent is penalised when taking actions that are either saturated (i.e., don’t
add new nodes) or lead to large memory overhead.

###  5.2 Spatio-temporal Reinforcement Learning

We propose an actor-critic RL agent Sutton & Barto (2018) to learn the
environment (Figure 6). We use knowledge from Graph Neural Networks (GNNs)
Zhou et al. (2020) and Recurrent Neural Networks (RNNs) Medsker & Jain (2001)
to learn the dynamics of expanding query e-graphs.

Figure 6: The architecture of the spatio-temporal RL agent. It features a
Graph Attention Network layer, a global add pool, a Long Short-Term Memory
layer and two linear heads for actor and critic, respectively.

####  5.2.1 Learning the Spatial Representation

In equality saturation, the e-graph serves two interleaved purposes: 1) it
efficiently represents multiple versions of compiler plans, and 2) it guides
the rewrite process. The layout of the e-graph guides which rewrite rules can
be applied over its structure based on pattern matching. Given some rewrite
rules are more beneficial than others, it is paramount to learn node
neighbourhood patterns, warranting an attention-based method.

GNNs are learning methods for graph-structured data Zhou et al. (2020). The
key idea behind GNNs is to use pairwise message passing between graph nodes to
update their latent representations. The Graph Attention Network (GAT),
introduced by Velickovic et al. Velickovic et al. (2017), is an attention-
based architecture for graph-structured data. A key advantage of GATs over
other types of graph learning methods, such as Graph Convolutional Networks
(GCN) Kipf & Welling (2016), is their ability to automatically infer the
importance of graph substructures. In our proposed agent architecture, a
3-layer GAT is introduced to model a spatial representation of the encoded
e-graph.

####  5.2.2 Learning the Temporal Link

Recurrent neural networks (RNN) Medsker & Jain (2001) are neural networks for
learning sequential data. The key innovation behind RNNs is the ability to
memorize previous steps to influence subsequent inference. We leverage an RNN
to model the trajectory in the e-graph construction MDP. The motivation is
two-fold. First, the RNN is employed to capture the purely additive nature of
rewriting e-graphs (e-nodes and some e-classes will not change). Second, we
aim to mitigate some of the expressiveness pitfalls of the graph encoder. As
an example, a global mean pool function may project two distinct e-graphs into
the same latent representation, leading to a sub-optimal policy.

Aurora uses a Long Short-Term Memory (LSTM) Sherstinsky (2020); Hochreiter &
Schmidhuber (1997) to prevent the vanishing gradient problem Hochreiter (1998)
over long equality saturation episodes. The LSTM is formally defined below:

$c_{t}=f_{t}\odot c_{t-1}+i_{t}\odot g_{t},\quad h_{t}=o_{t}\odot\tanh(c_{t})$
(1)

In Equation 1, $c_{t}$ collects information from the state space. There are
two gates: an input gate denoted $i_{t}$ and a forget gate denoted $f_{t}$. If
the input gate is on, then the newly learnt representation $g_{t}$ over the
current state is stored. If the forget gate is on, then some information from
the preceding cell $c_{t-1}$ is discarded. Finally, the information from the
cell is passed to the hidden state $h_{t}$, with a dependency on the output
gate $o_{t}$.

####  5.2.3 Spatio-temporal Proximal Policy Optimisation

Aurora extends PPO Schulman et al. (2017), a policy gradient RL method, to
perform gradient updates on sequential training data. Additionally, we use the
best practices from previous literature on RNNs for RL Huang et al. (2022);
Piao & Liu (2019) and instantiate the LSTM hidden and memory cells with values
of zero, which are reset at the beginning of each RL episode. The actor and
the critic networks ($V$ and $\pi$ in Figure 6) produce a probability
distribution of rewrite actions and the value of an e-graph state encoding,
respectively. In a closed form, these can be computed as follows:

$h_{t},c_{t}=lstm(pool(gat(s_{t}))),h_{t-1},c_{t-1})$ (2)
$\pi_{\theta_{k}}=f_{actor}(h_{t}),\quad V_{\theta_{k}}=f_{critic}(h_{t})$ (3)

The computations above, in which $\theta_{k}$ represents the parameters of the
entire deep model after $k$ training iterations, and $s_{t}$ the e-graph
encoding, respectively, can be used during inference to solve a new equality
saturation instance.

##  6 Query Plan Extraction

The final instrument in Aurora is the e-graph extraction routine (step 5 in
Figure 3), assigned to extract the optimal plan from the query e-graph. It is
important to remember the e-graph is an efficient encoding of multiple
equivalent query plans.

Extracting the optimal plan from the e-graph is non-trivial in the database
scenario. While previous work from the realm mainly focused on the size of the
expression (i.e., the number of nodes in the query plan) as a cost function
Singh (2022), we deem this provides no intuition regarding the performance of
a query plan. To this end, it is paramount to devise a cost function over the
query e-graph that can be linked back to the end-to-end latency of a query
plan. In our implementation, we try variations of the cost model proposed by
Risinglight (2022).

First, a proxy for the size of the data that flows through the query plan is
defined. Define $n$ the number of tuples and $m$ the number of columns in a
target table. The size of the data is proportional to the intermediate result
returned by an operator. For example, the intermediary output returned by a
scan operation over the table has a size of $n\cdot m$. If the scan operator
is followed by a filter operation, which without loss of generality we assume
reduces the number of tuples by $50\%$, then the new intermediary result will
have a size $\frac{n}{2}\cdot m$.

As shown in Section 5.1, the size of the intermediary results can be stored in
e-classes at runtime. This is because all operators (e-nodes) from an e-class
should produce equivalent outputs.

Second, an operator cost function is implemented over the size of the
intermediary result. Each e-node is either a plan operator, a column, or a
table. We use the asymptotic costs from DB literature to assign costs. For
instance, if a table has $n$ tuples and $m$ columns, then a scan operation is
assumed to cost $\mathcal{O}(nm)$. Finally, the plan is extracted from the
e-graph using integer linear programming to minimise the cost heuristic above,
a popular technique in equality saturation Yang et al. (2021).

##  7 Evaluation

This section describes our experimental evaluation of Aurora. Its performance
as an end-to-end query rewrite mechanism is evaluated on a complex suite of
analytical SQL queries. The results are evaluated against e-graphs good
Willsey et al. (2021), the state-of-the-art equality saturation solver,
Omelette Singh (2022), an RL-based equality saturation solver, and a pseudo-
random agent. Each solver has access to the same set of rewrite rules Aurora
uses.

###  7.1 Experimental Setup

####  7.1.1 Testbeds

There are two database management systems (DBMSs) used in our experimental
analysis. Evaluating Aurora as a competitive equality saturation solver is
done by rewriting queries in Risinglight against the equality saturation
baselines. To further validate Aurora’s utility as a query rewriting tool, we
also benchmark it against PostgreSQL’s in-built optimizer, reinforcing
Aurora’s value as a robust query optimiser.

  1. 1.

Risinglight Risinglight (2022) is a DBMS implemented in Rust. It exhibits a
database query engine that independently confines rewrite rules. Aurora
expands the set of rewrite rules from Risinglight and extends the engine to
support user-defined rewrite orders. The plans devised by Aurora can be
directly executed against Risinglight’s storage system without converting the
plan back into a SQL dialect. The usage of Risinglight as a testbed is further
motivated by the smooth integration with e-graphs good, which we show to be a
robust alternative to mainstream query engines.

  2. 2.

PostgreSQL PostgreSQL (1996) is a popular mainstream DBMS with wide
applications in both industry as well as research. To execute queries on
PostgreSQL, the query plans devised by Aurora are translated back into the
Postgres dialect.

####  7.1.2 Benchmarks & Metrics

Aurora has to be effective at rewriting complex database transactions, such as
multi-join analytical queries. To achieve the desired level of rewrite
difficulty, we experiment with the TPC-H database benchmark Poess & Floyd
(2000). The TPC-H benchmark is an online analytical processing (OLAP)
workload. The database schema contains 62 columns across eight tables. The
queries simulate decision support transactions associated with warehousing
environments. For our experimental analysis, we subtract a subset of queries
from the TPC-H suite that exhibits high rewrite potential and is reflective of
the intricacies of the TPC-H benchmark.

We evaluate Aurora as a query rewrite system through query planning and
execution latency. Furthermore, we evaluate the performance of our e-graph
cost model (see Section 6) by measuring the estimated cost against the real
latency.

####  7.1.3 Training & Hardware

For each combinatorial optimisation task, Aurora and Omelette are trained for
$200$k steps, with a horizon of 200 steps. Unless otherwise noted, all the
experimental results are averaged over $5$ seeds. The same seeds are used for
all algorithms. For any given seed, given the trained policies are stochastic,
the pre-trained agents are rolled out in the environment for $100$ roll-outs,
collecting only the best results. The same applies to the heuristic agent.
This provides reasonable statistical significance across the results. E-graphs
good is run only once given it is deterministic.

The system used for training exhibits a Tesla P100 GPU with 16GB of VRAM, 24
logical CPUs, and 64GB of memory.

###  7.2 Performance Results

Figure 7: The median cost of the extracted plan by algorithm over 5 seeds. The
whiskers represent the maximums. Lower cost is better. Figure 8: The relative
improvement over egg, the state-of-the-art equality saturation solver. The
negative outcome is bounded at -100%. Higher relative improvement is better.

Figure 9: Left: the execution latency of the extracted SQL plan with a timeout
set at 350 seconds. Right: the planning latency. All metrics are averaged over
5 trials. In both cases, lower is better.

We first prove Aurora is able to rewrite query plans exhibiting lower costs
than existing equality saturation solvers. We use Risinglight for direct
execution and set an upper bound on the number of nodes in the e-graph to
$1000$ nodes to serve as the memory constraint in the optimisation problem.

Figure 7 renders the ability to rewrite queries across our baselines. The
estimated cost is computed by extracting the final plan from the e-graph. It
is evident from the plot that Aurora discovers query plans that are orders of
magnitude lower in cost than the baselines.

We observe that a random application of rules yields higher-cost plans than
egg. The significance of this result is two-fold. First, it shows that some
structure is needed to solve the query e-graph construction problem. The lack
of structure in random search plays a detrimental role in solving the query
rewrite combinatorial problem, by making it infeasible for the agent to guess
a competitive rule ordering. Second, it warrants that RL is a competitive
choice for optimising query plans over the space spanned by e-graphs. This
follows from the performance gap between the RL agents and heuristic search.

Figure 10: From left to right: the execution latency, the planning latency,
and the estimated cost (PostgreSQL) of the queries from Table 2 (5 trials).
Lower is better in all cases.

Figure 8 quantifies the relative improvement of each algorithm over egg. On
one hand, we remark the RL methods (Aurora, Omelette) are competitive at
guiding the application of rules in equality saturation, with a relative
improvement of $99\%$ over egg for $33\%$ of the queries. On the other hand,
we remark that Omelette performs worse than egg in two cases. It is important
to note that while the relative improvement is reasonably similar between
algorithms in some cases (Query 9), the magnitude of improvement differs
(Figure 8).

Figure 9 shows the execution latency (left) and planning latency (right) when
using Aurora and egg to rewrite TPC-H queries on Risinglight. We observe that
a non-trivial application of rules within equality saturation, such as the one
learnt by Aurora, greatly improves upon the state-of-the-art equality
saturation solver, egg. We set a timeout of 350 seconds for the execution to
finish given the rather small size of the data pool (4.7GB). In 4 out of 6
cases, egg runs into a timeout. This is consistent across the 5 trials we test
over. The importance of this result is two-fold. First, it is clear that
Aurora can find more competitive plans when compared to egg. Second, the cost
model we define is generally a good proxy for the true latency. Aurora’s
planning latency is higher for Query 5 and Query 9. This is expected, as a
wise e-graph construction routine should allow for more rewrites while keeping
the e-graph within the node limit. However, we emphasize that for this
experiment planning operates on an insignificant scale (milliseconds), which
combined with our previous result on the execution latency paints a compelling
picture to prove that Aurora outputs query plans with better end-to-end query
latency.

####  7.2.1 Case Study: PostgreSQL

This section presents a study case on PostgreSQL. The motivation is to show
Aurora is DBMS-agnostic. We use Aurora to rewrite the Alternative query from
Table 2 into the Optimal form. Subsequently, the rewritten query is executed
in PostgreSQL.

Table 2: Equivalent queries with different levels of performance. Alias |  Query
---|---
Original |  select o_orderkey from orders where o_orderstatus = ’F’ and o_orderkey in (select o_orderkey from orders where o_totalprice $>$ 40000)
Optimal |  select o_orderkey from orders where o_orderstatus = ’F’ and o_totalprice $>$ 40000
Alternative |  select key from (select o_orderkey as key, o_orderstatus as status, o_totalprice as price from orders where o_totalprice $>$ 40000 ) as inter where status = ’F’ and price $>$ 40000

Figure 10 shows that Aurora can be generalised across disparate DBMSs. We
remark the plans devised by PostgreSQL incur higher latency than the query
plan returned by Aurora. In the experiment, Aurora receives the second query,
coined Alternative, and returns the Optimal. The rewritten query is $11\%$
faster. We observe the PostgreSQL optimiser does not take advantage of the
correlation between the FILTER operation in the outer query and the one in the
nested query when executing the former (i.e., “price $>$ 40000”), resulting in
higher latency. The Original query is added for reference to show Aurora can
theoretically bring large improvements over mainstream DBs. However, Aurora
currently does not support IN operators over nested queries, while PostgreSQL
does. This engineering extension is left for future work.

###  7.3 The Importance of RL in Aurora

Figure 11: Results for TPC-H Query 5 rewritten using egg. The relation between
the size of the e-graph and the estimated cost of the extracted SQL plan. The
number of nodes is limited to 3.5k. Red line: the growth trend for the number
of nodes in the e-graph. Blue line: the cost of the extracted plan after each
rewrite rule.

This section empirically motivates introducing RL in equality saturation to
avoid pressing challenges such as high memory and algorithmic latency.

Figure 11 shows the link between the size of the e-graph and the cost of the
extracted plan using egg on TPC-H Query 5, a 6-table join query. Aligned with
our expectations, as more equivalent query plans are added to the e-graph, the
cost of the extracted plan decreases, while the number of nodes in the e-graph
increases. A clear observation is that the node increase is not proportional
to the cost improvement. For instance, we remark how the size of the e-graph
explodes after the first $140$ rule applications, while the cost remains
rather constant during the same window.

Since egg applies the rewrite rules sequentially, we raise the question of
whether a plan with a lower cost can be found within the same node limit, by
changing the order in which the rules are applied as well as their frequency.
This is possible if: 1) the e-graph node limit is a concern and therefore
bounded (e.g. egg is set to stop upon reaching 10k nodes by default) and 2)
the e-graph does not get saturated within the node limit. Upon reaching
saturation, the e-graph encodes all possible plans so the extracted plans will
always be optimal.

Figure 12: Top: The relation between the size of the e-graph and the
extraction latency. Bottom: Planning and execution statistics for two
alternative plans for TPC-H Query 5 executed on Risinglight. The plans are
extracted with ILP.

Figure 12 shows the relation between the size of the e-graph and query
latency. We observe from the top plot that as the number of equivalent plans
grows in the e-graph, extracting the optimal query plan from it incurs a
higher cost. In the bottom plot, we measure the planning time (including
extraction), as well as the execution time for two Query 5 plans devised with
flavours of equality saturation. The execution time is approximately $20$
seconds with a standard deviation of $2s$, and a planning latency in the order
of milliseconds using Aurora. The second plan is extracted from an e-graph of
$7.3k$ nodes built using egg. In this case, the planning latency has a mean of
$78$ seconds, with a large percentage associated with extracting the optimal
query plan from the e-graph. The execution time is discarded as the planning
time is already over $3$ times higher than the total latency for the first
plan. This shows equality saturation cannot be trivially applied to optimise
queries, and RL is necessary as per our proposed system, Aurora.

##  8 Related Work

Query optimisation has a long history of research in the database community.
Recently, machine learning techniques have been introduced into the query
optimisation pipeline at different critical points, ranging from query plan
cost estimation Sun & Li (2019); Marcus & Papaemmanouil (2019) to the join
ordering problem Marcus et al. (2019); Yu et al. (2020); Kadkhodaei & Mahmoudi
(2011). Much of the existing work neglects other logical rewrite
opportunities. While current methods can efficiently estimate costs and
propose competitive join orders, the field is scarce in terms of rewriting
complex queries. The reason is two-fold: 1) the space of possible rewritten
queries can be enormous, and 2) the ordering problem is very difficult to
learn. Of particular importance, we identify the work of Zhou et al. (2021),
who seek to learn over the space of all query rewrite orders with RL. Our
work, which adapts equality saturation for query rewrite is orthogonal, as we
propose to reduce the search space by learning over the query e-graph space.

Equality saturation has seen a boost in popularity as a consequence of ever-
faster hardware and improvements to the algorithm. E-graphs good (egg) Willsey
et al. (2021) revitalises the methodology with asymptotic latency improvement
over the original solver by Tate et al. (2009). Omelette Singh (2022) extends
egg by treating the application of rules as a graph combinatorial problem.
Omelette shares the same aim as Aurora, to encode as many competitive plans in
the e-graph as possible within a memory limit. Omelette demonstrates graph
representation learning can be used to learn e-graphs for simple programming
languages, such as first-order logic, a core result which we use in Aurora to
tackle query rewrite.

##  9 Conclusion

In this research paper, we have demonstrated how to use equality saturation
for query rewrite and the advantages it brings over mainstream optimisers. We
have shown that equality saturation cannot trivially optimise queries of
moderate complexity due to high planning latency and memory considerations. To
address these challenges, we have introduced Aurora, a query optimiser that
guides the application of rewrite rules within the equality saturation
framework. By directing our attention to learning methods, we have devised an
expressive RL pipeline for solving the e-graph construction combinatorial
problem, which extends to other programming languages with minimal change.

###  9.1 Limitations & Future Work

The most pressing problem in Aurora is generalising across disparate queries.
Training an agent to solve NP-hard problems across sparse problem instances
(such as the space of possible query e-graphs) is not realistic. We leave as
future work the adoption of more advanced RL techniques such as population-
based RL Jaderberg et al. (2019) to improve generalisation.

## References

  * Begoli et al. (2018) Begoli, E., Camacho-Rodríguez, J., Hyde, J., Mior, M. J., and Lemire, D.  Apache calcite: A foundational framework for optimized query processing over heterogeneous data sources.  In _Proceedings of the 2018 International Conference on Management of Data_ , pp. 221–230, 2018.
  * Finance & Gardarin (1991) Finance, B. and Gardarin, G.  A rule-based query rewriter in an extensible dbms.  In _Proceedings. Seventh International Conference on Data Engineering_ , pp. 248–249. IEEE Computer Society, 1991.
  * Frydenberg (1990) Frydenberg, M.  The chain graph markov property.  _Scandinavian Journal of Statistics_ , pp. 333–353, 1990.
  * Hochreiter (1998) Hochreiter, S.  The vanishing gradient problem during learning recurrent neural nets and problem solutions.  _International Journal of Uncertainty, Fuzziness and Knowledge-Based Systems_ , 6(02):107–116, 1998.
  * Hochreiter & Schmidhuber (1997) Hochreiter, S. and Schmidhuber, J.  Long short-term memory.  _Neural computation_ , 9(8):1735–1780, 1997.
  * Huang et al. (2022) Huang, S., Dossa, R. F. J., Raffin, A., Kanervisto, A., and Wang, W.  The 37 implementation details of proximal policy optimization.  _The ICLR Blog Track 2023_ , 2022.
  * Jaderberg et al. (2019) Jaderberg, M., Czarnecki, W. M., Dunning, I., Marris, L., Lever, G., Castaneda, A. G., Beattie, C., Rabinowitz, N. C., Morcos, A. S., Ruderman, A., et al.  Human-level performance in 3d multiplayer games with population-based reinforcement learning.  _Science_ , 364(6443):859–865, 2019.
  * Kadkhodaei & Mahmoudi (2011) Kadkhodaei, H. and Mahmoudi, F.  A combination method for join ordering problem in relational databases using genetic algorithm and ant colony.  In _2011 IEEE International Conference on Granular Computing_ , pp. 312–317. IEEE, 2011.
  * Kipf & Welling (2016) Kipf, T. N. and Welling, M.  Semi-supervised classification with graph convolutional networks.  _arXiv preprint arXiv:1609.02907_ , 2016.
  * Marcus & Papaemmanouil (2019) Marcus, R. and Papaemmanouil, O.  Plan-structured deep neural network models for query performance prediction.  _arXiv preprint arXiv:1902.00132_ , 2019.
  * Marcus et al. (2019) Marcus, R., Negi, P., Mao, H., Zhang, C., Alizadeh, M., Kraska, T., Papaemmanouil, O., and Tatbul, N.  Neo: A learned query optimizer.  _arXiv preprint arXiv:1904.03711_ , 2019.
  * Medsker & Jain (2001) Medsker, L. R. and Jain, L.  Recurrent neural networks.  _Design and Applications_ , 5:64–67, 2001.
  * Nelson & Oppen (1980) Nelson, G. and Oppen, D. C.  Fast decision procedures based on congruence closure.  _Journal of the ACM (JACM)_ , 27(2):356–364, 1980.
  * Nieuwenhuis & Oliveras (2005) Nieuwenhuis, R. and Oliveras, A.  Proof-producing congruence closure.  In _Term Rewriting and Applications: 16th International Conference, RTA 2005, Nara, Japan, April 19-21, 2005. Proceedings 16_ , pp. 453–468. Springer, 2005.
  * Norris (1998) Norris, J. R.  _Markov chains_.  Number 2. Cambridge university press, 1998.
  * Piao & Liu (2019) Piao, C. and Liu, C. H.  Energy-efficient mobile crowdsensing by unmanned vehicles: A sequential deep reinforcement learning approach.  _IEEE Internet of Things Journal_ , 7(7):6312–6324, 2019.
  * Pirahesh et al. (1992) Pirahesh, H., Hellerstein, J. M., and Hasan, W.  Extensible/rule based query rewrite optimization in starburst.  _ACM Sigmod Record_ , 21(2):39–48, 1992.
  * Poess & Floyd (2000) Poess, M. and Floyd, C.  New tpc benchmarks for decision support and web commerce.  _ACM Sigmod Record_ , 29(4):64–71, 2000.
  * PostgreSQL (1996) PostgreSQL, B.  Postgresql.  _Web resource: http://www. PostgreSQL. org/about_ , 1996.
  * Puterman (2014) Puterman, M. L.  _Markov decision processes: discrete stochastic dynamic programming_.  John Wiley & Sons, 2014.
  * Risinglight (2022) Risinglight.  Risinglight.  https://github.com/risinglightdb/risinglight, 2022.  Accessed: 10 01, 2023.
  * Schulman et al. (2017) Schulman, J., Wolski, F., Dhariwal, P., Radford, A., and Klimov, O.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ , 2017.
  * Sherstinsky (2020) Sherstinsky, A.  Fundamentals of recurrent neural network (rnn) and long short-term memory (lstm) network.  _Physica D: Nonlinear Phenomena_ , 404:132306, 2020.
  * Singh (2022) Singh, Z.  Deep reinforcement learning for equality saturation.  https://www.cl.cam.ac.uk/~ey204/pubs/MPHIL_P3/2022_Zak.pdf, 2022.  Accessed: 06 09, 2022.
  * Sun & Li (2019) Sun, J. and Li, G.  An end-to-end learning-based cost estimator.  _arXiv preprint arXiv:1906.02560_ , 2019.
  * Sutton & Barto (2018) Sutton, R. S. and Barto, A. G.  _Reinforcement learning: An introduction_.  MIT press, 2018.
  * Tarjan (1975) Tarjan, R. E.  Efficiency of a good but not linear set union algorithm.  _Journal of the ACM (JACM)_ , 22(2):215–225, 1975.
  * Tate et al. (2009) Tate, R., Stepp, M., Tatlock, Z., and Lerner, S.  Equality saturation: a new approach to optimization.  In _Proceedings of the 36th annual ACM SIGPLAN-SIGACT symposium on Principles of programming languages_ , pp. 264–276, 2009.
  * Velickovic et al. (2017) Velickovic, P., Cucurull, G., Casanova, A., Romero, A., Lio, P., Bengio, Y., et al.  Graph attention networks.  _stat_ , 1050(20):10–48550, 2017.
  * Wang et al. (2020) Wang, Y. R., Hutchison, S., Leang, J., Howe, B., and Suciu, D.  Spores: sum-product optimization via relational equality saturation for large scale linear algebra.  _arXiv preprint arXiv:2002.07951_ , 2020.
  * Wang et al. (2022) Wang, Y. R., Khamis, M. A., Ngo, H. Q., Pichler, R., and Suciu, D.  Optimizing recursive queries with program synthesis.  _arXiv preprint arXiv:2202.10390_ , 2022.
  * Willsey et al. (2021) Willsey, M., Nandi, C., Wang, Y. R., Flatt, O., Tatlock, Z., and Panchekha, P.  Egg: Fast and extensible equality saturation.  _Proceedings of the ACM on Programming Languages_ , 5(POPL):1–29, 2021.
  * Yang et al. (2021) Yang, Y., Phothilimthana, P., Wang, Y., Willsey, M., Roy, S., and Pienaar, J.  Equality saturation for tensor graph superoptimization.  _Proceedings of Machine Learning and Systems_ , 3:255–268, 2021.
  * Yu et al. (2020) Yu, X., Li, G., Chai, C., and Tang, N.  Reinforcement learning with tree-lstm for join order selection.  In _2020 IEEE 36th International Conference on Data Engineering (ICDE)_ , pp. 1297–1308. IEEE, 2020.
  * Zhou et al. (2020) Zhou, J., Cui, G., Hu, S., Zhang, Z., Yang, C., Liu, Z., Wang, L., Li, C., and Sun, M.  Graph neural networks: A review of methods and applications.  _AI open_ , 1:57–81, 2020.
  * Zhou et al. (2021) Zhou, X., Li, G., Chai, C., and Feng, J.  A learned query rewrite system using monte carlo tree search.  _Proceedings of the VLDB Endowment_ , 15(1):46–58, 2021.