Repository: onnx/sklearn-onnx

Issue:
OneHotEncoder works only with int64
Converting a scikit pipeline, I am able to use _OneHotEncoder_ only on a _int64_ column, but not on a _int32_. It should work on all column types supported by ONNX.


Discussion:

Resolving patch:
diff --git a/skl2onnx/_parse.py b/skl2onnx/_parse.py
--- a/skl2onnx/_parse.py
+++ b/skl2onnx/_parse.py
@@ -21,6 +21,8 @@ class OutlierMixin:
 from sklearn.model_selection import GridSearchCV
 from sklearn.neighbors import NearestNeighbors
 from sklearn.mixture import GaussianMixture, BayesianGaussianMixture
+from sklearn.preprocessing import OneHotEncoder
+from sklearn.pipeline import Pipeline
 from sklearn.svm import LinearSVC, NuSVC, SVC
 try:
     from sklearn.compose import ColumnTransformer
@@ -40,6 +42,11 @@ class OutlierMixin:
 from .common.utils import get_column_indices
 
 
+do_not_merge_columns = tuple(
+    filter(lambda op: op is not None,
+           [OneHotEncoder, ColumnTransformer]))
+
+
 def _fetch_input_slice(scope, inputs, column_indices):
     if not isinstance(inputs, list):
         raise TypeError("Parameter inputs must be a list.")
