refactor: use sklearn OrdinalEncoder as source of truth for DataTransformer categorical encoding (#1564) (#1569)

* refactor: use sklearn OrdinalEncoder as source of truth for DataTransformer categorical encoding (#1564)

* address review: dict lookup for encoder columns, accurate transform comment

* fix: preserve categorical encoding order and compatibility

---------

Co-authored-by: Li Jiang <bnujli@gmail.com>

Author: immu4989

flaml/automl/data.py

@@ -315,7 +315,8 @@ def fit_transform(self, X: Union[DataFrame, np.ndarray], y, task: Union[str, "Ta
                     elif X[column].dtype.name == "category":
                         current_categories = X[column].cat.categories
                         if "__NAN__" not in current_categories:
-                            X[column] = X[column].cat.add_categories("__NAN__").fillna("__NAN__")
+                            X[column] = X[column].cat.add_categories("__NAN__")
+                        X[column] = X[column].fillna("__NAN__")
                         cat_columns.append(column)
                     else:
                         X[column] = X[column].fillna("__NAN__")
@@ -351,17 +352,16 @@ def fit_transform(self, X: Union[DataFrame, np.ndarray], y, task: Union[str, "Ta
                 X.insert(0, TS_TIMESTAMP_COL, ds_col)
             if cat_columns:
                 X[cat_columns] = X[cat_columns].astype("category")
-                # Pin the per-column category list seen at fit time so
-                # `transform()` produces the same integer codes for the same
-                # values regardless of what is passed at predict time (see
-                # issue #1101). "__NAN__" is reserved as the sentinel slot
-                # used for values unseen at fit time.
-                self._cat_categories = {}
-                for col in cat_columns:
-                    cats = list(X[col].cat.categories)
-                    if "__NAN__" not in cats:
-                        cats.append("__NAN__")
-                    self._cat_categories[col] = cats
+                from sklearn.preprocessing import OrdinalEncoder
+
+                categories = [X[column].cat.categories.to_numpy(dtype=object) for column in cat_columns]
+                self._ordinal_encoder = OrdinalEncoder(
+                    categories=[np.arange(len(values)) for values in categories],
+                    handle_unknown="use_encoded_value",
+                    unknown_value=-1,
+                )
+                self._ordinal_encoder.fit(np.column_stack([X[column].cat.codes for column in cat_columns]))
+                self._ordinal_encoder.categories_ = categories
             if num_columns:
                 X_num = X[num_columns]
                 try:
@@ -459,16 +459,44 @@ def transform(self, X: Union[DataFrame, np.array]):
                 elif X[column].dtype.name == "category":
                     current_categories = X[column].cat.categories
                     if "__NAN__" not in current_categories:
-                        X[column] = X[column].cat.add_categories("__NAN__").fillna("__NAN__")
+                        X[column] = X[column].cat.add_categories("__NAN__")
+                    X[column] = X[column].fillna("__NAN__")
             if cat_columns:
                 X[cat_columns] = X[cat_columns].astype("category")
                 # Pin codes to the categories seen at fit time so they do not
                 # drift when the predict-time column has a different value
                 # distribution than the fit-time column (see issue #1101).
-                # Older pickles without `_cat_categories` fall back to
-                # whatever `astype("category")` inferred above.
+                # Three-tier fallback for cross-version pickle compatibility:
+                #   1. `_ordinal_encoder` (post-#1564)  — sklearn OrdinalEncoder is
+                #      the source of truth for allowed categories;
+                #   2. `_cat_categories` (post-#1561)   — ad-hoc dict from the
+                #      defensive patch that landed before this refactor;
+                #   3. neither (pre-#1561 pickle)       — fall through to
+                #      whatever `astype("category")` inferred above.
+                encoder = getattr(self, "_ordinal_encoder", None)
                 saved_cats_map = getattr(self, "_cat_categories", None)
-                if saved_cats_map:
+                if encoder is not None:
+                    for col_idx, column in enumerate(cat_columns):
+                        known_cats = list(encoder.categories_[col_idx])
+                        # Include "__NAN__" as the sentinel slot even if the
+                        # fit-time data did not contain missing values.
+                        pinned_cats = list(known_cats)
+                        if "__NAN__" not in pinned_cats:
+                            pinned_cats.append("__NAN__")
+                        current = X[column].astype(object)
+                        unseen_mask = ~current.isin(pinned_cats) & current.notna()
+                        if unseen_mask.any():
+                            samples = sorted({str(v) for v in current[unseen_mask].unique()})[:5]
+                            warnings.warn(
+                                f"Column '{column}' contains values unseen at fit time "
+                                f"(e.g. {samples}); these rows will be encoded as '__NAN__' "
+                                "and predictions may be unreliable.",
+                                UserWarning,
+                                stacklevel=2,
+                            )
+                            current = current.where(~unseen_mask, "__NAN__")
+                        X[column] = pd.Categorical(current, categories=pinned_cats)
+                elif saved_cats_map:
                     for column in cat_columns:
                         saved_cats = saved_cats_map.get(column)
                         if saved_cats is None:

test/automl/test_preprocess_api.py

@@ -1,8 +1,10 @@
 """Tests for the public preprocessor APIs."""
 import unittest
+from unittest.mock import patch
 
 import numpy as np
 import pandas as pd
+import pytest
 from sklearn.datasets import load_breast_cancer, load_diabetes
 
 from flaml import AutoML
@@ -292,5 +294,151 @@ def test_unseen_categories_emit_warning_and_map_to_sentinel(self):
         self.assertTrue((unseen_rows == nan_code).all())
 
 
+class TestOrdinalEncoderBackedTransform(unittest.TestCase):
+    """Coverage for the categorical-encoding refactor in #1564: DataTransformer
+    uses sklearn's OrdinalEncoder as the source of truth for the per-column
+    category list at fit time, and the three-tier backward-compat fallback
+    (`_ordinal_encoder` → `_cat_categories` → legacy) at transform time so that
+    pickles from any recent FLAML version continue to load."""
+
+    def _fit_simple(self):
+        from flaml.automl.data import DataTransformer
+        from flaml.automl.task.factory import task_factory
+
+        rng = np.random.RandomState(0)
+        n = 100
+        fit_df = pd.DataFrame({"a": rng.randn(n), "gender": rng.choice(["M", "F"], n)})
+        fit_y = pd.Series(rng.randn(n))
+        transformer = DataTransformer()
+        task = task_factory("regression", fit_df, fit_y)
+        X_fit, _ = transformer.fit_transform(fit_df.copy(), fit_y, task)
+        return transformer, X_fit
+
+    def test_fit_transform_installs_ordinal_encoder(self):
+        from sklearn.preprocessing import OrdinalEncoder
+
+        transformer, _ = self._fit_simple()
+        self.assertTrue(hasattr(transformer, "_ordinal_encoder"))
+        self.assertIsInstance(transformer._ordinal_encoder, OrdinalEncoder)
+        self.assertIn("gender", transformer._cat_columns)
+
+    def test_ordinal_encoder_path_matches_1561_semantics(self):
+        """Refactored path preserves the observable behavior from #1561:
+        - known-category codes are stable across fit/predict distributions,
+        - unseen values are remapped to the "__NAN__" sentinel code,
+        - a `UserWarning` is emitted on unseen values."""
+        import warnings
+
+        transformer, X_fit = self._fit_simple()
+        fit_M_code = int(X_fit["gender"].cat.codes[X_fit["gender"] == "M"].iloc[0])
+
+        # (a) Stability under a smaller predict-time value set
+        predict_df = pd.DataFrame({"a": np.zeros(20), "gender": ["M"] * 20})
+        X_pred = transformer.transform(predict_df.copy())
+        pred_M_code = int(X_pred["gender"].cat.codes[X_pred["gender"] == "M"].iloc[0])
+        self.assertEqual(fit_M_code, pred_M_code)
+
+        # (b) Unseen-category warning + (c) sentinel remap
+        predict_df2 = pd.DataFrame({"a": np.zeros(5), "gender": ["M", "F", "X", "M", "Y"]})
+        with warnings.catch_warnings(record=True) as caught:
+            warnings.simplefilter("always")
+            X_pred2 = transformer.transform(predict_df2.copy())
+        unseen_warnings = [
+            w for w in caught if issubclass(w.category, UserWarning) and "unseen at fit time" in str(w.message)
+        ]
+        self.assertEqual(len(unseen_warnings), 1)
+        nan_code = list(X_pred2["gender"].cat.categories).index("__NAN__")
+        unseen_rows = X_pred2["gender"].cat.codes[predict_df2["gender"].isin(["X", "Y"]).values]
+        self.assertTrue((unseen_rows == nan_code).all())
+
+    def test_transform_falls_back_to_cat_categories_when_encoder_missing(self):
+        """Pickles produced between #1561 and #1564 only have `_cat_categories`.
+        `transform()` must still work correctly for them."""
+        transformer, X_fit = self._fit_simple()
+        # Simulate a #1561-era pickle by removing the encoder and installing the
+        # ad-hoc dict the older code produced.
+        del transformer._ordinal_encoder
+        transformer._cat_categories = {
+            "gender": list(X_fit["gender"].cat.categories) + ["__NAN__"],
+        }
+
+        fit_M_code = int(X_fit["gender"].cat.codes[X_fit["gender"] == "M"].iloc[0])
+        predict_df = pd.DataFrame({"a": np.zeros(20), "gender": ["M"] * 20})
+        X_pred = transformer.transform(predict_df.copy())
+        pred_M_code = int(X_pred["gender"].cat.codes[X_pred["gender"] == "M"].iloc[0])
+        self.assertEqual(fit_M_code, pred_M_code)
+
+    def test_transform_legacy_pickle_without_either_attribute(self):
+        """Pickles from before #1561 have neither `_ordinal_encoder` nor
+        `_cat_categories`. `transform()` must not raise; it falls through to the
+        legacy `astype("category")` path. Drift is possible on those pickles
+        (that is the bug that #1561 and #1564 fix), but load-and-predict must
+        continue to work without upgrading users' pickle files."""
+        transformer, _ = self._fit_simple()
+        del transformer._ordinal_encoder  # no `_cat_categories` was ever set
+
+        predict_df = pd.DataFrame({"a": np.zeros(20), "gender": ["M"] * 20})
+        # Legacy path just does astype("category") — no exception, no warning.
+        X_pred = transformer.transform(predict_df.copy())
+        self.assertEqual(str(X_pred["gender"].dtype), "category")
+
+
+@pytest.mark.parametrize("columns", [["first", "second"], [0, 1], [0, "second"]])
+@pytest.mark.parametrize("categories", [["z", "a", "unused"], [3, 1, 2], [1, "1", "unused"]])
+def test_category_order_and_column_labels(columns, categories):
+    from flaml.automl.data import DataTransformer
+    from flaml.automl.task.factory import task_factory
+
+    values = [categories[0], categories[1], None] * 4
+    X = pd.DataFrame(
+        {
+            columns[0]: pd.Categorical(values, categories=categories, ordered=True),
+            columns[1]: pd.Categorical(["b", "a"] * 6, categories=["b", "a"]),
+        }
+    )
+    y = pd.Series(np.arange(len(X), dtype=float))
+    transformer = DataTransformer()
+    fitted, _ = transformer.fit_transform(X.copy(), y, task_factory("regression", X, y))
+    predicted = transformer.transform(X.copy())
+    for column in columns:
+        assert list(fitted[column].cat.categories) == list(predicted[column].cat.categories)
+        np.testing.assert_array_equal(fitted[column].cat.codes, predicted[column].cat.codes)
+    assert list(fitted[columns[0]].cat.categories) == categories + ["__NAN__"]
+    np.testing.assert_array_equal(predicted[columns[0]].cat.codes, [0, 1, 3] * 4)
+    subset = X.iloc[[1, 0]].astype(object)
+    np.testing.assert_array_equal(transformer.transform(subset)[columns[0]].cat.codes, [1, 0])
+    subset.loc[subset.index[0], columns[0]] = "new"
+    with pytest.warns(UserWarning, match="unseen at fit time"):
+        unseen = transformer.transform(subset)
+    assert unseen[columns[0]].iloc[0] == "__NAN__"
+
+
+def test_encoder_uses_sklearn_10_constructor():
+    from sklearn.preprocessing import OrdinalEncoder
+
+    def legacy_encoder(*, categories="auto", dtype=np.float64, handle_unknown="error", unknown_value=None):
+        return OrdinalEncoder(
+            categories=categories, dtype=dtype, handle_unknown=handle_unknown, unknown_value=unknown_value
+        )
+
+    with patch("sklearn.preprocessing.OrdinalEncoder", side_effect=legacy_encoder):
+        transformer, fitted = TestOrdinalEncoderBackedTransform()._fit_simple()
+    predicted = transformer.transform(fitted.copy())
+    np.testing.assert_array_equal(fitted["gender"].cat.codes, predicted["gender"].cat.codes)
+
+
+def test_existing_missing_sentinel_is_filled():
+    from flaml.automl.data import DataTransformer
+    from flaml.automl.task.factory import task_factory
+
+    X = pd.DataFrame({"cat": pd.Categorical([1, "one", None] * 4, categories=[1, "one", "__NAN__"])})
+    y = pd.Series(np.arange(len(X), dtype=float))
+    transformer = DataTransformer()
+    fitted, _ = transformer.fit_transform(X.copy(), y, task_factory("regression", X, y))
+    predicted = transformer.transform(X.copy())
+    np.testing.assert_array_equal(fitted["cat"].cat.codes, [0, 1, 2] * 4)
+    np.testing.assert_array_equal(predicted["cat"].cat.codes, fitted["cat"].cat.codes)
+
+
 if __name__ == "__main__":
     unittest.main()