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