Resolve Parquet shard count via bucket index to optimize storage calls

CHANGELOG.md

@@ -39,6 +39,7 @@
 * [ENHANCEMENT] Distributor: Added `cortex_distributor_received_histogram_buckets` metric to track number of buckets in received native histogram samples before validation, per user. #7569
 * [ENHANCEMENT] Distributor: Add `WrappedHistogram` with configurable size limit (`-validation.max-native-histogram-size-bytes`) to cap native histogram protobuf size before unmarshalling. #7570
 * [ENHANCEMENT] Ingester: Add lazy regex evaluation on head postings cache miss. Defers expensive regex matchers on high-cardinality labels to per-series filtering when a selective equality matcher already narrows the result set. Configured via `-blocks-storage.expanded_postings_cache.head.lazy-matcher-max-cardinality` (disabled by default). #7553
+* [ENHANCEMENT] Store Gateway: Resolve the parquet shard count from the bucket index instead of reading the converter mark for each block, reducing object storage calls when the bucket index is enabled. A `component` label is added to the bucket index loader metrics to distinguish store-queryable and store-gateway. #7648
 * [ENHANCEMENT] Query Frontend: Improve the slow query log with `source`, `user_agent`, `engine_type`, `block_store_type`, and query stats fields to aid slow query diagnosis. #7601
 * [ENHANCEMENT] Ring: Add ring metric to count number of duplicate tokens. #7626
 * [ENHANCEMENT] Upgrade prometheus version to v3.9.1. #7535

docs/blocks-storage/querier.md

@@ -2006,13 +2006,14 @@ blocks_storage:
       [enabled: <boolean> | default = true]
 
       # How frequently a bucket index, which previously failed to load, should
-      # be tried to load again. This option is used only by querier.
+      # be tried to load again. This option is used by querier and store-gateway
+      # parquet mode.
       # CLI flag: -blocks-storage.bucket-store.bucket-index.update-on-error-interval
       [update_on_error_interval: <duration> | default = 1m]
 
       # How long a unused bucket index should be cached. Once this timeout
       # expires, the unused bucket index is removed from the in-memory cache.
-      # This option is used only by querier.
+      # This option is used by querier and store-gateway parquet mode.
       # CLI flag: -blocks-storage.bucket-store.bucket-index.idle-timeout
       [idle_timeout: <duration> | default = 1h]
 

docs/blocks-storage/store-gateway.md

@@ -2064,13 +2064,14 @@ blocks_storage:
       [enabled: <boolean> | default = true]
 
       # How frequently a bucket index, which previously failed to load, should
-      # be tried to load again. This option is used only by querier.
+      # be tried to load again. This option is used by querier and store-gateway
+      # parquet mode.
       # CLI flag: -blocks-storage.bucket-store.bucket-index.update-on-error-interval
       [update_on_error_interval: <duration> | default = 1m]
 
       # How long a unused bucket index should be cached. Once this timeout
       # expires, the unused bucket index is removed from the in-memory cache.
-      # This option is used only by querier.
+      # This option is used by querier and store-gateway parquet mode.
       # CLI flag: -blocks-storage.bucket-store.bucket-index.idle-timeout
       [idle_timeout: <duration> | default = 1h]
 

docs/configuration/config-file-reference.md

@@ -2656,13 +2656,14 @@ bucket_store:
     [enabled: <boolean> | default = true]
 
     # How frequently a bucket index, which previously failed to load, should be
-    # tried to load again. This option is used only by querier.
+    # tried to load again. This option is used by querier and store-gateway
+    # parquet mode.
     # CLI flag: -blocks-storage.bucket-store.bucket-index.update-on-error-interval
     [update_on_error_interval: <duration> | default = 1m]
 
     # How long a unused bucket index should be cached. Once this timeout
     # expires, the unused bucket index is removed from the in-memory cache. This
-    # option is used only by querier.
+    # option is used by querier and store-gateway parquet mode.
     # CLI flag: -blocks-storage.bucket-store.bucket-index.idle-timeout
     [idle_timeout: <duration> | default = 1h]
 

integration/parquet_querier_test.go

