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 efeab90

Browse filesBrowse files
feat: tensor_type for all DocVec serializations (#1679)
Signed-off-by: Johannes Messner <messnerjo@gmail.com>
1 parent 00e980d commit efeab90
Copy full SHA for efeab90

12 files changed

+751-363Lines changed: 751 additions & 363 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_list/doc_list.py‎

Copy file name to clipboardExpand all lines: docarray/array/doc_list/doc_list.py
+2-2Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from typing_inspect import is_union_type
2020

2121
from docarray.array.any_array import AnyDocArray
22-
from docarray.array.doc_list.io import IOMixinArray
22+
from docarray.array.doc_list.io import IOMixinDocList
2323
from docarray.array.doc_list.pushpull import PushPullMixin
2424
from docarray.array.list_advance_indexing import IndexIterType, ListAdvancedIndexing
2525
from docarray.base_doc import AnyDoc, BaseDoc
@@ -42,7 +42,7 @@
4242
class DocList(
4343
ListAdvancedIndexing[T_doc],
4444
PushPullMixin,
45-
IOMixinArray,
45+
IOMixinDocList,
4646
AnyDocArray[T_doc],
4747
):
4848
"""
Collapse file

‎docarray/array/doc_list/io.py‎

Copy file name to clipboardExpand all lines: docarray/array/doc_list/io.py
+46-19Lines changed: 46 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
Type,
2424
TypeVar,
2525
Union,
26+
cast,
2627
)
2728

2829
import orjson
@@ -40,9 +41,12 @@
4041
if TYPE_CHECKING:
4142
import pandas as pd
4243

44+
from docarray.array.doc_vec.doc_vec import DocVec
45+
from docarray.array.doc_vec.io import IOMixinDocVec
4346
from docarray.proto import DocListProto
47+
from docarray.typing.tensor.abstract_tensor import AbstractTensor
4448

45-
T = TypeVar('T', bound='IOMixinArray')
49+
T = TypeVar('T', bound='IOMixinDocList')
4650
T_doc = TypeVar('T_doc', bound=BaseDoc)
4751

4852
ARRAY_PROTOCOLS = {'protobuf-array', 'pickle-array', 'json-array'}
@@ -96,7 +100,7 @@ def __getitem__(self, item: slice):
96100
return self.content[item]
97101

98102

99-
class IOMixinArray(Iterable[T_doc]):
103+
class IOMixinDocList(Iterable[T_doc]):
100104
doc_type: Type[T_doc]
101105

102106
@abstractmethod
@@ -515,8 +519,6 @@ class Person(BaseDoc):
515519
doc_dict = _access_path_dict_to_nested_dict(access_path2val)
516520
docs.append(doc_type.parse_obj(doc_dict))
517521

518-
if not isinstance(docs, cls):
519-
return cls(docs)
520522
return docs
521523

522524
def to_dataframe(self) -> 'pd.DataFrame':
@@ -577,11 +579,13 @@ def _load_binary_all(
577579
protocol: Optional[str],
578580
compress: Optional[str],
579581
show_progress: bool,
582+
tensor_type: Optional[Type['AbstractTensor']] = None,
580583
):
581584
"""Read a `DocList` object from a binary file
582585
:param protocol: protocol to use. It can be 'pickle-array', 'protobuf-array', 'pickle' or 'protobuf'
583586
:param compress: compress algorithm to use between `lz4`, `bz2`, `lzma`, `zlib`, `gzip`
584587
:param show_progress: show progress bar, only works when protocol is `pickle` or `protobuf`
588+
:param tensor_type: only relevant for DocVec; tensor_type of the DocVec
585589
:return: a `DocList`
586590
"""
587591
with file_ctx as fp:
@@ -603,12 +607,20 @@ def _load_binary_all(
603607
proto = cls._get_proto_class()()
604608
proto.ParseFromString(d)
605609

606-
return cls.from_protobuf(proto)
610+
if tensor_type is not None:
611+
cls_ = cast('IOMixinDocVec', cls)
612+
return cls_.from_protobuf(proto, tensor_type=tensor_type)
613+
else:
614+
return cls.from_protobuf(proto)
607615
elif protocol is not None and protocol == 'pickle-array':
608616
return pickle.loads(d)
609617

610618
elif protocol is not None and protocol == 'json-array':
611-
return cls.from_json(d)
619+
if tensor_type is not None:
620+
cls_ = cast('IOMixinDocVec', cls)
621+
return cls_.from_json(d, tensor_type=tensor_type)
622+
else:
623+
return cls.from_json(d)
612624

613625
# Binary format for streaming case
614626
else:
@@ -658,6 +670,10 @@ def _load_binary_all(
658670
pbar.update(
659671
t, advance=1, total_size=str(filesize.decimal(_total_size))
660672
)
673+
if tensor_type is not None:
674+
cls__ = cast(Type['DocVec'], cls)
675+
# mypy doesn't realize that cls_ is callable
676+
return cls__(docs, tensor_type=tensor_type) # type: ignore
661677
return cls(docs)
662678

663679
@classmethod
@@ -724,6 +740,27 @@ def _load_binary_stream(
724740
t, advance=1, total_size=str(filesize.decimal(_total_size))
725741
)
726742

743+
@staticmethod
744+
def _get_file_context(
745+
file: Union[str, bytes, pathlib.Path, io.BufferedReader, _LazyRequestReader],
746+
protocol: str,
747+
compress: Optional[str] = None,
748+
) -> Tuple[Union[nullcontext, io.BufferedReader], Optional[str], Optional[str]]:
749+
load_protocol: Optional[str] = protocol
750+
load_compress: Optional[str] = compress
751+
file_ctx: Union[nullcontext, io.BufferedReader]
752+
if isinstance(file, (io.BufferedReader, _LazyRequestReader, bytes)):
753+
file_ctx = nullcontext(file)
754+
# by checking path existence we allow file to be of type Path, LocalPath, PurePath and str
755+
elif isinstance(file, (str, pathlib.Path)) and os.path.exists(file):
756+
load_protocol, load_compress = _protocol_and_compress_from_file_path(
757+
file, protocol, compress
758+
)
759+
file_ctx = open(file, 'rb')
760+
else:
761+
raise FileNotFoundError(f'cannot find file {file}')
762+
return file_ctx, load_protocol, load_compress
763+
727764
@classmethod
728765
def load_binary(
729766
cls: Type[T],
@@ -753,19 +790,9 @@ def load_binary(
753790
:return: a `DocList` object
754791
755792
"""
756-
load_protocol: Optional[str] = protocol
757-
load_compress: Optional[str] = compress
758-
file_ctx: Union[nullcontext, io.BufferedReader]
759-
if isinstance(file, (io.BufferedReader, _LazyRequestReader, bytes)):
760-
file_ctx = nullcontext(file)
761-
# by checking path existence we allow file to be of type Path, LocalPath, PurePath and str
762-
elif isinstance(file, (str, pathlib.Path)) and os.path.exists(file):
763-
load_protocol, load_compress = _protocol_and_compress_from_file_path(
764-
file, protocol, compress
765-
)
766-
file_ctx = open(file, 'rb')
767-
else:
768-
raise FileNotFoundError(f'cannot find file {file}')
793+
file_ctx, load_protocol, load_compress = cls._get_file_context(
794+
file, protocol, compress
795+
)
769796
if streaming:
770797
if load_protocol not in SINGLE_PROTOCOLS:
771798
raise ValueError(

0 commit comments

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