fix: Add `nbits` parameter check for IVF_PQ (#34451)

See also #34426

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
pull/34441/head^2
congqixia 2024-07-08 11:16:10 +08:00 committed by GitHub
parent efdaed4ac6
commit 145f4f919b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 11 additions and 1 deletions

View File

@ -33,10 +33,14 @@ func (c *ivfPQChecker) checkPQParams(params map[string]string) error {
// nbits can be set to default: 8 // nbits can be set to default: 8
nbitsStr, nbitsExist := params[NBITS] nbitsStr, nbitsExist := params[NBITS]
if nbitsExist { if nbitsExist {
_, err := strconv.Atoi(nbitsStr) nbits, err := strconv.Atoi(nbitsStr)
if err != nil { // invalid nbits if err != nil { // invalid nbits
return fmt.Errorf("invalid nbits: %s", nbitsStr) 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] mStr, ok := params[IVFM]

View File

@ -46,6 +46,10 @@ func Test_ivfPQChecker_CheckTrain(t *testing.T) {
invalidParamsNbits := copyParams(validParams) invalidParamsNbits := copyParams(validParams)
invalidParamsNbits[NBITS] = "NAN" invalidParamsNbits[NBITS] = "NAN"
invalidParamsNbitsLower := copyParams(validParams)
invalidParamsNbitsLower[NBITS] = "0"
invalidParamsNbitsUpper := copyParams(validParams)
invalidParamsNbitsUpper[NBITS] = "65"
invalidParamsWithoutIVF := map[string]string{ invalidParamsWithoutIVF := map[string]string{
DIM: strconv.Itoa(128), DIM: strconv.Itoa(128),
@ -123,6 +127,8 @@ func Test_ivfPQChecker_CheckTrain(t *testing.T) {
{validParamsWithoutDim, false}, {validParamsWithoutDim, false},
{invalidParamsDim, false}, {invalidParamsDim, false},
{invalidParamsNbits, false}, {invalidParamsNbits, false},
{invalidParamsNbitsLower, false},
{invalidParamsNbitsUpper, false},
{invalidParamsWithoutIVF, false}, {invalidParamsWithoutIVF, false},
{invalidParamsIVF, false}, {invalidParamsIVF, false},
{invalidParamsMzero, false}, {invalidParamsMzero, false},