parquet: optimize query sharding by leveraging materialized series hooks (#12136)

Leverage this feature
prometheus-community/parquet-common#83 to
optimize query sharding. Previously we were materializing all chunks and
post-filtering them before sending through the wire. Now we can push
down this filter to the materializer and avoid loading them entirely.

Author: npazosmendez

go.mod

@@ -30,7 +30,7 @@ require (
 	github.com/opentracing-contrib/go-stdlib v1.1.0 // indirect
 	github.com/opentracing/opentracing-go v1.2.1-0.20220228012449-10b1cf09e00b // indirect
 	github.com/pkg/errors v0.9.1
-	github.com/prometheus-community/parquet-common v0.0.0-20250714085052-3b8805a2f9ad
+	github.com/prometheus-community/parquet-common v0.0.0-20250716185251-4cfa597e936c
 	github.com/prometheus/alertmanager v0.28.1
 	github.com/prometheus/client_golang v1.23.0-rc.1
 	github.com/prometheus/client_model v0.6.2

go.sum

@@ -895,8 +895,8 @@ github.com/posener/complete v1.2.3/go.mod h1:WZIdtGGp+qx0sLrYKtIRAruyNpv6hFCicSg
 github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c h1:ncq/mPwQF4JjgDlrVEn3C11VoGHZN7m8qihwgMEtzYw=
 github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c/go.mod h1:OmDBASR4679mdNQnz2pUhc2G8CO2JrUAVFDRBDP/hJE=
 github.com/pquerna/ffjson v0.0.0-20190930134022-aa0246cd15f7/go.mod h1:YARuvh7BUWHNhzDq2OM5tzR2RiCcN2D7sapiKyCel/M=
-github.com/prometheus-community/parquet-common v0.0.0-20250714085052-3b8805a2f9ad h1:w3igWmqBUdyOpiKoh8eO1j+Skd/kO+T+sC/OHP0uK6I=
-github.com/prometheus-community/parquet-common v0.0.0-20250714085052-3b8805a2f9ad/go.mod h1:MbAv/yCv9GORLj0XvXgRF913R9Jc04+BvVq4VJpPCi0=
+github.com/prometheus-community/parquet-common v0.0.0-20250716185251-4cfa597e936c h1:yDtT3c2klcWJj6A0osq72qM8rd1ohtl/J3rHD3FHuNw=
+github.com/prometheus-community/parquet-common v0.0.0-20250716185251-4cfa597e936c/go.mod h1:MbAv/yCv9GORLj0XvXgRF913R9Jc04+BvVq4VJpPCi0=
 github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
 github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
 github.com/prometheus/client_golang v1.4.0/go.mod h1:e9GMxYsXl05ICDXkRhurwBS4Q3OK1iX/F2sw+iXX5zU=

pkg/storage/parquet/chunk_querier.go

@@ -23,19 +23,21 @@ import (
 var tracer = otel.Tracer("pkg/storage/parquet")
 
 type querierOpts struct {
-	concurrency                int
-	rowCountLimitFunc          search.QuotaLimitFunc
-	chunkBytesLimitFunc        search.QuotaLimitFunc
-	dataBytesLimitFunc         search.QuotaLimitFunc
-	materializedSeriesCallback search.MaterializedSeriesFunc
+	concurrency                      int
+	rowCountLimitFunc                search.QuotaLimitFunc
+	chunkBytesLimitFunc              search.QuotaLimitFunc
+	dataBytesLimitFunc               search.QuotaLimitFunc
+	materializedSeriesCallback       search.MaterializedSeriesFunc
+	materializedLabelsFilterCallback search.MaterializedLabelsFilterCallback
 }
 
 var DefaultQuerierOpts = querierOpts{
-	concurrency:                runtime.GOMAXPROCS(0),
-	rowCountLimitFunc:          search.NoopQuotaLimitFunc,
-	chunkBytesLimitFunc:        search.NoopQuotaLimitFunc,
-	dataBytesLimitFunc:         search.NoopQuotaLimitFunc,
-	materializedSeriesCallback: search.NoopMaterializedSeriesFunc,
+	concurrency:                      runtime.GOMAXPROCS(0),
+	rowCountLimitFunc:                search.NoopQuotaLimitFunc,
+	chunkBytesLimitFunc:              search.NoopQuotaLimitFunc,
+	dataBytesLimitFunc:               search.NoopQuotaLimitFunc,
+	materializedSeriesCallback:       search.NoopMaterializedSeriesFunc,
+	materializedLabelsFilterCallback: search.NoopMaterializedLabelsFilterCallback,
 }
 
 type QuerierOpts func(*querierOpts)
@@ -76,6 +78,14 @@ func WithMaterializedSeriesCallback(fn search.MaterializedSeriesFunc) QuerierOpt
 	}
 }
 
+// WithMaterializedLabelsFilterCallback sets a callback function to create a filter that can be used
+// to filter series based on their labels before materializing chunks.
+func WithMaterializedLabelsFilterCallback(cb search.MaterializedLabelsFilterCallback) QuerierOpts {
+	return func(opts *querierOpts) {
+		opts.materializedLabelsFilterCallback = cb
+	}
+}
+
 func NewParquetChunkQuerier(d *schema.PrometheusParquetChunksDecoder, shardFinder queryable.ShardsFinderFunction, opts ...QuerierOpts) (prom_storage.ChunkQuerier, error) {
 	cfg := DefaultQuerierOpts
 

pkg/storage/parquet/queryable_shard.go

@@ -28,7 +28,7 @@ func newQueryableShard(opts *querierOpts, block storage.ParquetShard, d *schema.
 	if err != nil {
 		return nil, err
 	}
-	m, err := search.NewMaterializer(s, d, block, opts.concurrency, rowCountQuota, chunkBytesQuota, dataBytesQuota, opts.materializedSeriesCallback)
+	m, err := search.NewMaterializer(s, d, block, opts.concurrency, rowCountQuota, chunkBytesQuota, dataBytesQuota, opts.materializedSeriesCallback, opts.materializedLabelsFilterCallback)
 	if err != nil {
 		return nil, err
 	}
@@ -66,7 +66,7 @@ func (b queryableShard) Query(ctx context.Context, sorted bool, mint, maxt int64
 				return nil
 			}
 
-			series, err := b.m.Materialize(ctx, rgi, mint, maxt, skipChunks, rr)
+			series, err := b.m.Materialize(ctx, nil, rgi, mint, maxt, skipChunks, rr)
 			if err != nil {
 				return err
 			}

pkg/storegateway/parquet_bucket_store.go

@@ -26,6 +26,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/prometheus-community/parquet-common/queryable"
 	"github.com/prometheus-community/parquet-common/schema"
+	"github.com/prometheus-community/parquet-common/search"
 	parquetStorage "github.com/prometheus-community/parquet-common/storage"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/storage"
@@ -595,7 +596,15 @@ func (s *ParquetBucketStore) createLabelsAndChunksIterators(
 	}
 
 	decoder := schema.NewPrometheusParquetChunksDecoder(chunkenc.NewPool())
-	q, err := parquet.NewParquetChunkQuerier(decoder, shardsFinder, s.querierOpts...)
+	opts := s.querierOpts
+	if shardSelector != nil {
+		shardFilter := func(_ context.Context, _ *storage.SelectHints) (search.MaterializedLabelsFilter, bool) {
+			return materializedLabelsShardFilter{shardSelector: shardSelector}, true
+		}
+		opts = append(opts, parquet.WithMaterializedLabelsFilterCallback(shardFilter))
+	}
+
+	q, err := parquet.NewParquetChunkQuerier(decoder, shardsFinder, opts...)
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "error creating parquet queryable")
 	}
@@ -618,7 +627,7 @@ func (s *ParquetBucketStore) createLabelsAndChunksIterators(
 	// not support streaming results, the iterator we get from q.Select is
 	// already backed by a slice. So we are not losing as much as it may seem.
 	// We are planning to implement proper streaming.
-	lbls, aggrChunks, err := toLabelsAndAggChunksSlice(chunkSeriesSet, shardSelector, req.SkipChunks)
+	lbls, aggrChunks, err := toLabelsAndAggChunksSlice(chunkSeriesSet, req.SkipChunks)
 	if err != nil {
 		return nil, nil, errors.Wrap(err, "error converting parquet series set to labels and chunks slice")
 	}
@@ -642,6 +651,31 @@ func (s *ParquetBucketStore) createLabelsAndChunksIterators(
 	return labelsIt, newConcreteIterator(aggrChunks), nil
 }
 
+type materializedLabelsShardFilter struct {
+	shardSelector *sharding.ShardSelector
+}
+
+func (f materializedLabelsShardFilter) Filter(lbls labels.Labels) bool {
+	return shardOwnedUncached(f.shardSelector, lbls)
+}
+
+func (f materializedLabelsShardFilter) Close() {
+}
+
+// shardOwnedUncached checks if the given labels belong to the shard specified
+// by the shard selector. As opposed to shardOwned & friends from the
+// non-Parquet path, this function does not cache hashes. This is because, at
+// least yet, we don't have easy access to an identifier for the series in the
+// block to use as a cache key.
+func shardOwnedUncached(shard *sharding.ShardSelector, lset labels.Labels) bool {
+	if shard == nil {
+		return true
+	}
+
+	hash := labels.StableHash(lset)
+	return hash%shard.ShardCount == shard.ShardIndex
+}
+
 func (s *ParquetBucketStore) sendStreamingSeriesLabelsAndStats(req *storepb.SeriesRequest, srv storegatewaypb.StoreGateway_SeriesServer, stats *safeQueryStats, labelsIt iterator[labels.Labels]) (numSeries int, err error) {
 	var (
 		encodeDuration = time.Duration(0)
@@ -1131,17 +1165,13 @@ func (s *ParquetBucketStore) cleanUpUnownedBlocks() error {
 // storage.ChunkSeriesSet and returns them as slices, converting the chunks to
 // storepb.AggrChunk format. If skipChunks is true, the chunks slice will be
 // empty.
-func toLabelsAndAggChunksSlice(chunkSeriesSet storage.ChunkSeriesSet, shardSelector *sharding.ShardSelector, skipChunks bool) ([]labels.Labels, [][]storepb.AggrChunk, error) {
+func toLabelsAndAggChunksSlice(chunkSeriesSet storage.ChunkSeriesSet, skipChunks bool) ([]labels.Labels, [][]storepb.AggrChunk, error) {
 	var seriesLabels []labels.Labels
 	var aggrChunks [][]storepb.AggrChunk
 
 	for chunkSeriesSet.Next() {
 		chunkSeries := chunkSeriesSet.At()
 		lbls := chunkSeries.Labels()
-		if !shardOwnedUncached(shardSelector, lbls) {
-			continue
-		}
-
 		seriesLabels = append(seriesLabels, lbls)
 
 		if skipChunks {
@@ -1168,20 +1198,6 @@ func toLabelsAndAggChunksSlice(chunkSeriesSet storage.ChunkSeriesSet, shardSelec
 	return seriesLabels, aggrChunks, chunkSeriesSet.Err()
 }
 
-// shardOwnedUncached checks if the given labels belong to the shard specified
-// by the shard selector. As opposed to shardOwned & friends from the
-// non-Parquet path, this function does not cache hashes. This is because, at
-// least yet, we don't have easy access to an identifier for the series in the
-// block to use as a cache key.
-func shardOwnedUncached(shard *sharding.ShardSelector, lset labels.Labels) bool {
-	if shard == nil {
-		return true
-	}
-
-	hash := labels.StableHash(lset)
-	return hash%shard.ShardCount == shard.ShardIndex
-}
-
 func prometheusChunkEncodingToStorePBChunkType(enc chunkenc.Encoding) storepb.Chunk_Encoding {
 	switch enc {
 	case chunkenc.EncXOR: