diff --git a/vectordb_bench/backend/clients/doris/config.py b/vectordb_bench/backend/clients/doris/config.py index bd7c91c7b..a15309922 100644 --- a/vectordb_bench/backend/clients/doris/config.py +++ b/vectordb_bench/backend/clients/doris/config.py @@ -35,9 +35,11 @@ def to_dict(self) -> dict: class DorisCaseConfig(BaseModel, DBCaseConfig): metric_type: MetricType | None = None - # Optional explicit HNSW params for convenience + # Optional explicit HNSW/IVF params for convenience + index_type: str | None = None m: int | None = None ef_construction: int | None = None + nlist: int | None = None # Arbitrary index properties and session variables index_properties: dict[str, str] | None = None session_vars: dict[str, str] | None = None @@ -59,15 +61,30 @@ def get_metric_fn(self) -> str: def index_param(self) -> dict: # Use exact metric function name for index creation by removing '_approximate' suffix - metric_fn = self.get_metric_fn() - if metric_fn.endswith("_approximate"): - metric_fn = metric_fn[: -len("_approximate")] - props = {"metric_fn": metric_fn} + metric_type = self.get_metric_fn() + if metric_type.endswith("_approximate"): + metric_type = metric_type[: -len("_approximate")] + props = {"metric_type": metric_type} + + if self.index_type is not None: + props.setdefault("index_type", self.index_type) + else: + props.setdefault("index_type", "hnsw") + # Merge optional HNSW params - if self.m is not None: - props.setdefault("max_degree", str(self.m)) - if self.ef_construction is not None: - props.setdefault("ef_construction", str(self.ef_construction)) + props["index_type"] = str.lower(props["index_type"]) + if props["index_type"] == "hnsw": + if self.m is not None: + props.setdefault("max_degree", str(self.m)) + if self.ef_construction is not None: + props.setdefault("ef_construction", str(self.ef_construction)) + elif props["index_type"] == "ivf": + if self.nlist is not None: + props.setdefault("nlist", str(self.nlist)) + else: + msg = f"Unsupported index type: {props['index_type']}" + raise ValueError(msg) + # Merge user provided index_properties if self.index_properties: props.update(self.index_properties) diff --git a/vectordb_bench/backend/clients/doris/doris.py b/vectordb_bench/backend/clients/doris/doris.py index 1805a60c9..82b3a12da 100644 --- a/vectordb_bench/backend/clients/doris/doris.py +++ b/vectordb_bench/backend/clients/doris/doris.py @@ -166,36 +166,26 @@ def _create_table(self): def _build_index_options(self) -> IndexOptions | None: index_param = self.case_config.index_param() - metric_type = index_param.get("metric_fn", "l2_distance") - index_options = IndexOptions( - index_type="hnsw", - metric_type=metric_type, - dim=self.dim, + index_options = IndexOptions() + + applied, not_applied = {}, {} + for key, value in index_param.items(): + attr_name = key + if hasattr(index_options, attr_name): + try: + setattr(index_options, attr_name, value) + applied[key] = value + except Exception: + not_applied[key] = value + else: + not_applied[key] = value + + log.info( + "Index options prepared: applied_props=%s not_applied_props=%s", + applied, + not_applied, ) - extra_props = {k: v for k, v in index_param.items() if k != "metric_fn"} - if extra_props: - applied, stored = {}, {} - for key, value in extra_props.items(): - attr_name = key - if hasattr(index_options, attr_name): - try: - setattr(index_options, attr_name, value) - applied[key] = value - except Exception: - stored[key] = value - else: - stored[key] = value - if stored: - index_options.properties = {**getattr(index_options, "properties", {}), **stored} - log.info( - "Index options prepared: metric=%s applied_props=%s stored_props=%s", - metric_type, - applied, - stored, - ) - else: - log.info("Index options prepared: metric=%s (no extra properties)", metric_type) log.info("Creating table %s with index %s", self.table_name, index_param) return index_options