cbrkit.indexable.lancedb

  1"""LanceDB storage backend."""
  2
  3from collections.abc import Collection
  4from dataclasses import dataclass, field
  5from typing import Any, Literal, cast, override
  6
  7import lancedb as ldb
  8import numpy as np
  9from lancedb.index import FTS, BTree
 10
 11from ..helpers import get_logger
 12from ..typing import BatchConversionFunc, Casebase, IndexableFunc, NumpyArray
 13from ._common import RowCodec, _normalize_patch_keys, _sql_in_clause, make_codec
 14
 15logger = get_logger(__name__)
 16
 17
 18@dataclass(slots=True)
 19class lancedb[K: int | str, V = str](IndexableFunc[Casebase[K, V], Collection[K]]):
 20    """LanceDB storage backend.
 21
 22    Manages an embedded LanceDB database on disk.  Supports dense
 23    (vector), sparse (FTS/BM25), and hybrid index types which
 24    control what data is stored and what indices are built.
 25
 26    Warning:
 27        Persisted vectors are tied to the
 28        :paramref:`conversion_func` used when they were written.
 29        Reopening a table backed by a different embedding model
 30        silently returns wrong results when the new model has the
 31        same dimension, and raises on `INSERT` when it does not —
 32        :meth:`put_index` only re-embeds entries whose text
 33        changed.  Drop the table (or use a fresh `table_name`)
 34        when changing models.
 35
 36    Args:
 37        uri: Path to the LanceDB database directory.
 38        table_name: Table name within the database.
 39        index_type: Determines what data is stored and which
 40            indices are created.  `"dense"` stores embeddings,
 41            `"sparse"` builds an FTS index, `"hybrid"` does
 42            both.
 43        conversion_func: Embedding function.  Required for
 44            `"dense"` and `"hybrid"` index types.
 45        key_column: Column name for case keys.
 46        value_column: Column holding the embeddable text.  With the
 47            default ``V = str`` the casebase value *is* this column;
 48            with a *model* it names the model field to embed.
 49        vector_column: Column name for dense embedding vectors.
 50        model: A dataclass or pydantic :class:`~pydantic.BaseModel`
 51            describing rows richer than plain text.  When set, ``V`` is
 52            the model type: every field becomes a stored column,
 53            ``value_column`` names the embeddable field, and reads
 54            reconstruct model instances.  This replaces any side-channel
 55            metadata — extra columns ride on the typed value itself.
 56    """
 57
 58    uri: str
 59    table_name: str
 60    index_type: Literal["dense", "sparse", "hybrid"] = "dense"
 61    conversion_func: BatchConversionFunc[str, NumpyArray] | None = None
 62    key_column: str = "key"
 63    value_column: str = "value"
 64    vector_column: str = "vector"
 65    model: type[V] | None = None
 66    _db: ldb.DBConnection = field(init=False, repr=False)
 67    _table: ldb.Table | None = field(default=None, init=False, repr=False)
 68    _codec: RowCodec[V] = field(init=False, repr=False)
 69
 70    def __post_init__(self) -> None:
 71        if self.index_type in ("dense", "hybrid") and self.conversion_func is None:
 72            raise ValueError(
 73                f"conversion_func is required for index_type={self.index_type!r}"
 74            )
 75
 76        self._codec = make_codec(
 77            self.model, self.value_column, key_column=self.key_column
 78        )
 79        self._db = ldb.connect(self.uri)
 80
 81        if self.table_name in self._db.list_tables().tables:
 82            self._table = self._db.open_table(self.table_name)
 83
 84    @override
 85    def has_index(self) -> bool:
 86        """Return whether a table exists in the database."""
 87        return self._table is not None
 88
 89    def search_limit(self) -> int | None:
 90        """Return the total number of rows, or `None` when empty."""
 91        if self._table is None:
 92            return None
 93
 94        return self._table.count_rows()
 95
 96    def _build_rows(self, casebase: Casebase[K, V]) -> list[dict[str, Any]]:
 97        """Build row dicts for LanceDB from a casebase."""
 98        codec = self._codec
 99        keys = list(casebase.keys())
