Skip to content

Refactor distribution optimizer #734

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fpetkovski merged 1 commit into from
Sep 1, 2026
+155 −132
Merged

Refactor distribution optimizer #734

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
210 changes: 78 additions & 132 deletions logicalplan/distribute.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,14 @@ type DistributedExecutionOptimizer struct {
SkipDedup bool
}

// Optimize distributes a plan in three phases:
//
// 1. Classify each subtree as non-distributive, distributable as-is, or
// distributable through an aggregation, avg, or absent rewrite. A parent
// can absorb only children that are distributable as-is.
// 2. Select the roots of maximal distributable subtrees. These are nodes with
// a distribution strategy whose parent has none.
// 3. Replace each selected subtree with the remote plan for its strategy.
func (m DistributedExecutionOptimizer) Optimize(plan Node, opts *query.Options) (Node, annotations.Annotations) {
engines := m.Endpoints.Engines(MinMaxTime(plan, opts))
sort.Slice(engines, func(i, j int) bool {
Expand All @@ -181,38 +189,25 @@ func (m DistributedExecutionOptimizer) Optimize(plan Node, opts *query.Options)
warns := annotations.New()

parents := computeParents(&plan)
distributionPoints := m.computeDistributionPoints(&plan, parents, engineLabels, warns)
distributionPoints := m.computeDistributionPoints(&plan, engineLabels, warns)

TraverseBottomUp(nil, &plan, func(parent, current *Node) (stop bool) {
if _, distributeNow := distributionPoints[current]; !distributeNow {
TraverseBottomUp(nil, &plan, func(_ *Node, current *Node) (stop bool) {
strategy, distributeNow := distributionPoints[current]
if !distributeNow {
return false
}

if isAvgAggregation(current) && !preservesPartitionLabels(*current, engineLabels) {
// avg without partition labels: rewrite as sum/count.
*current = m.distributeAvg(*current, engines, m.subqueryOpts(parents, current, opts), labelRanges)
return true
}

if isAbsent(current) {
*current = m.distributeAbsent(*current, engines, calculateStartOffset(current, opts.LookbackDelta), m.subqueryOpts(parents, current, opts))
return true
subqueryOpts := m.subqueryOpts(parents, current, opts)
switch strategy {
case rewriteAvg:
*current = m.distributeAvg(*current, engines, subqueryOpts, labelRanges)
case rewriteAbsent:
*current = m.distributeAbsent(*current, engines, calculateStartOffset(current, opts.LookbackDelta), subqueryOpts)
case rewriteAggregation:
*current = m.distributeAggregation((*current).(*Aggregation), engines, subqueryOpts, labelRanges)
case distributeAsIs:
*current = m.distributeQuery(current, engines, subqueryOpts, labelRanges)
}

if isAggregation(current) {
if preservesPartitionLabels(*current, engineLabels) {
// Partition-preserving aggregation: push as-is since each engine
// computes over disjoint partition values.
*current = m.distributeQuery(current, engines, m.subqueryOpts(parents, current, opts), labelRanges)
} else {
// Distributive aggregation that drops partition labels: use a
// two-level split with local_agg(remote_agg(X)).
*current = m.distributeAggregation((*current).(*Aggregation), engines, m.subqueryOpts(parents, current, opts), labelRanges)
}
return true
}

*current = m.distributeQuery(current, engines, m.subqueryOpts(parents, current, opts), labelRanges)
return true
})
return plan, *warns
Expand Down Expand Up @@ -243,105 +238,76 @@ func computeParents(plan *Node) map[*Node]*Node {
return parents
}

func (m DistributedExecutionOptimizer) computeDistributionPoints(plan *Node, parents map[*Node]*Node, engineLabels map[string]struct{}, warns *annotations.Annotations) map[*Node]struct{} {
marks := make(map[*Node]struct{})
// A rewrite strategy can absorb only children that distribute as-is.
type distributionStrategy uint8

// First pass: mark distribution points (aggregations, absent functions).
Traverse(plan, func(current *Node) {
// Skip subtrees that are already distributed (e.g. by a previous
// distributed optimizer). This lets multiple distributed optimizers
// be chained: once the plan is distributed, subsequent optimizers
// fall through instead of re-distributing.
if isDistributed(current) {
return
}
if isAbsent(current) {
if m.isDistributive(current, engineLabels, warns) {
marks[current] = struct{}{}
}
return
}
if isAggregation(current) {
// Non-distributive aggregations that don't preserve partition labels
// cannot be distributed, except for avg which gets rewritten as sum/count.
if !m.isDistributive(current, engineLabels, warns) {
if isAvgAggregation(current) {
marks[current] = struct{}{}
}
return
}
// Distributive aggregations (standard or partition-preserving):
// defer to ancestor if possible.
if preservesPartitionLabels(*current, engineLabels) {
if m.hasDistributiveAncestor(parents, current, engineLabels, warns) {
return
}
}
marks[current] = struct{}{}
return
}
const (
cannotDistribute distributionStrategy = iota
distributeAsIs
rewriteAggregation
rewriteAvg
rewriteAbsent
)

parent := parents[current]
if parent == nil || IsConstantExpr(*current) {
return
}
if _, parentMarked := marks[parent]; parentMarked {
return
}
if !m.isDistributive(parent, engineLabels, warns) && m.isDistributive(current, engineLabels, warns) {
marks[current] = struct{}{}
}
})
func (m DistributedExecutionOptimizer) computeDistributionPoints(plan *Node, engineLabels map[string]struct{}, warns *annotations.Annotations) map[*Node]distributionStrategy {
strategies := make(map[*Node]distributionStrategy)
m.classifyDistribution(plan, strategies, engineLabels, warns)

// Second pass: for nodes whose siblings have marks, mark them too so both
// sides of a binary expression get distributed.
Traverse(plan, func(current *Node) {
if _, ok := marks[current]; ok {
return
}
if isDistributed(current) {
return
}
if subtreeHasMark(current, marks) {
return
}
if !m.isDistributive(current, engineLabels, warns) {
return
// Select roots of maximal distributable subtrees.
points := make(map[*Node]distributionStrategy)
TraverseBottomUp(nil, plan, func(parent, current *Node) bool {
strategy := strategies[current]
if strategy == cannotDistribute || IsConstantExpr(*current) {
return false
}
parent := parents[current]
if parent != nil && (m.isDistributive(parent, engineLabels, warns) || isAvgAggregation(parent)) {
if !subtreeHasMark(parent, marks) {
return
}
if parent != nil && strategies[parent] != cannotDistribute {
return false
}
marks[current] = struct{}{}
points[current] = strategy
return false
})

return marks
return points
}

// isDistributed reports whether the subtree rooted at node has already been
// processed by a distributed optimizer, i.e. it contains a Deduplicate,
// RemoteExecution or Noop node (Noop is the terminal result of distributing a
// subtree that matched no engines). Such subtrees must not be distributed again.
func isDistributed(node *Node) bool {
// classifyDistribution records how each subtree can be distributed.
func (m DistributedExecutionOptimizer) classifyDistribution(node *Node, strategies map[*Node]distributionStrategy, engineLabels map[string]struct{}, warns *annotations.Annotations) distributionStrategy {
switch (*node).(type) {
case RemoteMerge, RemoteExecution, Noop:
return true
strategies[node] = cannotDistribute
return cannotDistribute
}
return slices.ContainsFunc((*node).Children(), isDistributed)
}

func subtreeHasMark(node *Node, marks map[*Node]struct{}) bool {
childrenCanPushDown := true
for _, child := range (*node).Children() {
if _, ok := marks[child]; ok {
return true
if m.classifyDistribution(child, strategies, engineLabels, warns) != distributeAsIs {
childrenCanPushDown = false
}
if subtreeHasMark(child, marks) {
return true
}

distributive := m.isDistributiveOperation(node, engineLabels, warns)
strategy := cannotDistribute
if childrenCanPushDown {
switch {
case isAbsent(node):
if distributive {
strategy = rewriteAbsent
}
case isAvgAggregation(node) && !preservesPartitionLabels(*node, engineLabels):
strategy = rewriteAvg
case isAggregation(node):
if distributive {
strategy = distributeAsIs
if !preservesPartitionLabels(*node, engineLabels) {
strategy = rewriteAggregation
}
}
case distributive:
strategy = distributeAsIs
}
}
return false

strategies[node] = strategy
return strategy
}

func (m DistributedExecutionOptimizer) subqueryOpts(parents map[*Node]*Node, current *Node, opts *query.Options) *query.Options {
Expand Down Expand Up @@ -734,7 +700,7 @@ func preservesPartitionLabels(expr Node, partitionLabels map[string]struct{}) bo
}
}

func (m DistributedExecutionOptimizer) isDistributive(expr *Node, engineLabels map[string]struct{}, warns *annotations.Annotations) bool {
func (m DistributedExecutionOptimizer) isDistributiveOperation(expr *Node, engineLabels map[string]struct{}, warns *annotations.Annotations) bool {
if expr == nil {
return false
}
Expand All @@ -746,10 +712,7 @@ func (m DistributedExecutionOptimizer) isDistributive(expr *Node, engineLabels m
if isBinaryExpressionWithOneScalarSide(e) {
return true
}
return !m.SkipBinaryPushdown &&
isBinaryExpressionWithDistributableMatching(e, engineLabels) &&
m.isDistributive(&e.LHS, engineLabels, warns) &&
m.isDistributive(&e.RHS, engineLabels, warns)
return !m.SkipBinaryPushdown && isBinaryExpressionWithDistributableMatching(e, engineLabels)
case *Aggregation:
switch e.Op {
// Mathematically distributive: can be split into local_agg(remote_agg(X))
Expand Down Expand Up @@ -939,23 +902,6 @@ func matchesExternalLabels(ms []*labels.Matcher, externalLabels labels.Labels) b
return true
}

// hasDistributiveAncestor checks if there's a distributive node somewhere up the
// parent chain from the current node that can handle distribution.
// We must have an unbroken chain of distributive nodes to the ancestor for it to
// be able to handle distribution on our behalf.
func (m DistributedExecutionOptimizer) hasDistributiveAncestor(parents map[*Node]*Node, current *Node, engineLabels map[string]struct{}, warns *annotations.Annotations) bool {
for p := parents[current]; p != nil; p = parents[p] {
if !m.isDistributive(p, engineLabels, warns) {
// We hit a non-distributive node, so we can't push through it.
// No ancestor can help us distribute.
return false
}
}
// All ancestors are distributive, so the root (or the point where we
// stop traversing) can handle distribution.
return parents[current] != nil
}

func maxTime(a, b time.Time) time.Time {
if a.After(b) {
return a
Expand Down
77 changes: 77 additions & 0 deletions logicalplan/distribute_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ import (
"github.com/prometheus/prometheus/promql"
"github.com/prometheus/prometheus/promql/parser"
"github.com/prometheus/prometheus/promql/promqltest"
"github.com/prometheus/prometheus/util/annotations"
)

var replacements = map[string]*regexp.Regexp{
Expand All @@ -28,6 +29,82 @@ var replacements = map[string]*regexp.Regexp{
")": closedParenthesis,
}

func TestComputeDistributionPoints(t *testing.T) {
type point struct {
expr string
strategy distributionStrategy
}

partitionLabels := map[string]struct{}{"region": {}}
cases := []struct {
name string
expr string
expected []point
}{
{
name: "distributive chain",
expr: `rate(metric_a[5m])`,
expected: []point{{`rate(metric_a[5m])`, distributeAsIs}},
},
{
name: "non-distributive function argument",
expr: `clamp_max(metric_a, -scalar(metric_b))`,
expected: []point{
{`metric_a`, distributeAsIs},
{`metric_b`, distributeAsIs},
},
},
{
name: "aggregation rewrite",
expr: `sum(rate(metric_a[5m]))`,
expected: []point{{`sum(rate(metric_a[5m]))`, rewriteAggregation}},
},
{
name: "average rewrite",
expr: `avg(metric_a)`,
expected: []point{{`avg(metric_a)`, rewriteAvg}},
},
{
name: "absent rewrite",
expr: `absent(metric_a)`,
expected: []point{{`absent(metric_a)`, rewriteAbsent}},
},
{
name: "nested aggregation rewrite",
expr: `max(sum by (instance) (metric_a))`,
expected: []point{{`sum by (instance) (metric_a)`, rewriteAggregation}},
},
{
name: "partition-preserving nested aggregation",
expr: `max(sum by (region, instance) (metric_a))`,
expected: []point{{`max(sum by (region, instance) (metric_a))`, rewriteAggregation}},
},
{
name: "non-distributive aggregation",
expr: `quantile(0.5, metric_a)`,
expected: []point{{`metric_a`, distributeAsIs}},
},
}

for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
expr, err := parser.ParseExpr(tc.expr)
testutil.Ok(t, err)
plan, err := NewFromAST(expr, &query.Options{}, PlanOptions{})
testutil.Ok(t, err)
root := plan.Root()
points := (DistributedExecutionOptimizer{}).computeDistributionPoints(&root, partitionLabels, annotations.New())
actual := make([]point, 0, len(points))
Traverse(&root, func(node *Node) {
if strategy, ok := points[node]; ok {
actual = append(actual, point{(*node).String(), strategy})
}
})
testutil.Equals(t, tc.expected, actual)
})
}
}

func TestDistributedExecution(t *testing.T) {
t.Parallel()
cases := []struct {
Expand Down
Loading