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

Commit 251eea7

Browse files
Charlotte GerhaherJohannesMessner
andauthored
feat(v2): load da from csv and save to csv (#1144)
* feat: load from and to csv Signed-off-by: anna-charlotte <[email protected]> * fix: from to csv Signed-off-by: anna-charlotte <[email protected]> * feat: add access path to dict Signed-off-by: anna-charlotte <[email protected]> * fix: from to csv Signed-off-by: anna-charlotte <[email protected]> * fix: clean up Signed-off-by: anna-charlotte <[email protected]> * docs: add docstring and update tmpdir in test Signed-off-by: anna-charlotte <[email protected]> * fix: merge nested dicts Signed-off-by: anna-charlotte <[email protected]> * fix: clean up Signed-off-by: anna-charlotte <[email protected]> * fix: clean up Signed-off-by: anna-charlotte <[email protected]> * test: update test Signed-off-by: anna-charlotte <[email protected]> * fix: apply samis suggestion from code review Signed-off-by: anna-charlotte <[email protected]> * fix: apply suggestions from code review wrt access paths Signed-off-by: anna-charlotte <[email protected]> * fix: apply johannes suggestion Co-authored-by: Johannes Messner <[email protected]> Signed-off-by: Charlotte Gerhaher <[email protected]> * fix: apply johannes suggestion Co-authored-by: Johannes Messner <[email protected]> Signed-off-by: Charlotte Gerhaher <[email protected]> * fix: apply suggestions from code review Signed-off-by: anna-charlotte <[email protected]> * fix: apply suggestions from code review Signed-off-by: anna-charlotte <[email protected]> * fix: typos Signed-off-by: anna-charlotte <[email protected]> * refactor: move helper functions to helper file Signed-off-by: anna-charlotte <[email protected]> * test: fix fixture Signed-off-by: anna-charlotte <[email protected]> --------- Signed-off-by: anna-charlotte <[email protected]> Signed-off-by: Charlotte Gerhaher <[email protected]> Co-authored-by: Johannes Messner <[email protected]>
1 parent 3c79073 commit 251eea7

7 files changed

Lines changed: 410 additions & 7 deletions

File tree

docarray/array/array/io.py

Lines changed: 103 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,45 @@
11
import base64
2+
import csv
23
import io
34
import json
45
import os
56
import pathlib
67
import pickle
78
from abc import abstractmethod
89
from contextlib import nullcontext
10+
from itertools import compress
911
from typing import (
1012
TYPE_CHECKING,
13+
Any,
1114
BinaryIO,
1215
ContextManager,
16+
Dict,
1317
Generator,
1418
Iterable,
1519
Optional,
20+
Sequence,
1621
Tuple,
1722
Type,
1823
TypeVar,
1924
Union,
2025
)
2126

22-
from docarray.base_document import BaseDocument
27+
from docarray.base_document import AnyDocument, BaseDocument
28+
from docarray.helper import (
29+
_access_path_to_dict,
30+
_dict_to_access_paths,
31+
_update_nested_dicts,
32+
is_access_path_valid,
33+
)
2334
from docarray.utils.compress import _decompress_bytes, _get_compress_ctx
2435

2536
if TYPE_CHECKING:
2637

38+
from docarray import DocumentArray
2739
from docarray.proto import DocumentArrayProto
2840

2941
T = TypeVar('T', bound='IOMixinArray')
3042

31-
3243
ARRAY_PROTOCOLS = {'protobuf-array', 'pickle-array'}
3344
SINGLE_PROTOCOLS = {'pickle', 'protobuf'}
3445
ALLOWED_PROTOCOLS = ARRAY_PROTOCOLS.union(SINGLE_PROTOCOLS)
@@ -291,6 +302,96 @@ def to_json(self) -> str:
291302
"""
292303
return json.dumps([doc.json() for doc in self])
293304

305+
@classmethod
306+
def from_csv(
307+
cls,
308+
file_path: str,
309+
encoding: str = 'utf-8',
310+
dialect: Union[str, csv.Dialect] = 'excel',
311+
) -> 'DocumentArray':
312+
"""
313+
Load a DocumentArray from a csv file following the schema defined in the
314+
:attr:`~docarray.DocumentArray.document_type` attribute.
315+
Every row of the csv file will be mapped to one document in the array.
316+
The column names (defined in the first row) have to match the field names
317+
of the Document type.
318+
For nested fields use "__"-separated access paths, such as 'image__url'.
319+
320+
List-like fields (including field of type DocumentArray) are not supported.
321+
322+
:param file_path: path to csv file to load DocumentArray from.
323+
:param encoding: encoding used to read the csv file. Defaults to 'utf-8'.
324+
:param dialect: defines separator and how to handle whitespaces etc.
325+
Can be a csv.Dialect instance or one string of:
326+
'excel' (for comma seperated values),
327+
'excel-tab' (for tab separated values),
328+
'unix' (for csv file generated on UNIX systems).
329+
:return: DocumentArray
330+
"""
331+
from docarray import DocumentArray
332+
333+
doc_type = cls.document_type
334+
if doc_type == AnyDocument:
335+
raise TypeError(
336+
'There is no document schema defined. '
337+
'To load from csv, please specify the DocumentArray\'s Document type using `DocumentArray[MyDoc]`.'
338+
)
339+
340+
da = DocumentArray.__class_getitem__(doc_type)()
341+
with open(file_path, 'r', encoding=encoding) as fp:
342+
rows = csv.DictReader(fp, dialect=dialect)
343+
field_names: Optional[Sequence[Any]] = rows.fieldnames
344+
345+
if field_names is None:
346+
raise TypeError("No field names are given.")
347+
348+
valid = [is_access_path_valid(doc_type, field) for field in field_names]
349+
if not all(valid):
350+
raise ValueError(
351+
f'Fields provided in the csv file do not match the schema of the DocumentArray\'s '
352+
f'document type ({doc_type.__name__}): {list(compress(field_names, [not v for v in valid]))}'
353+
)
354+
355+
for access_path2val in rows:
356+
doc_dict: Dict[Any, Any] = {}
357+
for access_path, value in access_path2val.items():
358+
field2val = _access_path_to_dict(
359+
access_path=access_path,
360+
value=value if value not in ['', 'None'] else None,
361+
)
362+
_update_nested_dicts(to_update=doc_dict, update_with=field2val)
363+
364+
da.append(doc_type.parse_obj(doc_dict))
365+
366+
return da
367+
368+
def to_csv(
369+
self, file_path: str, dialect: Union[str, csv.Dialect] = 'excel'
370+
) -> None:
371+
"""
372+
Save a DocumentArray to a csv file.
373+
The field names will be stored in the first row. Each row corresponds to the
374+
information of one Document.
375+
Columns for nested fields will be named after the "__"-seperated access paths,
376+
such as `"image__url"` for `image.url`.
377+
378+
:param file_path: path to a csv file.
379+
:param dialect: defines separator and how to handle whitespaces etc.
380+
Can be a csv.Dialect instance or one string of:
381+
'excel' (for comma seperated values),
382+
'excel-tab' (for tab separated values),
383+
'unix' (for csv file generated on UNIX systems).
384+
"""
385+
fields = self.document_type._get_access_paths()
386+
387+
with open(file_path, 'w') as csv_file:
388+
writer = csv.DictWriter(csv_file, fieldnames=fields, dialect=dialect)
389+
writer.writeheader()
390+
391+
for doc in self:
392+
doc_dict = _dict_to_access_paths(doc.dict())
393+
writer.writerow(doc_dict)
394+
294395
# Methods to load from/to files in different formats
295396
@property
296397
def _stream_header(self) -> bytes:

docarray/base_document/mixins/io.py

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,12 +7,15 @@
77
Callable,
88
Dict,
99
Iterable,
10+
List,
1011
Optional,
1112
Tuple,
1213
Type,
1314
TypeVar,
1415
)
1516

17+
from typing_inspect import is_union_type
18+
1619
from docarray.base_document.base_node import BaseNode
1720
from docarray.typing.proto_register import _PROTO_TYPE_NAME_TO_CLASS
1821
from docarray.utils.compress import _compress_bytes, _decompress_bytes
@@ -291,3 +294,23 @@ def _to_node_protobuf(self) -> 'NodeProto':
291294
:return: the nested item protobuf message
292295
"""
293296
return NodeProto(document=self.to_protobuf())
297+
298+
@classmethod
299+
def _get_access_paths(cls) -> List[str]:
300+
"""
301+
Get "__"-separated access paths of all fields, including nested ones.
302+
303+
:return: list of all access paths
304+
"""
305+
from docarray import BaseDocument
306+
307+
paths = []
308+
for field in cls.__fields__.keys():
309+
field_type = cls._get_field_type(field)
310+
if not is_union_type(field_type) and issubclass(field_type, BaseDocument):
311+
sub_paths = field_type._get_access_paths()
312+
for path in sub_paths:
313+
paths.append(f'{field}__{path}')
314+
else:
315+
paths.append(field)
316+
return paths

docarray/helper.py

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
from typing import TYPE_CHECKING, Any, Dict, Type
2+
3+
if TYPE_CHECKING:
4+
from docarray import BaseDocument
5+
6+
7+
def is_access_path_valid(doc: Type['BaseDocument'], access_path: str) -> bool:
8+
"""
9+
Check if a given access path ("__"-separated) is a valid path for a given Document class.
10+
"""
11+
from docarray import BaseDocument
12+
13+
field, _, remaining = access_path.partition('__')
14+
if len(remaining) == 0:
15+
return access_path in doc.__fields__.keys()
16+
else:
17+
valid_field = field in doc.__fields__.keys()
18+
if not valid_field:
19+
return False
20+
else:
21+
d = doc._get_field_type(field)
22+
if not issubclass(d, BaseDocument):
23+
return False
24+
else:
25+
return is_access_path_valid(d, remaining)
26+
27+
28+
def _access_path_to_dict(access_path: str, value) -> Dict[str, Any]:
29+
"""
30+
Convert an access path ("__"-separated) and its value to a (potentially) nested dict.
31+
32+
EXAMPLE USAGE
33+
.. code-block:: python
34+
assert access_path_to_dict('image__url', 'img.png') == {'image': {'url': 'img.png'}}
35+
"""
36+
fields = access_path.split('__')
37+
for field in reversed(fields):
38+
result = {field: value}
39+
value = result
40+
return result
41+
42+
43+
def _dict_to_access_paths(d: dict) -> Dict[str, Any]:
44+
"""
45+
Convert a (nested) dict to a Dict[access_path, value].
46+
Access paths are defined as a path of field(s) separated by "__".
47+
48+
EXAMPLE USAGE
49+
.. code-block:: python
50+
assert dict_to_access_paths({'image': {'url': 'img.png'}}) == {'image__url', 'img.png'}
51+
"""
52+
result = {}
53+
for k, v in d.items():
54+
if isinstance(v, dict):
55+
v = _dict_to_access_paths(v)
56+
for nested_k, nested_v in v.items():
57+
new_key = '__'.join([k, nested_k])
58+
result[new_key] = nested_v
59+
else:
60+
result[k] = v
61+
return result
62+
63+
64+
def _update_nested_dicts(
65+
to_update: Dict[Any, Any], update_with: Dict[Any, Any]
66+
) -> None:
67+
"""
68+
Update a dict with another one, while considering shared nested keys.
69+
70+
EXAMPLE USAGE:
71+
72+
.. code-block:: python
73+
74+
d1 = {'image': {'tensor': None}, 'title': 'hello'}
75+
d2 = {'image': {'url': 'some.png'}}
76+
77+
update_nested_dicts(d1, d2)
78+
assert d1 == {'image': {'tensor': None, 'url': 'some.png'}, 'title': 'hello'}
79+
80+
:param to_update: dict that should be updated
81+
:param update_with: dict to update with
82+
:return: merged dict
83+
"""
84+
for k, v in update_with.items():
85+
if k not in to_update.keys():
86+
to_update[k] = v
87+
else:
88+
_update_nested_dicts(to_update[k], update_with[k])

tests/toydata/docs_nested.csv

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
count,text,image,image2__url
2+
000,hello 0,image_0.png,image_10.png
3+
111,hello 1,image_1.png,None
4+
222,hello 2,image_2.png,
Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
import os
2+
from typing import Optional
3+
4+
import pytest
5+
6+
from docarray import BaseDocument, DocumentArray
7+
from docarray.documents import Image
8+
from tests import TOYDATA_DIR
9+
10+
11+
@pytest.fixture()
12+
def nested_doc_cls():
13+
class MyDoc(BaseDocument):
14+
count: Optional[int]
15+
text: str
16+
17+
class MyDocNested(MyDoc):
18+
image: Image
19+
image2: Image
20+
21+
return MyDocNested
22+
23+
24+
def test_to_from_csv(tmpdir, nested_doc_cls):
25+
da = DocumentArray[nested_doc_cls](
26+
[
27+
nested_doc_cls(
28+
count=0,
29+
text='hello',
30+
image=Image(url='aux.png'),
31+
image2=Image(url='aux.png'),
32+
),
33+
nested_doc_cls(text='hello world', image=Image(), image2=Image()),
34+
]
35+
)
36+
tmp_file = str(tmpdir / 'tmp.csv')
37+
da.to_csv(tmp_file)
38+
assert os.path.isfile(tmp_file)
39+
40+
da_from = DocumentArray[nested_doc_cls].from_csv(tmp_file)
41+
for doc1, doc2 in zip(da, da_from):
42+
assert doc1 == doc2
43+
44+
45+
def test_from_csv_nested(nested_doc_cls):
46+
da = DocumentArray[nested_doc_cls].from_csv(
47+
file_path=str(TOYDATA_DIR / 'docs_nested.csv')
48+
)
49+
assert len(da) == 3
50+
51+
for i, doc in enumerate(da):
52+
assert doc.count.__class__ == int
53+
assert doc.count == int(f'{i}{i}{i}')
54+
55+
assert doc.text.__class__ == str
56+
assert doc.text == f'hello {i}'
57+
58+
assert doc.image.__class__ == Image
59+
assert doc.image.tensor is None
60+
assert doc.image.embedding is None
61+
assert doc.image.bytes is None
62+
63+
assert doc.image2.__class__ == Image
64+
assert doc.image2.tensor is None
65+
assert doc.image2.embedding is None
66+
assert doc.image2.bytes is None
67+
68+
assert da[0].image2.url == 'image_10.png'
69+
assert da[1].image2.url is None
70+
assert da[2].image2.url is None
71+
72+
73+
@pytest.fixture()
74+
def nested_doc():
75+
class Inner(BaseDocument):
76+
img: Optional[Image]
77+
78+
class Middle(BaseDocument):
79+
img: Optional[Image]
80+
inner: Optional[Inner]
81+
82+
class Outer(BaseDocument):
83+
img: Optional[Image]
84+
middle: Optional[Middle]
85+
86+
doc = Outer(img=Image(), middle=Middle(img=Image(), inner=Inner(img=Image())))
87+
return doc
88+
89+
90+
def test_from_csv_without_schema_raise_exception():
91+
with pytest.raises(TypeError, match='no document schema defined'):
92+
DocumentArray.from_csv(file_path=str(TOYDATA_DIR / 'docs_nested.csv'))
93+
94+
95+
def test_from_csv_with_wrong_schema_raise_exception(nested_doc):
96+
with pytest.raises(
97+
ValueError, match='Fields provided in the csv file do not match the schema'
98+
):
99+
DocumentArray[nested_doc.__class__].from_csv(
100+
file_path=str(TOYDATA_DIR / 'docs.csv')
101+
)

tests/units/array/test_array_from_to_json.py

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,6 @@
1-
import pytest
2-
3-
from docarray import BaseDocument
4-
from docarray.typing import NdArray
1+
from docarray import BaseDocument, DocumentArray
52
from docarray.documents import Image
6-
from docarray import DocumentArray
3+
from docarray.typing import NdArray
74

85

96
class MyDoc(BaseDocument):

0 commit comments

Comments
 (0)