100        rows = [codec.encode(value) for value in casebase.values()]
101
102        if self.index_type != "sparse":
103            assert self.conversion_func is not None
104            texts = [row[self.value_column] for row in rows]
105            for row, vec in zip(rows, self.conversion_func(texts), strict=True):
106                row[self.vector_column] = np.asarray(vec).tolist()
107
108        for row, key in zip(rows, keys, strict=True):
109            row[self.key_column] = key
110
111        return rows
112
113    def _setup_indices(self, table: ldb.Table) -> None:
114        """Create scalar and optional FTS indices on a table."""
115        table.create_index(self.key_column, config=BTree())
116
117        if self.index_type in ("sparse", "hybrid"):
118            table.create_index(self.value_column, config=FTS())
119
120    @property
121    @override
122    def index(self) -> Casebase[K, V]:
123        """Return the indexed casebase from the LanceDB table."""
124        if self._table is None:
125            return {}
126        codec = self._codec
127        needed = codec.columns
128        table = self._table.to_arrow()
129        keys = table.column(self.key_column).to_pylist()
130        payloads = {c: table.column(c).to_pylist() for c in needed}
131        return {
132            key: codec.decode({c: payloads[c][i] for c in needed})
133            for i, key in enumerate(keys)
134        }
135
136    @override
137    def put_index(self, data: Casebase[K, V]) -> None:
138        """Replace the LanceDB table contents with *data*."""
139        if self._table is None:
140            if not data:
141                return
142
143            rows = self._build_rows(data)
144            self._table = self._db.create_table(
145                self.table_name,
146                rows,
147                mode="overwrite",
148            )
149            self._setup_indices(self._table)
150            return
151
152        if not data:
153            self._table.delete("true")
154            return
155
156        rows = self._build_rows(data)
157        (
158            self._table.merge_insert(self.key_column)
159            .when_matched_update_all()
160            .when_not_matched_insert_all()
161            .when_not_matched_by_source_delete()
162            .execute(rows)
163        )
164
165    @override
166    def upsert_index(self, data: Casebase[K, V]) -> None:
167        """Insert or replace rows in the LanceDB table.
168
169        If no table exists yet, delegates to :meth:`put_index`.
170        """
171        if self._table is None:
172            self.put_index(data)
173            return
174
175        if not data:
176            return
177
178        rows = self._build_rows(data)
179        (
180            self._table.merge_insert(self.key_column)
181            .when_matched_update_all()
182            .when_not_matched_insert_all()
183            .execute(rows)
184        )
185
186    @override
187    def delete_index(
188        self,
189        data: Collection[K],
190    ) -> None:
191        """Delete rows from the LanceDB table by key."""
192        if self._table is None or not data:
193            return
194
195        self._table.delete(_sql_in_clause(self.key_column, data))
196
197    @override
198    def patch_index(
199        self,
200        upsert: Casebase[K, V] | None = None,
201        delete: Collection[K] | None = None,
202    ) -> None:
203        """Apply inserts, replacements, and deletes as one LanceDB mutation."""
204        normalized = _normalize_patch_keys(upsert, delete)
205
206        if normalized is None:
207            return
208
209        _, delete_keys = normalized
210
211        if self._table is None:
212            if upsert:
213                self.put_index(upsert)
214            return
215
216        if not upsert:
217            self.delete_index(delete_keys)
218            return
219
220        rows = self._build_rows(upsert)
221        operation = (
222            self._table.merge_insert(self.key_column)
223            .when_matched_update_all()
224            .when_not_matched_insert_all()
225        )
226
227        if delete_keys:
228            operation = operation.when_not_matched_by_source_delete(
229                _sql_in_clause(self.key_column, delete_keys)
230            )
231
232        operation.execute(rows)
233
234    def keys_where(self, where: str | None = None) -> list[K]:
235        """Return keys matching a native LanceDB predicate."""
236        if self._table is None:
237            return []
238
239        query = self._table.search().select([self.key_column])
240
241        if where is not None:
242            query = query.where(where)
243
244        table = query.to_arrow()
245        return cast(list[K], table.column(self.key_column).to_pylist())
246
247    def delete_where(
248        self,
249        where: str,
250    ) -> list[K]:
251        """Delete rows matching a native LanceDB predicate and return their keys."""
252        if self._table is None:
253            return []
254
255        keys = self.keys_where(where)
256
257        if not keys:
258            return []
259
260        self._table.delete(where)
261        return keys
262
263    def replace_where(
264        self,
265        where: str,
266        data: Casebase[K, V],
267    ) -> list[K]:
268        """Replace rows matching a native LanceDB predicate with *data*."""
269        if self._table is None:
270            self.put_index(data)
271            return []
272
273        keys = self.keys_where(where)
274
275        if not data:
276            if keys:
277                self._table.delete(where)
278            return keys
279
280        rows = self._build_rows(data)
281        (
282            self._table.merge_insert(self.key_column)
283            .when_matched_update_all()
284            .when_not_matched_insert_all()
285            .when_not_matched_by_source_delete(where)
286            .execute(rows)
287        )
288        return keys
289
290
291__all__ = ["lancedb"]