Skip to content

Commit 2f206d9

Browse files
alwayslove2013alwayslove2013
committed
add more milvus index types: hnsw sq/pq/prq; ivf rabitq
Signed-off-by: min.tian <min.tian.cn@gmail.com>
1 parent 90879f7 commit 2f206d9

5 files changed

Lines changed: 234 additions & 9 deletions

File tree

‎vectordb_bench/backend/clients/api.py‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,14 @@ class MetricType(str, Enum):
1616

1717
class IndexType(str, Enum):
1818
HNSW = "HNSW"
19+
HNSW_SQ = "HNSW_SQ"
20+
HNSW_PQ = "HNSW_PQ"
21+
HNSW_PRQ = "HNSW_PRQ"
1922
DISKANN = "DISKANN"
2023
STREAMING_DISKANN = "DISKANN"
2124
IVFFlat = "IVF_FLAT"
2225
IVFSQ8 = "IVF_SQ8"
26+
IVF_RABITQ = "IVF_RABITQ"
2327
Flat = "FLAT"
2428
AUTOINDEX = "AUTOINDEX"
2529
ES_HNSW = "hnsw"
@@ -30,6 +34,12 @@ class IndexType(str, Enum):
3034
GPU_CAGRA = "GPU_CAGRA"
3135
SCANN = "scann"
3236

37+
class SQType(str, Enum):
38+
SQ6 = "SQ6"
39+
SQ8 = "SQ8"
40+
BF16 = "BF16"
41+
FP16 = "FP16"
42+
FP32 = "FP32"
3343

