#project#cmu#LLVM#compiler#database

Compilation of Cursor Loops by Realizing Aggify "Correctly": Project Final Report for Course 15-745 Optimizing Compilers

Published on
•
21 mins read
•
––– views

Aggify series: Home · Proposal · Milestone · Final report (this page)

Our project poster

Group Information

Haoyu Zhang: haoyuzha

Yuchen Liu: yuchenl6

URL for Project

https://www.haoyu.dev/blog/aggify-home

1. Introduction

This is the technical report for our final project in CMU 15-745: Optimizing Compilers. We got hands-on with classic optimizer techniques by implementing the paper Aggify: Lifting the Curse of Cursor Loops using Custom Aggregates in DuckDB [1]. This section describes the problem, gives an overview of our solution, and lists our contributions.

1.1 Problem Statement

User-defined functions have had a reputation as performance poison since INGRES introduced them in the 1980s [2]. One optimization, inlining, changes that picture. [3] rewrites the procedural logic of a UDF as declarative SQL so the query optimizer can work across the UDF boundary, which brings many queries with UDF calls close to pure-SQL performance. [4] later made inlining applicable to almost any database that implements the SQL standard. But inlining does not apply to UDFs with cursor loops, the construct that iterates over a query's result set one row at a time.

Aggify [5] rewrites cursor loops as custom aggregate functions, which [3] and [4] can handle, completing the puzzle. No system has put Aggify into production, though. We wanted to be the first to implement it and see whether it can help repair the reputation of UDFs. Aggify also helps when the UDF is not inlined: the paper reports up to a 1000x speedup from rewriting the cursor loop and running the UDF sequentially.

1.2 Our Approach

We target PL/pgSQL UDFs and DuckDB. PL/pgSQL is static and syntactically conservative, which makes dataflow analysis easier; DuckDB is a system we already know well, and its open-source community has been very helpful.

To realize Aggify's idea, we first build a control flow graph (CFG) for the UDF, then run dataflow analyses on it (reaching definitions, live variables). Following Aggify, we use the results to extract the loop's logic into a custom aggregate and replace the cursor loop with a call to it. Section 2 has the details. Since DuckDB custom aggregates must be written in C++, the last step generates C++ for the extracted logic.

1.3 Contributions

We have three main contributions in this project.

  1. We are the first to correctly implement Aggify in DuckDB or any database system. Combining our work with [3], we allow DuckDB to support all the analytical PL/pgSQL UDFs. During implementation, we modified DuckDB to support Custom Aggregates without Combine() functions.
  2. We pioneered the compilation path from PL/pgSQL cursor loops to C++ code. This path includes applying compiler optimization passes to CFG of PL/pgSQL and a framework to efficiently add support for expression compilation to a vectorized database system. This contribution is shared with another ongoing project led by Sam Arch at CMU.
  3. We discovered a fundamental flaw in Aggify's approach, and proposed an optimized solution. This finding makes Aggify correct in the general case, which is crucial for Aggify to be practically adopted.

2. Compilation of Cursor Loops

DuckDB only accepts custom aggregates defined as a C++ class implementing certain methods. Generating correct C++ takes several steps; this section walks through them in order.

2.1 Construction of CFG for PL/pgSQL

To rewrite a cursor loop we first need the PL/pgSQL function in an intermediate representation. That takes two steps: parsing the function into an AST, and building a control flow graph from it.

  1. Parse the input SQL functions into an AST. We use JSON for the AST because it is easy to emit straight from the parser.
  2. Build the CFG from the JSON AST with a dedicated compiler that walks the tree and handles each node type (assignments, if statements, while loops, and so on).

The two steps are described below.

2.1.1 Parsing

Parsing turns each input SQL function into an abstract syntax tree (AST), stored as JSON for its simple structure. The ASTs are a structured representation of the syntax of the PL/pgSQL code.

2.1.2 CFG Construction

With the ASTs in hand, a dedicated compiler traverses them and builds the CFG, handling the node types found in PL/pgSQL: assignments, conditionals, and while and for loops.

