resolved TODOs

Author: metonymic-smokey

index/scorch/optimize_knn.go

@@ -109,7 +109,7 @@ func (o *OptimizeVR) Finish() error {
 						// for each VR, populate postings list and iterators
 						// by passing the obtained vector index and getting similar vectors.
 						pl, err := vecIndex.Search(vr.vector, vr.k,
-							eligibleLocalDocNums, vr.searchParams)
+							vr.requireFiltering, eligibleLocalDocNums, vr.searchParams)
 						if err != nil {
 							errorsM.Lock()
 							errors = append(errors, err)

index/scorch/snapshot_index_vr.go

@@ -51,8 +51,9 @@ type IndexSnapshotVectorReader struct {
 	currID        index.IndexInternalID
 	ctx           context.Context
 
-	searchParams   json.RawMessage
-	eligibleDocIDs []index.IndexInternalID
+	searchParams     json.RawMessage
+	eligibleDocIDs   []index.IndexInternalID
+	requireFiltering bool
 }
 
 func (i *IndexSnapshotVectorReader) EligibleDocIDs() *roaring.Bitmap {
@@ -121,7 +122,7 @@ func (i *IndexSnapshotVectorReader) Advance(ID index.IndexInternalID,
 
 	if i.currPosting != nil && bytes.Compare(i.currID, ID) >= 0 {
 		i2, err := i.snapshot.VectorReader(i.ctx, i.vector, i.field, i.k,
-			i.searchParams, i.eligibleDocIDs)
+			i.searchParams, i.eligibleDocIDs, i.requireFiltering)
 		if err != nil {
 			return nil, err
 		}

index/scorch/snapshot_vector_index.go

@@ -26,16 +26,18 @@ import (
 )
 
 func (is *IndexSnapshot) VectorReader(ctx context.Context, vector []float32,
-	field string, k int64, searchParams json.RawMessage, filterIDs []index.IndexInternalID) (
+	field string, k int64, searchParams json.RawMessage,
+	filterIDs []index.IndexInternalID, requireFiltering bool) (
 	index.VectorReader, error) {
 
 	rv := &IndexSnapshotVectorReader{
-		vector:         vector,
-		field:          field,
-		k:              k,
-		snapshot:       is,
-		searchParams:   searchParams,
-		eligibleDocIDs: filterIDs,
+		vector:           vector,
+		field:            field,
+		k:                k,
+		snapshot:         is,
+		searchParams:     searchParams,
+		eligibleDocIDs:   filterIDs,
+		requireFiltering: requireFiltering,
 	}
 
 	if rv.postings == nil {

search/query/knn.go

@@ -35,9 +35,10 @@ type KNNQuery struct {
 	BoostVal    *Boost    `json:"boost,omitempty"`
 
 	// see KNNRequest.Params for description
-	Params        json.RawMessage `json:"params"`
-	FilterQuery   Query           `json:"filter,omitempty"`
-	FilterResults []index.IndexInternalID
+	Params           json.RawMessage `json:"params"`
+	FilterQuery      Query           `json:"filter,omitempty"`
+	FilterResults    []index.IndexInternalID
+	RequireFiltering bool
 }
 
 func NewKNNQuery(vector []float32) *KNNQuery {
@@ -69,8 +70,9 @@ func (q *KNNQuery) SetParams(params json.RawMessage) {
 	q.Params = params
 }
 
-func (q *KNNQuery) SetFilterQuery(f Query) {
+func (q *KNNQuery) SetFilterQuery(f Query, requireFiltering bool) {
 	q.FilterQuery = f
+	q.RequireFiltering = requireFiltering
 }
 
 func (q *KNNQuery) Searcher(ctx context.Context, i index.IndexReader,
@@ -89,5 +91,6 @@ func (q *KNNQuery) Searcher(ctx context.Context, i index.IndexReader,
 	}
 
 	return searcher.NewKNNSearcher(ctx, i, m, options, q.VectorField,
-		q.Vector, q.K, q.BoostVal.Value(), similarityMetric, q.Params, q.FilterResults)
+		q.Vector, q.K, q.BoostVal.Value(), similarityMetric, q.Params,
+		q.FilterResults, q.RequireFiltering)
 }

search/searcher/search_knn.go

@@ -50,12 +50,12 @@ type KNNSearcher struct {
 func NewKNNSearcher(ctx context.Context, i index.IndexReader, m mapping.IndexMapping,
 	options search.SearcherOptions, field string, vector []float32, k int64,
 	boost float64, similarityMetric string, searchParams json.RawMessage,
-	filterIDs []index.IndexInternalID) (
+	filterIDs []index.IndexInternalID, requireFiltering bool) (
 	search.Searcher, error) {
 
 	if vr, ok := i.(index.VectorIndexReader); ok {
 		vectorReader, err := vr.VectorReader(ctx, vector, field, k, searchParams,
-			filterIDs)
+			filterIDs, requireFiltering)
 		if err != nil {
 			return nil, err
 		}

search_knn.go

@@ -228,7 +228,8 @@ var (
 	knnOperatorOr  = knnOperator("or")
 )
 
-func createKNNQuery(req *SearchRequest, eligibleDocsMap map[int][]index.IndexInternalID) (
+func createKNNQuery(req *SearchRequest, eligibleDocsMap map[int][]index.IndexInternalID,
+	requiresFiltering map[int]bool) (
 	query.Query, []int64, int64, error) {
 	if requestHasKNN(req) {
 		// first perform validation
@@ -245,8 +246,9 @@ func createKNNQuery(req *SearchRequest, eligibleDocsMap map[int][]index.IndexInt
 			knnQuery.SetK(knn.K)
 			knnQuery.SetBoost(knn.Boost.Value())
 			knnQuery.SetParams(knn.Params)
-			knnQuery.SetFilterQuery(knn.FilterQuery)
-			if filterResults, exists := eligibleDocsMap[i]; exists {
+			knnQuery.SetFilterQuery(knn.FilterQuery, requiresFiltering[i])
+			filterResults, exists := eligibleDocsMap[i]
+			if exists && requiresFiltering[i] {
 				knnQuery.FilterResults = filterResults
 			}
 			subQueries = append(subQueries, knnQuery)
@@ -330,20 +332,26 @@ func (i *indexImpl) runKnnCollector(ctx context.Context, req *SearchRequest, rea
 	// maps the index of the KNN query in the req to the pre-filter hits aka
 	// eligible docs' internal IDs .
 	filterHitsMap := make(map[int][]index.IndexInternalID)
+	// Indicates if this query requires filtering downstream
+	// No filtering required if it's a match all query/no filters applied.
+	requiresFiltering := make(map[int]bool)
 
 	for idx, knnReq := range req.KNN {
 		// TODO Can use goroutines for this filter query stuff - do it if perf results
 		// show this to be significantly slow otherwise.
 		filterQ := knnReq.FilterQuery
-		// If there are no filters here, add a match all since that will ensure that
-		// all the live docs in the index are eligible.
-		// TODO See if running MatchAll queries can be skipped too - if perf shows
-		// them to be time-consuming in existing kNN tests?
 		if filterQ == nil {
-			filterQ = query.NewMatchAllQuery()
+			requiresFiltering[idx] = false
+		}
+
+		if _, ok := filterQ.(*query.MatchAllQuery); ok {
+			requiresFiltering[idx] = false
+			continue
 		}
 
 		if _, ok := filterQ.(*query.MatchNoneQuery); ok {
+			// Filtering required since no hits are eligible.
+			requiresFiltering[idx] = true
 			// a match none query just means none the documents are eligible
 			// hence, we can save on running the query.
 			continue
@@ -369,10 +377,11 @@ func (i *indexImpl) runKnnCollector(ctx context.Context, req *SearchRequest, rea
 		for _, docMatch := range filterHits {
 			filterHitsMap[idx] = append(filterHitsMap[idx], docMatch.IndexInternalID)
 		}
+		requiresFiltering[idx] = true
 	}
 
 	// Add the filter hits when creating the kNN query
-	KNNQuery, kArray, sumOfK, err := createKNNQuery(req, filterHitsMap)
+	KNNQuery, kArray, sumOfK, err := createKNNQuery(req, filterHitsMap, requiresFiltering)
 	if err != nil {
 		return nil, err
 	}