MB-65473: Refactor and Optimize Pre-Filtered Vector Search (#2169)

- Refactor pre-filtered vector search to enhance performance and reduce
memory footprint.
- Replace the current bitmap-based approach for calculating segment
local document numbers with a more direct method, where the local
document numbers are mapped directly to the segment ID during the
execution of the eligible collector.
- Requires:
    - blevesearch/bleve_index_api#63
    - blevesearch/bleve_index_api#66
    - blevesearch/zapx#317
    - blevesearch/go-faiss#41
    - blevesearch/faiss#49

---------

Co-authored-by: Abhinav Dangeti <abhinav@couchbase.com>

Author: CascadingRadium

go.mod

@@ -7,7 +7,7 @@ require (
 	github.com/bits-and-blooms/bitset v1.12.0
 	github.com/blevesearch/bleve_index_api v1.1.12
 	github.com/blevesearch/geo v0.1.20
-	github.com/blevesearch/go-faiss v1.0.24
+	github.com/blevesearch/go-faiss v1.0.25
 	github.com/blevesearch/go-metrics v0.0.0-20201227073835-cf1acfcdf475
 	github.com/blevesearch/go-porterstemmer v1.0.3
 	github.com/blevesearch/goleveldb v1.0.1

go.sum

@@ -6,8 +6,8 @@ github.com/blevesearch/bleve_index_api v1.1.12 h1:P4bw9/G/5rulOF7SJ9l4FsDoo7UFJ+
 github.com/blevesearch/bleve_index_api v1.1.12/go.mod h1:PbcwjIcRmjhGbkS/lJCpfgVSMROV6TRubGGAODaK1W8=
 github.com/blevesearch/geo v0.1.20 h1:paaSpu2Ewh/tn5DKn/FB5SzvH0EWupxHEIwbCk/QPqM=
 github.com/blevesearch/geo v0.1.20/go.mod h1:DVG2QjwHNMFmjo+ZgzrIq2sfCh6rIHzy9d9d0B59I6w=
-github.com/blevesearch/go-faiss v1.0.24 h1:K79IvKjoKHdi7FdiXEsAhxpMuns0x4fM0BO93bW5jLI=
-github.com/blevesearch/go-faiss v1.0.24/go.mod h1:OMGQwOaRRYxrmeNdMrXJPvVx8gBnvE5RYrr0BahNnkk=
+github.com/blevesearch/go-faiss v1.0.25 h1:lel1rkOUGbT1CJ0YgzKwC7k+XH0XVBHnCVWahdCXk4U=
+github.com/blevesearch/go-faiss v1.0.25/go.mod h1:OMGQwOaRRYxrmeNdMrXJPvVx8gBnvE5RYrr0BahNnkk=
 github.com/blevesearch/go-metrics v0.0.0-20201227073835-cf1acfcdf475 h1:kDy+zgJFJJoJYBvdfBSiZYBbdsUL0XcjHYWezpQBGPA=
 github.com/blevesearch/go-metrics v0.0.0-20201227073835-cf1acfcdf475/go.mod h1:9eJDeqxJ3E7WnLebQUlPD7ZjSce7AnDb9vjGmMCbD0A=
 github.com/blevesearch/go-porterstemmer v1.0.3 h1:GtmsqID0aZdCSNiY8SkuPJ12pD4jI+DdXTAn4YRcHCo=

index/scorch/optimize_knn.go

@@ -23,7 +23,6 @@ import (
 	"sync"
 	"sync/atomic"
 
-	"github.com/RoaringBitmap/roaring"
 	"github.com/blevesearch/bleve/v2/search"
 	index "github.com/blevesearch/bleve_index_api"
 	segment_api "github.com/blevesearch/scorch_segment_api/v2"
@@ -65,27 +64,6 @@ func (o *OptimizeVR) Finish() error {
 	var errorsM sync.Mutex
 	var errors []error
 
-	var snapshotGlobalDocNums map[int]*roaring.Bitmap
-	var eligibleDocIDsMap map[string]map[int]*roaring.Bitmap
-	fields := make([]string, len(o.vrs))
-	for field := range o.vrs {
-		fields = append(fields, field)
-	}
-
-	if o.requiresFiltering {
-		snapshotGlobalDocNums = o.snapshot.globalDocNums()
-		eligibleDocIDsMap = make(map[string]map[int]*roaring.Bitmap)
-		for _, field := range fields {
-			vrs := o.vrs[field]
-			eligibleDocIDsMap[field] = make(map[int]*roaring.Bitmap)
-			for index, vr := range vrs {
-				if vr.eligibleDocIDs != nil && len(vr.eligibleDocIDs) > 0 {
-					eligibleDocIDsMap[field][index] = vr.getEligibleDocIDs()
-				}
-			}
-		}
-	}
-
 	defer o.invokeSearcherEndCallback()
 
 	wg := sync.WaitGroup{}
@@ -100,8 +78,7 @@ func (o *OptimizeVR) Finish() error {
 					<-semaphore // Release the semaphore slot
 					wg.Done()
 				}()
-				for _, field := range fields {
-					vrs := o.vrs[field]
+				for field, vrs := range o.vrs {
 					vecIndex, err := segment.InterpretVectorIndex(field,
 						o.requiresFiltering, origSeg.deleted)
 					if err != nil {
@@ -114,34 +91,19 @@ func (o *OptimizeVR) Finish() error {
 					// update the vector index size as a meta value in the segment snapshot
 					vectorIndexSize := vecIndex.Size()
 					origSeg.cachedMeta.updateMeta(field, vectorIndexSize)
-					for vrIdx, vr := range vrs {
+					for _, vr := range vrs {
 						var pl segment_api.VecPostingsList
 						var err error
 
 						// for each VR, populate postings list and iterators
 						// by passing the obtained vector index and getting similar vectors.
 
-						// Only applies to filtered kNN.
-						if vr.eligibleDocIDs != nil && len(vr.eligibleDocIDs) > 0 {
-							eligibleVectorInternalIDs := eligibleDocIDsMap[field][vrIdx]
-							eligibleVectorInternalIDsClone := eligibleVectorInternalIDs.Clone()
-							if snapshotGlobalDocNums != nil {
-								// Only the eligible documents belonging to this segment
-								// will get filtered out.
-								// There is no way to determine which doc belongs to which segment
-								eligibleVectorInternalIDsClone.And(snapshotGlobalDocNums[index])
-							}
-
-							eligibleLocalDocNums := make([]uint64, 0)
-							// get the (segment-)local document numbers
-							for _, docNum := range eligibleVectorInternalIDsClone.ToArray() {
-								localDocNum := o.snapshot.localDocNumFromGlobal(index,
-									uint64(docNum))
-								eligibleLocalDocNums = append(eligibleLocalDocNums, localDocNum)
-							}
-
+						// check if the vector reader is configured to use a pre-filter
+						// to filter out ineligible documents before performing
+						// kNN search.
+						if vr.eligibleSelector != nil {
 							pl, err = vecIndex.SearchWithFilter(vr.vector, vr.k,
-								eligibleLocalDocNums, vr.searchParams)
+								vr.eligibleSelector.SegmentEligibleDocs(index), vr.searchParams)
 						} else {
 							pl, err = vecIndex.Search(vr.vector, vr.k, vr.searchParams)
 						}
@@ -195,7 +157,7 @@ func (s *IndexSnapshotVectorReader) VectorOptimize(ctx context.Context,
 	}
 	o.ctx = ctx
 	if !o.requiresFiltering {
-		o.requiresFiltering = len(s.eligibleDocIDs) > 0
+		o.requiresFiltering = s.eligibleSelector != nil
 	}
 
 	if o.snapshot != s.snapshot {

index/scorch/snapshot_index.go

@@ -473,31 +473,7 @@ func (is *IndexSnapshot) segmentIndexAndLocalDocNumFromGlobal(docNum uint64) (in
 			return is.offsets[x] > docNum
 		}) - 1
 
-	localDocNum := is.localDocNumFromGlobal(segmentIndex, docNum)
-	return int(segmentIndex), localDocNum
-}
-
-// This function returns the local docnum, given the segment index and global docnum
-func (is *IndexSnapshot) localDocNumFromGlobal(segmentIndex int, docNum uint64) uint64 {
-	return docNum - is.offsets[segmentIndex]
-}
-
-// Function to return a mapping of the segment index to the live global	doc nums
-// in the segment of the specified index snapshot.
-func (is *IndexSnapshot) globalDocNums() map[int]*roaring.Bitmap {
-	if len(is.segment) == 0 {
-		return nil
-	}
-
-	segmentIndexGlobalDocNums := make(map[int]*roaring.Bitmap)
-
-	for i := range is.segment {
-		segmentIndexGlobalDocNums[i] = roaring.NewBitmap()
-		for _, localDocNum := range is.segment[i].DocNumbersLive().ToArray() {
-			segmentIndexGlobalDocNums[i].Add(localDocNum + uint32(is.offsets[i]))
-		}
-	}
-	return segmentIndexGlobalDocNums
+	return int(segmentIndex), docNum - is.offsets[segmentIndex]
 }
 
 func (is *IndexSnapshot) ExternalID(id index.IndexInternalID) (string, error) {
@@ -518,6 +494,15 @@ func (is *IndexSnapshot) ExternalID(id index.IndexInternalID) (string, error) {
 	return string(v), nil
 }
 
+func (is *IndexSnapshot) segmentIndexAndLocalDocNum(id index.IndexInternalID) (int, uint64, error) {
+	docNum, err := docInternalToNumber(id)
+	if err != nil {
+		return 0, 0, err
+	}
+	segIdx, localDocNum := is.segmentIndexAndLocalDocNumFromGlobal(docNum)
+	return segIdx, localDocNum, nil
+}
+
 func (is *IndexSnapshot) InternalID(id string) (rv index.IndexInternalID, err error) {
 	// FIXME could be done more efficiently directly, but reusing for simplicity
 	tfr, err := is.TermFieldReader(nil, []byte(id), "_id", false, false, false)

index/scorch/snapshot_index_vr.go

@@ -24,7 +24,6 @@ import (
 	"fmt"
 	"reflect"
 
-	"github.com/RoaringBitmap/roaring"
 	"github.com/blevesearch/bleve/v2/size"
 	index "github.com/blevesearch/bleve_index_api"
 	segment_api "github.com/blevesearch/scorch_segment_api/v2"
@@ -51,32 +50,8 @@ type IndexSnapshotVectorReader struct {
 	currID        index.IndexInternalID
 	ctx           context.Context
 
-	searchParams json.RawMessage
-
-	// The following fields are only applicable for vector readers which will
-	// process pre-filtered kNN queries.
-	eligibleDocIDs []index.IndexInternalID
-}
-
-// Function to convert the internal IDs of the eligible documents to a type suitable
-// for addition to a bitmap.
-// Useful to have the eligible doc IDs in a bitmap to leverage the fast intersection
-// (AND) operations. Eg. finding the eligible doc IDs present in a segment.
-func (i *IndexSnapshotVectorReader) getEligibleDocIDs() *roaring.Bitmap {
-	res := roaring.NewBitmap()
-	if len(i.eligibleDocIDs) > 0 {
-		internalDocIDs := make([]uint32, 0, len(i.eligibleDocIDs))
-		// converts the doc IDs to uint32 and returns
-		for _, eligibleDocInternalID := range i.eligibleDocIDs {
-			internalDocID, err := docInternalToNumber(index.IndexInternalID(eligibleDocInternalID))
-			if err != nil {
-				continue
-			}
-			internalDocIDs = append(internalDocIDs, uint32(internalDocID))
-		}
-		res.AddMany(internalDocIDs)
-	}
-	return res
+	searchParams     json.RawMessage
+	eligibleSelector index.EligibleDocumentSelector
 }
 
 func (i *IndexSnapshotVectorReader) Size() int {
@@ -134,17 +109,8 @@ func (i *IndexSnapshotVectorReader) Advance(ID index.IndexInternalID,
 	preAlloced *index.VectorDoc) (*index.VectorDoc, error) {
 
 	if i.currPosting != nil && bytes.Compare(i.currID, ID) >= 0 {
-		var i2 index.VectorReader
-		var err error
-
-		if len(i.eligibleDocIDs) > 0 {
-			i2, err = i.snapshot.VectorReaderWithFilter(i.ctx, i.vector, i.field,
-				i.k, i.searchParams, i.eligibleDocIDs)
-		} else {
-			i2, err = i.snapshot.VectorReader(i.ctx, i.vector, i.field, i.k,
-				i.searchParams)
-		}
-
+		i2, err := i.snapshot.VectorReader(i.ctx, i.vector, i.field, i.k,
+			i.searchParams, i.eligibleSelector)
 		if err != nil {
 			return nil, err
 		}

index/scorch/snapshot_vector_index.go

@@ -20,21 +20,23 @@ package scorch
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 
 	index "github.com/blevesearch/bleve_index_api"
 	segment_api "github.com/blevesearch/scorch_segment_api/v2"
 )
 
 func (is *IndexSnapshot) VectorReader(ctx context.Context, vector []float32,
-	field string, k int64, searchParams json.RawMessage) (
+	field string, k int64, searchParams json.RawMessage,
+	eligibleSelector index.EligibleDocumentSelector) (
 	index.VectorReader, error) {
-
 	rv := &IndexSnapshotVectorReader{
-		vector:       vector,
-		field:        field,
-		k:            k,
-		snapshot:     is,
-		searchParams: searchParams,
+		vector:           vector,
+		field:            field,
+		k:                k,
+		snapshot:         is,
+		searchParams:     searchParams,
+		eligibleSelector: eligibleSelector,
 	}
 
 	if rv.postings == nil {
@@ -43,34 +45,41 @@ func (is *IndexSnapshot) VectorReader(ctx context.Context, vector []float32,
 	if rv.iterators == nil {
 		rv.iterators = make([]segment_api.VecPostingsIterator, len(is.segment))
 	}
-
 	// initialize postings and iterators within the OptimizeVR's Finish()
-
 	return rv, nil
 }
 
-func (is *IndexSnapshot) VectorReaderWithFilter(ctx context.Context, vector []float32,
-	field string, k int64, searchParams json.RawMessage,
-	filterIDs []index.IndexInternalID) (
-	index.VectorReader, error) {
+// eligibleDocumentSelector is used to filter out documents that are eligible for
+// the KNN search from a pre-filter query.
+type eligibleDocumentSelector struct {
+	// segment ID -> segment local doc nums
+	eligibleDocNums map[int][]uint64
+	is              *IndexSnapshot
+}
 
-	rv := &IndexSnapshotVectorReader{
-		vector:         vector,
-		field:          field,
-		k:              k,
-		snapshot:       is,
-		searchParams:   searchParams,
-		eligibleDocIDs: filterIDs,
-	}
+// SegmentEligibleDocs returns the list of eligible local doc numbers for the given segment.
+func (eds *eligibleDocumentSelector) SegmentEligibleDocs(segmentID int) []uint64 {
+	return eds.eligibleDocNums[segmentID]
+}
 
-	if rv.postings == nil {
-		rv.postings = make([]segment_api.VecPostingsList, len(is.segment))
+// AddEligibleDocumentMatch adds a document match to the list of eligible documents.
+func (eds *eligibleDocumentSelector) AddEligibleDocumentMatch(id index.IndexInternalID) error {
+	if eds.is == nil {
+		return fmt.Errorf("eligibleDocumentSelector is not initialized with IndexSnapshot")
 	}
-	if rv.iterators == nil {
-		rv.iterators = make([]segment_api.VecPostingsIterator, len(is.segment))
+	// Get the segment number and the local doc number for this document.
+	segIdx, docNum, err := eds.is.segmentIndexAndLocalDocNum(id)
+	if err != nil {
+		return err
 	}
+	// Add the local doc number to the list of eligible doc numbers for this segment.
+	eds.eligibleDocNums[segIdx] = append(eds.eligibleDocNums[segIdx], docNum)
+	return nil
+}
 
-	// initialize postings and iterators within the OptimizeVR's Finish()
-
-	return rv, nil
+func (is *IndexSnapshot) NewEligibleDocumentSelector() index.EligibleDocumentSelector {
+	return &eligibleDocumentSelector{
+		eligibleDocNums: map[int][]uint64{},
+		is:              is,
+	}
 }

search.go

@@ -589,3 +589,13 @@ func (r *SearchRequest) SortFunc() func(data sort.Interface) {
 
 	return sort.Sort
 }
+
+func isMatchNoneQuery(q query.Query) bool {
+	_, ok := q.(*query.MatchNoneQuery)
+	return ok
+}
+
+func isMatchAllQuery(q query.Query) bool {
+	_, ok := q.(*query.MatchAllQuery)
+	return ok
+}