Fixing Query Store After on parquet -> we need to change the hints to we filter out the right chunks

Signed-off-by: alanprot <[email protected]>

Author: alanprot

pkg/querier/parquet_queryable.go

@@ -356,7 +356,7 @@ func (q *parquetQuerierWithFallback) LabelNames(ctx context.Context, hints *stor
 	return result, rAnnotations, nil
 }
 
-func (q *parquetQuerierWithFallback) Select(ctx context.Context, sortSeries bool, hints *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
+func (q *parquetQuerierWithFallback) Select(ctx context.Context, sortSeries bool, h *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
 	userID, err := tenant.TenantID(ctx)
 	if err != nil {
 		storage.ErrSeriesSet(err)
@@ -366,68 +366,99 @@ func (q *parquetQuerierWithFallback) Select(ctx context.Context, sortSeries bool
 		uLogger := util_log.WithUserID(userID, q.logger)
 		level.Warn(uLogger).Log("msg", "parquet queryable enabled but vertical sharding > 1. Falling back to the block storage")
 
-		return q.blocksStoreQuerier.Select(ctx, sortSeries, hints, matchers...)
+		return q.blocksStoreQuerier.Select(ctx, sortSeries, h, matchers...)
 	}
 
-	mint, maxt, limit := q.minT, q.maxT, 0
+	hints := storage.SelectHints{
+		Start: q.minT,
+		End:   q.maxT,
+	}
 
-	if hints != nil {
+	mint, maxt, limit := q.minT, q.maxT, 0
+	if h != nil {
+		// let copy the hints here as we wanna potentially modify it
+		hints = *h
 		mint, maxt, limit = hints.Start, hints.End, hints.Limit
 	}
 
+	maxt = q.adjustMaxT(maxt)
+	hints.End = maxt
+
+	if maxt < mint {
+		return nil
+	}
+
 	remaining, parquet, err := q.getBlocks(ctx, mint, maxt)
 	if err != nil {
 		return storage.ErrSeriesSet(err)
 	}
 
-	serieSets := []storage.SeriesSet{}
-
 	// Lets sort the series to merge
 	if len(parquet) > 0 && len(remaining) > 0 {
 		sortSeries = true
 	}
 
+	promises := make([]chan storage.SeriesSet, 0, 2)
+
 	if len(parquet) > 0 {
-		serieSets = append(serieSets, q.parquetQuerier.Select(InjectBlocksIntoContext(ctx, parquet...), sortSeries, hints, matchers...))
+		p := make(chan storage.SeriesSet, 1)
+		promises = append(promises, p)
+		go func() {
+			p <- q.parquetQuerier.Select(InjectBlocksIntoContext(ctx, parquet...), sortSeries, &hints, matchers...)
+		}()
 	}
 
 	if len(remaining) > 0 {
-		serieSets = append(serieSets, q.blocksStoreQuerier.Select(InjectBlocksIntoContext(ctx, remaining...), sortSeries, hints, matchers...))
+		p := make(chan storage.SeriesSet, 1)
+		promises = append(promises, p)
+		go func() {
+			p <- q.blocksStoreQuerier.Select(InjectBlocksIntoContext(ctx, remaining...), sortSeries, &hints, matchers...)
+		}()
 	}
 
-	if len(serieSets) == 1 {
-		return serieSets[0]
+	if len(promises) == 1 {
+		return <-promises[0]
 	}
 
-	return storage.NewMergeSeriesSet(serieSets, limit, storage.ChainedSeriesMerge)
-}
+	seriesSets := make([]storage.SeriesSet, len(promises))
+	for i, promise := range promises {
+		seriesSets[i] = <-promise
+	}
 
-func (q *parquetQuerierWithFallback) Close() error {
-	mErr := multierror.MultiError{}
-	mErr.Add(q.parquetQuerier.Close())
-	mErr.Add(q.blocksStoreQuerier.Close())
-	return mErr.Err()
+	return storage.NewMergeSeriesSet(seriesSets, limit, storage.ChainedSeriesMerge)
 }
 
-func (q *parquetQuerierWithFallback) getBlocks(ctx context.Context, minT, maxT int64) ([]*bucketindex.Block, []*bucketindex.Block, error) {
+func (q *parquetQuerierWithFallback) adjustMaxT(maxt int64) int64 {
 	// If queryStoreAfter is enabled, we do manipulate the query maxt to query samples up until
 	// now - queryStoreAfter, because the most recent time range is covered by ingesters. This
 	// optimization is particularly important for the blocks storage because can be used to skip
 	// querying most recent not-compacted-yet blocks from the storage.
 	if q.queryStoreAfter > 0 {
 		now := time.Now()
-		maxT = min(maxT, util.TimeToMillis(now.Add(-q.queryStoreAfter)))
-
-		if maxT < minT {
-			return nil, nil, nil
-		}
+		maxt = min(maxt, util.TimeToMillis(now.Add(-q.queryStoreAfter)))
 	}
+	return maxt
+}
 
+func (q *parquetQuerierWithFallback) Close() error {
+	mErr := multierror.MultiError{}
+	mErr.Add(q.parquetQuerier.Close())
+	mErr.Add(q.blocksStoreQuerier.Close())
+	return mErr.Err()
+}
+
+func (q *parquetQuerierWithFallback) getBlocks(ctx context.Context, minT, maxT int64) ([]*bucketindex.Block, []*bucketindex.Block, error) {
 	userID, err := tenant.TenantID(ctx)
 	if err != nil {
 		return nil, nil, err
 	}
 
+	maxT = q.adjustMaxT(maxT)
+
+	if maxT < minT {
+		return nil, nil, nil
+	}
+	
 	blocks, _, err := q.finder.GetBlocks(ctx, userID, minT, maxT)
 	if err != nil {
 		return nil, nil, err

pkg/querier/parquet_queryable_test.go

@@ -3,6 +3,7 @@ package querier
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/oklog/ulid"
@@ -19,6 +20,7 @@ import (
 	"github.com/cortexproject/cortex/pkg/querier/series"
 	"github.com/cortexproject/cortex/pkg/storage/parquet"
 	"github.com/cortexproject/cortex/pkg/storage/tsdb/bucketindex"
+	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/cortexproject/cortex/pkg/util/flagext"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
@@ -27,7 +29,7 @@ func TestParquetQueryableFallbackLogic(t *testing.T) {
 	block1 := ulid.MustNew(1, nil)
 	block2 := ulid.MustNew(2, nil)
 	minT := int64(10)
-	maxT := int64(20)
+	maxT := util.TimeToMillis(time.Now())
 
 	createStore := func() *blocksStoreSetMock {
 		return &blocksStoreSetMock{mockedResponses: []interface{}{
@@ -124,13 +126,14 @@ func TestParquetQueryableFallbackLogic(t *testing.T) {
 			finder:                finder,
 			blocksStoreQuerier:    q,
 			parquetQuerier:        mParquetQuerier,
+			queryStoreAfter:       time.Hour,
 			metrics:               newParquetQueryableFallbackMetrics(prometheus.NewRegistry()),
 			limits:                defaultOverrides(t, 0),
 			logger:                log.NewNopLogger(),
 			defaultBlockStoreType: parquetBlockStore,
 		}
 
-		finder.On("GetBlocks", mock.Anything, "user-1", minT, maxT).Return(bucketindex.Blocks{
+		finder.On("GetBlocks", mock.Anything, "user-1", minT, mock.Anything).Return(bucketindex.Blocks{
 			&bucketindex.Block{ID: block1},
 			&bucketindex.Block{ID: block2},
 		}, map[ulid.ULID]*bucketindex.BlockDeletionMark(nil), nil)
@@ -246,30 +249,39 @@ func TestParquetQueryableFallbackLogic(t *testing.T) {
 		}
 
 		mParquetQuerier := &mockParquetQuerier{}
+		queryStoreAfter := time.Hour
 		pq := &parquetQuerierWithFallback{
 			minT:                  minT,
 			maxT:                  maxT,
 			finder:                finder,
 			blocksStoreQuerier:    q,
 			parquetQuerier:        mParquetQuerier,
+			queryStoreAfter:       queryStoreAfter,
 			metrics:               newParquetQueryableFallbackMetrics(prometheus.NewRegistry()),
 			limits:                defaultOverrides(t, 0),
 			logger:                log.NewNopLogger(),
 			defaultBlockStoreType: parquetBlockStore,
 		}
 
-		finder.On("GetBlocks", mock.Anything, "user-1", minT, maxT).Return(bucketindex.Blocks{
+		finder.On("GetBlocks", mock.Anything, "user-1", minT, mock.Anything).Return(bucketindex.Blocks{
 			&bucketindex.Block{ID: block1, Parquet: &parquet.ConverterMarkMeta{Version: 1}},
 			&bucketindex.Block{ID: block2, Parquet: &parquet.ConverterMarkMeta{Version: 1}},
 		}, map[ulid.ULID]*bucketindex.BlockDeletionMark(nil), nil)
 
 		t.Run("select", func(t *testing.T) {
 			stores.Reset()
 			mParquetQuerier.Reset()
-			ss := pq.Select(ctx, true, nil, matchers...)
+			hints := storage.SelectHints{
+				Start: minT,
+				End:   maxT,
+			}
+			ss := pq.Select(ctx, true, &hints, matchers...)
 			require.NoError(t, ss.Err())
 			require.Len(t, stores.queriedBlocks, 0)
 			require.Len(t, mParquetQuerier.queriedBlocks, 2)
+			require.Equal(t, mParquetQuerier.queriedHints.Start, minT)
+			queriedDelta := time.Duration(maxT-mParquetQuerier.queriedHints.End) * time.Millisecond
+			require.InDeltaf(t, queriedDelta.Minutes(), queryStoreAfter.Minutes(), 0.1, "query after not set")
 		})
 
 		t.Run("labelNames", func(t *testing.T) {
@@ -409,13 +421,14 @@ func defaultOverrides(t *testing.T, queryVerticalShardSize int) *validation.Over
 
 type mockParquetQuerier struct {
 	queriedBlocks []*bucketindex.Block
+	queriedHints  *storage.SelectHints
 }
 
 func (m *mockParquetQuerier) Select(ctx context.Context, sortSeries bool, sp *storage.SelectHints, matchers ...*labels.Matcher) storage.SeriesSet {
 	if blocks, ok := ExtractBlocksFromContext(ctx); ok {
 		m.queriedBlocks = append(m.queriedBlocks, blocks...)
 	}
-
+	m.queriedHints = sp
 	return series.NewConcreteSeriesSet(sortSeries, nil)
 }