-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathbackend.py
More file actions
363 lines (307 loc) · 12.2 KB
/
Copy pathbackend.py
File metadata and controls
363 lines (307 loc) · 12.2 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
import copy
import uuid
from typing import Optional, TYPE_CHECKING, Union, Dict, Iterable, List, Tuple
from dataclasses import dataclass, field
import re
import numpy as np
from pymilvus import (
connections,
Collection,
FieldSchema,
DataType,
CollectionSchema,
has_collection,
loading_progress,
)
from docarray import Document, DocumentArray
from docarray.array.storage.base.backend import BaseBackendMixin, TypeMap
from docarray.helper import dataclass_from_dict, _safe_cast_int
from docarray.score import NamedScore
if TYPE_CHECKING:
from docarray.typing import (
DocumentArraySourceType,
)
ID_VARCHAR_LEN = 1024
SERIALIZED_VARCHAR_LEN = (
65_535 # 65_535 is the maximum that Milvus allows for a VARCHAR field
)
COLUMN_VARCHAR_LEN = 1024
OFFSET_VARCHAR_LEN = 1024
def _always_true_expr(primary_key: str) -> str:
"""
Returns a Milvus expression that is always true, thus allowing for the retrieval of all entries in a Collection
Assumes that the primary key is of type DataType.VARCHAR
:param primary_key: the name of the primary key
:return: a Milvus expression that is always true for that primary key
"""
return f'({primary_key} in ["1"]) or ({primary_key} not in ["1"])'
def _ids_to_milvus_expr(ids):
ids = ['"' + _id + '"' for _id in ids]
return '[' + ','.join(ids) + ']'
def _batch_list(l: List, batch_size: int):
"""Iterates over a list in batches of size batch_size"""
if batch_size < 1:
yield l
return
l_len = len(l)
for ndx in range(0, l_len, batch_size):
yield l[ndx : min(ndx + batch_size, l_len)]
def _sanitize_collection_name(name):
"""Removes all chars that are not allowed in a Milvus collection name.
Thus, it removes all chars that are not alphanumeric or an underscore.
:param name: the collection name to sanitize
:return: the sanitized collection name.
"""
return ''.join(
re.findall('[a-zA-Z0-9_]', name)
) # remove everything that is not a letter, number or underscore
@dataclass
class MilvusConfig:
n_dim: int
collection_name: str = None
host: str = 'localhost'
port: Optional[Union[str, int]] = 19530 # 19530 for gRPC, 9091 for HTTP
distance: str = 'IP' # metric_type in milvus
index_type: str = 'HNSW'
index_params: Dict = field(
default_factory=lambda: {
'M': 4,
'efConstruction': 200,
}
) # passed to milvus at index creation time. The default assumes 'HNSW' index type
collection_config: Dict = field(
default_factory=dict
) # passed to milvus at collection creation time
serialize_config: Dict = field(default_factory=dict)
consistency_level: str = 'Session'
batch_size: int = -1
columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
list_like: bool = True
root_id: bool = True
class BackendMixin(BaseBackendMixin):
TYPE_MAP = {
'str': TypeMap(type=DataType.VARCHAR, converter=str),
'float': TypeMap(
type=DataType.DOUBLE, converter=float
), # it doesn't like DataType.FLOAT type, perhaps because python floats are double precision?
'double': TypeMap(type=DataType.DOUBLE, converter=float),
'int': TypeMap(type=DataType.INT64, converter=_safe_cast_int),
'bool': TypeMap(type=DataType.BOOL, converter=bool),
}
def _init_storage(
self,
_docs: Optional['DocumentArraySourceType'] = None,
config: Optional[Union[MilvusConfig, Dict]] = None,
**kwargs,
):
config = copy.deepcopy(config)
if not config:
raise ValueError('Empty config is not allowed for Milvus storage')
elif isinstance(config, dict):
config = dataclass_from_dict(MilvusConfig, config)
if config.collection_name is None:
id = uuid.uuid4().hex
config.collection_name = 'docarray__' + id
self._list_like = config.list_like
self._config = config
self._config.columns = self._normalize_columns(self._config.columns)
self._connection_alias = (
f'docarray_{config.host}_{config.port}_{uuid.uuid4().hex}'
)
connections.connect(
alias=self._connection_alias, host=config.host, port=config.port
)
self._collection = self._create_or_reuse_collection()
self._offset2id_collection = self._create_or_reuse_offset2id_collection()
self._build_index()
super()._init_storage(**kwargs)
# To align with Sqlite behavior; if `docs` is not `None` and table name
# is provided, :class:`DocumentArraySqlite` will clear the existing
# table and load the given `docs`
if _docs is None:
return
self.clear()
if isinstance(_docs, Iterable):
self.extend(_docs)
else:
if isinstance(_docs, Document):
self.append(_docs)
def _create_or_reuse_collection(self):
if has_collection(self._config.collection_name, using=self._connection_alias):
return Collection(
self._config.collection_name, using=self._connection_alias
)
document_id = FieldSchema(
name='document_id',
dtype=DataType.VARCHAR,
max_length=ID_VARCHAR_LEN,
is_primary=True,
)
embedding = FieldSchema(
name='embedding', dtype=DataType.FLOAT_VECTOR, dim=self._config.n_dim
)
serialized = FieldSchema(
name='serialized', dtype=DataType.VARCHAR, max_length=SERIALIZED_VARCHAR_LEN
)
additional_columns = []
for col, coltype in self._config.columns.items():
mapped_type = self._map_type(coltype)
if mapped_type == DataType.VARCHAR:
field_ = FieldSchema(
name=col, dtype=mapped_type, max_length=COLUMN_VARCHAR_LEN
)
else:
field_ = FieldSchema(name=col, dtype=mapped_type)
additional_columns.append(field_)
schema = CollectionSchema(
fields=[document_id, embedding, serialized, *additional_columns],
description='DocumentArray collection schema',
)
return Collection(
name=self._config.collection_name,
schema=schema,
using=self._connection_alias,
**self._config.collection_config,
)
def _build_index(self):
index_params = {
'metric_type': self._config.distance,
'index_type': self._config.index_type,
'params': self._config.index_params,
}
self._collection.create_index(field_name='embedding', index_params=index_params)
def _create_or_reuse_offset2id_collection(self):
if has_collection(
self._config.collection_name + '_offset2id', using=self._connection_alias
):
return Collection(
self._config.collection_name + '_offset2id',
using=self._connection_alias,
)
document_id = FieldSchema(
name='document_id', dtype=DataType.VARCHAR, max_length=ID_VARCHAR_LEN
)
offset = FieldSchema(
name='offset',
dtype=DataType.VARCHAR,
max_length=OFFSET_VARCHAR_LEN,
is_primary=True,
)
dummy_vector = FieldSchema(
name='dummy_vector', dtype=DataType.FLOAT_VECTOR, dim=1
)
schema = CollectionSchema(
fields=[offset, document_id, dummy_vector],
description='offset2id for DocumentArray',
)
return Collection(
name=self._config.collection_name + '_offset2id',
schema=schema,
using=self._connection_alias,
# **self._config.collection_config, # we probably don't want to apply the same config here
)
def _ensure_unique_config(
self,
config_root: dict,
config_subindex: dict,
config_joined: dict,
subindex_name: str,
) -> dict:
if 'collection_name' not in config_subindex:
config_joined['collection_name'] = _sanitize_collection_name(
config_joined['collection_name'] + '_subindex_' + subindex_name
)
return config_joined
def _doc_to_milvus_payload(self, doc):
return self._docs_to_milvus_payload([doc])
def _docs_to_milvus_payload(self, docs: 'Iterable[Document]'):
extra_columns = [
[self._map_column(doc.tags.get(col), col_type) for doc in docs]
for col, col_type in self._config.columns.items()
]
return [
[doc.id for doc in docs],
[self._map_embedding(doc.embedding) for doc in docs],
[doc.to_base64(**self._config.serialize_config) for doc in docs],
*extra_columns,
]
@staticmethod
def _docs_from_query_response(response):
return DocumentArray([Document.from_base64(d['serialized']) for d in response])
@staticmethod
def _docs_from_search_response(responses, distance: str) -> 'List[DocumentArray]':
das = []
for r in responses:
da = []
for hit in r:
doc = Document.from_base64(hit.entity.get('serialized'))
doc.scores[distance] = NamedScore(value=hit.score)
da.append(doc)
das.append(DocumentArray(da))
return das
def _update_kwargs_from_config(self, field_to_update, **kwargs):
kwargs_field_value = kwargs.get(field_to_update, None)
config_field_value = getattr(self._config, field_to_update, None)
if (
kwargs_field_value is not None or config_field_value is None
): # no need to update
return kwargs
kwargs[field_to_update] = config_field_value
return kwargs
def _map_embedding(self, embedding):
if embedding is not None:
from docarray.math.ndarray import to_numpy_array
embedding = to_numpy_array(embedding)
if embedding.ndim > 1:
embedding = np.asarray(embedding).squeeze()
else:
embedding = np.zeros(self._config.n_dim)
return embedding
def __getstate__(self):
d = dict(self.__dict__)
del d['_collection']
del d['_offset2id_collection']
return d
def __setstate__(self, state):
self.__dict__ = state
connections.connect(
alias=self._connection_alias, host=self._config.host, port=self._config.port
)
self._collection = self._create_or_reuse_collection()
self._offset2id_collection = self._create_or_reuse_offset2id_collection()
def __enter__(self):
_ = super().__enter__()
self._collection.load()
self._offset2id_collection.load()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self._collection.release()
self._offset2id_collection.release()
super().__exit__(exc_type, exc_val, exc_tb)
def loaded_collection(self, collection=None):
"""
Context manager to load a collection and release it after the context is exited.
If the collection is already loaded when entering, it will not be released while exiting.
:param collection: the collection to load. If None, the main collection of this indexer is used.
:return: Context manager for the provided collection.
"""
class LoadedCollectionManager:
def __init__(self, coll, connection_alias):
self._collection = coll
self._loaded_when_enter = False
self._connection_alias = connection_alias
def __enter__(self):
self._loaded_when_enter = (
loading_progress(
self._collection.name, using=self._connection_alias
)['loading_progress']
!= '0%'
)
self._collection.load()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if not self._loaded_when_enter:
self._collection.release()
return LoadedCollectionManager(
collection if collection else self._collection, self._connection_alias
)