feat(distributed_execution): add core query planning primitives and shard-count flag

Add the core building blocks for distributed query execution (fragment_metadata,
deduplicate_node/ShardedRemoteExecutions, distributed_optimizer, codec) and the
QFE-level -querier.distributed-exec-shard-count flag that drives how many shards
the optimizer fans a shardable aggregation into (0 disables it).

Signed-off-by: Essam Eldaly <eeldaly@amazon.com>

Author: eeldaly

pkg/distributed_execution/codec.go

@@ -191,6 +191,16 @@ func unmarshalNode(data []byte) (logicalplan.Node, error) {
 			return nil, err
 		}
 		return r, nil
+	case ShardedRemoteExecutionNode:
+		s := &ShardedRemoteExecutions{}
+		for _, c := range t.Children {
+			child, err := unmarshalNode(c)
+			if err != nil {
+				return nil, err
+			}
+			s.Expressions = append(s.Expressions, child)
+		}
+		return s, nil
 	}
 	return nil, nil
 }

pkg/distributed_execution/deduplicate_node.go

@@ -0,0 +1,77 @@
+package distributed_execution
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/prometheus/prometheus/promql/parser"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/thanos-io/promql-engine/execution"
+	"github.com/thanos-io/promql-engine/execution/exchange"
+	"github.com/thanos-io/promql-engine/execution/model"
+	"github.com/thanos-io/promql-engine/logicalplan"
+	"github.com/thanos-io/promql-engine/query"
+)
+
+const (
+	ShardedRemoteExecutionNode logicalplan.NodeType = "ShardedRemoteExecutionNode"
+)
+
+var _ logicalplan.UserDefinedExpr = (*ShardedRemoteExecutions)(nil)
+
+// ShardedRemoteExecutions is a custom logical-plan node that fans a single
+// aggregation out into one sub-expression per shard. At execution time each
+// sub-expression is built into its own operator and the partial results are
+// coalesced together, so a parent aggregation (e.g. sum) can combine them.
+//
+// NOTE: This is experimental. It is reconciled to compile against the current
+// Thanos engine API, but it is not yet wired end-to-end (see caveats in the
+// port notes): the node has no JSON (un)marshaling for codec round-tripping,
+// and the shard matcher injected by the optimizer does not match the merged
+// metric-name-hash storage sharding mechanism.
+type ShardedRemoteExecutions struct {
+	Expressions []logicalplan.Node `json:"-"`
+}
+
+// MakeExecutionOperator builds one operator per shard sub-expression and
+// coalesces them into a single operator stream for the parent aggregation.
+func (r *ShardedRemoteExecutions) MakeExecutionOperator(
+	ctx context.Context,
+	opts *query.Options,
+	hints storage.SelectHints,
+) (model.VectorOperator, error) {
+	operators := make([]model.VectorOperator, len(r.Expressions))
+	var err error
+	for i := range operators {
+		operators[i], err = execution.New(ctx, r.Expressions[i], nil, opts)
+		if err != nil {
+			return nil, err
+		}
+	}
+	coalesce := exchange.NewCoalesce(opts, 0, operators...)
+	return exchange.NewConcurrent(coalesce, 2, opts), nil
+}
+
+func (r *ShardedRemoteExecutions) Clone() logicalplan.Node {
+	clone := &ShardedRemoteExecutions{Expressions: make([]logicalplan.Node, len(r.Expressions))}
+	for i, e := range r.Expressions {
+		clone.Expressions[i] = e.Clone()
+	}
+	return clone
+}
+
+func (r *ShardedRemoteExecutions) Children() []*logicalplan.Node {
+	children := make([]*logicalplan.Node, len(r.Expressions))
+	for i := range r.Expressions {
+		children[i] = &r.Expressions[i]
+	}
+	return children
+}
+
+func (r *ShardedRemoteExecutions) String() string {
+	return fmt.Sprintf("shard(%d)", len(r.Expressions))
+}
+
+func (r *ShardedRemoteExecutions) ReturnType() parser.ValueType { return parser.ValueTypeVector }
+
+func (r *ShardedRemoteExecutions) Type() logicalplan.NodeType { return ShardedRemoteExecutionNode }

pkg/distributed_execution/distributed_optimizer.go

