[dbnode] Optimize index.Aggregate() for FieldQuery

src/dbnode/integration/peers_bootstrap_index_aggregate_test.go

@@ -192,7 +192,7 @@ func TestPeersBootstrapIndexAggregateQuery(t *testing.T) {
 	// Now test term filtering, match all new_*r*, filtering on `foo`
 	regexpQuery, err = idx.NewRegexpQuery([]byte("city"), []byte("new_.*r.*"))
 	require.NoError(t, err)
-	queryOpts.TermFilter = index.AggregateTermFilter([][]byte{[]byte("foo")})
+	queryOpts.FieldFilter = index.AggregateFieldFilter([][]byte{[]byte("foo")})
 	iter, exhaustive, err = session.Aggregate(ns1.ID(),
 		index.Query{regexpQuery}, queryOpts)
 	require.NoError(t, err)
@@ -212,7 +212,7 @@ func TestPeersBootstrapIndexAggregateQuery(t *testing.T) {
 	// Now test term filter and tag name filtering, match all new_*r*, names only, filtering on `city`
 	regexpQuery, err = idx.NewRegexpQuery([]byte("city"), []byte("new_.*r.*"))
 	require.NoError(t, err)
-	queryOpts.TermFilter = index.AggregateTermFilter([][]byte{[]byte("city")})
+	queryOpts.FieldFilter = index.AggregateFieldFilter([][]byte{[]byte("city")})
 	queryOpts.Type = index.AggregateTagNames
 	iter, exhaustive, err = session.Aggregate(ns1.ID(),
 		index.Query{regexpQuery}, queryOpts)

src/dbnode/network/server/tchannelthrift/convert/convert.go

@@ -306,9 +306,9 @@ func FromRPCAggregateQueryRequest(
 		return nil, index.Query{}, index.AggregationOptions{}, err
 	}
 
-	opts.TermFilter = make(index.AggregateTermFilter, 0, len(req.TagNameFilter))
+	opts.FieldFilter = make(index.AggregateFieldFilter, 0, len(req.TagNameFilter))
 	for _, f := range req.TagNameFilter {
-		opts.TermFilter = append(opts.TermFilter, []byte(f))
+		opts.FieldFilter = append(opts.FieldFilter, []byte(f))
 	}
 
 	if req.AggregateQueryType == rpc.AggregateQueryType_AGGREGATE_BY_TAG_NAME_VALUE {
@@ -351,7 +351,7 @@ func FromRPCAggregateQueryRawRequest(
 		return nil, index.Query{}, index.AggregationOptions{}, err
 	}
 
-	opts.TermFilter = index.AggregateTermFilter(req.TagNameFilter)
+	opts.FieldFilter = index.AggregateFieldFilter(req.TagNameFilter)
 	if req.AggregateQueryType == rpc.AggregateQueryType_AGGREGATE_BY_TAG_NAME_VALUE {
 		opts.Type = index.AggregateTagNamesAndValues
 	} else {
@@ -407,9 +407,9 @@ func ToRPCAggregateQueryRawRequest(
 		request.AggregateQueryType = rpc.AggregateQueryType_AGGREGATE_BY_TAG_NAME
 	}
 
-	// TODO(prateek): pool the []byte underlying opts.TermFilter
-	filters := make([][]byte, 0, len(opts.TermFilter))
-	for _, f := range opts.TermFilter {
+	// TODO(prateek): pool the []byte underlying opts.FieldFilter
+	filters := make([][]byte, 0, len(opts.FieldFilter))
+	for _, f := range opts.FieldFilter {
 		copied := append([]byte(nil), f...)
 		filters = append(filters, copied)
 	}

src/dbnode/network/server/tchannelthrift/convert/convert_test.go

@@ -174,7 +174,7 @@ func TestConvertAggregateRawQueryRequest(t *testing.T) {
 			Limit:          10,
 		},
 		Type: index.AggregateTagNamesAndValues,
-		TermFilter: index.AggregateTermFilter{
+		FieldFilter: index.AggregateFieldFilter{
 			[]byte("some"),
 			[]byte("string"),
 		},

src/dbnode/network/server/tchannelthrift/node/service_test.go

@@ -1661,7 +1661,7 @@ func TestServiceAggregate(t *testing.T) {
 				EndExclusive:   end,
 				Limit:          10,
 			},
-			TermFilter: index.AggregateTermFilter{
+			FieldFilter: index.AggregateFieldFilter{
 				[]byte("foo"), []byte("bar"),
 			},
 			Type: index.AggregateTagNamesAndValues,

src/dbnode/storage/index.go

@@ -909,17 +909,26 @@ func (i *nsIndex) AggregateQuery(
 ) (index.AggregateQueryResult, error) {
 	// Get results and set the filters, namespace ID and size limit.
 	results := i.aggregateResultsPool.Get()
-	results.Reset(i.nsMetadata.ID(), index.AggregateResultsOptions{
-		SizeLimit:  opts.Limit,
-		TermFilter: opts.TermFilter,
-		Type:       opts.Type,
-	})
+	aopts := index.AggregateResultsOptions{
+		SizeLimit:   opts.Limit,
+		FieldFilter: opts.FieldFilter,
+		Type:        opts.Type,
+	}
 	ctx.RegisterFinalizer(results)
 	// use appropriate fn to query underlying blocks.
+	// default to block.Query()
 	fn := i.execBlockQueryFn
+	// use block.Aggregate() when possible
 	if query.Equal(allQuery) {
 		fn = i.execBlockAggregateQueryFn
 	}
+	field, isField := idx.FieldQuery(query.Query)
+	if isField {
+		fn = i.execBlockAggregateQueryFn
+		aopts.FieldFilter = aopts.FieldFilter.AddIfMissing(field)
+	}
+	aopts.FieldFilter = aopts.FieldFilter.SortAndDedupe()
+	results.Reset(i.nsMetadata.ID(), aopts)
 	exhaustive, err := i.query(ctx, query, results, opts.QueryOptions, fn)
 	if err != nil {
 		return index.AggregateQueryResult{}, err

src/dbnode/storage/index/aggregate_results.go

@@ -199,7 +199,7 @@ func (r *aggregatedResults) addTermWithLock(
 
 	// if a term filter is provided, ensure this field matches the filter,
 	// otherwise ignore it.
-	filter := r.aggregateOpts.TermFilter
+	filter := r.aggregateOpts.FieldFilter
 	if filter != nil && !filter.Allow(term) {
 		return nil
 	}
@@ -244,7 +244,7 @@ func (r *aggregatedResults) addFieldWithLock(
 
 	// if a term filter is provided, ensure this field matches the filter,
 	// otherwise ignore it.
-	filter := r.aggregateOpts.TermFilter
+	filter := r.aggregateOpts.FieldFilter
 	if filter != nil && !filter.Allow(term) {
 		return nil
 	}

src/dbnode/storage/index/aggregate_results_test.go

@@ -213,13 +213,13 @@ func expectedTermsOnly(ex map[string][]string) map[string][]string {
 	return m
 }
 
-func toFilter(strs ...string) AggregateTermFilter {
+func toFilter(strs ...string) AggregateFieldFilter {
 	b := make([][]byte, len(strs))
 	for i, s := range strs {
 		b[i] = []byte(s)
 	}
 
-	return AggregateTermFilter(b)
+	return AggregateFieldFilter(b)
 }
 
 var mergeTests = []struct {
@@ -247,7 +247,7 @@ var mergeTests = []struct {
 	},
 	{
 		name: "no limit empty filter",
-		opts: AggregateResultsOptions{TermFilter: toFilter()},
+		opts: AggregateResultsOptions{FieldFilter: toFilter()},
 		expected: map[string][]string{
 			"foo":  []string{"bar", "biz", "baz"},
 			"fizz": []string{"bar"},
@@ -257,20 +257,20 @@ var mergeTests = []struct {
 	},
 	{
 		name:     "no limit matchless filter",
-		opts:     AggregateResultsOptions{TermFilter: toFilter("zig")},
+		opts:     AggregateResultsOptions{FieldFilter: toFilter("zig")},
 		expected: map[string][]string{},
 	},
 	{
 		name: "empty limit with filter",
-		opts: AggregateResultsOptions{TermFilter: toFilter("buzz")},
+		opts: AggregateResultsOptions{FieldFilter: toFilter("buzz")},
 		expected: map[string][]string{
 			"buzz": []string{"bar", "bag"},
 		},
 	},
 	{
 		name: "with limit with filter",
 		opts: AggregateResultsOptions{
-			SizeLimit: 2, TermFilter: toFilter("buzz", "qux", "fizz")},
+			SizeLimit: 2, FieldFilter: toFilter("buzz", "qux", "fizz")},
 		expected: map[string][]string{
 			"fizz": []string{"bar"},
 			"buzz": []string{"bar", "bag"},

src/dbnode/storage/index/aggregated_term_filter.go

@@ -20,10 +20,13 @@
 
 package index
 
-import "bytes"
+import (
+	"bytes"
+	"sort"
+)
 
 // Allow returns true if the given term satisfies the filter.
-func (f AggregateTermFilter) Allow(term []byte) bool {
+func (f AggregateFieldFilter) Allow(term []byte) bool {
 	if len(f) == 0 {
 		// NB: if filter is empty, all values are valid.
 		return true
@@ -37,3 +40,38 @@ func (f AggregateTermFilter) Allow(term []byte) bool {
 
 	return false
 }
+
+// AddIfMissing adds the provided field if it's missing from the filter.
+func (f AggregateFieldFilter) AddIfMissing(field []byte) AggregateFieldFilter {
+	for _, fi := range f {
+		if bytes.Equal(fi, field) {
+			return f
+		}
+	}
+	f = append(f, field)
+	return f
+}
+
+// SortAndDedupe sorts and de-dupes the fields in the filter.
+func (f AggregateFieldFilter) SortAndDedupe() AggregateFieldFilter {
+	// short-circuit if possible.
+	if len(f) <= 1 {
+		return f
+	}
+	// sort the provided filter fields in order.
+	sort.Slice(f, func(i, j int) bool {
+		return bytes.Compare(f[i], f[j]) < 0
+	})
+	// ensure successive elements are not the same.
+	deduped := make(AggregateFieldFilter, 0, len(f))
+	if len(f) > 0 {
+		deduped = append(deduped, f[0])
+	}
+	for i := 1; i < len(f); i++ {
+		if bytes.Equal(f[i-1], f[i]) {
+			continue
+		}
+		deduped = append(deduped, f[i])
+	}
+	return deduped
+}

src/dbnode/storage/index/aggregated_term_filter_test.go

@@ -0,0 +1,50 @@
+// Copyright (c) 2019 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package index
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestSortAndDedupe(t *testing.T) {
+	filter := AggregateFieldFilter{
+		[]byte("c"),
+		[]byte("a"),
+		[]byte("b"),
+		[]byte("a"),
+		[]byte("c"),
+		[]byte("a"),
+	}
+	observed := filter.SortAndDedupe()
+	require.Equal(t, AggregateFieldFilter{
+		[]byte("a"),
+		[]byte("b"),
+		[]byte("c"),
+	}, observed)
+}
+
+func TestSortAndDedupeEmpty(t *testing.T) {
+	filter := AggregateFieldFilter{}
+	observed := filter.SortAndDedupe()
+	require.Equal(t, AggregateFieldFilter{}, observed)
+}

src/dbnode/storage/index/block.go

@@ -902,7 +902,23 @@ func (b *block) Aggregate(
 			if bytes.Equal(field, doc.IDReservedFieldName) {
 				return false
 			}
-			return aggOpts.TermFilter.Allow(field)
+			return aggOpts.FieldFilter.Allow(field)
+		},
+		fieldIterFn: func(s segment.Segment) (segment.FieldsIterator, error) {
+			// NB(prateek): we default to using the regular (FST) fields iterator
+			// unless we have a predefined list of fields we know we need to restrict
+			// our search to, in which case we iterate that list and check if known values
+			// in the FST to restrict our search. This is going to be significantly faster
+			// while len(FieldsFilter) < 5-10 elements;
+			// but there will exist a ratio between the len(FieldFilter) v size(FST) after which
+			// iterating the entire FST is faster.
+			// Here, we chose to avoid factoring that in to our choice because almost all input
+			// to this function is expected to have (FieldsFilter) pretty small. If that changes
+			// in the future, we can revisit this.
+			if len(aggOpts.FieldFilter) == 0 {
+				return s.FieldsIterable().Fields()
+			}
+			return newFilterFieldsIterator(s, aggOpts.FieldFilter)
 		},
 	}
 

src/dbnode/storage/index/block_test.go

@@ -1728,15 +1728,35 @@ func TestBlockE2EInsertAggregate(t *testing.T) {
 		SizeLimit: 10,
 		Type:      AggregateTagNamesAndValues,
 	}, testOpts)
-
 	exhaustive, err := b.Aggregate(resource.NewCancellableLifetime(), QueryOptions{Limit: 10}, results)
 	require.NoError(t, err)
 	require.True(t, exhaustive)
-
 	assertAggregateResultsMapEquals(t, map[string][]string{
 		"bar":  []string{"baz", "qux"},
 		"some": []string{"more", "other"},
 	}, results)
+
+	results = NewAggregateResults(ident.StringID("ns"), AggregateResultsOptions{
+		SizeLimit:   10,
+		Type:        AggregateTagNamesAndValues,
+		FieldFilter: AggregateFieldFilter{[]byte("bar")},
+	}, testOpts)
+	exhaustive, err = b.Aggregate(resource.NewCancellableLifetime(), QueryOptions{Limit: 10}, results)
+	require.NoError(t, err)
+	require.True(t, exhaustive)
+	assertAggregateResultsMapEquals(t, map[string][]string{
+		"bar": []string{"baz", "qux"},
+	}, results)
+
+	results = NewAggregateResults(ident.StringID("ns"), AggregateResultsOptions{
+		SizeLimit:   10,
+		Type:        AggregateTagNamesAndValues,
+		FieldFilter: AggregateFieldFilter{[]byte("random")},
+	}, testOpts)
+	exhaustive, err = b.Aggregate(resource.NewCancellableLifetime(), QueryOptions{Limit: 10}, results)
+	require.NoError(t, err)
+	require.True(t, exhaustive)
+	assertAggregateResultsMapEquals(t, map[string][]string{}, results)
 }
 
 func assertAggregateResultsMapEquals(t *testing.T, expected map[string][]string, observed AggregateResults) {

src/dbnode/storage/index/fields_terms_iterator.go

@@ -29,6 +29,7 @@ import (
 type fieldsAndTermsIteratorOpts struct {
 	iterateTerms bool
 	allowFn      allowFn
+	fieldIterFn  newFieldIterFn
 }
 
 func (o fieldsAndTermsIteratorOpts) allow(f []byte) bool {
@@ -38,8 +39,17 @@ func (o fieldsAndTermsIteratorOpts) allow(f []byte) bool {
 	return o.allowFn(f)
 }
 
+func (o fieldsAndTermsIteratorOpts) newFieldIter(s segment.Segment) (segment.FieldsIterator, error) {
+	if o.fieldIterFn == nil {
+		return s.FieldsIterable().Fields()
+	}
+	return o.fieldIterFn(s)
+}
+
 type allowFn func(field []byte) bool
 
+type newFieldIterFn func(s segment.Segment) (segment.FieldsIterator, error)
+
 type fieldsAndTermsIter struct {
 	seg  segment.Segment
 	opts fieldsAndTermsIteratorOpts
@@ -81,7 +91,7 @@ func (fti *fieldsAndTermsIter) Reset(s segment.Segment, opts fieldsAndTermsItera
 	if s == nil {
 		return nil
 	}
-	fiter, err := s.FieldsIterable().Fields()
+	fiter, err := fti.opts.newFieldIter(s)
 	if err != nil {
 		return err
 	}