upstream-bleve/index/scorch/optimize_knn.go

//  Copyright (c) 2023 Couchbase, Inc.
//
// 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.

//go:build vectors
// +build vectors

package scorch

import (
	"context"
	"fmt"
	"sync"

	"github.com/blevesearch/bleve/v2/search"
	index "github.com/blevesearch/bleve_index_api"
	segment_api "github.com/blevesearch/scorch_segment_api/v2"
)

type OptimizeVR struct {
	ctx       context.Context
	snapshot  *IndexSnapshot
	totalCost uint64
	// maps field to vector readers
	vrs map[string][]*IndexSnapshotVectorReader
}

func (o *OptimizeVR) invokeSearcherEndCallback() {
	if o.ctx != nil {
		if cb := o.ctx.Value(search.SearcherEndCallbackKey); cb != nil {
			if cbF, ok := cb.(search.SearcherEndCallbackFn); ok {
				if o.totalCost > 0 {
					// notify the callback that the searcher creation etc. is finished
					// and report back the total cost for it to track and take actions
					// appropriately.
					_ = cbF(o.totalCost)
				}
			}
		}
	}
}

// search runs the configured kNN searches for every vector
// reader against a single segment, populating the per-segment postings and
// iterators on each reader.
func (o *OptimizeVR) search(segID int) error {
	snapshot := o.snapshot.segment[segID]
	vecSeg, ok := snapshot.segment.(segment_api.VectorSegment)
	if !ok {
		return nil
	}
	meta := snapshot.cachedMeta
	// for each field, get the vector index --> invoke the zap func.
	// for each VR, populate postings list and iterators
	// by passing the obtained vector index and getting similar vectors.
	for field, vrs := range o.vrs {
		// Early exit if the field is supposed to be completely deleted or
		// if it's index data has been deleted
		if info, ok := o.snapshot.updatedFields[field]; ok && (info.Deleted || info.Index) {
			continue
		}
		vecIndex, err := vecSeg.InterpretVectorIndex(field, snapshot.deleted)
		if err != nil {
			return err
		}
		// update the vector index size as a meta value in the segment snapshot
		// if not already present
		if !meta.contains(field) {
			meta.store(field, vecIndex.Size())
		}
		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.

			// 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,
					vr.eligibleSelector.SegmentEligibleDocuments(segID),
					vr.searchParams,
				)
			} else {
				pl, err = vecIndex.Search(
					vr.vector,
					vr.k,
					vr.searchParams,
				)
			}
			if err != nil {
				vecIndex.Close()
				return err
			}
			// postings and iterators are already alloc'ed when
			// IndexSnapshotVectorReader is created, here we are just
			// populating them with the obtained postings list.
			vr.postings[segID] = pl
			vr.iterators[segID] = pl.Iterator(vr.iterators[segID])
		}
		go vecIndex.Close()
	}
	return nil
}

func (o *OptimizeVR) Finish() error {
	defer o.invokeSearcherEndCallback()
	numSegments := len(o.snapshot.segment)

	var wg sync.WaitGroup
	wg.Add(numSegments)
	var errM sync.Mutex
	var searchErr error
	// launch goroutines to search the vector index for each segment
	for i := 0; i < numSegments; i++ {
		go func(segID int) {
			defer wg.Done()
			if err := o.search(segID); err != nil {
				errM.Lock()
				searchErr = err
				errM.Unlock()
			}
		}(i)
	}
	// wait until all the launched goroutines finish and collect errors if any
	wg.Wait()
	return searchErr
}

func (s *IndexSnapshotVectorReader) VectorOptimize(ctx context.Context,
	octx index.VectorOptimizableContext,
) (index.VectorOptimizableContext, error) {
	if s.snapshot.parent.segPlugin.Version() < VectorSearchSupportedSegmentVersion {
		return nil, fmt.Errorf("vector search not supported for this index, "+
			"index's segment version %v, supported segment version for vector search %v",
			s.snapshot.parent.segPlugin.Version(), VectorSearchSupportedSegmentVersion)
	}

	if octx == nil {
		octx = &OptimizeVR{
			snapshot: s.snapshot,
			vrs:      make(map[string][]*IndexSnapshotVectorReader),
		}
	}

	o, ok := octx.(*OptimizeVR)
	if !ok {
		return octx, nil
	}
	o.ctx = ctx

	if o.snapshot != s.snapshot {
		o.invokeSearcherEndCallback()
		return nil, fmt.Errorf("tried to optimize KNN across different snapshots")
	}

	// for every searcher creation, consult the segment snapshot to see
	// what's the vector index size and since you're anyways going
	// to use this vector index to perform the search etc. as part of the Finish()
	// perform a check as to whether we allow the searcher creation (the downstream)
	// Finish() logic to even occur or not.
	var sumVectorIndexSize uint64
	for _, seg := range o.snapshot.segment {
		if vecIndexSize, ok := seg.cachedMeta.load(s.field); ok {
			sumVectorIndexSize += vecIndexSize.(uint64)
		}
	}

	if o.ctx != nil {
		if cb := o.ctx.Value(search.SearcherStartCallbackKey); cb != nil {
			if cbF, ok := cb.(search.SearcherStartCallbackFn); ok {
				err := cbF(sumVectorIndexSize)
				if err != nil {
					// it's important to invoke the end callback at this point since
					// if the earlier searchers of this optimize struct were successful
					// the cost corresponding to it would be incremented and if the
					// current searcher fails the check then we end up erroring out
					// the overall optimized searcher creation, the cost needs to be
					// handled appropriately.
					o.invokeSearcherEndCallback()
					return nil, err
				}
			}
		}
	}

	// total cost is essentially the sum of the vector indexes' size across all the
	// searchers - all of them end up reading and maintaining a vector index.
	// misacconting this value would end up calling the "end" callback with a value
	// not equal to the value passed to "start" callback.
	o.totalCost += sumVectorIndexSize
	o.vrs[s.field] = append(o.vrs[s.field], s)
	return o, nil
}