Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit c835681

Browse filesBrowse files
JohannesMessnerJohannesMessner
authored
fix: protobuf (de)ser for docvec (#1639)
Signed-off-by: Johannes Messner <messnerjo@gmail.com>
1 parent f36c621 commit c835681
Copy full SHA for c835681

6 files changed

+360-43Lines changed: 360 additions & 43 deletions

File tree

Expand file treeCollapse file tree
Open diff view settings
Filter options
Expand file treeCollapse file tree
Open diff view settings
Collapse file

‎docarray/array/doc_vec/doc_vec.py‎

Copy file name to clipboardExpand all lines: docarray/array/doc_vec/doc_vec.py
+131-19Lines changed: 131 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
overload,
1818
)
1919

20+
import numpy as np
2021
from pydantic import BaseConfig, parse_obj_as
2122
from typing_inspect import typingGenericAlias
2223

@@ -34,7 +35,12 @@
3435
if TYPE_CHECKING:
3536
from pydantic.fields import ModelField
3637

37-
from docarray.proto import DocVecProto
38+
from docarray.proto import (
39+
DocVecProto,
40+
ListOfDocArrayProto,
41+
ListOfDocVecProto,
42+
NdArrayProto,
43+
)
3844

3945
torch_available = is_torch_available()
4046
if torch_available:
@@ -54,6 +60,56 @@
5460
T = TypeVar('T', bound='DocVec')
5561
IndexIterType = Union[slice, Iterable[int], Iterable[bool], None]
5662

63+
NONE_NDARRAY_PROTO_SHAPE = (0,)
64+
NONE_NDARRAY_PROTO_DTYPE = 'None'
65+
66+
67+
def _none_ndarray_proto() -> 'NdArrayProto':
68+
from docarray.proto import NdArrayProto
69+
70+
zeros_arr = parse_obj_as(NdArray, np.zeros(NONE_NDARRAY_PROTO_SHAPE))
71+
nd_proto = NdArrayProto()
72+
nd_proto.dense.buffer = zeros_arr.tobytes()
73+
nd_proto.dense.ClearField('shape')
74+
nd_proto.dense.shape.extend(list(zeros_arr.shape))
75+
nd_proto.dense.dtype = NONE_NDARRAY_PROTO_DTYPE
76+
77+
return nd_proto
78+
79+
80+
def _none_docvec_proto() -> 'DocVecProto':
81+
from docarray.proto import DocVecProto
82+
83+
return DocVecProto()
84+
85+
86+
def _none_list_of_docvec_proto() -> 'ListOfDocArrayProto':
87+
from docarray.proto import ListOfDocVecProto
88+
89+
return ListOfDocVecProto()
90+
91+
92+
def _is_none_ndarray_proto(proto: 'NdArrayProto') -> bool:
93+
return (
94+
proto.dense.shape == list(NONE_NDARRAY_PROTO_SHAPE)
95+
and proto.dense.dtype == NONE_NDARRAY_PROTO_DTYPE
96+
)
97+
98+
99+
def _is_none_docvec_proto(proto: 'DocVecProto') -> bool:
100+
return (
101+
proto.tensor_columns == {}
102+
and proto.doc_columns == {}
103+
and proto.docs_vec_columns == {}
104+
and proto.any_columns == {}
105+
)
106+
107+
108+
def _is_none_list_of_docvec_proto(proto: 'ListOfDocVecProto') -> bool:
109+
from docarray.proto import ListOfDocVecProto
110+
111+
return isinstance(proto, ListOfDocVecProto) and len(proto.data) == 0
112+
57113

