feat(sqlite-vector): add SQLite-vector backend

README.md

@@ -51,6 +51,7 @@ All the database client supported
 | redis                    | `pip install vectordb-bench[redis]`         |
 | memorydb                 | `pip install vectordb-bench[memorydb]`      |
 | chromadb                 | `pip install vectordb-bench[chromadb]`      |
+| sqlite-vector            | `pip install vectordb-bench[sqlite-vector]` |
 | cockroachdb              | `pip install vectordb-bench[cockroachdb]`   |
 | awsopensearch            | `pip install vectordb-bench[opensearch]` |
 | aliyun_opensearch        | `pip install vectordb-bench[aliyun_opensearch]` |

pyproject.toml

@@ -66,6 +66,7 @@ pgvecto_rs      = [ "pgvecto_rs[psycopg3]>=0.2.2" ]
 redis           = [ "redis" ]
 memorydb        = [ "memorydb" ]
 chromadb        = [ "chromadb" ]
+sqlite-vector   = [ "sqliteai-vector>=1.0.0,<2.0.0" ]
 opensearch      = [ "opensearch-py", "boto3", "requests-aws4auth" ]
 aliyun_opensearch = [ "alibabacloud_ha3engine_vector" ]
 mongodb         = [ "pymongo" ]

tests/test_sqlite_vector.py

@@ -0,0 +1,137 @@
+import multiprocessing as mp
+from collections.abc import Iterator
+from copy import deepcopy
+from pathlib import Path
+
+import pandas as pd
+import pytest
+
+pytest.importorskip("sqlite_vector")
+
+from vectordb_bench import config
+from vectordb_bench.backend.clients.api import MetricType
+from vectordb_bench.backend.clients.sqlite_vector.config import SQLiteVectorConfig, SQLiteVectorIndexConfig
+from vectordb_bench.backend.clients.sqlite_vector.sqlite_vector import SQLiteVector
+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,
+) -> SQLiteVector:
+    return SQLiteVector(
+        dim=3,
+        db_config=SQLiteVectorConfig(db_path=str(path)).to_dict(),
+        db_case_config=SQLiteVectorIndexConfig(metric_type=metric_type),
+        drop_old=drop_old,
+    )
+
+
+@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_sqlite_vector_exact_search_reopens_after_copy_round_trip(
+    tmp_path: Path,
+    metric_type: MetricType,
+    expected: list[int],
+) -> None:
+    client = make_client(tmp_path / "vectors.db", 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_sqlite_vector_loads_through_concurrent_runner(tmp_path: Path) -> None:
+    client = make_client(tmp_path / "runner.db")
+    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_sqlite_vector_serializes_streaming_insert_threads(tmp_path: Path) -> None:
+    client = make_client(tmp_path / "streaming.db")
+    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

vectordb_bench/backend/clients/__init__.py

@@ -37,6 +37,7 @@ class DB(Enum):
     Redis = "Redis"
     MemoryDB = "MemoryDB"
     Chroma = "Chroma"
+    SQLiteVector = "SQLiteVector"
     AWSOpenSearch = "OpenSearch"
     OSSOpenSearch = "OSSOpenSearch"
     AliyunElasticsearch = "AliyunElasticsearch"
@@ -139,6 +140,11 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return ChromaClient
 
+        if self == DB.SQLiteVector:
+            from .sqlite_vector.sqlite_vector import SQLiteVector
+
+            return SQLiteVector
+
         if self == DB.AWSOpenSearch:
             from .aws_opensearch.aws_opensearch import AWSOpenSearch
 
@@ -357,6 +363,11 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901, PLR0915
 
             return ChromaConfig
 
+        if self == DB.SQLiteVector:
+            from .sqlite_vector.config import SQLiteVectorConfig
+
+            return SQLiteVectorConfig
+
         if self == DB.AWSOpenSearch:
             from .aws_opensearch.config import AWSOpenSearchConfig
 
@@ -685,6 +696,11 @@ def case_config_cls(  # noqa: C901, PLR0911, PLR0912, PLR0915
 
             return ChromaIndexConfig
 
+        if self == DB.SQLiteVector:
+            from .sqlite_vector.config import SQLiteVectorIndexConfig
+
+            return SQLiteVectorIndexConfig
+
         if self == DB.Lindorm:
             from .lindorm.config import _lindorm_vector_case_config
 

vectordb_bench/backend/clients/sqlite_vector/cli.py

@@ -0,0 +1,35 @@
+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 SQLiteVectorConfig, SQLiteVectorIndexConfig
+
+
+class SQLiteVectorTypedDict(CommonTypedDict):
+    db_path: Annotated[
+        str,
+        click.option(
+            "--db-path",
+            type=click.Path(dir_okay=False),
+            help="Path to a dedicated SQLite-vector benchmark database file.",
+            required=True,
+        ),
+    ]
+
+
+@cli.command(name="sqlite-vector")
+@click_parameter_decorators_from_typed_dict(SQLiteVectorTypedDict)
+def SQLiteVector(**parameters: Unpack[SQLiteVectorTypedDict]) -> None:
+    run(
+        db=DB.SQLiteVector,
+        db_config=SQLiteVectorConfig(
+            db_label=parameters["db_label"],
+            version=version("sqliteai-vector"),
+            db_path=parameters["db_path"],
+        ),
+        db_case_config=SQLiteVectorIndexConfig(),
+        **parameters,
+    )

vectordb_bench/backend/clients/sqlite_vector/config.py

@@ -0,0 +1,21 @@
+from pydantic import BaseModel
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class SQLiteVectorConfig(DBConfig):
+    db_path: str
+
+    def to_dict(self) -> dict:
+        return {"db_path": self.db_path}
+
+
+class SQLiteVectorIndexConfig(BaseModel, DBCaseConfig):
+    index: IndexType = IndexType.Flat
+    metric_type: MetricType | None = None
+
+    def index_param(self) -> dict:
+        return {}
+
+    def search_param(self) -> dict:
+        return {}