feat: add from to dict io (#160)

davidbp · web-flow · commit 3b169d4dbbd3 · 2022-03-03T07:37:08.000+01:00
diff --git a/docarray/array/mixins/__init__.py b/docarray/array/mixins/__init__.py
@@ -12,6 +12,7 @@
 from .io.common import CommonIOMixin
 from .io.csv import CsvIOMixin
 from .io.dataframe import DataframeIOMixin
+from .io.dict import DictIOMixin
 from .io.from_gen import FromGeneratorMixin
 from .io.json import JsonIOMixin
 from .io.pushpull import PushPullMixin
@@ -41,6 +42,7 @@ class AllMixins(
     CsvIOMixin,
     JsonIOMixin,
     BinaryIOMixin,
+    DictIOMixin,
     CommonIOMixin,
     EmbedMixin,
     PushPullMixin,
diff --git a/docarray/array/mixins/io/dict.py b/docarray/array/mixins/io/dict.py
@@ -0,0 +1,33 @@
+from typing import TYPE_CHECKING, Type, List
+
+if TYPE_CHECKING:
+    from ....types import T
+
+
+class DictIOMixin:
+    """Save/load a DocumentArray into a dict of the form `{offset_0: doc_0, offset_1: doc_1, ...}`"""
+
+    def to_dict(self, protocol: str = 'jsonschema', **kwargs) -> List:
+        """Convert the object into a Python dict of the form `{offset_0: doc_0, offset_1: doc_1, ...}`
+
+        :param protocol: `jsonschema` or `protobuf`
+        :return: a Python list
+        """
+        return {k: d.to_dict(protocol=protocol, **kwargs) for k, d in enumerate(self)}
+
+    @classmethod
+    def from_dict(cls: Type['T'], input_dict: dict, *args, **kwargs) -> 'T':
+        """Import a :class:`DocumentArray` from a :class:`dict` object of the form `{offset_0: doc_0, offset_1: doc_1, ...}`
+
+        :param input_dict: a `dict` object.
+        :return: a :class:`DocumentArray` object
+        """
+        from .... import Document, DocumentArray
+
+        da = cls.empty(len(input_dict), *args, **kwargs)
+
+        for offset, d in input_dict.items():
+            da[offset] = Document(
+                {k: v for k, v in d.items() if (not isinstance(v, float) or v == v)}
+            )
+        return da
diff --git a/docarray/array/mixins/io/json.py b/docarray/array/mixins/io/json.py
@@ -99,7 +99,3 @@ def to_json(self, protocol: str = 'jsonschema', **kwargs) -> str:
         :return: a Python list
         """
         return json.dumps(self.to_list(protocol=protocol, **kwargs))
-
-    # to comply with Document interfaces but less semantically accurate
-    to_dict = to_list
-    from_dict = from_list
diff --git a/docs/fundamentals/documentarray/serialization.md b/docs/fundamentals/documentarray/serialization.md
@@ -12,6 +12,7 @@ Moreover, there is the ability to store/load `DocumentArray` objects to/from dis
 - Base64 (compressed): `.from_base64()`/`.to_base64()` 
 - Protobuf Message: `.from_protobuf()`/`.to_protobuf()`
 - Python List: `.from_list()`/`.to_list()`
+- Python Dict: `.from_dict()`/`.to_dict()`
 - Pandas Dataframe: `.from_dataframe()`/`.to_dataframe()`
 - Cloud: `.push()`/`.pull()`
 
@@ -324,6 +325,25 @@ da.to_list()
 More parameters and usages can be found in the Document-level {ref}`doc-dict`.
 ```
 
+
+## From/to dict
+
+
+Serializing to/from Python dict is less frequently used for the same reason as `Document.to_dict()`: it is often an intermediate step of serializing to JSON. You can do:
+
+```python
+from docarray import DocumentArray, Document
+
+da = DocumentArray([Document(text='hello'), Document(text='world')])
+da.to_dict()
+```
+
+```text
+{0: {'id': '3b31cb4c993f11ec8d12787b8ab3f5de', 'mime_type': 'text/plain', 'text': 'hello', 1: {'id': '3b31cca0993f11ec8d12787b8ab3f5de', 'mime_type': 'text/plain', 'text': 'world'}}```
+```
+
+
+
 ## From/to dataframe
 
 ```{important}
diff --git a/tests/unit/array/mixins/test_io.py b/tests/unit/array/mixins/test_io.py
@@ -172,6 +172,34 @@ def test_from_to_pd_dataframe(da_cls, config, start_storage):
     assert da2[1].tags == {}
 
 
+@pytest.mark.parametrize(
+    'da_cls, config',
+    [
+        (DocumentArrayInMemory, lambda: None),
+        (DocumentArraySqlite, lambda: None),
+        (DocumentArrayPqlite, lambda: PqliteConfig(n_dim=3)),
+        (DocumentArrayWeaviate, lambda: WeaviateConfig(n_dim=3)),
+        (DocumentArrayQdrant, lambda: QdrantConfig(n_dim=3)),
+    ],
+)
+def test_from_to_dict(da_cls, config, start_storage):
+    da_dict = da_cls.empty(2, config=config()).to_dict()
+    assert len(da_cls.from_dict(da_dict, config=config())) == 2
+
+    # more complicated
+    da = da_cls.empty(2, config=config())
+
+    da[:, 'embedding'] = [[1, 2, 3], [4, 5, 6]]
+    da[:, 'tensor'] = [[1, 2], [2, 1]]
+    da[0, 'tags'] = {'hello': 'world'}
+    da_dict = da.to_dict()
+
+    da2 = da_cls.from_dict(da_dict, config=config())
+
+    assert da2[0].tags == {'hello': 'world'}
+    assert da2[1].tags == {}
+
+
 @pytest.mark.parametrize(
     'da_cls, config',
     [
diff --git a/tests/unit/array/test_from_to_bytes.py b/tests/unit/array/test_from_to_bytes.py
@@ -95,14 +95,21 @@ def test_from_to_protobuf(target_da):
 
 @pytest.mark.parametrize('target', [DocumentArray.empty(10), random_docs(10)])
 @pytest.mark.parametrize('protocol', ['jsonschema', 'protobuf'])
-@pytest.mark.parametrize('to_fn', ['dict', 'json'])
+@pytest.mark.parametrize('to_fn', ['list', 'json'])
 def test_from_to_safe_list(target, protocol, to_fn):
     da_r = getattr(DocumentArray, f'from_{to_fn}')(
         getattr(target, f'to_{to_fn}')(protocol=protocol), protocol=protocol
     )
     assert da_r == target
 
 
+@pytest.mark.parametrize('target', [DocumentArray.empty(10), random_docs(10)])
+def test_from_to_safe_dict(target):
+    target_dict = getattr(target, f'to_dict')(target)
+    da_r = getattr(DocumentArray, f'from_dict')(target_dict)
+    assert da_r == target
+
+
 @pytest.mark.parametrize('protocol', ['protobuf', 'pickle'])
 @pytest.mark.parametrize('show_progress', [True, False])
 def test_push_pull_show_progress(show_progress, protocol):
