|
1 | 1 | import base64 |
| 2 | +import csv |
2 | 3 | import io |
3 | 4 | import json |
4 | 5 | import os |
5 | 6 | import pathlib |
6 | 7 | import pickle |
7 | 8 | from abc import abstractmethod |
8 | 9 | from contextlib import nullcontext |
| 10 | +from itertools import compress |
9 | 11 | from typing import ( |
10 | 12 | TYPE_CHECKING, |
| 13 | + Any, |
11 | 14 | BinaryIO, |
12 | 15 | ContextManager, |
| 16 | + Dict, |
13 | 17 | Generator, |
14 | 18 | Iterable, |
15 | 19 | Optional, |
| 20 | + Sequence, |
16 | 21 | Tuple, |
17 | 22 | Type, |
18 | 23 | TypeVar, |
19 | 24 | Union, |
20 | 25 | ) |
21 | 26 |
|
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 | +) |
23 | 34 | from docarray.utils.compress import _decompress_bytes, _get_compress_ctx |
24 | 35 |
|
25 | 36 | if TYPE_CHECKING: |
26 | 37 |
|
| 38 | + from docarray import DocumentArray |
27 | 39 | from docarray.proto import DocumentArrayProto |
28 | 40 |
|
29 | 41 | T = TypeVar('T', bound='IOMixinArray') |
30 | 42 |
|
31 | | - |
32 | 43 | ARRAY_PROTOCOLS = {'protobuf-array', 'pickle-array'} |
33 | 44 | SINGLE_PROTOCOLS = {'pickle', 'protobuf'} |
34 | 45 | ALLOWED_PROTOCOLS = ARRAY_PROTOCOLS.union(SINGLE_PROTOCOLS) |
@@ -291,6 +302,96 @@ def to_json(self) -> str: |
291 | 302 | """ |
292 | 303 | return json.dumps([doc.json() for doc in self]) |
293 | 304 |
|
| 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 | + |
294 | 395 | # Methods to load from/to files in different formats |
295 | 396 | @property |
296 | 397 | def _stream_header(self) -> bytes: |
|
0 commit comments