[Bug]: flaml.default.XGBRegressor does not preprocess eval_set

[fingoldo]

Describe the bug

Sometimes (when there are cat columns, for example) flaml (at least zero-shot) re-arranges columns. But for XGBRegressor/XGBClassifier, it misses to do that for the validation dataframe.

Steps to reproduce

import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from xgboost import XGBClassifier
from sklearn.metrics import accuracy_score

# --- 1. Create synthetic dataset with numeric + categorical features ---
np.random.seed(42)
n = 1000

df = pd.DataFrame({
    "num1": np.random.randn(n),
    "num2": np.random.rand(n) * 10,
    "cat1": np.random.choice(["A", "B", "C"], size=n),
    "cat2": np.random.choice(["X", "Y"], size=n),
    "target": np.random.choice([0, 1], size=n)
})

# --- 2. Split data ---
X = df.drop(columns="target")
y = df["target"]

X_train, X_valid, y_train, y_valid = train_test_split(X, y, test_size=0.2, random_state=0)

# --- 3. Convert categorical columns to pandas 'category' dtype ---
for col in X_train.select_dtypes(include="object").columns:
    X_train[col] = X_train[col].astype("category")
    X_valid[col] = X_valid[col].astype("category")


# --- 4. Define XGBoost model ---
model = XGBClassifier(
    tree_method="hist",              # Efficient, supports categorical features
    enable_categorical=True,         # Important!
    eval_metric="logloss",
    use_label_encoder=False,
    early_stopping_rounds=10,
    random_state=0
)

# --- 5. Fit model with early stopping ---
model.fit(
    X_train, y_train,
    eval_set=[(X_valid, y_valid)],   # validation set for early stopping    
    verbose=True
)

[0] validation_0-logloss:0.69096 [1] validation_0-logloss:0.69439 [2] validation_0-logloss:0.70184 [3] validation_0-logloss:0.70530 [4] validation_0-logloss:0.70542 [5] validation_0-logloss:0.70719 [6] validation_0-logloss:0.71508 [7] validation_0-logloss:0.71836 [8] validation_0-logloss:0.72136 [9] validation_0-logloss:0.72541

import flaml.default as flaml_zeroshot

model = flaml_zeroshot.XGBClassifier(
    tree_method="hist",              # Efficient, supports categorical features
    enable_categorical=True,         # Important!
    eval_metric="logloss",
    use_label_encoder=False,
    early_stopping_rounds=10,
    random_state=0
)

# --- 5. Fit model with early stopping ---
model.fit(
    X_train, y_train,
    eval_set=[(X_valid, y_valid)],   # validation set for early stopping    
    verbose=True
)

---

ValueError Traceback (most recent call last) Cell In[4], line 13 3 model = flaml_zeroshot.XGBClassifier( 4 tree_method="hist", # Efficient, supports categorical features 5 enable_categorical=True, # Important! (...) 9 random_state=0 10 ) 12 # --- 5. Fit model with early stopping --- ---> 13 model.fit( 14 X_train, y_train, 15 eval_set=[(X_valid, y_valid)], # validation set for early stopping
16 verbose=True 17 )

File /venv/main/lib/python3.12/site-packages/flaml/default/estimator.py:106, in flamlize_estimator..EstimatorClass.fit(self, X, y, *args, **params) 97 self.set_params(**hyperparams) 98 if self._label_transformer and estimator_name in [ 99 "rf", 100 "extra_tree", (...) 104 ]: 105 # rf and et have trouble in handling boolean labels; xgboost requires integer labels --> 106 fitted = super().fit(X, y_transformed, *args, **params) 107 # if hasattr(self, "_classes"): 108 # self._classes = self.label_transformer.classes 109 # else: 110 try:

File /venv/main/lib/python3.12/site-packages/xgboost/core.py:774, in require_keyword_args..throw_if..inner_f(*args, **kwargs) 772 for k, arg in zip(sig.parameters, args): 773 kwargs[k] = arg --> 774 return func(**kwargs)

