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
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]

View File

@ -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},