3444
class DBConfig(ABC, BaseModel):
3545
"""DBConfig contains the connection info of vector database

‎vectordb_bench/backend/clients/milvus/config.py‎

Lines changed: 112 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from pydantic import BaseModel, SecretStr, validator
22

3-
from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
3+
from ..api import DBCaseConfig, DBConfig, IndexType, MetricType, SQType
44

55

66
class MilvusConfig(DBConfig):
@@ -88,6 +88,88 @@ def search_param(self) -> dict:
8888
}
8989

9090

91+
class HNSWSQConfig(HNSWConfig, DBCaseConfig):
92+
index: IndexType = IndexType.HNSW_SQ
93+
sq_type: SQType = SQType.SQ8
94+
refine: bool = True
95+
refine_type: SQType = SQType.FP32
96+
refine_k: float = 1
97+
98+
def index_param(self) -> dict:
99+
return {
100+
"metric_type": self.parse_metric(),
101+
"index_type": self.index.value,
102+
"params": {
103+
"M": self.M,
104+
"efConstruction": self.efConstruction,
105+
"sq_type": self.sq_type.value,
106+
"refine": self.refine,
107+
"refine_type": self.refine_type.value,
108+
},
109+
}
110+
111+
def search_param(self) -> dict:
112+
return {
113+
"metric_type": self.parse_metric(),
114+
"params": {"ef": self.ef, "refine_k": self.refine_k},
115+
}
116+
117+
118+
class HNSWPQConfig(HNSWConfig):
119+
index: IndexType = IndexType.HNSW_PQ
120+
m: int = 32
121+
nbits: int = 8
122+
refine: bool = True
123+
refine_type: SQType = SQType.FP32
124+
refine_k: float = 1
125+
126+
def index_param(self) -> dict:
127+
return {
128+
"metric_type": self.parse_metric(),
129+
"index_type": self.index.value,
130+
"params": {
131+
"M": self.M,
132+
"efConstruction": self.efConstruction,
133+
"m": self.m,
134+
"nbits": self.nbits,
135+
"refine": self.refine,
136+
"refine_type": self.refine_type.value,
137+
},
138+
}
139+
140+
def search_param(self) -> dict:
141+
return {
142+
"metric_type": self.parse_metric(),
143+
"params": {"ef": self.ef, "refine_k": self.refine_k},
144+
}
145+
146+
147+
class HNSWPRQConfig(HNSWPQConfig):
148+
index: IndexType = IndexType.HNSW_PRQ
149+
nrq: int = 2
150+
151+
def index_param(self) -> dict:
152+
return {
153+
"metric_type": self.parse_metric(),
154+
"index_type": self.index.value,
155+
"params": {
156+
"M": self.M,
157+
"efConstruction": self.efConstruction,
158+
"m": self.m,
159+
"nbits": self.nbits,
160+
"nrq": self.nrq,
161+
"refine": self.refine,
162+
"refine_type": self.refine_type.value,
163+
},
164+
}
165+
166+
def search_param(self) -> dict:
167+
return {
168+
"metric_type": self.parse_metric(),
169+
"params": {"ef": self.ef, "refine_k": self.refine_k},
170+
}
171+
172+
91173
class DISKANNConfig(MilvusIndexConfig, DBCaseConfig):
92174
search_list: int | None = None
93175
index: IndexType = IndexType.DISKANN
@@ -144,6 +226,31 @@ def search_param(self) -> dict:
144226
}
145227

146228

229+
class IVFRABITQConfig(IVFSQ8Config):
230+
index: IndexType = IndexType.IVF_RABITQ
231+
rbq_bits_query: int = 0 # 0, 1, 2, ..., 8
232+
refine: bool = True
233+
refine_type: SQType = SQType.FP32
234+
refine_k: float = 1
235+
236+
def index_param(self) -> dict:
237+
return {
238+
"metric_type": self.parse_metric(),
239+
"index_type": self.index.value,
240+
"params": {
241+
"nlist": self.nlist,
242+
"refine": self.refine,
243+
"refine_type": self.refine_type.value,
244+
},
245+
}
246+
247+
def search_param(self) -> dict:
248+
return {
249+
"metric_type": self.parse_metric(),
250+
"params": {"nprobe": self.nprobe, "rbq_bits_query": self.rbq_bits_query, "refine_k": self.refine_k},
251+
}
252+
253+
147254
class FLATConfig(MilvusIndexConfig, DBCaseConfig):
148255
index: IndexType = IndexType.Flat
149256

@@ -285,9 +392,13 @@ def search_param(self) -> dict:
285392
_milvus_case_config = {
286393
IndexType.AUTOINDEX: AutoIndexConfig,
287394
IndexType.HNSW: HNSWConfig,
395+
IndexType.HNSW_SQ: HNSWSQConfig,
396+
IndexType.HNSW_PQ: HNSWPQConfig,
397+
IndexType.HNSW_PRQ: HNSWPRQConfig,
288398
IndexType.DISKANN: DISKANNConfig,
289399
IndexType.IVFFlat: IVFFlatConfig,
290400
IndexType.IVFSQ8: IVFSQ8Config,
401+
IndexType.IVF_RABITQ: IVFRABITQConfig,
291402
IndexType.Flat: FLATConfig,
292403
IndexType.GPU_IVF_FLAT: GPUIVFFlatConfig,
293404
IndexType.GPU_IVF_PQ: GPUIVFPQConfig,

‎vectordb_bench/backend/clients/milvus/milvus.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,7 @@ def __init__(
6161
consistency_level="Session",
6262
)
6363

64+
log.info(f"{self.name} create index: index_params: {self.case_config.index_param()}")
6465
col.create_index(
6566
self._vector_field,
6667
self.case_config.index_param(),
@@ -71,7 +72,7 @@ def __init__(
7172
connections.disconnect("default")
7273

7374
@contextmanager
74-
def init(self) -> None:
75+
def init(self):
7576
"""
7677
Examples:
7778
>>> with self.init():
@@ -126,6 +127,7 @@ def wait_index():
126127
try:
127128
self.col.compact()
128129
self.col.wait_for_compaction_completed()
130+
log.info("compactation completed. waiting for the rest of index buliding.")
129131
except Exception as e:
130132
log.warning(f"{self.name} compact error: {e}")
131133
if hasattr(e, "code"):

‎vectordb_bench/frontend/config/dbCaseConfigs.py‎

Lines changed: 103 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
from pydantic import BaseModel
44
from vectordb_bench.backend.cases import CaseLabel, CaseType
55
from vectordb_bench.backend.clients import DB
6-
from vectordb_bench.backend.clients.api import IndexType, MetricType
6+
from vectordb_bench.backend.clients.api import IndexType, MetricType, SQType
77
from vectordb_bench.frontend.components.custom.getCustomConfig import get_custom_configs
88

