Repository: onnx/sklearn-onnx
Pull request #1123 (2024-08-07 23:39:13)

Support fill_value for SimpleImputer with string data
Adds support for setting the `fill_value` of a SimpleImputer using the constant strategy on string data. This is done by moving the existing guard statement that limits string values of `fill_value` to be only the value "nan".

Based on some testing, likely this guard statement can be removed completely due to sci-kit learn doing it's own casting checks when fitting data to the SimpleImputer (see <https://github.com/scikit-learn/scikit-learn/blob/5600faab452b0adb1f4090219012aa3ae6208c66/sklearn/impute/_base.py#L375-L398>). I moved the statement in this PR, in case there was some edge case I missed while doing my own checks and there is some numeric SimpleImputer that allows `fill_value=="nan"`. But let me know if you rather I remove it outright.

Addresses #1122 




bodnarbm: @xadupre Sorry for the direct mention, but looking at the git history, you might be the best person to ask. Is there anything I can do to help with review on this pull request?
xadupre: Sorry for the delay, I was off. Your PR looks good.
bodnarbm: No worries. 

I will try to clean up the CodeQL reported issues in the PR tomorrow. The CI failures look to be related to other unchanged areas of code, but I am also willing to dive into those too as I have time this week.
bodnarbm: @xadupre I went ahead and re-ran the CI jobs on my own fork. It looks like the latest merge from main that I pulled into this branch addressed the remaining test flakiness. I believe the workflows should succeed now for this branch. Thank you for fixing those tests / dependencies in the other PRs you did!

