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

[Imran Ahamed (immu4989)]

Why are these changes needed?

Implements the refactor agreed on #1564 (-1 sentinel greenlit by Li Jiang (@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.

[Imran Ahamed (immu4989)]

Heads-up: CI is red on this PR (7 of 8 build jobs), and after digging in I think this refactor is not viable as scoped — reporting rather than patching around it.

What fails

test/automl/test_classification.py::TestClassification::test_preprocess and test/automl/test_score.py::TestScore::test_classification:

TypeError: Encoders require their input argument must be uniformly strings or numbers. Got ['float', 'str']
  sklearn/utils/_encode.py:182

Root cause

FLAML supports mixed-type categorical columns, and test_preprocess deliberately exercises that — its f3 column is ['a','b','a','c','c','b','b','b','b','b','a','b',1.0,1.0,'a'], mixing strings with a float. sklearn's OrdinalEncoder rejects that combination by design:

col = pd.Series(['a','b','a','c', 1.0, 1.0, 'a'])

pd.Categorical(col).categories          # → [1.0, 'a', 'b', 'c']   works
OrdinalEncoder().fit(col.to_frame())    # → TypeError: ... Got ['float', 'str']

The home-rolled pd.Categorical path (both pre-#1561 and the _cat_categories pinned-list version) handles this natively. OrdinalEncoder structurally cannot. Verified that the same test passes on current main.

Why I'm not patching around it

The two workarounds both cost more than the refactor is worth:

1. Per-column fallback — use OrdinalEncoder for uniformly-typed columns, keep the _cat_categories pinned list for mixed-type ones. That leaves us with two encoding mechanisms plus the "__NAN__" sentinel plus the three-tier pickle-compat chain — strictly more complex than the fix: DataTransformer pins categorical codes at fit time + warns on unseen values (#1101) #1561 state this was meant to simplify.
2. Coerce everything to str on both fit and transform. Consistent, but it changes the observable category values handed to downstream estimators (CatBoost's cat_features and LGBM's categorical path would see "1.0" instead of 1.0) — a silent behavior change on a correctness-sensitive path, which is exactly what Handling categorical variables on new data #1101 was about.

The stated benefits in #1564 were "standard sklearn idiom, drop the ad-hoc dict, drop the sentinel." Neither workaround delivers those, so the refactor stops paying for itself.

Suggestion

Close this PR and leave the #1561 defensive implementation in place — it already fixes the actual user-facing bug from #1101 (encoding drift + unseen-value warning) and it supports mixed-type columns, which OrdinalEncoder cannot.

The one durable thing worth keeping from this exercise is the constraint itself: any future refactor of DataTransformer's categorical handling has to preserve mixed-type column support, and that isn't documented anywhere today. Happy to send a small PR adding that as a comment plus an explicit regression test if you think it's worth pinning down.

Li Jiang (@thinkall) — your call; I'll close this unless you'd prefer one of the workarounds above.