perf(core): Add a count min sketch for eq filter

dgraph/cmd/alpha/run.go

@@ -659,7 +659,7 @@ func run() {
 		AuthToken:          security.GetString("token"),
 		Audit:              conf,
 		ChangeDataConf:     Alpha.Conf.GetString("cdc"),
-		TypeFilterUidLimit: x.Config.Limit.GetInt64("type-filter-uid-limit"),
+		TypeFilterUidLimit: x.Config.Limit.GetUint64("type-filter-uid-limit"),
 	}
 
 	keys, err := ee.GetKeys(Alpha.Conf)

go.mod

@@ -37,6 +37,7 @@ require (
 	github.com/google/uuid v1.6.0
 	github.com/gorilla/websocket v1.5.3
 	github.com/graph-gophers/graphql-go v1.5.0
+	github.com/harshil-goel/BoomFilters v0.0.0-20241110225754-c092867bb322
 	github.com/hashicorp/vault/api v1.15.0
 	github.com/minio/minio-go/v6 v6.0.57
 	github.com/mitchellh/panicwrap v1.0.0
@@ -84,6 +85,7 @@ require (
 	github.com/cenkalti/backoff/v4 v4.3.0 // indirect
 	github.com/cespare/xxhash/v2 v2.3.0 // indirect
 	github.com/chewxy/math32 v1.11.0 // indirect
+	github.com/d4l3k/messagediff v1.2.1 // indirect
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/distribution/reference v0.6.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect

go.sum

@@ -136,6 +136,8 @@ github.com/coreos/pkg v0.0.0-20180928190104-399ea9e2e55f/go.mod h1:E3G3o1h8I7cfc
 github.com/cpuguy83/go-md2man/v2 v2.0.0-20190314233015-f79a8a8ca69d/go.mod h1:maD7wRr/U5Z6m/iR4s+kqSMx2CaBsrgA7czyZG/E6dU=
 github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E=
+github.com/d4l3k/messagediff v1.2.1 h1:ZcAIMYsUg0EAp9X+tt8/enBE/Q8Yd5kzPynLyKptt9U=
+github.com/d4l3k/messagediff v1.2.1/go.mod h1:Oozbb1TVXFac9FtSIxHBMnBCq2qeH/2KkEQxENCrlLo=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc h1:U9qPSI2PIWSS1VwoXQT9A3Wy9MM3WgvqSxFWenqJduM=
@@ -340,6 +342,8 @@ github.com/grpc-ecosystem/grpc-gateway v1.9.0 h1:bM6ZAFZmc/wPFaRDi0d5L7hGEZEx/2u
 github.com/grpc-ecosystem/grpc-gateway v1.9.0/go.mod h1:vNeuVxBJEsws4ogUvrchl83t/GYV9WGTSLVdBhOQFDY=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0 h1:asbCHRVmodnJTuQ3qamDwqVOIjwqUPTYmYuemVOx+Ys=
 github.com/grpc-ecosystem/grpc-gateway/v2 v2.22.0/go.mod h1:ggCgvZ2r7uOoQjOyu2Y1NhHmEPPzzuhWgcza5M1Ji1I=
+github.com/harshil-goel/BoomFilters v0.0.0-20241110225754-c092867bb322 h1:pk5qge6xZUS4Vh3pZpInFVW/fCSQSPNAQgdOF7Qo1kg=
+github.com/harshil-goel/BoomFilters v0.0.0-20241110225754-c092867bb322/go.mod h1:ssEcAHqGeDHiNeR3QVo4iZPFSTzY4UympUKUf3Ia1Rg=
 github.com/hashicorp/consul/api v1.1.0/go.mod h1:VmuI/Lkw1nC05EYQWNKwWGbkg+FbDBtguAZLlVdkD9Q=
 github.com/hashicorp/consul/sdk v0.1.1/go.mod h1:VKf9jXwCTEY1QZP2MOLRhb5i/I/ssyNV1vwHyQBF0x8=
 github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=

posting/lists.go

@@ -45,6 +45,7 @@ var (
 
 // Init initializes the posting lists package, the in memory and dirty list hash.
 func Init(ps *badger.DB, cacheSize int64) {
+	GlobalStatsHolder = NewStatsHolder()
 	pstore = ps
 	closer = z.NewCloser(1)
 	go x.MonitorMemoryMetrics(closer)

posting/stats.go

@@ -0,0 +1,100 @@
+/*
+ * Copyright 2024 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package posting
+
+import (
+	"math"
+	"sync"
+
+	boom "github.com/harshil-goel/BoomFilters"
+)
+
+var (
+	GlobalStatsHolder *StatsHolder
+)
+
+type StatsHolder struct {
+	sync.RWMutex
+
+	predStats map[string]StatContainer
+	//computeRequests chan *ComputeStatRequest
+}
+
+//type ComputeStatRequest struct {
+//	pred string
+//}
+
+func NewStatsHolder() *StatsHolder {
+	return &StatsHolder{
+		predStats: make(map[string]StatContainer),
+	}
+
+}
+
+type StatContainer interface {
+	InsertRecord([]byte, uint64)
+	Estimate([]byte) uint64
+}
+
+type EqContainer struct {
+	sync.RWMutex
+
+	cmf *boom.CountMinSketch
+}
+
+func NewEqContainer() *EqContainer {
+	return &EqContainer{cmf: boom.NewCountMinSketch(0.001, 0.99)}
+}
+
+func (eq *EqContainer) Estimate(key []byte) uint64 {
+	eq.RLock()
+	defer eq.RUnlock()
+
+	return eq.cmf.Count(key)
+}
+
+func (eq *EqContainer) InsertRecord(key []byte, count uint64) {
+	eq.Lock()
+	defer eq.Unlock()
+	eq.cmf.Set(key, count)
+}
+
+func (sh *StatsHolder) InsertRecord(pred string, key []byte, count uint64) {
+	sh.RLock()
+	val, ok := sh.predStats[pred]
+	sh.RUnlock()
+	if !ok {
+		return
+	}
+
+	val.InsertRecord(key, count)
+}
+
+func (sh *StatsHolder) ProcessEqPredicate(pred string, key []byte) uint64 {
+	sh.RLock()
+	val, ok := sh.predStats[pred]
+	sh.RUnlock()
+	if ok {
+		return val.Estimate(key)
+	}
+
+	sh.Lock()
+	val = NewEqContainer()
+	sh.predStats[pred] = val
+	sh.Unlock()
+	return math.MaxUint64
+}

worker/config.go

@@ -73,7 +73,7 @@ type Options struct {
 	// vs searched via type index. If the number of elements are too low, then querying the
 	// index might be slower. This would allow people to set their limit according to
 	// their use case.
-	TypeFilterUidLimit int64
+	TypeFilterUidLimit uint64
 }
 
 // Config holds an instance of the server options..

worker/task.go

@@ -20,6 +20,7 @@
 	"bytes"
 	"context"
 	"fmt"
+	"math"
 	"sort"
 	"strconv"
 	"strings"
@@ -836,6 +837,10 @@
 				return err
 			}
 
+			if srcFn.fnType == compareAttrFn {
+				posting.GlobalStatsHolder.InsertRecord(q.Attr, []byte(srcFn.tokens[i]), uint64(pl.ApproxLen()))
+			}
+
 			switch {
 			case q.DoCount:
 				if i == 0 {
@@ -1797,6 +1802,48 @@
 	return langs[0]
 }
 
+func planForEqFilter(fc *functionContext, pred string, uidlist []uint64) {
+	checkUidEmpty := func(uids []uint64) bool {
+		for _, i := range uids {
+			if i == 0 {
+				return false
+			}
+		}
+		return true
+	}
+
+	if uint64(len(uidlist)) < Config.TypeFilterUidLimit && checkUidEmpty(uidlist) {
+		fc.tokens = fc.tokens[:0]
+		fc.n = len(uidlist)
+		return
+	}
+
+	estimatedCount := uint64(0)
+	gotEstimate := false
+	for _, eqToken := range fc.tokens {
+		count := posting.GlobalStatsHolder.ProcessEqPredicate(pred, []byte(eqToken))
+		if count != math.MaxUint64 {
+			estimatedCount += count
+			gotEstimate = true
+		} else {
+			break
+		}
+	}
+
+	if gotEstimate && estimatedCount == 0 {
+		gotEstimate = false
+	}
+
+	// TODO make a different config
+	if gotEstimate && uint64(len(uidlist)) < estimatedCount/Config.TypeFilterUidLimit && checkUidEmpty(uidlist) {
+		fc.tokens = fc.tokens[:0]
+		fc.n = len(uidlist)
+		return
+	}
+
+	fc.n = len(fc.tokens)
+}
+
 func parseSrcFn(ctx context.Context, q *pb.Query) (*functionContext, error) {
 	fnType, f := parseFuncType(q.SrcFunc)
 	attr := q.Attr
@@ -1886,15 +1933,6 @@
 			fc.tokens = append(fc.tokens, tokens...)
 		}
 
-		checkUidEmpty := func(uids []uint64) bool {
-			for _, i := range uids {
-				if i == 0 {
-					return false
-				}
-			}
-			return true
-		}
-
 		// In case of non-indexed predicate, there won't be any tokens. We will fetch value
 		// from data keys.
 		// If number of index keys is more than no. of uids to filter, so its better to fetch values
@@ -1908,10 +1946,10 @@
 		case q.UidList != nil && len(fc.tokens) > len(q.UidList.Uids) && fc.fname != eq:
 			fc.tokens = fc.tokens[:0]
 			fc.n = len(q.UidList.Uids)
-		case q.UidList != nil && fc.fname == eq && strings.HasSuffix(attr, "dgraph.type") &&
-			int64(len(q.UidList.Uids)) < Config.TypeFilterUidLimit && checkUidEmpty(q.UidList.Uids):
-			fc.tokens = fc.tokens[:0]
-			fc.n = len(q.UidList.Uids)
+		case q.UidList != nil && fc.fname == eq:
+			if len(fc.tokens) > 0 {
+				planForEqFilter(fc, attr, q.UidList.Uids)
+			}
 		default:
 			fc.n = len(fc.tokens)
 		}

worker/worker_test.go

@@ -248,6 +248,42 @@ func runQuery(dg *dgo.Dgraph, attr string, uids []uint64, srcFunc []string) (*ap
 	return resp, err
 }
 
+func BenchmarkEqFilter(b *testing.B) {
+	dg, err := testutil.DgraphClient(testutil.SockAddr)
+	if err != nil {
+		panic(err)
+	}
+	err = dg.Alter(context.Background(), &api.Operation{DropAll: true})
+	if err != nil {
+		panic(err)
+	}
+
+	n := 10000
+	for i := 0; i < n; i++ {
+		rdf := fmt.Sprintf("<%#x> <dgraph.type> \"abc\" .", i+1)
+		mu := &api.Mutation{SetNquads: []byte(rdf), CommitNow: true}
+		err := testutil.RetryMutation(dg, mu)
+		if err != nil {
+			panic(err)
+		}
+	}
+
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		query := `
+		{
+			q(func: uid(1, 2)) @filter(type(abc)) {
+				uid
+			}
+		}
+		`
+		_, err := testutil.RetryQuery(dg, query)
+		if err != nil {
+			panic(err)
+		}
+	}
+}
+
 // Index-related test. Similar to TestProcessTaskIndex but we call MergeLists only
 // at the end. In other words, everything is happening only in mutation layers,
 // and not committed to BadgerDB until near the end.