File /venv/main/lib/python3.12/site-packages/xgboost/sklearn.py:1803, in XGBClassifier.fit(self, X, y, sample_weight, base_margin, eval_set, verbose, xgb_model, sample_weight_eval_set, base_margin_eval_set, feature_weights) 1783 evals_result: EvalsLog = {} 1784 train_dmatrix, evals = _wrap_evaluation_matrices( 1785 missing=self.missing, 1786 X=X, (...) 1800 feature_types=feature_types, 1801 ) -> 1803 self._Booster = train( 1804 params, 1805 train_dmatrix, 1806 self.get_num_boosting_rounds(), 1807 evals=evals, 1808 early_stopping_rounds=self.early_stopping_rounds, 1809 evals_result=evals_result, 1810 obj=obj, 1811 custom_metric=metric, 1812 verbose_eval=verbose, 1813 xgb_model=model, 1814 callbacks=self.callbacks, 1815 ) 1817 if not callable(self.objective): 1818 self.objective = params["objective"]

File /venv/main/lib/python3.12/site-packages/xgboost/core.py:774, in require_keyword_args..throw_if..inner_f(*args, **kwargs) 772 for k, arg in zip(sig.parameters, args): 773 kwargs[k] = arg --> 774 return func(**kwargs)

File /venv/main/lib/python3.12/site-packages/xgboost/training.py:182, in train(params, dtrain, num_boost_round, evals, obj, maximize, early_stopping_rounds, evals_result, verbose_eval, xgb_model, callbacks, custom_metric) 175 if ( 176 isinstance(va, _RefMixIn) 177 and va.ref is not weakref.ref(dtrain) 178 and va is not dtrain 179 ): 180 raise ValueError(_RefError) --> 182 bst = Booster(params, [dtrain] + [d[0] for d in evals], model_file=xgb_model) 183 start_iteration = 0 185 if verbose_eval:

File /venv/main/lib/python3.12/site-packages/xgboost/core.py:2005, in Booster.init(self, params, cache, model_file) 1998 _check_call( 1999 _LIB.XGBoosterCreate( 2000 dmats, c_bst_ulong(len(cache)), ctypes.byref(self.handle) 2001 ) 2002 ) 2003 for d in cache: 2004 # Validate feature only after the feature names are saved into booster. -> 2005 self._assign_dmatrix_features(d) 2007 if isinstance(model_file, Booster): 2008 assert self.handle is not None

File /venv/main/lib/python3.12/site-packages/xgboost/core.py:3396, in Booster._assign_dmatrix_features(self, data) 3393 if self.feature_types is None: 3394 self.feature_types = ft -> 3396 self._validate_features(fn)

File /venv/main/lib/python3.12/site-packages/xgboost/core.py:3431, in Booster._validate_features(self, feature_names) 3425 if my_missing: 3426 msg += ( 3427 "\ntraining data did not have the following fields: " 3428 + ", ".join(str(s) for s in my_missing) 3429 ) -> 3431 raise ValueError(msg.format(self.feature_names, feature_names))

ValueError: feature_names mismatch: ['cat1', 'cat2', 'num1', 'num2'] ['num1', 'num2', 'cat1', 'cat2']

Model Used

No response

Expected Behavior

To process all validation sets.

Surprisingly, LGBMClassifier does that:

model = flaml_zeroshot.LGBMClassifier(
    early_stopping_rounds=10,
    verbose=2,
    random_state=0
)

# --- 5. Fit model with early stopping ---
model.fit(
    X_train, y_train,
    eval_set=(X_valid, y_valid),   # validation set for early stopping    
)

Early stopping, best iteration is: [1] valid_0's binary_logloss: 0.691358

I'm also not a fan of copying entire X dataframe. This makes FLAML not ready for big datasets. Can there be added some flag to tell FLAML to do its preprocessing inplace instead of copying entire X dataframe?

Screenshots and logs

No response

Additional Information

No response

[fingoldo]

In fact, why do XGB* estimators even need changing X? Xgboost supports auto handling of features of dtype=Category, I vote for dropping current logic & just ensuring that enable_categorical=True param is passed to the estimator.

[thinkall]

Thank you, @fingoldo , for the feedback! Would you like to raise a PR for the fix?