@@ -236,7 +243,15 @@ def _parse_sklearn_column_transformer(scope, model, inputs,
             tr_inputs = _fetch_input_slice(scope, [inputs[onnx_var]], onnx_is)
             transform_inputs.extend(tr_inputs)
 
+        merged_cols = False
         if len(transform_inputs) > 1:
+            if isinstance(op, Pipeline):
+                if not isinstance(op.steps[0][1], do_not_merge_columns):
+                    merged_cols = True
+            elif not isinstance(op, do_not_merge_columns):
+                merged_cols = True
+
+        if merged_cols:
             # Many ONNX operators expect one input vector,
             # the default behaviour is to merge columns.
             ty = transform_inputs[0].type.__class__([None, None])
diff --git a/skl2onnx/operator_converters/one_hot_encoder.py b/skl2onnx/operator_converters/one_hot_encoder.py
--- a/skl2onnx/operator_converters/one_hot_encoder.py
+++ b/skl2onnx/operator_converters/one_hot_encoder.py
@@ -6,75 +6,116 @@
 
 import numpy as np
 from ..common._apply_operation import apply_cast, apply_concat, apply_reshape
-from ..common.data_types import Int64TensorType, StringTensorType
+from ..common.data_types import (
+    Int64TensorType, StringTensorType, Int32TensorType,
+    FloatTensorType, DoubleTensorType
+)
 from ..common._registration import register_converter
 from ..proto import onnx_proto
 
 
 def convert_sklearn_one_hot_encoder(scope, operator, container):
+    """
+    Converts *OneHotEncoder* into ONNX.
+    It supports multiple inputs of types
+    string or int64.
+    """
     ohe_op = operator.raw_operator
-    result, categories_len = [], 0
-    concatenated_input_name = operator.inputs[0].full_name
-    concat_result_name = scope.get_unique_variable_name('concat_result')
 
     if len(operator.inputs) > 1:
-        concatenated_input_name = scope.get_unique_variable_name(
-            'concatenated_input')
-        if all(isinstance(inp.type, type(operator.inputs[0].type))
-               for inp in operator.inputs):
-            input_names = list(map(lambda x: x.full_name, operator.inputs))
-        else:
-            input_names = []
-            for inp in operator.inputs:
-                if isinstance(inp.type, Int64TensorType):
-                    input_names.append(scope.get_unique_variable_name(
-                        'cast_input'))
-                    apply_cast(scope, inp.full_name, input_names[-1],
-                               container, to=onnx_proto.TensorProto.STRING)
-                elif isinstance(inp.type, StringTensorType):
-                    input_names.append(inp.full_name)
-                else:
-                    raise NotImplementedError(
-                        "{} input datatype not yet supported. "
-                        "You may raise an issue at "
-                        "https://github.com/onnx/sklearn-onnx/issues"
-                        "".format(type(inp.type)))
-
-        apply_concat(scope, input_names,
-                     concatenated_input_name, container, axis=1)
-    for index, categories in enumerate(ohe_op.categories_):
+        all_shapes = [inp.type.shape[1] for inp in operator.inputs]
+        if any(map(lambda x: not isinstance(x, int) or x < 1, all_shapes)):
+            raise RuntimeError(
+                "Shapes must be known when OneHotEncoder is converted. "
+                "There are {} inputs with the following number of columns "
+                "{}.".format(len(operator.inputs), all_shapes))
+        total = sum(all_shapes)
+        if total != len(ohe_op.categories_):
+            raise RuntimeError(
+                "Mismatch between the number of sets of categories {} and "
+                "the total number of inputs columns {}.".format(
+                    len(ohe_op.categories_), total))
+
+        enum_cats = []
+        index_inputs = 0
+
+        for index, cats in enumerate(ohe_op.categories_):
+            while sum(all_shapes[:index_inputs+1]) <= index:
+                index_inputs += 1
+            index_in_input = index - sum(all_shapes[:index_inputs])
+
+            inp = operator.inputs[index_inputs]
+            if not isinstance(
+                    inp.type,
+                    (Int64TensorType, StringTensorType, Int32TensorType,
+                     FloatTensorType, DoubleTensorType)):
+                raise NotImplementedError(
+                    "{} input datatype not yet supported. "
+                    "You may raise an issue at "
+                    "https://github.com/onnx/sklearn-onnx/issues"
+                    "".format(type(inp.type)))
+
+            if all_shapes[index_inputs] == 1:
+                assert index_in_input == 0
+                afeat = False
+            else:
+                afeat = True
+            enum_cats.append(
+                (afeat, index_in_input, inp.full_name, cats, inp.type))
+    else:
+        inp = operator.inputs[0]
+        enum_cats = [(True, i, inp.full_name, cats, inp.type)
+                     for i, cats in enumerate(ohe_op.categories_)]
+
+    result, categories_len = [], 0
+    for index, enum_c in enumerate(enum_cats):
+        afeat, index_in, name, categories, inp_type = enum_c
+        if len(categories) == 0:
+            continue
+        if afeat:
+            index_name = scope.get_unique_variable_name(
+                name + str(index_in))
+            container.add_initializer(
+                index_name, onnx_proto.TensorProto.INT64, [], [index_in])
+            out_name = scope.get_unique_variable_name(
+                name + str(index_in))
+            container.add_node(
+                'ArrayFeatureExtractor', [name, index_name],
+                out_name, op_domain='ai.onnx.ml',
+                name=scope.get_unique_operator_name('ArrayFeatureExtractor'))
+            name = out_name
+
         attrs = {'name': scope.get_unique_operator_name('OneHotEncoder')}
         attrs['zeros'] = 1 if ohe_op.handle_unknown == 'ignore' else 0
         if hasattr(ohe_op, 'drop_idx_') and ohe_op.drop_idx_ is not None:
             categories = (categories[np.arange(len(categories)) !=
                                      ohe_op.drop_idx_[index]])
-        if len(categories) > 0:
-            if (np.issubdtype(categories.dtype, np.floating)
-                    or np.issubdtype(categories.dtype, np.signedinteger)):
-                attrs['cats_int64s'] = categories.astype(np.int64)
-            else:
-                attrs['cats_strings'] = np.array(
-                    [str(s).encode('utf-8') for s in categories])
 
-            index_name = scope.get_unique_variable_name('index')
-            feature_column_name = scope.get_unique_variable_name(
-                'feature_column')
-            result.append(scope.get_unique_variable_name('ohe_output'))
+        if isinstance(inp_type, (Int64TensorType, Int32TensorType)):
+            attrs['cats_int64s'] = categories.astype(np.int64)
+        else:
+            attrs['cats_strings'] = np.array(
+                [str(s).encode('utf-8') for s in categories])
 
-            container.add_initializer(
-                index_name, onnx_proto.TensorProto.INT64, [], [index])
+        ohe_output = scope.get_unique_variable_name(name + 'out')
+        result.append(ohe_output)
 
-            container.add_node(
-                'ArrayFeatureExtractor',
-                [concatenated_input_name, index_name],
-                feature_column_name, op_domain='ai.onnx.ml',
-                name=scope.get_unique_operator_name('ArrayFeatureExtractor'))
+        if 'cats_int64s' in attrs:
+            # Let's cast this input in int64.
+            cast_feature = scope.get_unique_variable_name(name + 'cast')
+            apply_cast(scope, name, cast_feature, container,
+                       to=onnx_proto.TensorProto.INT64)
+            name = cast_feature
+
+        container.add_node('OneHotEncoder', name,
+                           ohe_output, op_domain='ai.onnx.ml',
+                           **attrs)
+
+        categories_len += len(categories)
+
+    concat_result_name = scope.get_unique_variable_name('concat_result')
+    apply_concat(scope, result, concat_result_name, container, axis=2)
 
-            container.add_node('OneHotEncoder', feature_column_name,
-                               result[-1], op_domain='ai.onnx.ml', **attrs)
-            categories_len += len(categories)
-    apply_concat(scope, result,
-                 concat_result_name, container, axis=2)
     reshape_input = concat_result_name
     if np.issubdtype(ohe_op.dtype, np.signedinteger):
         reshape_input = scope.get_unique_variable_name('cast')
diff --git a/skl2onnx/shape_calculators/concat.py b/skl2onnx/shape_calculators/concat.py
--- a/skl2onnx/shape_calculators/concat.py
+++ b/skl2onnx/shape_calculators/concat.py
@@ -35,10 +35,12 @@ def calculate_sklearn_concat(operator):
             seen_types.append(nt)
         elif nt != seen_types[0]:
             inps = "\n".join(str(v) for v in operator.inputs)
-            raise RuntimeError("Columns must have the same type. "
-                               "C++ backends do not support mixed types. "
-                               "Inputs:\n"
-                               + inps)
+            outs = "\n".join(str(v) for v in operator.outputs)
+            raise RuntimeError(
+                "Columns must have the same type. "
+                "C++ backends do not support mixed types.\n"
+                "Inputs:\n{}\nOutputs:\n{}".format(
+                    inps, outs))
     operator.outputs[0].type.shape = [N, C]
 
 


Test patch:
diff --git a/tests/test_onnx_helper.py b/tests/test_onnx_helper.py
--- a/tests/test_onnx_helper.py
+++ b/tests/test_onnx_helper.py
@@ -59,8 +59,10 @@ def test_onnx_helper_load_save(self):
         reason="OneHotEncoder did not have categories_ before 0.20",
     )
     def test_onnx_helper_load_save_init(self):
