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

flaml/automl/data.py

@@ -351,17 +351,23 @@ 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
+                # Fit an OrdinalEncoder as the source of truth for the per-column
+                # allowed category list — see issue #1564. This replaces the
+                # ad-hoc `_cat_categories` dict from #1561 with sklearn's
+                # standard implementation. `transform()` reads the pinned lists
+                # from `encoder.categories_` (membership check; the encoder's
+                # own transform is not called) and remaps values outside the
+                # fit-time list to the "__NAN__" sentinel category, preserving
+                # the pandas categorical dtype that the downstream estimator
+                # wrappers (LGBM / CatBoost / sklearn) key off of.
+                from sklearn.preprocessing import OrdinalEncoder
+
+                self._ordinal_encoder = OrdinalEncoder(
+                    handle_unknown="use_encoded_value",
+                    unknown_value=-1,
+                    encoded_missing_value=-1,
+                )
+                self._ordinal_encoder.fit(X[cat_columns].astype(object))
             if num_columns:
                 X_num = X[num_columns]
                 try:
@@ -465,10 +471,41 @@ def transform(self, X: Union[DataFrame, np.array]):
                 # 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:
+                    encoder_col_idx = {name: i for i, name in enumerate(encoder.feature_names_in_)}
+                    for column in cat_columns:
+                        col_idx = encoder_col_idx.get(column)
+                        if col_idx is None:
+                            continue
+                        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

@@ -292,5 +292,94 @@ 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", list(transformer._ordinal_encoder.feature_names_in_))
+
+    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")
+
+
 if __name__ == "__main__":
     unittest.main()