feat: Implement balanced query planner

[bryanhuhta]

Overview

Introduces a new query planning algorithm to make up for some of the deficiencies of the old algorithm. This can be toggled by enabled -query-frontend.query-planner-strategy=balanced. By default the old algorithm will be used.

Problem

The current query planner has an algorithm which greedily assigns the maximum number of child nodes to a merge node. If a merge node can have a maximum of 20 children and the query creates 30 children, there will be one merge node that merges 20 children and one merge node that merges 10 children. This algorithm results in frequently imbalanced queries where there tends to be a merge node (or multiple) which contribute almost no work to the query execution.

Consider the following query parameters:

blocks    = 7
maxReads  = 2
maxMerges = 3

The current query planner will generate the following plan:

7 blocks,  maxReads=2 -> leaves: [1,2] [3,4] [5,6] [7]
4 leaves, maxMerges=3 -> merges: [L1,L2,L3]        [L4]

graph TD
    Root["MERGE<br/>2 children"]
    M1["MERGE<br/>3 children, 6 blocks"]
    M2["MERGE<br/>1 child, 1 block"]
    R1["READ<br/>id: 1,2"]
    R2["READ<br/>id: 3,4"]
    R3["READ<br/>id: 5,6"]
    R4["READ<br/>id: 7"]
    Root --> M1
    Root --> M2
    M1 --> R1
    M1 --> R2
    M1 --> R3
    M2 --> R4

Loading

In this case, one of the first level merge nodes will only have a single leaf node. We pay the costs of network hops to a merge node that provides no benefit to the execution of the query except to pass along the single leaf node that it was assigned. In this example, the bulk of the query work is being done by a single merge node.

Solution

This PR implements a new query planning algorithm that evenly balances out the work across all merge and leaf nodes. Instead of greedily assigning children to a merge node until the node is at its limit, the algorithm evenly distributes leaf nodes across all merge nodes at each level of the tree. Taking the earlier example:

7 blocks,  maxReads=2 -> leaves: [1] [2,3] [4,5] [6,7]
4 leaves, maxMerges=3 -> merges: [L1,L2]   [L3,L4]

graph TD
    Root["MERGE<br/>2 children"]
    M1["MERGE<br/>2 children, 3 blocks"]
    M2["MERGE<br/>2 children, 4 blocks"]
    R1["READ<br/>id: 1"]
    R2["READ<br/>id: 2,3"]
    R3["READ<br/>id: 4,5"]
    R4["READ<br/>id: 6,7"]
    Root --> M1
    Root --> M2
    M1 --> R1
    M1 --> R2
    M2 --> R3
    M2 --> R4

Loading

None of the merge nodes are at the capacity of 3, and the leaf nodes are split evenly between them. Note that the algorithm balances the number of leaf nodes assigned to each subtree, not the raw block count directly. Since blocks are already spread as evenly as possible across the leaves, this keeps block counts close too in practice, but a merge node's total block count can still end up a bit higher or lower than its siblings, especially in larger trees. What the algorithm actually guarantees is that every leaf node sits at the same depth, and every merge node's children are split into an equal, or as close to equal as possible, number of leaf nodes.

⚠️ Very detailed Claude summary ⚠️

BuildBalanced builds the tree in two passes. First it slices the block list into leaf READ nodes, spreading any remainder evenly across leaves instead of dumping it onto the last one. Second it hands those leaves to a recursive helper, buildMergeTree, which wraps them in MERGE nodes so that no branch of the tree ends up carrying much more work than its siblings.

The core idea

Each call to buildMergeTree receives a slice of nodes — leaves on the first call, subtrees on every call after that — and has exactly one job: build a single MERGE node whose Children field holds at most maxMerges entries.

If the incoming nodes already fit within maxMerges, there's nothing to compute: they all become direct children and the call returns. The interesting case is when they don't fit, e.g. 21 nodes with maxMerges = 20. The 21 nodes can't all be direct children of one MERGE node, so some of them first have to be grouped into buckets, each bucket wrapped in its own MERGE node, so that a bucket counts as a single child slot instead of many:

graph TD
    subgraph illegal ["✗ illegal — 21 direct children, too many"]
        R0["MERGE<br/>21 children"]
        R0 --> n0[n0]
        R0 --> n1[n1]
        R0 --> dots1[…]
        R0 --> n20[n20]
    end

