Skip to content

Commit

Permalink
avoiding type casts
Browse files Browse the repository at this point in the history
  • Loading branch information
@metonymic-smokey
metonymic-smokey committed Nov 4, 2024
1 parent e09d5ba commit 517ee06
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 26 deletions.
25 changes: 13 additions & 12 deletions faiss_vector_cache.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ func (vc *vectorIndexCache) Clear() {
// map. It's false otherwise.
func (vc *vectorIndexCache) loadOrCreate(fieldID uint16, mem []byte,
loadDocVecIDMap bool, except *roaring.Bitmap) (
index *faiss.IndexImpl, vecDocIDMap map[int64]uint32, docVecIDMap map[uint32][]int64,
index *faiss.IndexImpl, vecDocIDMap map[int64]uint32, docVecIDMap map[uint32][]uint32,
vecIDsToExclude []int64, err error) {
index, vecDocIDMap, docVecIDMap, vecIDsToExclude, err = vc.loadFromCache(
fieldID, loadDocVecIDMap, mem, except)
Expand All @@ -68,7 +68,7 @@ func (vc *vectorIndexCache) loadOrCreate(fieldID uint16, mem []byte,
// If not, it will create these and add them to the cache.
func (vc *vectorIndexCache) loadFromCache(fieldID uint16, loadDocVecIDMap bool,
mem []byte, except *roaring.Bitmap) (index *faiss.IndexImpl, vecDocIDMap map[int64]uint32,
docVecIDMap map[uint32][]int64, vecIDsToExclude []int64, err error) {
docVecIDMap map[uint32][]uint32, vecIDsToExclude []int64, err error) {

vc.m.RLock()

Expand Down Expand Up @@ -98,16 +98,16 @@ func (vc *vectorIndexCache) loadFromCache(fieldID uint16, loadDocVecIDMap bool,
return vc.createAndCacheLOCKED(fieldID, mem, loadDocVecIDMap, except)
}

func (vc *vectorIndexCache) addDocVecIDMapToCacheLOCKED(ce *cacheEntry) map[uint32][]int64 {
func (vc *vectorIndexCache) addDocVecIDMapToCacheLOCKED(ce *cacheEntry) map[uint32][]uint32 {
// Handle concurrent accesses (to avoid unnecessary work) by adding a
// check within the write lock here.
if ce.docVecIDMap != nil {
return ce.docVecIDMap
}

docVecIDMap := make(map[uint32][]int64)
docVecIDMap := make(map[uint32][]uint32)
for vecID, docID := range ce.vecDocIDMap {
docVecIDMap[docID] = append(docVecIDMap[docID], vecID)
docVecIDMap[docID] = append(docVecIDMap[docID], uint32(vecID))
}

ce.docVecIDMap = docVecIDMap
Expand All @@ -118,7 +118,7 @@ func (vc *vectorIndexCache) addDocVecIDMapToCacheLOCKED(ce *cacheEntry) map[uint
func (vc *vectorIndexCache) createAndCacheLOCKED(fieldID uint16, mem []byte,
loadDocVecIDMap bool, except *roaring.Bitmap) (
index *faiss.IndexImpl, vecDocIDMap map[int64]uint32,
docVecIDMap map[uint32][]int64, vecIDsToExclude []int64, err error) {
docVecIDMap map[uint32][]uint32, vecIDsToExclude []int64, err error) {

// Handle concurrent accesses (to avoid unnecessary work) by adding a
// check within the write lock here.
Expand All @@ -141,7 +141,7 @@ func (vc *vectorIndexCache) createAndCacheLOCKED(fieldID uint16, mem []byte,

vecDocIDMap = make(map[int64]uint32, numVecs)
if loadDocVecIDMap {
docVecIDMap = make(map[uint32][]int64, numVecs)
docVecIDMap = make(map[uint32][]uint32, numVecs)
}
isExceptNotEmpty := except != nil && !except.IsEmpty()
for i := 0; i < int(numVecs); i++ {
Expand All @@ -157,7 +157,8 @@ func (vc *vectorIndexCache) createAndCacheLOCKED(fieldID uint16, mem []byte,
}
vecDocIDMap[vecID] = docIDUint32
if loadDocVecIDMap {
docVecIDMap[docIDUint32] = append(docVecIDMap[docIDUint32], vecID)
docVecIDMap[uint32(docID)] = append(docVecIDMap[uint32(docID)],
uint32(vecID))
}
}

Expand All @@ -175,7 +176,7 @@ func (vc *vectorIndexCache) createAndCacheLOCKED(fieldID uint16, mem []byte,

func (vc *vectorIndexCache) insertLOCKED(fieldIDPlus1 uint16,
index *faiss.IndexImpl, vecDocIDMap map[int64]uint32, loadDocVecIDMap bool,
docVecIDMap map[uint32][]int64) {
docVecIDMap map[uint32][]uint32) {
// the first time we've hit the cache, try to spawn a monitoring routine
// which will reconcile the moving averages for all the fields being hit
if len(vc.cache) == 0 {
Expand Down Expand Up @@ -284,7 +285,7 @@ func (e *ewma) add(val uint64) {
// -----------------------------------------------------------------------------

func createCacheEntry(index *faiss.IndexImpl, vecDocIDMap map[int64]uint32,
loadDocVecIDMap bool, docVecIDMap map[uint32][]int64, alpha float64) *cacheEntry {
loadDocVecIDMap bool, docVecIDMap map[uint32][]uint32, alpha float64) *cacheEntry {
ce := &cacheEntry{
index: index,
vecDocIDMap: vecDocIDMap,
Expand All @@ -310,7 +311,7 @@ type cacheEntry struct {

index *faiss.IndexImpl
vecDocIDMap map[int64]uint32
docVecIDMap map[uint32][]int64
docVecIDMap map[uint32][]uint32
}

func (ce *cacheEntry) incHit() {
Expand All @@ -325,7 +326,7 @@ func (ce *cacheEntry) decRef() {
atomic.AddInt64(&ce.refs, -1)
}

func (ce *cacheEntry) load() (*faiss.IndexImpl, map[int64]uint32, map[uint32][]int64) {
func (ce *cacheEntry) load() (*faiss.IndexImpl, map[int64]uint32, map[uint32][]uint32) {
ce.incHit()
ce.addRef()
return ce.index, ce.vecDocIDMap, ce.docVecIDMap
Expand Down
27 changes: 13 additions & 14 deletions faiss_vector_posting.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -310,7 +310,7 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool
// Params needed for the closures
var vecIndex *faiss.IndexImpl
var vecDocIDMap map[int64]uint32
var docVecIDMap map[uint32][]int64
var docVecIDMap map[uint32][]uint32
var vectorIDsToExclude []int64
var fieldIDPlus1 uint16
var vecIndexSize uint64
Expand Down Expand Up @@ -400,7 +400,9 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool
// considered for the search
vectorIDsToInclude := make([]int64, 0, len(eligibleDocIDs))
for _, id := range eligibleDocIDs {
vectorIDsToInclude = append(vectorIDsToInclude, docVecIDMap[uint32(id)]...)
for _, vecID := range docVecIDMap[uint32(id)] {
vectorIDsToInclude = append(vectorIDsToInclude, int64(vecID))
}
}

// Retrieve the mapping of centroid IDs to vectors within
Expand Down Expand Up @@ -446,7 +448,7 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool
ineligibleVecIDsBitmap := roaring.NewBitmap()
eligibleDocIDsMap := make(map[uint64]struct{})
for _, eligibleDocID := range eligibleDocIDs {
eligibleDocIDsMap[(eligibleDocID)] = struct{}{}
eligibleDocIDsMap[uint64(eligibleDocID)] = struct{}{}
}

ineligibleVectorIDs := make([]int64, 0, len(vecDocIDMap)-
Expand All @@ -456,7 +458,7 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool
if _, exists := eligibleDocIDsMap[uint64(docID)]; !exists {
for _, vecID := range vecIDs {
ineligibleVecIDsBitmap.Add(uint32(vecID))
ineligibleVectorIDs = append(ineligibleVectorIDs, vecID)
ineligibleVectorIDs = append(ineligibleVectorIDs, int64(vecID))
}
}
}
Expand Down Expand Up @@ -487,20 +489,17 @@ func (sb *SegmentBase) InterpretVectorIndex(field string, requiresFiltering bool
// Eg. docID d1 -> vecID v1, for the first case
// d1 -> {v1,v2}, for the second case.
eligibleVecIDsBitmap := roaring.NewBitmap()
vecIDsUint32 := make([]uint32, 0)
vecIDs := make([]uint32, 0)
for _, eligibleDocID := range eligibleDocIDs {
vecIDs := docVecIDMap[uint32(eligibleDocID)]
for _, vecID := range vecIDs {
vecIDsUint32 = append(vecIDsUint32, uint32(vecID))
}
vecIDs = append(vecIDs, docVecIDMap[uint32(eligibleDocID)]...)
}
eligibleVecIDsBitmap.AddMany(vecIDsUint32)
for centroidID, vecIDs := range centroidVecIDMap {
vecIDs.And(eligibleVecIDsBitmap)
if !vecIDs.IsEmpty() {
eligibleVecIDsBitmap.AddMany(vecIDs)
for centroidID, vecIDsBM := range centroidVecIDMap {
vecIDsBM.And(eligibleVecIDsBitmap)
if !vecIDsBM.IsEmpty() {
// The mapping is now reduced to those vectors which
// are also eligible docs for the filter query.
centroidVecIDMap[centroidID] = vecIDs
centroidVecIDMap[centroidID] = vecIDsBM
eligibleCentroidIDs = append(eligibleCentroidIDs, centroidID)
} else {
// don't consider clusters with no eligible IDs.
Expand Down

0 comments on commit 517ee06

Please sign in to comment.