99
from vectordb_bench.models import CaseConfig, CaseConfigParamType
@@ -164,10 +164,13 @@ class CaseConfigInput(BaseModel):
164164
inputConfig={
165165
"options": [
166166
IndexType.HNSW.value,
167+
IndexType.HNSW_SQ.value,
168+
IndexType.HNSW_PQ.value,
169+
IndexType.HNSW_PRQ.value,
167170
IndexType.IVFFlat.value,
168171
IndexType.IVFSQ8.value,
172+
IndexType.IVF_RABITQ.value,
169173
IndexType.DISKANN.value,
170-
IndexType.STREAMING_DISKANN.value,
171174
IndexType.Flat.value,
172175
IndexType.AUTOINDEX.value,
173176
IndexType.GPU_IVF_FLAT.value,
@@ -346,9 +349,16 @@ class CaseConfigInput(BaseModel):
346349
"max": 64,
347350
"value": 30,
348351
},
349-
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) == IndexType.HNSW.value,
352+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
353+
in [
354+
IndexType.HNSW.value,
355+
IndexType.HNSW_SQ.value,
356+
IndexType.HNSW_PQ.value,
357+
IndexType.HNSW_PRQ.value,
358+
],
350359
)
351360

361+
352362
CaseConfigParamInput_m = CaseConfigInput(
353363
label=CaseConfigParamType.m,
354364
inputType=InputType.Number,
@@ -369,7 +379,62 @@ class CaseConfigInput(BaseModel):
369379
"max": 512,
370380
"value": 360,
371381
},
372-
isDisplayed=lambda config: config[CaseConfigParamType.IndexType] == IndexType.HNSW.value,
382+
isDisplayed=lambda config: config[CaseConfigParamType.IndexType]
383+
in [
384+
IndexType.HNSW.value,
385+
IndexType.HNSW_SQ.value,
386+
IndexType.HNSW_PQ.value,
387+
IndexType.HNSW_PRQ.value,
388+
],
389+
)
390+
391+
CaseConfigParamInput_SQType = CaseConfigInput(
392+
label=CaseConfigParamType.sq_type,
393+
inputType=InputType.Option,
394+
inputHelp="Scalar quantizer type.",
395+
inputConfig={
396+
"options": [SQType.SQ6.value, SQType.SQ8.value, SQType.BF16.value, SQType.FP16.value, SQType.FP32.value]
397+
},
398+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) in [IndexType.HNSW_SQ.value],
399+
)
400+
401+
CaseConfigParamInput_Refine = CaseConfigInput(
402+
label=CaseConfigParamType.refine,
403+
inputType=InputType.Option,
404+
inputHelp="Whether refined data is reserved during index building.",
405+
inputConfig={"options": [True, False]},
406+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
407+
in [IndexType.HNSW_SQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value, IndexType.IVF_RABITQ.value],
408+
)
409+
410+
CaseConfigParamInput_RefineType = CaseConfigInput(
411+
label=CaseConfigParamType.refine_type,
412+
inputType=InputType.Option,
413+
inputHelp="The data type of the refine index.",
414+
inputConfig={
415+
"options": [SQType.FP32.value, SQType.FP16.value, SQType.BF16.value, SQType.SQ8.value, SQType.SQ6.value]
416+
},
417+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
418+
in [IndexType.HNSW_SQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value, IndexType.IVF_RABITQ.value]
419+
and config.get(CaseConfigParamType.refine, True),
420+
)
421+
422+
CaseConfigParamInput_RefineK = CaseConfigInput(
423+
label=CaseConfigParamType.refine_k,
424+
inputType=InputType.Float,
425+
inputHelp="The magnification factor of refine compared to k.",
426+
inputConfig={"min": 1.0, "max": 10000.0, "value": 1.0},
427+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
428+
in [IndexType.HNSW_SQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value, IndexType.IVF_RABITQ.value]
429+
and config.get(CaseConfigParamType.refine, True),
430+
)
431+
432+
CaseConfigParamInput_RBQBitsQuery = CaseConfigInput(
433+
label=CaseConfigParamType.rbq_bits_query,
434+
inputType=InputType.Number,
435+
inputHelp="The magnification factor of refine compared to k.",
436+
inputConfig={"min": 0, "max": 8, "value": 0},
437+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) in [IndexType.IVF_RABITQ.value],
373438
)
374439

