Skip to content

Commit

Permalink
Add KNN Operator
Browse files Browse the repository at this point in the history
  • Loading branch information
metonymic-smokey committed Nov 17, 2023
1 parent d06a3a9 commit 00b5d9f
Show file tree
Hide file tree
Showing 3 changed files with 38 additions and 12 deletions.
10 changes: 6 additions & 4 deletions index_impl.go
Original file line number Diff line number Diff line change
Expand Up @@ -496,10 +496,12 @@ func (i *indexImpl) SearchInContext(ctx context.Context, req *SearchRequest) (sr
ctx = context.WithValue(ctx, search.GeoBufferPoolCallbackKey,
search.GeoBufferPoolCallbackFunc(getBufferPool))


// Using a disjunction query to get union of results from KNN query
// and the original query
searchQuery := disjunctQueryWithKNN(req)
// Using a query to get results from KNN queries
// and the original query based on the KNN operator.
searchQuery, err := queryWithKNN(req)
if err != nil {
return nil, err
}

searcher, err := searchQuery.Searcher(ctx, indexReader, i.m, search.SearcherOptions{
Explain: req.Explain,
Expand Down
36 changes: 30 additions & 6 deletions search_knn.go
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,8 @@ import (
"github.com/blevesearch/bleve/v2/search/query"
)

type knnOperator string

type SearchRequest struct {
Query query.Query `json:"query"`
Size int `json:"size"`
Expand All @@ -39,7 +41,8 @@ type SearchRequest struct {
SearchAfter []string `json:"search_after"`
SearchBefore []string `json:"search_before"`

KNN []*KNNRequest `json:"knn"`
KNN []*KNNRequest `json:"knn"`
KNNOperator knnOperator `json:"knn_operator,omitempty"`

sortFunc func(sort.Interface)
}
Expand All @@ -61,6 +64,10 @@ func (r *SearchRequest) AddKNN(field string, vector []float32, k int64, boost fl
})
}

func (r *SearchRequest) AddKNNOperator(operator knnOperator) {
r.KNNOperator = operator
}

// UnmarshalJSON deserializes a JSON representation of
// a SearchRequest
func (r *SearchRequest) UnmarshalJSON(input []byte) error {
Expand All @@ -78,6 +85,7 @@ func (r *SearchRequest) UnmarshalJSON(input []byte) error {
SearchAfter []string `json:"search_after"`
SearchBefore []string `json:"search_before"`
KNN []*KNNRequest `json:"knn"`
KNNOperator knnOperator `json:"knn_operator"`
}

err := json.Unmarshal(input, &temp)
Expand Down Expand Up @@ -120,6 +128,10 @@ func (r *SearchRequest) UnmarshalJSON(input []byte) error {
}

r.KNN = temp.KNN
r.KNNOperator = temp.KNNOperator
if r.KNNOperator == "" {
r.KNNOperator = knnOperatorOr
}

return nil

Expand All @@ -142,24 +154,36 @@ func copySearchRequest(req *SearchRequest) *SearchRequest {
SearchAfter: req.SearchAfter,
SearchBefore: req.SearchBefore,
KNN: req.KNN,
KNNOperator: req.KNNOperator,
}
return &rv

}

func disjunctQueryWithKNN(req *SearchRequest) query.Query {
var (
knnOperatorAnd = knnOperator("and")
knnOperatorOr = knnOperator("or")
)

func queryWithKNN(req *SearchRequest) (query.Query, error) {
if len(req.KNN) > 0 {
disjuncts := []query.Query{req.Query}
subQueries := []query.Query{req.Query}
for _, knn := range req.KNN {
if knn != nil {
knnQuery := query.NewKNNQuery(knn.Vector)
knnQuery.SetFieldVal(knn.Field)
knnQuery.SetK(knn.K)
knnQuery.SetBoost(knn.Boost.Value())
disjuncts = append(disjuncts, knnQuery)
subQueries = append(subQueries, knnQuery)
}
}
return query.NewDisjunctionQuery(disjuncts)
if req.KNNOperator == knnOperatorAnd {
return query.NewConjunctionQuery(subQueries), nil
} else if req.KNNOperator == knnOperatorOr {
return query.NewDisjunctionQuery(subQueries), nil
} else {
return nil, fmt.Errorf("unknown knn operator: %s", req.KNNOperator)
}
}
return req.Query
return req.Query, nil
}
4 changes: 2 additions & 2 deletions search_no_knn.go
Original file line number Diff line number Diff line change
Expand Up @@ -144,6 +144,6 @@ func copySearchRequest(req *SearchRequest) *SearchRequest {
return &rv
}

func disjunctQueryWithKNN(req *SearchRequest) query.Query {
return req.Query
func queryWithKNN(req *SearchRequest) (query.Query, error) {
return req.Query, nil
}

0 comments on commit 00b5d9f

Please sign in to comment.