Skip to content

Commit 2196082

Browse files
numb3r3numb3r3
authored
fix: rename pqlite to annlite (#151)
1 parent a4e5a38 commit 2196082

30 files changed

+246
-212
lines changed

‎docarray/array/annlite.py‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
from .document import DocumentArray
2+
from .storage.annlite import StorageMixins, AnnliteConfig
3+
4+
__all__ = ['AnnliteConfig', 'DocumentArrayAnnlite']
5+
6+
7+
class DocumentArrayAnnlite(StorageMixins, DocumentArray):
8+
def __new__(cls, *args, **kwargs):
9+
return super().__new__(cls)

‎docarray/array/document.py‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,10 +7,10 @@
77
from ..types import DocumentArraySourceType
88
from .memory import DocumentArrayInMemory
99
from .sqlite import DocumentArraySqlite
10-
from .pqlite import DocumentArrayPqlite
10+
from .annlite import DocumentArrayAnnlite
1111
from .weaviate import DocumentArrayWeaviate
1212
from .storage.sqlite import SqliteConfig
13-
from .storage.pqlite import PqliteConfig
13+
from .storage.annlite import AnnliteConfig
1414
from .storage.weaviate import WeaviateConfig
1515

1616

@@ -46,10 +46,10 @@ def __new__(
4646
def __new__(
4747
cls,
4848
_docs: Optional['DocumentArraySourceType'] = None,
49-
storage: str = 'pqlite',
50-
config: Optional[Union['PqliteConfig', Dict]] = None,
51-
) -> 'DocumentArrayPqlite':
52-
"""Create a PQLite-powered DocumentArray object."""
49+
storage: str = 'annlite',
50+
config: Optional[Union['AnnliteConfig', Dict]] = None,
51+
) -> 'DocumentArrayAnnlite':
52+
"""Create a AnnLite-powered DocumentArray object."""
5353
...
5454

5555
def __new__(cls, *args, storage: str = 'memory', **kwargs):
@@ -62,10 +62,10 @@ def __new__(cls, *args, storage: str = 'memory', **kwargs):
6262
from .sqlite import DocumentArraySqlite
6363

6464
instance = super().__new__(DocumentArraySqlite)
65-
elif storage == 'pqlite':
66-
from .pqlite import DocumentArrayPqlite
65+
elif storage == 'annlite':
66+
from .annlite import DocumentArrayAnnlite
6767

68-
instance = super().__new__(DocumentArrayPqlite)
68+
instance = super().__new__(DocumentArrayAnnlite)
6969
elif storage == 'weaviate':
7070
from .weaviate import DocumentArrayWeaviate
7171

‎docarray/array/mixins/find.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ def find(
105105
result: List['DocumentArray']
106106

107107
if isinstance(_result, list) and isinstance(_result[0], DocumentArray):
108-
# already auto-boxed by the storage backend, e.g. pqlite
108+
# already auto-boxed by the storage backend, e.g. annlite
109109
result = _result
110110
elif (
111111
isinstance(_result, tuple)

‎docarray/array/pqlite.py‎

Lines changed: 0 additions & 9 deletions
This file was deleted.

docarray/array/storage/pqlite/__init__.py renamed to docarray/array/storage/annlite/__init__.py

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,11 @@
11
from abc import ABC
2-
import numpy as np
32

4-
from typing import TYPE_CHECKING
5-
6-
if TYPE_CHECKING:
7-
from .... import Document
8-
9-
from .backend import BackendMixin, PqliteConfig
3+
from .backend import BackendMixin, AnnliteConfig
104
from .find import FindMixin
115
from .getsetdel import GetSetDelMixin
126
from .seqlike import SequenceLikeMixin
137

14-
__all__ = ['StorageMixins', 'PqliteConfig']
8+
__all__ = ['StorageMixins', 'AnnliteConfig']
159

1610

1711
class StorageMixins(FindMixin, BackendMixin, GetSetDelMixin, SequenceLikeMixin, ABC):

docarray/array/storage/pqlite/backend.py renamed to docarray/array/storage/annlite/backend.py

Lines changed: 25 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import itertools
2+
import numpy as np
23
from dataclasses import dataclass, asdict, field
34
from typing import (
45
Union,
@@ -10,9 +11,6 @@
1011
Iterator,
1112
)
1213

13-
import numpy as np
14-
from pqlite import PQLite
15-
1614
from ..base.backend import BaseBackendMixin
1715
from ....helper import dataclass_from_dict
1816

@@ -21,7 +19,7 @@
2119

2220

2321
@dataclass
24-
class PqliteConfig:
22+
class AnnliteConfig:
2523
n_dim: int
2624
metric: str = 'cosine'
2725
serialize_config: Dict = field(default_factory=dict)
@@ -31,16 +29,28 @@ class PqliteConfig:
3129
class BackendMixin(BaseBackendMixin):
3230
"""Provide necessary functions to enable this storage backend. """
3331

32+
def _map_embedding(self, embedding: 'ArrayType') -> 'ArrayType':
33+
if embedding is None:
34+
embedding = np.zeros(self.n_dim, dtype=np.float32)
35+
elif isinstance(embedding, list):
36+
from ....math.ndarray import to_numpy_array
37+
38+
embedding = to_numpy_array(embedding)
39+
40+
if embedding.ndim > 1:
41+
embedding = np.asarray(embedding).squeeze()
42+
return embedding
43+
3444
def _init_storage(
3545
self,
3646
_docs: Optional['DocumentArraySourceType'] = None,
37-
config: Optional[Union[PqliteConfig, Dict]] = None,
47+
config: Optional[Union[AnnliteConfig, Dict]] = None,
3848
**kwargs,
3949
):
4050
if not config:
4151
raise ValueError('Config object must be specified')
4252
elif isinstance(config, dict):
43-
config = dataclass_from_dict(PqliteConfig, config)
53+
config = dataclass_from_dict(AnnliteConfig, config)
4454

4555
self._persist = bool(config.data_path)
4656

@@ -52,9 +62,11 @@ def _init_storage(
5262
self._config = config
5363

5464
config = asdict(config)
55-
n_dim = config.pop('n_dim')
65+
self.n_dim = config.pop('n_dim')
5666

57-
self._pqlite = PQLite(n_dim, lock=False, **config)
67+
from annlite import AnnLite
68+
69+
self._annlite = AnnLite(self.n_dim, lock=False, **config)
5870
from ... import DocumentArray
5971
from .... import Document
6072

@@ -74,7 +86,7 @@ def _init_storage(
7486

7587
def __getstate__(self):
7688
state = dict(self.__dict__)
77-
del state['_pqlite']
89+
del state['_annlite']
7890
del state['_offsetmapping']
7991
return state
8092

@@ -85,21 +97,14 @@ def __setstate__(self, state):
8597
config = asdict(config)
8698
n_dim = config.pop('n_dim')
8799

88-
from pqlite import PQLite
100+
from annlite import AnnLite
89101

90-
self._pqlite = PQLite(n_dim, lock=False, **config)
102+
self._annlite = AnnLite(n_dim, lock=False, **config)
91103

92104
def _get_storage_infos(self) -> Dict:
93105
return {
94-
'Backend': 'PQLite',
95-
'Distance Metric': self._pqlite.metric.name,
106+
'Backend': 'AnnLite',
107+
'Distance Metric': self._annlite.metric.name,
96108
'Data Path': self._config.data_path,
97109
'Serialization Protocol': self._config.serialize_config.get('protocol'),
98110
}
99-
100-
def _map_embedding(self, embedding: 'ArrayType') -> 'ArrayType':
101-
if embedding is None:
102-
embedding = np.zeros(self._pqlite.dim, dtype=np.float32)
103-
elif isinstance(embedding, list):
104-
embedding = np.array(embedding, dtype=np.float32)
105-
return embedding

docarray/array/storage/pqlite/find.py renamed to docarray/array/storage/annlite/find.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ def _find(
3333
if n_rows == 1:
3434
query = query.reshape(1, -1)
3535

36-
_, match_docs = self._pqlite._search_documents(
36+
_, match_docs = self._annlite._search_documents(
3737
query, limit=limit, include_metadata=not only_id
3838
)
3939

docarray/array/storage/pqlite/getsetdel.py renamed to docarray/array/storage/annlite/getsetdel.py

Lines changed: 9 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -13,33 +13,32 @@ class GetSetDelMixin(BaseGetSetDelMixin):
1313
# essential methods start
1414

1515
def _get_doc_by_id(self, _id: str) -> 'Document':
16-
doc = self._pqlite.get_doc_by_id(_id)
16+
doc = self._annlite.get_doc_by_id(_id)
1717
if doc is None:
1818
raise KeyError(f'Can not find Document with id=`{_id}`')
1919
return doc
2020

2121
def _set_doc_by_id(self, _id: str, value: 'Document'):
2222
if _id != value.id:
23-
self._pqlite.delete([_id])
23+
self._del_doc_by_id(_id)
24+
2425
value.embedding = self._map_embedding(value.embedding)
2526
docs = DocumentArrayInMemory([value])
26-
self._pqlite.update(docs)
27+
self._annlite.update(docs)
2728

2829
def _del_doc_by_id(self, _id: str):
29-
self._pqlite.delete([_id])
30+
self._annlite.delete([_id])
3031

3132
def _clear_storage(self):
32-
self._pqlite.clear()
33+
self._annlite.clear()
3334

3435
def _set_docs_by_ids(self, ids, docs: Iterable['Document'], mismatch_ids: Dict):
35-
self._pqlite.delete(list(mismatch_ids.keys()))
36-
docs = DocumentArrayInMemory(docs)
37-
for doc in docs:
36+
for _id, doc in zip(ids, docs):
3837
doc.embedding = self._map_embedding(doc.embedding)
39-
self._pqlite.update(docs)
38+
self._set_doc_by_id(_id, doc)
4039

4140
def _del_docs_by_ids(self, ids):
42-
self._pqlite.delete(ids)
41+
self._annlite.delete(ids)
4342

4443
def _load_offset2ids(self):
4544
self._offsetmapping = OffsetMapping(

docarray/array/storage/pqlite/helper.py renamed to docarray/array/storage/annlite/helper.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
from pathlib import Path
22
from typing import Optional, List, Tuple
33

4-
from pqlite.storage.table import Table
4+
from annlite.storage.table import Table
55

66

77
class OffsetMapping(Table):

docarray/array/storage/pqlite/seqlike.py renamed to docarray/array/storage/annlite/seqlike.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,16 @@ def extend(self, values: Iterable['Document']) -> None:
1717
for doc in docs:
1818
doc.embedding = self._map_embedding(doc.embedding)
1919

20-
self._pqlite.index(docs)
20+
self._annlite.index(docs)
2121
self._offset2ids.extend([doc.id for doc in docs])
2222

2323
def __del__(self) -> None:
2424
if not self._persist:
2525
self._offset2ids.clear()
26-
self._pqlite.clear()
26+
self._annlite.clear()
2727

2828
def __eq__(self, other):
29-
"""In pqlite backend, data are considered as identical if configs point to the same database source"""
29+
"""In annlite backend, data are considered as identical if configs point to the same database source"""
3030
return (
3131
type(self) is type(other)
3232
and type(self._config) is type(other._config)
@@ -41,7 +41,7 @@ def __bool__(self):
4141
return len(self) > 0
4242

4343
def __repr__(self):
44-
return f'<DocumentArray[PQLite] (length={len(self)}) at {id(self)}>'
44+
return f'<DocumentArray[AnnLite] (length={len(self)}) at {id(self)}>'
4545

4646
def __add__(self, other: Union['Document', Sequence['Document']]):
4747
v = type(self)(self)
@@ -50,8 +50,8 @@ def __add__(self, other: Union['Document', Sequence['Document']]):
5050

5151
def __contains__(self, x: Union[str, 'Document']):
5252
if isinstance(x, str):
53-
return self._pqlite.get_doc_by_id(x) is not None
53+
return self._annlite.get_doc_by_id(x) is not None
5454
elif isinstance(x, Document):
55-
return self._pqlite.get_doc_by_id(x.id) is not None
55+
return self._annlite.get_doc_by_id(x.id) is not None
5656
else:
5757
return False

0 commit comments

Comments
 (0)