From d9b590a11184f9e61d47d393aaef378a18ba0a0b Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Wed, 12 Aug 2026 15:19:17 -0700
Subject: [PATCH 1/9] Implement balanced query planner

---
 .../readpath/queryfrontend/query_frontend.go  |   4 +-
 pkg/querybackend/queryplan/query_plan.go      | 139 +++++++++++++++++-
 pkg/querybackend/queryplan/query_plan_test.go |  59 +++++---
 .../{empty.txt => balanced_empty.txt}         |   0
 ...e_leaf.txt => balanced_exact_one_leaf.txt} |   0
 .../testdata/balanced_forced_equal_split.txt  |  35 +++++
 ..._depth_2.txt => balanced_full_depth_2.txt} |   0
 ...lanced_full_merge_vs_single_leaf_merge.txt |  45 ++++++
 ...es.txt => balanced_invalid_max_merges.txt} |   0
 ...ads.txt => balanced_invalid_max_reads.txt} |   0
 .../testdata/balanced_just_over_depth_2.txt   |  14 ++
 ...le_block.txt => balanced_single_block.txt} |   0
 .../testdata/balanced_three_way_split.txt     |  44 ++++++
 .../testdata/balanced_twenty_five_blocks.txt  |  46 ++++++
 ...two_leaves.txt => balanced_two_leaves.txt} |   0
 .../queryplan/testdata/build_empty.txt        |   0
 .../testdata/build_exact_one_leaf.txt         |   3 +
 .../testdata/build_forced_equal_split.txt     |  35 +++++
 .../queryplan/testdata/build_full_depth_2.txt |  10 ++
 .../build_full_merge_vs_single_leaf_merge.txt |  45 ++++++
 .../testdata/build_invalid_max_merges.txt     |   0
 .../testdata/build_invalid_max_reads.txt      |   0
 ...epth_2.txt => build_just_over_depth_2.txt} |   0
 .../queryplan/testdata/build_single_block.txt |   2 +
 .../testdata/build_three_way_split.txt        |  44 ++++++
 ...locks.txt => build_twenty_five_blocks.txt} |   0
 .../queryplan/testdata/build_two_leaves.txt   |   6 +
 .../full_merge_vs_single_leaf_merge.txt       |  45 ++++++
 28 files changed, 546 insertions(+), 30 deletions(-)
 rename pkg/querybackend/queryplan/testdata/{empty.txt => balanced_empty.txt} (100%)
 rename pkg/querybackend/queryplan/testdata/{exact_one_leaf.txt => balanced_exact_one_leaf.txt} (100%)
 create mode 100644 pkg/querybackend/queryplan/testdata/balanced_forced_equal_split.txt
 rename pkg/querybackend/queryplan/testdata/{full_depth_2.txt => balanced_full_depth_2.txt} (100%)
 create mode 100644 pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt
 rename pkg/querybackend/queryplan/testdata/{invalid_max_merges.txt => balanced_invalid_max_merges.txt} (100%)
 rename pkg/querybackend/queryplan/testdata/{invalid_max_reads.txt => balanced_invalid_max_reads.txt} (100%)
 create mode 100644 pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt
 rename pkg/querybackend/queryplan/testdata/{single_block.txt => balanced_single_block.txt} (100%)
 create mode 100644 pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt
 rename pkg/querybackend/queryplan/testdata/{two_leaves.txt => balanced_two_leaves.txt} (100%)
 create mode 100644 pkg/querybackend/queryplan/testdata/build_empty.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/build_exact_one_leaf.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/build_forced_equal_split.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/build_full_depth_2.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/build_full_merge_vs_single_leaf_merge.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/build_invalid_max_merges.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/build_invalid_max_reads.txt
 rename pkg/querybackend/queryplan/testdata/{just_over_depth_2.txt => build_just_over_depth_2.txt} (100%)
 create mode 100644 pkg/querybackend/queryplan/testdata/build_single_block.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/build_three_way_split.txt
 rename pkg/querybackend/queryplan/testdata/{twenty_five_blocks.txt => build_twenty_five_blocks.txt} (100%)
 create mode 100644 pkg/querybackend/queryplan/testdata/build_two_leaves.txt
 create mode 100644 pkg/querybackend/queryplan/testdata/full_merge_vs_single_leaf_merge.txt

