Skip to content
This repository was archived by the owner on Sep 9, 2026. It is now read-only.

Commit 45a35a5

Browse files
alaeddine-13alaeddine-13
authored
feat: DocumentArray pull out columns (#162)
* test: cover updating embeddings and find * test: assert serialized embedding equals pulled out * test: fix asserts * perf: set docs in batches in weaviate * refactor: refactor qdrant extend * test: test IDs synced * fix: make sure IDs are synced * refactor: unify ID modifier methods naming * refactor: add parent map methods * docs: reformat tables * chore: quote port mapping * fix: make static * refactor: unify embedding map methods naming * fix: lint * fix: remove wrong import * test: cover modify ID * fix: fix pqlite sync ID
1 parent c0ffa59 commit 45a35a5

26 files changed

Lines changed: 383 additions & 151 deletions

‎docarray/array/storage/base/backend.py‎

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,8 @@
11
from abc import ABC
22
from typing import Dict, Optional, TYPE_CHECKING
33

4-
from docarray.array.storage.base.helper import Offset2ID
5-
64
if TYPE_CHECKING:
7-
from ....types import (
8-
DocumentArraySourceType,
9-
)
5+
from ....types import DocumentArraySourceType, ArrayType
106

117

128
class BaseBackendMixin(ABC):
@@ -21,3 +17,11 @@ def _init_storage(
2117

2218
def _get_storage_infos(self) -> Optional[Dict]:
2319
...
20+
21+
def _map_id(self, _id: str) -> str:
22+
return _id
23+
24+
def _map_embedding(self, embedding: 'ArrayType') -> 'ArrayType':
25+
from ....math.ndarray import to_numpy_array
26+
27+
return to_numpy_array(embedding)

‎docarray/array/storage/base/getsetdel.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
Sequence,
55
Any,
66
Iterable,
7+
Dict,
78
)
89

910
from .helper import Offset2ID
@@ -151,7 +152,7 @@ def _set_doc(self, _id: str, value: 'Document'):
151152
def _set_doc_by_id(self, _id: str, value: 'Document'):
152153
...
153154

154-
def _set_docs_by_ids(self, ids, docs: Iterable['Document']):
155+
def _set_docs_by_ids(self, ids, docs: Iterable['Document'], mismatch_ids: Dict):
155156
"""This function is derived from :meth:`_set_doc_by_id`
156157
Override this function if there is a more efficient logic
157158
@@ -162,8 +163,8 @@ def _set_docs_by_ids(self, ids, docs: Iterable['Document']):
162163

163164
def _set_docs(self, ids, docs: Iterable['Document']):
164165
docs = list(docs)
165-
self._set_docs_by_ids(ids, docs)
166166
mismatch_ids = {_id: doc.id for _id, doc in zip(ids, docs) if _id != doc.id}
167+
self._set_docs_by_ids(ids, docs, mismatch_ids)
167168
self._offset2ids.update_ids(mismatch_ids)
168169

169170
def _set_docs_by_slice(self, _slice: slice, value: Sequence['Document']):

‎docarray/array/storage/pqlite/__init__.py‎

Lines changed: 0 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,3 @@
1616

1717
class StorageMixins(FindMixin, BackendMixin, GetSetDelMixin, SequenceLikeMixin, ABC):
1818
...
19-
20-
def _to_numpy_embedding(self, doc: 'Document'):
21-
if doc.embedding is None:
22-
doc.embedding = np.zeros(self._pqlite.dim, dtype=np.float32)
23-
elif isinstance(doc.embedding, list):
24-
doc.embedding = np.array(doc.embedding, dtype=np.float32)

‎docarray/array/storage/pqlite/backend.py‎

Lines changed: 10 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,15 @@
99
Generator,
1010
Iterator,
1111
)
12+
13+
import numpy as np
1214
from pqlite import PQLite
1315

1416
from ..base.backend import BaseBackendMixin
1517
from ....helper import dataclass_from_dict
1618

1719
if TYPE_CHECKING:
18-
from ....types import (
19-
DocumentArraySourceType,
20-
)
20+
from ....types import DocumentArraySourceType, ArrayType
2121

2222

2323
@dataclass
@@ -96,3 +96,10 @@ def _get_storage_infos(self) -> Dict:
9696
'Data Path': self._config.data_path,
9797
'Serialization Protocol': self._config.serialize_config.get('protocol'),
9898
}
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/getsetdel.py‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import Iterable
1+
from typing import Iterable, Dict
22

33
from .helper import OffsetMapping
44
from ..base.getsetdel import BaseGetSetDelMixin
@@ -19,7 +19,9 @@ def _get_doc_by_id(self, _id: str) -> 'Document':
1919
return doc
2020

2121
def _set_doc_by_id(self, _id: str, value: 'Document'):
22-
self._to_numpy_embedding(value)
22+
if _id != value.id:
23+
self._pqlite.delete([_id])
24+
value.embedding = self._map_embedding(value.embedding)
2325
docs = DocumentArrayInMemory([value])
2426
self._pqlite.update(docs)
2527

@@ -29,10 +31,11 @@ def _del_doc_by_id(self, _id: str):
2931
def _clear_storage(self):
3032
self._pqlite.clear()
3133

32-
def _set_docs_by_ids(self, ids, docs: Iterable['Document']):
34+
def _set_docs_by_ids(self, ids, docs: Iterable['Document'], mismatch_ids: Dict):
35+
self._pqlite.delete(list(mismatch_ids.keys()))
3336
docs = DocumentArrayInMemory(docs)
3437
for doc in docs:
35-
self._to_numpy_embedding(doc)
38+
doc.embedding = self._map_embedding(doc.embedding)
3639
self._pqlite.update(docs)
3740

3841
def _del_docs_by_ids(self, ids):

‎docarray/array/storage/pqlite/seqlike.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ def extend(self, values: Iterable['Document']) -> None:
1515
return
1616

1717
for doc in docs:
18-
self._to_numpy_embedding(doc)
18+
doc.embedding = self._map_embedding(doc.embedding)
1919

2020
self._pqlite.index(docs)
2121
self._offset2ids.extend([doc.id for doc in docs])

‎docarray/array/storage/qdrant/__init__.py‎

Lines changed: 2 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,16 @@
1-
from typing import Iterable, TYPE_CHECKING
1+
from typing import TYPE_CHECKING
22

33
from .backend import BackendMixin, QdrantConfig
4-
from .helper import DISTANCES
54
from .find import FindMixin
65
from .getsetdel import GetSetDelMixin
6+
from .helper import DISTANCES
77
from .seqlike import SequenceLikeMixin
88

99
__all__ = ['StorageMixins', 'QdrantConfig']
1010

1111
if TYPE_CHECKING:
1212
from qdrant_client import QdrantClient
1313
from qdrant_openapi_client.models.models import Distance
14-
from docarray import Document
1514

1615

1716
class StorageMixins(FindMixin, BackendMixin, GetSetDelMixin, SequenceLikeMixin):
@@ -23,11 +22,6 @@ def serialize_config(self) -> dict:
2322
def distance(self) -> 'Distance':
2423
return DISTANCES[self._config.distance]
2524

26-
def extend(self, docs: Iterable['Document']):
27-
docs = list(docs)
28-
self._upload_batch(docs)
29-
self._offset2ids.extend([doc.id for doc in docs])
30-
3125
@property
3226
def serialization_config(self) -> dict:
3327
return self._serialize_config

‎docarray/array/storage/qdrant/backend.py‎

Lines changed: 19 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
List,
1313
)
1414

15+
import numpy as np
1516
from qdrant_client import QdrantClient
1617
from qdrant_openapi_client.models.models import (
1718
Distance,
@@ -24,11 +25,10 @@
2425
from docarray.array.storage.base.backend import BaseBackendMixin
2526
from docarray.array.storage.qdrant.helper import DISTANCES
2627
from docarray.helper import dataclass_from_dict, random_identity
28+
from docarray.math.helper import EPSILON
2729

2830
if TYPE_CHECKING:
29-
from docarray.types import (
30-
DocumentArraySourceType,
31-
)
31+
from docarray.types import DocumentArraySourceType, ArrayType
3232

3333

3434
@dataclass
@@ -125,7 +125,7 @@ def _collection_exists(self, collection_name):
125125
return collection_name in collections
126126

127127
@staticmethod
128-
def _qmap(doc_id: str):
128+
def _map_id(doc_id: str):
129129
# if doc_id is a random ID in hex format, just translate back to UUID str
130130
# otherwise, create UUID5 from doc_id
131131
try:
@@ -181,3 +181,18 @@ def _get_storage_infos(self) -> Dict:
181181
'Distance': self._config.distance,
182182
'Serialization Protocol': self._config.serialize_config.get('protocol'),
183183
}
184+
185+
def _map_embedding(self, embedding: 'ArrayType') -> List[float]:
186+
if embedding is None:
187+
embedding = np.random.rand(self.n_dim)
188+
else:
189+
from ....math.ndarray import to_numpy_array
190+
191+
embedding = to_numpy_array(embedding)
192+
193+
if embedding.ndim > 1:
194+
embedding = np.asarray(embedding).squeeze()
195+
196+
if np.all(embedding == 0):
197+
embedding = embedding + EPSILON
198+
return embedding.tolist()

‎docarray/array/storage/qdrant/find.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
List,
88
)
99

10-
from .helper import QdrantStorageHelper
1110
from .... import Document, DocumentArray
1211
from ....math import ndarray
1312
from ....score import NamedScore
@@ -50,7 +49,7 @@ def distance(self) -> 'Distance':
5049
raise NotImplementedError()
5150

5251
def _find_similar_vectors(self, q: 'QdrantArrayType', limit=10):
53-
query_vector = QdrantStorageHelper.embedding_to_array(q, default_dim=0)
52+
query_vector = self._map_embedding(q)
5453

5554
search_result = self.client.search(
5655
self.collection_name,

‎docarray/array/storage/qdrant/getsetdel.py‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,6 @@
1313
from docarray import Document
1414
from docarray.array.storage.base.getsetdel import BaseGetSetDelMixin
1515
from docarray.array.storage.base.helper import Offset2ID
16-
from docarray.array.storage.qdrant.helper import QdrantStorageHelper
1716

1817

1918
class GetSetDelMixin(BaseGetSetDelMixin):
@@ -67,15 +66,15 @@ def _qdrant_to_document(self, qdrant_record: dict) -> 'Document':
6766

6867
def _document_to_qdrant(self, doc: 'Document') -> 'PointStruct':
6968
return PointStruct(
70-
id=self._qmap(doc.id),
69+
id=self._map_id(doc.id),
7170
payload=dict(_serialized=doc.to_base64(**self.serialization_config)),
72-
vector=QdrantStorageHelper.embedding_to_array(doc.embedding, self.n_dim),
71+
vector=self._map_embedding(doc.embedding),
7372
)
7473

7574
def _get_doc_by_id(self, _id: str) -> 'Document':
7675
try:
7776
resp = self.client.http.points_api.get_point(
78-
name=self.collection_name, id=self._qmap(_id)
77+
name=self.collection_name, id=self._map_id(_id)
7978
)
8079
return self._qdrant_to_document(resp.result.payload)
8180
except UnexpectedResponse as response_error:
@@ -86,7 +85,7 @@ def _del_doc_by_id(self, _id: str):
8685
self.client.http.points_api.delete_points(
8786
name=self.collection_name,
8887
wait=True,
89-
points_selector=PointIdsList(points=[self._qmap(_id)]),
88+
points_selector=PointIdsList(points=[self._map_id(_id)]),
9089
)
9190

9291
def _set_doc_by_id(self, _id: str, value: 'Document'):

0 commit comments

Comments
 (0)