test: add test for docvec json

Signed-off-by: Johannes Messner <[email protected]>
docarray · JohannesMessner · Jun 28, 2023 · May 22, 2023 · May 22, 2023 · May 22, 2023
commit 0603fc5edc9cf19f9e1535e1a89ca2d8a9ed9957
diff --git 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