feat(blooms): compute chunks once

cmd/migrate/main.go

@@ -309,7 +309,7 @@ func (m *chunkMover) moveChunks(ctx context.Context, threadID int, syncRangeCh <
 			start := time.Now()
 			var totalBytes uint64
 			var totalChunks uint64
-			schemaGroups, fetchers, err := m.source.GetChunks(m.ctx, m.sourceUser, model.TimeFromUnixNano(sr.from), model.TimeFromUnixNano(sr.to), chunk.NewPredicate(m.matchers, nil))
+			schemaGroups, fetchers, err := m.source.GetChunks(m.ctx, m.sourceUser, model.TimeFromUnixNano(sr.from), model.TimeFromUnixNano(sr.to), chunk.NewPredicate(m.matchers, nil), nil)
 			if err != nil {
 				log.Println(threadID, "Error querying index for chunk refs:", err)
 				errCh <- err

pkg/indexgateway/client.go

@@ -28,9 +28,7 @@ import (
 
 	"github.com/grafana/loki/v3/pkg/distributor/clientpool"
 	"github.com/grafana/loki/v3/pkg/logproto"
-	"github.com/grafana/loki/v3/pkg/logql/syntax"
 	"github.com/grafana/loki/v3/pkg/storage/stores/series/index"
-	"github.com/grafana/loki/v3/pkg/storage/stores/shipper/indexshipper/tsdb/sharding"
 	"github.com/grafana/loki/v3/pkg/util/constants"
 	"github.com/grafana/loki/v3/pkg/util/discovery"
 	util_math "github.com/grafana/loki/v3/pkg/util/math"
@@ -340,8 +338,7 @@ func (s *GatewayClient) GetShards(
 				if err != nil {
 					return errors.WithStack(err)
 				}
-				perReplicaResult.Shards = append(perReplicaResult.Shards, resp.Shards...)
-				perReplicaResult.Statistics.Merge(resp.Statistics)
+				perReplicaResult.Merge(resp)
 			}
 
 			// Since `poolDo` retries on error, we only want to set the response if we got a successful response.
@@ -355,48 +352,11 @@ func (s *GatewayClient) GetShards(
 			return errCt <= maxErrs
 		},
 	); err != nil {
-		if isUnimplementedCallError(err) {
-			return s.getShardsFromStatsFallback(ctx, in)
-		}
 		return nil, err
 	}
 	return res, nil
 }
 
-func (s *GatewayClient) getShardsFromStatsFallback(
-	ctx context.Context,
-	in *logproto.ShardsRequest,
-) (*logproto.ShardsResponse, error) {
-	userID, err := tenant.TenantID(ctx)
-	if err != nil {
-		return nil, errors.Wrap(err, "index gateway client get tenant ID")
-	}
-
-	p, err := ExtractShardRequestMatchersAndAST(in.Query)
-	if err != nil {
-		return nil, errors.Wrap(err, "failure while falling back to stats for shard calculation")
-
-	}
-
-	stats, err := s.GetStats(
-		ctx,
-		&logproto.IndexStatsRequest{
-			From:     in.From,
-			Through:  in.Through,
-			Matchers: (&syntax.MatchersExpr{Mts: p.Matchers}).String(),
-		},
-	)
-	if err != nil {
-		return nil, err
-	}
-
-	var strategy sharding.PowerOfTwoSharding
-	shards := strategy.ShardsFor(stats.Bytes, uint64(s.limits.TSDBMaxBytesPerShard(userID)))
-	return &logproto.ShardsResponse{
-		Shards: shards,
-	}, nil
-}
-
 // TODO(owen-d): this was copied from ingester_querier.go -- move it to a shared pkg
 // isUnimplementedCallError tells if the GRPC error is a gRPC error with code Unimplemented.
 func isUnimplementedCallError(err error) bool {

pkg/indexgateway/gateway.go

@@ -214,7 +214,7 @@ func (g *Gateway) GetChunkRef(ctx context.Context, req *logproto.GetChunkRefRequ
 	}
 
 	predicate := chunk.NewPredicate(matchers, &req.Plan)
-	chunks, _, err := g.indexQuerier.GetChunks(ctx, instanceID, req.From, req.Through, predicate)
+	chunks, _, err := g.indexQuerier.GetChunks(ctx, instanceID, req.From, req.Through, predicate, nil)
 	if err != nil {
 		return nil, err
 	}
@@ -417,7 +417,7 @@ func (g *Gateway) getShardsWithBlooms(
 	defer sp.Finish()
 
 	// 1) for all bounds, get chunk refs
-	grps, _, err := g.indexQuerier.GetChunks(ctx, instanceID, req.From, req.Through, p)
+	grps, _, err := g.indexQuerier.GetChunks(ctx, instanceID, req.From, req.Through, p, nil)
 	if err != nil {
 		return err
 	}
@@ -463,11 +463,12 @@ func (g *Gateway) getShardsWithBlooms(
 			},
 		}
 	} else {
-		shards, err := accumulateChunksToShards(ctx, instanceID, forSeries, req, p, filtered)
+		shards, chunkGrps, err := accumulateChunksToShards(ctx, instanceID, forSeries, req, p, filtered)
 		if err != nil {
 			return err
 		}
 		resp.Shards = shards
+		resp.ChunkGroups = chunkGrps
 	}
 
 	sp.LogKV("msg", "send shards response", "shards", len(resp.Shards))
@@ -525,7 +526,7 @@ func accumulateChunksToShards(
 	req *logproto.ShardsRequest,
 	p chunk.Predicate,
 	filtered []*logproto.ChunkRef,
-) ([]logproto.Shard, error) {
+) ([]logproto.Shard, []logproto.ChunkRefGroup, error) {
 	// map for looking up post-filtered chunks in O(n) while iterating the index again for sizing info
 	filteredM := make(map[model.Fingerprint][]refWithSizingInfo, 1024)
 	for _, ref := range filtered {
@@ -579,7 +580,7 @@ func accumulateChunksToShards(
 		},
 		p.Matchers...,
 	); err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	collectedSeries := sharding.SizedFPs(sharding.SizedFPsPool.Get(len(filteredM)))
@@ -597,7 +598,21 @@ func accumulateChunksToShards(
 	}
 	sort.Sort(collectedSeries)
 
-	return collectedSeries.ShardsFor(req.TargetBytesPerShard), nil
+	shards := collectedSeries.ShardsFor(req.TargetBytesPerShard)
+	chkGrps := make([]logproto.ChunkRefGroup, 0, len(shards))
+	for _, s := range shards {
+		from := sort.Search(len(filtered), func(i int) bool {
+			return filtered[i].Fingerprint >= uint64(s.Bounds.Min)
+		})
+		through := sort.Search(len(filtered), func(i int) bool {
+			return filtered[i].Fingerprint > uint64(s.Bounds.Max)
+		})
+		chkGrps = append(chkGrps, logproto.ChunkRefGroup{
+			Refs: filtered[from:through],
+		})
+	}
+
+	return shards, chkGrps, nil
 }
 
 type refWithSizingInfo struct {

pkg/indexgateway/gateway_test.go

@@ -532,7 +532,7 @@ func TestAccumulateChunksToShards(t *testing.T) {
 		},
 	}
 
-	shards, err := accumulateChunksToShards(
+	shards, grps, err := accumulateChunksToShards(
 		context.Background(),
 		"",
 		fsImpl(series),
@@ -543,6 +543,12 @@ func TestAccumulateChunksToShards(t *testing.T) {
 		filtered,
 	)
 
+	expectedChks := [][]*logproto.ChunkRef{
+		filtered[0:3],
+		filtered[3:6],
+		filtered[6:9],
+		filtered[9:10],
+	}
 	exp := []logproto.Shard{
 		{
 			Bounds: logproto.FPBounds{Min: 0, Max: 1},
@@ -586,6 +592,9 @@ func TestAccumulateChunksToShards(t *testing.T) {
 
 	for i := range shards {
 		require.Equal(t, exp[i], shards[i], "invalid shard at index %d", i)
+		for j := range grps[i].Refs {
+			require.Equal(t, expectedChks[i][j], grps[i].Refs[j], "invalid chunk in grp %d at index %d", i, j)
+		}
 	}
 	require.Equal(t, len(exp), len(shards))
 

pkg/ingester/flush_test.go

@@ -361,7 +361,7 @@ func (s *testStore) SelectSeries(_ context.Context, _ logql.SelectLogParams) ([]
 	return nil, nil
 }
 
-func (s *testStore) GetChunks(_ context.Context, _ string, _, _ model.Time, _ chunk.Predicate) ([][]chunk.Chunk, []*fetcher.Fetcher, error) {
+func (s *testStore) GetChunks(_ context.Context, _ string, _, _ model.Time, _ chunk.Predicate, _ *logproto.ChunkRefGroup) ([][]chunk.Chunk, []*fetcher.Fetcher, error) {
 	return nil, nil, nil
 }
 

pkg/ingester/ingester.go

@@ -1036,7 +1036,7 @@ func (i *Ingester) GetChunkIDs(ctx context.Context, req *logproto.GetChunkIDsReq
 	}
 
 	// get chunk references
-	chunksGroups, _, err := i.store.GetChunks(ctx, orgID, start, end, chunk.NewPredicate(matchers, nil))
+	chunksGroups, _, err := i.store.GetChunks(ctx, orgID, start, end, chunk.NewPredicate(matchers, nil), nil)
 	if err != nil {
 		return nil, err
 	}

pkg/ingester/ingester_test.go

@@ -457,7 +457,7 @@ func (s *mockStore) PutOne(_ context.Context, _, _ model.Time, _ chunk.Chunk) er
 	return nil
 }
 
-func (s *mockStore) GetChunks(_ context.Context, _ string, _, _ model.Time, _ chunk.Predicate) ([][]chunk.Chunk, []*fetcher.Fetcher, error) {
+func (s *mockStore) GetChunks(_ context.Context, _ string, _, _ model.Time, _ chunk.Predicate, _ *logproto.ChunkRefGroup) ([][]chunk.Chunk, []*fetcher.Fetcher, error) {
 	return nil, nil, nil
 }
 

pkg/logcli/client/file.go

@@ -77,6 +77,7 @@ func (f *FileClient) Query(q string, limit int, t time.Time, direction logproto.
 		direction,
 		uint32(limit),
 		nil,
+		nil,
 	)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse query: %w", err)
@@ -118,6 +119,7 @@ func (f *FileClient) QueryRange(queryStr string, limit int, start, end time.Time
 		direction,
 		uint32(limit),
 		nil,
+		nil,
 	)
 	if err != nil {
 		return nil, err

pkg/logcli/query/query.go

@@ -490,6 +490,7 @@ func (q *Query) DoLocalQuery(out output.LogOutput, statistics bool, orgID string
 			q.resultsDirection(),
 			uint32(q.Limit),
 			nil,
+			nil,
 		)
 		if err != nil {
 			return err
@@ -506,6 +507,7 @@ func (q *Query) DoLocalQuery(out output.LogOutput, statistics bool, orgID string
 			q.resultsDirection(),
 			uint32(q.Limit),
 			nil,
+			nil,
 		)
 		if err != nil {
 			return err

pkg/logcli/query/query_test.go

@@ -426,7 +426,7 @@ func (t *testQueryClient) Query(_ string, _ int, _ time.Time, _ logproto.Directi
 func (t *testQueryClient) QueryRange(queryStr string, limit int, from, through time.Time, direction logproto.Direction, step, interval time.Duration, _ bool) (*loghttp.QueryResponse, error) {
 	ctx := user.InjectOrgID(context.Background(), "fake")
 
-	params, err := logql.NewLiteralParams(queryStr, from, through, step, interval, direction, uint32(limit), nil)
+	params, err := logql.NewLiteralParams(queryStr, from, through, step, interval, direction, uint32(limit), nil, nil)
 	if err != nil {
 		return nil, err
 	}

pkg/logproto/extensions.go

@@ -182,3 +182,9 @@ func (r *QueryPatternsResponse) UnmarshalJSON(data []byte) error {
 func (d DetectedFieldType) String() string {
 	return string(d)
 }
+
+func (m *ShardsResponse) Merge(other *ShardsResponse) {
+	m.Shards = append(m.Shards, other.Shards...)
+	m.ChunkGroups = append(m.ChunkGroups, other.ChunkGroups...)
+	m.Statistics.Merge(other.Statistics)
+}