diff --git a/README.md b/README.md index 7851ef81b..a2c0ad892 100644 --- a/README.md +++ b/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]` | diff --git a/pyproject.toml b/pyproject.toml index 223291721..909458130 100644 --- a/pyproject.toml +++ b/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" ] diff --git a/tests/test_sqlite_vector.py b/tests/test_sqlite_vector.py new file mode 100644 index 000000000..26e540a4a --- /dev/null +++ b/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 diff --git a/vectordb_bench/backend/clients/__init__.py b/vectordb_bench/backend/clients/__init__.py index beacb37af..85fe8c93e 100644 --- a/vectordb_bench/backend/clients/__init__.py +++ b/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 diff --git a/vectordb_bench/backend/clients/sqlite_vector/cli.py b/vectordb_bench/backend/clients/sqlite_vector/cli.py new file mode 100644 index 000000000..1a4a0e36d --- /dev/null +++ b/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, + ) diff --git a/vectordb_bench/backend/clients/sqlite_vector/config.py b/vectordb_bench/backend/clients/sqlite_vector/config.py new file mode 100644 index 000000000..6225bb3b1 --- /dev/null +++ b/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 {} diff --git a/vectordb_bench/backend/clients/sqlite_vector/sqlite_vector.py b/vectordb_bench/backend/clients/sqlite_vector/sqlite_vector.py new file mode 100644 index 000000000..e9d46d79f --- /dev/null +++ b/vectordb_bench/backend/clients/sqlite_vector/sqlite_vector.py @@ -0,0 +1,185 @@ +import importlib.resources +import sqlite3 +import threading +from collections.abc import Iterator +from contextlib import contextmanager, suppress +from pathlib import Path + +import numpy as np +import sqlite_vector + +from ..api import MetricType, VectorDB +from .config import SQLiteVectorIndexConfig + +_DISTANCE_BY_METRIC = { + MetricType.COSINE: "COSINE", + MetricType.L2: "L2", + MetricType.IP: "DOT", + MetricType.DP: "DOT", +} + +_SEARCH_SQL = """ + SELECT rowid + FROM vector_full_scan('vectors', 'embedding', ?, ?) + """ + + +class SQLiteVector(VectorDB): + name = "SQLiteVector" + thread_safe = False + + def __init__( + self, + dim: int, + db_config: dict, + db_case_config: SQLiteVectorIndexConfig | None, + collection_name: str = "vector_bench_test", + drop_old: bool = False, + **kwargs, + ) -> None: + del collection_name, kwargs + self.dim = dim + self.db_path = Path(db_config["db_path"]).expanduser() + self.connection: sqlite3.Connection | None = None + 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 = _DISTANCE_BY_METRIC[metric_type] + except KeyError: + msg = f"Unsupported metric type: {metric_type}" + raise ValueError(msg) from None + self.vector_options = f"type=FLOAT32,dimension={dim},distance={distance}" + self.extension_path = str(importlib.resources.files(sqlite_vector) / "binaries" / "vector") + + 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.connection is not None: + raise RuntimeError("SQLite-vector connection is already open") + connection = self._connect() + self.connection = connection + try: + yield + finally: + with self._operation_lock: + try: + connection.close() + finally: + self.connection = None + + 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 + vectors = np.asarray(embeddings, dtype=" list[int]: + del kwargs + vector = np.asarray(query, dtype=" None: + del data_size + with self._operation_lock: + connection = self._connection() + connection.execute("PRAGMA optimize").fetchall() + connection.commit() + + def __getstate__(self) -> dict[str, object]: + state = self.__dict__.copy() + state.pop("_operation_lock", None) + state["connection"] = None + return state + + def __setstate__(self, state: dict[str, object]) -> None: + self.__dict__.update(state) + self._operation_lock = threading.Lock() + + def _connection(self) -> sqlite3.Connection: + if self.connection is None: + raise RuntimeError("Call init() before using the SQLite-vector client") + return self.connection + + def _create_table(self) -> None: + connection = self._open_connection() + try: + connection.execute(""" + CREATE TABLE IF NOT EXISTS vectors ( + id INTEGER PRIMARY KEY, + embedding BLOB NOT NULL + ) + """) + self._initialize_vector(connection) + connection.commit() + finally: + connection.close() + + def _remove_database(self) -> None: + for suffix in ("", "-journal", "-shm", "-wal"): + Path(f"{self.db_path}{suffix}").unlink(missing_ok=True) + + def _open_connection(self) -> sqlite3.Connection: + connection = sqlite3.connect(self.db_path, check_same_thread=False) + try: + connection.enable_load_extension(True) + try: + connection.load_extension(self.extension_path) + finally: + connection.enable_load_extension(False) + except Exception: + connection.close() + raise + return connection + + def _connect(self) -> sqlite3.Connection: + connection = self._open_connection() + try: + self._initialize_vector(connection) + except Exception: + connection.close() + raise + return connection + + def _initialize_vector(self, connection: sqlite3.Connection) -> None: + connection.execute( + "SELECT vector_init('vectors', 'embedding', ?)", + (self.vector_options,), + ).fetchone() diff --git a/vectordb_bench/cli/vectordbbench.py b/vectordb_bench/cli/vectordbbench.py index e0cb98652..9db3ebc58 100644 --- a/vectordb_bench/cli/vectordbbench.py +++ b/vectordb_bench/cli/vectordbbench.py @@ -43,6 +43,7 @@ from ..backend.clients.redis.cli import Redis from ..backend.clients.s3_vectors.cli import S3Vectors from ..backend.clients.seekdb.cli import SeekDBHNSW +from ..backend.clients.sqlite_vector.cli import SQLiteVector from ..backend.clients.tencent_elasticsearch.cli import TencentElasticsearch from ..backend.clients.test.cli import Test from ..backend.clients.tidb.cli import TiDB @@ -99,6 +100,7 @@ cli.add_command(TurboPuffer) cli.add_command(TurboPufferUnpin) cli.add_command(Chroma) +cli.add_command(SQLiteVector) cli.add_command(Zvec) cli.add_command(Endee) cli.add_command(LindormIVFPQ)