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

Commit d89ae6c

Browse files
author
Joan Fontanals Martinez
committed
feat: add update method to BaseDocument and fix reduce behavior
Signed-off-by: Joan Fontanals Martinez <joan.martinez@jina.ai>
1 parent 7c3d9f4 commit d89ae6c

4 files changed

Lines changed: 106 additions & 45 deletions

File tree

‎docarray/base_document/document.py‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,3 +46,39 @@ def __str__(self):
4646

4747
def _get_string_for_regex_filter(self):
4848
return str(self)
49+
50+
def update(self, other: 'BaseDocument'):
51+
"""
52+
Updates the content of this Document with the contents of other using
53+
:func:`~docarray.utils.reduce.reduce_docs`.
54+
55+
It behaves as an update operation for Dictionaries, except that since
56+
it is applied to a static schema type, the presence of the field is
57+
given by the field not having a None value.
58+
59+
EXAMPLE USAGE
60+
61+
.. code-block:: python
62+
63+
from docarray import BaseDocument
64+
from docarray.documents import Text
65+
66+
class MyDocument(BaseDocument):
67+
content: str
68+
title: Optional[str] = None
69+
tags_: List
70+
71+
doc1 = MyDocument(content='Core content of the document',
72+
title='Title', tags_=['python', 'AI'])
73+
doc2 = MyDocument(content='Core content updated', tags_=['docarray'])
74+
75+
doc1.update(doc2)
76+
assert doc1.content == 'Core content updated'
77+
assert doc1.title == 'Title'
78+
assert doc1.tags_ == ['python', 'AI', 'docarray']
79+
80+
:param other: The Document with which to update the contents of this
81+
"""
82+
from docarray.utils.reduce import reduce_docs
83+
84+
reduce_docs(self, other)

‎docarray/utils/reduce.py‎

Lines changed: 36 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
from docarray import DocumentArray
2-
from typing import List, Optional, Dict, _GenericAlias # type: ignore
2+
from typing import List, Optional, Dict, TypeVar, _GenericAlias # type: ignore
33
from docarray.base_document import BaseDocument
44

55
from collections import namedtuple
@@ -17,6 +17,8 @@
1717
],
1818
)
1919

20+
FORBIDDEN_FIELDS_TO_UPDATE = ['ID']
21+
2022

2123
def _group_fields(doc: 'BaseDocument') -> _FieldGroups:
2224
simple_non_empty_fields: List[str] = []
@@ -27,23 +29,29 @@ def _group_fields(doc: 'BaseDocument') -> _FieldGroups:
2729
nested_docarray_fields: List[str] = []
2830

2931
for field_name, field in doc.__fields__.items():
30-
field_type = doc._get_field_type(field_name)
31-
if not isinstance(field_type, _GenericAlias) and issubclass(
32-
field_type, DocumentArray
33-
):
34-
nested_docarray_fields.append(field_name)
35-
elif isinstance(field_type, _GenericAlias) and field_type.__origin__ is list:
36-
list_fields.append(field_name)
37-
elif isinstance(field_type, _GenericAlias) and field_type.__origin__ is set:
38-
set_fields.append(field_name)
39-
elif isinstance(field_type, _GenericAlias) and field_type.__origin__ is dict:
40-
dict_fields.append(field_name)
41-
v = getattr(doc, field_name)
42-
if v:
43-
if isinstance(v, BaseDocument):
44-
nested_docs_fields.append(field_name)
32+
if field_name not in FORBIDDEN_FIELDS_TO_UPDATE:
33+
field_type = doc._get_field_type(field_name)
34+
if not isinstance(field_type, _GenericAlias) and issubclass(
35+
field_type, DocumentArray
36+
):
37+
nested_docarray_fields.append(field_name)
38+
elif (
39+
isinstance(field_type, _GenericAlias) and field_type.__origin__ is list
40+
):
41+
list_fields.append(field_name)
42+
elif isinstance(field_type, _GenericAlias) and field_type.__origin__ is set:
43+
set_fields.append(field_name)
44+
elif (
45+
isinstance(field_type, _GenericAlias) and field_type.__origin__ is dict
46+
):
47+
dict_fields.append(field_name)
4548
else:
46-
simple_non_empty_fields.append(field_name)
49+
v = getattr(doc, field_name)
50+
if v:
51+
if isinstance(v, BaseDocument):
52+
nested_docs_fields.append(field_name)
53+
else:
54+
simple_non_empty_fields.append(field_name)
4755
return _FieldGroups(
4856
simple_non_empty_fields,
4957
list_fields,
@@ -54,12 +62,15 @@ def _group_fields(doc: 'BaseDocument') -> _FieldGroups:
5462
)
5563

