diff --git a/.github/workflows/pull_request.yml b/.github/workflows/pull_request.yml index e9de2c3b9..2049bf4b5 100644 --- a/.github/workflows/pull_request.yml +++ b/.github/workflows/pull_request.yml @@ -30,7 +30,7 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - pip install -e ".[test]" + pip install -e ".[test,adbpg]" - name: Run coding checks run: | diff --git a/Makefile b/Makefile index ef8207c55..4aeca25ae 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,6 @@ unittest: PYTHONPATH=`pwd` python3 -m pytest tests/test_dataset.py::TestDataSet::test_download_small -svv + PYTHONPATH=`pwd` python3 -m pytest tests/test_adbpg_search_config.py -svv format: PYTHONPATH=`pwd` python3 -m black vectordb_bench diff --git a/README.md b/README.md index 85d34d76d..3e8728abf 100644 --- a/README.md +++ b/README.md @@ -623,12 +623,16 @@ vectordbbench adbpgnova --case-type Performance1024D1M --k 10 \ --ef-search 130 --max-scan-points 5000 --quantize-rescore-amp 2.0 ``` -**Example: Run from config file** +**Example: Import Cohere 1M, build NOVAMR, run autotune, and benchmark top-10 search** ```shell -vectordbbench adbpgnova --config-file adbpg_bioasq1m_novamr.yml +vectordbbench adbpgnova \ + --config-file vectordb_bench/config-files/adbpg_cohere1m_autotune.yml ``` +The example config sets the benchmark `k` and the autotune `topk` independently. Autotune `topk` and +`target_recall` accept either a scalar value or an `ARRAY[...]` expression for tuning multiple targets. + To list the options for ADBPG, execute `vectordbbench adbpgnova --help`. The following are some ADBPG-specific command-line options. ```text @@ -638,13 +642,20 @@ To list the options for ADBPG, execute `vectordbbench adbpgnova --help`. The fol --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] + --hnsw-m INTEGER hnsw_m [default: 48] + --ef-construction INTEGER ef_construction [default: 600] + --ef-search INTEGER ef_search [default: 150] + --max-scan-points INTEGER max scan points [default: 20000] + --quantize-rescore-amp FLOAT fastann.quantize_rescore_amp [default: 0.0] --nova-adaptive-gamma FLOAT fastann.nova_adaptive_gamma [default: 0.0] --auto-reduction/--no-auto-reduction Index WITH auto_reduction=on [default: False] + --index-build-reloption TEXT CREATE INDEX WITH reloption as name=value; repeatable + --index-build-include TEXT CREATE INDEX INCLUDE column; repeatable [default: id] + --session-guc TEXT Search-session GUC as name=value; repeatable + --setup-sql TEXT SQL run once after the index exists; repeatable + --index-reset-reloption TEXT Post-build reloption; name=value sets and a bare name resets + --autotune-param TEXT NOVA autotune argument as name=SQL-expression; repeatable + --autotune-timeout INTEGER Seconds to wait for NOVA autotune [default: 43200] ``` ### Run PolarDB from command line diff --git a/tests/test_adbpg_search_config.py b/tests/test_adbpg_search_config.py new file mode 100644 index 000000000..261bf0c73 --- /dev/null +++ b/tests/test_adbpg_search_config.py @@ -0,0 +1,396 @@ +from pathlib import Path +from typing import Any + +import pytest +from click import BadParameter +from click.testing import CliRunner +from psycopg import adapters as psycopg_adapters + +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.adbpg import cli as adbpg_cli +from vectordb_bench.backend.clients.adbpg.adbpg import Adbpg +from vectordb_bench.backend.clients.adbpg.config import AdbpgConfig, AdbpgIndexConfig +from vectordb_bench.backend.clients.adbpg.options import parse_key_values, parse_reloptions +from vectordb_bench.backend.clients.api import MetricType + + +def _config(**kwargs: Any) -> AdbpgIndexConfig: + return AdbpgIndexConfig(metric_type=MetricType.COSINE, **kwargs) + + +def test_generic_session_gucs_override_builtin_values() -> None: + config = _config( + session_gucs={ + "fastann.hnsw_ef_search": "321", + "fastann.new_search_switch": "off", + } + ) + + options = { + item["parameter"]["setting_name"]: item["parameter"]["val"] + for item in config.session_param()["session_options"] + } + assert options["fastann.hnsw_ef_search"] == "321" + assert options["fastann.new_search_switch"] == "off" + + +def test_index_build_reloptions_override_dedicated_defaults() -> None: + config = _config( + algorithm="novamr", + index_build_reloptions={ + "algorithm": "novadflat", + "hnsw_m": "64", + "future_build_option": "enabled", + }, + ) + + options = config.index_param()["index_creation_with_options"] + by_name = {option["option_name"]: option["val"] for option in options} + assert by_name["algorithm"] == "novadflat" + assert by_name["hnsw_m"] == "64" + assert by_name["future_build_option"] == "enabled" + assert sum(option["option_name"] == "algorithm" for option in options) == 1 + + with pytest.raises(ValueError, match="cannot override dataset-derived options: dim"): + _config(index_build_reloptions={"dim": "384"}).index_param() + + +def test_index_build_reloptions_render_in_create_index(monkeypatch: pytest.MonkeyPatch) -> None: + executed: list[str] = [] + monkeypatch.setattr("vectordb_bench.backend.clients.adbpg.adbpg.log.debug", lambda *_args: None) + + class FakeCursor: + connection = None + adapters = psycopg_adapters + + def execute(self, query: Any) -> "FakeCursor": + executed.append(query if isinstance(query, str) else query.as_string(None)) + return self + + def fetchall(self) -> list[Any]: + return [] + + class FakeConnection: + def commit(self) -> None: + pass + + client = object.__new__(Adbpg) + client.name = "Adbpg" + client.dim = 768 + client.table_name = "vector" + client._index_name = "vector_novamr_index" + client._vector_field = "embedding" + client._primary_field = "id" + client.conn = FakeConnection() + client.cursor = FakeCursor() + client.case_config = _config( + index_build_includes=("id", "label", "tenant_id"), + index_build_reloptions={ + "algorithm": "novamr", + "hnsw_m": "48", + "hnsw_ef_construction": "600", + "rabitq_bits": "7", + "auto_reduction": "on", + }, + ) + + client._create_index() + + create_sql = next(statement for statement in executed if "CREATE INDEX" in statement) + assert "\"algorithm\" = 'novamr'" in create_sql + assert "\"hnsw_m\" = '48'" in create_sql + assert "\"hnsw_ef_construction\" = '600'" in create_sql + assert "\"rabitq_bits\" = '7'" in create_sql + assert "\"auto_reduction\" = 'on'" in create_sql + assert create_sql.count('"algorithm"') == 1 + assert 'INCLUDE ("id", "label", "tenant_id")' in create_sql + + +def test_index_build_include_defaults_and_validation() -> None: + assert _config().index_build_includes == ("id",) + assert _config(index_build_includes=("id", "label", "id")).index_build_includes == ("id", "label") + + with pytest.raises(ValueError, match="column cannot be empty"): + _config(index_build_includes=("id", " ")) + + +def test_cli_parses_generic_settings_and_reset() -> None: + assert parse_key_values(None, None, ("a=1", "b=2", "a=3")) == {"a": "3", "b": "2"} + assert parse_key_values(None, None, ("search_path=foo,bar",)) == {"search_path": "foo,bar"} + assert parse_reloptions(None, None, ("nova_ef_search=120", "nova_nprobe")) == { + "nova_ef_search": "120", + "nova_nprobe": None, + } + with pytest.raises(BadParameter): + parse_key_values(None, None, ("missing_value",)) + + +def test_cohere_autotune_example_loads_through_click(monkeypatch: pytest.MonkeyPatch) -> None: + captured: dict[str, Any] = {} + + def fake_run(*, db: DB, db_config: AdbpgConfig, db_case_config: AdbpgIndexConfig, **parameters: Any) -> None: + captured.update( + db=db, + db_config=db_config, + db_case_config=db_case_config, + parameters=parameters, + ) + + monkeypatch.setattr(adbpg_cli, "run", fake_run) + config_file = Path(__file__).parents[1] / "vectordb_bench" / "config-files" / "adbpg_cohere1m_autotune.yml" + + result = CliRunner().invoke( + adbpg_cli.AdbpgNova, + ["--config-file", str(config_file), "--dry-run"], + ) + + assert result.exit_code == 0, result.output + assert captured["db"] == DB.Adbpg + assert captured["db_config"].db_label == "cohere1m-novamr-autotune" + assert captured["parameters"]["case_type"] == "Performance768D1M" + assert captured["parameters"]["k"] == 10 + assert captured["parameters"]["drop_old"] is True + assert captured["parameters"]["load"] is True + case_config = captured["db_case_config"] + assert case_config.index_build_reloptions == { + "algorithm": "novamr", + "hnsw_m": "48", + "hnsw_ef_construction": "600", + "rabitq_bits": "7", + "auto_reduction": "on", + } + assert case_config.index_build_includes == ("id",) + assert case_config.session_gucs == {"fastann.nova_adaptive_gamma": "0"} + assert case_config.index_reset_reloptions == {} + assert case_config.autotune_params == { + "topk": "10", + "target_recall": "0.95", + "n_samples": "300", + "n_trials": "500", + "n_threads": "32", + } + assert case_config.autotune_timeout == 600 + assert case_config.setup_sql == ('ANALYZE "public"."vector"',) + assert "$" not in config_file.read_text() + + +def test_search_setup_uses_coordinator_and_preserves_sql() -> None: + executed: list[str] = [] + connected_with: dict[str, Any] = {} + + class FakeCursor: + def execute(self, query: Any) -> None: + executed.append(query if isinstance(query, str) else query.as_string(None)) + + def close(self) -> None: + pass + + class FakeConnection: + commits = 0 + rollbacks = 0 + + def commit(self) -> None: + self.commits += 1 + + def rollback(self) -> None: + self.rollbacks += 1 + + def close(self) -> None: + pass + + client = object.__new__(Adbpg) + client.name = "Adbpg" + client.table_name = "docs" + client._index_name = "docs_novam_index" + client.connect_config = {"host": "db", "options": "-c gp_session_role=utility"} + client.case_config = _config( + index_reset_reloptions={ + "nova_autotune_topk": "10", + "nova_autotune_recall": "0.95", + "nova_ef_search": None, + }, + setup_sql=( + 'ANALYZE "public"."docs"', + "DO $$ BEGIN NULL; END $$;", + "SELECT '$topk', '$table', '$index'", + ), + ) + connection = FakeConnection() + cursor = FakeCursor() + client._create_connection = lambda **kwargs: (connected_with.update(kwargs) or (connection, cursor)) + + client._apply_search_setup() + + assert connected_with == {"host": "db"} + expected_set_sql = "".join( + ( + 'ALTER INDEX "public"."docs_novam_index" SET ', + "(\"nova_autotune_topk\" = '10', \"nova_autotune_recall\" = '0.95')", + ) + ) + assert executed == [ + 'ANALYZE "public"."docs"', + "DO $$ BEGIN NULL; END $$;", + "SELECT '$topk', '$table', '$index'", + expected_set_sql, + 'ALTER INDEX "public"."docs_novam_index" RESET ("nova_ef_search")', + ] + assert connection.commits == 2 + assert connection.rollbacks == 0 + + +def test_autotune_waits_for_all_config_rows_and_worker_exit(monkeypatch: pytest.MonkeyPatch) -> None: + executed: list[tuple[str, tuple[Any, ...] | None]] = [] + + class FakeCursor: + current_query = "" + progress_polls = 0 + + def execute(self, query: Any, params: tuple[Any, ...] | None = None) -> "FakeCursor": + self.current_query = query if isinstance(query, str) else query.as_string(None) + executed.append((self.current_query, params)) + return self + + def fetchone(self) -> tuple[Any, ...] | None: + if "SELECT fastann.nova_autotune(" in self.current_query: + return (12345,) + if "nova_autotune_progress" in self.current_query: + self.progress_polls += 1 + if self.progress_polls == 1: + return (4321, "running", 10, 20, 4) + return None + if "nova_autotune_configs" in self.current_query: + return (4,) + raise AssertionError(self.current_query) + + def close(self) -> None: + pass + + class FakeConnection: + commits = 0 + rollbacks = 0 + + def commit(self) -> None: + self.commits += 1 + + def rollback(self) -> None: + self.rollbacks += 1 + + def close(self) -> None: + pass + + client = object.__new__(Adbpg) + client.name = "Adbpg" + client.table_name = "vector" + client._index_name = "vector_novamr_index" + client.connect_config = {"host": "db", "options": "-c gp_session_role=utility"} + client.case_config = _config( + setup_sql=('ANALYZE "public"."vector"',), + autotune_params={ + "topk": "ARRAY[10,100]", + "target_recall": "ARRAY[0.90,0.95]", + "n_samples": "300", + "synthetic_query_mode": "'anchored_gaussian'", + }, + ) + connection = FakeConnection() + cursor = FakeCursor() + client._create_connection = lambda **_kwargs: (connection, cursor) + monkeypatch.setattr("vectordb_bench.backend.clients.adbpg.adbpg.time.sleep", lambda _seconds: None) + + client._apply_search_setup() + + sql_text = [query for query, _params in executed] + assert sql_text[0] == 'ANALYZE "public"."vector"' + assert "SELECT fastann.nova_autotune(" in sql_text[1] + assert '"topk" => ARRAY[10,100]' in sql_text[1] + assert '"target_recall" => ARRAY[0.90,0.95]' in sql_text[1] + assert '"n_samples" => 300' in sql_text[1] + assert "\"synthetic_query_mode\" => 'anchored_gaussian'" in sql_text[1] + assert executed[1][1] == ("public.vector_novamr_index",) + assert "nova_autotune_progress" in sql_text[2] + assert "nova_autotune_configs" in sql_text[3] + assert "c.topk = ANY(%s::integer[])" in sql_text[3] + assert "c.target_recall = ANY(%s::real[])" in sql_text[3] + assert executed[3][1] == ( + "public.vector_novamr_index", + [10, 100], + [0.90, 0.95], + ) + assert sum("nova_autotune_progress" in query for query in sql_text) == 2 + assert sum("nova_autotune_configs" in query for query in sql_text) == 2 + assert all("nova_autotune_status" not in query for query in sql_text) + assert connection.commits == 4 + assert connection.rollbacks == 0 + + +def test_autotune_and_target_reloptions_are_separate_modes() -> None: + client = object.__new__(Adbpg) + client.case_config = _config( + autotune_params={"topk": "10", "target_recall": "0.95"}, + index_reset_reloptions={"nova_autotune_topk": "10", "nova_autotune_recall": "0.95"}, + ) + + with pytest.raises(ValueError, match="Do not combine autotune_param"): + client._validate_autotune_configuration(drop_old=True) + + +def test_autotune_fails_closed_when_config_row_count_is_incomplete(monkeypatch: pytest.MonkeyPatch) -> None: + class FakeCursor: + current_query = "" + + def execute(self, query: Any, _params: tuple[Any, ...] | None = None) -> "FakeCursor": + self.current_query = query if isinstance(query, str) else query.as_string(None) + return self + + def fetchone(self) -> tuple[Any, ...] | None: + if "SELECT fastann.nova_autotune(" in self.current_query: + return (12345,) + if "nova_autotune_progress" in self.current_query: + return None + if "nova_autotune_configs" in self.current_query: + return (0,) + raise AssertionError(self.current_query) + + class FakeConnection: + def commit(self) -> None: + pass + + client = object.__new__(Adbpg) + client._index_name = "vector_novamr_index" + client.case_config = _config(autotune_params={"topk": "10", "target_recall": "0.95"}) + monkeypatch.setattr("vectordb_bench.backend.clients.adbpg.adbpg.time.sleep", lambda _seconds: None) + + with pytest.raises(RuntimeError, match=r"finished with 0 config rows.*expected 1"): + client._run_nova_autotune(FakeConnection(), FakeCursor()) + + +def test_autotune_requires_new_index_and_explicit_topk() -> None: + client = object.__new__(Adbpg) + client.case_config = _config(autotune_params={"topk": "ARRAY[10,100]", "target_recall": "ARRAY[0.90,0.95]"}) + + with pytest.raises(ValueError, match="requires a load run"): + client._validate_autotune_configuration(drop_old=False) + + client.case_config = _config(autotune_params={"target_recall": "0.95"}) + with pytest.raises(ValueError, match="requires an explicit topk"): + client._validate_autotune_configuration(drop_old=True) + + client.case_config = _config(autotune_params={"topk": "ARRAY[10,100]", "target_recall": "ARRAY[0.90,0.95]"}) + assert client._autotune_topks() == (10, 100) + assert client._autotune_target_recalls() == (0.90, 0.95) + + client.case_config = _config(autotune_params={"topk": "ARRAY[10,10]", "target_recall": "ARRAY[0.95,0.95]"}) + assert client._autotune_topks() == (10,) + assert client._autotune_target_recalls() == (0.95,) + + +def test_optimize_applies_setup_after_index_work() -> None: + events: list[str] = [] + client = object.__new__(Adbpg) + client._post_insert = lambda: events.append("index") + client._apply_search_setup = lambda: events.append("setup") + + client.optimize() + + assert events == ["index", "setup"] diff --git a/vectordb_bench/backend/clients/adbpg/adbpg.py b/vectordb_bench/backend/clients/adbpg/adbpg.py index bc5ac3486..5c8e0dcb8 100644 --- a/vectordb_bench/backend/clients/adbpg/adbpg.py +++ b/vectordb_bench/backend/clients/adbpg/adbpg.py @@ -1,6 +1,8 @@ """Wrapper around the Aliyun ADBPG (AnalyticDB for PostgreSQL) vector database.""" import logging +import re +import time from collections.abc import Generator, Sequence from contextlib import contextmanager from typing import Any @@ -17,6 +19,23 @@ log = logging.getLogger(__name__) +NOVA_AUTOTUNE_POLL_SECONDS = 5.0 +NOVA_AUTOTUNE_MISSING_PROGRESS_LIMIT = 3 +NOVA_AUTOTUNE_NUMERIC_EXPRESSION = re.compile( + r"(?P[+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)" + r"(?:\s*::\s*(?:real|float4|float8|double\s+precision))?", + re.IGNORECASE, +) +NOVA_AUTOTUNE_INTEGER_EXPRESSION = re.compile( + r"(?P[+-]?\d+)(?:\s*::\s*(?:smallint|int2|integer|int4|bigint|int8))?", + re.IGNORECASE, +) +NOVA_AUTOTUNE_ARRAY_EXPRESSION = re.compile( + r"ARRAY\s*\[(.*)\](?:\s*::\s*(?:smallint|int2|integer|int4|bigint|int8|" + r"real|float4|float8|double\s+precision)\s*\[\s*\])?", + re.IGNORECASE | re.DOTALL, +) + class Adbpg(VectorDB): """ADBPG vector database client, using psycopg.""" @@ -61,6 +80,7 @@ def __init__( self._scalar_label_field = "label" # Index name derives from the table name + algorithm, e.g. vector_1024d_10m_novamr_index. self._index_name = f"{self.table_name}_{self.case_config.algorithm}_index" + self._validate_autotune_configuration(drop_old) self.where_clause = "" @@ -87,6 +107,9 @@ def __init__( self._create_table(dim) if self.case_config.create_index_before_load: self._create_index() + else: + # Search-only runs reuse an existing index and skip optimize(). + self._apply_search_setup() self.cursor.close() self.conn.close() @@ -144,7 +167,7 @@ def init(self) -> Generator[None, None, None]: 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"])), + val=sql.Literal(setting["parameter"]["val"]), ) log.debug(command.as_string(self.cursor)) self.cursor.execute(command) @@ -172,6 +195,7 @@ def _drop_table(self): def optimize(self, data_size: int | None = None): self._post_insert() + self._apply_search_setup() def _post_insert(self): log.info(f"{self.name} post insert before optimize") @@ -179,6 +203,82 @@ def _post_insert(self): self._drop_index() self._create_index() + def _validate_autotune_configuration(self, drop_old: bool) -> None: + params = self.case_config.autotune_params + if not params: + return + if not drop_old: + msg = "autotune_param requires a load run that creates a new index" + raise ValueError(msg) + reserved = {"index_relation"}.intersection(params) + if reserved: + names = ", ".join(sorted(reserved)) + msg = f"autotune_param cannot override {names}; it comes from the created index" + raise ValueError(msg) + empty = [name for name, value in params.items() if not value.strip()] + if empty: + msg = f"autotune_param requires a SQL expression for: {', '.join(sorted(empty))}" + raise ValueError(msg) + self._autotune_topks() + self._autotune_target_recalls() + + target_reloptions = {"nova_autotune_topk", "nova_autotune_recall"}.intersection( + self.case_config.index_reset_reloptions + ) + if target_reloptions: + names = ", ".join(sorted(target_reloptions)) + msg = f"Do not combine autotune_param with target-selection index_reset_reloption values: {names}" + raise ValueError(msg) + + def _autotune_numeric_values( + self, + name: str, + expression: str, + pattern: re.Pattern[str], + ) -> tuple[str, ...]: + """Parse a scalar or ARRAY of numeric SQL literals for completion checks.""" + expression = expression.strip() + scalar_match = pattern.fullmatch(expression) + if scalar_match: + return (scalar_match.group("number"),) + + array_match = NOVA_AUTOTUNE_ARRAY_EXPRESSION.fullmatch(expression) + if array_match: + items = [item.strip() for item in array_match.group(1).split(",")] + matches = [pattern.fullmatch(item) for item in items] + if items and all(matches): + return tuple(match.group("number") for match in matches if match is not None) + + msg = f"autotune {name} must be a numeric SQL literal or an ARRAY[...] of numeric SQL literals" + raise ValueError(msg) + + def _autotune_topks(self) -> tuple[int, ...]: + """Return the explicitly requested topK values.""" + expression = self.case_config.autotune_params.get("topk") + if expression is None: + msg = "autotune_param requires an explicit topk value" + raise ValueError(msg) + values = self._autotune_numeric_values("topk", expression, NOVA_AUTOTUNE_INTEGER_EXPRESSION) + topks = tuple(dict.fromkeys(int(value) for value in values)) + if any(value <= 0 for value in topks): + msg = "autotune topk values must be positive integers" + raise ValueError(msg) + return topks + + def _autotune_target_recalls(self) -> tuple[float, ...]: + """Return requested target recalls, including the server default.""" + expression = self.case_config.autotune_params.get("target_recall", "0.99") + values = self._autotune_numeric_values( + "target_recall", + expression, + NOVA_AUTOTUNE_NUMERIC_EXPRESSION, + ) + target_recalls = tuple(dict.fromkeys(float(value) for value in values)) + if any(value < 0.0 or value > 1.0 for value in target_recalls): + msg = "autotune target_recall values must be between 0 and 1" + raise ValueError(msg) + return target_recalls + def _drop_index(self): assert self.conn is not None, "Connection is not initialized" assert self.cursor is not None, "Cursor is not initialized" @@ -191,6 +291,165 @@ def _drop_index(self): self.cursor.execute(drop_index_sql) self.conn.commit() + def _apply_search_setup(self) -> None: + """Run setup SQL, optional autotune, and reloptions on the coordinator.""" + reloptions = self.case_config.index_reset_reloptions + statements = self.case_config.setup_sql + autotune_params = self.case_config.autotune_params + if not reloptions and not statements and not autotune_params: + return + + connect_config = dict(self.connect_config) + connect_config.pop("options", None) + conn, cursor = self._create_connection(**connect_config) + try: + # setup_sql is the post-index preparation phase. Commit it before + # launching the asynchronous worker so ANALYZE and similar changes + # are visible to autotune. + for statement in statements: + log.info("%s setup SQL: %s", self.name, statement) + cursor.execute(statement) + if statements: + conn.commit() + + self._run_nova_autotune(conn, cursor) + + index = sql.Identifier("public", self._index_name) + 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 + ) + cursor.execute( + sql.SQL("ALTER INDEX {index} SET ({assignments})").format( + index=index, + assignments=assignments, + ) + ) + if reset_options: + names = sql.SQL(", ").join(sql.Identifier(name) for name in reset_options) + cursor.execute(sql.SQL("ALTER INDEX {index} RESET ({names})").format(index=index, names=names)) + if set_options or reset_options: + conn.commit() + except Exception: + conn.rollback() + raise + finally: + cursor.close() + conn.close() + + def _run_nova_autotune(self, conn: Connection, cursor: Cursor) -> None: + """Submit NOVA autotune, then wait for all config rows and worker exit.""" + params = self.case_config.autotune_params + if not params: + return + + topks = self._autotune_topks() + target_recalls = self._autotune_target_recalls() + expected_config_count = len(topks) * len(target_recalls) + # This is bound as a query parameter, not interpolated into SQL. ADBPG + # currently derives index names from framework-controlled table and + # algorithm names, so the regclass text form is stable. + qualified_index = f"public.{self._index_name}" + arguments = [sql.SQL("index_relation => %s::text")] + arguments.extend( + sql.SQL("{name} => {value}").format( + name=sql.Identifier(name), + value=sql.SQL(value), + ) + for name, value in params.items() + ) + submit_sql = sql.SQL("SELECT fastann.nova_autotune({arguments})").format( + arguments=sql.SQL(", ").join(arguments) + ) + cursor.execute(submit_sql, (qualified_index,)) + row = cursor.fetchone() + if row is None: + msg = "nova_autotune did not return a task handle" + raise RuntimeError(msg) + handle = int(row[0]) + conn.commit() + log.info( + "NOVA_AUTOTUNE_SUBMITTED handle=%s index=%s topks=%s target_recalls=%s params=%s", + handle, + qualified_index, + topks, + target_recalls, + params, + ) + + deadline = time.monotonic() + self.case_config.autotune_timeout + missing_progress_polls = 0 + latest_progress = None + latest_config_count = 0 + while True: + cursor.execute( + """ + SELECT pid, stage, work_done, work_total, target_count + FROM fastann.nova_autotune_progress(%s) + """, + (handle,), + ) + latest_progress = cursor.fetchone() + cursor.execute( + """ + SELECT count(*) + FROM fastann.nova_autotune_configs c + JOIN pg_catalog.pg_class i + ON i.oid = c.index_relid + AND i.relfilenode::oid = c.index_relfilenode + WHERE i.oid = %s::regclass + AND c.topk = ANY(%s::integer[]) + AND c.target_recall = ANY(%s::real[]) + """, + (qualified_index, list(topks), list(target_recalls)), + ) + count_row = cursor.fetchone() + latest_config_count = int(count_row[0]) if count_row is not None else 0 + conn.commit() + + if latest_progress is None and latest_config_count == expected_config_count: + log.info( + "NOVA_AUTOTUNE_COMPLETED handle=%s index=%s configs=%s", + handle, + qualified_index, + latest_config_count, + ) + return + + if latest_progress is None: + missing_progress_polls += 1 + if missing_progress_polls >= NOVA_AUTOTUNE_MISSING_PROGRESS_LIMIT: + msg = ( + f"nova_autotune handle {handle} finished with {latest_config_count} config rows " + f"for topks={topks} and target_recalls={target_recalls}; " + f"expected {expected_config_count}" + ) + raise RuntimeError(msg) + else: + missing_progress_polls = 0 + + remaining = deadline - time.monotonic() + if remaining <= 0: + msg = ( + f"nova_autotune handle {handle} exceeded {self.case_config.autotune_timeout}s; " + f"progress={latest_progress}, configs={latest_config_count}/{expected_config_count}" + ) + raise TimeoutError(msg) + log.info( + "NOVA_AUTOTUNE_PROGRESS handle=%s progress=%s configs=%s/%s", + handle, + latest_progress, + latest_config_count, + expected_config_count, + ) + time.sleep(min(NOVA_AUTOTUNE_POLL_SECONDS, remaining)) + def _set_parallel_index_build_param(self): assert self.conn is not None, "Connection is not initialized" assert self.cursor is not None, "Cursor is not initialized" @@ -242,17 +501,25 @@ def _create_index(self): 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). + include_columns = self.case_config.index_build_includes + include_clause = ( + sql.SQL(" INCLUDE ({})").format( + sql.SQL(", ").join(sql.Identifier(column) for column in include_columns), + ) + if include_columns + else sql.Composed(()) + ) + index_create_sql = sql.SQL( """ CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} - USING ann ({vector_field}) INCLUDE ({primary_field}) + USING ann ({vector_field}){include_clause} """, ).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), + include_clause=include_clause, ) full_sql = (index_create_sql + with_clause).join(" ") diff --git a/vectordb_bench/backend/clients/adbpg/cli.py b/vectordb_bench/backend/clients/adbpg/cli.py index e579f9028..d7bdb7ed6 100644 --- a/vectordb_bench/backend/clients/adbpg/cli.py +++ b/vectordb_bench/backend/clients/adbpg/cli.py @@ -13,6 +13,7 @@ get_custom_case_config, run, ) +from .options import parse_key_values, parse_reloptions class AdbpgTypedDict(CommonTypedDict): @@ -168,6 +169,80 @@ class AdbpgTypedDict(CommonTypedDict): required=False, ), ] + session_guc: Annotated[ + dict[str, str], + click.option( + "--session-guc", + type=str, + multiple=True, + callback=parse_key_values, + help="Session GUC as name=value; repeat the option for multiple settings", + ), + ] + index_build_reloption: Annotated[ + dict[str, str], + click.option( + "--index-build-reloption", + type=str, + multiple=True, + callback=parse_key_values, + help="CREATE INDEX WITH reloption as name=value; repeat for multiple options", + ), + ] + index_build_include: Annotated[ + tuple[str, ...], + click.option( + "--index-build-include", + type=str, + multiple=True, + default=("id",), + show_default=True, + help="CREATE INDEX INCLUDE column; repeat for multiple columns", + ), + ] + index_reset_reloption: Annotated[ + dict[str, str | None], + click.option( + "--index-reset-reloption", + type=str, + multiple=True, + callback=parse_reloptions, + help="Post-build index reloption: name=value sets it; a bare name resets it", + ), + ] + setup_sql: Annotated[ + tuple[str, ...], + click.option( + "--setup-sql", + type=str, + multiple=True, + help="SQL run once after the index exists", + ), + ] + autotune_param: Annotated[ + dict[str, str], + click.option( + "--autotune-param", + type=str, + multiple=True, + callback=parse_key_values, + help=( + "NOVA autotune argument as name=SQL-expression; repeat for multiple arguments. " + "Set topk explicitly; scalar and ARRAY[...] values are supported. " + "The index_relation argument comes from the created index." + ), + ), + ] + autotune_timeout: Annotated[ + int, + click.option( + "--autotune-timeout", + type=click.IntRange(min=1), + default=43200, + show_default=True, + help="Seconds to wait for NOVA autotune to finish", + ), + ] @cli.command() @@ -179,6 +254,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"], @@ -200,6 +276,13 @@ def AdbpgNova(**parameters: Unpack[AdbpgTypedDict]): max_scan_points=parameters["max_scan_points"], index_scan_mode=parameters["index_scan_mode"], nprobe=parameters["nprobe"], + index_build_reloptions=parameters["index_build_reloption"], + index_build_includes=parameters["index_build_include"], + session_gucs=parameters["session_guc"], + index_reset_reloptions=parameters["index_reset_reloption"], + setup_sql=parameters["setup_sql"], + autotune_params=parameters["autotune_param"], + autotune_timeout=parameters["autotune_timeout"], ), **parameters, ) diff --git a/vectordb_bench/backend/clients/adbpg/config.py b/vectordb_bench/backend/clients/adbpg/config.py index cc69085bc..f622d1ff7 100644 --- a/vectordb_bench/backend/clients/adbpg/config.py +++ b/vectordb_bench/backend/clients/adbpg/config.py @@ -1,7 +1,7 @@ from collections.abc import Mapping, Sequence from typing import Any, TypedDict -from pydantic import BaseModel, SecretStr +from pydantic import BaseModel, Field, SecretStr, field_validator from ..api import DBCaseConfig, DBConfig, MetricType @@ -66,6 +66,37 @@ class AdbpgIndexConfig(BaseModel, DBCaseConfig): pca_dim: int | None = None # novad-specific search param (no-op for novamr/HNSW algorithms) nprobe: int = 5 + # Generic reloptions merged into CREATE INDEX ... WITH (...). These + # override the dedicated fields above when the same option is supplied. + index_build_reloptions: dict[str, str] = Field(default_factory=dict) + # Covering columns rendered as CREATE INDEX ... INCLUDE (...). + index_build_includes: tuple[str, ...] = ("id",) + # Generic ADBPG search setup. Product-specific parameters no longer need + # dedicated Python fields: GUCs apply per connection, while post-build + # reloptions and SQL run once after the index exists. + session_gucs: dict[str, str] = Field(default_factory=dict) + # Reloptions applied after the index exists. A name=value item emits + # ALTER INDEX ... SET; a bare name emits ALTER INDEX ... RESET. + index_reset_reloptions: dict[str, str | None] = Field(default_factory=dict) + setup_sql: tuple[str, ...] = () + # Optional NOVA autotune invocation performed after index creation and + # setup_sql. Values are SQL expressions so newly added server parameters + # can be used without adding another client field. + autotune_params: dict[str, str] = Field(default_factory=dict) + autotune_timeout: int = Field(default=43200, gt=0) + + @field_validator("index_build_includes") + @classmethod + def validate_index_build_includes(cls, columns: tuple[str, ...]) -> tuple[str, ...]: + normalized: list[str] = [] + for column in columns: + name = column.strip() + if not name: + msg = "index_build_include column cannot be empty" + raise ValueError(msg) + if name not in normalized: + normalized.append(name) + return tuple(normalized) def parse_metric(self) -> str: if self.metric_type == MetricType.L2: @@ -97,7 +128,7 @@ def index_param(self) -> dict: {"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. + # Maximum number of keys (ctids) stored per HNSW node. {"option_name": "max_key_len", "val": 1}, ] # Optional: auto_reduction=on — only include when True. @@ -108,6 +139,25 @@ def index_param(self) -> dict: if self.pca_dim is not None: with_options.append({"option_name": "pca_dim", "val": self.pca_dim}) + reserved_options = {"dim", "distancemeasure"} + invalid_options = reserved_options.intersection(self.index_build_reloptions) + if invalid_options: + names = ", ".join(sorted(invalid_options)) + msg = f"index_build_reloption cannot override dataset-derived options: {names}" + raise ValueError(msg) + + # Preserve the existing option order while allowing generic build + # reloptions to override dedicated defaults without emitting duplicate + # CREATE INDEX WITH keys. + option_indexes = {option["option_name"]: index for index, option in enumerate(with_options)} + for name, value in self.index_build_reloptions.items(): + option = {"option_name": name, "val": value} + if name in option_indexes: + with_options[option_indexes[name]] = option + else: + option_indexes[name] = len(with_options) + with_options.append(option) + return { "metric": self.parse_metric(), "build_parallel_processes": self.build_parallel_processes, @@ -134,4 +184,5 @@ def session_param(self) -> AdbpgSessionCommands: "optimizer": "off", "elog_process_parameters": "off", } + session_parameters.update(self.session_gucs) return {"session_options": self._build_forced_set_options(session_parameters)} diff --git a/vectordb_bench/backend/clients/adbpg/options.py b/vectordb_bench/backend/clients/adbpg/options.py new file mode 100644 index 000000000..5309424a6 --- /dev/null +++ b/vectordb_bench/backend/clients/adbpg/options.py @@ -0,0 +1,36 @@ +from typing import Any + +import click + + +def _normalize_config_items(values: tuple[str, ...]) -> list[str]: + """Normalize repeatable CLI/YAML values without changing value contents.""" + return [item for value in values if (item := str(value).strip())] + + +def parse_key_values(_ctx: Any, _param: Any, values: tuple[str, ...]) -> dict[str, str]: + """Parse repeatable ``name=value`` settings; the last value wins.""" + parsed: dict[str, str] = {} + for item in _normalize_config_items(values): + if "=" not in item: + message = f"Expected name=value, got: {item}" + raise click.BadParameter(message) + name, value = (part.strip() for part in item.split("=", 1)) + if not name: + message = f"Empty setting name in: {item}" + raise click.BadParameter(message) + parsed[name] = value + return parsed + + +def parse_reloptions(_ctx: Any, _param: Any, values: tuple[str, ...]) -> dict[str, str | None]: + """Parse reloptions; a bare name means ``ALTER INDEX ... RESET``.""" + parsed: dict[str, str | None] = {} + for item in _normalize_config_items(values): + name, separator, value = item.partition("=") + name = name.strip() + if not name: + message = f"Empty reloption name in: {item}" + raise click.BadParameter(message) + parsed[name] = value.strip() if separator else None + return parsed diff --git a/vectordb_bench/config-files/adbpg_cohere1m_autotune.yml b/vectordb_bench/config-files/adbpg_cohere1m_autotune.yml new file mode 100644 index 000000000..00e6a8130 --- /dev/null +++ b/vectordb_bench/config-files/adbpg_cohere1m_autotune.yml @@ -0,0 +1,59 @@ +# Cohere-1M import, NOVAMR build, autotune, and top-10 benchmark. +# Performance guide: https://help.aliyun.com/zh/analyticdb/analyticdb-for-postgresql/user-guide/nova-vector-index-performance-white-paper +adbpgnova: + db_label: cohere1m-novamr-autotune + task_label: cohere1m-top10-r095-autotune + + # Dataset and benchmark. + 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 + + # Database connection. Set the password with POSTGRES_PASSWORD. + host: + port: 5432 + db_name: postgres + user_name: + + # CREATE INDEX WITH reloptions. + index_build_reloption: + - "algorithm=novamr" + - "hnsw_m=48" + - "hnsw_ef_construction=600" + - "rabitq_bits=7" + - "auto_reduction=on" + + # CREATE INDEX INCLUDE columns. + index_build_include: + - "id" + + # Index build parallelism. + build_parallel_processes: 32 + + # Search GUCs. + session_guc: + - "fastann.nova_adaptive_gamma=0" + + # Run after the index is built and before autotune. + setup_sql: + - 'ANALYZE "public"."vector"' + + # index_relation uses the new index; set autotune topk explicitly. + autotune_param: + - "topk=10" + - "target_recall=0.95" + - "n_samples=300" + - "n_trials=500" + - "n_threads=32" + autotune_timeout: 600 + + # To reuse an already tuned index, set drop_old/load to false, remove the + # autotune settings above, and uncomment this block. + # index_reset_reloption: + # - "nova_autotune_topk=10" + # - "nova_autotune_recall=0.95"