Loading

graph TD
    subgraph legal ["✓ legal — 2 children, each a bucket"]
        R1["MERGE<br/>2 children"]
        B1["MERGE<br/>bucket: n0…n10"]
        B2["MERGE<br/>bucket: n11…n20"]
        R1 --> B1
        R1 --> B2
    end

Loading

To pick a bucket size, the function walks up the powers of maxMerges — 1, maxMerges, maxMerges², … — and stops at the first one small enough that grouping every node into buckets of that size needs maxMerges buckets or fewer:

capacity = 1   → 21 buckets needed → too many ✗
capacity = 20  →  2 buckets needed → fits    ✓

Once the bucket size decides how many buckets there are, the actual split across those buckets is the part that gives the function its name. It'd be easy, and wrong, to fill each bucket up to that size greedily and let the last one take whatever remainder is left over — that's exactly what the older, unbalanced Build function does. Instead, a small helper (balanceGroupItems) divides the nodes across the buckets as evenly as possible, so bucket sizes never differ by more than one, which is why the legal diagram above splits 11 and 10 rather than 20 and 1.

Each bucket then recurses through buildMergeTree on its own, with no memory of what its parent or siblings decided. Repeating this at every level is what keeps every leaf at the same depth and every branch carrying a close-to-even share of the leaves below it, all the way down the tree.

Working through the example above

Picking back up the 7-block example from the Solution section, the 4 leaves get handed to buildMergeTree like so:

buildMergeTree([L1,L2,L3,L4], maxMerges=3)
  capacity = 1 → 4 buckets needed → too many (4 > 3) ✗
  capacity = 3 → 2 buckets needed → fits            ✓
  split 4 nodes into 2 balanced buckets: [L1,L2] and [L3,L4]

graph TD
    Root["MERGE"]
    M1["MERGE<br/>bucket: L1, L2"]
    M2["MERGE<br/>bucket: L3, L4"]
    Root --> M1
    Root --> M2
    M1 --> L1["L1: id 1"]
    M1 --> L2["L2: id 2,3"]
    M2 --> L3["L3: id 4,5"]
    M2 --> L4["L4: id 6,7"]

Loading

Neither bucket ends up close to the maxMerges limit of 3, but both come out the same size, which is the point: the tree stays shallow and even rather than lopsided.

Why this holds up at scale

The same bucket-size search runs independently at every level, so large inputs don't need any special casing. 8000 leaves with maxMerges = 20, for example, resolves into 3 levels of 20-way merges (20 × 20 × 20 = 8000), with each level discovered by the exact same search and no coordination needed between levels. A more detailed walkthrough of this recursion, including traces at several leaf counts, was used as background research while writing this section.

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

Want higher recall? High effort reviews run extra passes and find more bugs. A team admin can switch effort levels in the Cursor dashboard.

^{Reviewed by Cursor Bugbot for commit 1f35a8b. Configure here.}

[aleks-p]

I like where this is going, the imbalanced-plan problem is real. I think we can get a similar result with a bit less code, though.

What if instead of front-loading the remaining blocks, we interleave them (using the block index)? For 37 blocks, maxReads=4:

method	leaves	prefix sums	split
front-load	[4 4 4 4 4 4 4 3 3 3]	4 8 12 16 [20] 24 28 31 34 37	20/17
interleave	[3 4 4 3 4 4 3 4 4 4]	[3 7 11 14 [18] 22 25 29 33 37]	18/19

Then we don't need weight tracking or the second pass, the split closest to
even by node count is automatically the split closest to even by block count.

Sketch at aleks/balanced-query-plan-alt, on top of your branch and not meant to be
merged. Let me know what you think!

[bryanhuhta]

I created a test to compare the maximum imbalance between merge nodes for the 3 query plan options from 1 to 64,000 blocks. Aleks's algorithm was by far the most stable, never drifting away from a query plan that has more than 4/5 block imbalance. Both of the balanced algorithms far outperform the current query plan builder, but the plan proposed in this PR begins to suffer as the number of blocks increases.

While 64k blocks is unlikely to occur in practice, I'll adopt @aleks-p's proposal since its both simpler to understand and produces more reliable results. Most importantly, this algorithm doesn't suffer from spikes of imbalance, which eliminates another component of unpredictability from our system.