From a1dbddfbfe1fb218e693e86e13cf8ebdcca07733 Mon Sep 17 00:00:00 2001 From: Congqi Xia Date: Fri, 11 Aug 2023 16:22:02 +0800 Subject: [PATCH] Support SCANN index type Signed-off-by: Congqi Xia --- entity/genidx/genidx.go | 35 ++++++++++++++++ entity/index.go | 1 + entity/indexes_gen.go | 56 +++++++++++++++++++++++++ entity/indexes_gen_test.go | 41 ++++++++++++++++++ entity/indexes_search_param_gen.go | 42 +++++++++++++++++++ entity/indexes_search_param_gen_test.go | 48 +++++++++++++++++++++ 6 files changed, 223 insertions(+) diff --git a/entity/genidx/genidx.go b/entity/genidx/genidx.go index 379979e5..88e3f45e 100644 --- a/entity/genidx/genidx.go +++ b/entity/genidx/genidx.go @@ -708,6 +708,41 @@ func main() { "nprobe = 65537", }, }, + { + IdxName: "SCANN", + IdxType: entity.IvfFlat, + ConstructParams: []idxParam{ + { + Name: "nlist", + ValidationRule: "[1, 65536]", + }, + }, + SearchParams: []idxParam{ + { + Name: "nprobe", + ValidationRule: "[1, 65536]", // [1, nlist], refer to index construct param, not supported yet + }, + { + Name: "reorder_k", + ValidationRule: "[1, 9223372036854775807]", // [topk, MAX_INT], refer to index construct param, not supported yet + }, + }, + ValidExamples: []string{ + "nlist = 100", + }, + InvalidExamples: []string{ + "nlist = 0", + "nlist = 65537", + }, + ValidSearchParams: []string{ + "nprobe, reorder_k = 10, 200", + }, + InvalidSearchParams: []string{ + "nprobe, reorder_k = 0, 200", + "nprobe, reorder_k = 65537, 200", + "nprobe, reorder_k = 10, -1", + }, + }, }, } diff --git a/entity/index.go b/entity/index.go index bc92bffd..1dae8e5f 100644 --- a/entity/index.go +++ b/entity/index.go @@ -36,6 +36,7 @@ const ( IvfHNSW IndexType = "IVF_HNSW" AUTOINDEX IndexType = "AUTOINDEX" DISKANN IndexType = "DISKANN" + SCANN IndexType = "SCANN" GPUIvfFlat IndexType = "GPU_IVF_FLAT" GPUIvfPQ IndexType = "GPU_IVF_PQ" diff --git a/entity/indexes_gen.go b/entity/indexes_gen.go index acbd696e..ebcd3c61 100755 --- a/entity/indexes_gen.go +++ b/entity/indexes_gen.go @@ -724,3 +724,59 @@ func NewIndexGPUIvfPQ(metricType MetricType, }, nil } + +var _ Index = &IndexSCANN{} + +// IndexSCANN idx type for IVF_FLAT +type IndexSCANN struct { //auto generated fields + nlist int + metricType MetricType +} + +// Name returns index type name, implementing Index interface +func(i *IndexSCANN) Name() string { + return "SCANN" +} + +// IndexType returns IndexType, implementing Index interface +func(i *IndexSCANN) IndexType() IndexType { + return IndexType("IVF_FLAT") +} + +// SupportBinary returns whether index type support binary vector +func(i *IndexSCANN) SupportBinary() bool { + return 0 & 2 > 0 +} + +// Params returns index construction params, implementing Index interface +func(i *IndexSCANN) Params() map[string]string { + params := map[string]string {//auto generated mapping + "nlist": fmt.Sprintf("%v",i.nlist), + } + bs, _ := json.Marshal(params) + return map[string]string { + "params": string(bs), + "index_type": string(i.IndexType()), + "metric_type": string(i.metricType), + } +} + +// NewIndexSCANN create index with construction parameters +func NewIndexSCANN(metricType MetricType, + nlist int, +) (*IndexSCANN, error) { + // auto generate parameters validation code, if any + if nlist < 1 { + return nil, errors.New("nlist has to be in range [1, 65536]") + } + if nlist > 65536 { + return nil, errors.New("nlist has to be in range [1, 65536]") + } + + return &IndexSCANN{ + //auto generated setting + nlist: nlist, + metricType: metricType, + }, nil +} + diff --git a/entity/indexes_gen_test.go b/entity/indexes_gen_test.go index 4de1bbc2..97fa4f77 100755 --- a/entity/indexes_gen_test.go +++ b/entity/indexes_gen_test.go @@ -570,3 +570,44 @@ func TestIndexGPUIvfPQ(t *testing.T){ }) } +func TestIndexSCANN(t *testing.T){ + + var nlist int + + mt := L2 + + + t.Run("valid usage case", func(t *testing.T){ + + nlist = 100 + idx0, err := NewIndexSCANN(mt, + nlist, + ) + assert.Nil(t, err) + assert.NotNil(t, idx0) + assert.Equal(t, "SCANN", idx0.Name()) + assert.EqualValues(t, "IVF_FLAT", idx0.IndexType()) + assert.NotNil(t, idx0.Params()) + assert.False(t, idx0.SupportBinary()) + + }) + + t.Run("invalid usage case", func(t *testing.T){ + + nlist = 0 + idx0, err := NewIndexSCANN(mt, + nlist, + ) + assert.NotNil(t, err) + assert.Nil(t, idx0) + + nlist = 65537 + idx1, err := NewIndexSCANN(mt, + nlist, + ) + assert.NotNil(t, err) + assert.Nil(t, idx1) + + }) +} + diff --git a/entity/indexes_search_param_gen.go b/entity/indexes_search_param_gen.go index bcff1018..a5a8148c 100755 --- a/entity/indexes_search_param_gen.go +++ b/entity/indexes_search_param_gen.go @@ -380,3 +380,45 @@ func NewIndexGPUIvfPQSearchParam( return sp, nil } +var _ SearchParam = &IndexSCANNSearchParam{} + +// IndexSCANNSearchParam search param struct for index type IVF_FLAT +type IndexSCANNSearchParam struct { //auto generated fields + baseSearchParams + + nprobe int + reorder_k int +} + +// NewIndexSCANNSearchParam create index search param +func NewIndexSCANNSearchParam( + nprobe int, + + reorder_k int, +) (*IndexSCANNSearchParam, error) { + // auto generate parameters validation code, if any + if nprobe < 1 { + return nil, errors.New("nprobe has to be in range [1, 65536]") + } + if nprobe > 65536 { + return nil, errors.New("nprobe has to be in range [1, 65536]") + } + + if reorder_k < 1 { + return nil, errors.New("reorder_k has to be in range [1, 9223372036854775807]") + } + if reorder_k > 9223372036854775807 { + return nil, errors.New("reorder_k has to be in range [1, 9223372036854775807]") + } + + sp := &IndexSCANNSearchParam{ + baseSearchParams: newBaseSearchParams(), + } + + //auto generated setting + sp.params["nprobe"] = nprobe + sp.params["reorder_k"] = reorder_k + + return sp, nil +} + diff --git a/entity/indexes_search_param_gen_test.go b/entity/indexes_search_param_gen_test.go index f196d338..67b9d599 100755 --- a/entity/indexes_search_param_gen_test.go +++ b/entity/indexes_search_param_gen_test.go @@ -449,3 +449,51 @@ func TestIndexGPUIvfPQSearchParam(t *testing.T) { } +func TestIndexSCANNSearchParam(t *testing.T) { + + var nprobe int + var reorder_k int + + t.Run("valid usage case", func(t *testing.T){ + + nprobe, reorder_k = 10, 200 + idx0, err := NewIndexSCANNSearchParam( + nprobe, + reorder_k, + ) + assert.Nil(t, err) + assert.NotNil(t, idx0) + assert.NotNil(t, idx0.Params()) + + }) + + t.Run("invalid usage case", func(t *testing.T){ + + nprobe, reorder_k = 0, 200 + idx0, err := NewIndexSCANNSearchParam( + nprobe, + reorder_k, + ) + assert.NotNil(t, err) + assert.Nil(t, idx0) + + nprobe, reorder_k = 65537, 200 + idx1, err := NewIndexSCANNSearchParam( + nprobe, + reorder_k, + ) + assert.NotNil(t, err) + assert.Nil(t, idx1) + + nprobe, reorder_k = 10, -1 + idx2, err := NewIndexSCANNSearchParam( + nprobe, + reorder_k, + ) + assert.NotNil(t, err) + assert.Nil(t, idx2) + + }) + +} +