MB-62985 - Add support for binary quantised vectors. #329
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -23,6 +23,7 @@ import ( | |
| "math" | ||
| "reflect" | ||
|
|
||
| "container/heap" | ||
| "github.com/RoaringBitmap/roaring/v2" | ||
| "github.com/RoaringBitmap/roaring/v2/roaring64" | ||
| "github.com/bits-and-blooms/bitset" | ||
|
|
@@ -301,6 +302,31 @@ func (i *vectorIndexWrapper) Size() uint64 { | |
| return i.size() | ||
| } | ||
|
|
||
| // distanceID represents a distance-ID pair for heap operations | ||
| type distanceID struct { | ||
| distance float32 | ||
| id int64 | ||
| } | ||
|
|
||
| // maxHeap implements heap.Interface for distanceID | ||
| type maxHeap []*distanceID | ||
|
|
||
| func (h maxHeap) Len() int { return len(h) } | ||
| func (h maxHeap) Less(i, j int) bool { return h[i].distance > h[j].distance } | ||
| func (h maxHeap) Swap(i, j int) { h[i], h[j] = h[j], h[i] } | ||
|
|
||
| func (h *maxHeap) Push(x interface{}) { | ||
| *h = append(*h, x.(*distanceID)) | ||
| } | ||
|
|
||
| func (h *maxHeap) Pop() interface{} { | ||
| old := *h | ||
| n := len(old) | ||
| x := old[n-1] | ||
| *h = old[0 : n-1] | ||
| return x | ||
| } | ||
|
|
||
metonymic-smokey marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| // InterpretVectorIndex returns a construct of closures (vectorIndexWrapper) | ||
| // that will allow the caller to - | ||
| // (1) search within an attached vector index | ||
|
|
@@ -312,6 +338,7 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool | |
| segment.VectorIndex, error) { | ||
| // Params needed for the closures | ||
| var vecIndex *faiss.IndexImpl | ||
| var binaryIndex *faiss.IndexImpl | ||
| var vecDocIDMap map[int64]uint32 | ||
| var docVecIDMap map[uint32][]int64 | ||
| var vectorIDsToExclude []int64 | ||
|
|
@@ -354,14 +381,49 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool | |
| return rv, nil | ||
| } | ||
|
|
||
| if binaryIndex != nil { | ||
| binaryQueryVector := convertToBinary(qVector) | ||
| _, binIDs, err := binaryIndex.SearchBinaryWithoutIDs(binaryQueryVector, k*4, | ||
| vectorIDsToExclude, params) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| distances := make([]float32, k*4) | ||
| err = vecIndex.DistCompute(qVector, binIDs, int(k*4), distances) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| // Need to map distances to the original IDs to get the top K. | ||
| // Use a heap to keep track of the top K. | ||
| h := &maxHeap{} | ||
| heap.Init(h) | ||
| for i := 0; i < len(binIDs); i++ { | ||
| heap.Push(h, &distanceID{distance: distances[i], id: binIDs[i]}) | ||
| if h.Len() > int(k) { | ||
| heap.Pop(h) | ||
| } | ||
| } | ||
|
|
||
| // Pop the top K in reverse order to get them in ascending order | ||
| ids := make([]int64, k) | ||
| scores := make([]float32, k) | ||
| for i := int(k) - 1; i >= 0; i-- { | ||
| distanceID := heap.Pop(h).(*distanceID) | ||
| scores[i] = distanceID.distance | ||
| ids[i] = distanceID.id | ||
| } | ||
|
Comment on lines
+400
to
+416
Member
There was a problem hiding this comment. At this stage, we already have the complete set of ids and their corresponding distances available for Top-K selection. This scenario differs from the K-sized heap approach used in In our case, since the entire dataset is available upfront, using a heap becomes unnecessarily expensive—particularly in terms of space complexity O(K). This overhead can compound significantly with the number of segments in a highly distributed index, as the heap becomes temporary query-specific state for each segment. Given that, it would be more efficient to use a |
||
|
|
||
| addIDsToPostingsList(rv, ids, scores) | ||
| } else { | ||
| scores, ids, err := vecIndex.SearchWithoutIDs(qVector, k, | ||
| vectorIDsToExclude, params) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| addIDsToPostingsList(rv, ids, scores) | ||
| } | ||
| return rv, nil | ||
| }, | ||
| searchWithFilter: func(qVector []float32, k int64, | ||
|
|
@@ -542,14 +604,27 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool | |
| // 1. doc values(first 2 iterations) - adhering to the sections format. never | ||
| // valid values for vector section | ||
| // 2. index optimization type. | ||
| for i := 0; i < 2; i++ { | ||
| _, n := binary.Uvarint(sb.mem[pos : pos+binary.MaxVarintLen64]) | ||
| pos += n | ||
| } | ||
|
|
||
| // Determining if a binary index is needed. | ||
| indexOptimisationType, n := binary.Uvarint(sb.mem[pos : pos+binary.MaxVarintLen64]) | ||
| pos += n | ||
|
|
||
| loadBinaryIndex := isBinaryIndex(indexOptimisationType) | ||
|
|
||
| vecIndexes, vecDocIDMap, docVecIDMap, vectorIDsToExclude, err := | ||
| sb.vecIndexCache.loadOrCreate(fieldIDPlus1, sb.mem[pos:], requiresFiltering, | ||
| loadBinaryIndex, except) | ||
|
|
||
| if loadBinaryIndex { | ||
| vecIndex = vecIndexes[1] | ||
| binaryIndex = vecIndexes[0] | ||
| } else { | ||
| vecIndex = vecIndexes[0] | ||
| } | ||
|
|
||
| if vecIndex != nil { | ||
| vecIndexSize = vecIndex.Size() | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
take pointer to the cacheEntryReqs struct, don't copy the struct