MQE: rename statsTracker to subsetStats, introduce support for serializing OperatorEvaluationStats (#15045)

#### What this PR does

This PR makes two changes to MQE's support for tracking query
statistics:

* it renames `statsTracker` to `subsetStats`, which better conveys its
purpose
* it introduces support for serializing `OperatorEvaluationStats` to /
from Protobuf, which is a prerequisite for supporting query stats with
remote execution and range vector splitting

#### Which issue(s) this PR fixes or relates to

(none)

#### Checklist

- [x] Tests updated.
- [n/a] Documentation added.
- [x] `CHANGELOG.md` updated - the order of entries should be
`[CHANGE]`, `[FEATURE]`, `[ENHANCEMENT]`, `[BUGFIX]`. If changelog entry
is not needed, please add the `changelog-not-needed` label to the PR.
- [n/a]
[`about-versioning.md`](https://github.com/grafana/mimir/blob/main/docs/sources/mimir/configure/about-versioning.md)
updated with experimental features.

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Medium Risk**
> Adds a new protobuf-serialized representation for
`OperatorEvaluationStats` and changes how per-subset stats are
allocated/decoded, which could affect query stats correctness and memory
accounting in MQE paths.
> 
> **Overview**
> Adds protobuf serialization support for MQE operator evaluation stats
by introducing `EncodedOperatorEvaluationStats`/`EncodedSubsetStats`
(new `stats.proto` + generated `stats.pb.go`) and
`OperatorEvaluationStats.Encode()` / `Decode()` methods.
> 
> Refactors internal tracking by renaming `statsTracker` to
`subsetStats`, avoiding subset slice allocation when `subsetCount==0`,
and adds tests covering round-trip encoding/decoding plus validation
errors and memory-consumption tracking. Updates the changelog entry to
include the additional PR reference.
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
55be141. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

Author: charleskorn

CHANGELOG.md

@@ -59,7 +59,7 @@
 * [FEATURE] Ingest storage: Add `-ingest-storage.kafka.tls*` flags to connect to Kafka using TLS. #14550
 * [FEATURE] Ingest storage: Add `-ingest-storage.ingestion-partition-tenant-write-shard-size` to limit the number of partitions used for writes independently from reads, allowing safely reducing the shard size without losing query coverage during the migration. #14780
 * [FEATURE] MQE: Add experimental support for splitting and caching intermediate results for functions over range vectors in instant queries. #13472 #14479 #14506 #14499 #14517 #14536 #14614 #14645 #14677 #14788
-* [FEATURE] MQE: Add experimental support for reporting the number of samples read per query. #14828 #14839 #14952
+* [FEATURE] MQE: Add experimental support for reporting the number of samples read per query. #14828 #14839 #14952 #15045
 * [FEATURE] Compactor: Add `-compactor.ooo-split-and-merge-shards` per-tenant limit to allow a separate shard count for blocks with the out-of-order external label. #14704
 * [FEATURE] Distributor: add experimental support for controlling OTLP metric name suffix addition and translation strategy via `X-Mimir-OTLP-AddSuffixes` and `X-Mimir-OTLP-TranslationStrategy` request headers on the OTLP push path, gated by `-api.otlp-translation-headers-enabled` (off by default). #14782
 * [ENHANCEMENT] Distributor: Add per-tenant `-distributor.active-series-limit-response-code` override to configure the HTTP response code returned when rejecting series due to the active series limit. Defaults to 429 (Too Many Requests). Set to 400 (Bad Request) to prevent clients from retrying rejected requests. #14981

pkg/streamingpromql/types/stats.go

@@ -50,37 +50,42 @@ type OperatorEvaluationStats struct {
 	timeRange                QueryTimeRange
 	memoryConsumptionTracker *limiter.MemoryConsumptionTracker
 
-	allSeries *statsTracker
-	subsets   []*statsTracker
+	allSeries *subsetStats
+	subsets   []*subsetStats
 }
 
 // NewOperatorEvaluationStats creates a new OperatorEvaluationStats for the given time range.
 //
 // subsetCount is the number of subsets to track. It is the caller's responsibility to track
 // which subset is which.
 func NewOperatorEvaluationStats(timeRange QueryTimeRange, memoryConsumptionTracker *limiter.MemoryConsumptionTracker, subsetCount int) (*OperatorEvaluationStats, error) {
-	allSeries, err := newStatsTracker(timeRange, memoryConsumptionTracker)
+	allSeries, err := newSubsetStats(timeRange, memoryConsumptionTracker)
 	if err != nil {
 		return nil, err
 	}
 
-	subsets := make([]*statsTracker, 0, subsetCount)
-	for range subsetCount {
-		stats, err := newStatsTracker(timeRange, memoryConsumptionTracker)
-		if err != nil {
-			return nil, err
-		}
-
-		subsets = append(subsets, stats)
-	}
-
-	return &OperatorEvaluationStats{
+	stats := &OperatorEvaluationStats{
 		timeRange:                timeRange,
 		memoryConsumptionTracker: memoryConsumptionTracker,
 
 		allSeries: allSeries,
-		subsets:   subsets,
-	}, nil
+	}
+
+	if subsetCount > 0 {
+		// Only bother allocating a slice if we actually need it.
+		stats.subsets = make([]*subsetStats, 0, subsetCount)
+
+		for range subsetCount {
+			subset, err := newSubsetStats(timeRange, memoryConsumptionTracker)
+			if err != nil {
+				return nil, err
+			}
+
+			stats.subsets = append(stats.subsets, subset)
+		}
+	}
+
+	return stats, nil
 }
 
 // TrackSampleForInstantVectorSelector records a sample for an instant vector selector at output timestamp stepT.
@@ -270,6 +275,54 @@ func (s *OperatorEvaluationStats) ComputeForSubquery(
 	return result, nil
 }
 
+// Encode returns the encoded form of this instance, suitable for serialization.
+// The encoded form may share memory with this instance, and so may be modified
+// if this instance is modified, and becomes invalid when this instance is closed.
+func (s *OperatorEvaluationStats) Encode() *EncodedOperatorEvaluationStats {
+	encoded := &EncodedOperatorEvaluationStats{
+		AllSeries: s.allSeries.Encode(),
+	}
+
+	if len(s.subsets) > 0 {
+		// Only bother allocating a slice for subsets if there are any.
+		encoded.Subsets = make([]EncodedSubsetStats, 0, len(s.subsets))
+		for _, subset := range s.subsets {
+			encoded.Subsets = append(encoded.Subsets, subset.Encode())
+		}
+	}
+
+	return encoded
+}
+
+func (e *EncodedOperatorEvaluationStats) Decode(timeRange QueryTimeRange, memoryConsumptionTracker *limiter.MemoryConsumptionTracker) (*OperatorEvaluationStats, error) {
+	allSeries, err := e.AllSeries.decode(timeRange, memoryConsumptionTracker)
+	if err != nil {
+		return nil, err
+	}
+
+	decoded := &OperatorEvaluationStats{
+		timeRange:                timeRange,
+		memoryConsumptionTracker: memoryConsumptionTracker,
+		allSeries:                allSeries,
+	}
+
+	if len(e.Subsets) > 0 {
+		// Only bother allocating a slice for subsets if there are any.
+		decoded.subsets = make([]*subsetStats, 0, len(e.Subsets))
+
+		for _, encodedSubset := range e.Subsets {
+			decodedSubset, err := encodedSubset.decode(timeRange, memoryConsumptionTracker)
+			if err != nil {
+				return nil, err
+			}
+
+			decoded.subsets = append(decoded.subsets, decodedSubset)
+		}
+	}
+
+	return decoded, nil
+}
+
 func (s *OperatorEvaluationStats) Close() {
 	s.allSeries.Close()
 
@@ -278,14 +331,14 @@ func (s *OperatorEvaluationStats) Close() {
 	}
 }
 
-type statsTracker struct {
+type subsetStats struct {
 	samplesProcessedPerStep []int64
 	newSamplesReadPerStep   []int64
 
 	memoryConsumptionTracker *limiter.MemoryConsumptionTracker
 }
 
-func newStatsTracker(timeRange QueryTimeRange, memoryConsumptionTracker *limiter.MemoryConsumptionTracker) (*statsTracker, error) {
+func newSubsetStats(timeRange QueryTimeRange, memoryConsumptionTracker *limiter.MemoryConsumptionTracker) (*subsetStats, error) {
 	samplesProcessed, err := Int64SlicePool.Get(timeRange.StepCount, memoryConsumptionTracker)
 	if err != nil {
 		return nil, err
@@ -296,19 +349,19 @@ func newStatsTracker(timeRange QueryTimeRange, memoryConsumptionTracker *limiter
 		return nil, err
 	}
 
-	return &statsTracker{
+	return &subsetStats{
 		samplesProcessedPerStep:  samplesProcessed[:timeRange.StepCount],
 		newSamplesReadPerStep:    newSamplesRead[:timeRange.StepCount],
 		memoryConsumptionTracker: memoryConsumptionTracker,
 	}, nil
 }
 
-func (s *statsTracker) Add(pointIndex int64, samplesProcessed int64, newSamplesRead int64) {
+func (s *subsetStats) Add(pointIndex int64, samplesProcessed int64, newSamplesRead int64) {
 	s.samplesProcessedPerStep[pointIndex] += samplesProcessed
 	s.newSamplesReadPerStep[pointIndex] += newSamplesRead
 }
 
-func (s *statsTracker) SetFromSubquery(source *statsTracker, parentIdx, firstInnerIdx, firstNewSamplesInnerIdx, lastIdx int) {
+func (s *subsetStats) SetFromSubquery(source *subsetStats, parentIdx, firstInnerIdx, firstNewSamplesInnerIdx, lastIdx int) {
 	for innerIdx := firstInnerIdx; innerIdx <= lastIdx; innerIdx++ {
 		s.samplesProcessedPerStep[parentIdx] += source.samplesProcessedPerStep[innerIdx]
 	}
@@ -318,19 +371,46 @@ func (s *statsTracker) SetFromSubquery(source *statsTracker, parentIdx, firstInn
 	}
 }
 
-func (s *statsTracker) SetFromStepInvariant(samplesProcessed int64, newSamplesRead int64) {
+func (s *subsetStats) SetFromStepInvariant(samplesProcessed int64, newSamplesRead int64) {
 	s.newSamplesReadPerStep[0] = newSamplesRead
 	for idx := range s.samplesProcessedPerStep {
 		s.samplesProcessedPerStep[idx] = samplesProcessed
 	}
 }
 
-func (s *statsTracker) CopyFrom(source *statsTracker) {
+func (s *subsetStats) CopyFrom(source *subsetStats) {
 	copy(s.samplesProcessedPerStep, source.samplesProcessedPerStep)
 	copy(s.newSamplesReadPerStep, source.newSamplesReadPerStep)
 }
 
-func (s *statsTracker) Close() {
+func (s *subsetStats) Encode() EncodedSubsetStats {
+	return EncodedSubsetStats{
+		SamplesProcessedPerStep: s.samplesProcessedPerStep,
+		NewSamplesReadPerStep:   s.newSamplesReadPerStep,
+	}
+}
+
+func (e *EncodedSubsetStats) decode(timeRange QueryTimeRange, memoryConsumptionTracker *limiter.MemoryConsumptionTracker) (*subsetStats, error) {
+	if len(e.SamplesProcessedPerStep) != timeRange.StepCount {
+		return nil, fmt.Errorf("number of samples processed steps in encoded form (%d) does not match expected (%d)", len(e.SamplesProcessedPerStep), timeRange.StepCount)
+	}
+
+	if len(e.NewSamplesReadPerStep) != timeRange.StepCount {
+		return nil, fmt.Errorf("number of new samples read steps in encoded form (%d) does not match expected (%d)", len(e.NewSamplesReadPerStep), timeRange.StepCount)
+	}
+
+	if err := memoryConsumptionTracker.IncreaseMemoryConsumption(uint64(cap(e.SamplesProcessedPerStep)+cap(e.NewSamplesReadPerStep))*Int64Size, limiter.Int64Slices); err != nil {
+		return nil, err
+	}
+
+	return &subsetStats{
+		samplesProcessedPerStep:  e.SamplesProcessedPerStep,
+		newSamplesReadPerStep:    e.NewSamplesReadPerStep,
+		memoryConsumptionTracker: memoryConsumptionTracker,
+	}, nil
+}
+
+func (s *subsetStats) Close() {
 	Int64SlicePool.Put(&s.samplesProcessedPerStep, s.memoryConsumptionTracker)
 	Int64SlicePool.Put(&s.newSamplesReadPerStep, s.memoryConsumptionTracker)
 }