diff --git a/README.md b/README.md index b6391fdee..5d3494021 100644 --- a/README.md +++ b/README.md @@ -46,7 +46,7 @@ All the database client supported | pinecone | `pip install vectordb-bench[pinecone]` | | weaviate | `pip install vectordb-bench[weaviate]` | | elastic, aliyun_elasticsearch| `pip install vectordb-bench[elastic]` | -| pgvector, pgvectorscale, pgdiskann, alloydb, vectorchord | `pip install vectordb-bench[pgvector]` | +| pgvector, pgvectorscale, pgdiskann, alloydb, vectorchord, lakebase_vector | `pip install vectordb-bench[pgvector]` | | pgvecto.rs | `pip install vectordb-bench[pgvecto_rs]` | | redis | `pip install vectordb-bench[redis]` | | memorydb | `pip install vectordb-bench[memorydb]` | diff --git a/pyproject.toml b/pyproject.toml index acd30275b..feee070ba 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ elastic = [ "elasticsearch" ] # For elastic and aliyun_elasticsearch pgvector = [ "psycopg", "psycopg-binary", "pgvector" ] -# for pgvector, pgvectorscale, pgdiskann, and, alloydb +# for pgvector, pgvectorscale, pgdiskann, alloydb, and lakebase_vector pgvecto_rs = [ "pgvecto_rs[psycopg3]>=0.2.2" ] redis = [ "redis" ] diff --git a/tests/test_lakebase_vector.py b/tests/test_lakebase_vector.py new file mode 100644 index 000000000..7990c2848 --- /dev/null +++ b/tests/test_lakebase_vector.py @@ -0,0 +1,321 @@ +"""Offline tests for the Lakebase Vector config, client, and metric assembly path. + +Usage: + pytest tests/test_lakebase_vector.py -v +""" + +from __future__ import annotations + +import importlib +import pickle +from typing import TYPE_CHECKING +from unittest.mock import MagicMock, call + +import numpy as np +import pytest +from pydantic import SecretStr + +from vectordb_bench.backend.assembler import Assembler +from vectordb_bench.backend.cases import CaseType +from vectordb_bench.backend.clients import DB +from vectordb_bench.backend.clients.api import MetricType +from vectordb_bench.backend.clients.lakebase_vector.config import LakebaseANNConfig, LakebaseVectorConfig +from vectordb_bench.backend.data_source import DatasetSource +from vectordb_bench.backend.filter import Filter, IntFilter, LabelFilter, NonFilter +from vectordb_bench.models import CaseConfig, TaskConfig + +if TYPE_CHECKING: + from vectordb_bench.backend.clients.lakebase_vector.lakebase_vector import LakebaseVector + +DB_CONFIG = { + "connect_config": { + "host": "localhost", + "port": 5432, + "dbname": "vectordb", + "user": "vectordb", + "password": "vectordb", + }, + "table_name": "test_lakebase_vector", +} + +DIM = 128 + + +# ── Helpers ────────────────────────────────────────────────────────────────── + + +def make_ann_config(**overrides) -> LakebaseANNConfig: + values = { + "metric_type": MetricType.COSINE, + "probes": None, + "epsilon": None, + "max_parallel_workers": None, + } + values.update(overrides) + return LakebaseANNConfig(**values) + + +def lakebase_client_cls(): + pytest.importorskip("psycopg") + pytest.importorskip("pgvector") + from vectordb_bench.backend.clients.lakebase_vector.lakebase_vector import LakebaseVector # noqa: PLC0415 + + return LakebaseVector + + +def patch_sql_rendering(monkeypatch: pytest.MonkeyPatch): + client_cls = lakebase_client_cls() + client_module = importlib.import_module(client_cls.__module__) + original_as_string = client_module.sql.Composed.as_string + monkeypatch.setattr( + client_module.sql.Composed, + "as_string", + lambda composed, _context=None: original_as_string(composed), + ) + return original_as_string + + +def make_db( + table_name: str = "test_lakebase_vector", + drop_old: bool = True, + *, + case_config: LakebaseANNConfig | None = None, + with_scalar_labels: bool = False, +) -> LakebaseVector: + config = dict(DB_CONFIG) + config["connect_config"] = dict(DB_CONFIG["connect_config"]) + config["table_name"] = table_name + return DB.LakebaseVector.init_cls( + dim=DIM, + db_config=config, + db_case_config=case_config or make_ann_config(), + drop_old=drop_old, + with_scalar_labels=with_scalar_labels, + ) + + +@pytest.fixture +def mocked_db_connection(monkeypatch: pytest.MonkeyPatch) -> tuple[MagicMock, MagicMock]: + client_cls = lakebase_client_cls() + conn = MagicMock(name="connection") + cursor = MagicMock(name="cursor") + monkeypatch.setattr(client_cls, "_create_connection", staticmethod(lambda **_kwargs: (conn, cursor))) + return conn, cursor + + +class TestLakebaseVectorConfig: + def test_connection_config(self): + config = LakebaseVectorConfig( + user_name=SecretStr("lakebase-user"), + password=SecretStr("lakebase-password"), + host="lakebase.example.com", + port=6432, + db_name="benchmark", + table_name="vectors", + ) + + assert config.to_dict() == { + "connect_config": { + "host": "lakebase.example.com", + "port": 6432, + "dbname": "benchmark", + "user": "lakebase-user", + "password": "lakebase-password", + }, + "table_name": "vectors", + } + + +@pytest.mark.parametrize( + ("case_type", "expected_metric_type"), + [ + (CaseType.Performance768D100M, MetricType.L2), + (CaseType.Performance1536D50K, MetricType.COSINE), + ], +) +def test_dataset_metric( + case_type: CaseType, + expected_metric_type: MetricType, +): + db_case_config = LakebaseANNConfig() + task = TaskConfig( + db=DB.LakebaseVector, + db_config=LakebaseVectorConfig(password=SecretStr("test-password"), db_name="test-db"), + db_case_config=db_case_config, + case_config=CaseConfig(case_id=case_type), + ) + + assert db_case_config.metric_type is None + + runner = Assembler.assemble("test-run", task, DatasetSource.S3) + + assert runner.config.db_case_config.metric_type == expected_metric_type + assert runner.config.db_case_config.index_param()["metric"] is not None + + +class TestLakebaseVectorClient: + @pytest.mark.parametrize( + ("metric_type", "operator_class", "search_operator"), + [ + (MetricType.L2, "vector_l2_ops", "<->"), + (MetricType.IP, "vector_ip_ops", "<#>"), + (MetricType.COSINE, "vector_cosine_ops", "<=>"), + ], + ) + def test_index_sql( + self, + monkeypatch: pytest.MonkeyPatch, + mocked_db_connection: tuple[MagicMock, MagicMock], + metric_type: MetricType, + operator_class: str, + search_operator: str, + ) -> None: + conn, cursor = mocked_db_connection + render_sql = patch_sql_rendering(monkeypatch) + db = make_db( + "test_create_index", + drop_old=False, + case_config=make_ann_config(metric_type=metric_type, max_parallel_workers=16), + ) + db.conn = conn + db.cursor = cursor + conn.reset_mock() + cursor.reset_mock() + monkeypatch.setattr(db, "_set_parallel_index_build_param", MagicMock()) + + db._create_index() + + query = cursor.execute.call_args.args[0] + assert render_sql(query) == ( + 'CREATE INDEX IF NOT EXISTS "lakebase_vector_index" ' + 'ON public."test_create_index" USING lakebase_ann ' + f'("embedding" {operator_class})' + ) + assert db.case_config.search_param() == {"metric_fun_op": search_operator} + assert db.case_config.index_param()["max_parallel_workers"] == 16 + conn.commit.assert_called_once_with() + + def test_optimize( + self, + monkeypatch: pytest.MonkeyPatch, + mocked_db_connection: tuple[MagicMock, MagicMock], + ) -> None: + db = make_db("test_optimize", drop_old=False) + lifecycle = MagicMock() + monkeypatch.setattr(db, "_drop_index", lifecycle.drop_index) + monkeypatch.setattr(db, "_create_index", lifecycle.create_index) + + db.optimize() + + assert lifecycle.mock_calls == [call.drop_index(), call.create_index()] + + # The two probes cases configure single-level and two-level IVF, respectively. + @pytest.mark.parametrize("probes", ["10", "10,20"]) + def test_session_guc( + self, + monkeypatch: pytest.MonkeyPatch, + mocked_db_connection: tuple[MagicMock, MagicMock], + probes: str, + ) -> None: + conn, cursor = mocked_db_connection + render_sql = patch_sql_rendering(monkeypatch) + case_config = make_ann_config(probes=probes, epsilon=1.5) + db = make_db( + "test_session_gucs", + drop_old=False, + case_config=case_config, + ) + conn.reset_mock() + cursor.reset_mock() + + with db.init(): + assert db.conn is conn + assert db.cursor is cursor + + commands = [render_sql(execute_call.args[0]) for execute_call in cursor.execute.call_args_list] + assert commands == [ + f'SET "lakebase_ann.probes" = "{probes}";', + 'SET "lakebase_ann.epsilon" = "1.5";', + ] + conn.commit.assert_called_once_with() + cursor.close.assert_called_once_with() + conn.close.assert_called_once_with() + assert db.conn is None + assert db.cursor is None + + def test_pickle(self, mocked_db_connection: tuple[MagicMock, MagicMock]) -> None: + db = make_db("test_pickle", drop_old=False) + + restored = pickle.loads(pickle.dumps(db)) # noqa: S301 + + assert restored.dim == DIM + assert restored.table_name == "test_pickle" + assert restored.case_config.metric_type == MetricType.COSINE + + @pytest.mark.parametrize( + ("filters", "expected_sql"), + [ + (NonFilter(), 'ORDER BY "embedding" <=>'), + (IntFilter(int_value=42, filter_rate=0.5), "WHERE id >= 42"), + (LabelFilter(label_percentage=0.2), "WHERE label = 'label_20p'"), + ], + ) + def test_filter( + self, + mocked_db_connection: tuple[MagicMock, MagicMock], + filters: Filter, + expected_sql: str, + ): + db = make_db("test_filter", drop_old=False) + + db.prepare_filter(filters) + + assert expected_sql in db._search.as_string() + + @pytest.mark.parametrize("with_scalar_labels", [False, True]) + def test_insert( + self, + mocked_db_connection: tuple[MagicMock, MagicMock], + with_scalar_labels: bool, + ): + conn, cursor = mocked_db_connection + db = make_db("test_insert", drop_old=False, with_scalar_labels=with_scalar_labels) + db.conn = conn + db.cursor = cursor + copy_writer = cursor.copy.return_value.__enter__.return_value + embeddings = [[0.1, 0.2], [0.3, 0.4]] + metadata = [7, 8] + labels = ["label-a", "label-b"] if with_scalar_labels else None + + count, error = db.insert_embeddings(embeddings, metadata, labels) + + assert error is None + assert count == 2 + expected_types = ["bigint", "vector", "varchar"] if with_scalar_labels else ["bigint", "vector"] + assert copy_writer.set_types.call_count == 1 + copy_writer.set_types.assert_called_with(expected_types) + rows = [call.args[0] for call in copy_writer.write_row.call_args_list] + assert [int(row[0]) for row in rows] == metadata + np.testing.assert_allclose([row[1] for row in rows], embeddings) + if with_scalar_labels: + assert [row[2] for row in rows] == labels + conn.commit.assert_called() + + def test_search( + self, + mocked_db_connection: tuple[MagicMock, MagicMock], + ): + conn, cursor = mocked_db_connection + db = make_db("test_search", drop_old=False) + db.conn = conn + db.cursor = cursor + db.prepare_filter(NonFilter()) + cursor.execute.return_value.fetchall.return_value = [(7,), (3,)] + + result = db.search_embedding([0.1, 0.2], k=2) + + assert result == [7, 3] + query_args = cursor.execute.call_args.args[1] + np.testing.assert_allclose(query_args[0], [0.1, 0.2]) + assert query_args[1] == 2 + assert cursor.execute.call_args.kwargs == {"prepare": True, "binary": True} diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index beacb37af..dbd83bb66 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -34,6 +34,7 @@ class DB(Enum): PgVectorScale = "PgVectorScale" PgDiskANN = "PgDiskANN" AlloyDB = "AlloyDB" + LakebaseVector = "LakebaseVector" Redis = "Redis" MemoryDB = "MemoryDB" Chroma = "Chroma" @@ -159,6 +160,11 @@ def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 return AlloyDB + if self == DB.LakebaseVector: + from .lakebase_vector.lakebase_vector import LakebaseVector + + return LakebaseVector + if self == DB.AliyunElasticsearch: from .aliyun_elasticsearch.aliyun_elasticsearch import AliyunElasticsearch @@ -377,6 +383,11 @@ def config_cls(self) -> type[DBConfig]: # noqa: PLR0911, PLR0912, C901, PLR0915 return AlloyDBConfig + if self == DB.LakebaseVector: + from .lakebase_vector.config import LakebaseVectorConfig + + return LakebaseVectorConfig + if self == DB.AliyunElasticsearch: from .aliyun_elasticsearch.config import AliyunElasticsearchConfig @@ -588,6 +599,11 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912, PLR0915 return _alloydb_case_config.get(index_type) + if self == DB.LakebaseVector: + from .lakebase_vector.config import _lakebase_search_case_config + + return _lakebase_search_case_config.get(index_type) + if self == DB.AliyunElasticsearch: from .elastic_cloud.config import ElasticCloudIndexConfig diff --git a/vectordb_bench/backend/clients/api.py b/vectordb_bench/backend/clients/api.py index 3828c6145..2365f72aa 100644 --- a/vectordb_bench/backend/clients/api.py +++ b/vectordb_bench/backend/clients/api.py @@ -47,6 +47,7 @@ class IndexType(StrEnum): GPU_IVF_PQ = "GPU_IVF_PQ" GPU_CAGRA = "GPU_CAGRA" SCANN = "scann" + LAKEBASE_ANN = "lakebase_ann" VCHORDRQ = "vchordrq" VCHORDG = "vchordg" SCANN_MILVUS = "SCANN_MILVUS" diff --git a/vectordb_bench/backend/clients/lakebase_vector/__init__.py b/vectordb_bench/backend/clients/lakebase_vector/__init__.py new file mode 100644 index 000000000..997c74a62 --- /dev/null +++ b/vectordb_bench/backend/clients/lakebase_vector/__init__.py @@ -0,0 +1 @@ +"""Lakebase Search client backed by the lakebase_vector extension.""" diff --git a/vectordb_bench/backend/clients/lakebase_vector/cli.py b/vectordb_bench/backend/clients/lakebase_vector/cli.py new file mode 100644 index 000000000..e772de0ee --- /dev/null +++ b/vectordb_bench/backend/clients/lakebase_vector/cli.py @@ -0,0 +1,83 @@ +import os +from typing import Annotated, Unpack + +import click +from pydantic import SecretStr + +from vectordb_bench.backend.clients import DB + +from ....cli.cli import CommonTypedDict, cli, click_parameter_decorators_from_typed_dict, run + + +class LakebaseVectorTypedDict(CommonTypedDict): + user_name: Annotated[str, click.option("--user-name", type=str, required=True)] + password: Annotated[ + str, + click.option( + "--password", + type=str, + default=lambda: os.environ.get("POSTGRES_PASSWORD", ""), + show_default="$POSTGRES_PASSWORD", + ), + ] + host: Annotated[str, click.option("--host", type=str, required=True)] + port: Annotated[int, click.option("--port", type=int, default=5432, show_default=True)] + db_name: Annotated[str, click.option("--db-name", type=str, help="Db name", required=True)] + table_name: Annotated[ + str, + click.option("--table-name", type=str, default="vdbbench_table_test", show_default=True), + ] + max_parallel_workers: Annotated[ + int | None, + click.option( + "--max-parallel-workers", + type=int, + help="Set max_parallel_maintenance_workers and max_parallel_workers for index creation", + ), + ] + probes: Annotated[ + str | None, + click.option( + "--probes", + type=str, + help=( + "Positive integer or comma-separated positive integers for lakebase_ann.probes " + "(for example: 10 or 54,380); omit to use index defaults" + ), + ), + ] + epsilon: Annotated[ + float | None, + click.option( + "--epsilon", + type=float, + default=None, + help="Lakebase ANN reranking margin; omit to use the index default", + ), + ] + + +@cli.command() +@click_parameter_decorators_from_typed_dict(LakebaseVectorTypedDict) +def LakebaseANN(**parameters: Unpack[LakebaseVectorTypedDict]): + from .config import LakebaseANNConfig, LakebaseVectorConfig + + run( + db=DB.LakebaseVector, + db_config=LakebaseVectorConfig( + db_label=parameters["db_label"], + user_name=SecretStr(parameters["user_name"]), + password=SecretStr(parameters["password"]), + host=parameters["host"], + port=parameters["port"], + db_name=parameters["db_name"], + table_name=parameters["table_name"], + ), + db_case_config=LakebaseANNConfig( + metric_type=None, + probes=parameters["probes"], + epsilon=parameters["epsilon"], + max_parallel_workers=parameters["max_parallel_workers"], + ), + **parameters, + ) diff --git a/vectordb_bench/backend/clients/lakebase_vector/config.py b/vectordb_bench/backend/clients/lakebase_vector/config.py new file mode 100644 index 000000000..9bd980323 --- /dev/null +++ b/vectordb_bench/backend/clients/lakebase_vector/config.py @@ -0,0 +1,107 @@ +from collections.abc import Sequence +from typing import Any, LiteralString, TypedDict + +from pydantic import BaseModel, SecretStr + +from ..api import DBCaseConfig, DBConfig, IndexType, MetricType + + +class LakebaseVectorConfigDict(TypedDict): + connect_config: dict[str, Any] + table_name: str + + +class LakebaseVectorParam(TypedDict): + metric_fun_op: LiteralString + + +class LakebaseSessionCommands(TypedDict): + session_options: Sequence[dict[str, Any]] + + +class LakebaseVectorConfig(DBConfig): + user_name: SecretStr = SecretStr("postgres") + password: SecretStr + host: str = "localhost" + port: int = 5432 + db_name: str + table_name: str = "vdbbench_table_test" + + def to_dict(self) -> LakebaseVectorConfigDict: + user_str = self.user_name.get_secret_value() + pwd_str = self.password.get_secret_value() + return { + "connect_config": { + "host": self.host, + "port": self.port, + "dbname": self.db_name, + "user": user_str, + "password": pwd_str, + }, + "table_name": self.table_name, + } + + +class LakebaseANNConfig(BaseModel, DBCaseConfig): + metric_type: MetricType | None = None + create_index_before_load: bool = False + create_index_after_load: bool = True + index: IndexType = IndexType.LAKEBASE_ANN + probes: str | None = None + epsilon: float | None = None + max_parallel_workers: int | None = None + + def parse_metric(self) -> str: + if self.metric_type == MetricType.L2: + return "vector_l2_ops" + if self.metric_type == MetricType.IP: + return "vector_ip_ops" + if self.metric_type == MetricType.COSINE: + return "vector_cosine_ops" + raise NotImplementedError + + def parse_metric_fun_op(self) -> LiteralString: + if self.metric_type == MetricType.L2: + return "<->" + if self.metric_type == MetricType.IP: + return "<#>" + if self.metric_type == MetricType.COSINE: + return "<=>" + raise NotImplementedError + + def index_param(self) -> dict[str, Any]: + return { + "metric": self.parse_metric(), + "index_type": self.index.value, + "max_parallel_workers": self.max_parallel_workers, + } + + def search_param(self) -> LakebaseVectorParam: + return {"metric_fun_op": self.parse_metric_fun_op()} + + def session_param(self) -> LakebaseSessionCommands: + session_options = [] + if self.probes is not None and self.probes.strip(): + session_options.append( + { + "parameter": { + "setting_name": "lakebase_ann.probes", + "val": self.probes, + }, + }, + ) + if self.epsilon is not None: + session_options.append( + { + "parameter": { + "setting_name": "lakebase_ann.epsilon", + "val": str(self.epsilon), + }, + }, + ) + return {"session_options": session_options} + + +_lakebase_search_case_config = { + IndexType.LAKEBASE_ANN: LakebaseANNConfig, +} diff --git a/vectordb_bench/backend/clients/lakebase_vector/lakebase_vector.py b/vectordb_bench/backend/clients/lakebase_vector/lakebase_vector.py new file mode 100644 index 000000000..18d3919fa --- /dev/null +++ b/vectordb_bench/backend/clients/lakebase_vector/lakebase_vector.py @@ -0,0 +1,317 @@ +"""Wrapper around the Lakebase vector database over VectorDB""" + +import logging +from collections.abc import Generator, Sequence +from contextlib import contextmanager +from typing import Any + +import numpy as np +import psycopg +from pgvector.psycopg import register_vector +from psycopg import Connection, Cursor, sql + +from vectordb_bench.backend.filter import Filter, FilterOp + +from ..api import VectorDB +from .config import LakebaseANNConfig, LakebaseVectorConfigDict + +log = logging.getLogger(__name__) + + +class LakebaseVector(VectorDB): + """Use psycopg instructions""" + + thread_safe: bool = False + supported_filter_types: list[FilterOp] = [ + FilterOp.NonFilter, + FilterOp.NumGE, + FilterOp.StrEqual, + ] + + conn: psycopg.Connection[Any] | None = None + cursor: psycopg.Cursor[Any] | None = None + + _search: sql.Composed + + def __init__( + self, + dim: int, + db_config: LakebaseVectorConfigDict, + db_case_config: LakebaseANNConfig, + drop_old: bool = False, + with_scalar_labels: bool = False, + **kwargs, + ): + self.name = "LakebaseVector" + self.case_config = db_case_config + self.table_name = db_config["table_name"] + self.connect_config = db_config["connect_config"] + self.dim = dim + self.with_scalar_labels = with_scalar_labels + self._index_name = "lakebase_vector_index" + self._primary_field = "id" + self._vector_field = "embedding" + self._scalar_label_field = "label" + + # construct basic units + self.conn, self.cursor = self._create_connection(**self.connect_config) + + # create lakebase_vector extension + self.cursor.execute("CREATE EXTENSION IF NOT EXISTS lakebase_vector CASCADE") + self.conn.commit() + + log.info(f"{self.name} case config: {self.case_config}") + if db_case_config.create_index_before_load or not db_case_config.create_index_after_load: + msg = "LakebaseVector supports only create_index_after_load" + log.error(msg) + raise RuntimeError(msg) + if drop_old: + self._drop_index() + self._drop_table() + self._create_table(dim) + self.cursor.close() + self.conn.close() + self.cursor = None + self.conn = None + + @staticmethod + def _create_connection(**kwargs) -> tuple[Connection, Cursor]: + conn = psycopg.connect(**kwargs) + register_vector(conn) + conn.autocommit = False + cursor = conn.cursor() + + assert conn is not None, "Connection is not initialized" + assert cursor is not None, "Cursor is not initialized" + + return conn, cursor + + @contextmanager + def init(self) -> Generator[None, None, None]: + self.conn, self.cursor = self._create_connection(**self.connect_config) + + # index configuration may have commands defined that we should set during each client session + 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: + self.cursor.close() + self.conn.close() + self.cursor = None + self.conn = None + + def _drop_table(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 table : {self.table_name}") + + self.cursor.execute( + sql.SQL("DROP TABLE IF EXISTS public.{table_name}").format( + table_name=sql.Identifier(self.table_name), + ), + ) + 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() + + 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( + 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): + assert self.conn is not None, "Connection is not initialized" + assert self.cursor is not None, "Cursor is not initialized" + + index_param = self.case_config.index_param() + + if index_param["max_parallel_workers"] is not None: + self.cursor.execute( + sql.SQL("SET max_parallel_maintenance_workers TO '{}';").format( + index_param["max_parallel_workers"], + ), + ) + self.cursor.execute( + sql.SQL("ALTER USER {} SET max_parallel_maintenance_workers TO '{}';").format( + sql.Identifier(self.connect_config["user"]), + index_param["max_parallel_workers"], + ), + ) + self.cursor.execute( + sql.SQL("SET max_parallel_workers TO '{}';").format( + index_param["max_parallel_workers"], + ), + ) + self.cursor.execute( + sql.SQL("ALTER USER {} SET max_parallel_workers TO '{}';").format( + sql.Identifier(self.connect_config["user"]), + index_param["max_parallel_workers"], + ), + ) + self.cursor.execute( + sql.SQL("ALTER TABLE {} SET (parallel_workers = {});").format( + sql.Identifier(self.table_name), + index_param["max_parallel_workers"], + ), + ) + self.conn.commit() + + results = self.cursor.execute(sql.SQL("SHOW max_parallel_maintenance_workers;")).fetchall() + results.extend(self.cursor.execute(sql.SQL("SHOW max_parallel_workers;")).fetchall()) + log.info(f"{self.name} parallel index creation parameters: {results}") + + def _create_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 create index : {self._index_name}") + + index_param = self.case_config.index_param() + self._set_parallel_index_build_param() + index_create_sql = sql.SQL( + "CREATE INDEX IF NOT EXISTS {index_name} ON public.{table_name} " + "USING {index_type} ({vector_field} {metric})" + ).format( + index_name=sql.Identifier(self._index_name), + table_name=sql.Identifier(self.table_name), + index_type=sql.SQL(index_param["index_type"]), + vector_field=sql.Identifier(self._vector_field), + metric=sql.SQL(index_param["metric"]), + ) + log.debug(index_create_sql.as_string(self.cursor)) + self.cursor.execute(index_create_sql) + self.conn.commit() + + def _create_table(self, dim: int): + assert self.conn is not None, "Connection is not initialized" + assert self.cursor is not None, "Cursor is not initialized" + + # create table + if self.with_scalar_labels: + self.cursor.execute( + sql.SQL( + "CREATE TABLE IF NOT EXISTS public.{table_name} " + "({primary_field} BIGINT PRIMARY KEY, {vector_field} vector({dim}), " + "{label_field} VARCHAR(64))" + ).format( + table_name=sql.Identifier(self.table_name), + primary_field=sql.Identifier(self._primary_field), + vector_field=sql.Identifier(self._vector_field), + dim=sql.Literal(dim), + label_field=sql.Identifier(self._scalar_label_field), + ) + ) + else: + self.cursor.execute( + sql.SQL( + "CREATE TABLE IF NOT EXISTS public.{table_name} " + "({primary_field} BIGINT PRIMARY KEY, {vector_field} vector({dim}))" + ).format( + table_name=sql.Identifier(self.table_name), + primary_field=sql.Identifier(self._primary_field), + vector_field=sql.Identifier(self._vector_field), + dim=sql.Literal(dim), + ) + ) + self.cursor.execute( + sql.SQL("ALTER TABLE public.{table_name} ALTER COLUMN {vector_field} SET STORAGE PLAIN").format( + table_name=sql.Identifier(self.table_name), + vector_field=sql.Identifier(self._vector_field), + ) + ) + self.conn.commit() + + def insert_embeddings( + self, + embeddings: list[list[float]], + metadata: list[int], + labels_data: list[str] | None = None, + **kwargs: Any, + ) -> tuple[int, Exception | None]: + assert self.conn is not None, "Connection is not initialized" + assert self.cursor is not None, "Cursor is not initialized" + if self.with_scalar_labels: + assert labels_data is not None, "labels_data should be provided if with_scalar_labels is set to True" + try: + metadata_arr = np.array(metadata) + embeddings_arr = np.array(embeddings) + with self.cursor.copy( + sql.SQL("COPY public.{table_name} FROM STDIN (FORMAT BINARY)").format( + table_name=sql.Identifier(self.table_name), + ) + ) as copy: + copy_types = ["bigint", "vector", "varchar"] if self.with_scalar_labels else ["bigint", "vector"] + copy.set_types(copy_types) + for i, row in enumerate(metadata_arr): + if self.with_scalar_labels: + copy.write_row((row, embeddings_arr[i], labels_data[i])) + else: + copy.write_row((row, embeddings_arr[i])) + self.conn.commit() + + return len(metadata), None + except Exception as e: + self.conn.rollback() + log.warning(f"Failed to insert data into lakebase_vector table ({self.table_name}), error: {e}") + return 0, e + + def prepare_filter(self, filters: Filter): + if filters.type == FilterOp.NonFilter: + where_clause = "" + elif filters.type == FilterOp.NumGE: + where_clause = f"WHERE {self._primary_field} >= {filters.int_value}" + elif filters.type == FilterOp.StrEqual: + where_clause = f"WHERE {self._scalar_label_field} = '{filters.label_value}'" + else: + msg = f"Not support Filter for lakebase_vector - {filters}" + raise ValueError(msg) + self._search = sql.Composed( + [ + sql.SQL( + "SELECT {primary_field} FROM public.{table_name} {where_clause} ORDER BY {vector_field} " + ).format( + primary_field=sql.Identifier(self._primary_field), + table_name=sql.Identifier(self.table_name), + where_clause=sql.SQL(where_clause), + vector_field=sql.Identifier(self._vector_field), + ), + sql.SQL(self.case_config.search_param()["metric_fun_op"]), + sql.SQL(" %s::vector LIMIT %s::int"), + ] + ) + + def search_embedding( + self, + query: list[float], + k: int = 100, + timeout: int | None = None, + **kwargs: Any, + ) -> list[int]: + assert self.conn is not None, "Connection is not initialized" + assert self.cursor is not None, "Cursor is not initialized" + result = self.cursor.execute(self._search, (np.asarray(query), k), prepare=True, binary=True) + return [int(i[0]) for i in result.fetchall()] diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index e0cb98652..1bbc462ef 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -1,6 +1,9 @@ +# ruff: noqa: I001 # Keep Lakebase immediately after AlloyDB in the PG-extension registration order. + from ..backend.clients.adbpg.cli import AdbpgNova from ..backend.clients.alisql.cli import AliSQLHNSW from ..backend.clients.alloydb.cli import AlloyDBScaNN +from ..backend.clients.lakebase_vector.cli import LakebaseANN from ..backend.clients.aws_opensearch.cli import AWSOpenSearch from ..backend.clients.chroma.cli import Chroma from ..backend.clients.clickhouse.cli import Clickhouse @@ -72,6 +75,7 @@ cli.add_command(PgVectorScaleDiskAnn) cli.add_command(PgDiskAnn) cli.add_command(AlloyDBScaNN) +cli.add_command(LakebaseANN) cli.add_command(OceanBaseHNSW) cli.add_command(OceanBaseIVF) cli.add_command(MariaDBHNSW) diff --git a/vectordb_bench/frontend/config/dbCaseConfigs.py b/vectordb_bench/frontend/config/dbCaseConfigs.py index b8acdf4c1..b8190a175 100644 --- a/vectordb_bench/frontend/config/dbCaseConfigs.py +++ b/vectordb_bench/frontend/config/dbCaseConfigs.py @@ -1683,6 +1683,56 @@ class CaseConfigInput(BaseModel): }, ) +CaseConfigParamInput_IndexType_LakebaseVector = CaseConfigInput( + label=CaseConfigParamType.IndexType, + inputHelp="Select Index Type", + inputType=InputType.Option, + inputConfig={ + "options": [ + IndexType.LAKEBASE_ANN.value, + ], + }, +) + +CaseConfigParamInput_max_parallel_workers_LakebaseVector = CaseConfigInput( + label=CaseConfigParamType.max_parallel_workers, + displayLabel="Max parallel workers", + inputHelp="Recommended value: (cpu cores - 1). This will set the parameters: max_parallel_maintenance_workers," + " max_parallel_workers & table(parallel_workers)", + inputType=InputType.Number, + inputConfig={ + "min": 0, + "max": 1024, + "value": 16, + }, +) + +CaseConfigParamInput_Probes_LakebaseVector = CaseConfigInput( + label=CaseConfigParamType.probes, + displayLabel="Probes", + inputHelp=( + "Optional positive integer or comma-separated positive integers for lakebase_ann.probes " + "(for example: 10 or 54,380); leave empty to use the server default" + ), + inputType=InputType.Text, + inputConfig={ + "value": "", + }, +) + +CaseConfigParamInput_Epsilon_LakebaseVector = CaseConfigInput( + label=CaseConfigParamType.epsilon, + displayLabel="Epsilon", + inputHelp="Optional lakebase_ann reranking margin; leave empty to use the server default", + inputType=InputType.Float, + inputConfig={ + "min": 0.0, + "max": 4.0, + "step": 0.1, + "value": None, + }, +) + CaseConfigParamInput_EFConstruction_AliES = CaseConfigInput( label=CaseConfigParamType.EFConstruction, inputType=InputType.Number, @@ -2549,6 +2599,18 @@ class CaseConfigInput(BaseModel): CaseConfigParamInput_max_parallel_workers_AlloyDB, ] +LakebaseVectorLoadingConfig = [ + CaseConfigParamInput_IndexType_LakebaseVector, + CaseConfigParamInput_max_parallel_workers_LakebaseVector, +] + +LakebaseVectorPerformanceConfig = [ + CaseConfigParamInput_IndexType_LakebaseVector, + CaseConfigParamInput_max_parallel_workers_LakebaseVector, + CaseConfigParamInput_Probes_LakebaseVector, + CaseConfigParamInput_Epsilon_LakebaseVector, +] + AliyunElasticsearchLoadingConfig = [ CaseConfigParamInput_IndexType_ES, CaseConfigParamInput_NumShards_ES, @@ -3354,6 +3416,10 @@ class FilterType(Enum): CaseLabel.Load: AlloyDBLoadConfig, CaseLabel.Performance: AlloyDBPerformanceConfig, }, + DB.LakebaseVector: { + CaseLabel.Load: LakebaseVectorLoadingConfig, + CaseLabel.Performance: LakebaseVectorPerformanceConfig, + }, DB.AliyunElasticsearch: { CaseLabel.Load: AliyunElasticsearchLoadingConfig, CaseLabel.Performance: AliyunElasticsearchPerformanceConfig, diff --git a/vectordb_bench/models.py b/vectordb_bench/models.py index 740e0ecdb..d36d958d1 100644 --- a/vectordb_bench/models.py +++ b/vectordb_bench/models.py @@ -70,6 +70,7 @@ class CaseConfigParamType(Enum): numCandidates = "num_candidates" lists = "lists" probes = "probes" + epsilon = "epsilon" quantizationType = "quantization_type" quantizationRatio = "quantization_ratio" tableQuantizationType = "table_quantization_type"