From b4c08afe5ca95dcd231bb40923baca7c6708444c Mon Sep 17 00:00:00 2001 From: UwU 420 <164377175+uwu-420@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:44:06 +0200 Subject: [PATCH] feat(duckdb): add exact-search backend --- README.md | 1 + pyproject.toml | 1 + tests/test_duckdb.py | 179 +++++++++++++++++ vectordb_bench/backend/clients/__init__.py | 16 ++ vectordb_bench/backend/clients/duckdb/cli.py | 46 +++++ .../backend/clients/duckdb/config.py | 22 ++ .../backend/clients/duckdb/duckdb.py | 189 ++++++++++++++++++ vectordb_bench/cli/vectordbbench.py | 2 + 8 files changed, 456 insertions(+) create mode 100644 tests/test_duckdb.py create mode 100644 vectordb_bench/backend/clients/duckdb/cli.py create mode 100644 vectordb_bench/backend/clients/duckdb/config.py create mode 100644 vectordb_bench/backend/clients/duckdb/duckdb.py diff --git a/README.md b/README.md index 85d34d76d..737802de4 100644 --- a/README.md +++ b/README.md @@ -57,6 +57,7 @@ All the database client supported | mongodb | `pip install vectordb-bench[mongodb]` | | tidb | `pip install vectordb-bench[tidb]` | | vespa | `pip install vectordb-bench[vespa]` | +| duckdb | `pip install vectordb-bench[duckdb]` | | oceanbase | `pip install vectordb-bench[oceanbase]` | | hologres | `pip install vectordb-bench[hologres]` | | tencent_es | `pip install vectordb-bench[tencent_es]` | diff --git a/pyproject.toml b/pyproject.toml index 223291721..1b94a321f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -75,6 +75,7 @@ cockroachdb = [ "psycopg[binary,pool]", "pgvector" ] clickhouse = [ "clickhouse-connect" ] vespa = [ "pyvespa" ] lancedb = [ "lancedb" ] +duckdb = [ "duckdb>=1.5.5,<2.0.0" ] oceanbase = [ "mysql-connector-python" ] alisql = [ "mysqlclient" ] polardb = [ "PyMySQL" ] diff --git a/tests/test_duckdb.py b/tests/test_duckdb.py new file mode 100644 index 000000000..4d82d9ccf --- /dev/null +++ b/tests/test_duckdb.py @@ -0,0 +1,179 @@ +import multiprocessing as mp +from collections.abc import Iterator +from concurrent.futures import ProcessPoolExecutor +from copy import deepcopy +from pathlib import Path + +import pandas as pd +import pytest + +pytest.importorskip("duckdb") + +from vectordb_bench import config +from vectordb_bench.backend.clients.api import MetricType +from vectordb_bench.backend.clients.duckdb.config import DuckDBConfig, DuckDBIndexConfig +from vectordb_bench.backend.clients.duckdb.duckdb import DuckDB +from vectordb_bench.backend.filter import non_filter +from vectordb_bench.backend.runner.concurrent_runner import ConcurrentInsertRunner +from vectordb_bench.backend.runner.rate_runner import RatedMultiThreadingInsertRunner + + +class SingleBatchDataset: + class Fields: + train_id_field = "id" + train_vector_field = "emb" + scalar_labels_file_separated = False + + data = Fields() + + def iter_batches(self, batch_size: int) -> Iterator[pd.DataFrame]: + del batch_size + yield pd.DataFrame( + { + "id": [40, 10, 30, 20], + "emb": [ + [10.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.8, 0.2, 0.0], + [-1.0, 0.0, 0.0], + ], + } + ) + + +class StreamingDataset: + def __init__(self) -> None: + self.batch_index = 0 + + def __iter__(self) -> "StreamingDataset": + return self + + def __next__(self) -> pd.DataFrame: + if self.batch_index == 4: + raise StopIteration + start = self.batch_index * config.NUM_PER_BATCH + self.batch_index += 1 + return pd.DataFrame( + { + "id": list(range(start, start + config.NUM_PER_BATCH)), + "emb": [[1.0, 0.0, 0.0]] * config.NUM_PER_BATCH, + } + ) + + +def make_client( + path: Path, + drop_old: bool = True, + metric_type: MetricType = MetricType.COSINE, +) -> DuckDB: + return DuckDB( + dim=3, + db_config=DuckDBConfig(db_path=str(path)).to_dict(), + db_case_config=DuckDBIndexConfig(metric_type=metric_type), + drop_old=drop_old, + ) + + +def search_in_process(client: DuckDB, query: list[float], k: int) -> list[int]: + client.prepare_filter(non_filter) + with client.init(): + return client.search_embedding(query, k=k) + + +@pytest.mark.parametrize( + ("metric_type", "expected"), + [ + (MetricType.COSINE, [40, 30, 10]), + (MetricType.L2, [30, 10, 20]), + (MetricType.IP, [40, 30, 10]), + (MetricType.DP, [40, 30, 10]), + ], +) +def test_duckdb_exact_search_reopens_after_copy_round_trip( + tmp_path: Path, + metric_type: MetricType, + expected: list[int], +) -> None: + client = make_client(tmp_path / "vectors.duckdb", metric_type=metric_type) + dataset = SingleBatchDataset() + batch = next(dataset.iter_batches(4)) + + with client.init(): + count, error = client.insert_embeddings(batch["emb"].tolist(), batch["id"].tolist()) + + assert error is None + assert count == 4 + + client = deepcopy(client) + with client.init(): + assert client.search_embedding([1.0, 0.0, 0.0], k=3) == expected + + +def test_duckdb_loads_through_concurrent_runner(tmp_path: Path) -> None: + client = make_client(tmp_path / "runner.duckdb") + runner = ConcurrentInsertRunner( + db=client, + dataset=SingleBatchDataset(), + normalize=False, + max_workers=4, + ) + + assert runner.max_workers == 1 + assert runner.task() == 4 + + with client.init(): + assert client.search_embedding([1.0, 0.0, 0.0], k=2) == [40, 30] + + +def test_duckdb_rolls_back_failed_load(tmp_path: Path) -> None: + client = make_client(tmp_path / "rollback.duckdb") + + def abort_load() -> None: + with client.init(): + count, error = client.insert_embeddings([[1.0, 0.0, 0.0]], [1]) + assert error is None + assert count == 1 + raise RuntimeError("abort load") + + with pytest.raises(RuntimeError, match="abort load"): + abort_load() + + with client.init(): + assert client.search_embedding([1.0, 0.0, 0.0], k=1) == [] + + +def test_duckdb_serializes_streaming_insert_threads(tmp_path: Path) -> None: + client = make_client(tmp_path / "streaming.duckdb") + runner = RatedMultiThreadingInsertRunner( + rate=config.NUM_PER_BATCH * 4, + db=client, + dataset_iter=StreamingDataset(), + ) + queue = mp.Queue() + + try: + runner.run_with_rate(queue) + finally: + queue.close() + queue.join_thread() + + expected_ids = set(range(config.NUM_PER_BATCH * 4)) + with client.init(): + assert set(client.search_embedding([1.0, 0.0, 0.0], k=len(expected_ids))) == expected_ids + + +def test_duckdb_supports_spawned_concurrent_searches(tmp_path: Path) -> None: + client = make_client(tmp_path / "concurrent.duckdb") + dataset = SingleBatchDataset() + batch = next(dataset.iter_batches(4)) + with client.init(): + count, error = client.insert_embeddings(batch["emb"].tolist(), batch["id"].tolist()) + + assert error is None + assert count == 4 + + context = mp.get_context("spawn") + with ProcessPoolExecutor(max_workers=2, mp_context=context) as executor: + futures = [executor.submit(search_in_process, client, [1.0, 0.0, 0.0], 2) for _ in range(2)] + + assert [future.result() for future in futures] == [[40, 30], [40, 30]] diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index beacb37af..ebfab23f0 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/vectordb_bench/backend/clients/__init__.py @@ -49,6 +49,7 @@ class DB(Enum): Clickhouse = "Clickhouse" Vespa = "Vespa" LanceDB = "LanceDB" + DuckDB = "DuckDB" OceanBase = "OceanBase" S3Vectors = "S3Vectors" Hologres = "Alibaba Cloud Hologres" @@ -217,6 +218,11 @@ def init_cls(self) -> type[VectorDB]: # noqa: PLR0911, PLR0912, C901, PLR0915 return LanceDB + if self == DB.DuckDB: + from .duckdb.duckdb import DuckDB + + return DuckDB + if self == DB.S3Vectors: from .s3_vectors.s3_vectors import S3Vectors @@ -435,6 +441,11 @@ def config_cls(self) -> type[DBConfig]: # noqa: PLR0911, PLR0912, C901, PLR0915 return LanceDBConfig + if self == DB.DuckDB: + from .duckdb.config import DuckDBConfig + + return DuckDBConfig + if self == DB.S3Vectors: from .s3_vectors.config import S3VectorsConfig @@ -637,6 +648,11 @@ def case_config_cls( # noqa: C901, PLR0911, PLR0912, PLR0915 return _lancedb_case_config.get(index_type) + if self == DB.DuckDB: + from .duckdb.config import DuckDBIndexConfig + + return DuckDBIndexConfig + if self == DB.S3Vectors: from .s3_vectors.config import S3VectorsIndexConfig diff --git a/vectordb_bench/backend/clients/duckdb/cli.py b/vectordb_bench/backend/clients/duckdb/cli.py new file mode 100644 index 000000000..1cc4adf4d --- /dev/null +++ b/vectordb_bench/backend/clients/duckdb/cli.py @@ -0,0 +1,46 @@ +from importlib.metadata import version +from typing import Annotated, Unpack + +import click + +from ....cli.cli import CommonTypedDict, cli, click_parameter_decorators_from_typed_dict, run +from .. import DB +from .config import DuckDBConfig, DuckDBIndexConfig + + +class DuckDBTypedDict(CommonTypedDict): + db_path: Annotated[ + str, + click.option( + "--db-path", + type=click.Path(dir_okay=False), + help="Path to a dedicated DuckDB benchmark database file.", + required=True, + ), + ] + threads: Annotated[ + int, + click.option( + "--threads", + type=click.IntRange(min=1), + default=1, + help="Number of DuckDB threads used by each benchmark process.", + show_default=True, + ), + ] + + +@cli.command(name="duckdb") +@click_parameter_decorators_from_typed_dict(DuckDBTypedDict) +def DuckDB(**parameters: Unpack[DuckDBTypedDict]) -> None: + run( + db=DB.DuckDB, + db_config=DuckDBConfig( + db_label=parameters["db_label"], + version=version("duckdb"), + db_path=parameters["db_path"], + threads=parameters["threads"], + ), + db_case_config=DuckDBIndexConfig(), + **parameters, + ) diff --git a/vectordb_bench/backend/clients/duckdb/config.py b/vectordb_bench/backend/clients/duckdb/config.py new file mode 100644 index 000000000..8238ed92d --- /dev/null +++ b/vectordb_bench/backend/clients/duckdb/config.py @@ -0,0 +1,22 @@ +from pydantic import BaseModel, Field + +from ..api import DBCaseConfig, DBConfig, IndexType, MetricType + + +class DuckDBConfig(DBConfig): + db_path: str + threads: int = Field(default=1, ge=1) + + def to_dict(self) -> dict: + return {"db_path": self.db_path, "threads": self.threads} + + +class DuckDBIndexConfig(BaseModel, DBCaseConfig): + index: IndexType = IndexType.Flat + metric_type: MetricType | None = None + + def index_param(self) -> dict: + return {} + + def search_param(self) -> dict: + return {} diff --git a/vectordb_bench/backend/clients/duckdb/duckdb.py b/vectordb_bench/backend/clients/duckdb/duckdb.py new file mode 100644 index 000000000..97c7ccd03 --- /dev/null +++ b/vectordb_bench/backend/clients/duckdb/duckdb.py @@ -0,0 +1,189 @@ +import threading +from collections.abc import Iterator +from contextlib import contextmanager, suppress +from pathlib import Path + +import duckdb +import numpy as np + +from ...filter import Filter, FilterOp +from ..api import MetricType, VectorDB +from .config import DuckDBIndexConfig + +_DISTANCE_FUNCTION_BY_METRIC = { + MetricType.COSINE: "array_cosine_distance", + MetricType.L2: "array_distance", + MetricType.IP: "array_negative_inner_product", + MetricType.DP: "array_negative_inner_product", +} + + +class DuckDB(VectorDB): + name = "DuckDB" + thread_safe = False + + def __init__( + self, + dim: int, + db_config: dict, + db_case_config: DuckDBIndexConfig | None, + collection_name: str = "vector_bench_test", + drop_old: bool = False, + **kwargs, + ) -> None: + del collection_name, kwargs + if dim <= 0: + msg = f"Embedding dimension must be positive, got {dim}" + raise ValueError(msg) + + self.dim = dim + self.db_path = Path(db_config["db_path"]).expanduser() + self.threads = int(db_config.get("threads", 1)) + self.connection: duckdb.DuckDBPyConnection | None = None + self._connection_read_only: bool | None = None + self._active = False + self._operation_lock = threading.Lock() + + metric_type = db_case_config.metric_type if db_case_config is not None else MetricType.COSINE + if metric_type is None: + metric_type = MetricType.COSINE + try: + distance_function = _DISTANCE_FUNCTION_BY_METRIC[metric_type] + except KeyError: + msg = f"Unsupported metric type: {metric_type}" + raise ValueError(msg) from None + + self._insert_sql = f"INSERT INTO vectors SELECT unnest(?::BIGINT[]), unnest(?::FLOAT[{self.dim}][])" + self._search_sql = ( + f"SELECT id FROM vectors ORDER BY {distance_function}(embedding, ?::FLOAT[{self.dim}]) LIMIT ?" + ) + + if self.db_path.exists() and self.db_path.is_dir(): + raise IsADirectoryError(self.db_path) + self.db_path.parent.mkdir(parents=True, exist_ok=True) + if drop_old: + self._remove_database() + self._create_table() + + @contextmanager + def init(self) -> Iterator[None]: + with self._operation_lock: + if self._active: + raise RuntimeError("DuckDB connection is already open") + self._active = True + try: + yield + except BaseException: + with self._operation_lock: + if self.connection is not None and self._connection_read_only is False: + with suppress(Exception): + self.connection.rollback() + raise + else: + with self._operation_lock: + if self.connection is not None and self._connection_read_only is False: + self.connection.commit() + finally: + with self._operation_lock: + try: + if self.connection is not None: + self.connection.close() + finally: + self.connection = None + self._connection_read_only = None + self._active = False + + def insert_embeddings( + self, + embeddings: list[list[float]], + metadata: list[int], + labels_data: list[str] | None = None, + tenant_labels_data: list[str] | None = None, + **kwargs, + ) -> tuple[int, Exception | None]: + del labels_data, tenant_labels_data, kwargs + try: + vectors = np.asarray(embeddings, dtype=np.float32) + ids = np.asarray(metadata, dtype=np.int64) + except Exception as error: + return 0, error + + expected_shape = (len(metadata), self.dim) + if vectors.shape != expected_shape: + msg = f"Expected embeddings with shape {expected_shape}, got {vectors.shape}" + return 0, ValueError(msg) + if ids.shape != (len(metadata),): + msg = f"Expected metadata with shape {(len(metadata),)}, got {ids.shape}" + return 0, ValueError(msg) + + try: + with self._operation_lock: + self._connection(read_only=False).execute(self._insert_sql, [ids, vectors]) + except Exception as error: + return 0, error + return len(metadata), None + + def search_embedding(self, query: list[float], k: int = 100, **kwargs) -> list[int]: + del kwargs + vector = np.asarray(query, dtype=np.float32) + expected_shape = (self.dim,) + if vector.shape != expected_shape: + msg = f"Expected query with shape {expected_shape}, got {vector.shape}" + raise ValueError(msg) + + with self._operation_lock: + rows = self._connection(read_only=True).execute(self._search_sql, [vector, k]).fetchall() + return [int(row[0]) for row in rows] + + def optimize(self, data_size: int | None = None) -> None: + del data_size + + def prepare_filter(self, filters: Filter) -> None: + if filters.type != FilterOp.NonFilter: + msg = f"Unsupported filter for DuckDB: {filters}" + raise ValueError(msg) + + def __getstate__(self) -> dict[str, object]: + state = self.__dict__.copy() + state.pop("_operation_lock", None) + state["connection"] = None + state["_connection_read_only"] = None + state["_active"] = False + return state + + def __setstate__(self, state: dict[str, object]) -> None: + self.__dict__.update(state) + self._operation_lock = threading.Lock() + + def _connection(self, read_only: bool) -> duckdb.DuckDBPyConnection: + if not self._active: + raise RuntimeError("Call init() before using the DuckDB client") + if self.connection is None: + connection = self._open_connection(read_only=read_only) + if not read_only: + connection.begin() + self.connection = connection + self._connection_read_only = read_only + elif not read_only and self._connection_read_only: + raise RuntimeError("Cannot write through a read-only DuckDB connection") + return self.connection + + def _create_table(self) -> None: + connection = self._open_connection(read_only=False) + try: + connection.execute( + f"CREATE TABLE IF NOT EXISTS vectors (id BIGINT PRIMARY KEY, embedding FLOAT[{self.dim}] NOT NULL)" + ) + finally: + connection.close() + + def _remove_database(self) -> None: + self.db_path.unlink(missing_ok=True) + Path(f"{self.db_path}.wal").unlink(missing_ok=True) + + def _open_connection(self, read_only: bool) -> duckdb.DuckDBPyConnection: + return duckdb.connect( + str(self.db_path), + read_only=read_only, + config={"threads": self.threads}, + ) diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index e0cb98652..96fb45732 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -6,6 +6,7 @@ from ..backend.clients.clickhouse.cli import Clickhouse from ..backend.clients.cockroachdb.cli import CockroachDB as CockroachDBCli from ..backend.clients.doris.cli import Doris +from ..backend.clients.duckdb.cli import DuckDB from ..backend.clients.elastic_cloud.cli import ( ElasticCloudHNSW, ElasticCloudHNSWBBQ, @@ -84,6 +85,7 @@ cli.add_command(LanceDBIVFPQ) cli.add_command(LanceDBIVFHNSWSQ) cli.add_command(LanceDBIVFHNSWPQ) +cli.add_command(DuckDB) cli.add_command(HologresHGraph) cli.add_command(QdrantCloud) cli.add_command(QdrantLocal)