diff --git a/docarray/__init__.py b/docarray/__init__.py index 85aecf4498d..89a327d0786 100644 --- a/docarray/__init__.py +++ b/docarray/__init__.py @@ -1,4 +1,4 @@ -__version__ = '0.2.1' +__version__ = '0.3.0' from .document import Document from .array import DocumentArray diff --git a/docarray/array/mixins/io/csv.py b/docarray/array/mixins/io/csv.py index 0552d4ad010..9edca8820d2 100644 --- a/docarray/array/mixins/io/csv.py +++ b/docarray/array/mixins/io/csv.py @@ -72,15 +72,12 @@ def save_csv( if with_header: writer.writeheader() - from .... import Document - for d in self: - _d = d - if exclude_fields: - _d = Document(d, copy=True) - _d.pop(*exclude_fields) - - pd = _d.to_dict() + pd = d.to_dict( + protocol='jsonschema', + exclude=set(exclude_fields) if exclude_fields else None, + exclude_none=True, + ) if flatten_tags: t = pd.pop('tags') pd.update({f'tag__{k}': v for k, v in t.items()}) diff --git a/docarray/array/mixins/io/json.py b/docarray/array/mixins/io/json.py index 4b962632cde..2740bfa7fa0 100644 --- a/docarray/array/mixins/io/json.py +++ b/docarray/array/mixins/io/json.py @@ -10,12 +10,15 @@ class JsonIOMixin: """Save/load a array into a JSON file.""" - def save_json(self, file: Union[str, TextIO]) -> None: + def save_json( + self, file: Union[str, TextIO], protocol: str = 'jsonschema', **kwargs + ) -> None: """Save array elements into a JSON file. Comparing to :meth:`save_binary`, it is human-readable but slower to save/load and the file size larger. :param file: File or filename to which the data is saved. + :param protocol: `jsonschema` or `protobuf` """ if hasattr(file, 'write'): file_ctx = nullcontext(file) @@ -24,15 +27,17 @@ def save_json(self, file: Union[str, TextIO]) -> None: with file_ctx as fp: for d in self: - json.dump(d.to_dict(), fp) + json.dump(d.to_dict(protocol=protocol, **kwargs), fp) fp.write('\n') @classmethod - def load_json(cls: Type['T'], file: Union[str, TextIO]) -> 'T': + def load_json( + cls: Type['T'], file: Union[str, TextIO], protocol: str = 'jsonschema', **kwargs + ) -> 'T': """Load array elements from a JSON file. :param file: File or filename or a JSON string to which the data is saved. - + :param protocol: `jsonschema` or `protobuf` :return: a DocumentArrayLike object """ @@ -48,31 +53,38 @@ def load_json(cls: Type['T'], file: Union[str, TextIO]) -> 'T': constructor = Document.from_dict with file_ctx as fp: - return cls(constructor(v) for v in fp) + return cls(constructor(v, protocol=protocol, **kwargs) for v in fp) @classmethod - def from_json(cls: Type['T'], file: Union[str, TextIO]) -> 'T': - return cls.load_json(file) + def from_json( + cls: Type['T'], file: Union[str, TextIO], protocol: str = 'jsonschema', **kwargs + ) -> 'T': + return cls.load_json(file, protocol=protocol, **kwargs) @classmethod - def from_list(cls: Type['T'], values: List) -> 'T': + def from_list( + cls: Type['T'], values: List, protocol: str = 'jsonschema', **kwargs + ) -> 'T': from .... import Document - return cls(Document.from_dict(v) for v in values) + return cls(Document.from_dict(v, protocol=protocol, **kwargs) for v in values) - def to_list(self, strict: bool = True) -> List: + def to_list(self, protocol: str = 'jsonschema', **kwargs) -> List: """Convert the object into a Python list. - .. note:: - Array like object such as :class:`numpy.ndarray` will be converted to Python list. - + :param protocol: `jsonschema` or `protobuf` :return: a Python list """ - return [d.to_dict(strict=strict) for d in self] + return [d.to_dict(protocol=protocol, **kwargs) for d in self] - def to_json(self) -> str: + def to_json(self, protocol: str = 'jsonschema', **kwargs) -> str: """Convert the object into a JSON string. Can be loaded via :meth:`.load_json`. + :param protocol: `jsonschema` or `protobuf` :return: a Python list """ - return json.dumps(self.to_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/docarray/base.py b/docarray/base.py index e6a05dce29f..cafb814ab6f 100644 --- a/docarray/base.py +++ b/docarray/base.py @@ -52,7 +52,7 @@ def __init__( if _unknown_kwargs and unknown_fields_handler == 'catch': getattr(self, self._unresolved_fields_dest).update(_unknown_kwargs) - if _obj is None and not kwargs and self._data is None: + if not _obj and not kwargs and self._data is None: self._data = self._data_class(self) if self._data is None: diff --git a/docarray/document/mixins/porting.py b/docarray/document/mixins/porting.py index 6cbf7c1b357..9880e9c6055 100644 --- a/docarray/document/mixins/porting.py +++ b/docarray/document/mixins/porting.py @@ -1,5 +1,6 @@ import dataclasses import pickle +import warnings from typing import Optional, TYPE_CHECKING, Type, Dict, Any import base64 @@ -11,32 +12,76 @@ class PortingMixin: @classmethod - def from_dict(cls: Type['T'], obj: Dict) -> 'T': - from google.protobuf import json_format - from ...proto.docarray_pb2 import DocumentProto + def from_dict( + cls: Type['T'], obj: Dict, protocol: str = 'jsonschema', **kwargs + ) -> 'T': + """Convert a dict object into a Document. + + :param obj: a Python dict object + :param protocol: `jsonschema` or `protobuf` + :param kwargs: extra key-value args pass to pydantic and protobuf parser. + :return: the parsed Document object + """ + if protocol == 'jsonschema': + from ..pydantic_model import PydanticDocument + + return cls.from_pydantic_model(PydanticDocument.parse_obj(obj, **kwargs)) + elif protocol == 'protobuf': + from google.protobuf import json_format + from ...proto.docarray_pb2 import DocumentProto - pb_msg = DocumentProto() - json_format.ParseDict(obj, pb_msg) - return cls.from_protobuf(pb_msg) + pb_msg = DocumentProto() + json_format.ParseDict(obj, pb_msg, **kwargs) + return cls.from_protobuf(pb_msg) + else: + raise ValueError(f'protocol=`{protocol}` is not supported') @classmethod - def from_json(cls: Type['T'], obj: str) -> 'T': - from google.protobuf import json_format - from ...proto.docarray_pb2 import DocumentProto + def from_json( + cls: Type['T'], obj: str, protocol: str = 'jsonschema', **kwargs + ) -> 'T': + """Convert a JSON string into a Document. + + :param obj: a valid JSON string + :param protocol: `jsonschema` or `protobuf` + :param kwargs: extra key-value args pass to pydantic and protobuf parser. + :return: the parsed Document object + """ + if protocol == 'jsonschema': + from ..pydantic_model import PydanticDocument + + return cls.from_pydantic_model(PydanticDocument.parse_raw(obj, **kwargs)) + elif protocol == 'protobuf': + from google.protobuf import json_format + from ...proto.docarray_pb2 import DocumentProto + + pb_msg = DocumentProto() + json_format.Parse(obj, pb_msg, **kwargs) + return cls.from_protobuf(pb_msg) + else: + raise ValueError(f'protocol=`{protocol}` is not supported') - pb_msg = DocumentProto() - json_format.Parse(obj, pb_msg) - return cls.from_protobuf(pb_msg) + def to_dict(self, protocol: str = 'jsonschema', **kwargs) -> Dict[str, Any]: + """Convert itself into a Python dict object. - def to_dict(self, strict: bool = True) -> Dict[str, Any]: - if strict: + :param protocol: `jsonschema` or `protobuf` + :param kwargs: extra key-value args pass to pydantic and protobuf dumper. + :return: the dumped Document as a dict object + """ + if protocol == 'jsonschema': + return self.to_pydantic_model().dict(**kwargs) + elif protocol == 'protobuf': from google.protobuf.json_format import MessageToDict return MessageToDict( self.to_protobuf(), - preserving_proto_field_name=True, + **kwargs, ) else: + warnings.warn( + f'protocol=`{protocol}` is not supported, ' + f'the result dict is a Python dynamic typing dict without any promise on the schema.' + ) return dataclasses.asdict(self._data) def to_bytes( @@ -68,25 +113,33 @@ def from_bytes( """ bstr = decompress_bytes(data, algorithm=compress) if protocol == 'pickle': - d = pickle.loads(bstr) + return pickle.loads(bstr) elif protocol == 'protobuf': from ...proto.docarray_pb2 import DocumentProto pb_msg = DocumentProto() pb_msg.ParseFromString(bstr) - d = cls.from_protobuf(pb_msg) + return cls.from_protobuf(pb_msg) else: raise ValueError( f'protocol={protocol} is not supported. Can be only `protobuf` or pickle protocols 0-5.' ) - return d - def to_json(self) -> str: - from google.protobuf.json_format import MessageToJson + def to_json(self, protocol: str = 'jsonschema', **kwargs) -> str: + """Convert itself into a JSON string. + + :param protocol: `jsonschema` or `protobuf` + :param kwargs: extra key-value args pass to pydantic and protobuf dumper. + :return: the dumped JSON string + """ + if protocol == 'jsonschema': + return self.to_pydantic_model().json(**kwargs) + elif protocol == 'protobuf': + from google.protobuf.json_format import MessageToJson - return MessageToJson( - self.to_protobuf(), preserving_proto_field_name=True, sort_keys=True - ) + return MessageToJson(self.to_protobuf(), **kwargs) + else: + raise ValueError(f'protocol={protocol} is not supported.') def to_base64( self, protocol: str = 'pickle', compress: Optional[str] = None diff --git a/docarray/document/mixins/pydantic.py b/docarray/document/mixins/pydantic.py index 892d0d0f3d1..a8e2abbb399 100644 --- a/docarray/document/mixins/pydantic.py +++ b/docarray/document/mixins/pydantic.py @@ -47,10 +47,11 @@ def from_pydantic_model(cls: Type['T'], model: 'BaseModel') -> 'T': from ... import Document fields = {} + _field_chunks, _field_matches = None, None if model.chunks: - fields['chunks'] = [Document.from_pydantic_model(d) for d in model.chunks] + _field_chunks = [Document.from_pydantic_model(d) for d in model.chunks] if model.matches: - fields['matches'] = [Document.from_pydantic_model(d) for d in model.matches] + _field_matches = [Document.from_pydantic_model(d) for d in model.matches] for (field, value) in model.dict( exclude_none=True, exclude={'chunks', 'matches'} @@ -66,4 +67,10 @@ def from_pydantic_model(cls: Type['T'], model: 'BaseModel') -> 'T': fields[f_name] = np.array(value) else: fields[f_name] = value - return Document(**fields) + + d = Document(**fields) + if _field_chunks: + d.chunks = _field_chunks + if _field_matches: + d.matches = _field_matches + return d diff --git a/docarray/document/pydantic_model.py b/docarray/document/pydantic_model.py index d90157038a1..bc4cb2a9aa7 100644 --- a/docarray/document/pydantic_model.py +++ b/docarray/document/pydantic_model.py @@ -14,7 +14,8 @@ def _convert_ndarray_to_list(v: 'ArrayType'): - return to_list(v) + if v is not None: + return to_list(v) class PydanticDocument(BaseModel): diff --git a/docs/fundamentals/document/serialization.md b/docs/fundamentals/document/serialization.md index 523476ffe5d..f6d4e6c51de 100644 --- a/docs/fundamentals/document/serialization.md +++ b/docs/fundamentals/document/serialization.md @@ -11,6 +11,7 @@ One should use {ref}`DocumentArray for serializing multiple Documents +``` + +By default, it uses {ref}`JSON Schema and pydantic model` for serialization, i.e. `protocol='jsonschema'`. You can switch the method to `protocol='protobuf'`, which leverages Protobuf as the JSON serialization backend. + +```python +from docarray import Document + +d = Document(text='hello, world', embedding=np.array([1, 2, 3])) +d.to_json(protocol='protobuf') +``` + +```json { + "id": "db66bc2e77b311eca5f51e008a366d49", + "text": "hello, world", + "mimeType": "text/plain", "embedding": { - "cls_name": "numpy", "dense": { "buffer": "AQAAAAAAAAACAAAAAAAAAAMAAAAAAAAA", - "dtype": " +Note that you can pass extra arguments to control the include/exclude fields, lower/uppercase of the resulted JSON. For example, we can remove those fields that are empty or `none` from JSON via: + +```python +from docarray import Document + +d = Document(text='hello, world', embedding=[1, 2, 3]) +d.to_json(exclude_none=True) +``` + +```json +{"id": "cdbc4f7a77b411ec96ad1e008a366d49", "mime_type": "text/plain", "text": "hello, world", "embedding": [1, 2, 3]} +``` + +It is easier to eyes. But when building REST API, you do not need to explicitly do this, pydantic model handle everything for you. More information can be found in {ref}`fastapi-support`. + +```{seealso} +To find out what extra parameters you can pass to `to_json()`/`to_dict`, please check out: +- [`protocol='jsonschema'`](https://pydantic-docs.helpmanual.io/usage/exporting_models/#modeljson) +- [`protocol='protobuf'`](https://googleapis.dev/python/protobuf/latest/google/protobuf/json_format.html#google.protobuf.json_format.MessageToJson) ``` (doc-in-bytes)= @@ -146,11 +183,11 @@ print(len(d.to_base64(protocol='protobuf', compress='lz4'))) Note that the same `protocol` and `compress` must be followed when using `.from_base64`. - +(doc-dict)= ## From/to dict ```{important} -This feature requires `protobuf` dependency. You can do `pip install "docarray[full]"` to install it. +This feature requires `protobuf` or `pydantic` dependency. You can do `pip install "docarray[full]"` to install it. ``` You can serialize a Document as a Python `dict` via {meth}`~docarray.document.mixins.porting.PortingMixin.to_dict`, and then read from it via {meth}`~docarray.document.mixins.porting.PortingMixin.from_dict`. @@ -167,17 +204,12 @@ print(d_as_dict, d) ``` ```text -{'id': 'b29d39066d5611ec87661e008a366d49', 'text': 'hello, world', 'mime_type': 'text/plain', 'embedding': {'dense': {'buffer': 'AQAAAAAAAAACAAAAAAAAAAMAAAAAAAAA', 'shape': [3], 'dtype': ' + ``` -(strict-arg-explain)= -```{note} -Note that the result dict is very "stricted" in the sense that all fields and values boil down to very basic data type such as `int`, `float`, `string`. This behavior is designed due to the "serialization to `dict`" is often an intermediate step of serializing into JSON/YAML. Hence all values in `dict` must be schema-friendly. After all, a Python `dict` object means nothing if you are not working in Python. - -You can use `to_dict(strict=False)` to override this behavior. This will preserve the original Python data type of every value, which may not be JSON-friendly. But hey, you want it. -``` +As the intermediate step of `to_json()`/`from_json()` it is unlikely to use dict IO directly. Nonetheless, you can pass the same `protocol` and `kwargs` as described in {ref}`doc-json` to control the serialization behavior. ## From/to Protobuf diff --git a/docs/fundamentals/documentarray/serialization.md b/docs/fundamentals/documentarray/serialization.md index 217a5ba4fd3..87ceb1b10af 100644 --- a/docs/fundamentals/documentarray/serialization.md +++ b/docs/fundamentals/documentarray/serialization.md @@ -24,7 +24,7 @@ If you are building a webservice and want to use JSON for passing DocArray objec ``` ```{important} -This feature requires `protobuf` dependency. You can do `pip install "docarray[full]"` to install it. +Depending on which protocol you use, this feature requires `pydantic` or `protobuf` dependency. You can do `pip install "docarray[full]"` to install it. ``` @@ -37,7 +37,7 @@ da.to_json() ``` ```text -[{"id": "72db9a7e6e3211ec97f51e008a366d49", "text": "hello", "mime_type": "text/plain"}, {"id": "72db9cb86e3211ec97f51e008a366d49", "text": "world", "mime_type": "text/plain"}] +[{"id": "a677577877b611eca3811e008a366d49", "parent_id": null, "granularity": null, "adjacency": null, "blob": null, "tensor": null, "mime_type": "text/plain", "text": "hello", "weight": null, "uri": null, "tags": null, "offset": null, "location": null, "embedding": null, "modality": null, "evaluations": null, "scores": null, "chunks": null, "matches": null}, {"id": "a67758f477b611eca3811e008a366d49", "parent_id": null, "granularity": null, "adjacency": null, "blob": null, "tensor": null, "mime_type": "text/plain", "text": "world", "weight": null, "uri": null, "tags": null, "offset": null, "location": null, "embedding": null, "modality": null, "evaluations": null, "scores": null, "chunks": null, "matches": null}] ``` @@ -65,6 +65,11 @@ da_r.summary() ``` +```{seealso} +More parameters and usages can be found in the Document-level {ref}`doc-json`. +``` + + ## From/to bytes ```{important} @@ -221,6 +226,10 @@ docs { ## From/to list +```{important} +This feature requires `protobuf` or `pydantic` dependency. You can do `pip install "docarray[full]"` to install it. +``` + Serializing to/from Python list 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 @@ -234,7 +243,9 @@ da.to_list() [{'id': 'ae55782a6e4d11ec803c1e008a366d49', 'text': 'hello', 'mime_type': 'text/plain'}, {'id': 'ae557a146e4d11ec803c1e008a366d49', 'text': 'world', 'mime_type': 'text/plain'}] ``` -There is an argument `strict` shares {ref}`the same semantic` as in `Document.to_dict()`. +```{seealso} +More parameters and usages can be found in the Document-level {ref}`doc-dict`. +``` ## From/to dataframe diff --git a/tests/unit/array/test_from_to_bytes.py b/tests/unit/array/test_from_to_bytes.py index a07d3152b3b..ffb0f5f3550 100644 --- a/tests/unit/array/test_from_to_bytes.py +++ b/tests/unit/array/test_from_to_bytes.py @@ -75,9 +75,14 @@ def test_from_to_protobuf(target_da): DocumentArray.from_protobuf(target_da.to_protobuf()) -@pytest.mark.parametrize('target_da', [DocumentArray.empty(100), random_docs(100)]) -def test_from_to_safe_list(target_da): - DocumentArray.from_list(target_da.to_list()) +@pytest.mark.parametrize('target', [DocumentArray.empty(10), random_docs(10)]) +@pytest.mark.parametrize('protocol', ['jsonschema', 'protobuf']) +@pytest.mark.parametrize('to_fn', ['dict', '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('protocol', ['protobuf', 'pickle']) diff --git a/tests/unit/document/test_porting.py b/tests/unit/document/test_porting.py index ff8abc455bd..42d78d2d9ee 100644 --- a/tests/unit/document/test_porting.py +++ b/tests/unit/document/test_porting.py @@ -15,11 +15,14 @@ def test_to_from_bytes(protocol, compress): @pytest.mark.parametrize('target', [DocumentArray.empty(10), random_docs(10)]) -def test_dict_json(target): +@pytest.mark.parametrize('protocol', ['jsonschema', 'protobuf']) +@pytest.mark.parametrize('to_fn', ['dict', 'json']) +def test_dict_json(target, protocol, to_fn): for d in target: - d1 = Document.from_dict(d.to_dict()) - d2 = Document.from_json(d.to_json()) - assert d1 == d2 + d_r = getattr(Document, f'from_{to_fn}')( + getattr(d, f'to_{to_fn}')(protocol=protocol), protocol=protocol + ) + assert d == d_r @pytest.mark.parametrize('protocol', ['protobuf', 'pickle']) diff --git a/tests/unit/document/test_protobuf.py b/tests/unit/document/test_protobuf.py index c219ae64375..67bd08ec958 100644 --- a/tests/unit/document/test_protobuf.py +++ b/tests/unit/document/test_protobuf.py @@ -44,6 +44,7 @@ def test_from_to_namescore_default_dict(attr, meth): assert isinstance(d.scores, defaultdict) assert isinstance(d.scores['random_score1'], NamedScore) - r_d = getattr(Document, f'from_{meth}')(getattr(d, f'to_{meth}')()) + r = getattr(d, f'to_{meth}')() + r_d = getattr(Document, f'from_{meth}')(r) assert isinstance(r_d.scores, defaultdict) assert isinstance(r_d.scores['random_score2'], NamedScore) diff --git a/tests/unit/test_pydantic.py b/tests/unit/test_pydantic.py index 82fad635367..b6b43216db6 100644 --- a/tests/unit/test_pydantic.py +++ b/tests/unit/test_pydantic.py @@ -111,3 +111,8 @@ def test_match_to_from_pydantic(): assert da_r[0].matches[0].scores['cosine'] assert isinstance(da_r[0].matches[0].scores, defaultdict) assert isinstance(da_r[0].matches[0].scores['random_score'], NamedScore) + + +def test_with_embedding_no_tensor(): + d = Document(embedding=np.random.rand(2, 2)) + PydanticDocument.parse_obj(d.to_pydantic_model().dict())