Make example input and PyFuncInput support csc csr matrix

mlflow/models/utils.py

@@ -9,8 +9,9 @@
 from mlflow.models import Model
 from mlflow.types.utils import TensorsNotSupportedException
 from mlflow.utils.proto_json_utils import NumpyEncoder, _dataframe_from_json, parse_tf_serving_input
+from scipy.sparse import csr_matrix, csc_matrix
 
-ModelInputExample = Union[pd.DataFrame, np.ndarray, dict, list]
+ModelInputExample = Union[pd.DataFrame, np.ndarray, dict, list, csr_matrix, csc_matrix]
 
 
 class _Example(object):
@@ -50,25 +51,33 @@ class _Example(object):
           encoded strings.
         - numpy types: Numpy types are converted to the corresponding python types or their closest
           equivalent.
+        - csc/csr matric: similar to 2 dims numpy array, csc/csr matric are converted to
+          corresponding python types or their closest equivalent.
     """
 
     def __init__(self, input_example: ModelInputExample):
         def _is_scalar(x):
             return np.isscalar(x) or x is None
 
         def _is_tensor(x):
-            return isinstance(x, np.ndarray) or (
-                isinstance(x, dict) and all([isinstance(ary, np.ndarray) for ary in x.values()])
+            return (
+                isinstance(x, np.ndarray)
+                or (
+                    isinstance(x, dict) and all([isinstance(ary, np.ndarray) for ary in x.values()])
+                )
+                or isinstance(x, (csr_matrix, csc_matrix))
             )
 
-        def _handle_tensor_input(input_tensor: Union[np.ndarray, dict]):
+        def _handle_tensor_input(input_tensor: Union[np.ndarray, dict, csr_matrix, csc_matrix]):
             if isinstance(input_tensor, dict):
                 result = {}
                 for name in input_tensor.keys():
                     result[name] = input_tensor[name].tolist()
                 return {"inputs": result}
-            else:
+            elif isinstance(input_tensor, np.ndarray):
                 return {"inputs": input_tensor.tolist()}
+            else:
+                return {"inputs": input_tensor.toarray().tolist()}
 
         def _handle_dataframe_input(input_ex):
             if isinstance(input_ex, dict):

mlflow/types/utils.py

@@ -30,8 +30,10 @@ def _get_tensor_shape(data: np.ndarray, variable_dimension: Optional[int] = 0) -
     :param variable_dimension: An optional integer representing a variable dimension.
     :return: tuple : Shape of the inputted data (including a variable dimension)
     """
-    if not isinstance(data, np.ndarray):
-        raise TypeError("Expected numpy.ndarray, got '{}'.".format(type(data)))
+    from scipy.sparse import csr_matrix, csc_matrix
+
+    if not isinstance(data, (np.ndarray, csr_matrix, csc_matrix)):
+        raise TypeError("Expected numpy.ndarray or csc/csr matrix, got '{}'.".format(type(data)))
     variable_input_data_shape = data.shape
     if variable_dimension is not None:
         try:
@@ -90,6 +92,7 @@ def _infer_schema(data: Any) -> Schema:
       - dictionary of { name -> numpy.ndarray}
       - numpy.ndarray
       - pyspark.sql.DataFrame
+      - csc/csr matrix
 
     The element types should be mappable to one of :py:class:`mlflow.models.signature.DataType` for
     dataframes and to one of numpy types for tensors.
