Reduce HLL memory usage

jeschkies · jeschkies · commit 5039b1d0602b · 2024-12-30T10:46:17.000+01:00
diff --git a/pkg/logcli/query/query_test.go b/pkg/logcli/query/query_test.go
@@ -5,10 +5,8 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"iter"
 	"os"
 	"path/filepath"
-	"slices"
 	"strings"
 	"testing"
 	"time"
@@ -413,7 +411,7 @@ type testQueryClient struct {
 }
 
 func newTestQueryClient(testStreams ...logproto.Stream) *testQueryClient {
-	q := logql.NewMockQuerier(0, func() iter.Seq2[int, logproto.Stream] { return slices.All(testStreams)})
+	q := logql.NewMockQuerier(0, testStreams)
 	e := logql.NewEngine(logql.EngineOpts{}, q, logql.NoLimits, log.NewNopLogger())
 	return &testQueryClient{
 		engine:          e,
diff --git a/pkg/logql/downstream_test.go b/pkg/logql/downstream_test.go
@@ -3,9 +3,7 @@ package logql
 import (
 	"context"
 	"fmt"
-	"iter"
 	"math"
-	"slices"
 	"testing"
 	"time"
 
@@ -87,7 +85,7 @@ func TestMappingEquivalence(t *testing.T) {
 	} {
 		q := NewMockQuerier(
 			shards,
-			func() iter.Seq2[int, logproto.Stream] { return slices.All(streams) },
+			streams,
 		)
 
 		opts := EngineOpts{}
@@ -195,7 +193,7 @@ func TestMappingEquivalenceSketches(t *testing.T) {
 	} {
 		q := NewMockQuerier(
 			shards,
-			func() iter.Seq2[int, logproto.Stream] { return slices.All(streams) },
+			streams,
 		)
 
 		opts := EngineOpts{
@@ -296,7 +294,7 @@ func TestApproxTopkSketches(t *testing.T) {
 		shardedQuery  string
 		regularQuery  string
 		realtiveError float64
-		// cardinalityEstimate int
+		//cardinalityEstimate int
 	}{
 		// Note:our data generation results in less spread between topk things for 10k streams than for 100k streams
 		// if we have 1k streams, we can get much more accurate results for topk 10 than topk 100
@@ -306,7 +304,7 @@ func TestApproxTopkSketches(t *testing.T) {
 			shardedQuery:  `approx_topk(3, sum by (a) (sum_over_time ({a=~".+"} | logfmt | unwrap value [1s])))`,
 			regularQuery:  `topk(3, sum by (a) (sum_over_time ({a=~".+"} | logfmt | unwrap value [1s])))`,
 			realtiveError: 0.0012,
-			// cardinalityEstimate: 3,
+			//cardinalityEstimate: 3,
 		},
 		{
 			labelShards:   10,
@@ -349,7 +347,7 @@ func TestApproxTopkSketches(t *testing.T) {
 
 			q := NewMockQuerier(
 				tc.labelShards,
-				func() iter.Seq2[int, logproto.Stream] { return slices.All(streams) },
+				streams,
 			)
 
 			opts := EngineOpts{
@@ -429,7 +427,7 @@ func TestShardCounter(t *testing.T) {
 	} {
 		q := NewMockQuerier(
 			shards,
-			func() iter.Seq2[int, logproto.Stream] { return slices.All(streams) },
+			streams,
 		)
 
 		opts := EngineOpts{}
@@ -688,7 +686,7 @@ func TestRangeMappingEquivalence(t *testing.T) {
 	} {
 		q := NewMockQuerier(
 			shards,
-			func() iter.Seq2[int, logproto.Stream] { return slices.All(streams) },
+			streams,
 		)
 
 		opts := EngineOpts{}
diff --git a/pkg/logql/engine_test.go b/pkg/logql/engine_test.go
@@ -5,9 +5,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	goiter "iter"
 	"math"
-	"math/rand"
 	"strings"
 	"testing"
 	"time"
@@ -2512,54 +2510,6 @@ func BenchmarkRangeQuery1000000(b *testing.B) {
 	benchmarkRangeQuery(int64(1000000), b)
 }
 
-func BenchmarkHighCardinalityInstantQuery(b *testing.B) {
-	b.ReportAllocs()
-
-	mockQuerier := NewMockQuerier(1, func() goiter.Seq2[int, logproto.Stream] { return generateLogs(1_000_000, 100_000) })
-
-	eng := NewEngine(EngineOpts{}, mockQuerier, NoLimits, log.NewNopLogger())
-	start := time.Unix(1, 0)
-	b.ResetTimer()
-
-	for i := 0; i < b.N; i++ {
-		params, err := NewLiteralParams(`sum by(cardinality) (sum_over_time({app="foo"} | logfmt | unwrap value [1s]))`, start, start, 0, 0, logproto.BACKWARD, 1000, nil, nil)
-		require.NoError(b, err)
-		q := eng.Query(params)
-
-		res, err := q.Exec(user.InjectOrgID(context.Background(), "fake"))
-		if err != nil {
-			b.Fatal(err)
-		}
-		result = res.Data
-		if result == nil {
-			b.Fatal("unexpected nil result")
-		}
-	}
-}
-
-func generateLogs(lines int, cardinality int) goiter.Seq2[int, logproto.Stream] {
-	batchSize := 10
-	return func(yield func(int, logproto.Stream) bool) {
-		for batch := range lines / batchSize {
-			entries := make([]logproto.Entry, batchSize)
-			for i := 0; i < batchSize; i++ {
-				line := batch*batchSize + i
-				entries[i] = logproto.Entry{
-					Timestamp: time.Unix(0, int64(batch)),
-					Line:      fmt.Sprintf("cardinality=%d batch=%d line=%d value=%f", i%cardinality, batch, line, rand.Float64()),
-				}
-			}
-			stream := logproto.Stream{
-				Entries: entries,
-				Labels:  `{app="foo"}`,
-			}
-			if !yield(batch, stream) {
-				return
-			}
-		}
-	}
-}
-
 var result promql_parser.Value
 
 func benchmarkRangeQuery(testsize int64, b *testing.B) {
diff --git a/pkg/logql/explain_test.go b/pkg/logql/explain_test.go
@@ -2,8 +2,6 @@ package logql
 
 import (
 	"context"
-	"iter"
-	"slices"
 	"testing"
 	"time"
 
@@ -12,7 +10,6 @@ import (
 
 	"github.com/grafana/dskit/user"
 
-	"github.com/grafana/loki/v3/pkg/logproto"
 	"github.com/grafana/loki/v3/pkg/logql/syntax"
 )
 
@@ -21,7 +18,7 @@ func TestExplain(t *testing.T) {
 
 	// TODO(karsten): Ideally the querier and downstreamer are not required
 	// to create the step evaluators.
-	querier := NewMockQuerier(4, func() iter.Seq2[int, logproto.Stream] { return slices.All(make([]logproto.Stream, 0)) })
+	querier := NewMockQuerier(4, nil)
 	opts := EngineOpts{}
 	regular := NewEngine(opts, querier, NoLimits, log.NewNopLogger())
 
diff --git a/pkg/logql/test_utils.go b/pkg/logql/test_utils.go
@@ -3,7 +3,6 @@ package logql
 import (
 	"context"
 	"fmt"
-	goiter "iter"
 	logger "log"
 	"math/rand"
 	"sort"
@@ -26,7 +25,7 @@ import (
 
 const ConCurrency = 100
 
-func NewMockQuerier(shards int, streams func() goiter.Seq2[int, logproto.Stream]) MockQuerier {
+func NewMockQuerier(shards int, streams []logproto.Stream) MockQuerier {
 	return MockQuerier{
 		shards:  shards,
 		streams: streams,
@@ -36,7 +35,7 @@ func NewMockQuerier(shards int, streams func() goiter.Seq2[int, logproto.Stream]
 // Shard aware mock querier
 type MockQuerier struct {
 	shards  int
-	streams func() goiter.Seq2[int, logproto.Stream]
+	streams []logproto.Stream
 }
 
 func (q MockQuerier) extractOldShard(xs []string) (*index.ShardAnnotation, error) {
@@ -75,7 +74,7 @@ func (q MockQuerier) SelectLogs(_ context.Context, req SelectLogParams) (iter.En
 	var matched []logproto.Stream
 
 outer:
-	for _, stream := range q.streams() {
+	for _, stream := range q.streams {
 		ls := mustParseLabels(stream.Labels)
 
 		// filter by shard if requested
@@ -113,7 +112,6 @@ outer:
 	return iter.NewSortEntryIterator(streamIters, req.Direction), nil
 }
 
-// TODO: check if this is the same as in storage
 func processStream(in []logproto.Stream, pipeline log.Pipeline) []logproto.Stream {
 	resByStream := map[string]*logproto.Stream{}
 
@@ -203,7 +201,7 @@ func (q MockQuerier) SelectSamples(_ context.Context, req SelectSampleParams) (i
 	var matched []logproto.Stream
 
 outer:
-	for _, stream := range q.streams() {
+	for _, stream := range q.streams {
 		ls := mustParseLabels(stream.Labels)
 
 		// filter by shard if requested
