Skip to content

Commit 3f9c498

Browse files
alwayslove2013XuanYang-cn
authored andcommitted
add more milvus index types: ivf_pq
Signed-off-by: min.tian <min.tian.cn@gmail.com>
1 parent b979e79 commit 3f9c498

3 files changed

Lines changed: 30 additions & 4 deletions

File tree

vectordb_bench/backend/clients/api.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ class IndexType(str, Enum):
2222
DISKANN = "DISKANN"
2323
STREAMING_DISKANN = "DISKANN"
2424
IVFFlat = "IVF_FLAT"
25+
IVFPQ = "IVF_PQ"
2526
IVFSQ8 = "IVF_SQ8"
2627
IVF_RABITQ = "IVF_RABITQ"
2728
Flat = "FLAT"

vectordb_bench/backend/clients/milvus/config.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,27 @@ def search_param(self) -> dict:
207207
}
208208

209209

210+
class IVFPQConfig(MilvusIndexConfig, DBCaseConfig):
211+
nlist: int
212+
nprobe: int | None = None
213+
m: int = 32
214+
nbits: int = 8
215+
index: IndexType = IndexType.IVFPQ
216+
217+
def index_param(self) -> dict:
218+
return {
219+
"metric_type": self.parse_metric(),
220+
"index_type": self.index.value,
221+
"params": {"nlist": self.nlist, "m": self.m, "nbits": self.nbits},
222+
}
223+
224+
def search_param(self) -> dict:
225+
return {
226+
"metric_type": self.parse_metric(),
227+
"params": {"nprobe": self.nprobe},
228+
}
229+
230+
210231
class IVFSQ8Config(MilvusIndexConfig, DBCaseConfig):
211232
nlist: int
212233
nprobe: int | None = None
@@ -397,6 +418,7 @@ def search_param(self) -> dict:
397418
IndexType.HNSW_PRQ: HNSWPRQConfig,
398419
IndexType.DISKANN: DISKANNConfig,
399420
IndexType.IVFFlat: IVFFlatConfig,
421+
IndexType.IVFPQ: IVFPQConfig,
400422
IndexType.IVFSQ8: IVFSQ8Config,
401423
IndexType.IVF_RABITQ: IVFRABITQConfig,
402424
IndexType.Flat: FLATConfig,

vectordb_bench/frontend/config/dbCaseConfigs.py

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,7 @@ class CaseConfigInput(BaseModel):
168168
IndexType.HNSW_PQ.value,
169169
IndexType.HNSW_PRQ.value,
170170
IndexType.IVFFlat.value,
171+
IndexType.IVFPQ.value,
171172
IndexType.IVFSQ8.value,
172173
IndexType.IVF_RABITQ.value,
173174
IndexType.DISKANN.value,
@@ -631,6 +632,7 @@ class CaseConfigInput(BaseModel):
631632
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
632633
in [
633634
IndexType.IVFFlat.value,
635+
IndexType.IVFPQ.value,
634636
IndexType.IVFSQ8.value,
635637
IndexType.IVF_RABITQ.value,
636638
IndexType.GPU_IVF_FLAT.value,
@@ -650,6 +652,7 @@ class CaseConfigInput(BaseModel):
650652
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
651653
in [
652654
IndexType.IVFFlat.value,
655+
IndexType.IVFPQ.value,
653656
IndexType.IVFSQ8.value,
654657
IndexType.IVF_RABITQ.value,
655658
IndexType.GPU_IVF_FLAT.value,
@@ -662,12 +665,12 @@ class CaseConfigInput(BaseModel):
662665
label=CaseConfigParamType.m,
663666
inputType=InputType.Number,
664667
inputConfig={
665-
"min": 0,
668+
"min": 1,
666669
"max": 65536,
667-
"value": 0,
670+
"value": 32,
668671
},
669672
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
670-
in [IndexType.GPU_IVF_PQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value],
673+
in [IndexType.GPU_IVF_PQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value, IndexType.IVFPQ.value],
671674
)
672675

673676

@@ -680,7 +683,7 @@ class CaseConfigInput(BaseModel):
680683
"value": 8,
681684
},
682685
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
683-
in [IndexType.GPU_IVF_PQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value],
686+
in [IndexType.GPU_IVF_PQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value, IndexType.IVFPQ.value],
684687
)
685688

686689
CaseConfigParamInput_NRQ = CaseConfigInput(

0 commit comments

Comments
 (0)