@@ -98,6 +101,8 @@ def _infer_schema(data: Any) -> Schema:
 
     :return: Schema
     """
+    from scipy.sparse import csr_matrix, csc_matrix
+
     if isinstance(data, dict):
         res = []
         for name in data.keys():
@@ -122,6 +127,10 @@ def _infer_schema(data: Any) -> Schema:
         schema = Schema(
             [TensorSpec(type=clean_tensor_type(data.dtype), shape=_get_tensor_shape(data))]
         )
+    elif isinstance(data, (csc_matrix, csr_matrix)):
+        schema = Schema(
+            [TensorSpec(type=clean_tensor_type(data.data.dtype), shape=_get_tensor_shape(data))]
+        )
     elif _is_spark_df(data):
         schema = Schema(
             [

setup.py

@@ -71,6 +71,7 @@ def package_files(directory):
     "Flask",
     "gunicorn; platform_system != 'Windows'",
     "numpy",
+    "scipy",
     "pandas",
     "prometheus-flask-exporter",
     "querystring_parser",

tests/models/test_model_input_examples.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from scipy.sparse import csr_matrix, csc_matrix
 
 from mlflow.models.signature import infer_signature
 from mlflow.models.utils import _Example, _read_tensor_input_from_json
@@ -44,6 +45,14 @@ def dict_of_ndarrays():
     }
 
 
+@pytest.fixture
+def dict_of_sparse_matrix():
+    return {
+        "csc": csc_matrix(np.arange(0, 12, 0.5).reshape(3, 8)),
+        "csr": csr_matrix(np.arange(0, 12, 0.5).reshape(3, 8)),
+    }
+
+
 def test_input_examples(pandas_df_with_all_types, dict_of_ndarrays):
     sig = infer_signature(pandas_df_with_all_types)
     # test setting example with data frame with all supported data types
@@ -117,3 +126,14 @@ def test_input_examples(pandas_df_with_all_types, dict_of_ndarrays):
         filename = x.info["artifact_path"]
         parsed_df = _dataframe_from_json(tmp.path(filename))
         assert example == parsed_df.to_dict(orient="records")[0]
+
+
+def test_sparse_matrix_input_examples(dict_of_sparse_matrix):
+    for col in dict_of_sparse_matrix:
+        input_example = dict_of_sparse_matrix[col]
+        with TempDir() as tmp:
+            example = _Example(input_example)
+            example.save(tmp.path())
+            filename = example.info["artifact_path"]
+            parsed_ary = _read_tensor_input_from_json(tmp.path(filename))
+            assert np.array_equal(parsed_ary, input_example.toarray())

tests/sklearn/test_sklearn_autolog.py

@@ -16,6 +16,7 @@
 import sklearn.pipeline
 import sklearn.model_selection
 from scipy.stats import uniform
+from scipy.sparse import csr_matrix, csc_matrix
 
 from mlflow.exceptions import MlflowException
 from mlflow.models import Model
@@ -849,13 +850,16 @@ def test_parameter_search_handles_large_volume_of_metric_outputs():
     assert len(child_run.data.metrics) >= metrics_size
 
 
-@pytest.mark.parametrize("data_type", [pd.DataFrame, np.array])
+@pytest.mark.parametrize("data_type", [pd.DataFrame, np.array, csr_matrix, csc_matrix])
 def test_autolog_logs_signature_and_input_example(data_type):
     mlflow.sklearn.autolog(log_input_examples=True, log_model_signatures=True)
 
     X, y = get_iris()
     X = data_type(X)
-    y = data_type(y)
+    if data_type in [csr_matrix, csc_matrix]:
+        y = np.array(y)
+    else:
+        y = data_type(y)
     model = sklearn.linear_model.LinearRegression()
 
     with mlflow.start_run() as run:

tests/types/test_schema.py

@@ -3,6 +3,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from scipy.sparse import csr_matrix, csc_matrix
 
 from mlflow.exceptions import MlflowException
 from mlflow.pyfunc import _enforce_tensor_spec
@@ -251,6 +252,22 @@ def test_get_tensor_shape(dict_of_ndarrays):
         _infer_schema({"x": 1})
 
 
+@pytest.fixture
+def dict_of_sparse_matrix():
+    return {
+        "csc": csc_matrix(np.arange(0, 12, 0.5).reshape(3, 8)),
+        "csr": csr_matrix(np.arange(0, 12, 0.5).reshape(3, 8)),
+    }
+
+
+def test_get_sparse_matrix_data_type_and_shape(dict_of_sparse_matrix):
+    for col in dict_of_sparse_matrix:
+        sparse_matrix = dict_of_sparse_matrix[col]
+        schema = _infer_schema(sparse_matrix)
+        schema.numpy_types() == ["float64"]
+        _get_tensor_shape(sparse_matrix) == (3, 8)
+
+
 def test_schema_inference_on_dictionary(dict_of_ndarrays):
     # test dictionary
     schema = _infer_schema(dict_of_ndarrays)