CFG construction has several steps:

  1. Basic block creation: A basic block is a sequence of consecutive statements that control enters at the top and leaves at the bottom, with no branching in between. Appendix A.1 shows the result of this step.

  2. Handling each statement type: The compiler processes assignments, returns, conditionals and loops according to how control flows through them.

    • Assignments and returns map directly from the AST to the CFG.
    • Conditionals (if, elseif, else) become branches between blocks.
    • Loops (while, for) become blocks for the loop's entry, condition check and exit.

  3. Continuations: A continuation is the next step of computation. They are what let the compiler represent nested loops and conditional branches cleanly.

    The continuation we are using is represented as a struct in our code:

    struct Continuations {
  Continuations(BasicBlock *fallthrough, BasicBlock *loopHeader,
                BasicBlock *loopExit, BasicBlock *functionExit)
      : fallthrough(fallthrough), loopHeader(loopHeader), loopExit(loopExit),
        functionExit(functionExit) {}

  BasicBlock *fallthrough;
  BasicBlock *loopHeader;
  BasicBlock *loopExit;
  BasicBlock *functionExit;
};

    The struct packages the blocks control can jump to at each point in the program: the fallthrough block, the loop header and exit, and the function exit.

  4. Merging basic blocks: Our construction creates a new basic block for every instruction, so a clean-up pass merges blocks wherever it is legal, maximizing the instructions per block. Appendix A.1 and A.2 show a CFG before and after merging.

2.2 Analysis and Rewrite on the CFG

Given a cursor loop (Q, Δ), the goal is a C++ custom aggregate equivalent to the loop body Δ. We follow Aggify's approach, which combines dataflow analysis with CFG rewriting.

There are three required methods for a Custom Aggregate to be registered within DuckDB:

  • initialize : Constructs the State in raw memory.
  • update : Accumulate the arguments into the corresponding State.
  • finalize : Convert a State into a final value.

The sections below cover the aggregate's fields and each of the three methods, using the example below throughout. The function finds the supplier selling a given part at the lowest cost.

/* Code Example 1: Cursor Loop Rewrite Target */

/* Postgres version of UDF */
CREATE OR REPLACE FUNCTION MinCostSupp(pk bigint)
    RETURNS char(25)
    STABLE AS
$$
DECLARE
    key         bigint         := pk;
    fetchedCost decimal(15, 2);
    fetchedName char(25);
    minCost     decimal(25, 2) := 100000;
    val         char(25);
BEGIN
    FOR fetchedCost, fetchedName IN (SELECT PS_SUPPLYCOST, S_NAME
                                     FROM partsupp,
                                          supplier
                                     WHERE PS_PARTKEY = key
                                       AND PS_SUPPKEY = S_SUPPKEY)
        LOOP
            IF fetchedCost < minCost THEN
                minCost := fetchedCost;
                val := fetchedName;
            END IF;
        END LOOP;
    RETURN val;
END;
$$ LANGUAGE plpgsql;

/* UDF call */
EXPLAIN ANALYZE
SELECT P_PARTKEY, MinCostSupp(P_PARTKEY)
FROM part;

Appendix A.4 has the complete rewritten version of this query.

Fields

The first step is to find the minimal set of fields the custom aggregate needs. Liveness analysis gives VlocalV_{local}, the variables local to the loop body; the use-def chain gives VfetchV_{fetch}, the variables assigned by the FETCH statement, and VΔV_{\Delta}, every variable referenced in the loop body Δ. The fields VFV_F of the custom aggregate are then:

𝑉F=(𝑉Δ−(𝑉fetch∪𝑉local))∪𝑖𝑠𝐼𝑛𝑖𝑡𝑖𝑎𝑙𝑖𝑧𝑒𝑑.𝑉_F = (𝑉_{\Delta} − (𝑉_{fetch} ∪ 𝑉_{local})) ∪ {𝑖𝑠𝐼𝑛𝑖𝑡𝑖𝑎𝑙𝑖𝑧𝑒𝑑 }.

For Code Example 1 shown above, after applying the analyses, we have:

VΔ={key,val,minCost,fetchedCost,fetchedName}Vfetch={fetchedCost,fetchedName}Vlocal={}VF={val,minCost}\begin{aligned} V_{\Delta}& = \{key, val, minCost, fetchedCost, fetchedName\} \\ V_{fetch}& = \{fetchedCost, fetchedName\} \\ V_{local}& = \{\} \\ V_{F}& = \{val, minCost\} \end{aligned}

