Fix IVF_SQ nbits check (#20183)

Signed-off-by: longjiquan <jiquan.long@zilliz.com>

Signed-off-by: longjiquan <jiquan.long@zilliz.com>
This commit is contained in:
Jiquan Long 2022-10-31 10:13:34 +08:00 committed by GitHub
parent a8e2a69aa2
commit 8d0cc4226c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 30 additions and 6 deletions

View File

@ -254,9 +254,21 @@ type IVFSQConfAdapter struct {
IVFConfAdapter
}
func (adapter *IVFSQConfAdapter) checkNBits(params map[string]string) bool {
// cgo will set this key to DefaultNBits (8), which is the only value Milvus supports.
_, exist := params[NBITS]
if exist {
// 8 is the only supported nbits.
return CheckIntByRange(params, NBITS, DefaultNBits, DefaultNBits)
}
return true
}
// CheckTrain returns true if the index can be built with the specific index parameters.
func (adapter *IVFSQConfAdapter) CheckTrain(params map[string]string) bool {
params[NBITS] = strconv.Itoa(DefaultNBits)
if !adapter.checkNBits(params) {
return false
}
return adapter.IVFConfAdapter.CheckTrain(params)
}

View File

@ -168,17 +168,29 @@ func TestIVFPQConfAdapter_CheckTrain(t *testing.T) {
}
func TestIVFSQConfAdapter_CheckTrain(t *testing.T) {
validParams := map[string]string{
DIM: strconv.Itoa(128),
NLIST: strconv.Itoa(100),
NBITS: strconv.Itoa(8),
Metric: L2,
getValidParams := func(withNBits bool) map[string]string {
validParams := map[string]string{
DIM: strconv.Itoa(128),
NLIST: strconv.Itoa(100),
NBITS: strconv.Itoa(8),
Metric: L2,
}
if withNBits {
validParams[NBITS] = strconv.Itoa(DefaultNBits)
}
return validParams
}
validParams := getValidParams(false)
validParamsWithNBits := getValidParams(true)
paramsWithInvalidNBits := getValidParams(false)
paramsWithInvalidNBits[NBITS] = strconv.Itoa(DefaultNBits + 1)
cases := []struct {
params map[string]string
want bool
}{
{validParams, true},
{validParamsWithNBits, true},
{paramsWithInvalidNBits, false},
{invalidIVFParamsMin(), false},
{invalidIVFParamsMax(), false},
}