diff --git a/pkg/frontend/readpath/queryfrontend/query_frontend.go b/pkg/frontend/readpath/queryfrontend/query_frontend.go
index eb2796c2ae..ef929aecae 100644
--- a/pkg/frontend/readpath/queryfrontend/query_frontend.go
+++ b/pkg/frontend/readpath/queryfrontend/query_frontend.go
@@ -228,7 +228,7 @@ func (q *QueryFrontend) doQuery(
 	endTime := time.UnixMilli(req.EndTime)
 	queryWindow := endTime.Sub(startTime).Round(time.Second)
 	traceID, _ := tracing.ExtractTraceID(ctx)
-	logArgs := []interface{}{
+	logArgs := []any{
 		"msg", "query weight",
 		"trace_id", traceID,
 		"tenant", strings.Join(tenants, ","),
@@ -256,7 +256,7 @@ func (q *QueryFrontend) doQuery(
 	})
 	xrandMutex.Unlock()
 	// TODO(kolesnikovae): Should be dynamic.
-	p := queryplan.Build(blocks, 4, 20)
+	p := queryplan.BuildBalanced(blocks, 4, 20)
 
 	backend := q.querybackend
 	if backendC != nil {
diff --git a/pkg/querybackend/queryplan/query_plan.go b/pkg/querybackend/queryplan/query_plan.go
index 475585d9d0..8e94f1d042 100644
--- a/pkg/querybackend/queryplan/query_plan.go
+++ b/pkg/querybackend/queryplan/query_plan.go
@@ -1,6 +1,8 @@
 package queryplan
 
 import (
+	"math"
+
 	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 )
@@ -36,10 +38,7 @@ func Build(
 	leafNodeCount := (len(blocks) + maxReads - 1) / maxReads
 	nodes := allocateContiguous[queryv1.QueryNode](leafNodeCount)
 	for start, idx := 0, 0; start < len(blocks); start, idx = start+maxReads, idx+1 {
-		end := start + maxReads
-		if end > len(blocks) {
-			end = len(blocks)
-		}
+		end := min(start+maxReads, len(blocks))
 		nodes[idx].Type = queryv1.QueryNode_READ
 		nodes[idx].Blocks = blocks[start:end]
 	}
@@ -50,10 +49,7 @@ func Build(
 		mergeNodes := allocateContiguous[queryv1.QueryNode](mergeNodeCount)
 
 		for start, idx := 0, 0; start < len(nodes); start, idx = start+maxMerges, idx+1 {
-			end := start + maxMerges
-			if end > len(nodes) {
-				end = len(nodes)
-			}
+			end := min(start+maxMerges, len(nodes))
 			mergeNodes[idx].Type = queryv1.QueryNode_MERGE
 			mergeNodes[idx].Children = nodes[start:end:end]
 		}
@@ -66,6 +62,133 @@ func Build(
 	}
 }
 
+func BuildBalanced(blocks []*metastorev1.BlockMeta, maxReads int, maxMerges int) *queryv1.QueryPlan {
+	if len(blocks) == 0 || maxReads < 1 || maxMerges < 2 {
+		return new(queryv1.QueryPlan)
+	}
+
+	readNodeCount := int(math.Ceil(float64(len(blocks)) / float64(maxReads)))
+	nodes := allocateContiguous[queryv1.QueryNode](readNodeCount)
+	weights := make([]int, readNodeCount)
+
+	// Build the read nodes, balancing blocks across all read nodes. We also
+	// record the number of blocks in each read node to be used later when
+	// assigning read nodes to merge nodes.
+	for start, idx := 0, 0; idx < readNodeCount; idx++ {
+		size := balancedGroupSize(len(blocks), readNodeCount, idx)
+		end := start + size
+
+		nodes[idx].Type = queryv1.QueryNode_READ
+		nodes[idx].Blocks = blocks[start:end:end]
+		weights[idx] = size
+		start = end
+	}
+
+	// Build the merge nodes. We assign merge children to merge nodes based on the
+	// number of blocks being merged at each level. We want each merge node to
+	// merge approximately the same number of blocks as its siblings.
+	for len(nodes) > 1 {
+		mergeNodeCount := int(math.Ceil(float64(len(nodes)) / float64(maxMerges)))
+		mergeNodes := allocateContiguous[queryv1.QueryNode](mergeNodeCount)
+		mergeWeights := make([]int, mergeNodeCount)
+
+		// workloadGroupSizes calculates how many children each merge node should
+		// have at this level.
+		groupSizes := distributeChildNodes(weights, mergeNodeCount, maxMerges)
+		start := 0
+		for idx, size := range groupSizes {
+			end := start + size
+			mergeNodes[idx].Type = queryv1.QueryNode_MERGE
+			mergeNodes[idx].Children = nodes[start:end:end]
+
+			for _, weight := range weights[start:end] {
+				mergeWeights[idx] += weight
+			}
+			start = end
+		}
+
+		nodes = mergeNodes
+		weights = mergeWeights
+	}
+
+	return &queryv1.QueryPlan{
+		Root: nodes[0],
+	}
+}
+
+func distributeChildNodes(childNodeBlocks []int, mergeNodeCount int, maxMergeNodeSize int) []int {
+	// This slice contains the number of child nodes each merge node should be
+	// allocated.
+	mergeNodeChildren := make([]int, mergeNodeCount)
+
+	remainingTotalBlockCount := 0
+	for _, count := range childNodeBlocks {
+		remainingTotalBlockCount += count
+	}
+
+	currentChildNodeIdx := 0
+	for mergeNodeIdx := range mergeNodeCount {
+		// Calculate the remaining merge nodes and child nodes we have left to
+		// allocate.
+		remainingMergeNodes := mergeNodeCount - mergeNodeIdx
+		remainingChildNodes := len(childNodeBlocks) - currentChildNodeIdx
+
+		// Calculate the minimum and maximum child nodes this merge node can have.
+		minChildNodeCount := max(1, remainingChildNodes-(remainingMergeNodes-1)*maxMergeNodeSize)
+		maxChildNodeCount := min(maxMergeNodeSize, remainingChildNodes-(remainingMergeNodes-1))
+
+		// Calculate the ideal number of blocks we could give to this merge node
+		// (and all subsequent merge nodes) to evenly distribute blocks amongst
+		// them. This number likely will not be a whole number.
+		targetBlockCount := float64(remainingTotalBlockCount) / float64(remainingMergeNodes)
+
+		bestChildNodeCount := minChildNodeCount
+		bestBlockCount := 0
+		bestDistance := float64(0)
+
+		// We have to allocate at least the minimum number of child nodes to this
+		// merge node, so we do that now. we then calcualte how far away we are from
+		// the ideal allocation.
+		for _, b := range childNodeBlocks[currentChildNodeIdx : currentChildNodeIdx+minChildNodeCount] {
+			bestBlockCount += b
+		}
+		bestDistance = math.Abs(float64(bestBlockCount) - targetBlockCount)
+
+		// Now we expand the number of child nodes we give to this merge node to see
+		// if we can get closer to the ideal block count.
+		candidateBlockCount := bestBlockCount
+		for nodeCount := minChildNodeCount + 1; nodeCount <= maxChildNodeCount; nodeCount++ {
+			candidateBlockCount += childNodeBlocks[currentChildNodeIdx+nodeCount-1]
+			distance := math.Abs(float64(candidateBlockCount) - targetBlockCount)
+			if distance < bestDistance {
+				bestChildNodeCount = nodeCount
+				bestBlockCount = candidateBlockCount
+				bestDistance = distance
+			}
+
+			if float64(candidateBlockCount) >= targetBlockCount {
+				// We are past the ideal block count, adding more child nodes won't get
+				// us closer to the ideal distance.
+				break
+			}
+		}
+
+		mergeNodeChildren[mergeNodeIdx] = bestChildNodeCount
+		currentChildNodeIdx += bestChildNodeCount
+		remainingTotalBlockCount -= bestBlockCount
+	}
+
+	return mergeNodeChildren
+}
+
+func balancedGroupSize(totalSize int, groupCount int, groupIdx int) int {
+	size := totalSize / groupCount
+	if groupIdx < totalSize%groupCount {
+		size++
+	}
+	return size
+}
+
 // allocateContiguous returns a []*T of length size where every element points
 // into a single backing []T allocation. This avoids the per-element heap
 // allocations from N separate &T{} expressions.
diff --git a/pkg/querybackend/queryplan/query_plan_test.go b/pkg/querybackend/queryplan/query_plan_test.go
index 318238bbdc..6187fa9d77 100644
--- a/pkg/querybackend/queryplan/query_plan_test.go
+++ b/pkg/querybackend/queryplan/query_plan_test.go
@@ -20,9 +20,14 @@ import (
 
 var update = flag.Bool("update", false, "rewrite golden files in testdata/ from the current plan output")
 
-// Test_Build verifies the shape of query plans produced by Build against
-// golden files in testdata/. Each subtest's golden file is named after the
-// subtest. E.g. Test_Build/single_block reads testdata/single_block.txt.
+// Test_Build verifies the shape of query plans produced by Build and
+// BuildBalanced against golden files in testdata/, running both builders over
+// the same input table so the two algorithms are easy to compare. Each
+// subtest's golden file is named after the subtest: Build's golden file is
+// named after the case (e.g. Test_Build/single_block reads
+// testdata/single_block.txt), and BuildBalanced's is named after the case
+// with a "balanced_" prefix (e.g. Test_Build/balanced_single_block reads
+// testdata/balanced_single_block.txt).
 //
 // To regenerate all golden files:
 //
@@ -30,7 +35,7 @@ var update = flag.Bool("update", false, "rewrite golden files in testdata/ from
 //
 // To regenerate a specific golden file:
 //
-//	go test ./pkg/querybackend/queryplan/ -run Test_Build/<name> -update
+//	go test ./pkg/querybackend/queryplan/ -run 'Test_Build/<name>$' -update
 func Test_Build(t *testing.T) {
 	tests := []struct {
 		name      string
@@ -47,35 +52,49 @@ func Test_Build(t *testing.T) {
 		{name: "full_depth_2", blocks: 6, maxReads: 2, maxMerges: 3},
 		{name: "just_over_depth_2", blocks: 7, maxReads: 2, maxMerges: 3},
 		{name: "twenty_five_blocks", blocks: 25, maxReads: 2, maxMerges: 3},
+		{name: "full_merge_vs_single_leaf_merge", blocks: 33, maxReads: 4, maxMerges: 8},
+		{name: "forced_equal_split", blocks: 16, maxReads: 1, maxMerges: 8},
+		{name: "three_way_split", blocks: 20, maxReads: 1, maxMerges: 8},
+	}
+
+	builders := []struct {
+		name   string
+		build  func(blocks []*metastorev1.BlockMeta, maxReads, maxMerges int) *queryv1.QueryPlan
+		prefix string
+	}{
+		{name: "Build", build: Build, prefix: "build_"},
+		{name: "BuildBalanced", build: BuildBalanced, prefix: "balanced_"},
 	}
 
 	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			blocks := makeBlocks(tt.blocks)
-			p := Build(blocks, tt.maxReads, tt.maxMerges)
+		for _, b := range builders {
+			t.Run(b.prefix+tt.name, func(t *testing.T) {
+				blocks := makeBlocks(tt.blocks)
+				p := b.build(blocks, tt.maxReads, tt.maxMerges)
 
-			var buf bytes.Buffer
-			writePlan(t, &buf, "", p.Root)
+				var buf bytes.Buffer
+				writePlan(t, &buf, "", p.Root)
 
-			// Ensure that the plan has not been modified during traversal.
-			assert.Equal(t, Build(blocks, tt.maxReads, tt.maxMerges), p)
+				// Ensure that the plan has not been modified during traversal.
+				assert.Equal(t, b.build(blocks, tt.maxReads, tt.maxMerges), p)
 
-			if *update {
-				require.NoError(t, os.WriteFile(goldenFile(t), buf.Bytes(), 0o644))
-				return
-			}
+				if *update {
+					require.NoError(t, os.WriteFile(goldenFile(t), buf.Bytes(), 0o644))
+					return
+				}
 
-			expected, err := os.ReadFile(goldenFile(t))
-			require.NoError(t, err)
-			assert.Equal(t, string(expected), buf.String())
-		})
+				expected, err := os.ReadFile(goldenFile(t))
+				require.NoError(t, err)
+				assert.Equal(t, string(expected), buf.String())
+			})
+		}
 	}
 }
 
 // makeBlocks creates n BlockMeta with sequential string IDs starting at "1".
 func makeBlocks(n int) []*metastorev1.BlockMeta {
 	blocks := make([]*metastorev1.BlockMeta, n)
-	for i := 0; i < n; i++ {
+	for i := range n {
 		blocks[i] = &metastorev1.BlockMeta{Id: strconv.Itoa(i + 1)}
 	}
 	return blocks
diff --git a/pkg/querybackend/queryplan/testdata/empty.txt b/pkg/querybackend/queryplan/testdata/balanced_empty.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/empty.txt
rename to pkg/querybackend/queryplan/testdata/balanced_empty.txt
diff --git a/pkg/querybackend/queryplan/testdata/exact_one_leaf.txt b/pkg/querybackend/queryplan/testdata/balanced_exact_one_leaf.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/exact_one_leaf.txt
rename to pkg/querybackend/queryplan/testdata/balanced_exact_one_leaf.txt
diff --git a/pkg/querybackend/queryplan/testdata/balanced_forced_equal_split.txt b/pkg/querybackend/queryplan/testdata/balanced_forced_equal_split.txt
new file mode 100644
index 0000000000..08aff02ba7
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/balanced_forced_equal_split.txt
@@ -0,0 +1,35 @@
+MERGE {children: 2, blocks: 0}
+	MERGE {children: 8, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"1"
+		READ {children: 0, blocks: 1}
+			id:"2"
+		READ {children: 0, blocks: 1}
+			id:"3"
+		READ {children: 0, blocks: 1}
+			id:"4"
+		READ {children: 0, blocks: 1}
+			id:"5"
+		READ {children: 0, blocks: 1}
+			id:"6"
+		READ {children: 0, blocks: 1}
+			id:"7"
+		READ {children: 0, blocks: 1}
+			id:"8"
+	MERGE {children: 8, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"9"
+		READ {children: 0, blocks: 1}
+			id:"10"
+		READ {children: 0, blocks: 1}
+			id:"11"
+		READ {children: 0, blocks: 1}
+			id:"12"
+		READ {children: 0, blocks: 1}
+			id:"13"
+		READ {children: 0, blocks: 1}
+			id:"14"
+		READ {children: 0, blocks: 1}
+			id:"15"
+		READ {children: 0, blocks: 1}
+			id:"16"
diff --git a/pkg/querybackend/queryplan/testdata/full_depth_2.txt b/pkg/querybackend/queryplan/testdata/balanced_full_depth_2.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/full_depth_2.txt
rename to pkg/querybackend/queryplan/testdata/balanced_full_depth_2.txt
diff --git a/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt b/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt
new file mode 100644
index 0000000000..a6187d0457
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt
@@ -0,0 +1,45 @@
+MERGE {children: 2, blocks: 0}
+	MERGE {children: 4, blocks: 0}
+		READ {children: 0, blocks: 4}
+			id:"1"
+			id:"2"
+			id:"3"
+			id:"4"
+		READ {children: 0, blocks: 4}
+			id:"5"
+			id:"6"
+			id:"7"
+			id:"8"
+		READ {children: 0, blocks: 4}
+			id:"9"
+			id:"10"
+			id:"11"
+			id:"12"
+		READ {children: 0, blocks: 4}
+			id:"13"
+			id:"14"
+			id:"15"
+			id:"16"
+	MERGE {children: 5, blocks: 0}
+		READ {children: 0, blocks: 4}
+			id:"17"
+			id:"18"
+			id:"19"
+			id:"20"
+		READ {children: 0, blocks: 4}
+			id:"21"
+			id:"22"
+			id:"23"
+			id:"24"
+		READ {children: 0, blocks: 3}
+			id:"25"
+			id:"26"
+			id:"27"
+		READ {children: 0, blocks: 3}
+			id:"28"
+			id:"29"
+			id:"30"
+		READ {children: 0, blocks: 3}
+			id:"31"
+			id:"32"
+			id:"33"
diff --git a/pkg/querybackend/queryplan/testdata/invalid_max_merges.txt b/pkg/querybackend/queryplan/testdata/balanced_invalid_max_merges.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/invalid_max_merges.txt
rename to pkg/querybackend/queryplan/testdata/balanced_invalid_max_merges.txt
diff --git a/pkg/querybackend/queryplan/testdata/invalid_max_reads.txt b/pkg/querybackend/queryplan/testdata/balanced_invalid_max_reads.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/invalid_max_reads.txt
rename to pkg/querybackend/queryplan/testdata/balanced_invalid_max_reads.txt
diff --git a/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt b/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt
new file mode 100644
index 0000000000..fb2eee45dd
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt
@@ -0,0 +1,14 @@
+MERGE {children: 2, blocks: 0}
+	MERGE {children: 2, blocks: 0}
+		READ {children: 0, blocks: 2}
+			id:"1"
+			id:"2"
+		READ {children: 0, blocks: 2}
+			id:"3"
+			id:"4"
+	MERGE {children: 2, blocks: 0}
+		READ {children: 0, blocks: 2}
+			id:"5"
+			id:"6"
+		READ {children: 0, blocks: 1}
+			id:"7"
diff --git a/pkg/querybackend/queryplan/testdata/single_block.txt b/pkg/querybackend/queryplan/testdata/balanced_single_block.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/single_block.txt
rename to pkg/querybackend/queryplan/testdata/balanced_single_block.txt
diff --git a/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt b/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt
new file mode 100644
index 0000000000..bd3cde2d67
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt
@@ -0,0 +1,44 @@
+MERGE {children: 3, blocks: 0}
+	MERGE {children: 7, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"1"
+		READ {children: 0, blocks: 1}
+			id:"2"
+		READ {children: 0, blocks: 1}
+			id:"3"
+		READ {children: 0, blocks: 1}
+			id:"4"
+		READ {children: 0, blocks: 1}
+			id:"5"
+		READ {children: 0, blocks: 1}
+			id:"6"
+		READ {children: 0, blocks: 1}
+			id:"7"
+	MERGE {children: 6, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"8"
+		READ {children: 0, blocks: 1}
+			id:"9"
+		READ {children: 0, blocks: 1}
+			id:"10"
+		READ {children: 0, blocks: 1}
+			id:"11"
+		READ {children: 0, blocks: 1}
+			id:"12"
+		READ {children: 0, blocks: 1}
+			id:"13"
+	MERGE {children: 7, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"14"
+		READ {children: 0, blocks: 1}
+			id:"15"
+		READ {children: 0, blocks: 1}
+			id:"16"
+		READ {children: 0, blocks: 1}
+			id:"17"
+		READ {children: 0, blocks: 1}
+			id:"18"
+		READ {children: 0, blocks: 1}
+			id:"19"
+		READ {children: 0, blocks: 1}
+			id:"20"
diff --git a/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt b/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt
new file mode 100644
index 0000000000..2744967ab5
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt
@@ -0,0 +1,46 @@
+MERGE {children: 2, blocks: 0}
+	MERGE {children: 3, blocks: 0}
+		MERGE {children: 2, blocks: 0}
+			READ {children: 0, blocks: 2}
+				id:"1"
+				id:"2"
+			READ {children: 0, blocks: 2}
+				id:"3"
+				id:"4"
+		MERGE {children: 3, blocks: 0}
+			READ {children: 0, blocks: 2}
+				id:"5"
+				id:"6"
+			READ {children: 0, blocks: 2}
+				id:"7"
+				id:"8"
+			READ {children: 0, blocks: 2}
+				id:"9"
+				id:"10"
+		MERGE {children: 2, blocks: 0}
+			READ {children: 0, blocks: 2}
+				id:"11"
+				id:"12"
+			READ {children: 0, blocks: 2}
+				id:"13"
+				id:"14"
+	MERGE {children: 2, blocks: 0}
+		MERGE {children: 3, blocks: 0}
+			READ {children: 0, blocks: 2}
+				id:"15"
+				id:"16"
+			READ {children: 0, blocks: 2}
+				id:"17"
+				id:"18"
+			READ {children: 0, blocks: 2}
+				id:"19"
+				id:"20"
+		MERGE {children: 3, blocks: 0}
+			READ {children: 0, blocks: 2}
+				id:"21"
+				id:"22"
+			READ {children: 0, blocks: 2}
+				id:"23"
+				id:"24"
+			READ {children: 0, blocks: 1}
+				id:"25"
diff --git a/pkg/querybackend/queryplan/testdata/two_leaves.txt b/pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/two_leaves.txt
rename to pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt
diff --git a/pkg/querybackend/queryplan/testdata/build_empty.txt b/pkg/querybackend/queryplan/testdata/build_empty.txt
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/pkg/querybackend/queryplan/testdata/build_exact_one_leaf.txt b/pkg/querybackend/queryplan/testdata/build_exact_one_leaf.txt
new file mode 100644
index 0000000000..a680a080ca
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/build_exact_one_leaf.txt
@@ -0,0 +1,3 @@
+READ {children: 0, blocks: 2}
+	id:"1"
+	id:"2"
diff --git a/pkg/querybackend/queryplan/testdata/build_forced_equal_split.txt b/pkg/querybackend/queryplan/testdata/build_forced_equal_split.txt
new file mode 100644
index 0000000000..08aff02ba7
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/build_forced_equal_split.txt
@@ -0,0 +1,35 @@
+MERGE {children: 2, blocks: 0}
+	MERGE {children: 8, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"1"
+		READ {children: 0, blocks: 1}
+			id:"2"
+		READ {children: 0, blocks: 1}
+			id:"3"
+		READ {children: 0, blocks: 1}
+			id:"4"
+		READ {children: 0, blocks: 1}
+			id:"5"
+		READ {children: 0, blocks: 1}
+			id:"6"
+		READ {children: 0, blocks: 1}
+			id:"7"
+		READ {children: 0, blocks: 1}
+			id:"8"
+	MERGE {children: 8, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"9"
+		READ {children: 0, blocks: 1}
+			id:"10"
+		READ {children: 0, blocks: 1}
+			id:"11"
+		READ {children: 0, blocks: 1}
+			id:"12"
+		READ {children: 0, blocks: 1}
+			id:"13"
+		READ {children: 0, blocks: 1}
+			id:"14"
+		READ {children: 0, blocks: 1}
+			id:"15"
+		READ {children: 0, blocks: 1}
+			id:"16"
diff --git a/pkg/querybackend/queryplan/testdata/build_full_depth_2.txt b/pkg/querybackend/queryplan/testdata/build_full_depth_2.txt
new file mode 100644
index 0000000000..bdb06b694c
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/build_full_depth_2.txt
@@ -0,0 +1,10 @@
+MERGE {children: 3, blocks: 0}
+	READ {children: 0, blocks: 2}
+		id:"1"
+		id:"2"
+	READ {children: 0, blocks: 2}
+		id:"3"
+		id:"4"
+	READ {children: 0, blocks: 2}
+		id:"5"
+		id:"6"
diff --git a/pkg/querybackend/queryplan/testdata/build_full_merge_vs_single_leaf_merge.txt b/pkg/querybackend/queryplan/testdata/build_full_merge_vs_single_leaf_merge.txt
new file mode 100644
index 0000000000..e8316ea5e7
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/build_full_merge_vs_single_leaf_merge.txt
@@ -0,0 +1,45 @@
+MERGE {children: 2, blocks: 0}
+	MERGE {children: 8, blocks: 0}
+		READ {children: 0, blocks: 4}
+			id:"1"
+			id:"2"
+			id:"3"
+			id:"4"
+		READ {children: 0, blocks: 4}
+			id:"5"
+			id:"6"
+			id:"7"
+			id:"8"
+		READ {children: 0, blocks: 4}
+			id:"9"
+			id:"10"
+			id:"11"
+			id:"12"
+		READ {children: 0, blocks: 4}
+			id:"13"
+			id:"14"
+			id:"15"
+			id:"16"
+		READ {children: 0, blocks: 4}
+			id:"17"
+			id:"18"
+			id:"19"
+			id:"20"
+		READ {children: 0, blocks: 4}
+			id:"21"
+			id:"22"
+			id:"23"
+			id:"24"
+		READ {children: 0, blocks: 4}
+			id:"25"
+			id:"26"
+			id:"27"
+			id:"28"
+		READ {children: 0, blocks: 4}
+			id:"29"
+			id:"30"
+			id:"31"
+			id:"32"
+	MERGE {children: 1, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"33"
diff --git a/pkg/querybackend/queryplan/testdata/build_invalid_max_merges.txt b/pkg/querybackend/queryplan/testdata/build_invalid_max_merges.txt
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/pkg/querybackend/queryplan/testdata/build_invalid_max_reads.txt b/pkg/querybackend/queryplan/testdata/build_invalid_max_reads.txt
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/pkg/querybackend/queryplan/testdata/just_over_depth_2.txt b/pkg/querybackend/queryplan/testdata/build_just_over_depth_2.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/just_over_depth_2.txt
rename to pkg/querybackend/queryplan/testdata/build_just_over_depth_2.txt
diff --git a/pkg/querybackend/queryplan/testdata/build_single_block.txt b/pkg/querybackend/queryplan/testdata/build_single_block.txt
new file mode 100644
index 0000000000..e0e4cf04c2
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/build_single_block.txt
@@ -0,0 +1,2 @@
+READ {children: 0, blocks: 1}
+	id:"1"
diff --git a/pkg/querybackend/queryplan/testdata/build_three_way_split.txt b/pkg/querybackend/queryplan/testdata/build_three_way_split.txt
new file mode 100644
index 0000000000..780dc17603
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/build_three_way_split.txt
@@ -0,0 +1,44 @@
+MERGE {children: 3, blocks: 0}
+	MERGE {children: 8, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"1"
+		READ {children: 0, blocks: 1}
+			id:"2"
+		READ {children: 0, blocks: 1}
+			id:"3"
+		READ {children: 0, blocks: 1}
+			id:"4"
+		READ {children: 0, blocks: 1}
+			id:"5"
+		READ {children: 0, blocks: 1}
+			id:"6"
+		READ {children: 0, blocks: 1}
+			id:"7"
+		READ {children: 0, blocks: 1}
+			id:"8"
+	MERGE {children: 8, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"9"
+		READ {children: 0, blocks: 1}
+			id:"10"
+		READ {children: 0, blocks: 1}
+			id:"11"
+		READ {children: 0, blocks: 1}
+			id:"12"
+		READ {children: 0, blocks: 1}
+			id:"13"
+		READ {children: 0, blocks: 1}
+			id:"14"
+		READ {children: 0, blocks: 1}
+			id:"15"
+		READ {children: 0, blocks: 1}
+			id:"16"
+	MERGE {children: 4, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"17"
+		READ {children: 0, blocks: 1}
+			id:"18"
+		READ {children: 0, blocks: 1}
+			id:"19"
+		READ {children: 0, blocks: 1}
+			id:"20"
diff --git a/pkg/querybackend/queryplan/testdata/twenty_five_blocks.txt b/pkg/querybackend/queryplan/testdata/build_twenty_five_blocks.txt
similarity index 100%
rename from pkg/querybackend/queryplan/testdata/twenty_five_blocks.txt
rename to pkg/querybackend/queryplan/testdata/build_twenty_five_blocks.txt
diff --git a/pkg/querybackend/queryplan/testdata/build_two_leaves.txt b/pkg/querybackend/queryplan/testdata/build_two_leaves.txt
new file mode 100644
index 0000000000..f6ade3b1e9
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/build_two_leaves.txt
@@ -0,0 +1,6 @@
+MERGE {children: 2, blocks: 0}
+	READ {children: 0, blocks: 2}
+		id:"1"
+		id:"2"
+	READ {children: 0, blocks: 1}
+		id:"3"
diff --git a/pkg/querybackend/queryplan/testdata/full_merge_vs_single_leaf_merge.txt b/pkg/querybackend/queryplan/testdata/full_merge_vs_single_leaf_merge.txt
new file mode 100644
index 0000000000..e8316ea5e7
--- /dev/null
+++ b/pkg/querybackend/queryplan/testdata/full_merge_vs_single_leaf_merge.txt
@@ -0,0 +1,45 @@
+MERGE {children: 2, blocks: 0}
+	MERGE {children: 8, blocks: 0}
+		READ {children: 0, blocks: 4}
+			id:"1"
+			id:"2"
+			id:"3"
+			id:"4"
+		READ {children: 0, blocks: 4}
+			id:"5"
+			id:"6"
+			id:"7"
+			id:"8"
+		READ {children: 0, blocks: 4}
+			id:"9"
+			id:"10"
+			id:"11"
+			id:"12"
+		READ {children: 0, blocks: 4}
+			id:"13"
+			id:"14"
+			id:"15"
+			id:"16"
+		READ {children: 0, blocks: 4}
+			id:"17"
+			id:"18"
+			id:"19"
+			id:"20"
+		READ {children: 0, blocks: 4}
+			id:"21"
+			id:"22"
+			id:"23"
+			id:"24"
+		READ {children: 0, blocks: 4}
+			id:"25"
+			id:"26"
+			id:"27"
+			id:"28"
+		READ {children: 0, blocks: 4}
+			id:"29"
+			id:"30"
+			id:"31"
+			id:"32"
+	MERGE {children: 1, blocks: 0}
+		READ {children: 0, blocks: 1}
+			id:"33"

From 80bb107993f80d303ae5de1e0d2d2d22a0c63d7a Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Wed, 12 Aug 2026 17:05:43 -0700
Subject: [PATCH 2/9] Add comments describing the query plan algorithm

---
 pkg/querybackend/queryplan/query_plan.go | 31 +++++++++++++++++++++++-
 1 file changed, 30 insertions(+), 1 deletion(-)

diff --git a/pkg/querybackend/queryplan/query_plan.go b/pkg/querybackend/queryplan/query_plan.go
index 8e94f1d042..3b6b4f0224 100644
--- a/pkg/querybackend/queryplan/query_plan.go
+++ b/pkg/querybackend/queryplan/query_plan.go
@@ -62,6 +62,8 @@ func Build(
 	}
 }
 
+// BuildBalanced builds a balanced query tree where each node of the tree has a
+// similar number of blocks it needs to process compared to its siblings.
 func BuildBalanced(blocks []*metastorev1.BlockMeta, maxReads int, maxMerges int) *queryv1.QueryPlan {
 	if len(blocks) == 0 || maxReads < 1 || maxMerges < 2 {
 		return new(queryv1.QueryPlan)
@@ -92,7 +94,7 @@ func BuildBalanced(blocks []*metastorev1.BlockMeta, maxReads int, maxMerges int)
 		mergeNodes := allocateContiguous[queryv1.QueryNode](mergeNodeCount)
 		mergeWeights := make([]int, mergeNodeCount)
 
-		// workloadGroupSizes calculates how many children each merge node should
+		// distributeChildNodes calculates how many children each merge node should
 		// have at this level.
 		groupSizes := distributeChildNodes(weights, mergeNodeCount, maxMerges)
 		start := 0
@@ -116,6 +118,19 @@ func BuildBalanced(blocks []*metastorev1.BlockMeta, maxReads int, maxMerges int)
 	}
 }
 
+// distributeChildNodes will take a collection of child nodes and evenly
+// distribute them across mergeNodeCount merge nodes, ensuring that each merge
+// node does not exceed maxMergeNodeSize children.
+//
+// The i-th element of childNodeBlocks is the number of blocks the i-th child
+// node contains (whether it is a merge node itself or a leaf node). The child
+// nodes are distributed such that each merge node has a similar number of
+// blocks and not necessarily a similar number of child nodes. A merge node
+// could have fewer child nodes than its peers, but it will always have a
+// similar number of blocks.
+//
+// It returns a slice of child node counts. The i-th element indicates how many
+// child nodes the i-th merge node should have.
 func distributeChildNodes(childNodeBlocks []int, mergeNodeCount int, maxMergeNodeSize int) []int {
 	// This slice contains the number of child nodes each merge node should be
 	// allocated.
@@ -181,6 +196,20 @@ func distributeChildNodes(childNodeBlocks []int, mergeNodeCount int, maxMergeNod
 	return mergeNodeChildren
 }
 
+// balancedGroupSize will take a totalSize and distribute it evenly across
+// groupCount groups. If there is a remainder R, that remainder is then spread
+// evenly across the first R groups.
+//
+// For example, given these parameters:
+//
+//	totalSize  = 5
+//	groupCount = 3
+//
+// We would get the following (for various values of groupIdx):
+//
+//	Group 0: 2
+//	Group 1: 2
+//	Group 2: 1
 func balancedGroupSize(totalSize int, groupCount int, groupIdx int) int {
 	size := totalSize / groupCount
 	if groupIdx < totalSize%groupCount {

From 1e3981eccc4b46461487acf520f573741e39b14f Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Wed, 12 Aug 2026 17:11:48 -0700
Subject: [PATCH 3/9] Update comment

---
 pkg/querybackend/queryplan/query_plan.go | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pkg/querybackend/queryplan/query_plan.go b/pkg/querybackend/queryplan/query_plan.go
index 3b6b4f0224..91b52499c9 100644
--- a/pkg/querybackend/queryplan/query_plan.go
+++ b/pkg/querybackend/queryplan/query_plan.go
@@ -94,8 +94,7 @@ func BuildBalanced(blocks []*metastorev1.BlockMeta, maxReads int, maxMerges int)
 		mergeNodes := allocateContiguous[queryv1.QueryNode](mergeNodeCount)
 		mergeWeights := make([]int, mergeNodeCount)
 
-		// distributeChildNodes calculates how many children each merge node should
-		// have at this level.
+		// Calculate how many children each merge node should have at this level.
 		groupSizes := distributeChildNodes(weights, mergeNodeCount, maxMerges)
 		start := 0
 		for idx, size := range groupSizes {

From 6bab408a1a1a59d5d7e5763cc8f3d4ab4fa19a58 Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Thu, 13 Aug 2026 09:54:37 -0700
Subject: [PATCH 4/9] Fix typo

---
 pkg/querybackend/queryplan/query_plan.go | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkg/querybackend/queryplan/query_plan.go b/pkg/querybackend/queryplan/query_plan.go
index 91b52499c9..f567663c92 100644
--- a/pkg/querybackend/queryplan/query_plan.go
+++ b/pkg/querybackend/queryplan/query_plan.go
@@ -161,7 +161,7 @@ func distributeChildNodes(childNodeBlocks []int, mergeNodeCount int, maxMergeNod
 		bestDistance := float64(0)
 
 		// We have to allocate at least the minimum number of child nodes to this
-		// merge node, so we do that now. we then calcualte how far away we are from
+		// merge node, so we do that now. we then calculate how far away we are from
 		// the ideal allocation.
 		for _, b := range childNodeBlocks[currentChildNodeIdx : currentChildNodeIdx+minChildNodeCount] {
 			bestBlockCount += b

From 412a05c15807592a775fffb33b8ee03d6818e45d Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Thu, 13 Aug 2026 10:37:03 -0700
Subject: [PATCH 5/9] Add config option to switch query planner types

---
 pkg/frontend/frontend.go                         | 16 ++++++++++++++++
 .../readpath/queryfrontend/query_frontend.go     | 13 +++++++++++--
 pkg/pyroscope/modules_experimental.go            |  2 ++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/pkg/frontend/frontend.go b/pkg/frontend/frontend.go
index 6d7818451b..96c0129468 100644
--- a/pkg/frontend/frontend.go
+++ b/pkg/frontend/frontend.go
@@ -54,6 +54,13 @@ type Config struct {
 	// on the request is rejected with Unimplemented.
 	AsyncQueriesEnabled bool `yaml:"async_queries_enabled" category:"experimental"`
 
+	// QueryPlannerStrategy sets the query planner strategy. By default this is
+	// "classic" which provides the legacy query planner behavior.
+	//
+	// Optionally this can be "balanced" which uses the balanced query planner
+	// algorithm.
+	QueryPlannerStrategy string `yaml:"query_planner_strategy" doc:"hidden"`
+
 	// Used to find local IP address, that is sent to scheduler and querier-worker.
 	InfNames   []string `yaml:"instance_interface_names" category:"advanced" doc:"default=[<private network interfaces>]"`
 	Addr       string   `yaml:"instance_addr" category:"advanced"`
@@ -79,6 +86,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet, logger log.Logger) {
 	f.BoolVar(&cfg.EnableIPv6, "query-frontend.instance-enable-ipv6", false, "Enable using a IPv6 instance address. (default false)")
 	f.IntVar(&cfg.Port, "query-frontend.instance-port", 0, "Port to advertise to query-scheduler and querier (defaults to -server.http-listen-port).")
 	f.BoolVar(&cfg.AsyncQueriesEnabled, "query-frontend.async-queries-enabled", false, "Enable the experimental asynchronous query path on SelectMergeStacktraces (default false)")
+	f.StringVar(&cfg.QueryPlannerStrategy, "query-frontend.query-planner-strategy", "classic", "Sets the query planner strategy, options: classic, balanced (default classic)")
 	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("query-frontend.grpc-client-config", f)
 }
 
@@ -87,6 +95,14 @@ func (cfg *Config) Validate() error {
 		return fmt.Errorf("scheduler address cannot be specified when query-scheduler service discovery mode is set to '%s'", cfg.QuerySchedulerDiscovery.Mode)
 	}
 
+	switch cfg.QueryPlannerStrategy {
+	case "":
+		cfg.QueryPlannerStrategy = "classic"
+	case "classic", "balanced":
+	default:
+		return fmt.Errorf("unknown query planner strategy: %q", cfg.QueryPlannerStrategy)
+	}
+
 	return cfg.GRPCClientConfig.Validate()
 }
 
diff --git a/pkg/frontend/readpath/queryfrontend/query_frontend.go b/pkg/frontend/readpath/queryfrontend/query_frontend.go
index ef929aecae..ec2b5abd90 100644
--- a/pkg/frontend/readpath/queryfrontend/query_frontend.go
+++ b/pkg/frontend/readpath/queryfrontend/query_frontend.go
@@ -66,6 +66,7 @@ type QueryFrontend struct {
 	symbolizer          Symbolizer
 	diagnosticsStore    DiagnosticsStore
 	now                 func() time.Time
+	queryPlanType       string
 
 	metrics *queryFrontendMetrics
 }
@@ -132,6 +133,7 @@ func newQueryFrontendMetrics(reg prometheus.Registerer) *queryFrontendMetrics {
 func NewQueryFrontend(
 	logger log.Logger,
 	limits frontend.Limits,
+	cfg frontend.Config,
 	metadataQueryClient metastorev1.MetadataQueryServiceClient,
 	tenantServiceClient metastorev1.TenantServiceClient,
 	querybackendClient QueryBackend,
@@ -148,6 +150,7 @@ func NewQueryFrontend(
 		symbolizer:          sym,
 		diagnosticsStore:    diagnosticsStore,
 		now:                 time.Now,
+		queryPlanType:       cfg.QueryPlannerStrategy,
 		metrics:             newQueryFrontendMetrics(reg),
 	}
 	return qf
@@ -255,8 +258,14 @@ func (q *QueryFrontend) doQuery(
 		blocks[i], blocks[j] = blocks[j], blocks[i]
 	})
 	xrandMutex.Unlock()
-	// TODO(kolesnikovae): Should be dynamic.
-	p := queryplan.BuildBalanced(blocks, 4, 20)
+
+	var p *queryv1.QueryPlan
+	switch q.queryPlanType {
+	case "balanced":
+		p = queryplan.BuildBalanced(blocks, 4, 20)
+	default: // Always fallback to the classic planner
+		p = queryplan.Build(blocks, 4, 20)
+	}
 
 	backend := q.querybackend
 	if backendC != nil {
diff --git a/pkg/pyroscope/modules_experimental.go b/pkg/pyroscope/modules_experimental.go
index 4b5e43b8e1..1278af6594 100644
--- a/pkg/pyroscope/modules_experimental.go
+++ b/pkg/pyroscope/modules_experimental.go
@@ -87,6 +87,7 @@ func (f *Pyroscope) initQueryFrontendV2() (services.Service, error) {
 	f.queryFrontend = queryfrontend.NewQueryFrontend(
 		queryFrontendLogger,
 		f.Overrides,
+		f.Cfg.Frontend,
 		f.metastoreClient,
 		f.metastoreClient,
 		f.queryBackendClient,
@@ -145,6 +146,7 @@ func (f *Pyroscope) initQueryFrontendV12() (services.Service, error) {
 	f.queryFrontend = queryfrontend.NewQueryFrontend(
 		queryFrontendLogger,
 		f.Overrides,
+		f.Cfg.Frontend,
 		f.metastoreClient,
 		f.metastoreClient,
 		f.queryBackendClient,

From 70ae7315723832e2ab1eded5765917ddba420fa9 Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Thu, 13 Aug 2026 10:38:13 -0700
Subject: [PATCH 6/9] Remove unused testdata

---
 .../full_merge_vs_single_leaf_merge.txt       | 45 -------------------
 1 file changed, 45 deletions(-)
 delete mode 100644 pkg/querybackend/queryplan/testdata/full_merge_vs_single_leaf_merge.txt

diff --git a/pkg/querybackend/queryplan/testdata/full_merge_vs_single_leaf_merge.txt b/pkg/querybackend/queryplan/testdata/full_merge_vs_single_leaf_merge.txt
deleted file mode 100644
index e8316ea5e7..0000000000
--- a/pkg/querybackend/queryplan/testdata/full_merge_vs_single_leaf_merge.txt
+++ /dev/null
@@ -1,45 +0,0 @@
-MERGE {children: 2, blocks: 0}
-	MERGE {children: 8, blocks: 0}
-		READ {children: 0, blocks: 4}
-			id:"1"
-			id:"2"
-			id:"3"
-			id:"4"
-		READ {children: 0, blocks: 4}
-			id:"5"
-			id:"6"
-			id:"7"
-			id:"8"
-		READ {children: 0, blocks: 4}
-			id:"9"
-			id:"10"
-			id:"11"
-			id:"12"
-		READ {children: 0, blocks: 4}
-			id:"13"
-			id:"14"
-			id:"15"
-			id:"16"
-		READ {children: 0, blocks: 4}
-			id:"17"
-			id:"18"
-			id:"19"
-			id:"20"
-		READ {children: 0, blocks: 4}
-			id:"21"
-			id:"22"
-			id:"23"
-			id:"24"
-		READ {children: 0, blocks: 4}
-			id:"25"
-			id:"26"
-			id:"27"
-			id:"28"
-		READ {children: 0, blocks: 4}
-			id:"29"
-			id:"30"
-			id:"31"
-			id:"32"
-	MERGE {children: 1, blocks: 0}
-		READ {children: 0, blocks: 1}
-			id:"33"

From ffe7477e92be9a97f738b4717facd73b1bfdbf16 Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Thu, 13 Aug 2026 11:04:18 -0700
Subject: [PATCH 7/9] Fix CI

---
 pkg/frontend/readpath/queryfrontend/query_frontend_test.go | 3 +++
 .../queryfrontend/query_select_merge_profile_test.go       | 3 +++
 .../queryfrontend/query_select_merge_span_profile_test.go  | 2 ++
 .../queryfrontend/query_select_merge_stacktraces_test.go   | 4 +++-
 .../queryfrontend/query_select_time_series_test.go         | 5 +++--
 .../queryfrontend/query_series_labels_compat_test.go       | 2 ++
 .../readpath/queryfrontend/symbol_ref_resolve_test.go      | 7 ++++---
 7 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/pkg/frontend/readpath/queryfrontend/query_frontend_test.go b/pkg/frontend/readpath/queryfrontend/query_frontend_test.go
index 66f0a1ee71..9b749d3b92 100644
--- a/pkg/frontend/readpath/queryfrontend/query_frontend_test.go
+++ b/pkg/frontend/readpath/queryfrontend/query_frontend_test.go
@@ -18,6 +18,7 @@ import (
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
 	"github.com/grafana/pyroscope/v2/pkg/block/metadata"
 	"github.com/grafana/pyroscope/v2/pkg/featureflags"
+	"github.com/grafana/pyroscope/v2/pkg/frontend"
 	"github.com/grafana/pyroscope/v2/pkg/tenant"
 	"github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend"
 	"github.com/grafana/pyroscope/v2/pkg/test/mocks/mockmetastorev1"
@@ -186,6 +187,7 @@ func Test_QueryFrontend_LabelNames_WithFiltering(t *testing.T) {
 			qf := NewQueryFrontend(
 				log.NewNopLogger(),
 				mockLimits,
+				frontend.Config{},
 				mockMetadataClient,
 				nil,
 				mockQueryBackend,
@@ -329,6 +331,7 @@ func Test_QueryFrontend_Series_WithLabelNameFiltering(t *testing.T) {
 			qf := NewQueryFrontend(
 				log.NewNopLogger(),
 				mockLimits,
+				frontend.Config{},
 				mockMetadataClient,
 				nil,
 				mockQueryBackend,
diff --git a/pkg/frontend/readpath/queryfrontend/query_select_merge_profile_test.go b/pkg/frontend/readpath/queryfrontend/query_select_merge_profile_test.go
index 7b446e9cf2..d20b11622c 100644
--- a/pkg/frontend/readpath/queryfrontend/query_select_merge_profile_test.go
+++ b/pkg/frontend/readpath/queryfrontend/query_select_merge_profile_test.go
@@ -19,6 +19,7 @@ import (
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
 	"github.com/grafana/pyroscope/v2/pkg/block/metadata"
+	"github.com/grafana/pyroscope/v2/pkg/frontend"
 	phlaremodel "github.com/grafana/pyroscope/v2/pkg/model"
 	"github.com/grafana/pyroscope/v2/pkg/pprof"
 	"github.com/grafana/pyroscope/v2/pkg/tenant"
@@ -41,6 +42,7 @@ func newSMPQueryFrontend(
 	return NewQueryFrontend(
 		log.NewNopLogger(),
 		limits,
+		frontend.Config{},
 		metaClient,
 		nil, // tenantServiceClient
 		backend,
@@ -682,6 +684,7 @@ func TestSelectMergeProfiles_Symbolization(t *testing.T) {
 			qf := NewQueryFrontend(
 				log.NewNopLogger(),
 				mockLimits,
+				frontend.Config{},
 				mockMetadataClient,
 				nil,
 				mockQueryBackend,
diff --git a/pkg/frontend/readpath/queryfrontend/query_select_merge_span_profile_test.go b/pkg/frontend/readpath/queryfrontend/query_select_merge_span_profile_test.go
index 9c22bff2d0..47dc57bd04 100644
--- a/pkg/frontend/readpath/queryfrontend/query_select_merge_span_profile_test.go
+++ b/pkg/frontend/readpath/queryfrontend/query_select_merge_span_profile_test.go
@@ -17,6 +17,7 @@ import (
 	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 	"github.com/grafana/pyroscope/v2/pkg/block/metadata"
+	"github.com/grafana/pyroscope/v2/pkg/frontend"
 	"github.com/grafana/pyroscope/v2/pkg/pprof"
 	"github.com/grafana/pyroscope/v2/pkg/tenant"
 	"github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend"
@@ -172,6 +173,7 @@ func TestSelectMergeSpanProfile_Symbolization(t *testing.T) {
 			qf := NewQueryFrontend(
 				log.NewNopLogger(),
 				mockLimits,
+				frontend.Config{},
 				mockMetadataClient,
 				nil,
 				mockQueryBackend,
diff --git a/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces_test.go b/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces_test.go
index 0a9c6df7b6..4aeefca28e 100644
--- a/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces_test.go
+++ b/pkg/frontend/readpath/queryfrontend/query_select_merge_stacktraces_test.go
@@ -16,6 +16,7 @@ import (
 	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 	"github.com/grafana/pyroscope/v2/pkg/block/metadata"
+	"github.com/grafana/pyroscope/v2/pkg/frontend"
 	phlaremodel "github.com/grafana/pyroscope/v2/pkg/model"
 	"github.com/grafana/pyroscope/v2/pkg/pprof"
 	"github.com/grafana/pyroscope/v2/pkg/tenant"
@@ -280,6 +281,7 @@ func TestSelectMergeStacktrace_Symbolization(t *testing.T) {
 			qf := NewQueryFrontend(
 				log.NewNopLogger(),
 				mockLimits,
+				frontend.Config{},
 				mockMetadataClient,
 				nil,
 				mockQueryBackend,
@@ -354,7 +356,7 @@ func TestSelectMergeStacktraces_DotFormat(t *testing.T) {
 			}},
 		}, nil)
 
-	qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, mockMetadataClient, nil, mockQueryBackend, nil, nil, nil)
+	qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, nil, nil, nil)
 
 	ctx := tenant.InjectTenantID(context.Background(), "tenant1")
 	start, end := smpValidTimeRange()
diff --git a/pkg/frontend/readpath/queryfrontend/query_select_time_series_test.go b/pkg/frontend/readpath/queryfrontend/query_select_time_series_test.go
index 19c359665c..770b9f2777 100644
--- a/pkg/frontend/readpath/queryfrontend/query_select_time_series_test.go
+++ b/pkg/frontend/readpath/queryfrontend/query_select_time_series_test.go
@@ -10,6 +10,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
+	"github.com/grafana/pyroscope/v2/pkg/frontend"
 	"github.com/grafana/pyroscope/v2/pkg/tenant"
 	"github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend"
 )
@@ -24,7 +25,7 @@ func TestSelectSeries_RejectsSubMillisecondStep(t *testing.T) {
 			limits.On("MaxQueryLookback", "test-tenant").Return(time.Duration(0)).Maybe()
 			limits.On("MaxQueryLength", "test-tenant").Return(time.Duration(0)).Maybe()
 
-			qf := NewQueryFrontend(log.NewNopLogger(), limits, nil, nil, nil, nil, nil, nil)
+			qf := NewQueryFrontend(log.NewNopLogger(), limits, frontend.Config{}, nil, nil, nil, nil, nil, nil)
 			ctx := tenant.InjectTenantID(context.Background(), "test-tenant")
 
 			_, err := qf.SelectSeries(ctx, connect.NewRequest(&querierv1.SelectSeriesRequest{
@@ -46,7 +47,7 @@ func TestSelectHeatmap_RejectsSubMillisecondStep(t *testing.T) {
 	for _, step := range []float64{0, 0.0001, 0.0005, 0.0009999} {
 		t.Run("step="+formatStep(step), func(t *testing.T) {
 			limits := mockfrontend.NewMockLimits(t)
-			qf := NewQueryFrontend(log.NewNopLogger(), limits, nil, nil, nil, nil, nil, nil)
+			qf := NewQueryFrontend(log.NewNopLogger(), limits, frontend.Config{}, nil, nil, nil, nil, nil, nil)
 			ctx := tenant.InjectTenantID(context.Background(), "test-tenant")
 
 			_, err := qf.SelectHeatmap(ctx, connect.NewRequest(&querierv1.SelectHeatmapRequest{
diff --git a/pkg/frontend/readpath/queryfrontend/query_series_labels_compat_test.go b/pkg/frontend/readpath/queryfrontend/query_series_labels_compat_test.go
index 89f46bc316..6be58c037e 100644
--- a/pkg/frontend/readpath/queryfrontend/query_series_labels_compat_test.go
+++ b/pkg/frontend/readpath/queryfrontend/query_series_labels_compat_test.go
@@ -14,6 +14,7 @@ import (
 	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
 	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
 	typesv1 "github.com/grafana/pyroscope/api/gen/proto/go/types/v1"
+	"github.com/grafana/pyroscope/v2/pkg/frontend"
 	phlaremodel "github.com/grafana/pyroscope/v2/pkg/model"
 	"github.com/grafana/pyroscope/v2/pkg/tenant"
 	"github.com/grafana/pyroscope/v2/pkg/test/mocks/mockfrontend"
@@ -136,6 +137,7 @@ func Test_QueryFrontend_Series_ProfileTypeQueryServedFromMetadata(t *testing.T)
 			qf := NewQueryFrontend(
 				log.NewNopLogger(),
 				mockLimits,
+				frontend.Config{},
 				mockMetadataClient,
 				nil,
 				mockQueryBackend,
diff --git a/pkg/frontend/readpath/queryfrontend/symbol_ref_resolve_test.go b/pkg/frontend/readpath/queryfrontend/symbol_ref_resolve_test.go
index 9c6259a87f..c039744d64 100644
--- a/pkg/frontend/readpath/queryfrontend/symbol_ref_resolve_test.go
+++ b/pkg/frontend/readpath/queryfrontend/symbol_ref_resolve_test.go
@@ -16,6 +16,7 @@ import (
 	querierv1 "github.com/grafana/pyroscope/api/gen/proto/go/querier/v1"
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 	"github.com/grafana/pyroscope/lidia"
+	"github.com/grafana/pyroscope/v2/pkg/frontend"
 	phlaremodel "github.com/grafana/pyroscope/v2/pkg/model"
 	"github.com/grafana/pyroscope/v2/pkg/model/symbolref"
 	"github.com/grafana/pyroscope/v2/pkg/tenant"
@@ -177,7 +178,7 @@ func TestSelectMergeStacktracesTree_SymbolRefFlagOn(t *testing.T) {
 		}},
 	}, nil).Once()
 
-	qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, mockMetadataClient, nil, mockQueryBackend, mockSymbolizer, nil, nil)
+	qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, mockSymbolizer, nil, nil)
 
 	ctx := tenant.InjectTenantID(context.Background(), "tenant1")
 	start, end := smpValidTimeRange()
@@ -235,7 +236,7 @@ func TestSelectMergeStacktracesTree_SymbolRefResolution(t *testing.T) {
 		Blocks: []*metastorev1.BlockMeta{{Id: "block_id"}},
 	}, nil).Once()
 
-	qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, mockMetadataClient, nil, mockQueryBackend, mockSymbolizer, nil, nil)
+	qf := NewQueryFrontend(log.NewNopLogger(), mockLimits, frontend.Config{}, mockMetadataClient, nil, mockQueryBackend, mockSymbolizer, nil, nil)
 
 	before := testutil.ToFloat64(qf.metrics.symbolRefLocationsTotal.WithLabelValues(symbolRefLocationResolved))
 
@@ -269,7 +270,7 @@ func TestSelectMergeStacktracesTree_SymbolRefResolution(t *testing.T) {
 // path (see TestRebuildInlineChainExpansionOrder for the Rebuild-side
 // contract).
 func TestBuildLookup_ReversesLidiaFrameOrder(t *testing.T) {
-	qf := NewQueryFrontend(log.NewNopLogger(), nil, nil, nil, nil, nil, nil, nil)
+	qf := NewQueryFrontend(log.NewNopLogger(), nil, frontend.Config{}, nil, nil, nil, nil, nil, nil)
 	lookup := qf.buildLookup([]binaryResolution{{
 		binary: symbolref.UnresolvedBinary{BuildID: "build-a", BinaryName: "libfoo.so", Addresses: []uint64{0x100}},
 		frames: [][]lidia.SourceInfoFrame{{

From 67e7b0784f6faae1446c87d1cf9d5355f4d6b259 Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Thu, 13 Aug 2026 11:14:52 -0700
Subject: [PATCH 8/9] Fix help strings

---
 cmd/pyroscope/help-all.txt.tmpl | 2 ++
 cmd/pyroscope/help.txt.tmpl     | 2 ++
 pkg/frontend/frontend.go        | 2 +-
 3 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/cmd/pyroscope/help-all.txt.tmpl b/cmd/pyroscope/help-all.txt.tmpl
index 3470e41c5c..d20783e9a0 100644
--- a/cmd/pyroscope/help-all.txt.tmpl
+++ b/cmd/pyroscope/help-all.txt.tmpl
@@ -935,6 +935,8 @@ Usage of ./pyroscope:
     	Port to advertise to query-scheduler and querier (defaults to -server.http-listen-port).
   -query-frontend.max-async-query-concurrency int
     	Maximum number of concurrent async queries per tenant. 0 to disable async queries. (default 5)
+  -query-frontend.query-planner-strategy string
+    	Sets the query planner strategy, options: classic, balanced (default "classic")
   -query-frontend.scheduler-worker-concurrency int
     	Number of concurrent workers forwarding queries to single query-scheduler. (default 5)
   -query-scheduler.grpc-client-config.backoff-max-period duration
diff --git a/cmd/pyroscope/help.txt.tmpl b/cmd/pyroscope/help.txt.tmpl
index e1ae759017..eab246e977 100644
--- a/cmd/pyroscope/help.txt.tmpl
+++ b/cmd/pyroscope/help.txt.tmpl
@@ -199,6 +199,8 @@ Usage of ./pyroscope:
     	Include profiles that were sampled out and stored with stacktraces stripped (marked __sampled__) in query results.
   -query-frontend.max-async-query-concurrency int
     	Maximum number of concurrent async queries per tenant. 0 to disable async queries. (default 5)
+  -query-frontend.query-planner-strategy string
+    	Sets the query planner strategy, options: classic, balanced (default "classic")
   -query-scheduler.max-outstanding-requests-per-tenant int
     	Maximum number of outstanding requests per tenant per query-scheduler. In-flight requests above this limit will fail with HTTP response status code 429. (default 100)
   -query-scheduler.ring.consul.hostname string
diff --git a/pkg/frontend/frontend.go b/pkg/frontend/frontend.go
index 96c0129468..7b33dc1b01 100644
--- a/pkg/frontend/frontend.go
+++ b/pkg/frontend/frontend.go
@@ -86,7 +86,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet, logger log.Logger) {
 	f.BoolVar(&cfg.EnableIPv6, "query-frontend.instance-enable-ipv6", false, "Enable using a IPv6 instance address. (default false)")
 	f.IntVar(&cfg.Port, "query-frontend.instance-port", 0, "Port to advertise to query-scheduler and querier (defaults to -server.http-listen-port).")
 	f.BoolVar(&cfg.AsyncQueriesEnabled, "query-frontend.async-queries-enabled", false, "Enable the experimental asynchronous query path on SelectMergeStacktraces (default false)")
-	f.StringVar(&cfg.QueryPlannerStrategy, "query-frontend.query-planner-strategy", "classic", "Sets the query planner strategy, options: classic, balanced (default classic)")
+	f.StringVar(&cfg.QueryPlannerStrategy, "query-frontend.query-planner-strategy", "classic", "Sets the query planner strategy, options: classic, balanced")
 	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("query-frontend.grpc-client-config", f)
 }
 

From b2a0b867841f48aa04825b7a68217c9056cf5bb6 Mon Sep 17 00:00:00 2001
From: Bryan Huhta <bryan.huhta@grafana.com>
Date: Fri, 14 Aug 2026 16:29:18 -0700
Subject: [PATCH 9/9] Implement simplier query tree balancing algorithm

---
 pkg/querybackend/queryplan/query_plan.go      | 246 +++++++++---------
 pkg/querybackend/queryplan/query_plan_test.go |  36 +++
 ...lanced_full_merge_vs_single_leaf_merge.txt |  22 +-
 .../testdata/balanced_just_over_depth_2.txt   |   8 +-
 .../testdata/balanced_three_way_split.txt     |   4 +-
 .../testdata/balanced_twenty_five_blocks.txt  |  30 +--
 .../testdata/balanced_two_leaves.txt          |   4 +-
 7 files changed, 195 insertions(+), 155 deletions(-)

diff --git a/pkg/querybackend/queryplan/query_plan.go b/pkg/querybackend/queryplan/query_plan.go
index f567663c92..7e7582e94a 100644
--- a/pkg/querybackend/queryplan/query_plan.go
+++ b/pkg/querybackend/queryplan/query_plan.go
@@ -1,8 +1,6 @@
 package queryplan
 
 import (
-	"math"
-
 	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
 	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
 )
@@ -69,152 +67,158 @@ func BuildBalanced(blocks []*metastorev1.BlockMeta, maxReads int, maxMerges int)
 		return new(queryv1.QueryPlan)
 	}
 
-	readNodeCount := int(math.Ceil(float64(len(blocks)) / float64(maxReads)))
-	nodes := allocateContiguous[queryv1.QueryNode](readNodeCount)
-	weights := make([]int, readNodeCount)
-
-	// Build the read nodes, balancing blocks across all read nodes. We also
-	// record the number of blocks in each read node to be used later when
-	// assigning read nodes to merge nodes.
-	for start, idx := 0, 0; idx < readNodeCount; idx++ {
-		size := balancedGroupSize(len(blocks), readNodeCount, idx)
-		end := start + size
+	leafNodeCount := (len(blocks) + maxReads - 1) / maxReads
+	nodes := allocateContiguous[queryv1.QueryNode](leafNodeCount)
 
+	// Uniformly assign blocks to leaf nodes.
+	var start int
+	for idx := range leafNodeCount {
+		end := (idx + 1) * len(blocks) / leafNodeCount
 		nodes[idx].Type = queryv1.QueryNode_READ
 		nodes[idx].Blocks = blocks[start:end:end]
-		weights[idx] = size
 		start = end
 	}
 
-	// Build the merge nodes. We assign merge children to merge nodes based on the
-	// number of blocks being merged at each level. We want each merge node to
-	// merge approximately the same number of blocks as its siblings.
-	for len(nodes) > 1 {
-		mergeNodeCount := int(math.Ceil(float64(len(nodes)) / float64(maxMerges)))
-		mergeNodes := allocateContiguous[queryv1.QueryNode](mergeNodeCount)
-		mergeWeights := make([]int, mergeNodeCount)
-
-		// Calculate how many children each merge node should have at this level.
-		groupSizes := distributeChildNodes(weights, mergeNodeCount, maxMerges)
-		start := 0
-		for idx, size := range groupSizes {
-			end := start + size
-			mergeNodes[idx].Type = queryv1.QueryNode_MERGE
-			mergeNodes[idx].Children = nodes[start:end:end]
-
-			for _, weight := range weights[start:end] {
-				mergeWeights[idx] += weight
-			}
-			start = end
-		}
-
-		nodes = mergeNodes
-		weights = mergeWeights
-	}
-
+	// Recursively build a balanced tree of merge nodes.
+	root := buildMergeTree(nodes, maxMerges)
 	return &queryv1.QueryPlan{
-		Root: nodes[0],
+		Root: root,
 	}
 }
 
-// distributeChildNodes will take a collection of child nodes and evenly
-// distribute them across mergeNodeCount merge nodes, ensuring that each merge
-// node does not exceed maxMergeNodeSize children.
-//
-// The i-th element of childNodeBlocks is the number of blocks the i-th child
-// node contains (whether it is a merge node itself or a leaf node). The child
-// nodes are distributed such that each merge node has a similar number of
-// blocks and not necessarily a similar number of child nodes. A merge node
-// could have fewer child nodes than its peers, but it will always have a
-// similar number of blocks.
-//
-// It returns a slice of child node counts. The i-th element indicates how many
-// child nodes the i-th merge node should have.
-func distributeChildNodes(childNodeBlocks []int, mergeNodeCount int, maxMergeNodeSize int) []int {
-	// This slice contains the number of child nodes each merge node should be
-	// allocated.
-	mergeNodeChildren := make([]int, mergeNodeCount)
-
-	remainingTotalBlockCount := 0
-	for _, count := range childNodeBlocks {
-		remainingTotalBlockCount += count
+// buildMergeTree will recursively create a query tree of merge nodes with at
+// most maxMerges number of children. At each level, the tree will maintain a
+// similar number of blocks assigned to each merge node compared to its
+// siblings.
+//
+// The algorithm is straightforward, but it has nuances. Since len(nodes) > 1,
+// we know we need to create a merge node and place nodes underneath it.
+//
+// When creating a merge node we want to partition the nodes evenly into groups
+// such that we have no more than maxMerges groups (since a given merge node
+// cannot exceed maxMerges children). Once we select a number of groups, we use
+// balanceGroupItems to evenly (as evenly as possible) spread the nodes across
+// all the groups.
+//
+// It's important to note that we want to select a groupSize which is a power of
+// maxMerges. This ensures that each subtree has the same depth as its siblings.
+//
+// As an example, consider the following input:
+//
+//	maxMerges = 3
+//	nodes     = [ n0 n1 n2 n3 n4 n5 n6 n7 n8 n9 ]
+//
+// We want to partition nodes such that each group has an even number of
+// nodes itself. Naively we could compute:
+//
+//	groupSize = ceil(len(nodes) / maxMerges)
+//
+// However, this would result in an imbalanced tree:
+//
+//	[                   M0                  ]
+//	[ n0 n1 n2 n3 ] [ n4 n5 n6 ] [ n7 n8 n9 ]
+//
+// From here, the first partition was given 4 nodes, which exceeds maxMerges,
+// so it needs to be broken down further. The other two partitions have 3 nodes,
+// so they do not need to be branched further. After splitting the first
+// partition, the tree becomes imbalanced.
+//
+//	[                     M0                    ]
+//	[        M1       ] [ n4 n5 n6 ] [ n7 n8 n9 ]
+//	[ n0 n1 ] [ n2 n3 ]
+//
+// Instead, we select the largest power of maxMerges K that's less than
+// len(nodes). In this case:
+//
+//	K = maxMerges^N
+//	  for max(N) and maxMerges^N < len(nodes)
+//	if N = 2, then 3^2 < 10 so K = 9
+//
+// If we allow K to represent the largest size of a partition, we can calculate
+// how many partitions we need:
+//
+//	# of partitions = ceil(len(nodes) / K)
+//	                = ceil(10 / 5)
+//	                = 2
+//
+// Now we can balance the nodes across 2 partitions:
+//
+//	[                 M0                ]
+//	[ n0 n1 n2 n3 n4 ] [ n5 n6 n7 n8 n9 ]
+//
+// Both partitions exceed maxMerges, so we repeat the algorithm:
+//
+//	[                     M0                    ]
+//	[         M1         ] [         M2         ]
+//	[ n0 n1 n2 ] [ n3 n4 ] [ n5 n6 n7 ] [ n8 n9 ]
+//
+// At this point, each subtree does the same amount of work as its siblings.
+//
+// As a theoretical note, this algorithm may produce an imbalanced depth tree
+// if the len(nodes) is not a pwer of maxMerges and maxMerges = 2. This case is
+// unlikely in production--as maxMerges = 20 is common--but also irrelevant to
+// performance overall since such a low maxMerges value would produce other
+// inefficiencies.
+func buildMergeTree(nodes []*queryv1.QueryNode, maxMerges int) *queryv1.QueryNode {
+	if len(nodes) == 1 {
+		// The base case, this is a leaf node.
+		return nodes[0]
 	}
 
-	currentChildNodeIdx := 0
-	for mergeNodeIdx := range mergeNodeCount {
-		// Calculate the remaining merge nodes and child nodes we have left to
-		// allocate.
-		remainingMergeNodes := mergeNodeCount - mergeNodeIdx
-		remainingChildNodes := len(childNodeBlocks) - currentChildNodeIdx
-
-		// Calculate the minimum and maximum child nodes this merge node can have.
-		minChildNodeCount := max(1, remainingChildNodes-(remainingMergeNodes-1)*maxMergeNodeSize)
-		maxChildNodeCount := min(maxMergeNodeSize, remainingChildNodes-(remainingMergeNodes-1))
-
-		// Calculate the ideal number of blocks we could give to this merge node
-		// (and all subsequent merge nodes) to evenly distribute blocks amongst
-		// them. This number likely will not be a whole number.
-		targetBlockCount := float64(remainingTotalBlockCount) / float64(remainingMergeNodes)
-
-		bestChildNodeCount := minChildNodeCount
-		bestBlockCount := 0
-		bestDistance := float64(0)
-
-		// We have to allocate at least the minimum number of child nodes to this
-		// merge node, so we do that now. we then calculate how far away we are from
-		// the ideal allocation.
-		for _, b := range childNodeBlocks[currentChildNodeIdx : currentChildNodeIdx+minChildNodeCount] {
-			bestBlockCount += b
-		}
-		bestDistance = math.Abs(float64(bestBlockCount) - targetBlockCount)
-
-		// Now we expand the number of child nodes we give to this merge node to see
-		// if we can get closer to the ideal block count.
-		candidateBlockCount := bestBlockCount
-		for nodeCount := minChildNodeCount + 1; nodeCount <= maxChildNodeCount; nodeCount++ {
-			candidateBlockCount += childNodeBlocks[currentChildNodeIdx+nodeCount-1]
-			distance := math.Abs(float64(candidateBlockCount) - targetBlockCount)
-			if distance < bestDistance {
-				bestChildNodeCount = nodeCount
-				bestBlockCount = candidateBlockCount
-				bestDistance = distance
-			}
-
-			if float64(candidateBlockCount) >= targetBlockCount {
-				// We are past the ideal block count, adding more child nodes won't get
-				// us closer to the ideal distance.
-				break
-			}
-		}
+	// We have len(nodes) number of nodes. We need to partition them into groups
+	// such that we have no more than maxMerges number of groups. Importantly, we
+	// want a group size such that we don't place all the nodes into a single
+	// group.
+	groupSize := 1
+	for groupSize*maxMerges < len(nodes) {
+		groupSize *= maxMerges
+	}
 
-		mergeNodeChildren[mergeNodeIdx] = bestChildNodeCount
-		currentChildNodeIdx += bestChildNodeCount
-		remainingTotalBlockCount -= bestBlockCount
+	// Given groupSize as the maximum number of nodes that can be assigned to
+	// each subtree, we calculate
+	//
+	//	ceil(len(nodes) / groupSize)
+	//
+	// to determine how many children this node will have.
+	childCount := (len(nodes) + groupSize - 1) / groupSize
+
+	parent := &queryv1.QueryNode{
+		Type:     queryv1.QueryNode_MERGE,
+		Children: make([]*queryv1.QueryNode, childCount),
+	}
+
+	// Evenly distribute all the nodes to each child of this merge node.
+	var start int
+	for idx := range childCount {
+		end := start + balanceGroupItems(len(nodes), childCount, idx)
+		parent.Children[idx] = buildMergeTree(nodes[start:end], maxMerges)
+		start = end
 	}
 
-	return mergeNodeChildren
+	return parent
 }
 
-// balancedGroupSize will take a totalSize and distribute it evenly across
-// groupCount groups. If there is a remainder R, that remainder is then spread
-// evenly across the first R groups.
+// balanceGroupItems will take numItems and distribute them evenly across
+// numGroups groups. If there is a remainder R, that remainder is then spread
+// evenly across the first R groups. It returns the number of items groupIdx
+// should have for all groups to remain balanced.
 //
 // For example, given these parameters:
 //
-//	totalSize  = 5
-//	groupCount = 3
+//	numItems  = 5
+//	numGroups = 3
 //
 // We would get the following (for various values of groupIdx):
 //
 //	Group 0: 2
 //	Group 1: 2
 //	Group 2: 1
-func balancedGroupSize(totalSize int, groupCount int, groupIdx int) int {
-	size := totalSize / groupCount
-	if groupIdx < totalSize%groupCount {
-		size++
+func balanceGroupItems(numItems int, numGroups int, groupIdx int) int {
+	groupSize := numItems / numGroups
+	if groupIdx < numItems%numGroups {
+		groupSize++
 	}
-	return size
+	return groupSize
 }
 
 // allocateContiguous returns a []*T of length size where every element points
diff --git a/pkg/querybackend/queryplan/query_plan_test.go b/pkg/querybackend/queryplan/query_plan_test.go
index 6187fa9d77..240b05ec7f 100644
--- a/pkg/querybackend/queryplan/query_plan_test.go
+++ b/pkg/querybackend/queryplan/query_plan_test.go
@@ -109,6 +109,42 @@ func goldenFile(t *testing.T) string {
 	return filepath.Join("testdata", parts[len(parts)-1]+".txt")
 }
 
+// Test_balanceGroupItems verifies that items are spread evenly across groups,
+// with any remainder distributed one-per-group starting from group 0.
+func Test_balanceGroupItems(t *testing.T) {
+	tests := []struct {
+		name      string
+		numItems  int
+		numGroups int
+		want      []int // want[i] is the expected size of group i
+	}{
+		{name: "zero_items", numItems: 0, numGroups: 3, want: []int{0, 0, 0}},
+		{name: "single_group", numItems: 5, numGroups: 1, want: []int{5}},
+		{name: "even_split", numItems: 6, numGroups: 3, want: []int{2, 2, 2}},
+		{name: "remainder_spread", numItems: 5, numGroups: 3, want: []int{2, 2, 1}},
+		{name: "remainder_almost_full", numItems: 8, numGroups: 3, want: []int{3, 3, 2}},
+		{name: "one_item_per_group", numItems: 3, numGroups: 3, want: []int{1, 1, 1}},
+		{name: "more_groups_than_items", numItems: 2, numGroups: 5, want: []int{1, 1, 0, 0, 0}},
+		{name: "no_items", numItems: 0, numGroups: 1, want: []int{0}},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			require.Len(t, tt.want, tt.numGroups, "test case is malformed: want must have numGroups elements")
+
+			got := make([]int, tt.numGroups)
+			sum := 0
+			for idx := range tt.numGroups {
+				got[idx] = balanceGroupItems(tt.numItems, tt.numGroups, idx)
+				sum += got[idx]
+			}
+
+			assert.Equal(t, tt.want, got)
+			assert.Equal(t, tt.numItems, sum, "group sizes must sum back to numItems")
+		})
+	}
+}
+
 // writePlan writes an indented textual representation of the plan rooted at
 // n to w. A nil root produces no output. The test fails on malformed nodes.
 func writePlan(t *testing.T, w io.Writer, pad string, n *queryv1.QueryNode) {
diff --git a/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt b/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt
index a6187d0457..44384abc9b 100644
--- a/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt
+++ b/pkg/querybackend/queryplan/testdata/balanced_full_merge_vs_single_leaf_merge.txt
@@ -1,45 +1,45 @@
 MERGE {children: 2, blocks: 0}
-	MERGE {children: 4, blocks: 0}
-		READ {children: 0, blocks: 4}
+	MERGE {children: 5, blocks: 0}
+		READ {children: 0, blocks: 3}
 			id:"1"
 			id:"2"
 			id:"3"
-			id:"4"
 		READ {children: 0, blocks: 4}
+			id:"4"
 			id:"5"
 			id:"6"
 			id:"7"
-			id:"8"
 		READ {children: 0, blocks: 4}
+			id:"8"
 			id:"9"
 			id:"10"
 			id:"11"
+		READ {children: 0, blocks: 3}
 			id:"12"
-		READ {children: 0, blocks: 4}
 			id:"13"
 			id:"14"
+		READ {children: 0, blocks: 4}
 			id:"15"
 			id:"16"
-	MERGE {children: 5, blocks: 0}
-		READ {children: 0, blocks: 4}
 			id:"17"
 			id:"18"
+	MERGE {children: 4, blocks: 0}
+		READ {children: 0, blocks: 4}
 			id:"19"
 			id:"20"
-		READ {children: 0, blocks: 4}
 			id:"21"
 			id:"22"
+		READ {children: 0, blocks: 3}
 			id:"23"
 			id:"24"
-		READ {children: 0, blocks: 3}
 			id:"25"
+		READ {children: 0, blocks: 4}
 			id:"26"
 			id:"27"
-		READ {children: 0, blocks: 3}
 			id:"28"
 			id:"29"
+		READ {children: 0, blocks: 4}
 			id:"30"
-		READ {children: 0, blocks: 3}
 			id:"31"
 			id:"32"
 			id:"33"
diff --git a/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt b/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt
index fb2eee45dd..53feaac55b 100644
--- a/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt
+++ b/pkg/querybackend/queryplan/testdata/balanced_just_over_depth_2.txt
@@ -1,14 +1,14 @@
 MERGE {children: 2, blocks: 0}
 	MERGE {children: 2, blocks: 0}
-		READ {children: 0, blocks: 2}
+		READ {children: 0, blocks: 1}
 			id:"1"
-			id:"2"
 		READ {children: 0, blocks: 2}
+			id:"2"
 			id:"3"
-			id:"4"
 	MERGE {children: 2, blocks: 0}
 		READ {children: 0, blocks: 2}
+			id:"4"
 			id:"5"
+		READ {children: 0, blocks: 2}
 			id:"6"
-		READ {children: 0, blocks: 1}
 			id:"7"
diff --git a/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt b/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt
index bd3cde2d67..6bf750c031 100644
--- a/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt
+++ b/pkg/querybackend/queryplan/testdata/balanced_three_way_split.txt
@@ -14,7 +14,7 @@ MERGE {children: 3, blocks: 0}
 			id:"6"
 		READ {children: 0, blocks: 1}
 			id:"7"
-	MERGE {children: 6, blocks: 0}
+	MERGE {children: 7, blocks: 0}
 		READ {children: 0, blocks: 1}
 			id:"8"
 		READ {children: 0, blocks: 1}
@@ -27,9 +27,9 @@ MERGE {children: 3, blocks: 0}
 			id:"12"
 		READ {children: 0, blocks: 1}
 			id:"13"
-	MERGE {children: 7, blocks: 0}
 		READ {children: 0, blocks: 1}
 			id:"14"
+	MERGE {children: 6, blocks: 0}
 		READ {children: 0, blocks: 1}
 			id:"15"
 		READ {children: 0, blocks: 1}
diff --git a/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt b/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt
index 2744967ab5..ad8ac16d21 100644
--- a/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt
+++ b/pkg/querybackend/queryplan/testdata/balanced_twenty_five_blocks.txt
@@ -1,46 +1,46 @@
 MERGE {children: 2, blocks: 0}
 	MERGE {children: 3, blocks: 0}
-		MERGE {children: 2, blocks: 0}
-			READ {children: 0, blocks: 2}
+		MERGE {children: 3, blocks: 0}
+			READ {children: 0, blocks: 1}
 				id:"1"
-				id:"2"
 			READ {children: 0, blocks: 2}
+				id:"2"
 				id:"3"
-				id:"4"
-		MERGE {children: 3, blocks: 0}
 			READ {children: 0, blocks: 2}
+				id:"4"
 				id:"5"
-				id:"6"
+		MERGE {children: 2, blocks: 0}
 			READ {children: 0, blocks: 2}
+				id:"6"
 				id:"7"
-				id:"8"
 			READ {children: 0, blocks: 2}
+				id:"8"
 				id:"9"
-				id:"10"
 		MERGE {children: 2, blocks: 0}
 			READ {children: 0, blocks: 2}
+				id:"10"
 				id:"11"
-				id:"12"
 			READ {children: 0, blocks: 2}
+				id:"12"
 				id:"13"
-				id:"14"
 	MERGE {children: 2, blocks: 0}
 		MERGE {children: 3, blocks: 0}
 			READ {children: 0, blocks: 2}
+				id:"14"
 				id:"15"
-				id:"16"
 			READ {children: 0, blocks: 2}
+				id:"16"
 				id:"17"
-				id:"18"
 			READ {children: 0, blocks: 2}
+				id:"18"
 				id:"19"
-				id:"20"
 		MERGE {children: 3, blocks: 0}
 			READ {children: 0, blocks: 2}
+				id:"20"
 				id:"21"
-				id:"22"
 			READ {children: 0, blocks: 2}
+				id:"22"
 				id:"23"
+			READ {children: 0, blocks: 2}
 				id:"24"
-			READ {children: 0, blocks: 1}
 				id:"25"
diff --git a/pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt b/pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt
index f6ade3b1e9..a5313fd486 100644
--- a/pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt
+++ b/pkg/querybackend/queryplan/testdata/balanced_two_leaves.txt
@@ -1,6 +1,6 @@
 MERGE {children: 2, blocks: 0}
-	READ {children: 0, blocks: 2}
+	READ {children: 0, blocks: 1}
 		id:"1"
+	READ {children: 0, blocks: 2}
 		id:"2"
-	READ {children: 0, blocks: 1}
 		id:"3"