5664

57-
def reduce_docs(doc1: 'BaseDocument', doc2: 'BaseDocument') -> 'BaseDocument':
65+
T = TypeVar('T', bound='BaseDocument')
66+
67+
68+
def reduce_docs(doc1: 'T', doc2: 'T') -> 'T':
5869
"""
5970
Reduces doc1 and doc2 into one Document in-place. Changes are applied to doc1.
6071
Reducing 2 Documents consists in the following:
6172
- setting data properties of the second Document to the first Document
62-
if they are empty (priority to the left-most Document)
73+
if they are not None
6374
- Concatenating lists and updating sets
6475
- Reducing recursively Documents and DocumentArrays
6576
- Updating Dictionaries of the left with the right
@@ -68,7 +79,6 @@ def reduce_docs(doc1: 'BaseDocument', doc2: 'BaseDocument') -> 'BaseDocument':
6879
:return: The reduced Document
6980
"""
7081
doc1_fields = _group_fields(doc1)
71-
doc1_simple_non_empty_fields = doc1_fields.simple_non_empty_fields
7282
doc1_list_fields = doc1_fields.list_fields
7383
doc1_set_fields = doc1_fields.set_fields
7484
doc1_dict_fields = doc1_fields.dict_fields
@@ -83,12 +93,7 @@ def reduce_docs(doc1: 'BaseDocument', doc2: 'BaseDocument') -> 'BaseDocument':
8393
doc2_nested_docarray_fields = doc2_fields.nested_docarray_fields
8494
doc2_nested_docs_fields = doc2_fields.nested_docs_fields
8595

86-
# update only fields that are set in doc2 and not set in doc1
87-
update_simple_fields = set(doc2_simple_non_empty_fields) - set(
88-
doc1_simple_non_empty_fields
89-
)
90-
91-
for field in update_simple_fields:
96+
for field in doc2_simple_non_empty_fields:
9297
setattr(doc1, field, getattr(doc2, field))
9398

