Repository: SuperDuperDB/superduperdb
File: superduperdb/vector_search/lance.py (Python)

Code before:
    @property
    def dataset(self):
        if not os.path.exists(self.dataset_path):
            self._create_or_append_to_dataset([], [])
        return lance.dataset(self.dataset_path)

    def __len__(self):
        return self.dataset.count_rows()

    def _create_or_append_to_dataset(self, vectors, ids, mode: str = 'create'):
        if not self._created:
            if not os.path.exists(self.dataset_path):
                mode = 'create'
                os.makedirs(self.dataset_path, exist_ok=True)
            else:
                self._created = True
                mode = "append"
        type = pa.list_(
            pa.field('values', pa.float32(), nullable=False), self.dimensions
        )
        vectors = self.to_list(vectors)
        _vecs = pa.array([v for v in vectors], type=type)
        _ids = pa.array(ids, type=pa.string())
        _table = pa.Table.from_arrays([_ids, _vecs], names=['id', 'vector'])
        lance.write_dataset(_table, self.dataset_path, mode=mode)
        self._created = True

    def add(self, items: t.Sequence[VectorItem]) -> None:
        ids = [item.id for item in items]
        vectors = [item.vector for item in items]
        self._create_or_append_to_dataset(vectors, ids, mode='append')

    def delete(self, ids: t.Sequence[str]) -> None:
        to_remove = ", ".join(f"'{str(id)}'" for id in ids)
        self.dataset.delete(f"id IN ({to_remove})")

    def find_nearest_from_id(
        self,
        _id,
        n: int = 100,
        within_ids: t.Sequence[str] = (),
    ) -> t.Tuple[t.List[str], t.List[float]]:
        # The ``lance`` file format has been specifically designed for fast
        # random access. The logic to take advantage of this is implemented
        # by the ``.take`` method.

        # See https://blog.lancedb.com/benchmarking-random-access-in-lance-ed690757a826
        # for further details.

        vector = self.dataset.take([int(_id)], columns=['vector']).to_pydict()[
            'vector'

Review comment (suggestion):
The details here are

### Case 1:

**`STATE_1_CREATE`: When there is no historical data during initialization**

`mode == "create"`

**`STATE_1_UPDATE`: As we continue to refresh or add new vectors**

`mode == "create" `

### Case2:

**`STATE_2_CREATE`: When there is historical data during initialization**

`mode == "upsert"`

**`STATE_2_UPDATE`: As we continue to refresh or add new vectors**

`mode == "upsert" `



In fact, the `mode` in state `STATE_1_UPDATE` and state `STATE_2_UPDATE` should be `upsert`.


Code after:

    @property
    def dataset(self):
        if not os.path.exists(self.dataset_path):
            self._create_or_append_to_dataset([], [], mode='create')
        return lance.dataset(self.dataset_path)

    def __len__(self):
        return self.dataset.count_rows()

    def _create_or_append_to_dataset(self, vectors, ids, mode: str = 'upsert'):
        type = pa.list_(
            pa.field('values', pa.float32(), nullable=False), self.dimensions
        )
        vectors = self.to_list(vectors)
        _vecs = pa.array([v for v in vectors], type=type)
        _ids = pa.array(ids, type=pa.string())
        _table = pa.Table.from_arrays([_ids, _vecs], names=['id', 'vector'])

        if mode == 'upsert':
            dataset = lance.dataset(self.dataset_path)
            dataset.merge_insert(
                "id"
            ).when_matched_update_all().when_not_matched_insert_all().execute(_table)
        else:
            lance.write_dataset(_table, self.dataset_path, mode=mode)

    def add(self, items: t.Sequence[VectorItem]) -> None:
        ids = [item.id for item in items]
        vectors = [item.vector for item in items]
        self._create_or_append_to_dataset(vectors, ids, mode='append')

    def delete(self, ids: t.Sequence[str]) -> None:
        to_remove = ", ".join(f"'{str(id)}'" for id in ids)
        self.dataset.delete(f"id IN ({to_remove})")

    def find_nearest_from_id(
        self,
        _id,
        n: int = 100,
        within_ids: t.Sequence[str] = (),
    ) -> t.Tuple[t.List[str], t.List[float]]:
        # The ``lance`` file format has been specifically designed for fast
        # random access. The logic to take advantage of this is implemented
        # by the ``.take`` method.

        # See https://blog.lancedb.com/benchmarking-random-access-in-lance-ed690757a826
        # for further details.

        vector = self.dataset.take([int(_id)], columns=['vector']).to_pydict()[
            'vector'