@@ -1,37 +1,84 @@
 package distributed_execution
 
 import (
+	"fmt"
+
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/prometheus/prometheus/util/annotations"
 	"github.com/thanos-io/promql-engine/logicalplan"
 	"github.com/thanos-io/promql-engine/query"
 )
 
-// This is a simplified implementation that only handles binary aggregation cases
-// Future versions of the distributed optimizer are expected to:
-// - Support more complex query patterns
-// - Incorporate diverse optimization strategies
-// - Extend support to node types beyond binary operations
-
-type DistributedOptimizer struct{}
+// This optimizer inserts Remote nodes so portions of the plan can be executed
+// remotely. It supports two strategies:
+//   - Binary-aggregation splitting: each operand of a binary expression that
+//     contains an aggregation is offloaded to a Remote fragment.
+//   - Sharded-aggregation splitting (experimental): a sum/count aggregation is
+//     rewritten into sum(ShardedRemoteExecutions{per-shard sub-aggregations}),
+//     so a single aggregation can be evaluated across shards and merged.
+//
+// NOTE: The sharded-aggregation path is experimental and not yet wired
+// end-to-end (no codec round-tripping for ShardedRemoteExecutions, and the
+// injected shard matcher does not match the merged metric-name-hash storage
+// sharding). It is reconciled here only to compile and to allow local
+// iteration.
+type DistributedOptimizer struct {
+	// ShardCount is the number of shards the sharded-aggregation strategy fans
+	// an aggregation out into. It should match the compactor's per-tenant
+	// metric-name-shard-size so per-shard subqueries route to disjoint blocks.
+	ShardCount int
+}
 
 func (d *DistributedOptimizer) Optimize(root logicalplan.Node, opts *query.Options) (logicalplan.Node, annotations.Annotations) {
 	warns := annotations.New()
 
-	// insert remote nodes
-	logicalplan.TraverseBottomUp(nil, &root, func(parent, current *logicalplan.Node) bool {
-
-		if (*current).Type() == logicalplan.BinaryNode && d.hasAggregation(current) {
-			ch := (*current).Children()
+	if root == nil {
+		return root, *warns
+	}
 
-			for _, child := range ch {
-				temp := (*child).Clone()
-				*child = NewRemoteNode(temp)
-				*(*child).Children()[0] = temp
+	// Strategy 1 (experimental): shard aggregations (sum/count) across shards.
+	// Only runs when a positive shard count is configured; otherwise it is left
+	// to the pre-existing binary-splitting strategy below.
+	sharded := false
+	if shardCount := d.ShardCount; shardCount > 0 {
+		logicalplan.TraverseBottomUp(nil, &root, func(parent, current *logicalplan.Node) bool {
+			if aggr, ok := (*current).(*logicalplan.Aggregation); ok {
+				switch aggr.Op {
+				case parser.SUM, parser.COUNT:
+					subqueries := newRemoteAggregation(aggr, shardCount)
+					*current = &logicalplan.Aggregation{
+						Op:       parser.SUM,
+						Expr:     &ShardedRemoteExecutions{Expressions: subqueries},
+						Param:    aggr.Param,
+						Grouping: aggr.Grouping,
+						Without:  aggr.Without,
+					}
+					sharded = true
+				}
+				return true
 			}
-		}
+			return false
+		})
+	}
 
-		return false
-	})
+	// Strategy 2 (pre-existing): offload binary-expression operands that
+	// contain an aggregation to Remote fragments. Skipped when the
+	// aggregation-sharding pass already distributed the plan, to avoid
+	// double-distributing the same subtrees.
+	if !sharded {
+		logicalplan.TraverseBottomUp(nil, &root, func(parent, current *logicalplan.Node) bool {
+			if (*current).Type() == logicalplan.BinaryNode && d.hasAggregation(current) {
+				ch := (*current).Children()
+				for _, child := range ch {
+					temp := (*child).Clone()
+					*child = NewRemoteNode(temp)
+					*(*child).Children()[0] = temp
+				}
+			}
+			return false
+		})
+	}
 
 	return root, *warns
 }
@@ -47,3 +94,37 @@ func (d *DistributedOptimizer) hasAggregation(root *logicalplan.Node) bool {
 	})
 	return isAggr
 }
