MB-66396: IndexSnapshot API to retrieve highest cardinality centroids

Author: abhinavdangeti

go.mod

@@ -9,12 +9,12 @@ require (
 	github.com/bits-and-blooms/bitset v1.22.0
 	github.com/blevesearch/bleve_index_api v1.2.8
 	github.com/blevesearch/geo v0.2.4
-	github.com/blevesearch/go-faiss v1.0.25
+	github.com/blevesearch/go-faiss v1.0.26-0.20250808221324-a1d3051837ea
 	github.com/blevesearch/go-metrics v0.0.0-20201227073835-cf1acfcdf475
 	github.com/blevesearch/go-porterstemmer v1.0.3
 	github.com/blevesearch/goleveldb v1.0.1
 	github.com/blevesearch/gtreap v0.1.1
-	github.com/blevesearch/scorch_segment_api/v2 v2.3.10
+	github.com/blevesearch/scorch_segment_api/v2 v2.3.11-0.20250812224933-05c6687cb15b
 	github.com/blevesearch/segment v0.9.1
 	github.com/blevesearch/snowball v0.6.1
 	github.com/blevesearch/snowballstem v0.9.0
@@ -26,7 +26,7 @@ require (
 	github.com/blevesearch/zapx/v13 v13.4.2
 	github.com/blevesearch/zapx/v14 v14.4.2
 	github.com/blevesearch/zapx/v15 v15.4.2
-	github.com/blevesearch/zapx/v16 v16.2.4
+	github.com/blevesearch/zapx/v16 v16.2.5-0.20250812225830-3e3f5b307941
 	github.com/couchbase/moss v0.2.0
 	github.com/spf13/cobra v1.8.1
 	go.etcd.io/bbolt v1.4.0

go.sum

@@ -7,8 +7,8 @@ github.com/blevesearch/bleve_index_api v1.2.8 h1:Y98Pu5/MdlkRyLM0qDHostYo7i+Vv1c
 github.com/blevesearch/bleve_index_api v1.2.8/go.mod h1:rKQDl4u51uwafZxFrPD1R7xFOwKnzZW7s/LSeK4lgo0=
 github.com/blevesearch/geo v0.2.4 h1:ECIGQhw+QALCZaDcogRTNSJYQXRtC8/m8IKiA706cqk=
 github.com/blevesearch/geo v0.2.4/go.mod h1:K56Q33AzXt2YExVHGObtmRSFYZKYGv0JEN5mdacJJR8=
-github.com/blevesearch/go-faiss v1.0.25 h1:lel1rkOUGbT1CJ0YgzKwC7k+XH0XVBHnCVWahdCXk4U=
-github.com/blevesearch/go-faiss v1.0.25/go.mod h1:OMGQwOaRRYxrmeNdMrXJPvVx8gBnvE5RYrr0BahNnkk=
+github.com/blevesearch/go-faiss v1.0.26-0.20250808221324-a1d3051837ea h1:xiz5c7fyBrgmt/TVWRM+8LVNFQS7hDjlyfAjliumDWM=
+github.com/blevesearch/go-faiss v1.0.26-0.20250808221324-a1d3051837ea/go.mod h1:OMGQwOaRRYxrmeNdMrXJPvVx8gBnvE5RYrr0BahNnkk=
 github.com/blevesearch/go-metrics v0.0.0-20201227073835-cf1acfcdf475 h1:kDy+zgJFJJoJYBvdfBSiZYBbdsUL0XcjHYWezpQBGPA=
 github.com/blevesearch/go-metrics v0.0.0-20201227073835-cf1acfcdf475/go.mod h1:9eJDeqxJ3E7WnLebQUlPD7ZjSce7AnDb9vjGmMCbD0A=
 github.com/blevesearch/go-porterstemmer v1.0.3 h1:GtmsqID0aZdCSNiY8SkuPJ12pD4jI+DdXTAn4YRcHCo=
@@ -20,8 +20,8 @@ github.com/blevesearch/gtreap v0.1.1/go.mod h1:QaQyDRAT51sotthUWAH4Sj08awFSSWzgY
 github.com/blevesearch/mmap-go v1.0.2/go.mod h1:ol2qBqYaOUsGdm7aRMRrYGgPvnwLe6Y+7LMvAB5IbSA=
 github.com/blevesearch/mmap-go v1.0.4 h1:OVhDhT5B/M1HNPpYPBKIEJaD0F3Si+CrEKULGCDPWmc=
 github.com/blevesearch/mmap-go v1.0.4/go.mod h1:EWmEAOmdAS9z/pi/+Toxu99DnsbhG1TIxUoRmJw/pSs=
-github.com/blevesearch/scorch_segment_api/v2 v2.3.10 h1:Yqk0XD1mE0fDZAJXTjawJ8If/85JxnLd8v5vG/jWE/s=
-github.com/blevesearch/scorch_segment_api/v2 v2.3.10/go.mod h1:Z3e6ChN3qyN35yaQpl00MfI5s8AxUJbpTR/DL8QOQ+8=
+github.com/blevesearch/scorch_segment_api/v2 v2.3.11-0.20250812224933-05c6687cb15b h1:dj5AYwlxpsqRXTmnO4nqUBiO5RvBdPdN5tVPJscZpTs=
+github.com/blevesearch/scorch_segment_api/v2 v2.3.11-0.20250812224933-05c6687cb15b/go.mod h1:Z3e6ChN3qyN35yaQpl00MfI5s8AxUJbpTR/DL8QOQ+8=
 github.com/blevesearch/segment v0.9.1 h1:+dThDy+Lvgj5JMxhmOVlgFfkUtZV2kw49xax4+jTfSU=
 github.com/blevesearch/segment v0.9.1/go.mod h1:zN21iLm7+GnBHWTao9I+Au/7MBiL8pPFtJBJTsk6kQw=
 github.com/blevesearch/snowball v0.6.1 h1:cDYjn/NCH+wwt2UdehaLpr2e4BwLIjN4V/TdLsL+B5A=
@@ -44,8 +44,8 @@ github.com/blevesearch/zapx/v14 v14.4.2 h1:2SGHakVKd+TrtEqpfeq8X+So5PShQ5nW6GNxT
 github.com/blevesearch/zapx/v14 v14.4.2/go.mod h1:rz0XNb/OZSMjNorufDGSpFpjoFKhXmppH9Hi7a877D8=
 github.com/blevesearch/zapx/v15 v15.4.2 h1:sWxpDE0QQOTjyxYbAVjt3+0ieu8NCE0fDRaFxEsp31k=
 github.com/blevesearch/zapx/v15 v15.4.2/go.mod h1:1pssev/59FsuWcgSnTa0OeEpOzmhtmr/0/11H0Z8+Nw=
-github.com/blevesearch/zapx/v16 v16.2.4 h1:tGgfvleXTAkwsD5mEzgM3zCS/7pgocTCnO1oyAUjlww=
-github.com/blevesearch/zapx/v16 v16.2.4/go.mod h1:Rti/REtuuMmzwsI8/C/qIzRaEoSK/wiFYw5e5ctUKKs=
+github.com/blevesearch/zapx/v16 v16.2.5-0.20250812225830-3e3f5b307941 h1:U1IBoeRBOtwrmyt3PrYJMTIZdsecdpS1epFfKdycqww=
+github.com/blevesearch/zapx/v16 v16.2.5-0.20250812225830-3e3f5b307941/go.mod h1:VWpjVlGQxtXK5dUEGqum0RhloAiHNYQYWKzrP+zFi68=
 github.com/couchbase/ghistogram v0.1.0 h1:b95QcQTCzjTUocDXp/uMgSNQi8oj1tGwnJ4bODWZnps=
 github.com/couchbase/ghistogram v0.1.0/go.mod h1:s1Jhy76zqfEecpNWJfWUiKZookAFaiGOEoyzgHt9i7k=
 github.com/couchbase/moss v0.2.0 h1:VCYrMzFwEryyhRSeI+/b3tRBSeTpi/8gn5Kf6dxqn+o=

index/scorch/scorch_knn_test.go

@@ -0,0 +1,134 @@
+//  Copyright (c) 2025 Couchbase, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// 		http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//go:build vectors
+// +build vectors
+
+package scorch
+
+import (
+	"fmt"
+	"log"
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/blevesearch/bleve/v2/document"
+	"github.com/blevesearch/bleve/v2/mapping"
+	index "github.com/blevesearch/bleve_index_api"
+)
+
+func TestIndexSnapshotHighestCardinalityCentroids(t *testing.T) {
+	cfg := CreateConfig("TestIndexSnapshotHighestCardinalityCentroids")
+	err := InitTest(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := DestroyTest(cfg)
+		if err != nil {
+			t.Log(err)
+		}
+	}()
+
+	testConfig := cfg
+	mp := mapping.NewIndexMapping()
+
+	vectorDims := 5
+
+	vecMapping := mapping.NewVectorFieldMapping()
+	vecMapping.Dims = vectorDims
+	vecMapping.Similarity = index.CosineSimilarity
+
+	docMapping := mapping.NewDocumentStaticMapping()
+	docMapping.AddFieldMappingsAt("vec", vecMapping)
+	mp.DefaultMapping = docMapping
+
+	analysisQueue := index.NewAnalysisQueue(1)
+	idx, err := NewScorch("storeName", testConfig, analysisQueue)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	err = idx.Open()
+	if err != nil {
+		t.Errorf("error opening index: %v", err)
+	}
+	defer func() {
+		err := idx.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	rand.Seed(time.Now().UnixNano())
+	min, max := float32(-10.0), float32(10.0)
+	genRandomVector := func() []float32 {
+		vec := make([]float32, vectorDims)
+		for i := range vec {
+			vec[i] = min + rand.Float32()*(max-min)
+		}
+		return vec
+	}
+
+	var batch *index.Batch
+	for i := 1; i <= 20000; i++ {
+		doc := document.NewDocument(fmt.Sprintf("doc-%d", i))
+		err = mp.MapDocument(doc, map[string]interface{}{
+			"vec": genRandomVector(),
+		})
+		if err != nil {
+			t.Errorf("error mapping doc: %v", err)
+		}
+		if batch == nil {
+			batch = index.NewBatch()
+		}
+		batch.Update(doc)
+
+		if i%200 == 0 {
+			err = idx.Batch(batch)
+			if err != nil {
+				t.Errorf("Error adding batch to index: %v", err)
+			}
+			batch = nil
+		}
+	}
+
+	if batch != nil {
+		// In case doc count is not a multiple of 200, we need to add the final batch
+		err = idx.Batch(batch)
+		if err != nil {
+			t.Errorf("Error adding final batch to index: %v", err)
+		}
+	}
+
+	reader, err := idx.Reader()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() {
+		err := reader.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	limit := 5
+	if snap, ok := reader.(*IndexSnapshot); ok {
+		centroids, err := snap.HighestCardinalityCentroids("vec", limit)
+		if err != nil {
+			t.Fatal(err)
+		}
+		fmt.Println(centroids)
+	}
+}

index/scorch/snapshot_index_vr.go

@@ -23,6 +23,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"reflect"
+	"sort"
 
 	"github.com/blevesearch/bleve/v2/size"
 	index "github.com/blevesearch/bleve_index_api"
@@ -163,3 +164,40 @@ func (i *IndexSnapshotVectorReader) Close() error {
 	// TODO Consider if any scope of recycling here.
 	return nil
 }
+
+func (i *IndexSnapshot) HighestCardinalityCentroids(field string, limit int) (
+	centroids []segment_api.CentroidCardinality, err error) {
+	if len(i.segment) == 0 {
+		return nil, fmt.Errorf("no segments available")
+	}
+
+	if limit <= 0 {
+		return nil, fmt.Errorf("limit must be positive")
+	}
+
+	rvCentroids := make([]segment_api.CentroidCardinality, limit)
+
+	for _, segment := range i.segment {
+		if sv, ok := segment.segment.(segment_api.VectorSegment); ok {
+			vecIndex, err := sv.InterpretVectorIndex(field,
+				false /* does not require filtering */, segment.deleted)
+			if err != nil {
+				return nil, fmt.Errorf("failed to interpret vector index for field %s in segment: %v", field, err)
+			}
+
+			centroidCardinalities, err := vecIndex.ObtainTopKCentroidCardinalitiesFromIVFIndex(limit)
+			if err != nil {
+				return nil, fmt.Errorf("failed to obtain top k centroid cardinalities for field %s in segment: %v", field, err)
+			}
+
+			centroidCardinalities = append(centroidCardinalities, rvCentroids...)
+			sort.Slice(centroidCardinalities, func(i, j int) bool {
+				return centroidCardinalities[i].Cardinality > centroidCardinalities[j].Cardinality
+			})
+
+			rvCentroids = centroidCardinalities[:limit]
+		}
+	}
+
+	return rvCentroids, nil
+}