375440
CaseConfigParamInput_EFConstruction_Weaviate = CaseConfigInput(
@@ -519,7 +584,13 @@ class CaseConfigInput(BaseModel):
519584
"max": MAX_STREAMLIT_INT,
520585
"value": 100,
521586
},
522-
isDisplayed=lambda config: config[CaseConfigParamType.IndexType] == IndexType.HNSW.value,
587+
isDisplayed=lambda config: config[CaseConfigParamType.IndexType]
588+
in [
589+
IndexType.HNSW.value,
590+
IndexType.HNSW_SQ.value,
591+
IndexType.HNSW_PQ.value,
592+
IndexType.HNSW_PRQ.value,
593+
],
523594
)
524595

525596
CaseConfigParamInput_EF_Weaviate = CaseConfigInput(
@@ -561,6 +632,7 @@ class CaseConfigInput(BaseModel):
561632
in [
562633
IndexType.IVFFlat.value,
563634
IndexType.IVFSQ8.value,
635+
IndexType.IVF_RABITQ.value,
564636
IndexType.GPU_IVF_FLAT.value,
565637
IndexType.GPU_IVF_PQ.value,
566638
IndexType.GPU_BRUTE_FORCE.value,
@@ -579,6 +651,7 @@ class CaseConfigInput(BaseModel):
579651
in [
580652
IndexType.IVFFlat.value,
581653
IndexType.IVFSQ8.value,
654+
IndexType.IVF_RABITQ.value,
582655
IndexType.GPU_IVF_FLAT.value,
583656
IndexType.GPU_IVF_PQ.value,
584657
IndexType.GPU_BRUTE_FORCE.value,
@@ -593,7 +666,8 @@ class CaseConfigInput(BaseModel):
593666
"max": 65536,
594667
"value": 0,
595668
},
596-
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) in [IndexType.GPU_IVF_PQ.value],
669+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
670+
in [IndexType.GPU_IVF_PQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value],
597671
)
598672

599673

@@ -605,7 +679,19 @@ class CaseConfigInput(BaseModel):
605679
"max": 65536,
606680
"value": 8,
607681
},
608-
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) in [IndexType.GPU_IVF_PQ.value],
682+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None)
683+
in [IndexType.GPU_IVF_PQ.value, IndexType.HNSW_PQ.value, IndexType.HNSW_PRQ.value],
684+
)
685+
686+
CaseConfigParamInput_NRQ = CaseConfigInput(
687+
label=CaseConfigParamType.nrq,
688+
inputType=InputType.Number,
689+
inputConfig={
690+
"min": 1,
691+
"max": 16,
692+
"value": 2,
693+
},
694+
isDisplayed=lambda config: config.get(CaseConfigParamType.IndexType, None) in [IndexType.HNSW_PRQ.value],
609695
)
610696

611697
CaseConfigParamInput_intermediate_graph_degree = CaseConfigInput(
@@ -1186,6 +1272,10 @@ class CaseConfigInput(BaseModel):
11861272
CaseConfigParamInput_graph_degree,
11871273
CaseConfigParamInput_build_algo,
11881274
CaseConfigParamInput_cache_dataset_on_device,
1275+
CaseConfigParamInput_SQType,
1276+
CaseConfigParamInput_Refine,
1277+
CaseConfigParamInput_RefineType,
1278+
CaseConfigParamInput_NRQ,
11891279
]
11901280
MilvusPerformanceConfig = [
11911281
CaseConfigParamInput_IndexType,
@@ -1197,6 +1287,8 @@ class CaseConfigInput(BaseModel):
11971287
CaseConfigParamInput_Nprobe,
11981288
CaseConfigParamInput_M_PQ,
11991289
CaseConfigParamInput_Nbits_PQ,
1290+
CaseConfigParamInput_RBQBitsQuery,
1291+
CaseConfigParamInput_NRQ,
12001292
CaseConfigParamInput_intermediate_graph_degree,
12011293
CaseConfigParamInput_graph_degree,
12021294
CaseConfigParamInput_itopk_size,
@@ -1207,6 +1299,10 @@ class CaseConfigInput(BaseModel):
12071299
CaseConfigParamInput_build_algo,
12081300
CaseConfigParamInput_cache_dataset_on_device,
12091301
CaseConfigParamInput_refine_ratio,
1302+
CaseConfigParamInput_SQType,
1303+
CaseConfigParamInput_Refine,
1304+
CaseConfigParamInput_RefineType,
1305+
CaseConfigParamInput_RefineK,
12101306
]
12111307

12121308
WeaviateLoadConfig = [

0 commit comments

Comments
 (0)