Source code for docarray.array.storage.annlite.getsetdel
from typing import Iterable, Dict
from docarray.array.storage.annlite.helper import OffsetMapping
from docarray.array.storage.base.getsetdel import BaseGetSetDelMixin
from docarray.array.storage.base.helper import Offset2ID
from docarray.array.memory import DocumentArrayInMemory
from docarray import Document, DocumentArray
[docs]class GetSetDelMixin(BaseGetSetDelMixin):
"""Implement required and derived functions that power `getitem`, `setitem`, `delitem`"""
# essential methods start
def _get_doc_by_id(self, _id: str) -> 'Document':
doc = self._annlite.get_doc_by_id(_id)
if doc is None:
raise KeyError(f'Can not find Document with id=`{_id}`')
return doc
def _set_doc_by_id(self, _id: str, value: 'Document'):
if _id != value.id:
self._del_doc_by_id(_id)
value.embedding = self._map_embedding(value.embedding)
docs = DocumentArrayInMemory([value])
self._annlite.update(docs)
def _del_doc_by_id(self, _id: str):
# delete the root document
self._del_docs_by_ids([_id])
def _clear_storage(self):
self._annlite.clear()
def _set_docs_by_ids(self, ids, docs: Iterable['Document'], mismatch_ids: Dict):
for _id, doc in zip(ids, docs):
doc.embedding = self._map_embedding(doc.embedding)
self._set_doc_by_id(_id, doc)
def _del_docs_by_ids(self, ids):
self._annlite.delete(ids)
def __del__(self) -> None:
self._annlite.close()
def _load_offset2ids(self):
if self._list_like:
self._offsetmapping = OffsetMapping(
data_path=self._config.data_path, in_memory=False
)
self._offsetmapping.create_table()
self._offset2ids = Offset2ID(
self._offsetmapping.get_all_ids(),
list_like=self._list_like,
)
else:
self._offset2ids = Offset2ID(
[],
list_like=self._list_like,
)
def _save_offset2ids(self):
if self._list_like:
self._offsetmapping.drop()
self._offsetmapping.create_table()
self._offsetmapping._insert(
[(i, doc_id) for i, doc_id in enumerate(self._offset2ids.ids)]
)