diff --git a/README.md b/README.md index b08458c79..56453f3e1 100644 --- a/README.md +++ b/README.md @@ -609,26 +609,65 @@ To list the options for Lindorm, execute `vectordbbench lindormhnsw --help`, The --reorder-factor INTEGER reorder factor ``` -### Run ADBPG (Aliyun AnalyticDB for PostgreSQL) from command line +### Run ADBPG (Aliyun AnalyticDB for PostgreSQL) ADBPG Nova uses the fastann/Nova vector index engine with `USING ann` syntax. -**Example: Run novamr index benchmark (BioASQ 1M, 1024-dim)** +**Example: Run from config file** ```shell -vectordbbench adbpgnova --case-type Performance1024D1M --k 10 \ ---host --port 5432 --db-name postgres \ ---user-name --password \ ---algorithm novamr --hnsw-m 48 --ef-construction 600 \ ---ef-search 130 --max-scan-points 5000 --quantize-rescore-amp 2.0 +vectordbbench adbpgnova --config-file vectordb_bench/config-files/adbpg_cohere1m_novamr.yml ``` -**Example: Run from config file** +**Example: Run from command line with structured parameters** ```shell -vectordbbench adbpgnova --config-file adbpg_bioasq1m_novamr.yml +vectordbbench adbpgnova \ + --case-type Performance1024D1M \ + --k 10 \ + --host \ + --port 5432 \ + --db-name postgres \ + --user-name \ + --password \ + --build-parameters \ +'{ + "reloption": { + "algorithm": "novamr", + "hnsw_m": 48, + "hnsw_ef_construction": 600 + }, + "guc": { + "fastann.build_parallel_processes": 32 + } +}' \ + --search-parameters \ +'{ + "reloption": { + "nova_autotune_topk": 10, + "nova_autotune_recall": 0.95 + }, + "guc": { + "fastann.hnsw_ef_search": 130, + "fastann.hnsw_max_scan_points": 5000, + "fastann.quantize_rescore_amp": 2.0 + } +}' \ + --autotune-parameters \ +'{ + "topk": [10, 100], + "target_recall": [0.90, 0.95, 0.99] +}' ``` +| Parameter | When it is applied | +| --- | --- | +| `build_parameters.reloption` | Added to the index `WITH (...)` clause during index construction. | +| `build_parameters.guc` | Set in the build session before `CREATE INDEX`. | +| `search_parameters.reloption` | Set on the index once after a new build and optimize, or once before search when reusing an existing index. | +| `search_parameters.guc` | Set when each search connection is initialized, before its first query. | +| `autotune_parameters` | Applied during optimization after a successful index build. Omit it to skip autotune. | + To list the options for ADBPG, execute `vectordbbench adbpgnova --help`. The following are some ADBPG-specific command-line options. ```text @@ -637,14 +676,9 @@ To list the options for ADBPG, execute `vectordbbench adbpgnova --help`. The fol --host TEXT Db host [required] --port INTEGER Postgres database port [default: 5432] --db-name TEXT Db name [required] - --algorithm TEXT algorithm [default: novamr] - --hnsw-m INTEGER hnsw_m [default: 16] - --ef-construction INTEGER ef_construction [default: 200] - --ef-search INTEGER ef_search [default: 100] - --max-scan-points INTEGER max scan points [default: 2000] - --quantize-rescore-amp FLOAT fastann.quantize_rescore_amp [default: 1.0] - --nova-adaptive-gamma FLOAT fastann.nova_adaptive_gamma [default: 0.0] - --auto-reduction/--no-auto-reduction Index WITH auto_reduction=on [default: False] + --build-parameters YAML_MAPPING Build reloption/GUC mapping as YAML or JSON + --search-parameters YAML_MAPPING Search reloption/GUC mapping as YAML or JSON + --autotune-parameters YAML_MAPPING Autotune mapping as YAML or JSON; omit to disable ``` ### Run PolarDB from command line diff --git a/tests/test_adbpg.py b/tests/test_adbpg.py index 465acf762..c203fcf57 100644 --- a/tests/test_adbpg.py +++ b/tests/test_adbpg.py @@ -1,66 +1,143 @@ -"""Unit tests for the ADB-PG Nova client config layer. - -These tests do not require a live database — they only exercise: - - AdbpgConfig defaults and connection-string assembly - - AdbpgIndexConfig.index_param() WITH-clause options (incl. raw auto_reduction) - - AdbpgIndexConfig.session_param() fastann GUC emission - - TestResult.read_file() round-trip when password is absent in saved JSON - (regression for the result-loading failure caused by polymorphic - serialization stripping subclass fields from DBConfig) - -Usage: - pytest tests/test_adbpg.py -v -""" - from __future__ import annotations import json -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any +import psycopg import pytest -from pydantic import SecretStr +from click.testing import CliRunner +from pydantic import SecretStr, ValidationError from vectordb_bench.backend.clients import DB -from vectordb_bench.backend.clients.adbpg.config import AdbpgConfig, AdbpgIndexConfig +from vectordb_bench.backend.clients.adbpg import cli as adbpg_cli +from vectordb_bench.backend.clients.adbpg.adbpg import Adbpg, AdbpgTimeoutError +from vectordb_bench.backend.clients.adbpg.config import ( + AdbpgAutotuneParameters, + AdbpgConfig, + AdbpgIndexConfig, + AdbpgParameterGroup, +) from vectordb_bench.backend.clients.api import MetricType -from vectordb_bench.models import TestResult +from vectordb_bench.backend.data_source import DatasetSource +from vectordb_bench.backend.filter import non_filter +from vectordb_bench.backend.task_runner import CaseRunner, RunningStatus +from vectordb_bench.cli.batch_cli import build_sub_cmd_args +from vectordb_bench.models import CaseConfig, CaseType, TaskConfig, TaskStage, TestResult if TYPE_CHECKING: from pathlib import Path -def make_index_config(**overrides) -> AdbpgIndexConfig: +def make_index_config(**overrides: Any) -> AdbpgIndexConfig: base = { "metric_type": MetricType.COSINE, - "hnsw_m": 32, - "ef_search": 100, - "ef_construction": 200, - "nlist": 1024, - "algorithm": "novamr", - "rabitq_bits": 7, - "quantize_rescore_amp": 1.0, - "nova_adaptive_gamma": 0.0, - "max_scan_points": 2000, - "index_scan_mode": "snapshot", - "auto_reduction": False, - "nprobe": 5, + "build_parameters": { + "reloption": { + "algorithm": "novamr", + "note": "quotes ', comma, and $$ stay exact", + }, + "guc": {}, + }, + "search_parameters": { + "reloption": { + "nova_autotune_topk": 10, + "nova_autotune_recall": 0.95, + }, + "guc": {"search_path": "foo,bar"}, + }, + "autotune_parameters": { + "topk": [10, 100], + "target_recall": [0.9, 0.95], + "n_threads": 4, + "future_optimizer": {"levels": [1, 2]}, + }, } base.update(overrides) return AdbpgIndexConfig(**base) +class FakeCursor: + def __init__( + self, + fail_on: str | None = None, + error: Exception | None = None, + fetchone_result: tuple[Any, ...] = (1, True, 2, True, True), + ): + self.executions: list[tuple[Any, Any]] = [] + self.fail_on = fail_on + self.error = error or RuntimeError("database error") + self.fetchone_result = fetchone_result + self.connection: FakeConnection | None = None + self.closed = False + + def execute(self, query: Any, params: Any = None, **kwargs: Any): + self.executions.append((query, params)) + if self.connection is not None and self.connection.aborted: + raise psycopg.errors.InFailedSqlTransaction + if self.fail_on and self.fail_on in str(query): + if self.connection is not None: + self.connection.aborted = True + raise self.error + return self + + def fetchall(self) -> list[tuple[str]]: + return [("ok",)] + + def fetchone(self) -> tuple[Any, ...]: + return self.fetchone_result + + def close(self) -> None: + self.closed = True + + +class FakeConnection: + def __init__(self, cursor: FakeCursor): + self._cursor = cursor + self.commit_count = 0 + self.rollback_count = 0 + self.aborted = False + self.closed = False + cursor.connection = self + + def cursor(self) -> FakeCursor: + return self._cursor + + def commit(self) -> None: + self.commit_count += 1 + + def rollback(self) -> None: + self.rollback_count += 1 + self.aborted = False + + def close(self) -> None: + self.closed = True + + +def make_client(*, with_scalar_labels: bool = False) -> Adbpg: + client = Adbpg.__new__(Adbpg) + client.name = "Adbpg" + client.case_config = make_index_config() + client.table_name = "vector$table" + client.connect_config = {} + client.dim = 768 + client.with_scalar_labels = with_scalar_labels + client._primary_field = "id" + client._vector_field = "embedding" + client._scalar_label_field = "label" + client._index_name = "vector$table_novamr_index" + client.where_clause = "" + return client + + class TestAdbpgConfig: def test_defaults_allow_construction_without_password(self): - # Regression: result JSON only contains DBConfig parent fields - # (db_label/version/note) because of pydantic polymorphic serialization. - # AdbpgConfig must therefore be constructible from that minimal dict. cfg = AdbpgConfig(db_label="", version="", note="") assert cfg.host == "localhost" assert cfg.port == 5432 assert cfg.db_name == "postgres" assert cfg.password.get_secret_value() == "" - def test_to_dict_carries_utility_session_option(self): + def test_to_dict_uses_normal_coordinator_session(self): cfg = AdbpgConfig( user_name=SecretStr("u"), password=SecretStr("pw"), @@ -68,95 +145,584 @@ def test_to_dict_carries_utility_session_option(self): port=5432, db_name="postgres", ) - d = cfg.to_dict() - assert d["table_name"] == "vector" - cc = d["connect_config"] - assert cc["host"] == "h.example.com" - assert cc["user"] == "u" - assert cc["password"] == "pw" # noqa: S105 - assert cc["dbname"] == "postgres" - assert cc["options"] == "-c gp_session_role=utility" - - -class TestAdbpgIndexConfigBuild: - def test_parse_metric(self): - assert make_index_config(metric_type=MetricType.L2).parse_metric() == "l2" - assert make_index_config(metric_type=MetricType.COSINE).parse_metric() == "cosine" - assert make_index_config(metric_type=MetricType.IP).parse_metric() == "ip" - - def test_parse_metric_unsupported_raises(self): - with pytest.raises(ValueError, match="Metric type"): - make_index_config(metric_type=None).parse_metric() - - def test_index_param_options_default(self): - params = make_index_config().index_param() - names = {opt["option_name"]: opt for opt in params["index_creation_with_options"]} - assert names["algorithm"]["val"] == "novamr" - assert names["hnsw_m"]["val"] == 32 - assert names["hnsw_ef_construction"]["val"] == 200 - assert names["nlist"]["val"] == 1024 - assert names["rabitq_bits"]["val"] == 7 - assert names["max_key_len"]["val"] == 1 - # auto_reduction is omitted when False - assert "auto_reduction" not in names - - def test_index_param_auto_reduction_emits_raw(self): - params = make_index_config(auto_reduction=True).index_param() - opt = next(o for o in params["index_creation_with_options"] if o["option_name"] == "auto_reduction") - # `raw=True` so the value is rendered as a bare SQL identifier (`on`) - # rather than a quoted literal. - assert opt["val"] == "on" - assert opt.get("raw") is True - - def test_index_param_pca_dim_omitted_when_none(self): - params = make_index_config(pca_dim=None).index_param() - names = {opt["option_name"] for opt in params["index_creation_with_options"]} - assert "pca_dim" not in names - - def test_index_param_pca_dim_emitted_when_set(self): - params = make_index_config(pca_dim=448).index_param() - opt = next(o for o in params["index_creation_with_options"] if o["option_name"] == "pca_dim") - assert opt["val"] == 448 - - -class TestAdbpgIndexConfigSession: - def test_session_param_emits_all_search_gucs(self): - cfg = make_index_config( - quantize_rescore_amp=0.6, - nova_adaptive_gamma=0.0, - ef_search=50, - max_scan_points=16000, - index_scan_mode="snapshot", - nprobe=64, + connect_config = cfg.to_dict()["connect_config"] + assert connect_config == { + "host": "h.example.com", + "port": 5432, + "dbname": "postgres", + "user": "u", + "password": "pw", + } + + +class TestAdbpgStructuredConfig: + def test_free_form_names_and_scalar_types_are_preserved(self): + group = AdbpgParameterGroup( + reloption={"future.knob": 7, "enabled": True, "reset_me": None}, + guc={"search_path": "foo,bar", "ratio": 0.5}, ) - opts = cfg.session_param()["session_options"] - emitted = {o["parameter"]["setting_name"]: o["parameter"]["val"] for o in opts} - assert emitted["fastann.quantize_rescore_amp"] == "0.6" - assert emitted["fastann.nova_adaptive_gamma"] == "0.0" - assert emitted["fastann.hnsw_ef_search"] == "50" - assert emitted["fastann.hnsw_max_scan_points"] == "16000" - assert emitted["fastann.index_scan_mode"] == "snapshot" - # novad-specific GUC is always emitted (no-op for HNSW algorithms) - assert emitted["fastann.nova_nprobe"] == "64" - - def test_session_param_emits_zero_values(self): - # Forcing 0 / 0.0 must still produce a SET command — callers rely on - # being able to pin a GUC to zero. - cfg = make_index_config(quantize_rescore_amp=0.0, nova_adaptive_gamma=0.0, nprobe=0) - opts = cfg.session_param()["session_options"] - emitted = {o["parameter"]["setting_name"]: o["parameter"]["val"] for o in opts} - assert emitted["fastann.quantize_rescore_amp"] == "0.0" - assert emitted["fastann.nova_adaptive_gamma"] == "0.0" - assert emitted["fastann.nova_nprobe"] == "0" + assert group.model_dump() == { + "reloption": {"future.knob": 7, "enabled": True, "reset_me": None}, + "guc": {"search_path": "foo,bar", "ratio": 0.5}, + } + @pytest.mark.parametrize( + "value", + [ + {"reloption": {"bad": [1, 2]}}, + {"guc": {"bad": {"nested": True}}}, + {"unsupported": {}}, + ], + ) + def test_invalid_parameter_structures_are_rejected(self, value: dict[str, Any]): + with pytest.raises(ValidationError): + AdbpgParameterGroup(**value) -class TestResultRoundTrip: - def test_read_file_with_minimal_db_config(self, tmp_path: Path): - """Saved result JSON keeps only DBConfig parent fields for db_config. + def test_build_defaults_are_merged_with_user_parameters(self): + build = AdbpgIndexConfig( + build_parameters={"reloption": {"algorithm": "novamr"}}, + ).build_parameters + assert build.model_dump() == { + "reloption": { + "algorithm": "novamr", + "hnsw_m": 48, + "hnsw_ef_construction": 600, + "rabitq_bits": 7, + "auto_reduction": False, + }, + "guc": {"fastann.build_parallel_processes": 32}, + } + + def test_default_algorithm_is_used_for_every_entry_point(self): + assert AdbpgIndexConfig().algorithm == "novamr" + + def test_autotune_defaults_and_preserves_multiple_targets(self): + autotune = AdbpgAutotuneParameters( + topk=[10, 100], + target_recall=[0.9, 0.95], + ) + assert autotune.topk == [10, 100] + assert autotune.target_recall == [0.9, 0.95] + assert autotune.n_samples == 200 + assert autotune.n_trials == 200 + assert autotune.timeout == 600 + + def test_autotune_preserves_extra_json_parameters_and_allows_zero_samples(self): + autotune = AdbpgAutotuneParameters( + topk=[10], + target_recall=[0.95], + n_samples=0, + n_trials=100, + timeout=600, + n_threads=4, + future_optimizer={"levels": [1, 2]}, + future_optional=None, + ) + assert autotune.model_dump(mode="json")["n_threads"] == 4 + assert autotune.model_dump(mode="json")["future_optimizer"] == {"levels": [1, 2]} + assert "future_optional" in autotune.model_dump(mode="json") + assert autotune.model_dump(mode="json")["future_optional"] is None + + def test_autotune_is_disabled_when_configuration_is_omitted(self): + assert AdbpgIndexConfig().autotune_parameters is None + assert AdbpgIndexConfig(autotune_parameters="").autotune_parameters is None + + def test_legacy_enable_switch_is_rejected(self): + with pytest.raises(ValidationError, match="omit autotune_parameters"): + AdbpgAutotuneParameters( + enable=False, + topk=[10], + target_recall=[0.95], + ) + + @pytest.mark.parametrize( + "value", + [ + {}, + {"topk": [0], "target_recall": [0.95], "n_samples": 1, "n_trials": 1, "timeout": 1}, + {"topk": [10], "target_recall": [1.1], "n_samples": 1, "n_trials": 1, "timeout": 1}, + {"topk": [10], "target_recall": [0.95], "n_samples": -1, "n_trials": 1, "timeout": 1}, + ], + ) + def test_invalid_autotune_is_rejected(self, value: dict[str, Any]): + with pytest.raises(ValidationError): + AdbpgAutotuneParameters(**value) + + def test_old_flattened_fields_are_migrated_to_structured_parameters(self): + config = AdbpgIndexConfig( + algorithm="novad", + hnsw_m=32, + ef_construction=256, + nlist=2048, + rabitq_bits=6, + auto_reduction=True, + pca_dim=384, + build_parallel_processes=16, + ef_search=77, + max_scan_points=900, + quantize_rescore_amp=2.0, + nova_adaptive_gamma=0.4, + index_scan_mode="streaming", + nprobe=9, + ) + assert config.build_parameters.reloption == { + "algorithm": "novad", + "hnsw_m": 32, + "hnsw_ef_construction": 256, + "rabitq_bits": 6, + "auto_reduction": True, + "nlist": 2048, + "pca_dim": 384, + } + assert config.build_parameters.guc == {"fastann.build_parallel_processes": 16} + assert config.search_parameters.guc == { + "fastann.hnsw_ef_search": 77, + "fastann.hnsw_max_scan_points": 900, + "fastann.quantize_rescore_amp": 2.0, + "fastann.nova_adaptive_gamma": 0.4, + "fastann.index_scan_mode": "streaming", + "fastann.nova_nprobe": 9, + } + + def test_structured_parameters_take_precedence_over_legacy_fields(self): + config = AdbpgIndexConfig( + algorithm="novad", + ef_search=77, + build_parameters={"reloption": {"algorithm": "novamr"}}, + search_parameters={"guc": {"fastann.hnsw_ef_search": 130}}, + ) + assert config.algorithm == "novamr" + assert config.search_parameters.guc == {"fastann.hnsw_ef_search": 130} + + def test_unknown_fields_are_still_rejected(self): + with pytest.raises(ValidationError, match="unknown"): + AdbpgIndexConfig(unknown=1) + + +class TestAdbpgUdfPayload: + def test_pure_vector_payload_is_complete_and_excludes_credentials(self): + client = make_client() + payload = client._build_udf_payload() + assert payload["api_version"] == 1 + assert {"case_type", "workload_type"}.isdisjoint(payload) + assert payload["relation"] == { + "schema": "public", + "table": "vector$table", + "index": "vector$table_novamr_index", + } + assert payload["columns"] == [ + {"name": "id", "type": "bigint", "role": "primary_key"}, + {"name": "embedding", "type": "vector(768)", "role": "vector"}, + ] + assert payload["metric"] == "cosine" + assert set(payload["parameters"]) == {"build", "search", "autotune"} + assert payload["parameters"]["autotune"]["topk"] == [10, 100] + assert payload["parameters"]["autotune"]["n_samples"] == 200 + assert payload["parameters"]["autotune"]["n_trials"] == 200 + assert payload["parameters"]["autotune"]["timeout"] == 600 + assert payload["parameters"]["autotune"]["n_threads"] == 4 + assert payload["parameters"]["autotune"]["future_optimizer"] == {"levels": [1, 2]} + assert payload["parameters"]["build"]["reloption"]["note"] == "quotes ', comma, and $$ stay exact" + assert payload["parameters"]["search"] == { + "reloption": {"nova_autotune_topk": 10, "nova_autotune_recall": 0.95}, + "guc": {"search_path": "foo,bar"}, + } + assert {"password", "user", "user_name", "host", "port"}.isdisjoint(payload) + + def test_payload_omits_autotune_when_configuration_is_omitted(self): + client = make_client() + client.case_config = make_index_config(autotune_parameters=None) + + assert "autotune" not in client._build_udf_payload()["parameters"] + + def test_hybrid_payload_describes_filter_column(self): + client = make_client(with_scalar_labels=True) + payload = client._build_udf_payload() + assert payload["columns"][-1] == {"name": "label", "type": "varchar(64)", "role": "filter"} + + +class TestAdbpgLifecycle: + def test_optimize_calls_build_then_optimize_once_with_same_bound_payload(self): + client = make_client() + cursor = FakeCursor() + conn = FakeConnection(cursor) + client.conn = conn + client.cursor = cursor + + client.optimize() + + udf_calls = [(query, params) for query, params in cursor.executions if "vectordbbench_" in str(query)] + assert ["vectordbbench_build" in str(query) for query, _ in udf_calls] == [True, False] + assert ["vectordbbench_optimize" in str(query) for query, _ in udf_calls] == [False, True] + assert udf_calls[0][1][0].obj == udf_calls[1][1][0].obj + statements = [str(query) for query, _ in cursor.executions] + assert next(i for i, value in enumerate(statements) if "vectordbbench_optimize" in value) < next( + i for i, value in enumerate(statements) if "ALTER INDEX" in value + ) + assert sum("ALTER INDEX" in statement for statement in statements) == 1 + assert conn.commit_count == 3 + assert conn.rollback_count == 0 + assert any(params == ("630s",) for _, params in cursor.executions) + + def test_build_failure_rolls_back_and_skips_optimize(self): + client = make_client() + cursor = FakeCursor(fail_on="vectordbbench_build") + conn = FakeConnection(cursor) + client.conn = conn + client.cursor = cursor + + with pytest.raises(RuntimeError, match=r"fastann\.vectordbbench_build"): + client.optimize() + + assert conn.rollback_count == 1 + assert not any("vectordbbench_optimize" in str(query) for query, _ in cursor.executions) + + def test_optimize_failure_rolls_back_and_drops_incomplete_index(self): + client = make_client() + cursor = FakeCursor(fail_on="vectordbbench_optimize") + conn = FakeConnection(cursor) + client.conn = conn + client.cursor = cursor + + with pytest.raises(RuntimeError, match=r"fastann\.vectordbbench_optimize"): + client.optimize() + + assert conn.commit_count == 2 + assert conn.rollback_count == 2 + assert not any("ALTER INDEX" in str(query) for query, _ in cursor.executions) + assert sum("DROP INDEX" in str(query) for query, _ in cursor.executions) == 1 + + def test_optimize_statement_timeout_is_reported(self): + client = make_client() + cursor = FakeCursor( + fail_on="vectordbbench_optimize", + error=psycopg.errors.QueryCanceled("statement timeout"), + ) + conn = FakeConnection(cursor) + client.conn = conn + client.cursor = cursor + + with pytest.raises(AdbpgTimeoutError, match="630 seconds"): + client.optimize() - TestResult.read_file must still rehydrate the AdbpgConfig instance - without raising a Field-required pydantic ValidationError. - """ + assert conn.rollback_count == 2 + assert sum("DROP INDEX" in str(query) for query, _ in cursor.executions) == 1 + + def test_missing_udfs_fail_before_destructive_build_setup(self, monkeypatch: pytest.MonkeyPatch): + cursor = FakeCursor(fetchone_result=(None, None, 2, True, True)) + conn = FakeConnection(cursor) + monkeypatch.setattr(Adbpg, "_create_connection", staticmethod(lambda **_: (conn, cursor))) + + with pytest.raises(RuntimeError, match=r"vectordbbench_build\(jsonb\)"): + Adbpg( + dim=768, + db_config={"table_name": "vector", "connect_config": {}}, + db_case_config=make_index_config(), + drop_old=True, + ) + + assert len(cursor.executions) == 1 + assert "to_regprocedure" in str(cursor.executions[0][0]) + assert not any("DROP" in str(query) or "CREATE" in str(query) for query, _ in cursor.executions) + assert cursor.closed + assert conn.closed + + def test_missing_udf_privilege_fails_before_destructive_build_setup(self, monkeypatch: pytest.MonkeyPatch): + cursor = FakeCursor(fetchone_result=(1, False, 2, True, True)) + conn = FakeConnection(cursor) + monkeypatch.setattr(Adbpg, "_create_connection", staticmethod(lambda **_: (conn, cursor))) + + with pytest.raises(RuntimeError, match=r"vectordbbench_build\(jsonb\)"): + Adbpg( + dim=768, + db_config={"table_name": "vector", "connect_config": {}}, + db_case_config=make_index_config(), + drop_old=True, + ) + + assert len(cursor.executions) == 1 + assert cursor.closed + assert conn.closed + + def test_reloption_failure_rolls_back_before_dropping_incomplete_index(self): + client = make_client() + cursor = FakeCursor(fail_on="ALTER INDEX") + conn = FakeConnection(cursor) + client.conn = conn + client.cursor = cursor + + with pytest.raises(RuntimeError, match="database error"): + client.optimize() + + assert conn.rollback_count == 1 + assert sum("DROP INDEX" in str(query) for query, _ in cursor.executions) == 1 + assert conn.commit_count == 3 + + def test_query_cancellation_without_configured_timeout_is_not_mislabeled(self): + client = make_client() + cursor = FakeCursor( + fail_on="vectordbbench_build", + error=psycopg.errors.QueryCanceled("canceled by user"), + ) + conn = FakeConnection(cursor) + client.conn = conn + client.cursor = cursor + + with pytest.raises(RuntimeError, match="failed: canceled by user"): + client.optimize() + + def test_non_timeout_cancellation_with_configured_timeout_is_not_mislabeled(self): + client = make_client() + cursor = FakeCursor( + fail_on="vectordbbench_optimize", + error=psycopg.errors.QueryCanceled("canceling statement due to user request"), + ) + conn = FakeConnection(cursor) + client.conn = conn + client.cursor = cursor + + with pytest.raises(RuntimeError, match="failed: canceling statement due to user request"): + client.optimize() + + def test_optimize_without_autotune_still_calls_both_udfs(self): + client = make_client() + client.case_config = make_index_config(autotune_parameters=None) + cursor = FakeCursor() + conn = FakeConnection(cursor) + client.conn = conn + client.cursor = cursor + + client.optimize() + + udf_calls = [(query, params) for query, params in cursor.executions if "vectordbbench_" in str(query)] + assert ["vectordbbench_build" in str(query) for query, _ in udf_calls] == [True, False] + assert ["vectordbbench_optimize" in str(query) for query, _ in udf_calls] == [False, True] + assert "autotune" not in udf_calls[0][1][0].obj["parameters"] + assert not any("statement_timeout" in str(query) for query, _ in cursor.executions) + + def test_reuse_constructor_applies_reloptions_once_and_calls_no_udf(self, monkeypatch: pytest.MonkeyPatch): + cursor = FakeCursor() + conn = FakeConnection(cursor) + monkeypatch.setattr(Adbpg, "_create_connection", staticmethod(lambda **_: (conn, cursor))) + + Adbpg( + dim=768, + db_config={"table_name": "vector", "connect_config": {}}, + db_case_config=make_index_config( + search_parameters={ + "reloption": {"nova_autotune_topk": 10, "nova_autotune_recall": None}, + "guc": {}, + } + ), + drop_old=False, + ) + + statements = [str(query) for query, _ in cursor.executions] + assert sum("ALTER INDEX" in statement for statement in statements) == 2 + assert not any("vectordbbench_" in statement for statement in statements) + assert conn.commit_count == 1 + assert cursor.closed + assert conn.closed + + def test_search_connections_only_apply_search_gucs(self, monkeypatch: pytest.MonkeyPatch): + client = make_client() + connections = [] + + def create_connection(**_kwargs): + cursor = FakeCursor() + conn = FakeConnection(cursor) + connections.append((conn, cursor)) + return conn, cursor + + monkeypatch.setattr(Adbpg, "_create_connection", staticmethod(create_connection)) + for _ in range(2): + with client.init(): + client.prepare_filter(non_filter) + + assert len(connections) == 2 + for conn, cursor in connections: + assert [(query, params) for query, params in cursor.executions] == [ + ("SELECT set_config(%s, %s, false)", ("search_path", "foo,bar")) + ] + assert conn.commit_count == 1 + assert not any( + "ALTER INDEX" in str(query) or "vectordbbench_" in str(query) for query, _ in cursor.executions + ) + + def test_non_search_connections_do_not_apply_search_gucs(self, monkeypatch: pytest.MonkeyPatch): + client = make_client() + cursor = FakeCursor() + conn = FakeConnection(cursor) + monkeypatch.setattr(Adbpg, "_create_connection", staticmethod(lambda **_: (conn, cursor))) + + with client.init(): + pass + + assert cursor.executions == [] + assert conn.commit_count == 0 + + def test_udf_failure_prevents_search_runner_initialization(self, monkeypatch: pytest.MonkeyPatch): + task = TaskConfig( + db=DB.Adbpg, + db_config=AdbpgConfig(), + db_case_config=make_index_config(), + case_config=CaseConfig(case_id=CaseType.Performance768D1M, k=10), + stages=[TaskStage.DROP_OLD, TaskStage.LOAD, TaskStage.SEARCH_SERIAL], + ) + runner = CaseRunner( + run_id="run-id", + config=task, + ca=task.case_config.case, + status=RunningStatus.PENDING, + dataset_source=DatasetSource.S3, + ) + search_started = [] + monkeypatch.setattr(CaseRunner, "_load_data", lambda _self: (1_000_000, 1.0)) + + def fail_optimize(_self: CaseRunner) -> float: + raise RuntimeError("fastann.vectordbbench_build failed") + + monkeypatch.setattr(CaseRunner, "_optimize", fail_optimize) + monkeypatch.setattr(CaseRunner, "_init_search_runners", lambda _self: search_started.append(True)) + + with pytest.raises(RuntimeError, match=r"vectordbbench_build failed"): + runner._run_perf_case(drop_old=True) + assert search_started == [] + + +def test_single_and_batch_cli_preserve_equivalent_structured_models(monkeypatch: pytest.MonkeyPatch, tmp_path: Path): + captured = [] + monkeypatch.setattr( + adbpg_cli, + "run", + lambda **kwargs: captured.append((kwargs["db_config"], kwargs["db_case_config"])), + ) + values = { + "build_parameters": {"reloption": {"algorithm": "novamr", "hnsw_m": 48}, "guc": {}}, + "search_parameters": {"reloption": {"reset_me": None}, "guc": {"search_path": "foo,bar"}}, + "autotune_parameters": { + "topk": [10, 100], + "target_recall": [0.9, 0.95], + }, + } + config_path = tmp_path / "adbpg.yml" + config_path.write_text( + "adbpgnova:\n" + " db_label: cohere1m-novamr\n" + " user_name: tester\n" + " host: localhost\n" + " db_name: postgres\n" + f" build_parameters: {json.dumps(values['build_parameters'])}\n" + f" search_parameters: {json.dumps(values['search_parameters'])}\n" + f" autotune_parameters: {json.dumps(values['autotune_parameters'])}\n" + ) + result = CliRunner().invoke(adbpg_cli.AdbpgNova, ["--config-file", str(config_path), "--dry-run"]) + assert result.exit_code == 0, result.output + + batch_args = build_sub_cmd_args( + { + "adbpgnova": [ + { + "db_label": "cohere1m-novamr", + "user_name": "tester", + "host": "localhost", + "db_name": "postgres", + "dry_run": True, + **values, + } + ] + } + )[0] + result = CliRunner().invoke(adbpg_cli.AdbpgNova, batch_args[1:]) + assert result.exit_code == 0, result.output + assert captured[0][0].db_label == captured[1][0].db_label == "cohere1m-novamr" + assert captured[0][1].model_dump(mode="json") == captured[1][1].model_dump(mode="json") + + +def test_cli_omits_autotune_to_disable_it(monkeypatch: pytest.MonkeyPatch): + captured = [] + monkeypatch.setattr(adbpg_cli, "run", lambda **kwargs: captured.append(kwargs["db_case_config"])) + + result = CliRunner().invoke( + adbpg_cli.AdbpgNova, + ["--user-name", "tester", "--host", "localhost", "--db-name", "postgres"], + ) + + assert result.exit_code == 0, result.output + assert captured[0].autotune_parameters is None + + batch_args = build_sub_cmd_args( + { + "adbpgnova": [ + { + "user_name": "tester", + "host": "localhost", + "db_name": "postgres", + } + ] + } + )[0] + result = CliRunner().invoke(adbpg_cli.AdbpgNova, batch_args[1:]) + assert result.exit_code == 0, result.output + assert captured[1].autotune_parameters is None + + +def test_cli_accepts_formatted_multiline_json_parameters(monkeypatch: pytest.MonkeyPatch): + captured = [] + monkeypatch.setattr(adbpg_cli, "run", lambda **kwargs: captured.append(kwargs["db_case_config"])) + build_parameters = """{ + "reloption": { + "algorithm": "novamr", + "hnsw_m": 48, + "hnsw_ef_construction": 600 + }, + "guc": {} +}""" + search_parameters = """{ + "reloption": {}, + "guc": { + "fastann.hnsw_ef_search": 130, + "fastann.hnsw_max_scan_points": 5000, + "fastann.quantize_rescore_amp": 2.0 + } +}""" + + result = CliRunner().invoke( + adbpg_cli.AdbpgNova, + [ + "--case-type", + "Performance1024D1M", + "--k", + "10", + "--host", + "localhost", + "--db-name", + "postgres", + "--user-name", + "tester", + "--password", + "password", + "--build-parameters", + build_parameters, + "--search-parameters", + search_parameters, + ], + ) + + assert result.exit_code == 0, result.output + config = captured[0] + assert config.build_parameters.reloption["algorithm"] == "novamr" + assert config.build_parameters.reloption["hnsw_m"] == 48 + assert config.build_parameters.reloption["hnsw_ef_construction"] == 600 + assert config.search_parameters.guc == { + "fastann.hnsw_ef_search": 130, + "fastann.hnsw_max_scan_points": 5000, + "fastann.quantize_rescore_amp": 2.0, + } + + +class TestResultRoundTrip: + def test_read_file_rehydrates_structured_adbpg_config(self, tmp_path: Path): result_dir = tmp_path / "AnalyticDB for PostgreSQL" result_dir.mkdir() result_file = result_dir / "result_test_run.json" @@ -184,14 +750,7 @@ def test_read_file_with_minimal_db_config(self, tmp_path: Path): "task_config": { "db": DB.Adbpg.value, "db_config": {"db_label": "", "version": "", "note": ""}, - "db_case_config": { - "metric_type": "COSINE", - "algorithm": "novamr", - "hnsw_m": 16, - "ef_search": 100, - "ef_construction": 200, - "nlist": 1024, - }, + "db_case_config": make_index_config().model_dump(mode="json"), "case_config": {"case_id": 5, "custom_case": {}, "k": 10}, "stages": ["search_serial"], "load_concurrency": 0, @@ -203,8 +762,7 @@ def test_read_file_with_minimal_db_config(self, tmp_path: Path): } result_file.write_text(json.dumps(payload)) - tr = TestResult.read_file(result_file, trans_unit=False) - assert len(tr.results) == 1 - rehydrated = tr.results[0].task_config.db_config - assert isinstance(rehydrated, AdbpgConfig) - assert rehydrated.host == "localhost" # came from default + result = TestResult.read_file(result_file, trans_unit=False) + rehydrated = result.results[0].task_config.db_case_config + assert isinstance(rehydrated, AdbpgIndexConfig) + assert rehydrated.model_dump(mode="json") == make_index_config().model_dump(mode="json") diff --git a/vectordb_bench/backend/clients/adbpg/adbpg.py b/vectordb_bench/backend/clients/adbpg/adbpg.py index c2e9bcc04..95ab928b1 100644 --- a/vectordb_bench/backend/clients/adbpg/adbpg.py +++ b/vectordb_bench/backend/clients/adbpg/adbpg.py @@ -1,7 +1,7 @@ """Wrapper around the Aliyun ADBPG (AnalyticDB for PostgreSQL) vector database.""" import logging -from collections.abc import Generator, Sequence +from collections.abc import Generator from contextlib import contextmanager from copy import copy from typing import Any @@ -10,6 +10,7 @@ import psycopg from pgvector.psycopg import register_vector from psycopg import Connection, Cursor, sql +from psycopg.types.json import Jsonb from vectordb_bench.backend.filter import Filter, FilterOp @@ -18,6 +19,12 @@ log = logging.getLogger(__name__) +UDF_CLEANUP_GRACE_SECONDS = 30 + + +class AdbpgTimeoutError(RuntimeError): + """ADBPG canceled a database statement after its configured timeout.""" + class Adbpg(VectorDB): """ADBPG vector database client, using psycopg.""" @@ -68,31 +75,20 @@ def __init__( # construct basic units self.conn, self.cursor = self._create_connection(**self.connect_config) - log.info(f"{self.name} config values: {self.connect_config}\n{self.case_config}") - if not any( - ( - self.case_config.create_index_before_load, - self.case_config.create_index_after_load, - ), - ): - msg = ( - f"{self.name} config must create an index using create_index_before_load or create_index_after_load" - f"{self.name} config values: {self.connect_config}\n{self.case_config}" - ) - log.error(msg) - raise RuntimeError(msg) - - if drop_old: - self._drop_index() - self._drop_table() - self._create_table(dim) - if self.case_config.create_index_before_load: - self._create_index() - - self.cursor.close() - self.conn.close() - self.cursor = None - self.conn = None + log.info("%s case config: %s", self.name, self.case_config) + try: + if drop_old: + self._check_required_udfs() + self._drop_index() + self._drop_table() + self._create_table(dim) + else: + self._apply_search_reloptions() + finally: + self.cursor.close() + self.conn.close() + self.cursor = None + self.conn = None @staticmethod def _create_connection(**kwargs) -> tuple[Connection, Cursor]: @@ -144,21 +140,9 @@ def _generate_search_query(self) -> sql.Composed: @contextmanager def init(self) -> Generator[None, None, None]: - """Open a session, apply GUCs, yield, then close.""" + """Open a database session, yield, then close it.""" self.conn, self.cursor = self._create_connection(**self.connect_config) - session_options: Sequence[dict[str, Any]] = self.case_config.session_param()["session_options"] - - if len(session_options) > 0: - for setting in session_options: - command = sql.SQL("SET {setting_name} = {val};").format( - setting_name=sql.Identifier(setting["parameter"]["setting_name"]), - val=sql.Identifier(str(setting["parameter"]["val"])), - ) - log.debug(command.as_string(self.cursor)) - self.cursor.execute(command) - self.conn.commit() - try: yield finally: @@ -180,94 +164,166 @@ def _drop_table(self): self.conn.commit() def optimize(self, data_size: int | None = None): - self._post_insert() - - def _post_insert(self): - log.info(f"{self.name} post insert before optimize") - if self.case_config.create_index_after_load: - self._drop_index() - self._create_index() + payload = self._build_udf_payload() + self._call_udf("vectordbbench_build", payload) + try: + autotune = self.case_config.autotune_parameters + statement_timeout = autotune.timeout + UDF_CLEANUP_GRACE_SECONDS if autotune is not None else None + self._call_udf("vectordbbench_optimize", payload, timeout=statement_timeout) + self._apply_search_reloptions() + except Exception: + try: + self.conn.rollback() + self._drop_index() + except Exception: + log.exception("Failed to remove incomplete ADBPG index %s", self._index_name) + raise + + def _check_required_udfs(self) -> None: + assert self.cursor is not None, "Cursor is not initialized" + signatures = ( + "fastann.vectordbbench_build(jsonb)", + "fastann.vectordbbench_optimize(jsonb)", + ) + row = self.cursor.execute( + """ + SELECT to_regprocedure(%s), has_function_privilege(to_regprocedure(%s), 'EXECUTE'), + to_regprocedure(%s), has_function_privilege(to_regprocedure(%s), 'EXECUTE'), + has_schema_privilege(to_regnamespace('fastann'), 'USAGE') + """, + (signatures[0], signatures[0], signatures[1], signatures[1]), + ).fetchone() + access = ((row[0], row[1]), (row[2], row[3])) + unavailable = [ + signature + for signature, (oid, allowed) in zip(signatures, access, strict=True) + if oid is None or not allowed + ] + if unavailable or not row[-1]: + details = list(unavailable) + if not row[-1]: + details.append("USAGE on schema fastann") + msg = f"Required ADBPG UDF access is unavailable: {', '.join(details)}" + raise RuntimeError(msg) def _drop_index(self): assert self.conn is not None, "Connection is not initialized" assert self.cursor is not None, "Cursor is not initialized" log.info(f"{self.name} client drop index : {self._index_name}") - drop_index_sql = sql.SQL("DROP INDEX IF EXISTS {index_name}").format( + drop_index_sql = sql.SQL("DROP INDEX IF EXISTS {schema}.{index_name}").format( + schema=sql.Identifier("public"), index_name=sql.Identifier(self._index_name), ) - log.debug(drop_index_sql.as_string(self.cursor)) self.cursor.execute(drop_index_sql) self.conn.commit() - def _set_parallel_index_build_param(self): + def _build_udf_payload(self) -> dict[str, Any]: + columns = [ + {"name": self._primary_field, "type": "bigint", "role": "primary_key"}, + {"name": self._vector_field, "type": f"vector({self.dim})", "role": "vector"}, + ] + if self.with_scalar_labels: + columns.append( + {"name": self._scalar_label_field, "type": "varchar(64)", "role": "filter"}, + ) + parameters = { + "build": self.case_config.build_parameters.model_dump(mode="json"), + "search": self.case_config.search_parameters.model_dump(mode="json"), + } + if self.case_config.autotune_parameters is not None: + parameters["autotune"] = self.case_config.autotune_parameters.model_dump(mode="json") + return { + "api_version": 1, + "relation": { + "schema": "public", + "table": self.table_name, + "index": self._index_name, + }, + "columns": columns, + "metric": self.case_config.parse_metric(), + "parameters": parameters, + } + + def _call_udf(self, function_name: str, payload: dict[str, Any], timeout: int | None = None) -> list[Any]: assert self.conn is not None, "Connection is not initialized" assert self.cursor is not None, "Cursor is not initialized" + qualified_name = f"fastann.{function_name}" + command = sql.SQL("SELECT * FROM {schema}.{function}(%s::jsonb)").format( + schema=sql.Identifier("fastann"), + function=sql.Identifier(function_name), + ) + try: + if timeout is not None: + self.cursor.execute( + "SELECT set_config('statement_timeout', %s, true)", + (f"{timeout}s",), + ) + rows = self.cursor.execute(command, (Jsonb(payload),)).fetchall() + self.conn.commit() + except Exception as exc: + self.conn.rollback() + cancel_reason = getattr(getattr(exc, "diag", None), "message_primary", None) or str(exc) + if ( + timeout is not None + and isinstance(exc, psycopg.errors.QueryCanceled) + and "statement timeout" in cancel_reason.lower() + ): + msg = f"{qualified_name} timed out after {timeout} seconds" + raise AdbpgTimeoutError(msg) from exc + msg = f"{qualified_name} failed: {exc}" + raise RuntimeError(msg) from exc + else: + log.info("%s returned diagnostic rows: %s", qualified_name, rows) + return rows - index_param = self.case_config.index_param() - - if index_param["build_parallel_processes"] is not None: - self.cursor.execute( - sql.SQL("SET fastann.build_parallel_processes TO {};").format( - index_param["build_parallel_processes"], - ), + def _apply_search_reloptions(self) -> None: + assert self.conn is not None, "Connection is not initialized" + assert self.cursor is not None, "Cursor is not initialized" + reloptions = self.case_config.search_parameters.reloption + set_options = [(name, value) for name, value in reloptions.items() if value is not None] + reset_options = [name for name, value in reloptions.items() if value is None] + if set_options: + assignments = sql.SQL(", ").join( + sql.SQL("{name} = {value}").format( + name=sql.Identifier(name), + value=sql.Literal(value), + ) + for name, value in set_options ) + command = sql.SQL("ALTER INDEX {schema}.{index} SET ({assignments})").format( + schema=sql.Identifier("public"), + index=sql.Identifier(self._index_name), + assignments=assignments, + ) + self.cursor.execute(command) + if reset_options: + names = sql.SQL(", ").join(sql.Identifier(name) for name in reset_options) + command = sql.SQL("ALTER INDEX {schema}.{index} RESET ({names})").format( + schema=sql.Identifier("public"), + index=sql.Identifier(self._index_name), + names=names, + ) + self.cursor.execute(command) + if set_options or reset_options: self.conn.commit() - results = self.cursor.execute(sql.SQL("SHOW fastann.build_parallel_processes;")).fetchall() - log.info(f"{self.name} parallel index creation parameters: {results}") - - def _create_index(self): + def _apply_search_gucs(self) -> None: assert self.conn is not None, "Connection is not initialized" assert self.cursor is not None, "Cursor is not initialized" - log.info(f"{self.name} client create index : {self._index_name}") - - index_param = self.case_config.index_param() - self._set_parallel_index_build_param() - - # Pre-build GUC: raise optimizer level before creating the ANN index. - self.cursor.execute(sql.SQL("SET fastann.nova_build_optimize_level = 3;")) - self.conn.commit() - - options = [] - options.append(sql.SQL("dim = {dim}").format(dim=sql.Literal(self.dim))) - options.append( - sql.SQL("distancemeasure = {measure}").format( - measure=sql.Identifier(index_param["metric"]), - ), - ) - - for option in index_param["index_creation_with_options"]: - if option["val"] is not None: - # When `raw` is set, emit the value as a bare SQL token - # (e.g. auto_reduction=on) instead of a quoted literal. - rendered_val = sql.SQL(str(option["val"])) if option.get("raw") else sql.Literal(option["val"]) - options.append( - sql.SQL("{option_name} = {val}").format( - option_name=sql.Identifier(option["option_name"]), - val=rendered_val, - ), + search_gucs = self.case_config.search_parameters.guc + for name, value in search_gucs.items(): + if value is None: + command = sql.SQL("RESET {setting_name}").format(setting_name=sql.Identifier(name)) + log.debug(command.as_string(self.cursor)) + self.cursor.execute(command) + else: + self.cursor.execute( + "SELECT set_config(%s, %s, false)", + (name, str(value)), ) - - with_clause = sql.SQL("WITH ({});").format(sql.SQL(", ").join(options)) if options else sql.Composed(()) - - # Covering index: always INCLUDE the primary field (e.g. id). - index_create_sql = sql.SQL( - """ - CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} - USING ann ({vector_field}) INCLUDE ({primary_field}) - """, - ).format( - index_name=sql.Identifier(self._index_name), - table_name=sql.Identifier(self.table_name), - vector_field=sql.Identifier(self._vector_field), - primary_field=sql.Identifier(self._primary_field), - ) - - full_sql = (index_create_sql + with_clause).join(" ") - log.debug(full_sql.as_string(self.cursor)) - self.cursor.execute(full_sql) - self.conn.commit() + if search_gucs: + self.conn.commit() def _create_table(self, dim: int): assert self.conn is not None, "Connection is not initialized" @@ -360,6 +416,7 @@ def prepare_filter(self, filters: Filter): msg = f"Not support Filter for Adbpg - {filters}" raise ValueError(msg) + self._apply_search_gucs() self._search = self._generate_search_query() def search_embedding( diff --git a/vectordb_bench/backend/clients/adbpg/cli.py b/vectordb_bench/backend/clients/adbpg/cli.py index e579f9028..fab0ee61a 100644 --- a/vectordb_bench/backend/clients/adbpg/cli.py +++ b/vectordb_bench/backend/clients/adbpg/cli.py @@ -1,5 +1,5 @@ import os -from typing import Annotated, Unpack +from typing import Annotated, Any, Unpack import click from pydantic import SecretStr @@ -13,6 +13,20 @@ get_custom_case_config, run, ) +from .config import parse_adbpg_parameter_group + + +class AdbpgParameterGroupType(click.ParamType): + name = "YAML_MAPPING" + + def convert(self, value: Any, param: click.Parameter | None, ctx: click.Context | None) -> dict[str, Any]: + try: + return dict(parse_adbpg_parameter_group(value, param.name if param is not None else self.name)) + except ValueError as exc: + self.fail(str(exc), param, ctx) + + +ADBPG_PARAMETER_GROUP = AdbpgParameterGroupType() class AdbpgTypedDict(CommonTypedDict): @@ -39,133 +53,36 @@ class AdbpgTypedDict(CommonTypedDict): help="Postgres database port", default=5432, show_default=True, - required=False, ), ] db_name: Annotated[str, click.option("--db-name", type=str, help="Db name", required=True)] - hnsw_m: Annotated[ - int, - click.option("--hnsw-m", type=int, help="hnsw_m", default=48, show_default=True, required=False), - ] - ef_search: Annotated[ - int, - click.option("--ef-search", type=int, help="ef_search", default=150, show_default=True, required=False), - ] - ef_construction: Annotated[ - int, - click.option( - "--ef-construction", - type=int, - help="ef_construction", - default=600, - show_default=True, - required=False, - ), - ] - nlist: Annotated[ - int, - click.option("--nlist", type=int, help="nlist", default=1024, show_default=True, required=False), - ] - rabitq_bits: Annotated[ - int, - click.option("--rabitq-bits", type=int, help="rabitq_bits", default=7, show_default=True, required=False), - ] - quantize_rescore_amp: Annotated[ - float, - click.option( - "--quantize-rescore-amp", - type=float, - help="fastann.quantize_rescore_amp", - default=0.0, - show_default=True, - required=False, - ), - ] - nova_adaptive_gamma: Annotated[ - float, - click.option( - "--nova-adaptive-gamma", - type=float, - help="fastann.nova_adaptive_gamma", - default=0.0, - show_default=True, - required=False, - ), - ] - auto_reduction: Annotated[ - bool, + build_parameters: Annotated[ + dict[str, Any], click.option( - "--auto-reduction/--no-auto-reduction", - "auto_reduction", - type=bool, - help="Index WITH auto_reduction=on when enabled", - default=False, + "--build-parameters", + type=ADBPG_PARAMETER_GROUP, + default="{}", + help="Build reloption/GUC mapping as YAML or JSON", show_default=True, - required=False, ), ] - max_scan_points: Annotated[ - int, - click.option( - "--max-scan-points", - type=int, - help="max_scan_points", - default=20000, - show_default=True, - required=False, - ), - ] - index_scan_mode: Annotated[ - str, + search_parameters: Annotated[ + dict[str, Any], click.option( - "--index-scan-mode", - type=str, - help="fastann.index_scan_mode", - default="snapshot", + "--search-parameters", + type=ADBPG_PARAMETER_GROUP, + default="{}", + help="Search reloption/GUC mapping as YAML or JSON", show_default=True, - required=False, ), ] - algorithm: Annotated[ - str, + autotune_parameters: Annotated[ + dict[str, Any] | None, click.option( - "--algorithm", - type=str, - help="algorithm", - default="novamr", - show_default=True, - required=False, - ), - ] - build_parallel_processes: Annotated[ - int, - click.option( - "--build-parallel-processes", - type=int, - help="Sets the maximum process to build index", - required=False, - ), - ] - pca_dim: Annotated[ - int | None, - click.option( - "--pca-dim", - type=int, - help="PCA dimension for index dimensionality reduction", + "--autotune-parameters", + type=ADBPG_PARAMETER_GROUP, + help="Autotune mapping as YAML or JSON; omit to disable", default=None, - show_default=True, - required=False, - ), - ] - nprobe: Annotated[ - int, - click.option( - "--nprobe", - type=int, - help="fastann.nova_nprobe (novad search)", - default=5, - show_default=True, - required=False, ), ] @@ -179,6 +96,7 @@ def AdbpgNova(**parameters: Unpack[AdbpgTypedDict]): run( db=DB.Adbpg, db_config=AdbpgConfig( + db_label=parameters["db_label"], user_name=SecretStr(parameters["user_name"]), password=SecretStr(parameters["password"]), host=parameters["host"], @@ -186,20 +104,9 @@ def AdbpgNova(**parameters: Unpack[AdbpgTypedDict]): db_name=parameters["db_name"], ), db_case_config=AdbpgIndexConfig( - hnsw_m=parameters["hnsw_m"], - ef_search=parameters["ef_search"], - ef_construction=parameters["ef_construction"], - nlist=parameters["nlist"], - algorithm=parameters["algorithm"], - build_parallel_processes=parameters["build_parallel_processes"], - rabitq_bits=parameters["rabitq_bits"], - quantize_rescore_amp=parameters["quantize_rescore_amp"], - nova_adaptive_gamma=parameters["nova_adaptive_gamma"], - auto_reduction=parameters["auto_reduction"], - pca_dim=parameters["pca_dim"], - max_scan_points=parameters["max_scan_points"], - index_scan_mode=parameters["index_scan_mode"], - nprobe=parameters["nprobe"], + build_parameters=parameters["build_parameters"], + search_parameters=parameters["search_parameters"], + autotune_parameters=parameters["autotune_parameters"], ), **parameters, ) diff --git a/vectordb_bench/backend/clients/adbpg/config.py b/vectordb_bench/backend/clients/adbpg/config.py index cc69085bc..2606bce9c 100644 --- a/vectordb_bench/backend/clients/adbpg/config.py +++ b/vectordb_bench/backend/clients/adbpg/config.py @@ -1,18 +1,48 @@ -from collections.abc import Mapping, Sequence -from typing import Any, TypedDict +from ast import literal_eval +from collections.abc import Mapping +from typing import Any, TypeAlias, TypedDict -from pydantic import BaseModel, SecretStr +from pydantic import BaseModel, ConfigDict, Field, JsonValue, SecretStr, field_validator, model_validator +from yaml import safe_load from ..api import DBCaseConfig, DBConfig, MetricType - -class AdbpgSessionCommands(TypedDict): - session_options: Sequence[dict[str, Any]] +AdbpgParameterValue: TypeAlias = str | int | float | bool | None + +DEFAULT_BUILD_RELOPTIONS: dict[str, AdbpgParameterValue] = { + "algorithm": "novamr", + "hnsw_m": 48, + "hnsw_ef_construction": 600, + "rabitq_bits": 7, + "auto_reduction": False, +} +DEFAULT_BUILD_GUCS: dict[str, AdbpgParameterValue] = { + "fastann.build_parallel_processes": 32, +} +DEFAULT_OPTIMIZE_UDF_TIMEOUT = 600 + +LEGACY_BUILD_RELOPTIONS = { + "algorithm": "algorithm", + "hnsw_m": "hnsw_m", + "ef_construction": "hnsw_ef_construction", + "nlist": "nlist", + "rabitq_bits": "rabitq_bits", + "auto_reduction": "auto_reduction", + "pca_dim": "pca_dim", +} +LEGACY_BUILD_GUCS = {"build_parallel_processes": "fastann.build_parallel_processes"} +LEGACY_SEARCH_GUCS = { + "ef_search": "fastann.hnsw_ef_search", + "max_scan_points": "fastann.hnsw_max_scan_points", + "quantize_rescore_amp": "fastann.quantize_rescore_amp", + "nova_adaptive_gamma": "fastann.nova_adaptive_gamma", + "index_scan_mode": "fastann.index_scan_mode", + "nprobe": "fastann.nova_nprobe", +} class AdbpgConfigDict(TypedDict): - """These keys will be directly used as kwargs in psycopg connection string, - so the names must match exactly psycopg API.""" + """Keys passed directly to psycopg.connect().""" user: str password: str @@ -39,33 +69,166 @@ def to_dict(self) -> dict: "dbname": self.db_name, "user": user_str, "password": pwd_str, - "options": "-c gp_session_role=utility", }, } +def parse_adbpg_parameter_group(value: Any, field_name: str) -> Mapping[str, Any]: + """Parse a structured CLI/UI value while preserving YAML scalar types.""" + if isinstance(value, str): + try: + value = literal_eval(value) + except (SyntaxError, ValueError): + try: + value = safe_load(value) + except Exception as exc: + msg = f"{field_name} must be valid YAML or JSON" + raise ValueError(msg) from exc + if not isinstance(value, Mapping): + msg = f"{field_name} must be a mapping" + raise ValueError(msg) # noqa: TRY004 - Pydantic validators must raise ValueError. + return value + + +class AdbpgParameterGroup(BaseModel): + model_config = ConfigDict(extra="forbid") + + reloption: dict[str, AdbpgParameterValue] = Field(default_factory=dict) + guc: dict[str, AdbpgParameterValue] = Field(default_factory=dict) + + @model_validator(mode="before") + @classmethod + def parse_mapping(cls, value: Any) -> Any: + return parse_adbpg_parameter_group(value, cls.__name__) + + @field_validator("reloption", "guc", mode="before") + @classmethod + def validate_parameter_map(cls, value: Any, info) -> Any: # noqa: ANN001 + if not isinstance(value, Mapping): + msg = f"{info.field_name} must be a mapping" + raise ValueError(msg) # noqa: TRY004 - Pydantic validators must raise ValueError. + for name, setting in value.items(): + if not isinstance(name, str) or not name: + msg = f"{info.field_name} parameter names must be non-empty strings" + raise ValueError(msg) + if setting is not None and not isinstance(setting, (str, int, float, bool)): + msg = f"{info.field_name}.{name} must be a YAML scalar" + raise ValueError(msg) + return dict(value) + + +class AdbpgBuildParameters(AdbpgParameterGroup): + reloption: dict[str, AdbpgParameterValue] = Field(default_factory=DEFAULT_BUILD_RELOPTIONS.copy) + guc: dict[str, AdbpgParameterValue] = Field(default_factory=DEFAULT_BUILD_GUCS.copy) + + @model_validator(mode="after") + def merge_defaults(self): + self.reloption = {**DEFAULT_BUILD_RELOPTIONS, **self.reloption} + self.guc = {**DEFAULT_BUILD_GUCS, **self.guc} + return self + + +class AdbpgAutotuneParameters(BaseModel): + __pydantic_extra__: dict[str, JsonValue] = Field(init=False) + model_config = ConfigDict(extra="allow") + + topk: list[int] + target_recall: list[float] + n_samples: int = 200 + n_trials: int = 200 + timeout: int = DEFAULT_OPTIMIZE_UDF_TIMEOUT + + @model_validator(mode="before") + @classmethod + def parse_mapping(cls, value: Any) -> Any: + value = parse_adbpg_parameter_group(value, cls.__name__) + if "enable" in value: + msg = "omit autotune_parameters to disable autotune; enable is not supported" + raise ValueError(msg) + return value + + @model_validator(mode="after") + def validate_autotune(self): + if not self.topk or any(value <= 0 for value in self.topk): + msg = "topk must contain positive integers" + raise ValueError(msg) + if not self.target_recall or any(value <= 0 or value > 1 for value in self.target_recall): + msg = "target_recall must contain values in the interval (0, 1]" + raise ValueError(msg) + for name in ("n_samples", "n_trials", "timeout"): + value = getattr(self, name) + minimum = 0 if name == "n_samples" else 1 + if value < minimum: + msg = f"{name} must be at least {minimum}" + raise ValueError(msg) + return self + + class AdbpgIndexConfig(BaseModel, DBCaseConfig): + model_config = ConfigDict(extra="forbid") + metric_type: MetricType | None = None - create_index_before_load: bool = False - create_index_after_load: bool = True - - # ADB PG specific parameters - hnsw_m: int = 48 - ef_search: int = 150 - ef_construction: int = 600 - nlist: int = 1024 - algorithm: str = "novamr" - build_parallel_processes: int | None = None - # rabitq quantization params - rabitq_bits: int = 7 - quantize_rescore_amp: float = 0.0 - nova_adaptive_gamma: float = 0.0 - max_scan_points: int = 20000 - index_scan_mode: str = "snapshot" - auto_reduction: bool = False - pca_dim: int | None = None - # novad-specific search param (no-op for novamr/HNSW algorithms) - nprobe: int = 5 + build_parameters: AdbpgBuildParameters = Field(default_factory=AdbpgBuildParameters) + search_parameters: AdbpgParameterGroup = Field(default_factory=AdbpgParameterGroup) + autotune_parameters: AdbpgAutotuneParameters | None = None + + @model_validator(mode="before") + @classmethod + def migrate_legacy_parameters(cls, value: Any) -> Any: + """Keep existing UI payloads working while accepting structured config.""" + if not isinstance(value, Mapping): + return value + data = dict(value) + cls._move_legacy_parameters(data, "build_parameters", "reloption", LEGACY_BUILD_RELOPTIONS) + cls._move_legacy_parameters(data, "build_parameters", "guc", LEGACY_BUILD_GUCS) + cls._move_legacy_parameters(data, "search_parameters", "guc", LEGACY_SEARCH_GUCS) + return data + + @staticmethod + def _move_legacy_parameters( + data: dict[str, Any], + group_name: str, + parameter_type: str, + mapping: Mapping[str, str], + ) -> None: + legacy = {name: data[name] for name in mapping if name in data} + if not legacy: + return + + group = data.get(group_name) + if group is None: + group = {} + elif isinstance(group, BaseModel): + group = group.model_dump() + elif isinstance(group, Mapping): + group = dict(group) + else: + return + + parameters = group.get(parameter_type) + if parameters is None: + parameters = {} + elif isinstance(parameters, Mapping): + parameters = dict(parameters) + else: + return + + for name, target in mapping.items(): + if name in legacy: + parameters.setdefault(target, legacy[name]) + data.pop(name) + group[parameter_type] = parameters + data[group_name] = group + + @field_validator("autotune_parameters", mode="before") + @classmethod + def empty_autotune_is_disabled(cls, value: Any) -> Any: + return None if value == "" else value + + @property + def algorithm(self) -> str: + value = self.build_parameters.reloption["algorithm"] + return str(value) def parse_metric(self) -> str: if self.metric_type == MetricType.L2: @@ -77,61 +240,14 @@ def parse_metric(self) -> str: msg = f"Metric type {self.metric_type} is not supported!" raise ValueError(msg) - @staticmethod - def _build_forced_set_options(set_mapping: Mapping[str, Any]) -> Sequence[dict[str, Any]]: - """Always emit SET commands regardless of value (including 0 / 0.0).""" - return [ - { - "parameter": { - "setting_name": name, - "val": str(value), - }, - } - for name, value in set_mapping.items() - ] - def index_param(self) -> dict: - with_options = [ - {"option_name": "algorithm", "val": self.algorithm}, - {"option_name": "hnsw_m", "val": self.hnsw_m}, - {"option_name": "hnsw_ef_construction", "val": self.ef_construction}, - {"option_name": "nlist", "val": self.nlist}, - {"option_name": "rabitq_bits", "val": self.rabitq_bits}, - # Covering index key length. - {"option_name": "max_key_len", "val": 1}, - ] - # Optional: auto_reduction=on — only include when True. - # Uses raw=True so the value 'on' is emitted as a bare identifier - # instead of a quoted string literal. - if self.auto_reduction: - with_options.append({"option_name": "auto_reduction", "val": "on", "raw": True}) - if self.pca_dim is not None: - with_options.append({"option_name": "pca_dim", "val": self.pca_dim}) - return { "metric": self.parse_metric(), - "build_parallel_processes": self.build_parallel_processes, - "create_index_before_load": self.create_index_before_load, - "create_index_after_load": self.create_index_after_load, - "index_creation_with_options": with_options, + **self.build_parameters.model_dump(mode="json"), } def search_param(self) -> dict: return { "metric": self.parse_metric(), + **self.search_parameters.model_dump(mode="json"), } - - def session_param(self) -> AdbpgSessionCommands: - # All CLI-driven search GUCs are always sent, regardless of value, - # so that callers can explicitly tune any parameter — including to 0. - session_parameters = { - "fastann.quantize_rescore_amp": self.quantize_rescore_amp, - "fastann.nova_adaptive_gamma": self.nova_adaptive_gamma, - "fastann.hnsw_ef_search": self.ef_search, - "fastann.hnsw_max_scan_points": self.max_scan_points, - "fastann.index_scan_mode": self.index_scan_mode, - "fastann.nova_nprobe": self.nprobe, - "optimizer": "off", - "elog_process_parameters": "off", - } - return {"session_options": self._build_forced_set_options(session_parameters)} diff --git a/vectordb_bench/config-files/adbpg_cohere1m_novamr.yml b/vectordb_bench/config-files/adbpg_cohere1m_novamr.yml new file mode 100644 index 000000000..356c26e2c --- /dev/null +++ b/vectordb_bench/config-files/adbpg_cohere1m_novamr.yml @@ -0,0 +1,43 @@ +adbpgnova: + # Performance reference: + # https://help.aliyun.com/zh/analyticdb/analyticdb-for-postgresql/user-guide/nova-vector-index-performance-white-paper + db_label: cohere1m-novamr + task_label: cohere1m-top10-r095 + + case_type: Performance768D1M + k: 10 + drop_old: true + load: true + search_serial: true + search_concurrent: true + num_concurrency: "1,32,48,64,128" + concurrency_duration: 100 + + host: + port: 5432 + db_name: postgres + user_name: + # Set the password with POSTGRES_PASSWORD. + + # Passed to fastann.vectordbbench_build(jsonb). + build_parameters: + reloption: + algorithm: novamr + hnsw_m: 48 + hnsw_ef_construction: 600 + rabitq_bits: 7 + auto_reduction: on + guc: + fastann.build_parallel_processes: 32 + + # Persistent index options and per-search-connection GUCs. + search_parameters: + reloption: + nova_autotune_topk: 10 + nova_autotune_recall: 0.95 + guc: {} + + # Executed by fastann.vectordbbench_optimize(jsonb) after a new build. + autotune_parameters: + topk: [10, 100] + target_recall: [0.90, 0.95]