So val and minCost become the fields of the generated aggregate.

Initialize()

Initialization only sets the aggregate's is_initialized flag to false. The other fields depend on runtime values, so their initialization is deferred to the first call of update().

Update()

This method is the core of the aggregate: it accumulates across rows. Its parameters and body come from further dataflow analysis.

Parameters. PaccumP_{accum} is the set of variables used inside the loop body that have at least one reaching definition outside the loop, computed from reaching-definitions analysis. Formally, let VuseV_{use} be the variables used inside the loop body. For each v∈Vusev \in V_{use}, let UCL(v)U_{CL}(v) be the uses of vv inside the cursor loop CL, and for each use u∈UCL(v)u \in U_{CL}(v) let RD(u)RD(u) be the definitions of vv that reach uu. Define R(v)R(v) as follows [2]:

R(v)={1,if ∃d∈RD(u)∣d is not in the loop.0,otherwise.R(v) = \begin{cases} 1, & \text{if } \exists d \in RD(u) \mid d \text{ is not in the loop.} \\ 0, & \text{otherwise}.\end{cases}

The parameters of update() are then:

Paccum={v∣v∈Vuse∧R(v)==1}P_{\text{accum}} = \{ v \mid v \in V_{\text{use}} \land R(v) == 1 \}

Method body. The body of update() has two parts: field initialization and the loop body. The fields that need initializing are:

𝑉init=𝑃accum−𝑉fetch𝑉_{init} = 𝑃_{accum} − 𝑉_{fetch}

Sample code for the field-initialization block:

if (!state.is_initialized) {
    state.min_cost = pmin_cost;
    state.is_initialized = true;
}

As noted above, the is_initialized flag (statically false) guards the runtime initialization: the first call to update() assigns the runtime-provided values, and later calls skip the block.

Again, for the same example above, we have:

Paccum={minCost,fetchedCost,fetchedName}Vinit={minCost}\begin{aligned} P_{accum}& = \{minCost, fetchedCost, fetchedName\} \\ V_{init}& = \{minCost\} \\ \end{aligned}

The loop body becomes the equivalent accumulation step:

if (current_cost < state.min_cost) {
    state.min_cost = current_cost;
    state.supp_name = current_name;
}

Finalize()

This method returns the part of the state that is the query result, here the supplier's name.

std::string finalize() {
	return state.supp_name;
}

The generated aggregate is compiled into a dynamic library and linked into DuckDB at runtime, after which queries can use it like any built-in aggregate.

Modify original UDF CFG to use the Custom Aggregate

The last step wraps the custom aggregate in a UDF. The wrapper receives the VFV_F variables as local parameters and invokes the aggregate with them; once registered in DuckDB, callers use the rewritten aggregate through this wrapper. The complete code is in Appendix A.4.

2.3 Compilation of the CFG to C++

This section describes how the loop-body CFG from Section 2.2 is compiled. There are two kinds of compilation: the CFG structure at the basic-block level is procedural code, and the instructions inside the blocks are SQL statements. The two are handled differently.

2.3.1 Procedural Construct Compilation

Once the UDF is a CFG (basic blocks ending in branch instructions), translating it to C++ is straightforward.

entry:
	jmp L0
L0:
	INT x = y;
	jmp L1;
L1:
	INT z = y;
	br (z > 0) L2, L0;
L2:
	jmp exit
exit:

Each basic block becomes a label in the C++ program. Each instruction in the block becomes a C++ statement via SQL statement compilation (next section), and the block's terminating branch becomes a goto or conditional goto. The result is a C++ program equivalent to the CFG.

2.3.2 SQL Statement Compilation

Every procedural construct in PL/pgSQL has a C++ counterpart, but a SQL statement that uses built-in operators or expressions has to call DuckDB's implementation of that logic.

INT a;
LONG b;
DOUBLE c;
...
c := a + b

For the assignment c := a + b above, we want the generated C++ to read c = duckdb::AddLongDouble(a, b). The trick that made DuckDB compilation-friendly quickly was to attach code-generation metadata to the logical plan, the tree DuckDB builds when it parses and binds a query. For this example the query is SELECT a + b FROM tmp_table, where tmp_table exists only to give DuckDB the types of a and b.

We modified the DuckDB binder so that when a vectorized operator is bound to a computation node, it records extra information about that operator: the name of the corresponding scalar function, its template instantiation if it has one, and so on.

Finally, we walk back through the logical plan and use the metadata to generate C++. The generated code calls DuckDB's internal functions, so it is guaranteed to produce the same result DuckDB would. This is a form of partial evaluation, in the spirit of the Futamura projections [7].

3. Fix to Aggify's Design Flaw

While trying to reproduce Aggify's experimental results, we found a major design flaw. Take one of the paper's own test cases, Figure 1(a); Figure 1(b) shows the result of applying Aggify to it. custom_count behaves like the native count except that it has no Combine() method, so it cannot run in parallel. Aggify produces the correct result for this test case as written. Change the initial value of val to a non-zero number, though, and the transformation breaks. If the cursor (pink) query returns nothing, running the loop normally leaves val untouched and the UDF returns 100 (Figure 1(c)), but the custom aggregate overwrites val with 0 (Figure 1(d)), which disagrees with the original UDF.

Figure 1. (b) is the result of Aggify for (a); (d) is the result of Aggify for (c)

The fix is simple. The problem only appears when the cursor query returns nothing, so we guard the custom aggregate with a precondition that the query is non-empty. Figure 2 shows the correct transformation of Figure 1(c). It introduces a second query that repeats work.

Figure 2. Examples of unoptimized fix to Aggify v.s. optimized fix to Aggify

We hoped DuckDB's common subexpression elimination, which finds repeated subqueries and runs them once, would absorb the duplicate. Experiments said otherwise. After trying several equivalent but syntactically different forms, the best correction turned out to be a CASE WHEN in place of IF EXISTS. Figure 2(b) is this optimized form of Figure 2(a); it needs one fewer table scan. Appendix A.5 has a complete example.

4. Experiments

To run the paper's test cases at all, we had to modify DuckDB to run custom aggregates without Combine() functions.

4.1 Preparation of DuckDB for the Experiment

A custom aggregate's Combine() method defines how to merge the results of aggregating two groups into one. With it, the execution engine can split a large group into partitions, aggregate each in parallel, and combine the results. Inferring Combine() from a cursor loop is out of scope for the Aggify paper, and we suspect it is impossible in general. DuckDB, however, requires every custom aggregate to provide one. So we modified query execution to run custom aggregates single-threaded, which removes the need for Combine().

4.2 Setup

The experiment uses the TPC-H dataset at scale factor 10 on a MacBook with an Apple M1 Pro and 32 GB of RAM. We ran each test ten times and report the average.

4.3 Results and Evaluation

We first compare the original Aggify (without our fix) against built-in aggregates, which represent the best possible performance.

Figure 3. Performance comparison of Aggify using custom aggregate (workload1) or built-in aggregate (workload2), with multi-threading enabled.

Figure 4. Performance comparison of Aggify using custom aggregate (workload1) or built-in aggregate (workload2), with only one thread used.

As expected, custom aggregates generally do not perform as well as built-in ones, up to 0.18 times slower on the DiscountRevenue workload in Figure 3. The gap comes from the lack of parallelism: Aggify's custom aggregates have no Combine() and so run sequentially, while built-in aggregates use every available thread. To check this, we ran both workloads again with DuckDB limited to one thread, and Aggify performed about the same as the built-in aggregate (Figure 4). Using more sophisticated program analysis to infer Combine() for cursor loops would be an interesting research direction.

With the fix applied, Aggify gets slower: the naive fix runs the same subquery twice, so its time roughly doubles. Replacing IF EXISTS with CASE WHEN narrows the gap (Figure 5). On PromoRevenue and DiscountRevenue the corrected version is even faster than the original, because the CASE WHEN condition filters out many groups before the expensive custom aggregate runs.

Figure 5. Performance comparison of the original Aggify with two of its fixes

5. Conclusion

In this project, we implemented Aggify from end to end in DuckDB. We also addressed a bug in Aggify. Future efforts in this topic can focus on applying well-known loop optimization techniques like Loop Invariant Code Motion to the cursor loop. It is also interesting to see SQL specific optimizations for loops.

Learnings

The report above is the group's. These takeaways are mine.

  • Implementing a paper is the best way to review it. The accumulator-initialization flaw in Section 3 is invisible when you read Aggify and obvious the moment you run its own example with a non-zero initial value. Reproduction found something peer review did not.
  • The target system's constraints drive the design more than the algorithm does. DuckDB's requirement that custom aggregates define Combine(), and our decision to relax it rather than fake a combiner, shaped the generated code, the parallelism story and the benchmark methodology far more than anything in the Aggify paper.
  • Textbook dataflow analysis transfers directly. Reaching definitions and live-variable analysis on the PL/pgSQL CFG did exactly what the compiler course promised. The work was in building a trustworthy CFG (basic-block merging, loop structure) before running them, not in the analyses themselves.
  • A correctness fix has to be benchmarked like a feature. Our first fix, an IF EXISTS precondition, was correct but doubled the scan; the CASE WHEN form gave the same guarantee for one scan. "Correct" and "not a regression" are two separate experiments.
  • Generated code needs a specification, not just tests. Mapping procedural constructs and SQL statements to C++ went smoothly only after we wrote down the translation rules (Section 2.3). Before that, every new test case surprised us.

6. References

[1] M. Raasveldt and H. Mühleisen, "DuckDB: an Embeddable Analytical Database," in Proceedings of the 2019 International Conference on Management of Data, Amsterdam Netherlands: ACM, Jun. 2019, pp. 1981–1984. doi: 10.1145/3299869.3320212.

[2] J. Ong et al. Implementation of data abstraction in the relational database system ingres. SIGMOD Record, 1983.

[3] S. Gupta, S. Purandare, and K. Ramachandra, "Aggify: Lifting the Curse of Cursor Loops using Custom Aggregates," in Proceedings of the 2020 ACM SIGMOD International Conference on Management of Data, Portland OR USA: ACM, Jun. 2020, pp. 559–573. doi: 10.1145/3318464.3389736.

[4] K. Ramachandra, K. Park, K. V. Emani, A. Halverson, C. Galindo-Legaria, and C. Cunningham, "Froid: optimization of imperative programs in a relational database," Proc. VLDB Endow., vol. 11, no. 4, pp. 432–444, Dec. 2017, doi: 10.1145/3186728.3164140.

[5] D. Hirn and T. Grust, "One WITH RECURSIVE is Worth Many GOTOs," in Proceedings of the 2021 International Conference on Management of Data, Virtual Event China: ACM, Jun. 2021, pp. 723–735. doi: 10.1145/3448016.3457272.

[6] M. Sichert and T. Neumann, "User-defined operators: efficiently integrating custom algorithms into modern databases," Proc. VLDB Endow., vol. 15, no. 5, pp. 1119–1131, Jan. 2022, doi: 10.14778/3510397.3510408.

[7] Y. Futamura, "Partial Evaluation of Computation Process --- An approach to a Compiler-Compiler", Transactions of the Institute of Electronics and Communications Engineers of Japan, 54-C: 721–728, 1971

7. Appendix

A.1 - CFG Before Basic Block Merging


A.2 - CFG After Basic Block Merging


A.3 - Cursor Loop Query Before Rewriting

/* Postgres version of UDF */
CREATE OR REPLACE FUNCTION MinCostSupp(pk bigint)
    RETURNS char(25)
    STABLE AS
$$
DECLARE
    key         bigint         := pk;
    fetchedCost decimal(15, 2);
    fetchedName char(25);
    minCost     decimal(25, 2) := 100000;
    val         char(25);
BEGIN
    FOR fetchedCost, fetchedName IN (SELECT PS_SUPPLYCOST, S_NAME
                                     FROM partsupp,
                                          supplier
                                     WHERE PS_PARTKEY = key
                                       AND PS_SUPPKEY = S_SUPPKEY)
        LOOP
            IF fetchedCost < minCost THEN
                minCost := fetchedCost;
                val := fetchedName;
            END IF;
        END LOOP;
    RETURN val;
END;
$$ LANGUAGE plpgsql;

/* UDF call */
EXPLAIN ANALYZE
SELECT P_PARTKEY, MinCostSupp(P_PARTKEY)
FROM part;

A.4 - Cursor Loop Query After Rewriting

-- Internal states
CREATE TYPE min_cost_state AS (
    min_cost DECIMAL(15, 2),
    supp_name CHAR(25),
    is_initialized BOOLEAN
);

-- Update function
CREATE OR REPLACE FUNCTION f_update(
    current_cost DECIMAL(15, 2), 
    current_name CHAR(25), 
    pmin_cost DECIMAL(15, 2), 
)
RETURNS min_cost_state
LANGUAGE plpgsql AS $$
BEGIN
    IF NOT state.is_initialized THEN
        state.min_cost := pmin_cost;
        state.is_initialized := true;
    END IF;
    IF current_cost < state.min_cost THEN
        state.min_cost := current_cost;
        state.supp_name := current_name;
    END IF;
    RETURN state;
END;
$$;

-- Finalize function
CREATE OR REPLACE FUNCTION f_finalize()
RETURNS CHAR(25)
LANGUAGE plpgsql AS $$
BEGIN
    RETURN state.supp_name;
END;
$$;

-- Aggregate
CREATE AGGREGATE MinCostSuppAggregate ()
(
    BASETYPE = DECIMAL(15, 2),
    STYPE = min_cost_state,
    SFUNC = f_update,
    FINALFUNC = f_finalize
    stype.supp_name = '',
    stype.is_initialized = false
);

-- UDF, Wrapper
CREATE OR REPLACE FUNCTION MinCostSuppWithCustomAgg(pk BIGINT)
RETURNS CHAR(25)
LANGUAGE plpgsql AS $$
DECLARE
    key         bigint         := pk;
    val         CHAR(25);
    pmin_cost   DECIMAL(25, 2) := 100000;
BEGIN
    SELECT MinCostSuppAggregate(PS_SUPPLYCOST, S_NAME, pmin_cost)
    FROM (SELECT PS_SUPPLYCOST, S_NAME
                                     FROM partsupp,
                                          supplier
                                     WHERE PS_PARTKEY = key
                                       AND PS_SUPPKEY = S_SUPPKEY) S
    INTO val;
    RETURN val;
END;
$$;

-- UDF Call
SELECT P_PARTKEY, MinCostSuppWithCustomAgg(P_PARTKEY)
FROM part;

A.5 - Cursor Loop Query After Rewriting and Correction

-- Type your code here, or load an example.
CREATE OR REPLACE FUNCTION MinCostSuppWithCustomAgg(pk BIGINT)
RETURNS CHAR(25)
LANGUAGE plpgsql AS $$
DECLARE
    key         bigint         := pk;
    val         CHAR(25);
    pmin_cost   DECIMAL(25, 2) := 100000;
BEGIN
    if exists (SELECT PS_SUPPLYCOST, S_NAME
                                     FROM partsupp,
                                          supplier
                                     WHERE PS_PARTKEY = key
                                       AND PS_SUPPKEY = S_SUPPKEY) then
        val := (SELECT MinCostSuppAggregate(PS_SUPPLYCOST, S_NAME, pmin_cost)
                FROM (SELECT PS_SUPPLYCOST, S_NAME
                                        FROM partsupp,
                                            supplier
                                        WHERE PS_PARTKEY = key
                                        AND PS_SUPPKEY = S_SUPPKEY) tmp);
    end if;                                    
    RETURN val;
END;
$$;

-- Type your code here, or load an example.
CREATE OR REPLACE FUNCTION MinCostSuppWithCustomAgg(pk BIGINT)
RETURNS CHAR(25)
LANGUAGE plpgsql AS $$
DECLARE
    key         bigint         := pk;
    val         CHAR(25);
    pmin_cost   DECIMAL(25, 2) := 100000;
BEGIN
    val := (SELECT case when count(*) > 0 then
            MinCostSuppAggregate(PS_SUPPLYCOST, S_NAME, pmin_cost) else val end
            FROM (SELECT PS_SUPPLYCOST, S_NAME
                                    FROM partsupp,
                                        supplier
                                    WHERE PS_PARTKEY = key
                                    AND PS_SUPPKEY = S_SUPPKEY) tmp);                         
    RETURN val;
END;
$$;

-- Q:
SELECT P_PARTKEY, MinCostSuppWithCustomAgg(P_PARTKEY)
FROM part;