diff --git a/pkg/util/indexparamcheck/ivf_pq_checker.go b/pkg/util/indexparamcheck/ivf_pq_checker.go index 51da64e0ff..4c35f193c4 100644 --- a/pkg/util/indexparamcheck/ivf_pq_checker.go +++ b/pkg/util/indexparamcheck/ivf_pq_checker.go @@ -33,10 +33,14 @@ func (c *ivfPQChecker) checkPQParams(params map[string]string) error { // nbits can be set to default: 8 nbitsStr, nbitsExist := params[NBITS] if nbitsExist { - _, err := strconv.Atoi(nbitsStr) + nbits, err := strconv.Atoi(nbitsStr) if err != nil { // invalid nbits return fmt.Errorf("invalid nbits: %s", nbitsStr) } + + if nbits < 1 || nbits > 64 { + return fmt.Errorf("parameter `nbits` out of range, expect range [1,64], current value: %d", nbits) + } } mStr, ok := params[IVFM] diff --git a/pkg/util/indexparamcheck/ivf_pq_checker_test.go b/pkg/util/indexparamcheck/ivf_pq_checker_test.go index b4de37579f..4a22d45542 100644 --- a/pkg/util/indexparamcheck/ivf_pq_checker_test.go +++ b/pkg/util/indexparamcheck/ivf_pq_checker_test.go @@ -46,6 +46,10 @@ func Test_ivfPQChecker_CheckTrain(t *testing.T) { invalidParamsNbits := copyParams(validParams) invalidParamsNbits[NBITS] = "NAN" + invalidParamsNbitsLower := copyParams(validParams) + invalidParamsNbitsLower[NBITS] = "0" + invalidParamsNbitsUpper := copyParams(validParams) + invalidParamsNbitsUpper[NBITS] = "65" invalidParamsWithoutIVF := map[string]string{ DIM: strconv.Itoa(128), @@ -123,6 +127,8 @@ func Test_ivfPQChecker_CheckTrain(t *testing.T) { {validParamsWithoutDim, false}, {invalidParamsDim, false}, {invalidParamsNbits, false}, + {invalidParamsNbitsLower, false}, + {invalidParamsNbitsUpper, false}, {invalidParamsWithoutIVF, false}, {invalidParamsIVF, false}, {invalidParamsMzero, false},