@@ -619,9 +619,20 @@ func TestParquetMultiShardQuery(t *testing.T) {
 			if tc.viaStoreGateway {
 				require.NoError(t, cortex.WaitSumMetricsWithOptions(e2e.Greater(0), []string{"cortex_querier_storegateway_instances_hit_per_query"}, e2e.WithMetricCount, e2e.SkipMissingMetrics))
 				require.NoError(t, cortex.WaitSumMetricsWithOptions(e2e.Greater(0), []string{"cortex_querier_blocks_consistency_checks_total"}, e2e.SkipMissingMetrics))
+				// The store-gateway parquet bucket store resolves the shard count via the
+				// bucket index loader, which is tagged with component="store-gateway".
+				if tc.extraStoreGateways == 0 {
+					// Only assert on the single binary when there are no extra replicas, since
+					// with replication the query may be served by another store-gateway.
+					require.NoError(t, cortex.WaitSumMetricsWithOptions(e2e.Greater(0), []string{"cortex_bucket_index_loads_total"}, e2e.WithLabelMatchers(
+						labels.MustNewMatcher(labels.MatchEqual, "component", "store-gateway"))))
+				}
 			} else {
 				require.NoError(t, cortex.WaitSumMetricsWithOptions(e2e.Greater(0), []string{"cortex_parquet_queryable_blocks_queried_total"}, e2e.WithLabelMatchers(
 					labels.MustNewMatcher(labels.MatchEqual, "type", "parquet"))))
+				// The querier's bucket index loader is tagged with component="store-queryable".
+				require.NoError(t, cortex.WaitSumMetricsWithOptions(e2e.Greater(0), []string{"cortex_bucket_index_loads_total"}, e2e.WithLabelMatchers(
+					labels.MustNewMatcher(labels.MatchEqual, "component", "store-queryable"))))
 			}
 		})
 	}

pkg/querier/blocks_store_queryable.go

