From 33e65884107558996dd7c6c002c8018c73951731 Mon Sep 17 00:00:00 2001 From: Imran Ahamed Date: Wed, 1 Jul 2026 19:11:18 -0500 Subject: [PATCH 1/2] refactor: use sklearn OrdinalEncoder as source of truth for DataTransformer categorical encoding (#1564) --- flaml/automl/data.py | 66 +++++++++++++++++----- test/automl/test_preprocess_api.py | 89 ++++++++++++++++++++++++++++++ 2 files changed, 141 insertions(+), 14 deletions(-) diff --git a/flaml/automl/data.py b/flaml/automl/data.py index c27adc2879..e6eb38b507 100644 --- a/flaml/automl/data.py +++ b/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. Unknown values and missing values + # both encode to -1 internally; `transform()` uses that to + # detect drift and remap those rows 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,42 @@ 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_columns = list(encoder.feature_names_in_) + for column in cat_columns: + try: + col_idx = encoder_columns.index(column) + except ValueError: + 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: diff --git a/test/automl/test_preprocess_api.py b/test/automl/test_preprocess_api.py index 5b207595e7..cfb6a60172 100644 --- a/test/automl/test_preprocess_api.py +++ b/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() From ddca0cd68ab42c01de80f2d2ca4fa7c68d8effd7 Mon Sep 17 00:00:00 2001 From: Imran Ahamed Date: Mon, 20 Jul 2026 16:11:15 -0500 Subject: [PATCH 2/2] address review: dict lookup for encoder columns, accurate transform comment --- flaml/automl/data.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/flaml/automl/data.py b/flaml/automl/data.py index e6eb38b507..cd34e98c32 100644 --- a/flaml/automl/data.py +++ b/flaml/automl/data.py @@ -354,12 +354,12 @@ def fit_transform(self, X: Union[DataFrame, np.ndarray], y, task: Union[str, "Ta # 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. Unknown values and missing values - # both encode to -1 internally; `transform()` uses that to - # detect drift and remap those rows to the "__NAN__" sentinel - # category, preserving the pandas categorical dtype that the - # downstream estimator wrappers (LGBM / CatBoost / sklearn) - # key off of. + # 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( @@ -481,11 +481,10 @@ def transform(self, X: Union[DataFrame, np.array]): encoder = getattr(self, "_ordinal_encoder", None) saved_cats_map = getattr(self, "_cat_categories", None) if encoder is not None: - encoder_columns = list(encoder.feature_names_in_) + encoder_col_idx = {name: i for i, name in enumerate(encoder.feature_names_in_)} for column in cat_columns: - try: - col_idx = encoder_columns.index(column) - except ValueError: + 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