github.com/weaviate/weaviate@v1.24.6/adapters/repos/db/vector/hnsw/search_with_max_dist.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package hnsw
    13  
    14  import (
    15  	"fmt"
    16  
    17  	"github.com/pkg/errors"
    18  	"github.com/weaviate/weaviate/adapters/repos/db/helpers"
    19  	"github.com/weaviate/weaviate/adapters/repos/db/priorityqueue"
    20  )
    21  
    22  func (h *hnsw) KnnSearchByVectorMaxDist(searchVec []float32, dist float32,
    23  	ef int, allowList helpers.AllowList,
    24  ) ([]uint64, error) {
    25  	entryPointID := h.entryPointID
    26  	entryPointDistance, ok, err := h.distBetweenNodeAndVec(entryPointID, searchVec)
    27  	if err != nil {
    28  		return nil, errors.Wrap(err, "knn search: distance between entrypoint and query node")
    29  	}
    30  
    31  	if !ok {
    32  		return nil, fmt.Errorf("entrypoint was deleted in the object store, " +
    33  			"it has been flagged for cleanup and should be fixed in the next cleanup cycle")
    34  	}
    35  
    36  	// stop at layer 1, not 0!
    37  	for level := h.currentMaximumLayer; level >= 1; level-- {
    38  		eps := priorityqueue.NewMin[any](1)
    39  		eps.Insert(entryPointID, entryPointDistance)
    40  		// ignore allowList on layers > 0
    41  		res, err := h.searchLayerByVector(searchVec, eps, 1, level, nil)
    42  		if err != nil {
    43  			return nil, errors.Wrapf(err, "knn search: search layer at level %d", level)
    44  		}
    45  		if res.Len() > 0 {
    46  			best := res.Pop()
    47  			entryPointID = best.ID
    48  			entryPointDistance = best.Dist
    49  		}
    50  
    51  		h.pools.pqResults.Put(res)
    52  	}
    53  
    54  	eps := priorityqueue.NewMin[any](1)
    55  	eps.Insert(entryPointID, entryPointDistance)
    56  	res, err := h.searchLayerByVector(searchVec, eps, ef, 0, allowList)
    57  	if err != nil {
    58  		return nil, errors.Wrapf(err, "knn search: search layer at level %d", 0)
    59  	}
    60  
    61  	all := make([]priorityqueue.Item[any], res.Len())
    62  	i := res.Len() - 1
    63  	for res.Len() > 0 {
    64  		all[i] = res.Pop()
    65  		i--
    66  	}
    67  
    68  	out := make([]uint64, len(all))
    69  	i = 0
    70  	for _, elem := range all {
    71  		if elem.Dist > dist {
    72  			break
    73  		}
    74  		out[i] = elem.ID
    75  		i++
    76  	}
    77  
    78  	h.pools.pqResults.Put(res)
    79  	return out[:i], nil
    80  }