@@ -210,7 +210,7 @@ func NewBlocksStoreQueryableFromConfig(querierCfg Config, gatewayCfg storegatewa
 			MaxStalePeriod:           storageCfg.BucketStore.BucketIndex.MaxStalePeriod,
 			IgnoreDeletionMarksDelay: storageCfg.BucketStore.IgnoreDeletionMarksDelay,
 			IgnoreBlocksWithin:       storageCfg.BucketStore.IgnoreBlocksWithin,
-		}, bucketClient, limits, logger, reg)
+		}, bucketClient, limits, logger, extprom.WrapRegistererWith(prometheus.Labels{"component": "store-queryable"}, reg))
 	} else {
 		usersScanner, err := users.NewScanner(storageCfg.UsersScanner, bucketClient, logger, extprom.WrapRegistererWith(prometheus.Labels{"component": "querier"}, reg))
 		if err != nil {

pkg/storage/tsdb/config.go

@@ -457,8 +457,8 @@ type BucketIndexConfig struct {
 
 func (cfg *BucketIndexConfig) RegisterFlagsWithPrefix(f *flag.FlagSet, prefix string) {
 	f.BoolVar(&cfg.Enabled, prefix+"enabled", true, "True to enable querier and store-gateway to discover blocks in the storage via bucket index instead of bucket scanning. Disabling the bucket index is not recommended for production.")
-	f.DurationVar(&cfg.UpdateOnErrorInterval, prefix+"update-on-error-interval", time.Minute, "How frequently a bucket index, which previously failed to load, should be tried to load again. This option is used only by querier.")
-	f.DurationVar(&cfg.IdleTimeout, prefix+"idle-timeout", time.Hour, "How long a unused bucket index should be cached. Once this timeout expires, the unused bucket index is removed from the in-memory cache. This option is used only by querier.")
+	f.DurationVar(&cfg.UpdateOnErrorInterval, prefix+"update-on-error-interval", time.Minute, "How frequently a bucket index, which previously failed to load, should be tried to load again. This option is used by querier and store-gateway parquet mode.")
+	f.DurationVar(&cfg.IdleTimeout, prefix+"idle-timeout", time.Hour, "How long a unused bucket index should be cached. Once this timeout expires, the unused bucket index is removed from the in-memory cache. This option is used by querier and store-gateway parquet mode.")
 	f.DurationVar(&cfg.MaxStalePeriod, prefix+"max-stale-period", time.Hour, "The maximum allowed age of a bucket index (last updated) before queries start failing because the bucket index is too old. The bucket index is periodically updated by the compactor, while this check is enforced in the querier (at query time).")
 }
 

pkg/storegateway/bucket_stores.go

@@ -49,6 +49,7 @@ type BucketStores interface {
 	storepb.StoreServer
 	SyncBlocks(ctx context.Context) error
 	InitialSync(ctx context.Context) error
+	Stop() error
 }
 
 // ThanosBucketStores is a multi-tenant wrapper of Thanos BucketStore.
@@ -230,6 +231,9 @@ func (u *ThanosBucketStores) SyncBlocks(ctx context.Context) error {
 	})
 }
 
+// Stop implements BucketStores. ThanosBucketStores has no background services to stop.
+func (u *ThanosBucketStores) Stop() error { return nil }
+
 func (u *ThanosBucketStores) syncUsersBlocksWithRetries(ctx context.Context, f func(context.Context, *store.BucketStore) error) error {
 	retries := backoff.New(ctx, backoff.Config{
 		MinBackoff: 1 * time.Second,

pkg/storegateway/gateway.go

@@ -391,6 +391,11 @@ func (g *StoreGateway) running(ctx context.Context) error {
 }
 
 func (g *StoreGateway) stopping(_ error) error {
+	if g.stores != nil {
+		if err := g.stores.Stop(); err != nil {
+			level.Warn(g.logger).Log("msg", "failed to stop bucket stores", "err", err)
+		}
+	}
 	if g.subservices != nil {
 		return services.StopManagerAndAwaitStopped(context.Background(), g.subservices)
 	}

pkg/storegateway/parquet_bucket_store.go

@@ -4,8 +4,10 @@ import (
 	"context"
 	"fmt"
 	"strings"
+	"sync"
 
 	"github.com/go-kit/log"
+	"github.com/go-kit/log/level"
 	"github.com/gogo/protobuf/types"
 	"github.com/oklog/ulid/v2"
 	"github.com/pkg/errors"
@@ -27,22 +29,32 @@ import (
 	"google.golang.org/grpc/status"
 
 	cortex_parquet "github.com/cortexproject/cortex/pkg/storage/parquet"
+	"github.com/cortexproject/cortex/pkg/storage/tsdb/bucketindex"
 	"github.com/cortexproject/cortex/pkg/util/parquetutil"
 	"github.com/cortexproject/cortex/pkg/util/spanlogger"
 	"github.com/cortexproject/cortex/pkg/util/validation"
 )
 
 type parquetBucketStore struct {
-	logger      log.Logger
-	bucket      objstore.InstrumentedBucket
-	limits      *validation.Overrides
-	concurrency int
+	logger             log.Logger
+	bucket             objstore.InstrumentedBucket
+	indexLoader        *bucketindex.Loader // nil when bucket index disabled
+	limits             *validation.Overrides
+	userID             string
+	bucketIndexEnabled bool
+	concurrency        int
 
 	chunksDecoder *schema.PrometheusParquetChunksDecoder
 
 	matcherCache      storecache.MatchersCache
 	parquetShardCache parquetutil.CacheInterface[parquet_storage.ParquetShard]
 	rowRangesCache    search.RowRangesForConstraintsCache
+
+	shardCountsMu sync.Mutex
+	// cachedShardCounts maps a block ID to its parquet shard count. The map is rebuilt
+	// only when the bucket index is refreshed (detected via cachedIndexUpdatedAt).
+	cachedShardCounts    map[string]int
+	cachedIndexUpdatedAt int64
 }
 
 func (p *parquetBucketStore) Close() error {
@@ -67,20 +79,25 @@ func (p *parquetBucketStore) findParquetBlocks(ctx context.Context, blockMatcher
 	bucketOpener := parquet_storage.NewParquetBucketOpener(p.bucket)
 	noopQuota := search.NewQuota(search.NoopQuotaLimitFunc(ctx))
 
-	// Read converter marks and expand to per-shard (blockID, shardID) lists.
-	// TODO(Sungjin1212): Read the shard count from the bucket index instead of reading the converter mark for each block.
+	shardCounts, err := p.resolveShardCounts(ctx, blockIDs)
+	if err != nil {
+		return nil, err
+	}
+
 	var shardBlockIDs []string
 	var shardIDs []int
 	for _, blockID := range blockIDs {
-		uid, err := ulid.Parse(blockID)
-		if err != nil {
-			return nil, errors.Wrapf(err, "failed to parse block ID %s", blockID)
-		}
-		marker, err := cortex_parquet.ReadConverterMark(ctx, uid, p.bucket, p.logger)
-		if err != nil {
-			return nil, errors.Wrapf(err, "failed to read converter mark for block %s", blockID)
+		numShards, ok := shardCounts[blockID]
+		if !ok {
+			// The bucket index may not carry this block's shard count yet, for example
+			// right after the block has been converted but before the index has been
+			// refreshed. In that case, read the converter mark directly to get the
+			// actual shard count instead of assuming a single shard.
+			numShards, err = p.shardCountFromConverterMark(ctx, blockID)
+			if err != nil {
+				return nil, err
+			}
 		}
-		numShards := marker.Shards
 		if numShards <= 0 {
 			// backward compatibility: blocks without a shard count have one shard
 			numShards = 1
@@ -112,6 +129,78 @@ func (p *parquetBucketStore) findParquetBlocks(ctx context.Context, blockMatcher
 	return parquetBlocks, nil
 }
 
+// resolveShardCounts returns the number of parquet shards for each requested block ID.
+//
+// When the bucket index is enabled, the shard count is read from the bucket index.
+// When the bucket index is disabled, it falls back to reading the converter mark
+// for each block.
+func (p *parquetBucketStore) resolveShardCounts(ctx context.Context, blockIDs []string) (map[string]int, error) {
+	shardCounts := make(map[string]int, len(blockIDs))
+
+	if p.bucketIndexEnabled && p.indexLoader != nil {
+		idx, _, err := p.indexLoader.GetIndex(ctx, p.userID)
+		if err == nil {
+			return p.shardCountsFromIndex(idx), nil
+		}
+		level.Warn(p.logger).Log("msg", "failed to get bucket index, falling back to converter marks for parquet shard count", "err", err)
+	}
+
+	// Fallback: read the converter mark for each block.
+	for _, blockID := range blockIDs {
+		numShards, err := p.shardCountFromConverterMark(ctx, blockID)
+		if err != nil {
+			return nil, err
+		}
+		shardCounts[blockID] = numShards
+	}
+	return shardCounts, nil
+}
+
+// shardCountFromConverterMark reads the parquet converter mark for the given block and
+// returns the number of shards recorded in it.
+func (p *parquetBucketStore) shardCountFromConverterMark(ctx context.Context, blockID string) (int, error) {
+	uid, err := ulid.Parse(blockID)
+	if err != nil {
+		return 0, errors.Wrapf(err, "failed to parse block ID %s", blockID)
+	}
+	marker, err := cortex_parquet.ReadConverterMark(ctx, uid, p.bucket, p.logger)
+	if err != nil {
+		return 0, errors.Wrapf(err, "failed to read converter mark for block %s", blockID)
+	}
+	return marker.Shards, nil
+}
+
+// shardCountsFromIndex returns a block ID to shard count map built from the bucket index.
+//
+// The parquet shard count of a block is immutable once the block is converted, so the map
+// is rebuilt only when the bucket index has been refreshed (detected via UpdatedAt). In
+// between refreshes the previously built map is reused, avoiding a full iteration over all
+// blocks in the index on every request.
+func (p *parquetBucketStore) shardCountsFromIndex(idx *bucketindex.Index) map[string]int {
+	p.shardCountsMu.Lock()
+	defer p.shardCountsMu.Unlock()
+
+	if p.cachedShardCounts != nil && p.cachedIndexUpdatedAt == idx.UpdatedAt {
+		return p.cachedShardCounts
+	}
+
+	shardCounts := make(map[string]int, len(idx.Blocks))
+	for _, b := range idx.Blocks {
+		if b.Parquet == nil {
+			continue
+		}
+		numShards := b.Parquet.Shards
+		if numShards <= 0 {
+			numShards = 1
+		}
+		shardCounts[b.ID.String()] = numShards
+	}
+
+	p.cachedShardCounts = shardCounts
+	p.cachedIndexUpdatedAt = idx.UpdatedAt
+	return shardCounts
+}
+
 // Series implements the store interface for a single parquet bucket store
 func (p *parquetBucketStore) Series(req *storepb.SeriesRequest, seriesSrv storepb.Store_SeriesServer) (err error) {
 	spanLog, ctx := spanlogger.New(seriesSrv.Context(), "ParquetBucketStore.Series")

pkg/storegateway/parquet_bucket_stores.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"sort"
 	"sync"
+	"time"
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
@@ -31,9 +32,11 @@ import (
 	"github.com/cortexproject/cortex/pkg/querysharding"
 	"github.com/cortexproject/cortex/pkg/storage/bucket"
 	"github.com/cortexproject/cortex/pkg/storage/tsdb"
+	"github.com/cortexproject/cortex/pkg/storage/tsdb/bucketindex"
 	cortex_util "github.com/cortexproject/cortex/pkg/util"
 	cortex_errors "github.com/cortexproject/cortex/pkg/util/errors"
 	"github.com/cortexproject/cortex/pkg/util/parquetutil"
+	"github.com/cortexproject/cortex/pkg/util/services"
 	"github.com/cortexproject/cortex/pkg/util/spanlogger"
 	"github.com/cortexproject/cortex/pkg/util/users"
 	"github.com/cortexproject/cortex/pkg/util/validation"
@@ -60,6 +63,10 @@ type ParquetBucketStores struct {
 	rowRangesCache    search.RowRangesForConstraintsCache
 
 	inflightRequests *cortex_util.InflightRequestTracker
+
+	// indexLoader lazily loads and caches the per-tenant bucket index in memory
+	// It is non-nil only when BucketIndex.Enabled.
+	indexLoader *bucketindex.Loader
 }
 
 // newParquetBucketStores creates a new multi-tenant parquet bucket stores
@@ -108,6 +115,15 @@ func newParquetBucketStores(cfg tsdb.BlocksStorageConfig, bucketClient objstore.
 		u.matcherCache = storecache.NoopMatchersCache
 	}
 
+	if cfg.BucketStore.BucketIndex.Enabled {
+		u.indexLoader = bucketindex.NewLoader(bucketindex.LoaderConfig{
+			CheckInterval:         time.Minute,
+			UpdateOnStaleInterval: cfg.BucketStore.SyncInterval,
+			UpdateOnErrorInterval: cfg.BucketStore.BucketIndex.UpdateOnErrorInterval,
+			IdleTimeout:           cfg.BucketStore.BucketIndex.IdleTimeout,
+		}, bucketClient, limits, logger, reg)
+	}
+
 	return u, nil
 }
 
@@ -218,6 +234,18 @@ func (u *ParquetBucketStores) SyncBlocks(ctx context.Context) error {
 
 // InitialSync implements BucketStores
 func (u *ParquetBucketStores) InitialSync(ctx context.Context) error {
+	if u.indexLoader != nil {
+		// Start indexLoader
+		return services.StartAndAwaitRunning(ctx, u.indexLoader)
+	}
+	return nil
+}
+
+// Stop implements BucketStores
+func (u *ParquetBucketStores) Stop() error {
+	if u.indexLoader != nil {
+		return services.StopAndAwaitTerminated(context.Background(), u.indexLoader)
+	}
 	return nil
 }
 
@@ -270,14 +298,17 @@ func (u *ParquetBucketStores) createParquetBucketStore(userID string, userLogger
 	userBucket := bucket.NewUserBucketClient(userID, u.bucket, u.limits)
 
 	store := &parquetBucketStore{
-		logger:            userLogger,
-		bucket:            userBucket,
-		limits:            u.limits,
-		concurrency:       u.cfg.BucketStore.ParquetQueryConcurrency,
-		chunksDecoder:     u.chunksDecoder,
-		matcherCache:      u.matcherCache,
-		parquetShardCache: u.parquetShardCache,
-		rowRangesCache:    u.rowRangesCache,
+		logger:             userLogger,
+		bucket:             userBucket,
+		indexLoader:        u.indexLoader,
+		limits:             u.limits,
+		userID:             userID,
+		bucketIndexEnabled: u.cfg.BucketStore.BucketIndex.Enabled,
+		concurrency:        u.cfg.BucketStore.ParquetQueryConcurrency,
+		chunksDecoder:      u.chunksDecoder,
+		matcherCache:       u.matcherCache,
+		parquetShardCache:  u.parquetShardCache,
+		rowRangesCache:     u.rowRangesCache,
 	}
 
 	return store, nil