Skip to content

Commit cf39b77

Browse files
uchenilyAkhil-Pathivada
authored andcommitted
doris: update index properties (zilliztech#664)
1 parent cd53c25 commit cf39b77

File tree

2 files changed

+44
-37
lines changed

2 files changed

+44
-37
lines changed

vectordb_bench/backend/clients/doris/config.py

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -35,9 +35,11 @@ def to_dict(self) -> dict:
3535

3636
class DorisCaseConfig(BaseModel, DBCaseConfig):
3737
metric_type: MetricType | None = None
38-
# Optional explicit HNSW params for convenience
38+
# Optional explicit HNSW/IVF params for convenience
39+
index_type: str | None = None
3940
m: int | None = None
4041
ef_construction: int | None = None
42+
nlist: int | None = None
4143
# Arbitrary index properties and session variables
4244
index_properties: dict[str, str] | None = None
4345
session_vars: dict[str, str] | None = None
@@ -59,15 +61,30 @@ def get_metric_fn(self) -> str:
5961

6062
def index_param(self) -> dict:
6163
# Use exact metric function name for index creation by removing '_approximate' suffix
62-
metric_fn = self.get_metric_fn()
63-
if metric_fn.endswith("_approximate"):
64-
metric_fn = metric_fn[: -len("_approximate")]
65-
props = {"metric_fn": metric_fn}
64+
metric_type = self.get_metric_fn()
65+
if metric_type.endswith("_approximate"):
66+
metric_type = metric_type[: -len("_approximate")]
67+
props = {"metric_type": metric_type}
68+
69+
if self.index_type is not None:
70+
props.setdefault("index_type", self.index_type)
71+
else:
72+
props.setdefault("index_type", "hnsw")
73+
6674
# Merge optional HNSW params
67-
if self.m is not None:
68-
props.setdefault("max_degree", str(self.m))
69-
if self.ef_construction is not None:
70-
props.setdefault("ef_construction", str(self.ef_construction))
75+
props["index_type"] = str.lower(props["index_type"])
76+
if props["index_type"] == "hnsw":
77+
if self.m is not None:
78+
props.setdefault("max_degree", str(self.m))
79+
if self.ef_construction is not None:
80+
props.setdefault("ef_construction", str(self.ef_construction))
81+
elif props["index_type"] == "ivf":
82+
if self.nlist is not None:
83+
props.setdefault("nlist", str(self.nlist))
84+
else:
85+
msg = f"Unsupported index type: {props['index_type']}"
86+
raise ValueError(msg)
87+
7188
# Merge user provided index_properties
7289
if self.index_properties:
7390
props.update(self.index_properties)

vectordb_bench/backend/clients/doris/doris.py

Lines changed: 18 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -166,36 +166,26 @@ def _create_table(self):
166166

167167
def _build_index_options(self) -> IndexOptions | None:
168168
index_param = self.case_config.index_param()
169-
metric_type = index_param.get("metric_fn", "l2_distance")
170-
index_options = IndexOptions(
171-
index_type="hnsw",
172-
metric_type=metric_type,
173-
dim=self.dim,
169+
index_options = IndexOptions()
170+
171+
applied, not_applied = {}, {}
172+
for key, value in index_param.items():
173+
attr_name = key
174+
if hasattr(index_options, attr_name):
175+
try:
176+
setattr(index_options, attr_name, value)
177+
applied[key] = value
178+
except Exception:
179+
not_applied[key] = value
180+
else:
181+
not_applied[key] = value
182+
183+
log.info(
184+
"Index options prepared: applied_props=%s not_applied_props=%s",
185+
applied,
186+
not_applied,
174187
)
175188

176-
extra_props = {k: v for k, v in index_param.items() if k != "metric_fn"}
177-
if extra_props:
178-
applied, stored = {}, {}
179-
for key, value in extra_props.items():
180-
attr_name = key
181-
if hasattr(index_options, attr_name):
182-
try:
183-
setattr(index_options, attr_name, value)
184-
applied[key] = value
185-
except Exception:
186-
stored[key] = value
187-
else:
188-
stored[key] = value
189-
if stored:
190-
index_options.properties = {**getattr(index_options, "properties", {}), **stored}
191-
log.info(
192-
"Index options prepared: metric=%s applied_props=%s stored_props=%s",
193-
metric_type,
194-
applied,
195-
stored,
196-
)
197-
else:
198-
log.info("Index options prepared: metric=%s (no extra properties)", metric_type)
199189
log.info("Creating table %s with index %s", self.table_name, index_param)
200190
return index_options
201191

0 commit comments

Comments
 (0)