+
+// newRemoteAggregation produces one Remote-wrapped per-shard copy of the given
+// aggregation, tagging each copy's selectors with its shard identity.
+func newRemoteAggregation(rootAggregation *logicalplan.Aggregation, shardNum int) []logicalplan.Node {
+	nodes := make([]logicalplan.Node, 0, shardNum)
+	for i := 0; i < shardNum; i++ {
+		rc := rootAggregation.Expr.Clone()
+		node := insertShardNum(&Remote{
+			Expr: &logicalplan.Aggregation{
+				Op:       rootAggregation.Op,
+				Expr:     rc,
+				Param:    rootAggregation.Param,
+				Grouping: rootAggregation.Grouping,
+				Without:  rootAggregation.Without,
+			},
+		}, shardNum, i)
+		nodes = append(nodes, node)
+	}
+	return nodes
+}
+
+// insertShardNum tags every vector selector in the subtree with the shard it
+// belongs to. NOTE: the label used here does not match the merged
+// metric-name-hash storage sharding; this is experimental.
+func insertShardNum(root logicalplan.Node, shardCount int, shardIdx int) logicalplan.Node {
+	logicalplan.TraverseBottomUp(nil, &root, func(parent, current *logicalplan.Node) bool {
+		if (*current).Type() == logicalplan.VectorSelectorNode {
+			cur := (*current).(*logicalplan.VectorSelector)
+			cur.LabelMatchers = append(cur.LabelMatchers, labels.MustNewMatcher(labels.MatchEqual, "__cortex_ingester_shard__", fmt.Sprintf("%d_%d", shardCount, shardIdx)))
+		}
+		return false
+	})
+	return root
+}

pkg/distributed_execution/fragment_metadata.go

@@ -0,0 +1,37 @@
+package distributed_execution
+
+import (
+	"context"
+)
+
+type fragmentMetadataKey struct{}
+
+type fragmentMetadata struct {
+	queryID       uint64
+	fragmentID    uint64
+	childIDToAddr map[uint64]string
+	isRoot        bool
+}
+
+// InjectFragmentMetaData stores the distributed execution metadata for the current
+// fragment into the context. This metadata is propagated from the query-scheduler to
+// the querier so that the querier knows which fragment it is executing, where to pull
+// its child fragments' results from, and whether it is the root (coordinator) fragment.
+func InjectFragmentMetaData(ctx context.Context, fragmentID uint64, queryID uint64, isRoot bool, childIDToAddr map[uint64]string) context.Context {
+	return context.WithValue(ctx, fragmentMetadataKey{}, fragmentMetadata{
+		queryID:       queryID,
+		fragmentID:    fragmentID,
+		childIDToAddr: childIDToAddr,
+		isRoot:        isRoot,
+	})
+}
+
+// ExtractFragmentMetaData retrieves the distributed execution metadata for the current
+// fragment from the context. The final return value reports whether metadata was present.
+func ExtractFragmentMetaData(ctx context.Context) (isRoot bool, queryID uint64, fragmentID uint64, childAddrs map[uint64]string, ok bool) {
+	metadata, ok := ctx.Value(fragmentMetadataKey{}).(fragmentMetadata)
+	if !ok {
+		return false, 0, 0, nil, false
+	}
+	return metadata.isRoot, metadata.queryID, metadata.fragmentID, metadata.childIDToAddr, true
+}

pkg/distributed_execution/fragment_metadata_test.go