58114
class DocVec(AnyDocArray[T_doc]):
59115
"""
@@ -243,7 +299,7 @@ def _check_doc_field_not_none(field_name, doc):
243299
elif issubclass(field_type, AnyDocArray):
244300
if first_doc_is_none:
245301
_verify_optional_field_of_docs(docs)
246-
doc_columns[field_name] = None
302+
docs_vec_columns[field_name] = None
247303
else:
248304
docs_list = list()
249305
for doc in docs:
@@ -534,48 +590,104 @@ def __len__(self):
534590

535591
@classmethod
536592
def from_protobuf(cls: Type[T], pb_msg: 'DocVecProto') -> T:
537-
"""create a Document from a protobuf message"""
593+
"""create a DocVec from a protobuf message"""
594+
595+
tensor_columns: Dict[str, Optional[AbstractTensor]] = {}
596+
doc_columns: Dict[str, Optional['DocVec']] = {}
597+
docs_vec_columns: Dict[str, Optional[ListAdvancedIndexing['DocVec']]] = {}
598+
any_columns: Dict[str, ListAdvancedIndexing] = {}
599+
600+
for tens_col_name, tens_col_proto in pb_msg.tensor_columns.items():
601+
if _is_none_ndarray_proto(tens_col_proto):
602+
# handle values that were None before serialization
603+
tensor_columns[tens_col_name] = None
604+
else:
605+
# TODO(johannes): handle torch, tf, numpy
606+
tensor_columns[tens_col_name] = NdArray.from_protobuf(tens_col_proto)
607+
608+
for doc_col_name, doc_col_proto in pb_msg.doc_columns.items():
609+
if _is_none_docvec_proto(doc_col_proto):
610+
# handle values that were None before serialization
611+
doc_columns[doc_col_name] = None
612+
else:
613+
col_doc_type: Type = cls.doc_type._get_field_type(doc_col_name)
614+
doc_columns[doc_col_name] = DocVec.__class_getitem__(
615+
col_doc_type
616+
).from_protobuf(doc_col_proto)
617+
618+
for docs_vec_col_name, docs_vec_col_proto in pb_msg.docs_vec_columns.items():
619+
vec_list: Optional[ListAdvancedIndexing]
620+
if _is_none_list_of_docvec_proto(docs_vec_col_proto):
621+
# handle values that were None before serialization
622+
vec_list = None
623+
else:
624+
vec_list = ListAdvancedIndexing()
625+
for doc_list_proto in docs_vec_col_proto.data:
626+
col_doc_type = cls.doc_type._get_field_type(
627+
docs_vec_col_name
628+
).doc_type
629+
vec_list.append(
630+
DocVec.__class_getitem__(col_doc_type).from_protobuf(
631+
doc_list_proto
632+
)
633+
)
634+
docs_vec_columns[docs_vec_col_name] = vec_list
635+
636+
for any_col_name, any_col_proto in pb_msg.any_columns.items():
637+
any_column: ListAdvancedIndexing = ListAdvancedIndexing()
638+
for node_proto in any_col_proto.data:
639+
content = cls.doc_type._get_content_from_node_proto(
640+
node_proto, any_col_name
641+
)
642+
any_column.append(content)
643+
any_columns[any_col_name] = any_column
644+
538645
storage = ColumnStorage(
539-
pb_msg.tensor_columns,
540-
pb_msg.doc_columns,
541-
pb_msg.docs_vec_columns,
542-
pb_msg.any_columns,
646+
tensor_columns=tensor_columns,
647+
doc_columns=doc_columns,
648+
docs_vec_columns=docs_vec_columns,
649+
any_columns=any_columns,
543650
)
544651

545652
return cls.from_columns_storage(storage)
546653

547654
def to_protobuf(self) -> 'DocVecProto':
548655
"""Convert DocVec into a Protobuf message"""
549656
from docarray.proto import (
550-
DocListProto,
551657
DocVecProto,
552658
ListOfAnyProto,
553659
ListOfDocArrayProto,
660+
ListOfDocVecProto,
554661
NdArrayProto,
555662
)
556663

557-
da_proto = DocListProto()
558-
for doc in self:
559-
da_proto.docs.append(doc.to_protobuf())
560-
561664
doc_columns_proto: Dict[str, DocVecProto] = dict()
562665
tensor_columns_proto: Dict[str, NdArrayProto] = dict()
563666
da_columns_proto: Dict[str, ListOfDocArrayProto] = dict()
564667
any_columns_proto: Dict[str, ListOfAnyProto] = dict()
565668

566669
for field, col_doc in self._storage.doc_columns.items():
567-
doc_columns_proto[field] = (
568-
col_doc.to_protobuf() if col_doc is not None else None
569-
)
670+
if col_doc is None:
671+
# put dummy empty DocVecProto for serialization
672+
doc_columns_proto[field] = _none_docvec_proto()
673+
else:
674+
doc_columns_proto[field] = col_doc.to_protobuf()
570675
for field, col_tens in self._storage.tensor_columns.items():
571-
tensor_columns_proto[field] = (
572-
col_tens.to_protobuf() if col_tens is not None else None
573-
)
676+
if col_tens is None:
677+
# put dummy empty NdArrayProto for serialization
678+
tensor_columns_proto[field] = _none_ndarray_proto()
679+
else:
680+
tensor_columns_proto[field] = (
681+
col_tens.to_protobuf() if col_tens is not None else None
682+
)
574683
for field, col_da in self._storage.docs_vec_columns.items():
575-
list_proto = ListOfDocArrayProto()
684+
list_proto = ListOfDocVecProto()
576685
if col_da:
577686
for docs in col_da:
578687
list_proto.data.append(docs.to_protobuf())
688+
else:
689+
# put dummy empty ListOfDocVecProto for serialization
690+
list_proto = _none_list_of_docvec_proto()
579691
da_columns_proto[field] = list_proto
580692
for field, col_any in self._storage.any_columns.items():
581693
list_proto = ListOfAnyProto()
Collapse file

‎docarray/proto/__init__.py‎

Copy file name to clipboardExpand all lines: docarray/proto/__init__.py
+3Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
DocVecProto,
1818
ListOfAnyProto,
1919
ListOfDocArrayProto,
20+
ListOfDocVecProto,
2021
NdArrayProto,
2122
NodeProto,
2223
)
@@ -28,6 +29,7 @@
2829
DocVecProto,
2930
ListOfAnyProto,
3031
ListOfDocArrayProto,
32+
ListOfDocVecProto,
3133
NdArrayProto,
3234
NodeProto,
3335
)
@@ -40,6 +42,7 @@
4042
'DocVecProto',
4143
'DocListProto',
4244
'ListOfDocArrayProto',
45+
'ListOfDocVecProto',
4346
'ListOfAnyProto',
4447
'DictOfAnyProto',
4548
]
Collapse file

‎docarray/proto/docarray.proto‎

Copy file name to clipboardExpand all lines: docarray/proto/docarray.proto
+5-1Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,9 +100,13 @@ message ListOfDocArrayProto {
100100
repeated DocListProto data = 1;
101101
}
102102

103+
message ListOfDocVecProto {
104+
repeated DocVecProto data = 1;
105+
}
106+
103107
message DocVecProto{
104108
map<string, NdArrayProto> tensor_columns = 1; // a dict of document columns
105109
map<string, DocVecProto> doc_columns = 2; // a dict of tensor columns
106-
map<string, ListOfDocArrayProto> docs_vec_columns = 3; // a dict of document array columns
110+
map<string, ListOfDocVecProto> docs_vec_columns = 3; // a dict of document array columns
107111
map<string, ListOfAnyProto> any_columns = 4; // a dict of any columns. Used for the rest of the data
108112
}
Collapse file

‎docarray/proto/pb/docarray_pb2.py‎

Copy file name to clipboardExpand all lines: docarray/proto/pb/docarray_pb2.py
+13-11Lines changed: 13 additions & 11 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
0 (0)
Morty Proxy This is a proxified and sanitized view of the page, visit original site.