feat: Introduce shardable probabilistic topk for instant queries.

pkg/logql/count_min_sketch.go

@@ -17,8 +17,8 @@ const (
 	CountMinSketchVectorType = "CountMinSketchVector"
 
 	epsilon = 0.0001
-	// delta of 0.05 results in a sketch size of 325KB, 0.01 gives a sketch size of 543KB
-	delta = 0.05
+	// delta of 0.01 results in a sketch size of 27183 * 7 * 4 bytes = 761,124 bytes, 0.05 would yield 543,660 bytes
+	delta = 0.01
 )
 
 // CountMinSketchVector tracks the count or sum of values of a metric, ie list of label value pairs. It's storage for

pkg/logql/downstream_test.go

@@ -283,6 +283,12 @@ func TestApproxTopkSketches(t *testing.T) {
 		limit  = 100
 	)
 
+	limits := &fakeLimits{
+		maxSeries: math.MaxInt64,
+		// timeout:   5 * time.Minute,
+		timeout: time.Hour,
+	}
+
 	for _, tc := range []struct {
 		labelShards   int
 		totalStreams  int
@@ -334,8 +340,14 @@ func TestApproxTopkSketches(t *testing.T) {
 			regularQuery:  `topk(100, sum by (a) (sum_over_time ({a=~".+"} | logfmt | unwrap value [1s])))`,
 			realtiveError: 0.0015,
 		},
+		{
+			labelShards:   10, // increasing this will make the test too slow
+			totalStreams:  1_000_000,
+			shardedQuery:  `approx_topk(100, sum by (a) (sum_over_time ({a=~".+"} | logfmt | unwrap value [1s])))`,
+			regularQuery:  `topk(100, sum by (a) (sum_over_time ({a=~".+"} | logfmt | unwrap value [1s])))`,
+			realtiveError: 0.0015,
+		},
 	} {
-
 		t.Run(fmt.Sprintf("%s/%d/%d", tc.shardedQuery, tc.labelShards, tc.totalStreams), func(t *testing.T) {
 			streams := randomStreams(tc.totalStreams, rounds+1, tc.labelShards, []string{"a", "b", "c", "d"}, true)
 
@@ -347,8 +359,8 @@ func TestApproxTopkSketches(t *testing.T) {
 			opts := EngineOpts{
 				MaxCountMinSketchHeapSize: 10_000,
 			}
-			regular := NewEngine(opts, q, NoLimits, log.NewNopLogger())
-			sharded := NewDownstreamEngine(opts, MockDownstreamer{regular}, NoLimits, log.NewNopLogger())
+			regular := NewEngine(opts, q, limits, log.NewNopLogger())
+			sharded := NewDownstreamEngine(opts, MockDownstreamer{regular}, limits, log.NewNopLogger())
 
 			// for an instant query we set the start and end to the same timestamp
 			// plus set step and interval to 0
@@ -365,7 +377,9 @@ func TestApproxTopkSketches(t *testing.T) {
 			)
 			require.NoError(t, err)
 			qry := regular.Query(params.Copy())
-			ctx := user.InjectOrgID(context.Background(), "fake")
+			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+			defer cancel()
+			ctx = user.InjectOrgID(ctx, "fake")
 
 			strategy := NewPowerOfTwoStrategy(ConstantShards(tc.labelShards))
 			mapper := NewShardMapper(strategy, nilShardMetrics, []string{ShardQuantileOverTime, SupportApproxTopk})

pkg/logql/sketch/cms.go

@@ -19,9 +19,12 @@ func NewCountMinSketch(w, d uint32) (*CountMinSketch, error) {
 	}, nil
 }
 
+// NewCountMinSketchFromErrorAndProbability creates a new CMS for a given epsilon and delta. The sketch width and depth
+// are calculated according to the RedisBloom implementation.
+// See https://github.com/RedisBloom/RedisBloom/blob/7bc047d1ea4113419b60eb6446ac3d4e61877a7b/src/cms.c#L38-L39
 func NewCountMinSketchFromErrorAndProbability(epsilon float64, delta float64) (*CountMinSketch, error) {
 	width := math.Ceil(math.E / epsilon)
-	depth := math.Ceil(math.Log(1.0 / delta))
+	depth := math.Ceil(math.Log(delta) / math.Log(0.5))
 	return NewCountMinSketch(uint32(width), uint32(depth))
 }
 

pkg/logql/test_utils.go

@@ -7,10 +7,12 @@
 	"math/rand"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 	"unicode"
 
 	"github.com/cespare/xxhash/v2"
+	"github.com/grafana/dskit/concurrency"
 	"github.com/prometheus/prometheus/model/labels"
 	promql_parser "github.com/prometheus/prometheus/promql/parser"
 
@@ -21,6 +23,8 @@
 	"github.com/grafana/loki/v3/pkg/storage/stores/shipper/indexshipper/tsdb/index"
 )
 
+const CON_CURRENCY = 100
+
 func NewMockQuerier(shards int, streams []logproto.Stream) MockQuerier {
 	return MockQuerier{
 		shards:  shards,
@@ -45,7 +49,6 @@
 	}
 
 	return parsed[0].PowerOfTwo, nil
-
 }
 
 func (q MockQuerier) SelectLogs(_ context.Context, req SelectLogParams) (iter.EntryIterator, error) {
@@ -113,8 +116,8 @@
 	resByStream := map[string]*logproto.Stream{}
 
 	for _, stream := range in {
+		sp := pipeline.ForStream(mustParseLabels(stream.Labels))
 		for _, e := range stream.Entries {
-			sp := pipeline.ForStream(mustParseLabels(stream.Labels))
 			if l, out, matches := sp.Process(e.Timestamp.UnixNano(), []byte(e.Line)); matches {
 				var s *logproto.Stream
 				var found bool
@@ -137,12 +140,12 @@
 	return streams
 }
 
-func processSeries(in []logproto.Stream, ex log.SampleExtractor) []logproto.Series {
+func processSeries(in []logproto.Stream, ex log.SampleExtractor) ([]logproto.Series, error) {
 	resBySeries := map[string]*logproto.Series{}
 
 	for _, stream := range in {
+		exs := ex.ForStream(mustParseLabels(stream.Labels))
 		for _, e := range stream.Entries {
-			exs := ex.ForStream(mustParseLabels(stream.Labels))
 			if f, lbs, ok := exs.Process(e.Timestamp.UnixNano(), []byte(e.Line)); ok {
 				var s *logproto.Series
 				var found bool
@@ -151,6 +154,7 @@
 					s = &logproto.Series{Labels: lbs.String(), StreamHash: exs.BaseLabels().Hash()}
 					resBySeries[lbs.String()] = s
 				}
+
 				s.Samples = append(s.Samples, logproto.Sample{
 					Timestamp: e.Timestamp.UnixNano(),
 					Value:     f,
@@ -159,15 +163,16 @@
 			}
 		}
 	}
+
 	series := []logproto.Series{}
 	for _, s := range resBySeries {
 		sort.Sort(s)
 		series = append(series, *s)
 	}
-	return series
+	return series, nil
 }
 
-func (q MockQuerier) SelectSamples(_ context.Context, req SelectSampleParams) (iter.SampleIterator, error) {
+func (q MockQuerier) SelectSamples(ctx context.Context, req SelectSampleParams) (iter.SampleIterator, error) {
 	selector, err := req.LogSelector()
 	if err != nil {
 		return nil, err
@@ -212,7 +217,10 @@
 		matched = append(matched, stream)
 	}
 
-	filtered := processSeries(matched, extractor)
+	filtered, err := processSeries(matched, extractor)
+	if err != nil {
+		return nil, err
+	}
 
 	return iter.NewTimeRangedSampleIterator(
 		iter.NewMultiSeriesIterator(filtered),
@@ -228,15 +236,19 @@
 func (m MockDownstreamer) Downstreamer(_ context.Context) Downstreamer { return m }
 
 func (m MockDownstreamer) Downstream(ctx context.Context, queries []DownstreamQuery, acc Accumulator) ([]logqlmodel.Result, error) {
-	for i, query := range queries {
-		res, err := m.Query(query.Params).Exec(ctx)
+	mu := sync.Mutex{}
+	err := concurrency.ForEachJob(ctx, len(queries), CON_CURRENCY, func(ctx context.Context, idx int) error {
+		res, err := m.Query(queries[idx].Params).Exec(ctx)
 		if err != nil {
-			return nil, err
-		}
-		err = acc.Accumulate(ctx, res, i)
-		if err != nil {
-			return nil, err
+			return err
 		}
+		mu.Lock()
+		defer mu.Unlock()
+		err = acc.Accumulate(ctx, res, idx)
+		return err
+	})
+	if err != nil {
+		return nil, err
 	}
 
 	return acc.Result(), nil