-        model = make_pipeline(Binarizer(), OneHotEncoder(sparse=False),
-                              StandardScaler())
+        model = make_pipeline(
+            Binarizer(),
+            OneHotEncoder(sparse=False, handle_unknown='ignore'),
+            StandardScaler())
         X = numpy.array([[0.1, 1.1], [0.2, 2.2], [0.4, 2.2], [0.2, 2.4]])
         model.fit(X)
         model_onnx = convert_sklearn(model, "pipe3",
diff --git a/tests/test_sklearn_one_hot_encoder_converter.py b/tests/test_sklearn_one_hot_encoder_converter.py
--- a/tests/test_sklearn_one_hot_encoder_converter.py
+++ b/tests/test_sklearn_one_hot_encoder_converter.py
@@ -5,8 +5,11 @@
 from onnxruntime import __version__ as ort_version
 from sklearn import __version__ as sklearn_version
 from sklearn.preprocessing import OneHotEncoder
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import RobustScaler
 from skl2onnx import convert_sklearn
 from skl2onnx.common.data_types import (
+    Int32TensorType,
     Int64TensorType,
     StringTensorType,
 )
@@ -26,6 +29,8 @@ def one_hot_encoder_supports_drop():
 
 
 class TestSklearnOneHotEncoderConverter(unittest.TestCase):
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
     @unittest.skipIf(
         not one_hot_encoder_supports_string(),
         reason="OneHotEncoder did not have categories_ before 0.20",
@@ -48,6 +53,53 @@ def test_model_one_hot_encoder(self):
             basename="SklearnOneHotEncoderInt64-SkipDim1",
         )
 
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
+    @unittest.skipIf(
+        not one_hot_encoder_supports_string(),
+        reason="OneHotEncoder did not have categories_ before 0.20",
+    )
+    def test_model_one_hot_encoder_int32(self):
+        model = OneHotEncoder(categories='auto')
+        data = numpy.array([[1, 2, 3], [4, 3, 0], [0, 1, 4], [0, 5, 6]],
+                           dtype=numpy.int32)
+        model.fit(data)
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn one-hot encoder",
+            [("input", Int32TensorType([None, 3]))],
+        )
+        str_model_onnx = str(model_onnx)
+        assert "int64_data" in str_model_onnx
+        self.assertTrue(model_onnx is not None)
+        dump_data_and_model(
+            data, model, model_onnx,
+            basename="SklearnOneHotEncoderInt32-SkipDim1")
+
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
+    @unittest.skipIf(
+        not one_hot_encoder_supports_string(),
+        reason="OneHotEncoder did not have categories_ before 0.20",
+    )
+    def test_model_one_hot_encoder_int32_scaler(self):
+        model = make_pipeline(OneHotEncoder(categories='auto', sparse=False),
+                              RobustScaler())
+        data = numpy.array([[1, 2, 3], [4, 3, 0], [0, 1, 4], [0, 5, 6]],
+                           dtype=numpy.int32)
+        model.fit(data)
+        model_onnx = convert_sklearn(
+            model,
+            "scikit-learn one-hot encoder",
+            [("input", Int32TensorType([None, 3]))],
+        )
+        str_model_onnx = str(model_onnx)
+        assert "int64_data" in str_model_onnx
+        self.assertTrue(model_onnx is not None)
+        dump_data_and_model(
+            data, model, model_onnx,
+            basename="SklearnOneHotEncoderInt32Scaler-SkipDim1")
+
     @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
                      reason="issues with shapes")
     @unittest.skipIf(
@@ -70,18 +122,15 @@ def test_one_hot_encoder_mixed_string_int_drop(self):
             ("input1", StringTensorType([None, 2])),
             ("input2", Int64TensorType([None, 1])),
         ]
