Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 6881351

Browse filesBrowse files
alaeddine-13alaeddine-13
authored
fix: cast column to its correct type before insertion in database (#392)
1 parent ae4d63b commit 6881351
Copy full SHA for 6881351

6 files changed

+148-11Lines changed: 148 additions & 11 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

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

Copy file name to clipboardExpand all lines: docarray/array/storage/annlite/backend.py
+7-3Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,8 @@
1111

1212
import numpy as np
1313

14-
from ..base.backend import BaseBackendMixin
15-
from ....helper import dataclass_from_dict, filter_dict
14+
from ..base.backend import BaseBackendMixin, TypeMap
15+
from ....helper import dataclass_from_dict, filter_dict, _safe_cast_int
1616

1717
if TYPE_CHECKING:
1818
from ....typing import DocumentArraySourceType, ArrayType
@@ -33,7 +33,11 @@ class AnnliteConfig:
3333
class BackendMixin(BaseBackendMixin):
3434
"""Provide necessary functions to enable this storage backend."""
3535

36-
TYPE_MAP = {'str': 'TEXT', 'float': 'float', 'int': 'integer'}
36+
TYPE_MAP = {
37+
'str': TypeMap(type='TEXT', converter=str),
38+
'float': TypeMap(type='float', converter=float),
39+
'int': TypeMap(type='integer', converter=_safe_cast_int),
40+
}
3741

3842
def _map_embedding(self, embedding: 'ArrayType') -> 'ArrayType':
3943
if embedding is None:
Collapse file

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

Copy file name to clipboardExpand all lines: docarray/array/storage/base/backend.py
+8-2Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,16 @@
11
from abc import ABC
2+
from collections import namedtuple
23
from dataclasses import is_dataclass, asdict
34
from typing import Dict, Optional, TYPE_CHECKING
45

56
if TYPE_CHECKING:
67
from ....typing import DocumentArraySourceType, ArrayType
78

9+
TypeMap = namedtuple('TypeMap', ['type', 'converter'])
10+
811

912
class BaseBackendMixin(ABC):
10-
TYPE_MAP: Dict
13+
TYPE_MAP: Dict[str, TypeMap]
1114

1215
def _init_storage(
1316
self,
@@ -25,10 +28,13 @@ def _get_storage_infos(self) -> Optional[Dict]:
2528
def _map_id(self, _id: str) -> str:
2629
return _id
2730

31+
def _map_column(self, value, col_type) -> str:
32+
return self.TYPE_MAP[col_type].converter(value)
33+
2834
def _map_embedding(self, embedding: 'ArrayType') -> 'ArrayType':
2935
from ....math.ndarray import to_numpy_array
3036

3137
return to_numpy_array(embedding)
3238

3339
def _map_type(self, col_type: str) -> str:
34-
return self.TYPE_MAP[col_type]
40+
return self.TYPE_MAP[col_type].type
Collapse file

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

Copy file name to clipboardExpand all lines: docarray/array/storage/weaviate/backend.py
+12-4Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,8 @@
1414
import weaviate
1515

1616
from .... import Document
17-
from ....helper import dataclass_from_dict, filter_dict
18-
from ..base.backend import BaseBackendMixin
17+
from ....helper import dataclass_from_dict, filter_dict, _safe_cast_int
18+
from ..base.backend import BaseBackendMixin, TypeMap
1919
from ..registry import _REGISTRY
2020

2121
if TYPE_CHECKING:
@@ -52,7 +52,11 @@ class WeaviateConfig:
5252
class BackendMixin(BaseBackendMixin):
5353
"""Provide necessary functions to enable this storage backend."""
5454

55-
TYPE_MAP = {'str': 'string', 'float': 'number', 'int': 'int'}
55+
TYPE_MAP = {
56+
'str': TypeMap(type='string', converter=str),
57+
'float': TypeMap(type='number', converter=float),
58+
'int': TypeMap(type='int', converter=_safe_cast_int),
59+
}
5660

5761
def _init_storage(
5862
self,
@@ -310,7 +314,11 @@ def _doc2weaviate_create_payload(self, value: 'Document'):
310314
:param value: document to create a payload for
311315
:return: the payload dictionary
312316
"""
313-
extra_columns = {col: value.tags.get(col) for col, _ in self._config.columns}
317+
columns_dict = {key: val for [key, val] in self._config.columns}
318+
extra_columns = {
319+
col: self._map_column(value.tags.get(col), columns_dict[col])
320+
for col, _ in self._config.columns
321+
}
314322

315323
return dict(
316324
data_object={
Collapse file

‎docarray/helper.py‎

Copy file name to clipboardExpand all lines: docarray/helper.py
+12-1Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import sys
66
import uuid
77
import warnings
8-
from typing import Any, Dict, Optional, Sequence, Tuple
8+
from typing import Any, Dict, Optional, Sequence, Tuple, Union
99

1010
__resources_path__ = os.path.join(
1111
os.path.dirname(
@@ -441,3 +441,14 @@ def filter_dict(d: Dict) -> Dict:
441441
:return: filtered dict
442442
"""
443443
return dict(filter(lambda item: item[1] is not None, d.items()))
444+
445+
446+
def _safe_cast_int(value: Union[str, int, float]) -> int:
447+
"""Safely cast string and float to an integer
448+
It mainly avoids silently rounding down the float value
449+
:param value: value to be cast
450+
:return: cast integer
451+
"""
452+
if isinstance(value, float) and not value.is_integer():
453+
raise ValueError(f"Can't safely cast {value} to an int")
454+
return int(value)
Collapse file

‎tests/unit/array/test_backend_configuration.py‎

Copy file name to clipboardExpand all lines: tests/unit/array/test_backend_configuration.py
+97-1Lines changed: 97 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,10 @@
1+
from typing import Tuple, Iterator
2+
3+
import pytest
14
import requests
5+
import itertools
26

3-
from docarray import DocumentArray
7+
from docarray import DocumentArray, Document
48

59

610
def test_weaviate_hnsw(start_storage):
@@ -45,3 +49,95 @@ def test_weaviate_hnsw(start_storage):
4549
assert main_class.get('vectorIndexConfig', {}).get('cleanupIntervalSeconds') == 1000
4650
assert main_class.get('vectorIndexConfig', {}).get('skip') is True
4751
assert main_class.get('vectorIndexConfig', {}).get('distance') == 'l2-squared'
52+
53+
54+
def test_weaviate_da_w_protobuff(start_storage):
55+
56+
N = 10
57+
58+
index = DocumentArray(
59+
storage='weaviate',
60+
config={
61+
'name': 'Test',
62+
'columns': [('price', 'int')],
63+
},
64+
)
65+
66+
docs = DocumentArray([Document(tags={'price': i}) for i in range(N)])
67+
docs = DocumentArray.from_protobuf(
68+
docs.to_protobuf()
69+
) # same as streaming the da in jina
70+
71+
index.extend(docs)
72+
73+
assert len(index) == N
74+
75+
76+
@pytest.mark.parametrize('type_da', [int, float, str])
77+
@pytest.mark.parametrize('type_column', ['int', 'float', 'str'])
78+
def test_cast_columns_weaviate(start_storage, type_da, type_column, request):
79+
80+
test_id = request.node.callspec.id.replace(
81+
'-', ''
82+
) # remove '-' from the test id for the weaviate name
83+
N = 10
84+
85+
index = DocumentArray(
86+
storage='weaviate',
87+
config={
88+
'name': f'Test{test_id}',
89+
'columns': [('price', type_column)],
90+
},
91+
)
92+
93+
docs = DocumentArray([Document(tags={'price': type_da(i)}) for i in range(10)])
94+
95+
index.extend(docs)
96+
97+
assert len(index) == N
98+
99+
100+
@pytest.mark.parametrize('type_da', [int, float, str])
101+
@pytest.mark.parametrize('type_column', ['int', 'float', 'str'])
102+
def test_cast_columns_annlite(start_storage, type_da, type_column):
103+
104+
N = 10
105+
106+
index = DocumentArray(
107+
storage='annlite',
108+
config={
109+
'n_dim': 3,
110+
'columns': [('price', type_column)],
111+
},
112+
)
113+
114+
docs = DocumentArray([Document(tags={'price': type_da(i)}) for i in range(10)])
115+
116+
index.extend(docs)
117+
118+
assert len(index) == N
119+
120+
121+
@pytest.mark.parametrize('type_da', [int, float, str])
122+
@pytest.mark.parametrize('type_column', ['int', 'float', 'str'])
123+
def test_cast_columns_qdrant(start_storage, type_da, type_column, request):
124+
125+
test_id = request.node.callspec.id.replace(
126+
'-', ''
127+
) # remove '-' from the test id for the weaviate name
128+
N = 10
129+
130+
index = DocumentArray(
131+
storage='qdrant',
132+
config={
133+
'collection_name': f'test{test_id}',
134+
'n_dim': 3,
135+
'columns': [('price', type_column)],
136+
},
137+
)
138+
139+
docs = DocumentArray([Document(tags={'price': type_da(i)}) for i in range(10)])
140+
141+
index.extend(docs)
142+
143+
assert len(index) == N
Collapse file

‎tests/unit/test_helper.py‎

Copy file name to clipboardExpand all lines: tests/unit/test_helper.py
+12Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
add_protocol_and_compress_to_file_path,
99
filter_dict,
1010
get_full_version,
11+
_safe_cast_int,
1112
)
1213

1314

@@ -61,3 +62,14 @@ def test_filter_dict():
6162
def test_ci_vendor():
6263
if 'GITHUB_WORKFLOW' in os.environ:
6364
assert get_full_version()['ci-vendor'] == 'GITHUB_ACTIONS'
65+
66+
67+
@pytest.mark.parametrize('input,output', [(1, 1), (1.0, 1), ('1', 1)])
68+
def test_safe_cast(input, output):
69+
assert output == _safe_cast_int(input)
70+
71+
72+
@pytest.mark.parametrize('wrong_input', [1.5, 1.001, 2 / 3])
73+
def test_safe_cast_raise_error(wrong_input):
74+
with pytest.raises(ValueError):
75+
_safe_cast_int(wrong_input)

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.