-
Notifications
You must be signed in to change notification settings - Fork 243
Expand file tree
/
Copy pathbackend.py
More file actions
273 lines (233 loc) · 8.78 KB
/
Copy pathbackend.py
File metadata and controls
273 lines (233 loc) · 8.78 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
import copy
import uuid
from abc import abstractmethod
from dataclasses import dataclass, field, asdict
from typing import (
Optional,
TYPE_CHECKING,
Union,
Dict,
Iterable,
List,
Tuple,
)
import numpy as np
from qdrant_client import QdrantClient
from qdrant_client.http import models
from qdrant_client.http.models.models import (
Distance,
CreateCollection,
PointsList,
PointStruct,
HnswConfigDiff,
VectorParams,
)
from docarray import Document
from docarray.array.storage.base.backend import BaseBackendMixin, TypeMap
from docarray.array.storage.qdrant.helper import DISTANCES
from docarray.helper import dataclass_from_dict, random_identity
from docarray.math.helper import EPSILON
if TYPE_CHECKING: # pragma: no cover
from docarray.typing import DocumentArraySourceType, ArrayType
@dataclass
class QdrantConfig:
n_dim: int
distance: str = 'cosine'
collection_name: Optional[str] = None
list_like: bool = True
host: Optional[str] = field(default="localhost")
port: Optional[int] = field(default=6333)
grpc_port: Optional[int] = field(default=6334)
prefer_grpc: Optional[bool] = field(default=False)
api_key: Optional[str] = field(default=None)
https: Optional[bool] = field(default=None)
serialize_config: Dict = field(default_factory=dict)
scroll_batch_size: int = 64
ef_construct: Optional[int] = None
full_scan_threshold: Optional[int] = None
m: Optional[int] = None
columns: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
root_id: bool = True
class BackendMixin(BaseBackendMixin):
@property
@abstractmethod
def client(self) -> 'QdrantClient':
raise NotImplementedError()
@property
@abstractmethod
def collection_name(self) -> str:
raise NotImplementedError()
@property
@abstractmethod
def distance(self) -> 'Distance':
raise NotImplementedError()
@classmethod
def _tmp_collection_name(cls) -> str:
return uuid.uuid4().hex
TYPE_MAP = {
'int': TypeMap(type='integer', converter=int),
'float': TypeMap(type='float', converter=float),
'bool': TypeMap(type='int', converter=bool),
'str': TypeMap(type='keyword', converter=str),
'text': TypeMap(type='text', converter=str),
'geo': TypeMap(type='geo', converter=dict),
}
def _init_storage(
self,
docs: Optional['DocumentArraySourceType'] = None,
config: Optional[Union[QdrantConfig, Dict]] = None,
**kwargs,
):
"""Initialize qdrant storage.
:param docs: the list of documents to initialize to
:param config: the config object used to initialize connection to qdrant server
:param kwargs: extra keyword arguments
:raises ValueError: only one of name or docs can be used for initialization,
raise an error if both are provided
"""
config = copy.deepcopy(config)
self._schemas = None
if not config:
raise ValueError('Empty config is not allowed for Qdrant storage')
elif isinstance(config, dict):
config = dataclass_from_dict(QdrantConfig, config)
if config.distance not in DISTANCES.keys():
raise ValueError(
f'Invalid distance parameter, must be one of: {", ".join(DISTANCES.keys())}'
)
if not config.collection_name:
config.collection_name = self._tmp_collection_name()
self._n_dim = config.n_dim
self._distance = config.distance
self._serialize_config = config.serialize_config
self._client = QdrantClient(
host=config.host,
port=config.port,
prefer_grpc=config.prefer_grpc,
grpc_port=config.grpc_port,
api_key=config.api_key,
https=config.https,
)
self._config = config
self._list_like = config.list_like
self._config.columns = self._normalize_columns(self._config.columns)
self._config.collection_name = (
self.__class__.__name__ + random_identity()
if self._config.collection_name is None
else self._config.collection_name
)
self._initialize_qdrant_schema()
super()._init_storage(**kwargs)
if docs is None and config.collection_name:
return
# 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`
self.clear()
if isinstance(docs, Iterable):
self.extend(docs)
elif isinstance(docs, Document):
self.append(docs)
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'] = (
config_joined['collection_name'] + '_subindex_' + subindex_name
)
return config_joined
def _initialize_qdrant_schema(self):
if not self._collection_exists(self.collection_name):
hnsw_config = HnswConfigDiff(
ef_construct=self._config.ef_construct,
full_scan_threshold=self._config.full_scan_threshold,
m=self._config.m,
)
self.client.recreate_collection(
collection_name=self.collection_name,
vectors_config=VectorParams(
size=self.n_dim,
distance=self.distance,
),
hnsw_config=hnsw_config,
)
for col, coltype in self._config.columns.items():
if coltype == 'text':
self.client.create_payload_index(
collection_name=self.collection_name,
field_name=col,
field_schema=models.TextIndexParams(
type="text",
tokenizer=models.TokenizerType.WORD,
),
)
else:
self.client.create_payload_index(
collection_name=self.collection_name,
field_name=col,
field_schema=self._map_type(coltype),
)
def _collection_exists(self, collection_name):
resp = self.client.get_collections()
collections = [collection.name for collection in resp.collections]
return collection_name in collections
@staticmethod
def _map_id(doc_id: str):
# if doc_id is a random ID in hex format, just translate back to UUID str
# otherwise, create UUID5 from doc_id
try:
return str(uuid.UUID(hex=doc_id))
except ValueError:
return str(uuid.uuid5(uuid.NAMESPACE_URL, doc_id))
def __getstate__(self):
d = dict(self.__dict__)
del d['_client']
return d
def __setstate__(self, state):
self.__dict__ = state
self._client = QdrantClient(
host=state['_config'].host,
port=state['_config'].port,
prefer_grpc=state['_config'].prefer_grpc,
grpc_port=state['_config'].grpc_port,
api_key=state['_config'].api_key,
https=state['_config'].https,
)
def _get_offset2ids_meta(self) -> List[str]:
if not self._collection_exists(self.collection_name_meta):
return []
return self.client.retrieve(self.collection_name_meta, ids=[1])[0].payload.get(
'offset2id', []
)
def _update_offset2ids_meta(self):
if not self._collection_exists(self.collection_name_meta):
self.client.recreate_collection(
collection_name=self.collection_name_meta,
vectors_config={}, # no vectors
)
self.client.upsert(
collection_name=self.collection_name_meta,
points=[
PointStruct(
id=1, payload={"offset2id": self._offset2ids.ids}, vector={}
)
],
wait=True,
)
def _map_embedding(self, embedding: 'ArrayType') -> List[float]:
if embedding is None:
embedding = np.random.rand(self.n_dim)
else:
from docarray.math.ndarray import to_numpy_array
embedding = to_numpy_array(embedding)
if embedding.ndim > 1:
embedding = np.asarray(embedding).squeeze()
if embedding.ndim == 0: # scalar
embedding = np.array([embedding])
if np.all(embedding == 0):
embedding = embedding + EPSILON
return embedding.astype(float).tolist()