Fix pandas codec for object type columns (#2080)

* Implemented JSONCodec

* Encode object type with JSONCodec for pandas

* Included more types for mlflow

* Implemented can_encode function for JSON codec

* Included test for JSON codec

* Fixed pandas codec to treat objects differently

* Included tests for mlflow metadata

* Included test for mlflow runtime

* Fixed can_encode for JSONCodec and a pandas test

* Fixed code registry but need to remove duplicated json encoding code from hf runtime

* Fixed linting issues

* Refactored duplicated code and simplified logic for can_encode

* Renamed util function

* Fixed support to encode numpy and reverted can_encode function

* Fixed bug in pandas encoding

* Reverted comment in pandas codec

* Moved env and cli tests to sequential.

* Included sleep before retrieving metrics

* Run metrics tests seq

* Simplified code by removing JSONCodec

* Included more tests for pandas codecs

* Moved PandasJsonContentType inside PandasCodec.

Author: RobertSamoilescu

mlserver/codecs/json.py

@@ -1,14 +1,17 @@
 # seperate file to side step circular dependency on the decode_str function
 
-from typing import Any, Union
 import json
+import numpy as np
+from typing import Any, List, Union
 
 try:
     import orjson
 except ImportError:
     orjson = None  # type: ignore
 
 from .string import decode_str
+from .lists import as_list
+from .utils import InputOrOutput
 
 
 # originally taken from: mlserver/rest/responses.py
@@ -57,5 +60,36 @@ def encode_to_json_bytes(v: Any) -> bytes:
 def decode_from_bytelike_json_to_dict(v: Union[bytes, str]) -> dict:
     if orjson is None:
         return json.loads(v)
-
     return orjson.loads(v)
+
+
+class JSONEncoderWithArray(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        elif isinstance(obj, np.floating):
+            return float(obj)
+        elif isinstance(obj, np.integer):
+            return int(obj)
+        else:
+            return json.JSONEncoder.default(self, obj)
+
+
+def encode_to_json(v: Any, use_bytes: bool = True) -> Union[str, bytes]:
+    enc_v = json.dumps(
+        v,
+        ensure_ascii=False,
+        allow_nan=False,
+        indent=None,
+        separators=(",", ":"),
+        cls=JSONEncoderWithArray,
+    )
+    if use_bytes:
+        enc_v = enc_v.encode("utf-8")  # type: ignore[assignment]
+    return enc_v
+
+
+def decode_json_input_or_output(input_or_output: InputOrOutput) -> List[Any]:
+    packed = input_or_output.data.root
+    unpacked = map(json.loads, as_list(packed))
+    return list(unpacked)

mlserver/codecs/pandas.py

@@ -4,7 +4,8 @@
 from typing import Optional, Any, List, Tuple
 
 from .base import RequestCodec, register_request_codec
-from .numpy import to_datatype, to_dtype, convert_nan
+from .numpy import to_dtype, convert_nan, to_datatype
+from .json import decode_json_input_or_output, encode_to_json
 from .string import encode_str, StringCodec
 from .utils import get_decoded_or_raw, InputOrOutput, inject_batch_dimension
 from .lists import ListElement
@@ -19,8 +20,12 @@
 
 
 def _to_series(input_or_output: InputOrOutput) -> pd.Series:
-    payload = get_decoded_or_raw(input_or_output)
+    parameters = input_or_output.parameters
+
+    if parameters and parameters.content_type == PandasCodec.JsonContentType:
+        return pd.Series(decode_json_input_or_output(input_or_output))
 
+    payload = get_decoded_or_raw(input_or_output)
     if Datatype(input_or_output.datatype) == Datatype.BYTES:
         # Don't convert the dtype of BYTES
         return pd.Series(payload)
@@ -43,7 +48,13 @@ def _to_response_output(series: pd.Series, use_bytes: bool = True) -> ResponseOu
 
     content_type = None
     if datatype == Datatype.BYTES:
-        data, content_type = _process_bytes(data, use_bytes)
+        processed_data, content_type = _process_bytes(data, use_bytes)
+
+        if content_type is None:
+            data = [encode_to_json(elem, use_bytes) for elem in data]
+            content_type = PandasCodec.JsonContentType
+        else:
+            data = processed_data
 
     shape = inject_batch_dimension(list(series.shape))
     parameters = None
@@ -90,6 +101,7 @@ class PandasCodec(RequestCodec):
     """
 
     ContentType = "pd"
+    JsonContentType = "pd_json"
     TypeHint = pd.DataFrame
 
     @classmethod

runtimes/huggingface/mlserver_huggingface/codecs/utils.py

@@ -5,20 +5,15 @@
 import numpy as np
 from PIL import Image, ImageChops
 from transformers.pipelines import Conversation
+from mlserver.codecs.json import JSONEncoderWithArray
 
 IMAGE_PREFIX = "data:image/"
 DEFAULT_IMAGE_FORMAT = "PNG"
 
 
-class HuggingfaceJSONEncoder(json.JSONEncoder):
+class HuggingfaceJSONEncoder(JSONEncoderWithArray):
     def default(self, obj):
-        if isinstance(obj, np.ndarray):
-            return obj.tolist()
-        elif isinstance(obj, np.floating):
-            return float(obj)
-        elif isinstance(obj, np.integer):
-            return int(obj)
-        elif isinstance(obj, Image.Image):
+        if isinstance(obj, Image.Image):
             buf = io.BytesIO()
             if not obj.format:
                 obj.format = DEFAULT_IMAGE_FORMAT
@@ -37,7 +32,7 @@ def default(self, obj):
                 "new_user_input": obj.new_user_input,
             }
         else:
-            return json.JSONEncoder.default(self, obj)
+            return super().default(obj)
 
 
 def json_encode(payload: Any, use_bytes: bool = False):

runtimes/mlflow/mlserver_mlflow/metadata.py

@@ -1,6 +1,15 @@
 from typing import Union, Tuple, List, Optional
 
-from mlflow.types.schema import Schema, ColSpec, TensorSpec, DataType
+from mlflow.types.schema import (
+    Schema,
+    ColSpec,
+    TensorSpec,
+    Array,
+    Object,
+    Map,
+    AnyType,
+    DataType,
+)
 
 from mlserver.types import MetadataTensor, Parameters
 from mlserver.types import Datatype as MDatatype
@@ -35,8 +44,10 @@
 def _get_content_type(input_spec: InputSpec) -> Tuple[MDatatype, str]:
     if isinstance(input_spec, TensorSpec):
         datatype = to_datatype(input_spec.type)
-        content_type = NumpyCodec.ContentType
-        return datatype, content_type
+        return datatype, NumpyCodec.ContentType
+
+    if isinstance(input_spec.type, (Array, Object, Map, AnyType)):
+        return MDatatype.BYTES, PandasCodec.JsonContentType
 
     # TODO: Check if new type, which may not exist
     return _MLflowToContentType[input_spec.type]

runtimes/mlflow/tests/test_metadata.py

@@ -3,7 +3,17 @@
 
 from typing import Tuple, List
 
-from mlflow.types.schema import ColSpec, TensorSpec, DataType, Schema
+from mlflow.types.schema import (
+    ColSpec,
+    TensorSpec,
+    DataType,
+    Schema,
+    Array,
+    Map,
+    Object,
+    Property,
+    AnyType,
+)
 from mlflow.pyfunc import _enforce_schema
 from mlserver.codecs import (
     NumpyCodec,
@@ -46,6 +56,30 @@
             ColSpec(name="foo", type=DataType.binary),
             (MDatatype.BYTES, Base64Codec.ContentType),
         ),
+        (
+            ColSpec(name="foo", type=Array(dtype=DataType.long)),
+            (MDatatype.BYTES, PandasCodec.JsonContentType),
+        ),
+        (
+            ColSpec(name="foo", type=Map(Array(dtype=DataType.long))),
+            (MDatatype.BYTES, PandasCodec.JsonContentType),
+        ),
+        (
+            ColSpec(
+                name="foo",
+                type=Object(
+                    properties=[
+                        Property("a", DataType.long),
+                        Property("b", DataType.string),
+                    ]
+                ),
+            ),
+            (MDatatype.BYTES, PandasCodec.JsonContentType),
+        ),
+        (
+            ColSpec(name="foo", type=AnyType()),
+            (MDatatype.BYTES, PandasCodec.JsonContentType),
+        ),
     ],
 )
 def test_get_content_type(input_spec: InputSpec, expected: Tuple[MDatatype, str]):
@@ -142,6 +176,49 @@ def test_get_shape(input_spec: InputSpec, expected: List[int]):
                 ),
             ],
         ),
+        (
+            Schema(
+                inputs=[
+                    ColSpec(type=Array(dtype=DataType.long)),
+                    ColSpec(type=Map(Array(dtype=DataType.long))),
+                    ColSpec(
+                        type=Object(
+                            properties=[
+                                Property("a", DataType.long),
+                                Property("b", DataType.string),
+                            ]
+                        )
+                    ),
+                    ColSpec(type=AnyType()),
+                ]
+            ),
+            [
+                MetadataTensor(
+                    name="input-0",
+                    datatype="BYTES",
+                    shape=[-1, 1],
+                    parameters=Parameters(content_type=PandasCodec.JsonContentType),
+                ),
+                MetadataTensor(
+                    name="input-1",
+                    datatype="BYTES",
+                    shape=[-1, 1],
+                    parameters=Parameters(content_type=PandasCodec.JsonContentType),
+                ),
+                MetadataTensor(
+                    name="input-2",
+                    datatype="BYTES",
+                    shape=[-1, 1],
+                    parameters=Parameters(content_type=PandasCodec.JsonContentType),
+                ),
+                MetadataTensor(
+                    name="input-3",
+                    datatype="BYTES",
+                    shape=[-1, 1],
+                    parameters=Parameters(content_type=PandasCodec.JsonContentType),
+                ),
+            ],
+        ),
     ],
 )
 def test_to_metadata_tensors(schema: Schema, expected: List[MetadataTensor]):
@@ -193,6 +270,54 @@ def test_to_metadata_tensors(schema: Schema, expected: List[MetadataTensor]):
                 data=[b"2021-08-24T15:01:19"],
             ),
         ),
+        (
+            ColSpec(name="foo", type=Array(dtype=DataType.long)),
+            RequestInput(
+                name="foo",
+                datatype="BYTES",
+                parameters=Parameters(content_type=PandasCodec.JsonContentType),
+                shape=[2],
+                data=[b"[1,2]", b"[3,4]"],
+            ),
+        ),
+        (
+            ColSpec(name="foo", type=Map(Array(dtype=DataType.long))),
+            RequestInput(
+                name="foo",
+                datatype="BYTES",
+                parameters=Parameters(content_type=PandasCodec.JsonContentType),
+                shape=[3],
+                data=[b'{"a":[1,2]}', b'{"b":[3,4]}', b'{"c":[5,6]}'],
+            ),
+        ),
+        (
+            ColSpec(
+                name="foo",
+                type=Object(
+                    properties=[
+                        Property("a", DataType.long),
+                        Property("b", DataType.string),
+                    ]
+                ),
+            ),
+            RequestInput(
+                name="foo",
+                datatype="BYTES",
+                parameters=Parameters(content_type=PandasCodec.JsonContentType),
+                shape=[2],
+                data=[b'{"a":1,"b":"hello"}', b'{"a":2,"b":"world"}'],
+            ),
+        ),
+        (
+            ColSpec(name="foo", type=AnyType()),
+            RequestInput(
+                name="foo",
+                datatype="BYTES",
+                parameters=Parameters(content_type=PandasCodec.JsonContentType),
+                shape=[3],
+                data=[b'"a"', b"[1,2]", b'{"b":2}'],
+            ),
+        ),
     ],
 )
 def test_content_types(tensor_spec: TensorSpec, request_input: RequestInput):
@@ -221,6 +346,25 @@ def test_content_types(tensor_spec: TensorSpec, request_input: RequestInput):
             ),
             PandasCodec.ContentType,
         ),
+        (
+            # Expect DataFrame for named column inputs
+            Schema(
+                inputs=[
+                    ColSpec(type=Array(dtype=DataType.long)),
+                    ColSpec(type=Map(Array(dtype=DataType.long))),
+                    ColSpec(
+                        type=Object(
+                            properties=[
+                                Property("a", DataType.long),
+                                Property("b", DataType.string),
+                            ]
+                        ),
+                    ),
+                    ColSpec(type=AnyType()),
+                ],
+            ),
+            PandasCodec.ContentType,
+        ),
         (
             # Expect tensor dictionary for named tensor inputs
             Schema(