Skip to content

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

Open
immu4989 wants to merge 1 commit into
microsoft:mainfrom
immu4989:flaml-feature-1564-ordinal-encoder-cat-refactor
Open

refactor: use sklearn OrdinalEncoder as source of truth for DataTransformer categorical encoding (#1564)#1569
immu4989 wants to merge 1 commit into
microsoft:mainfrom
immu4989:flaml-feature-1564-ordinal-encoder-cat-refactor

Conversation

@immu4989

@immu4989 immu4989 commented Jul 2, 2026

Copy link
Copy Markdown
Contributor

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.pyDataTransformer.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.pyDataTransformer.transform:

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

Checks

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant