feat(duckdb): add exact-search backend

Author: uwu-420

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]`    |

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" ]

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]]

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
 

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,
+    )

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 {}