Patch:
diff --git a/skl2onnx/operator_converters/imputer_op.py b/skl2onnx/operator_converters/imputer_op.py
index 8b466f6..fc227af 100644
--- a/skl2onnx/operator_converters/imputer_op.py
+++ b/skl2onnx/operator_converters/imputer_op.py
@@ -17,14 +17,6 @@ def convert_sklearn_imputer(
     op_type = "Imputer"
     attrs = {"name": scope.get_unique_operator_name(op_type)}
     op = operator.raw_operator
-    if (
-        hasattr(op, "fill_value")
-        and isinstance(op.fill_value, str)
-        and op.fill_value.lower() != "nan"
-    ):
-        raise RuntimeError(
-            "Imputer cannot fill missing values with a string '%s'." % op.fill_value
-        )
     if not hasattr(op, "statistics_"):
         raise RuntimeError("Member statistics_ is not present, was the model fitted?")
 
@@ -86,6 +78,14 @@ def convert_sklearn_imputer(
 
         apply_concat(scope, names, operator.outputs[0].full_name, container, axis=1)
     else:
+        if (
+            hasattr(op, "fill_value")
+            and isinstance(op.fill_value, str)
+            and op.fill_value.lower() != "nan"
+        ):
+            raise RuntimeError(
+                "Imputer cannot fill missing values with a string '%s'." % op.fill_value
+            )
         if isinstance(operator.inputs[0].type, Int64TensorType):
             attrs["imputed_value_int64s"] = op.statistics_.astype(np.int64)
             use_int = True


Test patch:
diff --git a/tests/test_sklearn_imputer_converter.py b/tests/test_sklearn_imputer_converter.py
index 331698d..5e62aff 100644
--- a/tests/test_sklearn_imputer_converter.py
+++ b/tests/test_sklearn_imputer_converter.py
@@ -51,6 +51,15 @@ class TestSklearnImputerConverter(unittest.TestCase):
         exp = model.transform(data)
         assert_almost_equal(res, exp)
 
+    def _check_outputs_floats(self, model, model_onnx, data):
+        sess = InferenceSession(
+            model_onnx.SerializeToString(), providers=["CPUExecutionProvider"]
+        )
+        idata = {"input": np.array(data).astype(np.float32)}
+        res = sess.run(None, idata)[0]
+        exp = model.transform(data)
+        assert_almost_equal(res, exp)
+
     def _check_outputs_strings(self, model, model_onnx, data, verbose=0):
         idata = {"input": np.array(data).astype(np.str_)}
         sess = InferenceSession(
@@ -206,6 +215,140 @@ class TestSklearnImputerConverter(unittest.TestCase):
         self.assertEqual(len(model_onnx.graph.output), 1)
         self._check_outputs_strings(model, model_onnx, data)
 
+    @unittest.skipIf(SimpleImputer is None, reason="SimpleImputer changed in 0.20")
+    def test_simple_imputer_float_constant_default_fill_value(self):
+        model = SimpleImputer(strategy="constant")
+        data = [[1, 2], [np.nan, 3], [7, 6]]
+        model.fit(data)
+
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn simple imputer",
+            [("input", FloatTensorType([None, 2]))],
+            target_opset=TARGET_OPSET,
+        )
+        self.assertIsNotNone(model_onnx.graph.node)
+
+        # should contain only node
+        self.assertEqual(len(model_onnx.graph.node), 1)
+
+        # last node should contain the Imputer
+        outputs = model_onnx.graph.output
+        self.assertEqual(len(outputs), 1)
+        self._check_outputs_floats(model, model_onnx, data)
+
+    @unittest.skipIf(SimpleImputer is None, reason="SimpleImputer changed in 0.20")
+    def test_simple_imputer_float_constant_provided_fill_value(self):
+        model = SimpleImputer(strategy="constant", fill_value=99.0)
+        data = [[1, 2], [np.nan, 3], [7, 6]]
+        model.fit(data)
+
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn simple imputer",
+            [("input", FloatTensorType([None, 2]))],
+            target_opset=TARGET_OPSET,
+        )
+        self.assertIsNotNone(model_onnx.graph.node)
+
+        # should contain only node
+        self.assertEqual(len(model_onnx.graph.node), 1)
+
+        # last node should contain the Imputer
+        outputs = model_onnx.graph.output
+        self.assertEqual(len(outputs), 1)
+        self._check_outputs_floats(model, model_onnx, data)
+
+    @unittest.skipIf(SimpleImputer is None, reason="SimpleImputer changed in 0.20")
+    def test_simple_imputer_int_constant_default_fill_value(self):
+        model = SimpleImputer(strategy="constant")
+        data = [[1, 2], [np.nan, 3], [7, 6], [8, np.nan]]
+        model.fit(data)
+
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn simple imputer",
+            [("input", Int64TensorType([None, 2]))],
+            target_opset=TARGET_OPSET,
+        )
+        self.assertIsNotNone(model_onnx.graph.node)
+
+        # should contain only node
+        self.assertEqual(len(model_onnx.graph.node), 1)
+
+        # last node should contain the Imputer
+        outputs = model_onnx.graph.output
+        self.assertEqual(len(outputs), 1)
+        self._check_outputs_ints(model, model_onnx, data)
+
+    @unittest.skipIf(SimpleImputer is None, reason="SimpleImputer changed in 0.20")
+    def test_simple_imputer_int_constant_provided_fill_value(self):
+        model = SimpleImputer(strategy="constant", fill_value=99)
+        data = [[1, 2], [np.nan, 3], [7, 6], [8, np.nan]]
+        model.fit(data)
+
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn simple imputer",
+            [("input", Int64TensorType([None, 2]))],
+            target_opset=TARGET_OPSET,
+        )
+        self.assertIsNotNone(model_onnx.graph.node)
+
+        # should contain only node
+        self.assertEqual(len(model_onnx.graph.node), 1)
+
+        # last node should contain the Imputer
+        outputs = model_onnx.graph.output
+        self.assertEqual(len(outputs), 1)
+        self._check_outputs_ints(model, model_onnx, data)
+
+    @unittest.skipIf(SimpleImputer is None, reason="SimpleImputer changed in 0.20")
+    @unittest.skipIf(
+        pv.Version(skl_ver) < pv.Version("0.24"),
+        reason="SimpleImputer does not support strings",
+    )
+    def test_simple_imputer_string_inputs_constant_provided_fill_value(self):
+        model = SimpleImputer(
+            strategy="constant", missing_values="", fill_value="missing"
+        )
+        data = pd.DataFrame(
+            [["s1", "s2"], ["s1", "s2"], ["", "s3"], ["s7", "s6"], ["s8", ""]]
+        )
+        model.fit(data)
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn simple imputer",
+            [("input", StringTensorType([None, 2]))],
+            target_opset=TARGET_OPSET,
+        )
+        self.assertIn("ai.onnx.ml", str(model_onnx))
+        self.assertIsNotNone(model_onnx.graph.node)
+        self.assertEqual(len(model_onnx.graph.output), 1)
+        self._check_outputs_strings(model, model_onnx, data)
+
+    @unittest.skipIf(SimpleImputer is None, reason="SimpleImputer changed in 0.20")
+    @unittest.skipIf(
+        pv.Version(skl_ver) < pv.Version("0.24"),
+        reason="SimpleImputer does not support strings",
+    )
+    def test_simple_imputer_string_inputs_constant_default_fill_value(self):
+        model = SimpleImputer(strategy="constant", missing_values="")
+        data = pd.DataFrame(
+            [["s1", "s2"], ["s1", "s2"], ["", "s3"], ["s7", "s6"], ["s8", ""]]
+        )
+        model.fit(data)
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn simple imputer",
+            [("input", StringTensorType([None, 2]))],
+            target_opset=TARGET_OPSET,
+        )
+        self.assertIn("ai.onnx.ml", str(model_onnx))
+        self.assertIsNotNone(model_onnx.graph.node)
+        self.assertEqual(len(model_onnx.graph.output), 1)
+        self._check_outputs_strings(model, model_onnx, data)
+
 
 if __name__ == "__main__":
     unittest.main()