Add lancedb

Author: s-h-a-d-o-w

README.md

@@ -267,13 +267,13 @@ pip install -e '.[pinecone]'
 ```
 ### Run test server
 ```
-$ python -m vectordb_bench
+python -m vectordb_bench
 ```
 
 OR:
 
 ```shell
-$ init_bench
+init_bench
 ```
 
 OR:
@@ -290,13 +290,13 @@ After reopen the repository in container, run `python -m vectordb_bench` in the
 
 ### Check coding styles
 ```shell
-$ make lint
+make lint
 ```
 
 To fix the coding styles automatically
 
 ```shell
-$ make format
+make format
 ```
 
 ## How does it work?

pyproject.toml

@@ -72,6 +72,7 @@ all = [
     "PyMySQL",
     "clickhouse-connect",
     "pyvespa",
+    "lancedb",
 ]
 
 qdrant          = [ "qdrant-client" ]
@@ -94,6 +95,7 @@ mariadb         = [ "mariadb" ]
 tidb            = [ "PyMySQL" ]
 clickhouse      = [ "clickhouse-connect" ]
 vespa           = [ "pyvespa" ]
+lancedb         = [ "lancedb" ]
 
 [project.urls]
 "repository" = "https://github.com/zilliztech/VectorDBBench"

vectordb_bench/backend/clients/__init__.py

@@ -45,9 +45,10 @@ class DB(Enum):
     TiDB = "TiDB"
     Clickhouse = "Clickhouse"
     Vespa = "Vespa"
+    LanceDB = "LanceDB"
 
     @property
-    def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
+    def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901, PLR0915
         """Import while in use"""
         if self == DB.Milvus:
             from .milvus.milvus import Milvus
@@ -164,11 +165,16 @@ def init_cls(self) -> type[VectorDB]:  # noqa: PLR0911, PLR0912, C901
 
             return Vespa
 
+        if self == DB.LanceDB:
+            from .lancedb.lancedb import LanceDB
+
+            return LanceDB
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
     @property
-    def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901
+    def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901, PLR0915
         """Import while in use"""
         if self == DB.Milvus:
             from .milvus.config import MilvusConfig
@@ -285,6 +291,11 @@ def config_cls(self) -> type[DBConfig]:  # noqa: PLR0911, PLR0912, C901
 
             return VespaConfig
 
+        if self == DB.LanceDB:
+            from .lancedb.config import LanceDBConfig
+
+            return LanceDBConfig
+
         msg = f"Unknown DB: {self.name}"
         raise ValueError(msg)
 
@@ -382,6 +393,11 @@ def case_config_cls(  # noqa: C901, PLR0911, PLR0912
 
             return VespaHNSWConfig
 
+        if self == DB.LanceDB:
+            from .lancedb.config import _lancedb_case_config
+
+            return _lancedb_case_config.get(index_type)
+
         # DB.Pinecone, DB.Chroma, DB.Redis
         return EmptyDBCaseConfig
 

vectordb_bench/backend/clients/api.py

@@ -34,6 +34,7 @@ class IndexType(str, Enum):
     GPU_IVF_PQ = "GPU_IVF_PQ"
     GPU_CAGRA = "GPU_CAGRA"
     SCANN = "scann"
+    NONE = "NONE"
 
 
 class SQType(str, Enum):

vectordb_bench/backend/clients/lancedb/cli.py

@@ -0,0 +1,92 @@
+from typing import Annotated, Unpack
+
+import click
+from pydantic import SecretStr
+
+from ....cli.cli import (
+    CommonTypedDict,
+    cli,
+    click_parameter_decorators_from_typed_dict,
+    run,
+)
+from .. import DB
+from ..api import IndexType
+
+
+class LanceDBTypedDict(CommonTypedDict):
+    uri: Annotated[
+        str,
+        click.option("--uri", type=str, help="URI connection string", required=True),
+    ]
+    token: Annotated[
+        str | None,
+        click.option("--token", type=str, help="Authentication token", required=False),
+    ]
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(LanceDBTypedDict)
+def LanceDB(**parameters: Unpack[LanceDBTypedDict]):
+    from .config import LanceDBConfig, _lancedb_case_config
+
+    run(
+        db=DB.LanceDB,
+        db_config=LanceDBConfig(
+            db_label=parameters["db_label"],
+            uri=parameters["uri"],
+            token=SecretStr(parameters["token"]) if parameters.get("token") else None,
+        ),
+        db_case_config=_lancedb_case_config.get("NONE")(),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(LanceDBTypedDict)
+def LanceDBAutoIndex(**parameters: Unpack[LanceDBTypedDict]):
+    from .config import LanceDBConfig, _lancedb_case_config
+
+    run(
+        db=DB.LanceDB,
+        db_config=LanceDBConfig(
+            db_label=parameters["db_label"],
+            uri=parameters["uri"],
+            token=SecretStr(parameters["token"]) if parameters.get("token") else None,
+        ),
+        db_case_config=_lancedb_case_config.get(IndexType.AUTOINDEX)(),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(LanceDBTypedDict)
+def LanceDBIVFPQ(**parameters: Unpack[LanceDBTypedDict]):
+    from .config import LanceDBConfig, _lancedb_case_config
+
+    run(
+        db=DB.LanceDB,
+        db_config=LanceDBConfig(
+            db_label=parameters["db_label"],
+            uri=parameters["uri"],
+            token=SecretStr(parameters["token"]) if parameters.get("token") else None,
+        ),
+        db_case_config=_lancedb_case_config.get(IndexType.IVFPQ)(),
+        **parameters,
+    )
+
+
+@cli.command()
+@click_parameter_decorators_from_typed_dict(LanceDBTypedDict)
+def LanceDBHNSW(**parameters: Unpack[LanceDBTypedDict]):
+    from .config import LanceDBConfig, _lancedb_case_config
+
+    run(
+        db=DB.LanceDB,
+        db_config=LanceDBConfig(
+            db_label=parameters["db_label"],
+            uri=parameters["uri"],
+            token=SecretStr(parameters["token"]) if parameters.get("token") else None,
+        ),
+        db_case_config=_lancedb_case_config.get(IndexType.HNSW)(),
+        **parameters,
+    )

vectordb_bench/backend/clients/lancedb/config.py

@@ -0,0 +1,103 @@
+from pydantic import BaseModel, SecretStr
+
+from ..api import DBCaseConfig, DBConfig, IndexType, MetricType
+
+
+class LanceDBConfig(DBConfig):
+    """LanceDB connection configuration."""
+
+    db_label: str
+    uri: str
+    token: SecretStr | None = None
+
+    def to_dict(self) -> dict:
+        return {
+            "uri": self.uri,
+            "token": self.token.get_secret_value() if self.token else None,
+        }
+
+
+class LanceDBIndexConfig(BaseModel, DBCaseConfig):
+    index: IndexType = IndexType.IVFPQ
+    metric_type: MetricType = MetricType.L2
+    num_partitions: int = 0
+    num_sub_vectors: int = 0
+    nbits: int = 8  # Must be 4 or 8
+    sample_rate: int = 256
+    max_iterations: int = 50
+
+    def index_param(self) -> dict:
+        if self.index not in [
+            IndexType.IVFPQ,
+            IndexType.HNSW,
+            IndexType.AUTOINDEX,
+            IndexType.NONE,
+        ]:
+            msg = f"Index type {self.index} is not supported for LanceDB!"
+            raise ValueError(msg)
+
+        # See https://lancedb.github.io/lancedb/python/python/#lancedb.table.Table.create_index
+        params = {
+            "metric": self.parse_metric(),
+            "num_bits": self.nbits,
+            "sample_rate": self.sample_rate,
+            "max_iterations": self.max_iterations,
+        }
+
+        if self.num_partitions > 0:
+            params["num_partitions"] = self.num_partitions
+        if self.num_sub_vectors > 0:
+            params["num_sub_vectors"] = self.num_sub_vectors
+
+        return params
+
+    def search_param(self) -> dict:
+        pass
+
+    def parse_metric(self) -> str:
+        if self.metric_type in [MetricType.L2, MetricType.COSINE]:
+            return self.metric_type.value.lower()
+        if self.metric_type in [MetricType.IP, MetricType.DP]:
+            return "dot"
+        msg = f"Metric type {self.metric_type} is not supported for LanceDB!"
+        raise ValueError(msg)
+
+
+class LanceDBNoIndexConfig(LanceDBIndexConfig):
+    index: IndexType = IndexType.NONE
+
+    def index_param(self) -> dict:
+        return {}
+
+
+class LanceDBAutoIndexConfig(LanceDBIndexConfig):
+    index: IndexType = IndexType.AUTOINDEX
+
+    def index_param(self) -> dict:
+        return {}
+
+
+class LanceDBHNSWIndexConfig(LanceDBIndexConfig):
+    index: IndexType = IndexType.HNSW
+    m: int = 0
+    ef_construction: int = 0
+
+    def index_param(self) -> dict:
+        params = LanceDBIndexConfig.index_param(self)
+
+        # See https://lancedb.github.io/lancedb/python/python/#lancedb.index.HnswSq
+        params["index_type"] = "IVF_HNSW_SQ"
+        if self.m > 0:
+            params["m"] = self.m
+        if self.ef_construction > 0:
+            params["ef_construction"] = self.ef_construction
+
+        return params
+
+
+_lancedb_case_config = {
+    IndexType.IVFPQ: LanceDBIndexConfig,
+    IndexType.AUTOINDEX: LanceDBAutoIndexConfig,
+    IndexType.HNSW: LanceDBHNSWIndexConfig,
+    IndexType.NONE: LanceDBNoIndexConfig,
+}

vectordb_bench/backend/clients/lancedb/lancedb.py

@@ -0,0 +1,91 @@
+import logging
+from contextlib import contextmanager
+
+import lancedb
+import pyarrow as pa
+from lancedb.pydantic import LanceModel
+
+from ..api import IndexType, VectorDB
+from .config import LanceDBConfig, LanceDBIndexConfig
+
+log = logging.getLogger(__name__)
+
+
+class VectorModel(LanceModel):
+    id: int
+    vector: list[float]
+
+
+class LanceDB(VectorDB):
+    def __init__(
+        self,
+        dim: int,
+        db_config: LanceDBConfig,
+        db_case_config: LanceDBIndexConfig,
+        collection_name: str = "vector_bench_test",
+        drop_old: bool = False,
+        **kwargs,
+    ):
+        self.name = "LanceDB"
+        self.db_config = db_config
+        self.case_config = db_case_config
+        self.table_name = collection_name
+        self.dim = dim
+        self.uri = db_config["uri"]
+
+        db = lancedb.connect(self.uri)
+
+        if drop_old:
+            try:
+                db.drop_table(self.table_name)
+            except Exception as e:
+                log.warning(f"Failed to drop table {self.table_name}: {e}")
+
+        try:
+            db.open_table(self.table_name)
+        except Exception:
+            schema = pa.schema(
+                [pa.field("id", pa.int64()), pa.field("vector", pa.list_(pa.float64(), list_size=self.dim))]
+            )
+            db.create_table(self.table_name, schema=schema, mode="overwrite")
+
+    @contextmanager
+    def init(self):
+        self.db = lancedb.connect(self.uri)
+        self.table = self.db.open_table(self.table_name)
+        yield
+        self.db = None
+        self.table = None
+
+    def insert_embeddings(
+        self,
+        embeddings: list[list[float]],
+        metadata: list[int],
+    ) -> tuple[int, Exception | None]:
+        try:
+            data = [{"id": meta, "vector": emb} for meta, emb in zip(metadata, embeddings, strict=False)]
+            self.table.add(data)
+            return len(metadata), None
+        except Exception as e:
+            log.warning(f"Failed to insert data into LanceDB table ({self.table_name}), error: {e}")
+            return 0, e
+
+    def search_embedding(
+        self,
+        query: list[float],
+        k: int = 100,
+        filters: dict | None = None,
+    ) -> list[int]:
+        if filters:
+            results = self.table.search(query).where(f"id >= {filters['id']}", prefilter=True).limit(k).to_list()
+        else:
+            results = self.table.search(query).limit(k).to_list()
+        return [int(result["id"]) for result in results]
+
+    def optimize(self, data_size: int | None = None):
+        if self.table and hasattr(self, "case_config") and self.case_config.index != IndexType.NONE:
+            log.info(f"Creating index for LanceDB table ({self.table_name})")
+            self.table.create_index(**self.case_config.index_param())
+            # Better recall with IVF_PQ (though still bad) but breaks HNSW: https://github.com/lancedb/lancedb/issues/2369
+            if self.case_config.index in (IndexType.IVFPQ, IndexType.AUTOINDEX):
+                self.table.optimize()