Support config small index nprobe and nlist (#16731)

Signed-off-by: xiaofan-luan <xiaofan.luan@zilliz.com>

configs/milvus.yaml

@@ -177,10 +177,12 @@ queryNode:
       pulsarBufSize: 512 # pulsar channel buffer size
     searchResult:
       recvBufSize: 64 # msgPack channel buffer size
-  # Segcore will divide a segment into multiple chunks.
+  # Segcore will divide a segment into multiple chunks to enbale small index
   segcore:
     chunkRows: 32768 # The number of vectors in a chunk.
-
+    smallIndex:
+      nlist: 256 # small index nlist, recommend to set sqrt(chunkRows), must smaller than chunkRows/8
+      nprobe: 16 # nprobe to search small index, based on your accuracy requirement, must smaller than nlist
   cache:
     enabled: true
     memoryLimit: 2147483648 # 2 GB, 2 * 1024 *1024 *1024

internal/core/src/segcore/SegcoreConfig.h

@@ -31,8 +31,8 @@ class SegcoreConfig {
     SegcoreConfig() {
         // hard code configurations for small index
         SmallIndexConf sub_conf;
-        sub_conf.build_params["nlist"] = 100;
-        sub_conf.search_params["nprobe"] = 4;
+        sub_conf.build_params["nlist"] = nlist_;
+        sub_conf.search_params["nprobe"] = nprobe_;
         sub_conf.index_type = "IVF";
         table_[MetricType::METRIC_L2] = sub_conf;
         table_[MetricType::METRIC_INNER_PRODUCT] = sub_conf;
@@ -65,6 +65,16 @@ class SegcoreConfig {
         chunk_rows_ = chunk_rows;
     }
 
+    void
+    set_nlist(int64_t nlist) {
+        nlist_ = nlist;
+    }
+
+    void
+    set_nprobe(int64_t nprobe) {
+        nprobe_ = nprobe;
+    }
+
     void
     set_small_index_config(MetricType metric_type, const SmallIndexConf& small_index_conf) {
         table_[metric_type] = small_index_conf;
@@ -72,6 +82,8 @@ class SegcoreConfig {
 
  private:
     int64_t chunk_rows_ = 32 * 1024;
+    int64_t nlist_ = 100;
+    int64_t nprobe_ = 4;
     std::map<MetricType, SmallIndexConf> table_;
 };
 

internal/core/src/segcore/segcore_init_c.cpp

@@ -20,11 +20,23 @@ SegcoreInit(const char* conf_file) {
     milvus::config::KnowhereInitImpl(conf_file);
 }
 
+// TODO merge small index config into one config map, including enable/disable small_index
 extern "C" void
 SegcoreSetChunkRows(const int64_t value) {
     milvus::segcore::SegcoreConfig& config = milvus::segcore::SegcoreConfig::default_config();
     config.set_chunk_rows(value);
-    LOG_SEGCORE_DEBUG_ << "set config chunk_size: " << config.get_chunk_rows();
+}
+
+extern "C" void
+SegcoreSetNlist(const int64_t value) {
+    milvus::segcore::SegcoreConfig& config = milvus::segcore::SegcoreConfig::default_config();
+    config.set_nlist(value);
+}
+
+extern "C" void
+SegcoreSetNprobe(const int64_t value) {
+    milvus::segcore::SegcoreConfig& config = milvus::segcore::SegcoreConfig::default_config();
+    config.set_nprobe(value);
 }
 
 // return value must be freed by the caller

internal/core/src/segcore/segcore_init_c.h

@@ -21,6 +21,12 @@ SegcoreInit(const char*);
 void
 SegcoreSetChunkRows(const int64_t);
 
+void
+SegcoreSetNlist(const int64_t);
+
+void
+SegcoreSetNprobe(const int64_t);
+
 // return value must be freed by the caller
 char*
 SegcoreSetSimdType(const char*);

internal/querynode/query_node.go

@@ -185,6 +185,12 @@ func (node *QueryNode) InitSegcore() {
 	cChunkRows := C.int64_t(Params.QueryNodeCfg.ChunkRows)
 	C.SegcoreSetChunkRows(cChunkRows)
 
+	nlist := C.int64_t(Params.QueryNodeCfg.SmallIndexNlist)
+	C.SegcoreSetNlist(nlist)
+
+	nprobe := C.int64_t(Params.QueryNodeCfg.SmallIndexNProbe)
+	C.SegcoreSetNprobe(nprobe)
+
 	// override segcore SIMD type
 	cSimdType := C.CString(Params.CommonCfg.SimdType)
 	cRealSimdType := C.SegcoreSetSimdType(cSimdType)

internal/util/paramtable/component_param.go

@@ -12,6 +12,7 @@
 package paramtable
 
 import (
+	"math"
 	"os"
 	"path"
 	"strconv"
@@ -702,7 +703,9 @@ type queryNodeConfig struct {
 	SliceIndex   int
 
 	// segcore
-	ChunkRows int64
+	ChunkRows        int64
+	SmallIndexNlist  int64
+	SmallIndexNProbe int64
 
 	CreatedTime time.Time
 	UpdatedTime time.Time
@@ -730,7 +733,7 @@ func (p *queryNodeConfig) init(base *BaseTable) {
 
 	p.initStatsPublishInterval()
 
-	p.initSegcoreChunkRows()
+	p.initSmallIndexParams()
 
 	p.initOverloadedMemoryThresholdPercentage()
 
@@ -797,8 +800,34 @@ func (p *queryNodeConfig) initGracefulTime() {
 	log.Debug("query node init gracefulTime", zap.Any("gracefulTime", p.GracefulTime))
 }
 
-func (p *queryNodeConfig) initSegcoreChunkRows() {
+func (p *queryNodeConfig) initSmallIndexParams() {
 	p.ChunkRows = p.Base.ParseInt64WithDefault("queryNode.segcore.chunkRows", 32768)
+	if p.ChunkRows < 1024 {
+		log.Warn("chunk rows can not be less than 1024, force set to 1024", zap.Any("current", p.ChunkRows))
+		p.ChunkRows = 1024
+	}
+
+	// default NList is the first nlist
+	var defaultNList int64
+	for i := int64(0); i < p.ChunkRows; i++ {
+		if math.Pow(2.0, float64(i)) > math.Sqrt(float64(p.ChunkRows)) {
+			defaultNList = int64(math.Pow(2, float64(i)))
+			break
+		}
+	}
+
+	p.SmallIndexNlist = p.Base.ParseInt64WithDefault("queryNode.segcore.smallIndex.nlist", defaultNList)
+	if p.SmallIndexNlist > p.ChunkRows/8 {
+		log.Warn("small index nlist must smaller than chunkRows/8, force set to", zap.Any("nliit", p.ChunkRows/8))
+		p.SmallIndexNlist = p.ChunkRows / 8
+	}
+
+	defaultNprobe := p.SmallIndexNlist / 16
+	p.SmallIndexNProbe = p.Base.ParseInt64WithDefault("queryNode.segcore.smallIndex.nprobe", defaultNprobe)
+	if p.SmallIndexNProbe > p.SmallIndexNlist {
+		log.Warn("small index nprobe must smaller than nlist, force set to", zap.Any("nprobe", p.SmallIndexNlist))
+		p.SmallIndexNProbe = p.SmallIndexNlist
+	}
 }
 
 func (p *queryNodeConfig) initOverloadedMemoryThresholdPercentage() {

internal/util/paramtable/component_param_test.go

@@ -236,6 +236,43 @@ func TestComponentParam(t *testing.T) {
 
 		maxParallelism := Params.FlowGraphMaxParallelism
 		assert.Equal(t, int32(1024), maxParallelism)
+
+		// test query side config
+		chunkRows := Params.ChunkRows
+		assert.Equal(t, int64(32768), chunkRows)
+
+		nlist := Params.SmallIndexNlist
+		assert.Equal(t, int64(256), nlist)
+
+		nprobe := Params.SmallIndexNProbe
+		assert.Equal(t, int64(16), nprobe)
+
+		// test small indexNlist/NProbe default
+		Params.Base.Remove("queryNode.segcore.smallIndex.nlist")
+		Params.Base.Remove("queryNode.segcore.smallIndex.nprobe")
+		Params.Base.Save("queryNode.segcore.chunkRows", "8192")
+		Params.initSmallIndexParams()
+		chunkRows = Params.ChunkRows
+		assert.Equal(t, int64(8192), chunkRows)
+
+		nlist = Params.SmallIndexNlist
+		assert.Equal(t, int64(128), nlist)
+
+		nprobe = Params.SmallIndexNProbe
+		assert.Equal(t, int64(8), nprobe)
+
+		Params.Base.Remove("queryNode.segcore.smallIndex.nlist")
+		Params.Base.Remove("queryNode.segcore.smallIndex.nprobe")
+		Params.Base.Save("queryNode.segcore.chunkRows", "64")
+		Params.initSmallIndexParams()
+		chunkRows = Params.ChunkRows
+		assert.Equal(t, int64(1024), chunkRows)
+
+		nlist = Params.SmallIndexNlist
+		assert.Equal(t, int64(64), nlist)
+
+		nprobe = Params.SmallIndexNProbe
+		assert.Equal(t, int64(4), nprobe)
 	})
 
 	t.Run("test dataCoordConfig", func(t *testing.T) {