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

Commit 1c15390

Browse files
Charlotte GerhaherJackmin801
Charlotte Gerhaher
authored and
Jackmin801
committed
feat: nested attribute access in find() (#1176)
* test: add test for find on nested attribute Signed-off-by: anna-charlotte <[email protected]> * test: add test for nested attr in stacked da Signed-off-by: anna-charlotte <[email protected]> * fix: nested access Signed-off-by: anna-charlotte <[email protected]> * fix: mypy Signed-off-by: anna-charlotte <[email protected]> * fix: field type by access path for da Signed-off-by: anna-charlotte <[email protected]> * fix: clean up Signed-off-by: anna-charlotte <[email protected]> * fix: move get field type by access path Signed-off-by: anna-charlotte <[email protected]> * fix: imports Signed-off-by: anna-charlotte <[email protected]> * fix: apply suggestion and fix import Signed-off-by: anna-charlotte <[email protected]> * fix: apply suggestion from code review Signed-off-by: anna-charlotte <[email protected]> * fix: apply suggestion from code review Signed-off-by: anna-charlotte <[email protected]> * fix: apply suggestions Signed-off-by: anna-charlotte <[email protected]> --------- Signed-off-by: anna-charlotte <[email protected]>
1 parent a29b2b1 commit 1c15390

4 files changed

Lines changed: 92 additions & 37 deletions

File tree

‎docarray/helper.py‎

Lines changed: 32 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from typing import TYPE_CHECKING, Any, Dict, List, Type
1+
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
22

33
if TYPE_CHECKING:
44
from docarray import BaseDocument
@@ -8,21 +8,9 @@ def _is_access_path_valid(doc_type: Type['BaseDocument'], access_path: str) -> b
88
"""
99
Check if a given access path ("__"-separated) is a valid path for a given Document class.
1010
"""
11-
from docarray import BaseDocument
1211

13-
field, _, remaining = access_path.partition('__')
14-
if len(remaining) == 0:
15-
return access_path in doc_type.__fields__.keys()
16-
else:
17-
valid_field = field in doc_type.__fields__.keys()
18-
if not valid_field:
19-
return False
20-
else:
21-
d = doc_type._get_field_type(field)
22-
if not issubclass(d, BaseDocument):
23-
return False
24-
else:
25-
return _is_access_path_valid(d, remaining)
12+
field_type = _get_field_type_by_access_path(doc_type, access_path)
13+
return field_type is not None
2614

2715

2816
def _all_access_paths_valid(
@@ -121,3 +109,32 @@ def _update_nested_dicts(
121109
to_update[k] = v
122110
else:
123111
_update_nested_dicts(to_update[k], update_with[k])
112+
113+
114+
def _get_field_type_by_access_path(
115+
doc_type: Type['BaseDocument'], access_path: str
116+
) -> Optional[Type]:
117+
"""
118+
Get field type by "__"-separated access path.
119+
:param doc_type: type of document
120+
:param access_path: "__"-separated access path
121+
:return: field type of accessed attribute. If access path is invalid, return None.
122+
"""
123+
from docarray import BaseDocument, DocumentArray
124+
125+
field, _, remaining = access_path.partition('__')
126+
field_valid = field in doc_type.__fields__.keys()
127+
128+
if field_valid:
129+
if len(remaining) == 0:
130+
return doc_type._get_field_type(field)
131+
else:
132+
d = doc_type._get_field_type(field)
133+
if issubclass(d, DocumentArray):
134+
return _get_field_type_by_access_path(d.document_type, remaining)
135+
elif issubclass(d, BaseDocument):
136+
return _get_field_type_by_access_path(d, remaining)
137+
else:
138+
return None
139+
else:
140+
return None

‎docarray/utils/find.py‎

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from docarray.array.array.array import DocumentArray
77
from docarray.array.stacked.array_stacked import DocumentArrayStacked
88
from docarray.base_document import BaseDocument
9+
from docarray.helper import _get_field_type_by_access_path
910
from docarray.typing import AnyTensor
1011
from docarray.typing.tensor.abstract_tensor import AbstractTensor
1112

@@ -192,8 +193,8 @@ class MyDocument(BaseDocument):
192193
comp_backend = embedding_type.get_comp_backend()
193194

194195
# extract embeddings from query and index
195-
index_embeddings = _extraxt_embeddings(index, embedding_field, embedding_type)
196-
query_embeddings = _extraxt_embeddings(query, embedding_field, embedding_type)
196+
index_embeddings = _extract_embeddings(index, embedding_field, embedding_type)
197+
query_embeddings = _extract_embeddings(query, embedding_field, embedding_type)
197198

198199
# compute distances and return top results
199200
metric_fn = getattr(comp_backend.Metrics, metric)
@@ -225,7 +226,7 @@ def _extract_embedding_single(
225226
:return: the embeddings
226227
"""
227228
if isinstance(data, BaseDocument):
228-
emb = getattr(data, embedding_field)
229+
emb = next(AnyDocumentArray._traverse(data, embedding_field))
229230
else: # treat data as tensor
230231
emb = data
231232
if len(emb.shape) == 1:
@@ -235,7 +236,7 @@ def _extract_embedding_single(
235236
return emb
236237

237238

238-
def _extraxt_embeddings(
239+
def _extract_embeddings(
239240
data: Union[AnyDocumentArray, BaseDocument, AnyTensor],
240241
embedding_field: str,
241242
embedding_type: Type,
@@ -247,40 +248,41 @@ def _extraxt_embeddings(
247248
:param embedding_type: type of the embedding: torch.Tensor, numpy.ndarray etc.
248249
:return: the embeddings
249250
"""
250-
251+
emb: AnyTensor
251252
if isinstance(data, DocumentArray):
252-
emb = getattr(data, embedding_field)
253-
emb = embedding_type._docarray_stack(emb)
254-
elif isinstance(data, DocumentArrayStacked):
255-
emb = getattr(data, embedding_field)
256-
elif isinstance(data, BaseDocument):
257-
emb = getattr(data, embedding_field)
253+
emb_list = list(AnyDocumentArray._traverse(data, embedding_field))
254+
emb = embedding_type._docarray_stack(emb_list)
255+
elif isinstance(data, (DocumentArrayStacked, BaseDocument)):
256+
emb = next(AnyDocumentArray._traverse(data, embedding_field))
258257
else: # treat data as tensor
259-
emb = data
258+
emb = cast(AnyTensor, data)
260259

261260
if len(emb.shape) == 1:
262-
# all currently supported frameworks provide `.reshape()`. Onc this is not true
263-
# anymore, we need to add a `.reshape()` method to the computational backend
264-
emb = emb.reshape(1, -1)
261+
emb = emb.get_comp_backend().reshape(array=emb, shape=(1, -1))
265262
return emb
266263

267264

268-
def _da_attr_type(da: AnyDocumentArray, attr: str) -> Type[AnyTensor]:
265+
def _da_attr_type(da: AnyDocumentArray, access_path: str) -> Type[AnyTensor]:
269266
"""Get the type of the attribute according to the Document type
270267
(schema) of the DocumentArray.
271268
272269
:param da: the DocumentArray
273-
:param attr: the attribute name
270+
:param access_path: the "__"-separated access path
274271
:return: the type of the attribute
275272
"""
276-
field_type = da.document_type._get_field_type(attr)
273+
field_type: Optional[Type] = _get_field_type_by_access_path(
274+
da.document_type, access_path
275+
)
276+
if field_type is None:
277+
raise ValueError(f"Access path is not valid: {access_path}")
278+
277279
if is_union_type(field_type):
278280
# determine type based on the fist element
279-
field_type = type(getattr(da[0], attr))
281+
field_type = type(next(AnyDocumentArray._traverse(da[0], access_path)))
280282

281283
if not issubclass(field_type, AbstractTensor):
282284
raise ValueError(
283-
f'attribute {attr} is not a tensor-like type, '
285+
f'attribute {access_path} is not a tensor-like type, '
284286
f'but {field_type.__class__.__name__}'
285287
)
286288

‎tests/units/test_helper.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import pytest
44

5-
from docarray import BaseDocument
5+
from docarray import BaseDocument, DocumentArray
66
from docarray.documents import Image
77
from docarray.helper import (
88
_access_path_dict_to_nested_dict,
@@ -25,8 +25,13 @@ class Middle(BaseDocument):
2525
class Outer(BaseDocument):
2626
img: Optional[Image]
2727
middle: Optional[Middle]
28+
da: DocumentArray[Inner]
2829

29-
doc = Outer(img=Image(), middle=Middle(img=Image(), inner=Inner(img=Image())))
30+
doc = Outer(
31+
img=Image(),
32+
middle=Middle(img=Image(), inner=Inner(img=Image())),
33+
da=DocumentArray[Inner]([Inner(img=Image(url='test.png'))]),
34+
)
3035
return doc
3136

3237

@@ -35,6 +40,7 @@ def test_is_access_path_valid(nested_doc):
3540
assert _is_access_path_valid(nested_doc.__class__, 'middle__img')
3641
assert _is_access_path_valid(nested_doc.__class__, 'middle__inner__img')
3742
assert _is_access_path_valid(nested_doc.__class__, 'middle')
43+
assert _is_access_path_valid(nested_doc.__class__, 'da__img__url')
3844

3945

4046
def test_is_access_path_not_valid(nested_doc):

‎tests/units/util/test_find.py‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -292,6 +292,36 @@ class MyDoc(BaseDocument):
292292
assert (torch.stack(sorted(scores, reverse=True)) == scores).all()
293293

294294

295+
@pytest.mark.parametrize('stack', [False, True])
296+
def test_find_nested(stack):
297+
class InnerDoc(BaseDocument):
298+
title: str
299+
embedding: TorchTensor
300+
301+
class MyDoc(BaseDocument):
302+
inner: InnerDoc
303+
304+
query = MyDoc(inner=InnerDoc(title='query', embedding=torch.rand(2)))
305+
index = DocumentArray[MyDoc](
306+
[
307+
MyDoc(inner=InnerDoc(title=f'doc {i}', embedding=torch.rand(2)))
308+
for i in range(10)
309+
]
310+
)
311+
if stack:
312+
index = index.stack()
313+
314+
top_k, scores = find(
315+
index,
316+
query,
317+
embedding_field='inner__embedding',
318+
limit=7,
319+
)
320+
assert len(top_k) == 7
321+
assert len(scores) == 7
322+
assert (torch.stack(sorted(scores, reverse=True)) == scores).all()
323+
324+
295325
def test_find_nested_union_optional():
296326
class MyDoc(BaseDocument):
297327
embedding: Union[Optional[TorchTensor], Optional[NdArray]]

0 commit comments

Comments
 (0)