-        model_onnx = convert_sklearn(model,
-                                     "one-hot encoder",
-                                     inputs)
+        model_onnx = convert_sklearn(
+            model, "one-hot encoder", inputs)
         self.assertTrue(model_onnx is not None)
         dump_data_and_model(
-            test,
-            model,
-            model_onnx,
-            basename="SklearnOneHotEncoderMixedStringIntDrop",
-            verbose=False,
-        )
+            test, model, model_onnx, verbose=False,
+            basename="SklearnOneHotEncoderMixedStringIntDrop")
 
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
     @unittest.skipIf(
         not one_hot_encoder_supports_string(),
         reason="OneHotEncoder does not support strings in 0.19",
@@ -101,6 +150,8 @@ def test_one_hot_encoder_onecat(self):
             basename="SklearnOneHotEncoderOneStringCat",
         )
 
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
     @unittest.skipIf(
         not one_hot_encoder_supports_string(),
         reason="OneHotEncoder does not support strings in 0.19",
@@ -146,6 +197,8 @@ def test_one_hot_encoder_string_drop_first(self):
             basename="SklearnOneHotEncoderStringDropFirst",
         )
 
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
     @unittest.skipIf(
         not one_hot_encoder_supports_string(),
         reason="OneHotEncoder does not support this in 0.19",
@@ -171,6 +224,8 @@ def test_model_one_hot_encoder_list_sparse(self):
             basename="SklearnOneHotEncoderCatSparse-SkipDim1",
         )
 
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
     @unittest.skipIf(
         not one_hot_encoder_supports_string(),
         reason="OneHotEncoder does not support this in 0.19",
@@ -196,6 +251,8 @@ def test_model_one_hot_encoder_list_dense(self):
             basename="SklearnOneHotEncoderCatDense-SkipDim1",
         )
 
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
     @unittest.skipIf(
         not one_hot_encoder_supports_drop(),
         reason="OneHotEncoder does not support drop in scikit versions < 0.21",
@@ -228,6 +285,8 @@ def test_one_hot_encoder_int_drop(self):
             verbose=False,
         )
 
