microsoft / microsoft/FLAML

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

Open
#1,564 5 comments 1 reaction 0 assignees View on GitHub
Dominant language
Jupyter Notebook
Stars
4.4k
Forks
565
Avg merge
5d 8m
Merged PRs (30d)
17

Description

## 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](https://github.com/microsoft/FLAML/issues/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](https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.OrdinalEncoder.html) 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:

```python
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-#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-#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 #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 #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.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in flaml/automl/data.py at DataTransformer.fit_transform and transform, including the task.is_ts_forecast() path. Review the existing _cat_categories, pd.Categorical, and .cat.codes flow alongside sklearn OrdinalEncoder, then define behavioral-equivalence coverage for common, unseen, and missing values. Also assess pickle fallback behavior, estimator paths, and benchmark results before treating the refactor as complete.

Written by the indexing model from the issue text.

Assessment

Tech stack
pandas, python, scikit-learn
Domain
data, machine-learning
Issue type
Refactor
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.