Regex set optimisation when looking series in ingester

CHANGELOG.md

@@ -23,7 +23,7 @@
 * [ENHANCEMENT] FIFO cache to support eviction based on memory usage. The `-<prefix>.fifocache.size` CLI flag has been renamed to `-<prefix>.fifocache.max-size-items` as well as its YAML config option `size` renamed to `max_size_items`. Added `-<prefix>.fifocache.max-size-bytes` CLI flag and YAML config option `max_size_bytes` to specify memory limit of the cache. #2319
 * [ENHANCEMENT] Single Binary: Added query-frontend to the single binary.  Single binary users will now benefit from various query-frontend features.  Primarily: sharding, parallelization, load shedding, additional caching (if configured), and query retries. #2437
 * [ENHANCEMENT] Allow 1w (where w denotes week) and 1y (where y denotes year) when setting `-store.cache-lookups-older-than` and `-store.max-look-back-period`. #2454
-* [ENHANCEMENT] Optimize index queries for matchers using "a|b|c"-type regex. #2446
+* [ENHANCEMENT] Optimize index queries for matchers using "a|b|c"-type regex. #2446 #2475
 * [BUGFIX] Fixes #2411, Ensure requests are properly routed to the prometheus api embedded in the query if `-server.path-prefix` is set. #2372
 * [BUGFIX] Experimental TSDB: fixed chunk data corruption when querying back series using the experimental blocks storage. #2400
 * [BUGFIX] Cassandra Storage: Fix endpoint TLS host verification. #2109

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))
 	}