+    @unittest.skipIf(StrictVersion(ort_version) <= StrictVersion("0.4.0"),
+                     reason="issues with shapes")
     @unittest.skipIf(
         not one_hot_encoder_supports_drop(),
         reason="OneHotEncoder does not support drop in scikit versions < 0.21",
diff --git a/tests/test_sklearn_pipeline.py b/tests/test_sklearn_pipeline.py
--- a/tests/test_sklearn_pipeline.py
+++ b/tests/test_sklearn_pipeline.py
@@ -289,9 +289,8 @@ def test_pipeline_column_transformer_titanic(self):
         for cat in ["embarked", "sex", "pclass"]:
             X[cat].fillna("missing", inplace=True)
 
-        X_train, X_test, y_train, y_test = train_test_split(X,
-                                                            y,
-                                                            test_size=0.2)
+        X_train, X_test, y_train, y_test = train_test_split(
+            X, y, test_size=0.2)
 
         numeric_features = ["age", "fare"]
         numeric_transformer = Pipeline(steps=[
@@ -315,7 +314,7 @@ def test_pipeline_column_transformer_titanic(self):
 
         clf = Pipeline(steps=[
             ("preprocessor", preprocessor),
-            ("classifier", LogisticRegression(solver="lbfgs")),
+            # ("classifier", LogisticRegression(solver="lbfgs")),
         ])
 
         # inputs
@@ -325,7 +324,7 @@ def convert_dataframe_schema(df, drop=None):
             for k, v in zip(df.columns, df.dtypes):
                 if drop is not None and k in drop:
                     continue
-                if v == "int64":
+                if v == 'int64':
                     t = Int64TensorType([None, 1])
                 elif v == "float64":
                     t = FloatTensorType([None, 1])
@@ -344,24 +343,33 @@ def convert_dataframe_schema(df, drop=None):
             "home.dest",
             "boat",
         }
+
+        X_train = X_train.copy()
+        X_test = X_test.copy()
+        X_train['pclass'] = X_train['pclass'].astype(numpy.int64)
+        X_test['pclass'] = X_test['pclass'].astype(numpy.int64)
         X_train = X_train.drop(to_drop, axis=1)
         X_test = X_test.drop(to_drop, axis=1)
+
         clf.fit(X_train, y_train)
-        X_train["pclass"] = X_train["pclass"].astype(str)
-        X_test["pclass"] = X_test["pclass"].astype(str)
         inputs = convert_dataframe_schema(X_train, to_drop)
         model_onnx = convert_sklearn(clf, "pipeline_titanic", inputs)
 
-        dump_data_and_model(
-            X_test[:5],
-            clf,
-            model_onnx,
-            basename="SklearnPipelineColumnTransformerPipelinerTitanic-DF",
-            allow_failure="StrictVersion(onnx.__version__)"
-                          " < StrictVersion('1.3') or "
-                          "StrictVersion(onnxruntime.__version__)"
-                          " <= StrictVersion('0.2.1')",
-        )
+        data = X_test[:5]
+        pred = clf.transform(data)
+        data_types = {
+            'pclass': numpy.int64,
+            'age': numpy.float32,
+            'sex': numpy.str,
+            'fare': numpy.float32,
+            'embarked': numpy.str,
+        }
+        inputs = {k: data[k].values.astype(data_types[k]).reshape(-1, 1)
+                  for k in data.columns}
+        sess = InferenceSession(model_onnx.SerializeToString())
+        run = sess.run(None, inputs)
+        got = run[-1]
+        assert_almost_equal(pred, got, decimal=5)
 
     @unittest.skipIf(
         ColumnTransformer is None,