Regex set optimisation when looking series in ingester

pkg/chunk/chunk_store.go

@@ -441,8 +441,8 @@ func (c *baseStore) lookupIdsByMetricNameMatcher(ctx context.Context, from, thro
 	} else if matcher.Type == labels.MatchEqual {
 		labelName = matcher.Name
 		queries, err = c.schema.GetReadQueriesForMetricLabelValue(from, through, userID, metricName, matcher.Name, matcher.Value)
-	} else if matcher.Type == labels.MatchRegexp && len(findSetMatches(matcher.Value)) > 0 {
-		set := findSetMatches(matcher.Value)
+	} else if matcher.Type == labels.MatchRegexp && len(FindSetMatches(matcher.Value)) > 0 {
+		set := FindSetMatches(matcher.Value)
 		for _, v := range set {
 			var qs []IndexQuery
 			qs, err = c.schema.GetReadQueriesForMetricLabelValue(from, through, userID, metricName, matcher.Name, v)

pkg/chunk/opts.go

@@ -19,9 +19,9 @@ func init() {
 	}
 }
 
+// FindSetMatches returns list of values that can be equality matched on.
 // copied from Prometheus querier.go, removed check for Prometheus wrapper.
-// Returns list of values that can regex matches.
-func findSetMatches(pattern string) []string {
+func FindSetMatches(pattern string) []string {
 	escaped := false
 	sets := []*strings.Builder{{}}
 	for i := 0; i < len(pattern); i++ {

pkg/chunk/opts_test.go

@@ -44,7 +44,7 @@ func TestFindSetMatches(t *testing.T) {
 	}
 
 	for _, c := range cases {
-		matches := findSetMatches(c.pattern)
+		matches := FindSetMatches(c.pattern)
 		require.Equal(t, c.exp, matches)
 	}
 }

pkg/ingester/index/index.go

@@ -8,6 +8,7 @@ import (
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/pkg/labels"
 
+	"github.com/cortexproject/cortex/pkg/chunk"
 	"github.com/cortexproject/cortex/pkg/ingester/client"
 	"github.com/cortexproject/cortex/pkg/util"
 )
@@ -168,6 +169,13 @@ func (shard *indexShard) lookup(matchers []*labels.Matcher) []model.Fingerprint
 		if matcher.Type == labels.MatchEqual {
 			fps := values.fps[matcher.Value]
 			toIntersect = append(toIntersect, fps.fps...) // deliberate copy
+		} else if matcher.Type == labels.MatchRegexp && len(chunk.FindSetMatches(matcher.Value)) > 0 {
+			// The lookup is of the form `=~"a|b|c|d"`
+			set := chunk.FindSetMatches(matcher.Value)
+			for _, value := range set {
+				toIntersect = append(toIntersect, values.fps[value].fps...)
+			}
+			sort.Sort(toIntersect)
 		} else {
 			// accumulate the matching fingerprints (which are all distinct)
 			// then sort to maintain the invariant

pkg/ingester/index/index_test.go

@@ -42,6 +42,18 @@ func TestIndex(t *testing.T) {
 		{mustParseMatcher(`{foo="bar", flip="flap"}`), []model.Fingerprint{2}},
 		{mustParseMatcher(`{foo="baz", flip="flop"}`), []model.Fingerprint{1}},
 		{mustParseMatcher(`{foo="baz", flip="flap"}`), []model.Fingerprint{0}},
+
+		{mustParseMatcher(`{fizz=~"b.*"}`), []model.Fingerprint{}},
+
+		{mustParseMatcher(`{foo=~"bar.*"}`), []model.Fingerprint{2, 3}},
+		{mustParseMatcher(`{foo=~"ba.*"}`), []model.Fingerprint{0, 1, 2, 3}},
+		{mustParseMatcher(`{flip=~"flop|flap"}`), []model.Fingerprint{0, 1, 2, 3}},
+		{mustParseMatcher(`{flip=~"flaps"}`), []model.Fingerprint{}},
+
+		{mustParseMatcher(`{foo=~"bar|bax", flip="flop"}`), []model.Fingerprint{3}},
+		{mustParseMatcher(`{foo=~"bar|baz", flip="flap"}`), []model.Fingerprint{0, 2}},
+		{mustParseMatcher(`{foo=~"baz.+", flip="flop"}`), []model.Fingerprint{}},
+		{mustParseMatcher(`{foo=~"baz", flip="flap"}`), []model.Fingerprint{0}},
 	} {
 		assert.Equal(t, tc.fps, index.Lookup(tc.matchers))
 	}