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

[immu4989]

Why are these changes needed?

Implements the refactor agreed on #1564 (-1 sentinel greenlit by @thinkall). Replaces the ad-hoc _cat_categories dict introduced in #1561 with a fitted sklearn.preprocessing.OrdinalEncoder(handle_unknown="use_encoded_value", unknown_value=-1, encoded_missing_value=-1) as the source of truth for the per-column category list.

The observable behavior is preserved (see the existing #1561 tests): known categories still get stable integer codes across fit/predict, unseen values still emit a UserWarning and get remapped to the "__NAN__" sentinel category, and pandas categorical dtype is still returned so that LGBM's categorical_feature detection, CatBoost's cat_features inference, and the sklearn/KNeighbors estimator wrappers all continue to work unchanged.

Scope note

My original RFC on #1564 speculated that the refactor could switch to a numeric integer output (with -1 visible in the returned DataFrame). While implementing, I audited the downstream consumers in flaml/automl/model.py:

• LGBMEstimator and KNeighborsEstimator _preprocess paths (select_dtypes(include=["category"]) + .cat.codes)
• CatBoostEstimator cat_features = X.select_dtypes(include="category").columns
• Multiple sklearn-flavored estimator paths that expect category-dtype columns

Switching to a numeric output would silently break CatBoost's categorical-column detection and would need a coordinated per-estimator update. That's a much bigger surgery than the RFC scoped, and the observable value — better sklearn idiom internally — is the same either way. So this PR keeps the category-dtype output and uses OrdinalEncoder as the internal source of truth. If we want the numeric-output form later, I'll open a follow-up RFC.

What the PR does

flaml/automl/data.py — DataTransformer.fit_transform:

• After the existing X[cat_columns].astype("category") call, fits a per-column OrdinalEncoder on X[cat_columns].astype(object) and stores it as self._ordinal_encoder.
• Drops the write side of self._cat_categories (the _ordinal_encoder's categories_ array is the equivalent, sklearn-standard representation).

flaml/automl/data.py — DataTransformer.transform:

• Three-tier fallback for cross-version pickle compatibility:

  1. _ordinal_encoder present (post-[Proposal]: Replace home-rolled categorical encoding with sklearn OrdinalEncoder in DataTransformer (follow-up to #1101 / #1561) #1564) — use it as the source-of-truth category list, pin pd.Categorical(..., categories=known_cats + ["__NAN__"]).
  2. _cat_categories present (post-fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561 but pre-[Proposal]: Replace home-rolled categorical encoding with sklearn OrdinalEncoder in DataTransformer (follow-up to #1101 / #1561) #1564) — fall back to the existing dict-based path.
  3. Neither present (pre-fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561 pickle) — fall through to astype("category") as before.

• Unseen-value handling (UserWarning + remap to "__NAN__") preserved in both _ordinal_encoder and _cat_categories branches.

test/automl/test_preprocess_api.py — new TestOrdinalEncoderBackedTransform class:

1. test_fit_transform_installs_ordinal_encoder — asserts _ordinal_encoder is a fitted sklearn OrdinalEncoder after fit_transform.
2. test_ordinal_encoder_path_matches_1561_semantics — behavioral equivalence with the fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561 defensive patch: stable known-cat codes, UserWarning on unseen values, unseen rows remapped to the sentinel code.
3. test_transform_falls_back_to_cat_categories_when_encoder_missing — simulates a fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561-era pickle (removes _ordinal_encoder, installs _cat_categories) and asserts transform() still produces stable codes.
4. test_transform_legacy_pickle_without_either_attribute — simulates a pre-fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561 pickle (neither attribute present) and asserts transform() doesn't raise; the legacy astype("category") path is exercised.

Verified locally

• New tests: pytest test/automl/test_preprocess_api.py::TestOrdinalEncoderBackedTransform — 4/4 pass.
• The existing fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561 tests (TestCategoricalEncodingStability) — 2/2 pass unchanged.
• Full test_preprocess_api.py file — 14/14 pass.
• Adjacent regression check: test_split.py + test_multioutput + test_ensemble_component_predict_via_public_preprocess — 10/10 pass.
• pre-commit run --files flaml/automl/data.py test/automl/test_preprocess_api.py — all hooks pass.

Related issue / PRs

• Closes [Proposal]: Replace home-rolled categorical encoding with sklearn OrdinalEncoder in DataTransformer (follow-up to #1101 / #1561) #1564.
• Builds on the defensive patch in fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561.
• No API surface changes; consumers of automl.preprocess(X) see identical output.

Checks

• I've used pre-commit to lint the changes in this PR (note the same in integrated in our CI checks).
• I've included any doc changes needed for https://microsoft.github.io/FLAML/. See https://microsoft.github.io/FLAML/docs/Contribute#documentation to build and test documentation locally. (The Production-Deployment guide already describes the current behavior in a way that survives the refactor — no doc update needed.)
• I've added tests (if relevant) corresponding to the changes introduced in this PR.
• I've made sure all auto checks have passed.

[Copilot]

Pull request overview

Refactors DataTransformer categorical encoding internals to use sklearn.preprocessing.OrdinalEncoder as the persisted “source of truth” for per-column category lists, while preserving the external behavior (pandas category output, stable codes across fit/transform, unseen-value warnings + "__NAN__" sentinel handling) and adding backward-compatible fallbacks for older pickles.

Changes:

• Store a fitted OrdinalEncoder during DataTransformer.fit_transform() and use its categories_ during transform() to pin categorical dtypes.
• Keep cross-version pickle compatibility via a three-tier transform fallback (_ordinal_encoder → _cat_categories → legacy).
• Add a new test class validating the new encoder-backed path and both fallback behaviors.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
flaml/automl/data.py	Replaces _cat_categories write-path with fitted OrdinalEncoder, and adds encoder-first fallback logic during transform().
test/automl/test_preprocess_api.py	Adds tests covering the encoder-backed behavior and backward-compatibility fallbacks.