fix: improve feature name resolution (#175)

* fix: improve feature name resolution

* chore: bump version to 1.2.2

* chore: reduce example iterations

Author: marcellodebernardi

Makefile

@@ -276,7 +276,7 @@ run-titanic: build
 			--user-id dev_user \
 			--intent "predict whether a passenger was transported" \
 			--experiment-id titanic \
-			--max-iterations 10 \
+			--max-iterations 5 \
 			--work-dir /workdir/titanic/$(TIMESTAMP) \
 			--spark-mode local \
 			--enable-final-evaluation
@@ -302,7 +302,7 @@ run-house-prices: build
 			--user-id dev_user \
 			--intent "predict house sale price" \
 			--experiment-id house_prices \
-			--max-iterations 10 \
+			--max-iterations 5 \
 			--work-dir /workdir/house_prices/$(TIMESTAMP) \
 			--spark-mode local \
 			--enable-final-evaluation

plexe/CODE_INDEX.md

@@ -1,6 +1,6 @@
 # Code Index: plexe
 
-> Generated on 2026-02-26 15:06:43
+> Generated on 2026-02-26 19:02:04
 
 Code structure and public interface documentation for the **plexe** package.
 

plexe/templates/features/pipeline_runner.py

@@ -9,7 +9,7 @@
 import logging
 from collections.abc import Iterator
 from textwrap import shorten
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import cloudpickle
 import pandas as pd
@@ -21,6 +21,61 @@
 logger = logging.getLogger(__name__)
 
 
+def _generic_feature_names(num_output_features: int) -> list[str]:
+    return [f"feature_{i}" for i in range(num_output_features)]
+
+
+def _call_get_feature_names_out(transformer: Any, feature_columns: list[str]) -> list[str]:
+    try:
+        names = transformer.get_feature_names_out(input_features=feature_columns)
+    except TypeError:
+        names = transformer.get_feature_names_out()
+    return list(names)
+
+
+def _try_get_feature_names_out(
+    fitted_pipeline: Pipeline, feature_columns: list[str]
+) -> tuple[list[str] | None, str | None]:
+    strategies: list[tuple[str, Any]] = []
+
+    if hasattr(fitted_pipeline, "get_feature_names_out"):
+        strategies.append(("pipeline.get_feature_names_out()", fitted_pipeline))
+
+    if isinstance(fitted_pipeline, Pipeline) and fitted_pipeline.steps:
+        last_step_name, last_step = fitted_pipeline.steps[-1]
+
+        if not hasattr(last_step, "get_feature_names_out") and len(fitted_pipeline.steps) > 1:
+            strategies.append(
+                (f"pipeline[:-1].get_feature_names_out() (skipped {last_step_name})", fitted_pipeline[:-1])
+            )
+
+        if hasattr(last_step, "get_feature_names_out"):
+            strategies.append((f"{last_step_name}.get_feature_names_out()", last_step))
+
+    for label, transformer in strategies:
+        try:
+            names = _call_get_feature_names_out(transformer, feature_columns)
+        except Exception:
+            continue
+        return names, label
+
+    return None, None
+
+
+def _resolve_transformed_feature_names(
+    fitted_pipeline: Pipeline, feature_columns: list[str], num_output_features: int
+) -> tuple[list[str], str]:
+    names, source = _try_get_feature_names_out(fitted_pipeline, feature_columns)
+
+    if names is None:
+        return _generic_feature_names(num_output_features), "generic"
+
+    if len(names) != num_output_features:
+        return _generic_feature_names(num_output_features), "generic_mismatch"
+
+    return names, source or "pipeline.get_feature_names_out()"
+
+
 def transform_dataset_via_spark(
     spark: SparkSession,
     dataset_uri: str,
@@ -101,20 +156,19 @@ def transform_dataset_via_spark(
         sample_transformed_features = sample_transformed
     else:
         # Pipeline returned numpy array - need to assign column names
-        try:
-            # Try sklearn's get_feature_names_out() (sklearn 1.0+)
-            if hasattr(fitted_pipeline, "get_feature_names_out"):
-                transformed_feature_cols = fitted_pipeline.get_feature_names_out(
-                    input_features=feature_columns
-                ).tolist()
-                logger.info("Using feature names from pipeline.get_feature_names_out()")
-            else:
-                raise AttributeError("No get_feature_names_out method")
-        except Exception as e:
-            # Fallback to generic names
-            num_output_features = sample_transformed.shape[1]
-            transformed_feature_cols = [f"feature_{i}" for i in range(num_output_features)]
-            logger.warning(f"Using generic feature names ({e})")
+        num_output_features = sample_transformed.shape[1]
+        transformed_feature_cols, feature_name_source = _resolve_transformed_feature_names(
+            fitted_pipeline=fitted_pipeline,
+            feature_columns=feature_columns,
+            num_output_features=num_output_features,
+        )
+
+        if feature_name_source == "generic":
+            logger.debug("Falling back to generic feature names (pipeline has no get_feature_names_out).")
+        elif feature_name_source == "generic_mismatch":
+            logger.debug("Falling back to generic feature names (pipeline feature count mismatch).")
+        else:
+            logger.info("Using feature names from %s", feature_name_source)
 
         sample_transformed_features = pd.DataFrame(sample_transformed, columns=transformed_feature_cols)
 

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "plexe"
-version = "1.2.1"
+version = "1.2.2"
 description = "An agentic framework for building ML models from natural language"
 authors = [
     "Marcello De Bernardi <mdebernardi@plexe.ai>",

tests/CODE_INDEX.md

@@ -1,6 +1,6 @@
 # Code Index: tests
 
-> Generated on 2026-02-26 15:06:43
+> Generated on 2026-02-26 19:02:04
 
 Test suite structure and test case documentation.
 
@@ -49,6 +49,23 @@ Unit tests for SearchJournal.
 - `test_journal_improvement_trend_improving()` - Test improvement trend with steadily improving solutions.
 - `test_journal_improvement_trend_insufficient_data()` - Test improvement trend with fewer than 2 successful solutions.
 
+---
+## `unit/templates/features/test_pipeline_runner.py`
+Unit tests for pipeline_runner feature name resolution.
+
+**`NoFeatureNamesTransformer`** - Transformer without get_feature_names_out.
+- `fit(self, x, y)` - No description
+- `transform(self, x)` - No description
+
+**`SelectFirstColumnTransformer`** - Transformer that reduces output to a single column.
+- `fit(self, x, y)` - No description
+- `transform(self, x)` - No description
+
+**Functions:**
+- `test_resolve_feature_names_uses_pipeline_minus_last()` - Falls back to pipeline[:-1] when last step lacks get_feature_names_out.
+- `test_resolve_feature_names_falls_back_on_mismatch()` - Returns generic names when resolved names don't match output count.
+- `test_resolve_feature_names_falls_back_when_unavailable()` - Returns generic names when no get_feature_names_out is available.
+
 ---
 ## `unit/test_config.py`
 Unit tests for config helpers.

tests/unit/templates/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for templates."""

tests/unit/templates/features/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for feature templates."""

tests/unit/templates/features/test_pipeline_runner.py

@@ -0,0 +1,83 @@
+"""
+Unit tests for pipeline_runner feature name resolution.
+"""
+
+import pandas as pd
+from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.pipeline import Pipeline
+from sklearn.preprocessing import StandardScaler
+
+from plexe.templates.features.pipeline_runner import _resolve_transformed_feature_names
+
+
+class NoFeatureNamesTransformer(BaseEstimator, TransformerMixin):
+    """Transformer without get_feature_names_out."""
+
+    def fit(self, x, y=None):
+        return self
+
+    def transform(self, x):
+        return x
+
+
+class SelectFirstColumnTransformer(BaseEstimator, TransformerMixin):
+    """Transformer that reduces output to a single column."""
+
+    def fit(self, x, y=None):
+        return self
+
+    def transform(self, x):
+        if hasattr(x, "iloc"):
+            return x.iloc[:, [0]]
+        return x[:, [0]]
+
+
+def test_resolve_feature_names_uses_pipeline_minus_last():
+    """Falls back to pipeline[:-1] when last step lacks get_feature_names_out."""
+    df = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]})
+    pipeline = Pipeline([("scale", StandardScaler()), ("noop", NoFeatureNamesTransformer())])
+
+    pipeline.fit(df)
+
+    names, source = _resolve_transformed_feature_names(
+        fitted_pipeline=pipeline,
+        feature_columns=["a", "b"],
+        num_output_features=2,
+    )
+
+    assert names == ["a", "b"]
+    assert "pipeline[:-1]" in source
+
+
+def test_resolve_feature_names_falls_back_on_mismatch():
+    """Returns generic names when resolved names don't match output count."""
+    df = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]})
+    pipeline = Pipeline([("scale", StandardScaler()), ("select", SelectFirstColumnTransformer())])
+
+    pipeline.fit(df)
+
+    names, source = _resolve_transformed_feature_names(
+        fitted_pipeline=pipeline,
+        feature_columns=["a", "b"],
+        num_output_features=1,
+    )
+
+    assert names == ["feature_0"]
+    assert source == "generic_mismatch"
+
+
+def test_resolve_feature_names_falls_back_when_unavailable():
+    """Returns generic names when no get_feature_names_out is available."""
+    df = pd.DataFrame({"a": [1.0, 2.0], "b": [3.0, 4.0]})
+    pipeline = Pipeline([("noop", NoFeatureNamesTransformer())])
+
+    pipeline.fit(df)
+
+    names, source = _resolve_transformed_feature_names(
+        fitted_pipeline=pipeline,
+        feature_columns=["a", "b"],
+        num_output_features=2,
+    )
+
+    assert names == ["feature_0", "feature_1"]
+    assert source == "generic"