@@ -0,0 +1,97 @@
+package distributed_execution
+
+import (
+	"context"
+	"reflect"
+	"testing"
+)
+
+func TestFragmentMetadata(t *testing.T) {
+	tests := []struct {
+		name      string
+		queryID   uint64
+		fragID    uint64
+		isRoot    bool
+		childIDs  []uint64
+		childAddr []string
+	}{
+		{
+			name:      "basic test",
+			queryID:   123,
+			fragID:    456,
+			isRoot:    true,
+			childIDs:  []uint64{1, 2, 3},
+			childAddr: []string{"addr1", "addr2", "addr3"},
+		},
+		{
+			name:      "empty children",
+			queryID:   789,
+			fragID:    101,
+			isRoot:    false,
+			childIDs:  []uint64{},
+			childAddr: []string{},
+		},
+		{
+			name:      "single child",
+			queryID:   999,
+			fragID:    888,
+			isRoot:    true,
+			childIDs:  []uint64{42},
+			childAddr: []string{"10.0.0.1:8080"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			// injection
+			ctx := context.Background()
+
+			childIDToAddr := make(map[uint64]string)
+			for i, childID := range tt.childIDs {
+				childIDToAddr[childID] = tt.childAddr[i]
+			}
+			newCtx := InjectFragmentMetaData(ctx, tt.fragID, tt.queryID, tt.isRoot, childIDToAddr)
+
+			// extraction
+			isRoot, queryID, fragmentID, childAddrs, ok := ExtractFragmentMetaData(newCtx)
+
+			// verify results
+			if !ok {
+				t.Error("ExtractFragmentMetaData failed, ok = false")
+			}
+
+			if isRoot != tt.isRoot {
+				t.Errorf("isRoot = %v, want %v", isRoot, tt.isRoot)
+			}
+
+			if queryID != tt.queryID {
+				t.Errorf("queryID = %v, want %v", queryID, tt.queryID)
+			}
+
+			if fragmentID != tt.fragID {
+				t.Errorf("fragmentID = %v, want %v", fragmentID, tt.fragID)
+			}
+
+			// create expected childIDToAddr map
+			expectedChildAddrs := make(map[uint64]string)
+			for i, childID := range tt.childIDs {
+				expectedChildAddrs[childID] = tt.childAddr[i]
+			}
+
+			if !reflect.DeepEqual(childAddrs, expectedChildAddrs) {
+				t.Errorf("childAddrs = %v, want %v", childAddrs, expectedChildAddrs)
+			}
+		})
+	}
+}
+
+func TestExtractFragmentMetaDataWithEmptyContext(t *testing.T) {
+	ctx := context.Background()
+	isRoot, queryID, fragmentID, childAddrs, ok := ExtractFragmentMetaData(ctx)
+	if ok {
+		t.Error("ExtractFragmentMetaData should return ok=false for empty context")
+	}
+	if isRoot || queryID != 0 || fragmentID != 0 || childAddrs != nil {
+		t.Error("ExtractFragmentMetaData should return zero values for empty context")
+	}
+}

pkg/querier/querier.go

@@ -101,6 +101,8 @@ type Config struct {
 
 	DistributedExecEnabled bool `yaml:"distributed_exec_enabled" doc:"hidden"`
 
+	DistributedExecShardCount int `yaml:"distributed_exec_shard_count" doc:"hidden"`
+
 	HonorProjectionHints bool `yaml:"honor_projection_hints"`
 
 	// Timeout classification flags for converting 5XX to 4XX on expensive queries.
@@ -164,6 +166,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.StringVar(&cfg.ParquetQueryableDefaultBlockStore, "querier.parquet-queryable-default-block-store", string(parquetBlockStore), "[Experimental] Parquet queryable's default block store to query. Valid options are tsdb and parquet. If it is set to tsdb, parquet queryable always fallback to store gateway.")
 	f.BoolVar(&cfg.HonorProjectionHints, "querier.honor-projection-hints", false, "[Experimental] If true, querier will honor projection hints and only materialize requested labels. Today, projection is only effective when Parquet Queryable is enabled. Projection is only applied when not querying mixed block types (parquet and non-parquet) and not querying ingesters.")
 	f.BoolVar(&cfg.DistributedExecEnabled, "querier.distributed-exec-enabled", false, "Experimental: Enables distributed execution of queries by passing logical query plan fragments to downstream components.")
+	f.IntVar(&cfg.DistributedExecShardCount, "querier.distributed-exec-shard-count", 0, "Experimental: Number of shards the distributed-execution optimizer fans a shardable aggregation into. Only effective when -querier.distributed-exec-enabled is true. 0 (default) disables aggregation sharding.")
 	f.BoolVar(&cfg.ParquetQueryableFallbackDisabled, "querier.parquet-queryable-fallback-disabled", false, "[Experimental] Disable Parquet queryable to fallback queries to Store Gateway if the block is not available as Parquet files but available in TSDB. Setting this to true will disable the fallback and users can remove Store Gateway. But need to make sure Parquet files are created before it is queryable.")
 	f.BoolVar(&cfg.TimeoutClassificationEnabled, "querier.timeout-classification-enabled", false, "If true, classify query timeouts as 4XX (user error) or 5XX (system error) based on phase timing.")
 	f.DurationVar(&cfg.TimeoutClassificationDeadline, "querier.timeout-classification-deadline", time.Minute+59*time.Second, "The total time before the querier proactively cancels a query for timeout classification. Set this a few seconds less than the querier timeout.")