From e0fa8213ef24ea93432dc9375ee7a8c1fedf71da Mon Sep 17 00:00:00 2001
From: Imran Ahamed <immu4989@gmail.com>
Date: Tue, 9 Jun 2026 00:22:11 -0500
Subject: [PATCH] docs: ensemble component predict pattern + regression test
 (#1136)

---
 flaml/automl/automl.py         |  8 +++++
 test/automl/test_regression.py | 62 ++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+)

diff --git a/flaml/automl/automl.py b/flaml/automl/automl.py
index cb3fe37857..44926bdc4d 100644
--- a/flaml/automl/automl.py
+++ b/flaml/automl/automl.py
@@ -890,6 +890,14 @@ def preprocess(
 
             # Apply task-level preprocessing to new data
             X_test_preprocessed = automl.preprocess(X_test)
+
+            # Required when calling a single ensemble component directly, since
+            # `automl.model.estimators_[i]` was fitted on already-preprocessed
+            # data and cannot consume raw input (see issue #1136):
+            automl_ensemble = AutoML()
+            automl_ensemble.fit(X_train, y_train, task="classification", ensemble=True)
+            X_test_preprocessed = automl_ensemble.preprocess(X_test)
+            component_pred = automl_ensemble.model.estimators_[0].predict(X_test_preprocessed)
             ```
         """
         if not hasattr(self, "_state") or self._state is None:
diff --git a/test/automl/test_regression.py b/test/automl/test_regression.py
index 74c6111274..6bfd75a898 100644
--- a/test/automl/test_regression.py
+++ b/test/automl/test_regression.py
@@ -244,6 +244,68 @@ def test_multioutput():
     print(model.predict(X_test))
 
 
+def test_ensemble_component_predict_via_public_preprocess():
+    """Regression coverage for #1136 — ensemble component models trained on data with
+    categorical features cannot consume raw input; consumers must apply the public
+    `automl.preprocess(X)` method (added in #1497) before delegating to a single
+    component picked out of `automl.model.estimators_`."""
+    import pandas as pd
+
+    rng = np.random.RandomState(42)
+    n = 400
+    df = pd.DataFrame(
+        {
+            "age": rng.randint(20, 70, n),
+            "income": rng.normal(50000, 15000, n),
+            "gender": rng.choice(["M", "F"], n),
+            "education": rng.choice(["HS", "BS", "MS", "PhD"], n),
+        }
+    )
+    y_true = (
+        0.02 * df["age"]
+        + 0.00001 * df["income"]
+        + (df["gender"] == "M").astype(int) * 0.5
+        + df["education"].map({"HS": 0, "BS": 0.3, "MS": 0.6, "PhD": 1.0}).values
+        + rng.normal(0, 0.1, n)
+    )
+
+    automl = AutoML()
+    automl.fit(
+        df,
+        y_true,
+        task="regression",
+        ensemble=True,
+        n_jobs=1,
+        time_budget=-1,
+        max_iter=12,
+        estimator_list=["lgbm", "xgboost", "rf"],
+        verbose=0,
+    )
+
+    components = getattr(automl.model, "estimators_", None)
+    if components is None:
+        pytest.skip("ensemble did not build with this configuration")
+
+    # Public predict on the top-level AutoML continues to work (sanity check).
+    top_pred = automl.predict(df)
+    assert len(top_pred) == n
+
+    # Component models cannot consume raw input — this is the original #1136 failure.
+    raised = 0
+    for est in components:
+        try:
+            est.predict(df)
+        except Exception:
+            raised += 1
+    assert raised >= 1, "expected at least one component to fail on raw categorical input"
+
+    # The public `preprocess(X)` API (added in #1497) is the supported workaround.
+    df_preprocessed = automl.preprocess(df)
+    for est in components:
+        pred = est.predict(df_preprocessed)
+        assert len(pred) == n
+
+
 @pytest.mark.parametrize(
     "estimator",
     [