MB-62182: golang bindings for fastmerge

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) {

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)
+}