From d43057b91872eca2b174d24af9d9f1d3f4ae7925 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 22 May 2023 16:19:11 +0200 Subject: [PATCH 01/21] feat: json and dict for docvec Signed-off-by: Johannes Messner --- docarray/array/doc_vec/doc_vec.py | 6 +++++- docarray/base_doc/doc.py | 26 ++++++++++++++------------ 2 files changed, 19 insertions(+), 13 deletions(-) diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index e27f6882fe9..fdd267f1322 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -21,6 +21,7 @@ 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 @@ -54,7 +55,7 @@ IndexIterType = Union[slice, Iterable[int], Iterable[bool], None] -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 @@ -509,6 +510,9 @@ def __len__(self): # IO related # #################### + def _docarray_to_json_compatible(self) -> List[Dict]: + return [doc._docarray_to_json_compatible() for doc in self] + @classmethod def from_protobuf(cls: Type[T], pb_msg: 'DocVecProto') -> T: """create a Document from a protobuf message""" 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 From c45bfcaadb3bbca2464d2ee6bbd250474381585f Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 22 May 2023 16:19:28 +0200 Subject: [PATCH 02/21] test: add tests Signed-off-by: Johannes Messner --- tests/units/array/test_array_from_to_json.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/units/array/test_array_from_to_json.py b/tests/units/array/test_array_from_to_json.py index c36b8af92a9..2623beb7348 100644 --- a/tests/units/array/test_array_from_to_json.py +++ b/tests/units/array/test_array_from_to_json.py @@ -1,3 +1,5 @@ +import pytest + from docarray import BaseDoc, DocList from docarray.documents import ImageDoc from docarray.typing import NdArray @@ -9,7 +11,8 @@ class MyDoc(BaseDoc): image: ImageDoc -def test_from_to_json(): +@pytest.mark.parametrize('doc_vec', [True]) +def test_from_to_json(doc_vec): da = DocList[MyDoc]( [ MyDoc( @@ -18,6 +21,8 @@ def test_from_to_json(): MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), ] ) + if doc_vec: + da = da.to_doc_vec() json_da = da.to_json() da2 = DocList[MyDoc].from_json(json_da) assert len(da2) == 2 From 564d14441bba5abc3b3ca37225c036bafdee0a0e Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 22 May 2023 16:21:44 +0200 Subject: [PATCH 03/21] test: add docvec to dict test Signed-off-by: Johannes Messner --- tests/units/document/test_base_document.py | 23 +++++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) 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() From 76f9c8e97c196209ce7398ccae6b16346f642004 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 22 May 2023 16:44:22 +0200 Subject: [PATCH 04/21] feat: to from dataframe for docvec Signed-off-by: Johannes Messner --- docarray/array/doc_list/io.py | 4 +++- docarray/array/doc_vec/doc_vec.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/docarray/array/doc_list/io.py b/docarray/array/doc_list/io.py index c2b531c2550..aa1d8b3e98d 100644 --- a/docarray/array/doc_list/io.py +++ b/docarray/array/doc_list/io.py @@ -443,7 +443,7 @@ def to_csv( writer.writerow(doc_dict) @classmethod - def from_dataframe(cls, df: 'pd.DataFrame') -> 'DocList': + def from_dataframe(cls, df: 'pd.DataFrame') -> 'T': """ Load a `DocList` from a `pandas.DataFrame` following the schema defined in the [`.doc_type`][docarray.DocList] attribute. @@ -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': diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index fdd267f1322..0d1c058bf87 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -100,7 +100,7 @@ class DocVec(IOMixinArray, AnyDocArray[T_doc]): AnyTensor or Union of NdArray and TorchTensor """ - doc_type: Type[T_doc] + doc_type: Type[T_doc] = AnyDoc def __init__( self: T, From 73a1ac7417b1192562d240a4413300e78953f9a7 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 22 May 2023 16:44:46 +0200 Subject: [PATCH 05/21] test: dataframe docvec tests Signed-off-by: Johannes Messner --- .../units/array/test_array_from_to_pandas.py | 25 +++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/tests/units/array/test_array_from_to_pandas.py b/tests/units/array/test_array_from_to_pandas.py index 612bf826b7b..d2fe017723c 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 @@ -19,7 +19,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( @@ -30,6 +31,8 @@ def test_to_from_pandas_df(nested_doc_cls): nested_doc_cls(text='hello world', image=ImageDoc()), ] ) + if doc_vec: + da = da.to_doc_vec() df = da.to_dataframe() assert isinstance(df, pd.DataFrame) assert len(df) == 2 @@ -47,7 +50,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 @@ -71,26 +79,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() From f83fb4f7ac73bd1cd7a1ae573cbacea83e7affa5 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 22 May 2023 17:13:06 +0200 Subject: [PATCH 06/21] feat: to from csv for docvec Signed-off-by: Johannes Messner --- docarray/array/doc_list/io.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/docarray/array/doc_list/io.py b/docarray/array/doc_list/io.py index aa1d8b3e98d..e2360de2c40 100644 --- a/docarray/array/doc_list/io.py +++ b/docarray/array/doc_list/io.py @@ -377,13 +377,12 @@ def from_csv( @classmethod def _from_csv_file( cls, file: Union[StringIO, TextIOWrapper], dialect: Union[str, csv.Dialect] - ) -> 'DocList': - from docarray import DocList + ) -> '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 +404,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' From ca8dc12d19653a66d3b06f6fd5f3510fb9d5ac5c Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 22 May 2023 17:13:32 +0200 Subject: [PATCH 07/21] test: test csv with docvec Signed-off-by: Johannes Messner --- tests/units/array/test_array_from_to_csv.py | 35 ++++++++++++++------- 1 file changed, 23 insertions(+), 12 deletions(-) diff --git a/tests/units/array/test_array_from_to_csv.py b/tests/units/array/test_array_from_to_csv.py index d00ea172c4e..2ab0d66c230 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 @@ -21,8 +21,9 @@ class MyDocNested(MyDoc): return MyDocNested -def test_to_from_csv(tmpdir, nested_doc_cls): - da = DocList[nested_doc_cls]( +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_to_from_csv(tmpdir, nested_doc_cls, array_cls): + da = array_cls[nested_doc_cls]( [ nested_doc_cls( count=0, @@ -37,15 +38,18 @@ def test_to_from_csv(tmpdir, nested_doc_cls): da.to_csv(tmp_file) assert os.path.isfile(tmp_file) - da_from = DocList[nested_doc_cls].from_csv(tmp_file) + da_from = array_cls[nested_doc_cls].from_csv(tmp_file) + assert isinstance(da_from, array_cls) for doc1, doc2 in zip(da, da_from): assert doc1 == doc2 -def test_from_csv_nested(nested_doc_cls): - da = DocList[nested_doc_cls].from_csv( +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_csv_nested(nested_doc_cls, array_cls): + da = array_cls[nested_doc_cls].from_csv( file_path=str(TOYDATA_DIR / 'docs_nested.csv') ) + assert isinstance(da, array_cls) assert len(da) == 3 for i, doc in enumerate(da): @@ -89,17 +93,22 @@ class Outer(BaseDoc): return doc -def test_from_csv_without_schema_raise_exception(): +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_csv_without_schema_raise_exception(array_cls): with pytest.raises(TypeError, match='no document schema defined'): - DocList.from_csv(file_path=str(TOYDATA_DIR / 'docs_nested.csv')) + array_cls.from_csv(file_path=str(TOYDATA_DIR / 'docs_nested.csv')) -def test_from_csv_with_wrong_schema_raise_exception(nested_doc): +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_csv_with_wrong_schema_raise_exception(nested_doc, array_cls): with pytest.raises(ValueError, match='Column names do not match the schema'): - DocList[nested_doc.__class__].from_csv(file_path=str(TOYDATA_DIR / 'docs.csv')) + array_cls[nested_doc.__class__].from_csv( + file_path=str(TOYDATA_DIR / 'docs.csv') + ) -def test_from_remote_csv_file(): +@pytest.mark.parametrize('array_cls', [DocList, DocVec]) +def test_from_remote_csv_file(array_cls): remote_url = 'https://github.com/docarray/docarray/blob/main/tests/toydata/books.csv?raw=true' class Book(BaseDoc): @@ -107,7 +116,8 @@ class Book(BaseDoc): author: str year: int - books = DocList[Book].from_csv(file_path=remote_url) + books = array_cls[Book].from_csv(file_path=remote_url) + assert isinstance(books, array_cls) assert len(books) == 3 @@ -116,6 +126,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): From b115637d00a16134cf0a6158341ddff0610f3797 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Jun 2023 16:46:47 +0200 Subject: [PATCH 08/21] feat: pickle serialization for docvec Signed-off-by: Johannes Messner --- docarray/array/doc_vec/column_storage.py | 5 +++++ docarray/array/doc_vec/doc_vec.py | 2 +- docarray/base_doc/mixins/io.py | 3 ++- tests/units/array/test_array_from_to_bytes.py | 9 +++++---- 4 files changed, 13 insertions(+), 6 deletions(-) diff --git a/docarray/array/doc_vec/column_storage.py b/docarray/array/doc_vec/column_storage.py index bd098ae8f34..feeb0a996d3 100644 --- a/docarray/array/doc_vec/column_storage.py +++ b/docarray/array/doc_vec/column_storage.py @@ -123,6 +123,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 0ed9f2b8d6f..ab58ee4df78 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -157,7 +157,7 @@ class DocVec(IOMixinArray, AnyDocArray[T_doc]): AnyTensor or Union of NdArray and TorchTensor """ - doc_type: Type[T_doc] = AnyDoc + doc_type: Type[T_doc] = AnyDoc # TODO(johannes): should this be BaseDoc? def __init__( self: T, diff --git a/docarray/base_doc/mixins/io.py b/docarray/base_doc/mixins/io.py index b19747d7a9b..252408cb195 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, @@ -170,7 +171,7 @@ def to_bytes( def from_bytes( cls: Type[T], data: bytes, - protocol: str = 'protobuf', + protocol: Literal['protobuf', 'pickle'] = 'protobuf', compress: Optional[str] = None, ) -> T: """Build Document object from binary bytes diff --git a/tests/units/array/test_array_from_to_bytes.py b/tests/units/array/test_array_from_to_bytes.py index 7cd9f0dfd8c..6e4868adec7 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 From bd86985eac44a5a2707141011def112e1e541461 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Jun 2023 17:24:20 +0200 Subject: [PATCH 09/21] feat: protbuf array serialization for docvec Signed-off-by: Johannes Messner --- docarray/array/doc_list/doc_list.py | 6 ++++++ docarray/array/doc_list/io.py | 13 ++++++++----- docarray/array/doc_vec/doc_vec.py | 6 ++++++ 3 files changed, 20 insertions(+), 5 deletions(-) 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 e2360de2c40..6252f4376a7 100644 --- a/docarray/array/doc_list/io.py +++ b/docarray/array/doc_list/io.py @@ -564,6 +564,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], @@ -594,12 +599,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/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index ab58ee4df78..ab83de8ecf8 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -589,6 +589,12 @@ def __len__(self): # IO related # #################### + @classmethod + def _get_proto_class(cls: Type[T]): + from docarray.proto import DocVecProto + + return DocVecProto + def _docarray_to_json_compatible(self) -> List[Dict]: return [doc._docarray_to_json_compatible() for doc in self] From c280ff28ad37f1e7ea89742409c7de27a1d3ba31 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Jun 2023 17:27:12 +0200 Subject: [PATCH 10/21] test: test base64 deser for docvec Signed-off-by: Johannes Messner --- tests/units/array/test_array_from_to_bytes.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/units/array/test_array_from_to_bytes.py b/tests/units/array/test_array_from_to_bytes.py index 6e4868adec7..519264cbfbc 100644 --- a/tests/units/array/test_array_from_to_bytes.py +++ b/tests/units/array/test_array_from_to_bytes.py @@ -47,8 +47,9 @@ def test_from_to_bytes(protocol, compress, show_progress, array_cls): ) @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') @@ -59,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 From ad881cf60f52992c1cb0af97b6a9c65bd773565d Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Wed, 14 Jun 2023 17:34:24 +0200 Subject: [PATCH 11/21] test: test save and load for docvec Signed-off-by: Johannes Messner --- tests/units/array/test_array_save_load.py | 25 ++++++++++++++--------- 1 file changed, 15 insertions(+), 10 deletions(-) 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 From 4b1b533ccf299eca6e37a3571ea71ed161e5a053 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 19 Jun 2023 17:33:35 +0200 Subject: [PATCH 12/21] feat: docvec json column wise Signed-off-by: Johannes Messner --- docarray/array/doc_vec/column_storage.py | 27 ++++++++ docarray/array/doc_vec/doc_vec.py | 80 +++++++++++++++++++++++- 2 files changed, 105 insertions(+), 2 deletions(-) diff --git a/docarray/array/doc_vec/column_storage.py b/docarray/array/doc_vec/column_storage.py index feeb0a996d3..8d0107a5028 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 + ) + class ColumnStorageView(dict, MutableMapping[str, Any]): index: int diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index ab83de8ecf8..07f3789c4ac 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -18,6 +18,7 @@ ) import numpy as np +from orjson import orjson from pydantic import BaseConfig, parse_obj_as from typing_inspect import typingGenericAlias @@ -595,8 +596,83 @@ def _get_proto_class(cls: Type[T]): return DocVecProto - def _docarray_to_json_compatible(self) -> List[Dict]: - return [doc._docarray_to_json_compatible() for doc in self] + 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, 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( From 0603fc5edc9cf19f9e1535e1a89ca2d8a9ed9957 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 20 Jun 2023 17:00:23 +0200 Subject: [PATCH 13/21] test: add test for docvec json Signed-off-by: Johannes Messner --- tests/units/array/test_array_from_to_json.py | 110 ++++++++++++++++++- 1 file changed, 104 insertions(+), 6 deletions(-) diff --git a/tests/units/array/test_array_from_to_json.py b/tests/units/array/test_array_from_to_json.py index 466f279d806..7ffad3464db 100644 --- a/tests/units/array/test_array_from_to_json.py +++ b/tests/units/array/test_array_from_to_json.py @@ -1,8 +1,12 @@ +from typing import Optional + +import numpy as np import pytest +import torch -from docarray import BaseDoc, DocList +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): @@ -11,8 +15,7 @@ class MyDoc(BaseDoc): image: ImageDoc -@pytest.mark.parametrize('doc_vec', [True]) -def test_from_to_json(doc_vec): +def test_from_to_json_doclist(): da = DocList[MyDoc]( [ MyDoc( @@ -21,8 +24,6 @@ def test_from_to_json(doc_vec): MyDoc(embedding=[5, 4, 3, 2, 1], text='hello world', image=ImageDoc()), ] ) - if doc_vec: - da = da.to_doc_vec() json_da = da.to_json() da2 = DocList[MyDoc].from_json(json_da) assert len(da2) == 2 @@ -35,6 +36,103 @@ def test_from_to_json(doc_vec): 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]) + 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) + ] + ) + 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 From c6ace8e6337da8ea7d53257c1c333c40df446763 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 20 Jun 2023 17:29:13 +0200 Subject: [PATCH 14/21] test: add tensor type arg Signed-off-by: Johannes Messner --- tests/units/array/test_array_from_to_json.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/units/array/test_array_from_to_json.py b/tests/units/array/test_array_from_to_json.py index 7ffad3464db..030f023d82c 100644 --- a/tests/units/array/test_array_from_to_json.py +++ b/tests/units/array/test_array_from_to_json.py @@ -106,7 +106,7 @@ class MyDoc(BaseDoc): inner_vec_none: Optional[DocVec[InnerDoc]] inner = InnerDoc(tens=np.random.rand(5)) - inner_vec = DocVec[InnerDoc]([inner, inner]) + inner_vec = DocVec[InnerDoc]([inner, inner], tensor_type=TensorFlowTensor) vec = DocVec[MyDoc]( [ MyDoc( @@ -119,7 +119,8 @@ class MyDoc(BaseDoc): inner_vec_none=None, ) for i in range(5) - ] + ], + tensor_type=TensorFlowTensor, ) return vec From 51719b2a60df308597e38b00a249bf7d463d5db6 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 26 Jun 2023 14:16:26 +0200 Subject: [PATCH 15/21] fix: mypy stuff Signed-off-by: Johannes Messner --- docarray/array/doc_list/io.py | 11 ++++++----- docarray/array/doc_vec/doc_vec.py | 10 +++++++--- docarray/base_doc/mixins/io.py | 2 +- 3 files changed, 14 insertions(+), 9 deletions(-) diff --git a/docarray/array/doc_list/io.py b/docarray/array/doc_list/io.py index 6252f4376a7..6bba4afd0ba 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. @@ -376,7 +375,9 @@ def from_csv( @classmethod def _from_csv_file( - cls, file: Union[StringIO, TextIOWrapper], dialect: Union[str, csv.Dialect] + cls: Type['T'], + file: Union[StringIO, TextIOWrapper], + dialect: Union[str, csv.Dialect], ) -> 'T': rows = csv.DictReader(file, dialect=dialect) @@ -442,7 +443,7 @@ def to_csv( writer.writerow(doc_dict) @classmethod - def from_dataframe(cls, df: 'pd.DataFrame') -> 'T': + 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. diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index 84d03d0788e..4a325e38201 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -18,7 +18,7 @@ ) import numpy as np -from orjson import orjson +import orjson from pydantic import BaseConfig, parse_obj_as from typing_inspect import typingGenericAlias @@ -60,6 +60,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,) @@ -158,7 +160,7 @@ class DocVec(IOMixinArray, AnyDocArray[T_doc]): AnyTensor or Union of NdArray and TorchTensor """ - doc_type: Type[T_doc] = AnyDoc # TODO(johannes): should this be BaseDoc? + doc_type: Type[T_doc] = BaseDoc # type: ignore def __init__( self: T, @@ -630,7 +632,9 @@ def from_json( @classmethod def _from_json_col_dict( - cls, json_columns: Dict[str, Any], tensor_type: Type[AbstractTensor] = NdArray + cls: Type[T], + json_columns: Dict[str, Any], + tensor_type: Type[AbstractTensor] = NdArray, ) -> T: tensor_cols = json_columns['tensor_columns'] diff --git a/docarray/base_doc/mixins/io.py b/docarray/base_doc/mixins/io.py index bc8ff8cf61c..aa000391d33 100644 --- a/docarray/base_doc/mixins/io.py +++ b/docarray/base_doc/mixins/io.py @@ -213,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 From ad5f5bdc478a4e745365e9ffc2628e6de92e7d98 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 26 Jun 2023 14:34:35 +0200 Subject: [PATCH 16/21] fix: raising of error when needed Signed-off-by: Johannes Messner --- docarray/array/doc_list/io.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docarray/array/doc_list/io.py b/docarray/array/doc_list/io.py index 6bba4afd0ba..92816c44799 100644 --- a/docarray/array/doc_list/io.py +++ b/docarray/array/doc_list/io.py @@ -357,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'): From 200dbaccd8cd0ac13644b4c19ccb79d4cf2fcc20 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 26 Jun 2023 14:59:50 +0200 Subject: [PATCH 17/21] fix: more exception raising Signed-off-by: Johannes Messner --- docarray/array/doc_list/io.py | 4 ++-- docarray/array/doc_vec/doc_vec.py | 6 +++++- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/docarray/array/doc_list/io.py b/docarray/array/doc_list/io.py index 92816c44799..dfad809083e 100644 --- a/docarray/array/doc_list/io.py +++ b/docarray/array/doc_list/io.py @@ -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 diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index 4a325e38201..cc45b328515 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -168,7 +168,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)' From 8d1f446d80cd5b8e884aa920002cfe1b2da25cfd Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 26 Jun 2023 15:35:51 +0200 Subject: [PATCH 18/21] fix: mypy Signed-off-by: Johannes Messner --- docarray/base_doc/mixins/io.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docarray/base_doc/mixins/io.py b/docarray/base_doc/mixins/io.py index aa000391d33..8eb9751af43 100644 --- a/docarray/base_doc/mixins/io.py +++ b/docarray/base_doc/mixins/io.py @@ -173,7 +173,7 @@ def to_bytes( def from_bytes( cls: Type[T], data: bytes, - protocol: Literal['protobuf', 'pickle'] = 'protobuf', + protocol: str = 'protobuf', compress: Optional[str] = None, ) -> T: """Build Document object from binary bytes From 6815720104f60d2b96444dc17ac94124189db388 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 26 Jun 2023 17:12:19 +0200 Subject: [PATCH 19/21] refactor: don't expose to/from csv for docvec Signed-off-by: Johannes Messner --- docarray/array/doc_list/io.py | 6 +++--- docarray/array/doc_vec/doc_vec.py | 34 +++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/docarray/array/doc_list/io.py b/docarray/array/doc_list/io.py index dfad809083e..9a2f3e756ed 100644 --- a/docarray/array/doc_list/io.py +++ b/docarray/array/doc_list/io.py @@ -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() diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index cc45b328515..528e713ac23 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -35,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 ( @@ -878,3 +880,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()`.' + ) From 6b5ddc76d4f45a686469c6ed02be63b344ce024b Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Mon, 26 Jun 2023 17:23:43 +0200 Subject: [PATCH 20/21] test: adjust tests Signed-off-by: Johannes Messner --- docarray/array/doc_vec/doc_vec.py | 2 +- tests/units/array/test_array_from_to_csv.py | 52 ++++++++++++--------- 2 files changed, 32 insertions(+), 22 deletions(-) diff --git a/docarray/array/doc_vec/doc_vec.py b/docarray/array/doc_vec/doc_vec.py index 528e713ac23..7eb8cff199b 100644 --- a/docarray/array/doc_vec/doc_vec.py +++ b/docarray/array/doc_vec/doc_vec.py @@ -900,7 +900,7 @@ def from_csv( cls: Type['T'], file_path: str, encoding: str = 'utf-8', - dialect: Union[str, csv.Dialect] = 'excel', + dialect: Union[str, 'csv.Dialect'] = 'excel', ) -> 'T': """ DocVec does not support `.from_csv()`. This is because CSV is a row-based format diff --git a/tests/units/array/test_array_from_to_csv.py b/tests/units/array/test_array_from_to_csv.py index c81523fb2fc..968967279d5 100644 --- a/tests/units/array/test_array_from_to_csv.py +++ b/tests/units/array/test_array_from_to_csv.py @@ -21,9 +21,8 @@ class MyDocNested(MyDoc): return MyDocNested -@pytest.mark.parametrize('array_cls', [DocList, DocVec]) -def test_to_from_csv(tmpdir, nested_doc_cls, array_cls): - da = array_cls[nested_doc_cls]( +def test_to_from_csv(tmpdir, nested_doc_cls): + da = DocList[nested_doc_cls]( [ nested_doc_cls( count=0, @@ -38,18 +37,17 @@ def test_to_from_csv(tmpdir, nested_doc_cls, array_cls): da.to_csv(tmp_file) assert os.path.isfile(tmp_file) - da_from = array_cls[nested_doc_cls].from_csv(tmp_file) - assert isinstance(da_from, array_cls) + 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 -@pytest.mark.parametrize('array_cls', [DocList, DocVec]) -def test_from_csv_nested(nested_doc_cls, array_cls): - da = array_cls[nested_doc_cls].from_csv( +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, array_cls) + assert isinstance(da, DocList) assert len(da) == 3 for i, doc in enumerate(da): @@ -93,22 +91,17 @@ class Outer(BaseDoc): return doc -@pytest.mark.parametrize('array_cls', [DocList, DocVec]) -def test_from_csv_without_schema_raise_exception(array_cls): +def test_from_csv_without_schema_raise_exception(): with pytest.raises(TypeError, match='no document schema defined'): - array_cls.from_csv(file_path=str(TOYDATA_DIR / 'docs_nested.csv')) + DocList.from_csv(file_path=str(TOYDATA_DIR / 'docs_nested.csv')) -@pytest.mark.parametrize('array_cls', [DocList, DocVec]) -def test_from_csv_with_wrong_schema_raise_exception(nested_doc, array_cls): +def test_from_csv_with_wrong_schema_raise_exception(nested_doc): with pytest.raises(ValueError, match='Column names do not match the schema'): - array_cls[nested_doc.__class__].from_csv( - file_path=str(TOYDATA_DIR / 'docs.csv') - ) + DocList[nested_doc.__class__].from_csv(file_path=str(TOYDATA_DIR / 'docs.csv')) -@pytest.mark.parametrize('array_cls', [DocList, DocVec]) -def test_from_remote_csv_file(array_cls): +def test_from_remote_csv_file(): remote_url = 'https://github.com/docarray/docarray/blob/main/tests/toydata/books.csv?raw=true' class Book(BaseDoc): @@ -116,8 +109,8 @@ class Book(BaseDoc): author: str year: int - books = array_cls[Book].from_csv(file_path=remote_url) - assert isinstance(books, array_cls) + books = DocList[Book].from_csv(file_path=remote_url) + assert isinstance(books, DocList) assert len(books) == 3 @@ -154,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') From 587c20a7d916283c0c08062b33dc01fa7a5309e4 Mon Sep 17 00:00:00 2001 From: Johannes Messner Date: Tue, 27 Jun 2023 14:49:24 +0200 Subject: [PATCH 21/21] docs: add documentation for docvec io Signed-off-by: Johannes Messner --- docs/user_guide/sending/serialization.md | 173 ++++++++++++++++++++++- 1 file changed, 167 insertions(+), 6 deletions(-) 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) +``` +