diff --git a/docarray/array/doc_list/doc_list.py b/docarray/array/doc_list/doc_list.py index 951256ef2ce..6dbfca3fa76 100644 --- a/docarray/array/doc_list/doc_list.py +++ b/docarray/array/doc_list/doc_list.py @@ -306,6 +306,12 @@ def from_protobuf(cls: Type[T], pb_msg: 'DocListProto') -> T: """ return super().from_protobuf(pb_msg) + @classmethod + def _get_proto_class(cls: Type[T]): + from docarray.proto import DocListProto + + return DocListProto + @overload def __getitem__(self, item: SupportsIndex) -> T_doc: ... diff --git a/docarray/array/doc_list/io.py b/docarray/array/doc_list/io.py index 9667c673c09..7dc5d5979d6 100644 --- a/docarray/array/doc_list/io.py +++ b/docarray/array/doc_list/io.py @@ -40,7 +40,6 @@ if TYPE_CHECKING: import pandas as pd - from docarray import DocList from docarray.proto import DocListProto T = TypeVar('T', bound='IOMixinArray') @@ -332,11 +331,11 @@ def to_json(self) -> bytes: @classmethod def from_csv( - cls, + cls: Type['T'], file_path: str, encoding: str = 'utf-8', dialect: Union[str, csv.Dialect] = 'excel', - ) -> 'DocList': + ) -> 'T': """ Load a DocList from a csv file following the schema defined in the [`.doc_type`][docarray.DocList] attribute. @@ -358,10 +357,10 @@ def from_csv( :return: `DocList` object """ - if cls.doc_type == AnyDoc: + if cls.doc_type == AnyDoc or cls.doc_type == BaseDoc: raise TypeError( 'There is no document schema defined. ' - 'Please specify the DocList\'s Document type using `DocList[MyDoc]`.' + f'Please specify the {cls}\'s Document type using `{cls}[MyDoc]`.' ) if file_path.startswith('http'): @@ -376,14 +375,15 @@ def from_csv( @classmethod def _from_csv_file( - cls, file: Union[StringIO, TextIOWrapper], dialect: Union[str, csv.Dialect] - ) -> 'DocList': - from docarray import DocList + cls: Type['T'], + file: Union[StringIO, TextIOWrapper], + dialect: Union[str, csv.Dialect], + ) -> 'T': rows = csv.DictReader(file, dialect=dialect) doc_type = cls.doc_type - docs = DocList.__class_getitem__(doc_type)() + docs = [] field_names: List[str] = ( [] if rows.fieldnames is None else [str(f) for f in rows.fieldnames] @@ -405,7 +405,7 @@ def _from_csv_file( doc_dict: Dict[Any, Any] = _access_path_dict_to_nested_dict(access_path2val) docs.append(doc_type.parse_obj(doc_dict)) - return docs + return cls(docs) def to_csv( self, file_path: str, dialect: Union[str, csv.Dialect] = 'excel' @@ -426,11 +426,11 @@ def to_csv( `'unix'` (for csv file generated on UNIX systems). """ - if self.doc_type == AnyDoc: + if self.doc_type == AnyDoc or self.doc_type == BaseDoc: raise TypeError( - 'DocList must be homogeneous to be converted to a csv.' + f'{type(self)} must be homogeneous to be converted to a csv.' 'There is no document schema defined. ' - 'Please specify the DocList\'s Document type using `DocList[MyDoc]`.' + f'Please specify the {type(self)}\'s Document type using `{type(self)}[MyDoc]`.' ) fields = self.doc_type._get_access_paths() @@ -443,7 +443,7 @@ def to_csv( writer.writerow(doc_dict) @classmethod - def from_dataframe(cls, df: 'pd.DataFrame') -> 'DocList': + def from_dataframe(cls: Type['T'], df: 'pd.DataFrame') -> 'T': """ Load a `DocList` from a `pandas.DataFrame` following the schema defined in the [`.doc_type`][docarray.DocList] attribute. @@ -486,10 +486,10 @@ class Person(BaseDoc): """ from docarray import DocList - if cls.doc_type == AnyDoc: + if cls.doc_type == AnyDoc or cls.doc_type == BaseDoc: raise TypeError( 'There is no document schema defined. ' - 'Please specify the DocList\'s Document type using `DocList[MyDoc]`.' + f'Please specify the {cls}\'s Document type using `{cls}[MyDoc]`.' ) doc_type = cls.doc_type @@ -515,6 +515,8 @@ class Person(BaseDoc): doc_dict = _access_path_dict_to_nested_dict(access_path2val) docs.append(doc_type.parse_obj(doc_dict)) + if not isinstance(docs, cls): + return cls(docs) return docs def to_dataframe(self) -> 'pd.DataFrame': @@ -563,6 +565,11 @@ def _stream_header(self) -> bytes: num_docs_as_bytes = len(self).to_bytes(8, 'big', signed=False) return version_byte + num_docs_as_bytes + @classmethod + @abstractmethod + def _get_proto_class(cls: Type[T]): + ... + @classmethod def _load_binary_all( cls: Type[T], @@ -593,12 +600,10 @@ def _load_binary_all( compress = None if protocol is not None and protocol == 'protobuf-array': - from docarray.proto import DocListProto - - dap = DocListProto() - dap.ParseFromString(d) + proto = cls._get_proto_class()() + proto.ParseFromString(d) - return cls.from_protobuf(dap) + return cls.from_protobuf(proto) elif protocol is not None and protocol == 'pickle-array': return pickle.loads(d) diff --git a/docarray/array/doc_vec/column_storage.py b/docarray/array/doc_vec/column_storage.py index 539e9fd42af..46895841da2 100644 --- a/docarray/array/doc_vec/column_storage.py +++ b/docarray/array/doc_vec/column_storage.py @@ -6,6 +6,7 @@ ItemsView, Iterable, MutableMapping, + NamedTuple, Optional, Type, TypeVar, @@ -26,6 +27,13 @@ T = TypeVar('T', bound='ColumnStorage') +class ColumnsJsonCompatible(NamedTuple): + tensor_columns: Dict[str, Any] + doc_columns: Dict[str, Any] + docs_vec_columns: Dict[str, Any] + any_columns: Dict[str, Any] + + class ColumnStorage: """ ColumnStorage is a container to store the columns of the @@ -91,6 +99,25 @@ def __getitem__(self: T, item: IndexIterType) -> T: self.tensor_type, ) + def columns_json_compatible(self) -> ColumnsJsonCompatible: + tens_cols = { + key: value._docarray_to_json_compatible() if value is not None else value + for key, value in self.tensor_columns.items() + } + doc_cols = { + key: value._docarray_to_json_compatible() if value is not None else value + for key, value in self.doc_columns.items() + } + doc_vec_cols = { + key: [vec._docarray_to_json_compatible() for vec in value] + if value is not None + else value + for key, value in self.docs_vec_columns.items() + } + return ColumnsJsonCompatible( + tens_cols, doc_cols, doc_vec_cols, self.any_columns + ) + def __eq__(self, other: Any) -> bool: if not isinstance(other, ColumnStorage): return False @@ -146,6 +173,11 @@ def __getitem__(self, name: str) -> Any: return None return col[self.index] + def __reduce__(self): + # implementing __reduce__ to solve a pickle issue when subclassing dict + # see here: https://stackoverflow.com/questions/21144845/how-can-i-unpickle-a-subclass-of-dict-that-validates-with-setitem-in-pytho + return (ColumnStorageView, (self.index, self.storage)) + def __setitem__(self, name, value) -> None: if self.storage.columns[name] is None: raise ValueError( diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index 3025c52a09b..fab2c7313bc 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -18,11 +18,13 @@ ) import numpy as np +import orjson from pydantic import BaseConfig, parse_obj_as from typing_inspect import typingGenericAlias from docarray.array.any_array import AnyDocArray from docarray.array.doc_list.doc_list import DocList +from docarray.array.doc_list.io import IOMixinArray from docarray.array.doc_vec.column_storage import ColumnStorage, ColumnStorageView from docarray.array.list_advance_indexing import ListAdvancedIndexing from docarray.base_doc import AnyDoc, BaseDoc @@ -33,6 +35,8 @@ from docarray.utils._internal.misc import is_tf_available, is_torch_available if TYPE_CHECKING: + import csv + from pydantic.fields import ModelField from docarray.proto import ( @@ -58,6 +62,8 @@ T_doc = TypeVar('T_doc', bound=BaseDoc) T = TypeVar('T', bound='DocVec') +T_io_mixin = TypeVar('T_io_mixin', bound='IOMixinArray') + IndexIterType = Union[slice, Iterable[int], Iterable[bool], None] NONE_NDARRAY_PROTO_SHAPE = (0,) @@ -111,7 +117,7 @@ def _is_none_list_of_docvec_proto(proto: 'ListOfDocVecProto') -> bool: return isinstance(proto, ListOfDocVecProto) and len(proto.data) == 0 -class DocVec(AnyDocArray[T_doc]): +class DocVec(IOMixinArray, AnyDocArray[T_doc]): """ DocVec is a container of Documents appropriates to perform computation that require batches of data (ex: matrix multiplication, distance @@ -156,7 +162,7 @@ class DocVec(AnyDocArray[T_doc]): AnyTensor or Union of NdArray and TorchTensor """ - doc_type: Type[T_doc] + doc_type: Type[T_doc] = BaseDoc # type: ignore def __init__( self: T, @@ -164,7 +170,11 @@ def __init__( tensor_type: Type['AbstractTensor'] = NdArray, ): - if not hasattr(self, 'doc_type') or self.doc_type == AnyDoc: + if ( + not hasattr(self, 'doc_type') + or self.doc_type == AnyDoc + or self.doc_type == BaseDoc + ): raise TypeError( f'{self.__class__.__name__} does not precise a doc_type. You probably should do' f'docs = DocVec[MyDoc](docs) instead of DocVec(docs)' @@ -600,6 +610,92 @@ def __eq__(self, other: Any) -> bool: # IO related # #################### + @classmethod + def _get_proto_class(cls: Type[T]): + from docarray.proto import DocVecProto + + return DocVecProto + + def _docarray_to_json_compatible(self) -> Dict[str, Dict[str, Any]]: + tup = self._storage.columns_json_compatible() + return tup._asdict() + + @classmethod + def from_json( + cls: Type[T], + file: Union[str, bytes, bytearray], + tensor_type: Type[AbstractTensor] = NdArray, + ) -> T: + """Deserialize JSON strings or bytes into a `DocList`. + + :param file: JSON object from where to deserialize a `DocList` + :param tensor_type: the tensor type to use for the tensor columns. + Could be NdArray, TorchTensor, or TensorFlowTensor. Defaults to NdArray. + All tensors of the output DocVec will be of this type. + :return: the deserialized `DocList` + """ + json_columns = orjson.loads(file) + return cls._from_json_col_dict(json_columns, tensor_type=tensor_type) + + @classmethod + def _from_json_col_dict( + cls: Type[T], + json_columns: Dict[str, Any], + tensor_type: Type[AbstractTensor] = NdArray, + ) -> T: + + tensor_cols = json_columns['tensor_columns'] + doc_cols = json_columns['doc_columns'] + docs_vec_cols = json_columns['docs_vec_columns'] + any_cols = json_columns['any_columns'] + + for key, col in tensor_cols.items(): + if col is not None: + tensor_cols[key] = parse_obj_as(tensor_type, col) + else: + tensor_cols[key] = None + + for key, col in doc_cols.items(): + if col is not None: + col_doc_type = cls.doc_type._get_field_type(key) + doc_cols[key] = DocVec.__class_getitem__( + col_doc_type + )._from_json_col_dict(col, tensor_type=tensor_type) + else: + doc_cols[key] = None + + for key, col in docs_vec_cols.items(): + if col is not None: + col_doc_type = cls.doc_type._get_field_type(key).doc_type + col_ = ListAdvancedIndexing( + DocVec.__class_getitem__(col_doc_type)._from_json_col_dict( + vec, tensor_type=tensor_type + ) + for vec in col + ) + docs_vec_cols[key] = col_ + else: + docs_vec_cols[key] = None + + for key, col in any_cols.items(): + if col is not None: + col_type = cls.doc_type._get_field_type(key) + col_type = ( + col_type + if cls.doc_type.__fields__[key].required + else Optional[col_type] + ) + col_ = ListAdvancedIndexing(parse_obj_as(col_type, val) for val in col) + any_cols[key] = col_ + else: + any_cols[key] = None + + return cls.from_columns_storage( + ColumnStorage( + tensor_cols, doc_cols, docs_vec_cols, any_cols, tensor_type=tensor_type + ) + ) + @classmethod def from_protobuf( cls: Type[T], pb_msg: 'DocVecProto', tensor_type: Type[AbstractTensor] = NdArray @@ -792,3 +888,35 @@ def traverse_flat( return flattened[0] else: return flattened + + def to_csv( + self, file_path: str, dialect: Union[str, 'csv.Dialect'] = 'excel' + ) -> None: + """ + DocVec does not support `.to_csv()`. This is because CSV is a row-based format + while DocVec has a column-based data layout. + To overcome this, do: `doc_vec.to_doc_list().to_csv(...)`. + """ + raise NotImplementedError( + f'{type(self)} does not support `.to_csv()`. This is because CSV is a row-based format' + f'while {type(self)} has a column-based data layout. ' + f'To overcome this, do: `doc_vec.to_doc_list().to_csv(...)`.' + ) + + @classmethod + def from_csv( + cls: Type['T'], + file_path: str, + encoding: str = 'utf-8', + dialect: Union[str, 'csv.Dialect'] = 'excel', + ) -> 'T': + """ + DocVec does not support `.from_csv()`. This is because CSV is a row-based format + while DocVec has a column-based data layout. + To overcome this, do: `DocList[MyDoc].from_csv(...).to_doc_vec()`. + """ + raise NotImplementedError( + f'{cls} does not support `.from_csv()`. This is because CSV is a row-based format while' + f'{cls} has a column-based data layout. ' + f'To overcome this, do: `DocList[MyDoc].from_csv(...).to_doc_vec()`.' + ) diff --git a/docarray/base_doc/doc.py b/docarray/base_doc/doc.py index 6747b269cfe..0dfbabd57c0 100644 --- a/docarray/base_doc/doc.py +++ b/docarray/base_doc/doc.py @@ -228,7 +228,7 @@ def json( `encoder` is an optional function to supply as `default` to json.dumps(), other arguments as per `json.dumps()`. """ - exclude, original_exclude, doclist_exclude_fields = self._exclude_doclist( + exclude, original_exclude, doclist_exclude_fields = self._exclude_docarray( exclude=exclude ) @@ -315,7 +315,7 @@ def dict( """ - exclude, original_exclude, doclist_exclude_fields = self._exclude_doclist( + exclude, original_exclude, docarray_exclude_fields = self._exclude_docarray( exclude=exclude ) @@ -329,7 +329,7 @@ def dict( exclude_none=exclude_none, ) - for field in doclist_exclude_fields: + for field in docarray_exclude_fields: # we need to do this because pydantic will not recognize DocList correctly original_exclude = original_exclude or {} if field not in original_exclude: @@ -338,30 +338,32 @@ def dict( return data - def _exclude_doclist( + def _exclude_docarray( self, exclude: ExcludeType ) -> Tuple[ExcludeType, ExcludeType, List[str]]: - doclist_exclude_fields = [] + docarray_exclude_fields = [] for field in self.__fields__.keys(): - from docarray import DocList + from docarray import DocList, DocVec type_ = self._get_field_type(field) - if isinstance(type_, type) and issubclass(type_, DocList): - doclist_exclude_fields.append(field) + if isinstance(type_, type) and ( + issubclass(type_, DocList) or issubclass(type_, DocVec) + ): + docarray_exclude_fields.append(field) original_exclude = exclude if exclude is None: - exclude = set(doclist_exclude_fields) + exclude = set(docarray_exclude_fields) elif isinstance(exclude, AbstractSet): - exclude = set([*exclude, *doclist_exclude_fields]) + exclude = set([*exclude, *docarray_exclude_fields]) elif isinstance(exclude, Mapping): exclude = dict(**exclude) - exclude.update({field: ... for field in doclist_exclude_fields}) + exclude.update({field: ... for field in docarray_exclude_fields}) return ( exclude, original_exclude, - doclist_exclude_fields, + docarray_exclude_fields, ) to_json = json diff --git a/docarray/base_doc/mixins/io.py b/docarray/base_doc/mixins/io.py index bc49ea8cdb5..8eb9751af43 100644 --- a/docarray/base_doc/mixins/io.py +++ b/docarray/base_doc/mixins/io.py @@ -8,6 +8,7 @@ Dict, Iterable, List, + Literal, Optional, Tuple, Type, @@ -212,7 +213,7 @@ def to_base64( def from_base64( cls: Type[T], data: str, - protocol: str = 'pickle', + protocol: Literal['pickle', 'protobuf'] = 'pickle', compress: Optional[str] = None, ) -> T: """Build Document object from binary bytes diff --git a/docs/user_guide/sending/serialization.md b/docs/user_guide/sending/serialization.md index 6cc2b64f35e..1986f12532b 100644 --- a/docs/user_guide/sending/serialization.md +++ b/docs/user_guide/sending/serialization.md @@ -54,7 +54,8 @@ assert doc == new_doc # True ## DocList -When sending or storing [`DocList`][docarray.array.doc_list.doc_list.DocList], you need to use serialization. [`DocList`][docarray.array.doc_list.doc_list.DocList] supports multiple ways to serialize the data. +When sending or storing [`DocList`][docarray.array.doc_list.doc_list.DocList], you need to use serialization. +[`DocList`][docarray.array.doc_list.doc_list.DocList] supports multiple ways to serialize the data. ### JSON @@ -193,7 +194,7 @@ dl_from_bytes = DocList[SimpleDoc].from_bytes( ) ``` -## CSV +### CSV - [`to_csv()`][docarray.array.doc_list.io.IOMixinArray.to_csv] serializes a [`DocList`][docarray.array.doc_list.doc_list.DocList] to a CSV file. - [`from_csv()`][docarray.array.doc_list.io.IOMixinArray.from_csv] deserializes a [`DocList`][docarray.array.doc_list.doc_list.DocList] from a CSV file. @@ -215,7 +216,7 @@ dl_from_csv = DocList[SimpleDoc].from_csv('simple-dl.csv') print(dl_from_csv) ``` -## Pandas.Dataframe +### Pandas.Dataframe - [`from_dataframe()`][docarray.array.doc_list.io.IOMixinArray.from_dataframe] loads a [`DocList`][docarray.array.doc_list.doc_list.DocList] from a [Pandas Dataframe](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html). - [`to_dataframe()`][docarray.array.doc_list.io.IOMixinArray.to_dataframe] saves a [`DocList`][docarray.array.doc_list.doc_list.DocList] to a [Pandas Dataframe](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html). @@ -237,10 +238,40 @@ print(dl_from_dataframe) ## DocVec -When sending or storing [`DocVec`][docarray.array.doc_list.doc_list.DocVec], you need to use protobuf serialization. +For sending or storing [`DocVec`][docarray.array.doc_vec.doc_vec.DocVec] it offers a very similar interface to that of +[`DocList`][docarray.array.doc_list.doc_list.DocList]. -!!! note - We plan to add more serialization formats in the future, notably JSON. +### JSON + +- [`to_json()`][docarray.array.doc_list.io.IOMixinArray.to_json] serializes a [`DocVec`][docarray.array.doc_vec.doc_vec.DocVec] to JSON. It returns the binary representation of the JSON object. +- [`from_json()`][docarray.array.doc_list.io.IOMixinArray.from_json] deserializes a [`DocList`][docarray.array.doc_vec.doc_vec.DocVec] from JSON. It can load from either a `str` or `binary` representation of the JSON object. + +In contrast to [DocList's JSON format](#json-1), `DocVec.to_json()` outputs a column oriented JSON file: + +```python +from docarray import BaseDoc, DocVec + + +class SimpleDoc(BaseDoc): + text: str + + +dv = DocVec[SimpleDoc]([SimpleDoc(text=f'doc {i}') for i in range(2)]) + +with open('simple-dv.json', 'wb') as f: + json_dv = dv.to_json() + print(json_dv) + f.write(json_dv) + +with open('simple-dv.json', 'r') as f: + dv_load_from_json = DocVec[SimpleDoc].from_json(f.read()) + print(dv_load_from_json) +``` + +```output +b'{"tensor_columns":{},"doc_columns":{},"docs_vec_columns":{},"any_columns":{"id":["005a208a0a9a368c16bf77913b710433","31d65f02cb94fc9756c57b0dbaac3a2c"],"text":["doc 0","doc 1"]}}' + +``` ### Protobuf @@ -306,5 +337,135 @@ assert isinstance(dv_from_proto_numpy.tensor, NdArray) !!! note Serialization to protobuf is not supported for union types involving `BaseDoc` types. +### Base64 + +When transferring data over the network, use `Base64` format to serialize the [DocVec][docarray.array.doc_list.doc_list.DocVec]. +Serializing a [DocVec][docarray.array.doc_list.doc_list.DocVec] in Base64 supports both the `pickle` and `protobuf` protocols. +You can also choose different compression methods. + +- [`to_base64()`][docarray.array.doc_list.io.IOMixinArray.to_base64] serializes a [DocVec][docarray.array.doc_list.doc_list.DocVec] to Base64 +- [`from_base64()`][docarray.array.doc_list.io.IOMixinArray.from_base64] deserializes a [DocVec][docarray.array.doc_list.doc_list.DocVec] from Base64: + +You can multiple compression methods: `lz4`, `bz2`, `lzma`, `zlib`, and `gzip`. + +```python +from docarray import BaseDoc, DocVec + + +class SimpleDoc(BaseDoc): + text: str + + +dv = DocVec[SimpleDoc]([SimpleDoc(text=f'doc {i}') for i in range(2)]) + +base64_repr_dv = dv.to_base64(compress=None, protocol='pickle') + +dl_from_base64 = DocVec[SimpleDoc].from_base64( + base64_repr_dv, compress=None, protocol='pickle' +) +``` + +### Save binary + +These methods **serialize and save** your data: + +- [`save_binary()`][docarray.array.doc_list.io.IOMixinArray.save_binary] saves a [DocVec][docarray.array.doc_list.doc_list.DocVec] to a binary file. +- [`load_binary()`][docarray.array.doc_list.io.IOMixinArray.load_binary] loads a [DocVec][docarray.array.doc_list.doc_list.DocVec] from a binary file. + +You can choose between multiple compression methods: `lz4`, `bz2`, `lzma`, `zlib`, and `gzip`. + +```python +from docarray import BaseDoc, DocVec + + +class SimpleDoc(BaseDoc): + text: str + + +dv = DocVec[SimpleDoc]([SimpleDoc(text=f'doc {i}') for i in range(2)]) + +dv.save_binary('simple-dl.pickle', compress=None, protocol='pickle') + +dv_from_binary = DocVec[SimpleDoc].load_binary( + 'simple-dv.pickle', compress=None, protocol='pickle' +) +``` + +In the above snippet, the [DocVec][docarray.array.doc_list.doc_list.DocVec] is stored as the file `simple-dv.pickle`. + +### Bytes + +These methods just serialize your data, without saving it to a file: + +- [to_bytes()][docarray.array.doc_list.io.IOMixinArray.to_bytes] saves a [DocVec][docarray.array.doc_list.doc_list.DocVec] to a byte object. +- [from_bytes()][docarray.array.doc_list.io.IOMixinArray.from_bytes] loads a [DocVec][docarray.array.doc_list.doc_list.DocVec] from a byte object. + +!!! note + These methods are used under the hood by [save_binary()][docarray.array.doc_list.io.IOMixinArray.to_base64] and [`load_binary()`][docarray.array.doc_list.io.IOMixinArray.load_binary] to prepare/load/save to a binary file. You can also use them directly to work with byte files. + +Like working with binary files: + +- You can use `protocol` to choose between `pickle` and `protobuf`. +- You can use multiple compression methods: `lz4`, `bz2`, `lzma`, `zlib`, and `gzip`. + +```python +from docarray import BaseDoc, DocVec + + +class SimpleDoc(BaseDoc): + text: str + + +dv = DocVec[SimpleDoc]([SimpleDoc(text=f'doc {i}') for i in range(2)]) + +bytes_dv = dv.to_bytes(protocol='pickle', compress=None) + +dv_from_bytes = DocVec[SimpleDoc].from_bytes(bytes_dv, compress=None, protocol='pickle') +``` + +### CSV + +!!! warning + [`DocVec`][docarray.array.doc_vec.doc_vec.DocVec] does not support `.to_csv()` or `from_csv()`. + This is because CSV is a row-based format while DocVec has a column-based data layout. + To overcome this, you can convert your [`DocVec`][docarray.array.doc_vec.doc_vec.DocVec] + to a [`DocList`][docarray.array.doc_list.doc_list.DocList]. + + ```python + from docarray import BaseDoc, DocList, DocVec + + + class SimpleDoc(BaseDoc): + text: str + + + dv = DocVec[SimpleDoc]([SimpleDoc(text=f'doc {i}') for i in range(2)]) + + dv.to_doc_list().to_csv('simple-dl.csv') + dv_from_csv = DocList[SimpleDoc].from_csv('simple-dl.csv').to_doc_vec() + ``` + + For more details you can check the [DocList section on CSV serialization](#csv) + +### Pandas.Dataframe + +- [`from_dataframe()`][docarray.array.doc_list.io.IOMixinArray.from_dataframe] loads a [DocVec][docarray.array.doc_list.doc_list.DocVec] from a [Pandas Dataframe](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html). +- [`to_dataframe()`][docarray.array.doc_list.io.IOMixinArray.to_dataframe] saves a [DocVec][docarray.array.doc_list.doc_list.DocVec] to a [Pandas Dataframe](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html). + +```python +from docarray import BaseDoc, DocVec + + +class SimpleDoc(BaseDoc): + text: str + + +dv = DocVec[SimpleDoc]([SimpleDoc(text=f'doc {i}') for i in range(2)]) + +df = dv.to_dataframe() +dv_from_dataframe = DocVec[SimpleDoc].from_dataframe(df) +print(dv_from_dataframe) +``` + diff --git a/tests/units/array/test_array_from_to_bytes.py b/tests/units/array/test_array_from_to_bytes.py index 7530b2b82be..fba84cdfa80 100644 --- a/tests/units/array/test_array_from_to_bytes.py +++ b/tests/units/array/test_array_from_to_bytes.py @@ -1,6 +1,6 @@ import pytest -from docarray import BaseDoc, DocList +from docarray import BaseDoc, DocList, DocVec from docarray.documents import ImageDoc from docarray.typing import NdArray @@ -16,8 +16,9 @@ class MyDoc(BaseDoc): ) @pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) @pytest.mark.parametrize('show_progress', [False, True]) -def test_from_to_bytes(protocol, compress, show_progress): - da = DocList[MyDoc]( +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_to_bytes(protocol, compress, show_progress, array_cls): + da = array_cls[MyDoc]( [ MyDoc( embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') @@ -28,7 +29,7 @@ def test_from_to_bytes(protocol, compress, show_progress): bytes_da = da.to_bytes( protocol=protocol, compress=compress, show_progress=show_progress ) - da2 = DocList[MyDoc].from_bytes( + da2 = array_cls[MyDoc].from_bytes( bytes_da, protocol=protocol, compress=compress, show_progress=show_progress ) assert len(da2) == 2 @@ -46,8 +47,9 @@ def test_from_to_bytes(protocol, compress, show_progress): ) @pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) @pytest.mark.parametrize('show_progress', [False, True]) -def test_from_to_base64(protocol, compress, show_progress): - da = DocList[MyDoc]( +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_to_base64(protocol, compress, show_progress, array_cls): + da = array_cls[MyDoc]( [ MyDoc( embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') @@ -58,7 +60,7 @@ def test_from_to_base64(protocol, compress, show_progress): bytes_da = da.to_base64( protocol=protocol, compress=compress, show_progress=show_progress ) - da2 = DocList[MyDoc].from_base64( + da2 = array_cls[MyDoc].from_base64( bytes_da, protocol=protocol, compress=compress, show_progress=show_progress ) assert len(da2) == 2 diff --git a/tests/units/array/test_array_from_to_csv.py b/tests/units/array/test_array_from_to_csv.py index b04323f63d4..968967279d5 100644 --- a/tests/units/array/test_array_from_to_csv.py +++ b/tests/units/array/test_array_from_to_csv.py @@ -3,7 +3,7 @@ import pytest -from docarray import BaseDoc, DocList +from docarray import BaseDoc, DocList, DocVec from docarray.documents import ImageDoc from tests import TOYDATA_DIR @@ -38,6 +38,7 @@ def test_to_from_csv(tmpdir, nested_doc_cls): assert os.path.isfile(tmp_file) da_from = DocList[nested_doc_cls].from_csv(tmp_file) + assert isinstance(da_from, DocList) for doc1, doc2 in zip(da, da_from): assert doc1 == doc2 @@ -46,6 +47,7 @@ def test_from_csv_nested(nested_doc_cls): da = DocList[nested_doc_cls].from_csv( file_path=str(TOYDATA_DIR / 'docs_nested.csv') ) + assert isinstance(da, DocList) assert len(da) == 3 for i, doc in enumerate(da): @@ -108,6 +110,7 @@ class Book(BaseDoc): year: int books = DocList[Book].from_csv(file_path=remote_url) + assert isinstance(books, DocList) assert len(books) == 3 @@ -116,6 +119,7 @@ def test_doc_list_error(tmpdir): class Book(BaseDoc): title: str + # not testing DocVec bc it already fails here (as it should!) docs = DocList([Book(title='hello'), Book(title='world')]) tmp_file = str(tmpdir / 'tmp.csv') with pytest.raises(TypeError): @@ -143,3 +147,20 @@ class BasisUnion(BaseDoc): docs_basic.to_csv(str(tmp_path) + ".csv") docs_copy = DocList[BasisUnion].from_csv(str(tmp_path) + ".csv") assert docs_copy == docs_basic + + +def test_to_from_csv_docvec_raises(): + class Book(BaseDoc): + title: str + author: str + year: int + + books = DocVec[Book]( + [Book(title='It\'s me, hi', author='I\'m the problem it\'s me', year=2022)] + ) + + with pytest.raises(NotImplementedError): + books.to_csv('dummy/file/path') + + with pytest.raises(NotImplementedError): + DocVec[Book].from_csv('dummy/file/path') diff --git a/tests/units/array/test_array_from_to_json.py b/tests/units/array/test_array_from_to_json.py index 5767fe9de11..030f023d82c 100644 --- a/tests/units/array/test_array_from_to_json.py +++ b/tests/units/array/test_array_from_to_json.py @@ -1,6 +1,12 @@ -from docarray import BaseDoc, DocList +from typing import Optional + +import numpy as np +import pytest +import torch + +from docarray import BaseDoc, DocList, DocVec from docarray.documents import ImageDoc -from docarray.typing import NdArray +from docarray.typing import NdArray, TorchTensor class MyDoc(BaseDoc): @@ -9,7 +15,7 @@ class MyDoc(BaseDoc): image: ImageDoc -def test_from_to_json(): +def test_from_to_json_doclist(): da = DocList[MyDoc]( [ MyDoc( @@ -30,6 +36,104 @@ def test_from_to_json(): assert da2[1].image.url is None +@pytest.mark.parametrize('tensor_type', [TorchTensor, NdArray]) +def test_from_to_json_docvec(tensor_type): + def generate_docs(tensor_type): + class InnerDoc(BaseDoc): + tens: tensor_type + + class MyDoc(BaseDoc): + text: str + num: Optional[int] + tens: tensor_type + tens_none: Optional[tensor_type] + inner: InnerDoc + inner_none: Optional[InnerDoc] + inner_vec: DocVec[InnerDoc] + inner_vec_none: Optional[DocVec[InnerDoc]] + + def _rand_vec_gen(tensor_type): + arr = np.random.rand(5) + if tensor_type == TorchTensor: + arr = torch.from_numpy(arr).to(torch.float32) + return arr + + inner = InnerDoc(tens=_rand_vec_gen(tensor_type)) + inner_vec = DocVec[InnerDoc]([inner, inner], tensor_type=tensor_type) + vec = DocVec[MyDoc]( + [ + MyDoc( + text=str(i), + num=None, + tens=_rand_vec_gen(tensor_type), + inner=inner, + inner_none=None, + inner_vec=inner_vec, + inner_vec_none=None, + ) + for i in range(5) + ], + tensor_type=tensor_type, + ) + return vec + + v = generate_docs(tensor_type) + bytes_ = v.to_json() + + v_after = DocVec[v.doc_type].from_json(bytes_, tensor_type=tensor_type) + + assert v_after.tensor_type == v.tensor_type + assert set(v_after._storage.columns.keys()) == set(v._storage.columns.keys()) + assert v_after._storage == v._storage + + +@pytest.mark.tensorflow +def test_from_to_json_docvec_tf(): + from docarray.typing import TensorFlowTensor + + def generate_docs(): + class InnerDoc(BaseDoc): + tens: TensorFlowTensor + + class MyDoc(BaseDoc): + text: str + num: Optional[int] + tens: TensorFlowTensor + tens_none: Optional[TensorFlowTensor] + inner: InnerDoc + inner_none: Optional[InnerDoc] + inner_vec: DocVec[InnerDoc] + inner_vec_none: Optional[DocVec[InnerDoc]] + + inner = InnerDoc(tens=np.random.rand(5)) + inner_vec = DocVec[InnerDoc]([inner, inner], tensor_type=TensorFlowTensor) + vec = DocVec[MyDoc]( + [ + MyDoc( + text=str(i), + num=None, + tens=np.random.rand(5), + inner=inner, + inner_none=None, + inner_vec=inner_vec, + inner_vec_none=None, + ) + for i in range(5) + ], + tensor_type=TensorFlowTensor, + ) + return vec + + v = generate_docs() + bytes_ = v.to_json() + + v_after = DocVec[v.doc_type].from_json(bytes_, tensor_type=TensorFlowTensor) + + assert v_after.tensor_type == v.tensor_type + assert set(v_after._storage.columns.keys()) == set(v._storage.columns.keys()) + assert v_after._storage == v._storage + + def test_union_type(): from typing import Union diff --git a/tests/units/array/test_array_from_to_pandas.py b/tests/units/array/test_array_from_to_pandas.py index bef4427ca6d..463b3c62e86 100644 --- a/tests/units/array/test_array_from_to_pandas.py +++ b/tests/units/array/test_array_from_to_pandas.py @@ -3,7 +3,7 @@ import pandas as pd import pytest -from docarray import BaseDoc, DocList +from docarray import BaseDoc, DocList, DocVec from docarray.documents import ImageDoc @@ -20,7 +20,8 @@ class MyDocNested(MyDoc): return MyDocNested -def test_to_from_pandas_df(nested_doc_cls): +@pytest.mark.parametrize('doc_vec', [False, True]) +def test_to_from_pandas_df(nested_doc_cls, doc_vec): da = DocList[nested_doc_cls]( [ nested_doc_cls( @@ -34,6 +35,8 @@ def test_to_from_pandas_df(nested_doc_cls): ), ] ) + if doc_vec: + da = da.to_doc_vec() df = da.to_dataframe() assert isinstance(df, pd.DataFrame) assert len(df) == 2 @@ -52,7 +55,12 @@ def test_to_from_pandas_df(nested_doc_cls): ] ).all() - da_from_df = DocList[nested_doc_cls].from_dataframe(df) + if doc_vec: + da_from_df = DocVec[nested_doc_cls].from_dataframe(df) + assert isinstance(da_from_df, DocVec) + else: + da_from_df = DocList[nested_doc_cls].from_dataframe(df) + assert isinstance(da_from_df, DocList) for doc1, doc2 in zip(da, da_from_df): assert doc1 == doc2 @@ -76,26 +84,29 @@ class Outer(BaseDoc): return doc -def test_from_pandas_without_schema_raise_exception(): +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_pandas_without_schema_raise_exception(array_cls): with pytest.raises(TypeError, match='no document schema defined'): df = pd.DataFrame( columns=['title', 'count'], data=[['title 0', 0], ['title 1', 1]] ) - DocList.from_dataframe(df=df) + array_cls.from_dataframe(df=df) -def test_from_pandas_with_wrong_schema_raise_exception(nested_doc): +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_pandas_with_wrong_schema_raise_exception(nested_doc, array_cls): with pytest.raises(ValueError, match='Column names do not match the schema'): df = pd.DataFrame( columns=['title', 'count'], data=[['title 0', 0], ['title 1', 1]] ) - DocList[nested_doc.__class__].from_dataframe(df=df) + array_cls[nested_doc.__class__].from_dataframe(df=df) def test_doc_list_error(): class Book(BaseDoc): title: str + # not testing DocVec bc it already fails here (as it should!) docs = DocList([Book(title='hello'), Book(title='world')]) with pytest.raises(TypeError): docs.to_dataframe() diff --git a/tests/units/array/test_array_save_load.py b/tests/units/array/test_array_save_load.py index a56ad13064a..fc9a3a22609 100644 --- a/tests/units/array/test_array_save_load.py +++ b/tests/units/array/test_array_save_load.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from docarray import BaseDoc, DocList +from docarray import BaseDoc, DocList, DocVec from docarray.documents import ImageDoc from docarray.typing import NdArray @@ -20,10 +20,11 @@ class MyDoc(BaseDoc): ) @pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) @pytest.mark.parametrize('show_progress', [False, True]) -def test_array_save_load_binary(protocol, compress, tmp_path, show_progress): +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_array_save_load_binary(protocol, compress, tmp_path, show_progress, array_cls): tmp_file = os.path.join(tmp_path, 'test') - da = DocList[MyDoc]( + da = array_cls[MyDoc]( [ MyDoc( embedding=[1, 2, 3, 4, 5], text='hello', image=ImageDoc(url='aux.png') @@ -36,7 +37,7 @@ def test_array_save_load_binary(protocol, compress, tmp_path, show_progress): tmp_file, protocol=protocol, compress=compress, show_progress=show_progress ) - da2 = DocList[MyDoc].load_binary( + da2 = array_cls[MyDoc].load_binary( tmp_file, protocol=protocol, compress=compress, show_progress=show_progress ) @@ -56,8 +57,12 @@ def test_array_save_load_binary(protocol, compress, tmp_path, show_progress): ) @pytest.mark.parametrize('compress', ['lz4', 'bz2', 'lzma', 'zlib', 'gzip', None]) @pytest.mark.parametrize('show_progress', [False, True]) -def test_array_save_load_binary_streaming(protocol, compress, tmp_path, show_progress): +@pytest.mark.parametrize('to_doc_vec', [True, False]) +def test_array_save_load_binary_streaming( + protocol, compress, tmp_path, show_progress, to_doc_vec +): tmp_file = os.path.join(tmp_path, 'test') + array_cls = DocVec if to_doc_vec else DocList da = DocList[MyDoc]() @@ -74,20 +79,20 @@ def _extend_da(num_docs=100): ) _extend_da() + if to_doc_vec: + da = da.to_doc_vec() da.save_binary( tmp_file, protocol=protocol, compress=compress, show_progress=show_progress ) - da2 = DocList[MyDoc]() - da_generator = DocList[MyDoc].load_binary( + da_after = array_cls[MyDoc].load_binary( tmp_file, protocol=protocol, compress=compress, show_progress=show_progress ) - for i, doc in enumerate(da_generator): + for i, doc in enumerate(da_after): assert doc.id == da[i].id assert doc.text == da[i].text assert doc.image.url == da[i].image.url - da2.append(doc) - assert len(da2) == 100 + assert i == 99 diff --git a/tests/units/document/test_base_document.py b/tests/units/document/test_base_document.py index 475c03b07df..d02d052c33d 100644 --- a/tests/units/document/test_base_document.py +++ b/tests/units/document/test_base_document.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from docarray import DocList +from docarray import DocList, DocVec from docarray.base_doc.doc import BaseDoc from docarray.typing import NdArray @@ -66,11 +66,32 @@ class NestedDoc(BaseDoc): return nested_docs +@pytest.fixture +def nested_docs_docvec(): + class SimpleDoc(BaseDoc): + simple_tens: NdArray[10] + + class NestedDoc(BaseDoc): + docs: DocVec[SimpleDoc] + hello: str = 'world' + + nested_docs = NestedDoc( + docs=DocList[SimpleDoc]([SimpleDoc(simple_tens=np.ones(10)) for j in range(2)]), + ) + + return nested_docs + + def test_nested_to_dict(nested_docs): d = nested_docs.dict() assert (d['docs'][0]['simple_tens'] == np.ones(10)).all() +def test_nested_docvec_to_dict(nested_docs_docvec): + d = nested_docs_docvec.dict() + assert (d['docs'][0]['simple_tens'] == np.ones(10)).all() + + def test_nested_to_dict_exclude(nested_docs): d = nested_docs.dict(exclude={'docs'}) assert 'docs' not in d.keys()