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

Commit 0307634

Browse files
hanxiaohanxiao
committed
feat: add field_resolver and fallback to tags
1 parent cee955f commit 0307634

8 files changed

Lines changed: 415 additions & 11 deletions

File tree

‎docarray/base.py‎

Lines changed: 27 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import copy as cp
2-
from typing import TYPE_CHECKING, Optional, Tuple
2+
from dataclasses import fields
3+
from typing import TYPE_CHECKING, Optional, Tuple, Dict
34
from .helper import typename
45

56
if TYPE_CHECKING:
@@ -19,7 +20,13 @@ def non_empty_fields(self) -> Tuple[str]:
1920
"""
2021
return self._data._non_empty_fields
2122

22-
def __init__(self: 'T', _obj: Optional['T'] = None, copy: bool = False, **kwargs):
23+
def __init__(
24+
self: 'T',
25+
_obj: Optional['T'] = None,
26+
copy: bool = False,
27+
field_resolver: Optional[Dict[str, str]] = None,
28+
**kwargs,
29+
):
2330
self._data = None
2431
if isinstance(_obj, type(self)):
2532
if copy:
@@ -28,8 +35,26 @@ def __init__(self: 'T', _obj: Optional['T'] = None, copy: bool = False, **kwargs
2835
self._data = _obj._data
2936
elif isinstance(_obj, dict):
3037
kwargs.update(_obj)
38+
3139
if kwargs:
40+
if field_resolver:
41+
kwargs = {field_resolver.get(k, k): v for k, v in kwargs.items()}
42+
43+
_unknown_kwargs = None
44+
if hasattr(self, '_unresolved_fields_dest'):
45+
_unresolved = set(kwargs.keys()).difference(
46+
{f.name for f in fields(self._data_class)}
47+
)
48+
if _unresolved:
49+
_unknown_kwargs = {k: kwargs[k] for k in _unresolved}
50+
for k in _unresolved:
51+
kwargs.pop(k)
52+
3253
self._data = self._data_class(self, **kwargs)
54+
55+
if _unknown_kwargs:
56+
getattr(self._data, self._unresolved_fields_dest).update(_unknown_kwargs)
57+
3358
if _obj is None and not kwargs:
3459
self._data = self._data_class(self)
3560

‎docarray/document/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,4 +4,5 @@
44

55

66
class Document(AllMixins, BaseDCType):
7-
_data_class = DocumentData
7+
_data_class = DocumentData
8+
_unresolved_fields_dest = 'tags'

‎docarray/helper.py‎

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import sys
44
import uuid
55
import warnings
6-
from typing import Iterable, Any, Dict, Optional
6+
from typing import Any, Dict, Optional, Sequence
77

88
__windows__ = sys.platform == 'win32'
99

@@ -67,6 +67,9 @@ def dunder_get(_dict: Any, key: str) -> Any:
6767
:return: (mixed) value corresponding to the key
6868
"""
6969

70+
if not _dict:
71+
return None
72+
7073
try:
7174
part1, part2 = key.split('__', 1)
7275
except ValueError:
@@ -77,19 +80,14 @@ def dunder_get(_dict: Any, key: str) -> Any:
7780
except ValueError:
7881
pass
7982

80-
from google.protobuf.struct_pb2 import ListValue
81-
from google.protobuf.struct_pb2 import Struct
82-
from google.protobuf.pyext._message import MessageMapContainer
83-
from .simple.struct import StructView
84-
8583
if isinstance(part1, int):
8684
result = _dict[part1]
87-
elif isinstance(_dict, (dict, Struct, MessageMapContainer, StructView)):
85+
elif isinstance(_dict, dict):
8886
if part1 in _dict:
8987
result = _dict[part1]
9088
else:
9189
result = None
92-
elif isinstance(_dict, (Iterable, ListValue)):
90+
elif isinstance(_dict, Sequence):
9391
result = _dict[part1]
9492
else:
9593
result = getattr(_dict, part1)

‎tests/unit/document/docs.csv‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
source,url,question,answer,wrong_answer
2+
testsrc,https://jina.ai,What are the symptoms?,Symptoms ...,As cases ...
3+
testsrc,https://jina.ai,When should I get tested?,Your doctor ...,If you ...

‎tests/unit/document/docs.jsonlines‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
{"text": "a"}
2+
{"text": "b"}

‎tests/unit/document/docs_groundtruth.jsonlines‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
{"document": {"text": "a"}, "groundtruth": {"text": "b"}}
2+
{"document": {"text": "c"}, "groundtruth": {"text": "d"}}

‎tests/unit/document/test_docdata.py‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from docarray import Document, DocumentArray
55
from docarray.array.chunk import ChunkArray
66
from docarray.array.match import MatchArray
7+
from docarray.score import NamedScore
78

89

910
@pytest.mark.parametrize('init_args', [None, dict(id=123), Document()])
@@ -140,3 +141,54 @@ def test_exclusive_content_2():
140141
d = Document(content=[1,2,3])
141142
assert d.content_type == 'blob'
142143

144+
def test_get_attr_values():
145+
d = Document(
146+
**{
147+
'id': '123',
148+
'text': 'document',
149+
'feature1': 121,
150+
'name': 'name',
151+
'tags': {'id': 'identity', 'a': 'b', 'c': 'd', 'e': [0, 1, {'f': 'g'}]},
152+
}
153+
)
154+
d.scores['metric'] = NamedScore(value=42)
155+
156+
required_keys = [
157+
'id',
158+
'text',
159+
'tags__name',
160+
'tags__feature1',
161+
'scores__metric__value',
162+
'tags__c',
163+
'tags__id',
164+
'tags__inexistant',
165+
'tags__e__2__f',
166+
'inexistant',
167+
]
168+
res = d.get_attributes(*required_keys)
169+
assert len(res) == len(required_keys)
170+
assert res[required_keys.index('id')] == '123'
171+
assert res[required_keys.index('tags__feature1')] == 121
172+
assert res[required_keys.index('tags__name')] == 'name'
173+
assert res[required_keys.index('text')] == 'document'
174+
assert res[required_keys.index('tags__c')] == 'd'
175+
assert res[required_keys.index('tags__id')] == 'identity'
176+
assert res[required_keys.index('scores__metric__value')] == 42
177+
assert res[required_keys.index('tags__inexistant')] is None
178+
assert res[required_keys.index('inexistant')] is None
179+
assert res[required_keys.index('tags__e__2__f')] == 'g'
180+
181+
required_keys_2 = ['tags', 'text']
182+
res2 = d.get_attributes(*required_keys_2)
183+
assert len(res2) == 2
184+
assert res2[required_keys_2.index('text')] == 'document'
185+
assert res2[required_keys_2.index('tags')] == d.tags
186+
187+
d = Document({'id': '123', 'tags': {'outterkey': {'innerkey': 'real_value'}}})
188+
required_keys_3 = ['tags__outterkey__innerkey']
189+
res3 = d.get_attributes(*required_keys_3)
190+
assert res3 == 'real_value'
191+
192+
d = Document(content=np.array([1, 2, 3]))
193+
res4 = np.stack(d.get_attributes(*['blob']))
194+
np.testing.assert_equal(res4, np.array([1, 2, 3]))

0 commit comments

Comments
 (0)