9499
for field in set(doc1_nested_docs_fields + doc2_nested_docs_fields):
@@ -157,7 +162,7 @@ def reduce(
157162

158163
for doc in right:
159164
if doc.id in left_id_map:
160-
reduce_docs(left[left_id_map[doc.id]], doc)
165+
left[left_id_map[doc.id]].update(doc)
161166
else:
162167
left.append(doc)
163168

@@ -187,7 +192,10 @@ def reduce_all(docarrays: List[DocumentArray]) -> DocumentArray:
187192
:return: the resulting DocumentArray
188193
"""
189194
if len(docarrays) <= 1:
190-
raise Exception('In order to reduce DocumentArrays we should have more than one DocumentArray')
195+
raise Exception(
196+
'In order to reduce DocumentArrays'
197+
' we should have more than one DocumentArray'
198+
)
191199
left = docarrays[0]
192200
others = docarrays[1:]
193201
left_id_map = {doc.id: i for i, doc in enumerate(left)}

‎tests/units/document/test_base_document.py‎

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,25 @@
1+
from typing import Optional, List
12
from docarray.base_document.document import BaseDocument
23

34

45
def test_base_document_init():
5-
66
doc = BaseDocument()
77

88
assert doc.id is not None
9+
10+
11+
def test_update():
12+
class MyDocument(BaseDocument):
13+
content: str
14+
title: Optional[str] = None
15+
tags_: List
16+
17+
doc1 = MyDocument(
18+
content='Core content of the document', title='Title', tags_=['python', 'AI']
19+
)
20+
doc2 = MyDocument(content='Core content updated', tags_=['docarray'])
21+
22+
doc1.update(doc2)
23+
assert doc1.content == 'Core content updated'
24+
assert doc1.title == 'Title'
25+
assert doc1.tags_ == ['python', 'AI', 'docarray']

‎tests/units/util/test_reduce.py‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ def doc1():
3535
),
3636
test_set={'a', 'a'},
3737
inner_doc=InnerDoc(integer=2, l=['c', 'd']),
38-
test_dict={'a': 0, 'b': 2, 'd': 4}
38+
test_dict={'a': 0, 'b': 2, 'd': 4, 'z': 3},
3939
)
4040

4141

@@ -53,36 +53,36 @@ def doc2(doc1):
5353
),
5454
test_set={'a', 'b'},
5555
inner_doc=InnerDoc(integer=3, l=['a', 'b']),
56-
test_dict={'a': 10, 'b': 10, 'c': 3}
56+
test_dict={'a': 10, 'b': 10, 'c': 3, 'z': None},
5757
)
5858

5959

6060
def test_reduce_docs(doc1, doc2):
6161
result = reduce_docs(doc1, doc2)
62-
assert result.text == 'hey here'
62+
assert result.text == 'hey here 2'
6363
assert len(result.matches) == 2
6464
assert result.categories == ['a', 'b', 'c', 'd', 'e', 'f']
6565
assert result.opt_int == 5
66-
assert result.price == 10
66+
assert result.price == 5
6767
assert result.test_set == {'a', 'b'}
6868
assert len(result.matches_with_same_id) == 1
6969
assert len(result.matches_with_same_id[0].matches) == 2
70-
assert result.inner_doc.integer == 2
70+
assert result.inner_doc.integer == 3
7171
assert result.inner_doc.l == ['c', 'd', 'a', 'b']
72-
assert result.test_dict == {'a': 10, 'b': 10, 'c': 3, 'd': 4}
72+
assert result.test_dict == {'a': 10, 'b': 10, 'c': 3, 'd': 4, 'z': None}
7373

7474
# doc1 is changed in place (no extra memory)
75-
assert doc1.text == 'hey here'
75+
assert doc1.text == 'hey here 2'
7676
assert doc1.categories == ['a', 'b', 'c', 'd', 'e', 'f']
7777
assert len(doc1.matches) == 2
7878
assert doc1.opt_int == 5
79-
assert doc1.price == 10
79+
assert doc1.price == 5
8080
assert doc1.test_set == {'a', 'b'}
8181
assert len(doc1.matches_with_same_id) == 1
8282
assert len(doc1.matches_with_same_id[0].matches) == 2
83-
assert doc1.inner_doc.integer == 2
83+
assert doc1.inner_doc.integer == 3
8484
assert doc1.inner_doc.l == ['c', 'd', 'a', 'b']
85-
assert doc1.test_dict == {'a': 10, 'b': 10, 'c': 3, 'd': 4}
85+
assert doc1.test_dict == {'a': 10, 'b': 10, 'c': 3, 'd': 4, 'z': None}
8686

8787

8888
def test_reduce_different_ids():
@@ -102,15 +102,15 @@ def test_reduce(doc1, doc2):
102102
# da1 is changed in place (no extra memory)
103103
assert len(da1) == 3
104104
merged_doc = result[0]
105-
assert merged_doc.text == 'hey here'
105+
assert merged_doc.text == 'hey here 2'
106106
assert merged_doc.categories == ['a', 'b', 'c', 'd', 'e', 'f']
107107
assert len(merged_doc.matches) == 2
108108
assert merged_doc.opt_int == 5
109-
assert merged_doc.price == 10
109+
assert merged_doc.price == 5
110110
assert merged_doc.test_set == {'a', 'b'}
111111
assert len(merged_doc.matches_with_same_id) == 1
112112
assert len(merged_doc.matches_with_same_id[0].matches) == 2
113-
assert merged_doc.inner_doc.integer == 2
113+
assert merged_doc.inner_doc.integer == 3
114114
assert merged_doc.inner_doc.l == ['c', 'd', 'a', 'b']
115115

116116

@@ -123,7 +123,7 @@ def test_reduce_all(doc1, doc2):
123123
# da1 is changed in place (no extra memory)
124124
assert len(da1) == 5
125125
merged_doc = result[0]
126-
assert merged_doc.text == 'hey here'
126+
assert merged_doc.text == 'hey here 2'
127127
assert merged_doc.categories == [
128128
'a',
129129
'b',
@@ -140,9 +140,9 @@ def test_reduce_all(doc1, doc2):
140140
]
141141
assert len(merged_doc.matches) == 2
142142
assert merged_doc.opt_int == 5
143-
assert merged_doc.price == 10
143+
assert merged_doc.price == 5
144144
assert merged_doc.test_set == {'a', 'b'}
145145
assert len(merged_doc.matches_with_same_id) == 1
146146
assert len(merged_doc.matches_with_same_id[0].matches) == 2
147-
assert merged_doc.inner_doc.integer == 2
147+
assert merged_doc.inner_doc.integer == 3
148148
assert merged_doc.inner_doc.l == ['c', 'd', 'a', 'b', 'c', 'd', 'a', 'b']

0 commit comments

Comments
 (0)