Merge pull request #15 from bgokden/refactor

Filters added.

data/search.go

@@ -15,6 +15,7 @@ import (
 	pb "github.com/bgokden/veri/veriservice"
 	badger "github.com/dgraph-io/badger/v2"
 	bpb "github.com/dgraph-io/badger/v2/pb"
+	"github.com/tidwall/gjson"
 )
 
 // type SearchConfig struct {
@@ -49,6 +50,8 @@ type Collector struct {
 	DatumKey       *pb.DatumKey
 	N              uint32
 	HigherIsBetter bool
+	Filters        []string
+	GroupFilters   []string
 }
 
 // // ScoredDatum helps to keep Data ordered
@@ -152,6 +155,38 @@ func (c *Collector) ToList(key []byte, itr *badger.Iterator) (*bpb.KVList, error
 		datumScore := &gencoder.DatumScore{
 			Score: c.ScoreFunc(datumKey.Feature, c.DatumKey.Feature),
 		}
+
+		if len(c.GroupFilters) > 0 {
+			filterFailed := false
+			jsonLabel := string(datumKey.GroupLabel)
+			for _, filter := range c.Filters {
+				value := gjson.Get(jsonLabel, filter)
+				if !value.Exists() {
+					filterFailed = true
+					break
+				}
+			}
+			if filterFailed {
+				break
+			}
+		}
+
+		if len(c.Filters) > 0 {
+			filterFailed := false
+			datumValue, _ := ToDatumValue(keyCopy)
+			jsonLabel := string(datumValue.Label)
+			for _, filter := range c.Filters {
+				value := gjson.Get(jsonLabel, filter)
+				if !value.Exists() {
+					filterFailed = true
+					break
+				}
+			}
+			if filterFailed {
+				break
+			}
+		}
+
 		datumScoreBytes, _ := datumScore.Marshal()
 		kv := &bpb.KV{
 			Key:       keyCopy,
@@ -192,6 +227,8 @@ func (dt *Data) Search(datum *pb.Datum, config *pb.SearchConfig) *Collector {
 	c.ScoreFunc = GetVectorComparisonFunction(config.ScoreFuncName)
 	c.HigherIsBetter = config.HigherIsBetter
 	c.N = config.Limit
+	c.Filters = config.Filters
+	c.GroupFilters = config.GroupFilters
 	stream := dt.DB.NewStream()
 	// db.NewStreamAt(readTs) for managed mode.
 

go.mod

@@ -16,6 +16,7 @@ require (
 	github.com/spf13/cobra v1.1.1
 	github.com/spf13/viper v1.7.0
 	github.com/stretchr/testify v1.4.0
+	github.com/tidwall/gjson v1.6.3
 	google.golang.org/grpc v1.27.0
 	google.golang.org/protobuf v1.25.0
 )

go.sum

@@ -252,6 +252,12 @@ github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJy
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 github.com/subosito/gotenv v1.2.0 h1:Slr1R9HxAlEKefgq5jn9U+DnETlIUa6HfgEzj0g5d7s=
 github.com/subosito/gotenv v1.2.0/go.mod h1:N0PQaV/YGNqwC0u51sEeR/aUtSLEXKX9iv69rRypqCw=
+github.com/tidwall/gjson v1.6.3 h1:aHoiiem0dr7GHkW001T1SMTJ7X5PvyekH5WX0whWGnI=
+github.com/tidwall/gjson v1.6.3/go.mod h1:BaHyNc5bjzYkPqgLq7mdVzeiRtULKULXLgZFKsxEHI0=
+github.com/tidwall/match v1.0.1 h1:PnKP62LPNxHKTwvHHZZzdOAOCtsJTjo6dZLCwpKm5xc=
+github.com/tidwall/match v1.0.1/go.mod h1:LujAq0jyVjBy028G1WhWfIzbpQfMO8bBZ6Tyb0+pL9E=
+github.com/tidwall/pretty v1.0.2 h1:Z7S3cePv9Jwm1KwS0513MRaoUe3S01WPbLNV40pwWZU=
+github.com/tidwall/pretty v1.0.2/go.mod h1:XNkn88O1ChpSDQmQeStsy+sBenx6DDtFZJxhVysOjyk=
 github.com/tmc/grpc-websocket-proxy v0.0.0-20190109142713-0ad062ec5ee5/go.mod h1:ncp9v5uamzpCO7NfCPTXjqaC+bZgJeR0sMTm6dMHP7U=
 github.com/ugorji/go/codec v0.0.0-20181204163529-d75b2dcb6bc8/go.mod h1:VFNgLljTbGfSG7qAOspJ7OScBnGdDN/yBr0sguwnwf0=
 github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU=