diff --git a/index.go b/index.go index 9e6efa2..7903b2c 100644 --- a/index.go +++ b/index.go @@ -61,6 +61,7 @@ type Index interface { // Applicable only to IVF indexes: Returns the top k centroid cardinalities and // their vectors in chosen order (descending or ascending) ObtainKCentroidCardinalitiesFromIVFIndex(limit int, descending bool) ([]uint64, [][]float32, error) + Nlist() int // Search queries the index with the vectors in x. // Returns the IDs of the k nearest neighbors for each query vector and the @@ -101,6 +102,8 @@ type Index interface { Size() uint64 cPtr() *C.FaissIndex + + SetQuantizers(coarseQuantizer Index) error } type faissIndex struct { @@ -284,6 +287,14 @@ func getIndicesOfKCentroidCardinalities(cardinalities []C.size_t, k int, descend return indices[:k] } +func (idx *faissIndex) Nlist() int { + if !idx.IsIVFIndex() { + return 0 + } else { + fmt.Println("nlist", int(C.faiss_IndexIVF_nlist(idx.idx))) + } + return int(C.faiss_IndexIVF_nlist(idx.idx)) +} func (idx *faissIndex) SearchClustersFromIVFIndex(eligibleCentroidIDs []int64, centroidDis []float32, centroidsToProbe int, x []float32, k int64, include Selector, params json.RawMessage) ([]float32, []int64, error) { diff --git a/index_ivf.go b/index_ivf.go index 797aeaf..485f637 100644 --- a/index_ivf.go +++ b/index_ivf.go @@ -60,3 +60,26 @@ func (idx *IndexImpl) IVFParams() (nprobe, nlist int) { return int(C.faiss_IndexIVF_nprobe(ivfPtr)), int(C.faiss_IndexIVF_nlist(ivfPtr)) } + +func (idx *faissIndex) SetQuantizers(srcIndex Index) error { + ivfPtr := C.faiss_IndexIVF_cast(idx.idx) + if ivfPtr == nil { + return fmt.Errorf("index is not of ivf type") + } + + srcIndexPtr := srcIndex.cPtr() + if srcIndexPtr == nil { + return fmt.Errorf("coarse quantizer is not valid") + } + + err := C.faiss_Set_coarse_quantizers(ivfPtr, srcIndexPtr) + if err != 0 { + return fmt.Errorf("couldn't set the coarse quantizer") + } + + return nil +} + +func (idx *IndexImpl) SetQuantizers(srcIndex Index) error { + return idx.Index.SetQuantizers(srcIndex) +}