[Proposal]: Replace home-rolled categorical encoding with sklearn OrdinalEncoder in DataTransformer (follow-up to #1101 / #1561)

[immu4989]

Background

PR #1561 (merged 2026-06-24) closed the silent-correctness bug in DataTransformer that #1101 reported: predict-time categorical codes used to drift relative to fit-time codes whenever the predict-time DataFrame happened to contain a different subset of the categorical values. #1561 fixed this defensively, by stashing the per-column category list at fit time and re-pinning it at predict time via pd.Categorical(..., categories=saved_cats), with a "__NAN__" sentinel slot for unseen values.

That fix was deliberately the minimum-change patch (~36 LoC in flaml/automl/data.py) — it preserves the existing .cat.codes-based mechanism and adds a defensive wrapper around it. The original reporter (@krolikowskib in #1101) actually asked for something more substantive:

"SKLearn and XGBoost estimators use ordinal encodings of categorical features. But it seems the categorical codes are extracted during inference. Doesn't it mean that the encodings will be different when running on a different dataset, thus mixing the categories passed to the model? If so, then sklearn's OrdinalEncoder would be a better choice here (persisting correct category codes)."

This proposal opens the question of whether it's worth doing the proper architectural refactor next.

Proposal

Replace the home-rolled pd.Categorical / .cat.codes mechanism inside DataTransformer with sklearn.preprocessing.OrdinalEncoder, configured as:

OrdinalEncoder(
    handle_unknown="use_encoded_value",
    unknown_value=-1,           # or another sentinel — see below
    encoded_missing_value=-1,
)

What this buys us over the #1561 defensive patch:

• Industry-standard idiom. Anyone familiar with sklearn knows what OrdinalEncoder(handle_unknown="use_encoded_value") does; no need to read FLAML internals to understand the encoding contract.
• Drops the "__NAN__" sentinel — handled by unknown_value/encoded_missing_value natively, no string-magic string in the pipeline.
• Removes the _cat_categories dict in favor of OrdinalEncoder.categories_ (already maintained by sklearn).
• Better integration story for users who want to pre-encode upstream of FLAML — they can use the same OrdinalEncoder configuration and FLAML will pass through unchanged.
• Cleaner test surface — sklearn's encoder is heavily tested; a smaller surface for FLAML to maintain.

Migration plan (proposed)

1. Behavioral equivalence test — write a parametric test that fits both DataTransformer (post-fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561) and the proposed DataTransformer (post-refactor) on the same DataFrame, asserts that for each common categorical value the integer code is identical, and that unseen values are encoded consistently in both. Lands as part of the PR.
2. Backward compatibility for pickled DataTransformer instances. Old pickles serialize _cat_categories (a dict). After the refactor, transform() checks for the presence of _ordinal_encoder first, falls back to the pre-fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561 _cat_categories path, then falls back to the legacy astype("category") path. Three-tier fallback covers any combination of pickled-from / loaded-by FLAML versions.
3. Sentinel choice. sklearn's OrdinalEncoder uses -1 as the default unknown_value. The current FLAML code reserves a "__NAN__" string slot. The migration would use unknown_value=-1, encoded_missing_value=-1 (both surface as -1) and keep emitting the same UserWarning introduced in fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561.
4. Estimator-side considerations. LightGBM and XGBoost accept -1 in their categorical-feature paths without complaint. Sklearn estimators that treat columns as numeric will see -1 as a low outlier; the existing FLAML normalizer/preprocessing already handles per-estimator scaling, so this should be neutral. Will verify per-estimator during implementation.

Risks / open questions

• Performance. OrdinalEncoder.transform involves a search over categories_; current pd.Categorical is C-level. On wide DataFrames with many cat columns this could be slower. Will benchmark in the PR.
• NaN semantics. Today the FLAML code maps both NaN and unseen categories to "__NAN__". With sklearn's encoder these would both map to -1 (we'd set both unknown_value and encoded_missing_value). Same observable behavior, but worth confirming in tests.
• Time-series path. DataTransformer has a separate code path for task.is_ts_forecast() — need to verify the refactor doesn't disturb the timestamp-column handling at the top of fit_transform.

What I'm asking for

Maintainer signal before I write any code. Specifically:

• Is this refactor wanted, or do you prefer to keep fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561's defensive shape and revisit only if a follow-up bug surfaces?
• If wanted: any preferences on the sentinel value (-1 vs another) or the backward-compat fallback strategy?
• Anything else in the DataTransformer codepath I should account for that isn't obvious from outside?

Happy to take this on once aligned. Will not start coding until you (or another maintainer) chime in here.

Cross-references: closes #1101's longer-form ask; built on top of the defensive patch in #1561.

[jeet2005]

Following this thread — solid writeup on the OrdinalEncoder migration tradeoffs. I work with sklearn/XGBoost/LightGBM regularly (built an AutoML platform with a 56-algorithm registry across these), so the behavioral-equivalence testing and backward-compat fallback approach you outlined makes sense to me.
If @immu4989 doesn't end up picking this up after maintainer sign-off, happy to take it on. In the meantime, I'd be glad to help review/test once a PR lands, or pair on the per-estimator -1 sentinel verification (#4 in your migration plan) if that's useful.

[immu4989]

Thanks @jeet2005 — appreciated.

Planning to take this on once a maintainer signs off on the direction (defensive #1561 patch is in main; the question for this issue is whether to do the proper rebuild on OrdinalEncoder). I'd rather not start coding on a refactor that might get redirected.

If/when there's a PR, a peer review from someone who's used OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1) with sklearn / XGBoost / LightGBM together in a real pipeline would be welcome — the per-estimator -1 sentinel verification (#4 in the migration plan above) is the part I'd most appreciate a second set of eyes on.

[thinkall]

Thanks a lot for the proposal and the discussion, @immu4989, @jeet2005. The proposal sounds good to me. I agree with using "-1" as the sentinel value. Looking forward to your PRs!

[immu4989]

Thanks @thinkall — will start on it. Sentinel = -1 confirmed for both unknown_value and encoded_missing_value. Plan is to open a PR with the migration in a single commit, plus the behavioral-equivalence test and the three-tier backward-compat fallback (post-refactor _ordinal_encoder → post-#1561 _cat_categories → legacy astype("category")) so pickled DataTransformer instances from any recent FLAML version continue to load.