feat: protbuf array serialization for docvec

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