From b440426e90e08e4be7b5359ec244d563e813592f Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Fri, 16 Oct 2026 14:14:36 +0200
Subject: [PATCH] fix(preprocessor): keep ColumnTransformer step names valid
 for any label

Preprocessor named each per-column step of its internal ColumnTransformer
f"num_{label}" / f"cat_{label}". scikit-learn rejects estimator names that
contain its "__" parameter separator, so fit crashed for any column label
starting with "_" or containing "__" (e.g. "_score", pyarrow's
"__index_level_0__", or "imp__age" from an upstream ColumnTransformer with
pandas output).

Steps now keep the num_<label> / cat_<label> name whenever scikit-learn
accepts it and otherwise get a unique positional name (num_col0, ...).
Block keys, feature names, lineage and the verbose=3 log are derived from
the step's kind and its column label, so they are unchanged for every
label that already worked and read num_<label> / cat_<label> for the new
ones.

Fixes #71
---
 pretab/compose/factory.py               | 47 +++++++++++----
 pretab/compose/inspection.py            | 60 ++++++++++++++-----
 pretab/preprocessor.py                  | 10 +++-
 tests/integration/test_column_labels.py | 79 +++++++++++++++++++++++++
 tests/integration/test_verbosity.py     | 11 ++++
 5 files changed, 180 insertions(+), 27 deletions(-)
 create mode 100644 tests/integration/test_column_labels.py

diff --git a/pretab/compose/factory.py b/pretab/compose/factory.py
index 8ced4d6..94ab082 100644
--- a/pretab/compose/factory.py
+++ b/pretab/compose/factory.py
@@ -305,6 +305,31 @@ def create_transformer(method: str, *, is_numerical: bool, config: PreprocessorC
     return pipeline
 
 
+def _step_names(blocks) -> list[str]:
+    """Return one valid, unique ColumnTransformer step name per ``(kind, feature)`` block.
+
+    A step is named ``f"{kind}_{feature}"`` whenever scikit-learn accepts that
+    name. A label that would put scikit-learn's ``__`` parameter separator into it
+    (one starting with ``_`` or containing ``__``) gets a positional
+    ``f"{kind}_col{i}"`` name instead, made unique against every other step. The
+    public block / feature names never use these internal names: they are derived
+    from the step's kind and its column label (see :mod:`pretab.compose.inspection`).
+    """
+    natural = [f"{kind}_{feature}" for kind, feature in blocks]
+    taken = {name for name in natural if "__" not in name}
+    names = []
+    for position, name in enumerate(natural):
+        if "__" in name:
+            kind = blocks[position][0]
+            name, suffix = f"{kind}_col{position}", 0
+            while name in taken:
+                suffix += 1
+                name = f"{kind}_col{position}_{suffix}"
+            taken.add(name)
+        names.append(name)
+    return names
+
+
 def build_column_transformer(
     config: PreprocessorConfig,
     numerical_features,
@@ -314,19 +339,19 @@ def build_column_transformer(
 ) -> ColumnTransformer:
     """Assemble the per-column pipelines into the final ColumnTransformer.
 
-    Numerical features are prefixed ``num_`` and categorical features ``cat_`` to
-    match the transformer names the Preprocessor exposes; untransformed columns
-    pass through via ``remainder="passthrough"``.
+    Numerical steps are prefixed ``num_`` and categorical steps ``cat_`` to match
+    the block names the Preprocessor exposes (see :func:`_step_names` for labels
+    that cannot be embedded in a step name); untransformed columns pass through via
+    ``remainder="passthrough"``.
     """
+    blocks = [("num", feature) for feature in numerical_features]
+    blocks += [("cat", feature) for feature in categorical_features]
     transformers = []
-    for feature in numerical_features:
-        method = config.method_for(feature, is_numerical=True)
-        pipeline = create_transformer(method, is_numerical=True, config=config)
-        transformers.append((f"num_{feature}", pipeline, [feature]))
-    for feature in categorical_features:
-        method = config.method_for(feature, is_numerical=False)
-        pipeline = create_transformer(method, is_numerical=False, config=config)
-        transformers.append((f"cat_{feature}", pipeline, [feature]))
+    for (kind, feature), name in zip(blocks, _step_names(blocks), strict=True):
+        is_numerical = kind == "num"
+        method = config.method_for(feature, is_numerical=is_numerical)
+        pipeline = create_transformer(method, is_numerical=is_numerical, config=config)
+        transformers.append((name, pipeline, [feature]))
     return ColumnTransformer(
         transformers=transformers,
         remainder="passthrough",
diff --git a/pretab/compose/inspection.py b/pretab/compose/inspection.py
index 7f3b9b3..d1a2786 100644
--- a/pretab/compose/inspection.py
+++ b/pretab/compose/inspection.py
@@ -16,6 +16,7 @@
 logger = get_logger(__name__)
 
 __all__ = [
+    "block_name",
     "build_feature_info",
     "build_feature_lineage",
     "build_transformer_summary",
@@ -24,15 +25,38 @@
 ]
 
 
+def _block(step_name, columns):
+    """Return ``(kind, feature)`` for a per-column ``num_*`` / ``cat_*`` step, else ``None``."""
+    if step_name == "remainder" or len(columns) != 1:
+        return None
+    kind, sep, _ = str(step_name).partition("_")
+    if not sep or kind not in ("num", "cat"):
+        return None
+    return kind, str(columns[0])
+
+
+def block_name(step_name, columns):
+    """Return the public block name (``"num_<label>"`` / ``"cat_<label>"``) of a step.
+
+    The name is built from the step's kind and its column label rather than read
+    from the step name itself, which differs for labels scikit-learn cannot embed
+    in a step name (see :func:`pretab.compose.factory._step_names`). Steps that are
+    not per-column representation blocks keep their own name.
+    """
+    block = _block(step_name, columns)
+    return step_name if block is None else f"{block[0]}_{block[1]}"
+
+
 def get_output_slices(column_transformer):
     """Return ordered ``(name, start, width)`` spans for each output block.
 
     Reads widths from ``output_indices_`` — the fitted index map that
     ``ColumnTransformer`` already maintains — so no second transform is needed.
+    ``name`` is the public block name (see :func:`block_name`).
     """
     indices = column_transformer.output_indices_
     slices = []
-    for name, transformer, _columns in column_transformer.transformers_:
+    for name, transformer, columns in column_transformer.transformers_:
         if transformer == "drop":
             continue
         span = indices.get(name)
@@ -41,22 +65,23 @@ def get_output_slices(column_transformer):
         width = span.stop - span.start
         if width == 0:
             continue
-        slices.append((name, span.start, width))
+        slices.append((block_name(name, columns), span.start, width))
     return slices
 
 
 def clean_feature_names(column_transformer, names):
     """Collapse the per-feature name that sklearn's ColumnTransformer duplicates.
 
-    Each per-column step is named ``f"{kind}_{feature}"`` (see ``compose/factory.py``),
-    and every PreTab transformer's own ``get_feature_names_out`` already bakes the
-    input feature name into each output column, so sklearn's default
-    ``f"{step}__{inner}"`` naming doubles it, e.g. ``"num_age__age_bs0"``. This
-    collapses that back to ``"num_age_bs0"``, leaving passthrough/remainder columns
-    and any name it cannot confidently match unchanged.
+    Each per-column step is a ``num_*`` / ``cat_*`` block for one feature (see
+    ``compose/factory.py``), and every PreTab transformer's own
+    ``get_feature_names_out`` already bakes the input feature name into each output
+    column, so sklearn's default ``f"{step}__{inner}"`` naming doubles it, e.g.
+    ``"num_age__age_bs0"``. This collapses that back to ``"num_age_bs0"``. Other
+    names of a block keep the public block name as their prefix, and
+    passthrough/remainder columns are left unchanged.
     """
-    step_to_feature = {
-        name: columns[0]
+    step_to_columns = {
+        name: columns
         for name, _transformer, columns in column_transformer.transformers_
         if name != "remainder" and len(columns) == 1
     }
@@ -64,15 +89,22 @@ def clean_feature_names(column_transformer, names):
     for raw in names:
         raw = str(raw)
         step_name, sep, inner_name = raw.partition("__")
-        feature = step_to_feature.get(step_name)
-        if not sep or feature is None:
+        columns = step_to_columns.get(step_name)
+        if not sep or columns is None:
             cleaned.append(raw)
             continue
-        if inner_name == feature or inner_name.startswith(f"{feature}_"):
+        block = _block(step_name, columns)
+        if block is not None:
+            kind_prefix, feature = block
+            public_name = f"{kind_prefix}_{feature}"
+        else:
+            feature = str(columns[0])
             kind_prefix = step_name[: -(len(feature) + 1)] if step_name.endswith(f"_{feature}") else ""
+            public_name = step_name
+        if inner_name == feature or inner_name.startswith(f"{feature}_"):
             cleaned.append(f"{kind_prefix}_{inner_name}" if kind_prefix else inner_name)
         else:
-            cleaned.append(raw)
+            cleaned.append(f"{public_name}__{inner_name}")
     return cleaned
 
 
diff --git a/pretab/preprocessor.py b/pretab/preprocessor.py
index a8193b8..d08cb8b 100644
--- a/pretab/preprocessor.py
+++ b/pretab/preprocessor.py
@@ -15,6 +15,7 @@
 from .compose.factory import build_column_transformer
 from .compose.feature_detection import detect_column_types, to_dataframe
 from .compose.inspection import (
+    block_name,
     build_feature_info,
     build_feature_lineage,
     build_transformer_summary,
@@ -305,7 +306,11 @@ class Preprocessor(TransformerMixin, BaseEstimator):
     ----------
     column_transformer\_ : ColumnTransformer
         The internal scikit-learn column transformer that handles feature-wise preprocessing.
-        Set when the preprocessor is fitted.
+        Set when the preprocessor is fitted. Each column gets one step, named ``num_<col>``
+        / ``cat_<col>`` whenever scikit-learn accepts that name; a label scikit-learn
+        cannot embed in a step name (one starting with ``_`` or containing ``__``) gets a
+        positional step name such as ``num_col0`` instead. Block keys, feature names and
+        lineage always use the column label.
     n_features_in\_ : int
         Number of input features seen during ``fit``.
     total_output_dim\_ : int
@@ -981,7 +986,8 @@ def get_feature_info(self, verbose=True):
 
     def _log_internal_decisions(self):
         """Log fitted internal decisions (bins / knots / centers) at DEBUG."""
-        for name, transformer, _columns in self.column_transformer_.transformers_:
+        for step_name, transformer, columns in self.column_transformer_.transformers_:
+            name = block_name(step_name, columns)
             last_step = transformer.steps[-1][1] if hasattr(transformer, "steps") else transformer
             for attr in (
                 "thresholds_",
diff --git a/tests/integration/test_column_labels.py b/tests/integration/test_column_labels.py
new file mode 100644
index 0000000..df61362
--- /dev/null
+++ b/tests/integration/test_column_labels.py
@@ -0,0 +1,79 @@
+"""Column labels the ColumnTransformer cannot take verbatim.
+
+Regression guards for issue #71: labels starting with ``_`` or containing ``__``
+were embedded into estimator names that scikit-learn rejects.
+"""
+
+from typing import cast
+
+import numpy as np
+import pandas as pd
+import pytest
+from sklearn.compose import ColumnTransformer
+from sklearn.impute import SimpleImputer
+from sklearn.pipeline import make_pipeline
+
+from pretab import Preprocessor
+
+
+def _blocks(out):
+    return cast("dict[str, np.ndarray]", out)
+
+
+@pytest.mark.parametrize("label", ["_score", "income__log", "__index_level_0__", "trailing__"])
+def test_numerical_labels_with_underscores_fit(label):
+    values = np.random.default_rng(0).normal(size=50)
+    frame = pd.DataFrame({label: values})
+    pre = Preprocessor(numerical_method="minmax").fit(frame)
+    assert pre.get_feature_names_out().tolist() == [f"num_{label}"]
+    blocks = _blocks(pre.transform(frame, return_array=False))
+    assert list(blocks) == [f"num_{label}"]
+    np.testing.assert_allclose(blocks[f"num_{label}"].ravel(), (values - values.min()) / np.ptp(values))
+    assert pre.output_dims_ == {label: 1}
+
+
+def test_categorical_and_expanded_labels_with_underscores_fit():
+    rng = np.random.default_rng(0)
+    frame = pd.DataFrame({"_city": rng.choice(["a", "b"], 50), "x__y": rng.normal(size=50)})
+    pre = Preprocessor(categorical_method="one-hot", output_dim=3).fit(frame, rng.normal(size=50))
+    names = pre.get_feature_names_out().tolist()
+    assert names == ["num_x__y_ple0", "num_x__y_ple1", "num_x__y_ple2", "cat__city_a", "cat__city_b"]
+    assert [record.output_feature for record in pre.get_feature_lineage()] == names
+    assert list(_blocks(pre.transform(frame, return_array=False))) == ["num_x__y", "cat__city"]
+    assert pre.output_dims_ == {"x__y": 3, "_city": 2}
+
+
+def test_positional_step_names_do_not_leak_into_names_or_block_keys():
+    rng = np.random.default_rng(0)
+    frame = pd.DataFrame({"_city": rng.choice(["a", "b"], 50), "x": rng.normal(size=50)})
+    pre = Preprocessor(numerical_method="minmax", categorical_method="one-hot", missing_policy="separate_state")
+    pre.fit(frame)
+    # "_city" runs under an internal positional step name; every public name uses the label.
+    assert [name for name, _, _ in pre.column_transformer_.transformers_] == ["num_x", "cat_col1"]
+    names = [
+        "num_x__representation__x",
+        "num_x__missing__x__missing",
+        "cat__city__representation___city_a",
+        "cat__city__representation___city_b",
+        "cat__city__missing___city__missing",
+    ]
+    assert pre.get_feature_names_out().tolist() == names
+    assert [record.output_feature for record in pre.get_feature_lineage()] == names
+    assert list(_blocks(pre.transform(frame, return_array=False))) == ["num_x", "cat__city"]
+
+
+def test_underscore_labels_do_not_collide_with_other_steps():
+    rng = np.random.default_rng(0)
+    # "_col0" would map to the fallback name of the first step without collision handling.
+    frame = pd.DataFrame({"_a": rng.normal(size=30), "col0": rng.normal(size=30)})
+    pre = Preprocessor(numerical_method="minmax").fit(frame)
+    assert pre.get_feature_names_out().tolist() == ["num__a", "num_col0"]
+
+
+def test_preprocessor_chains_after_a_pandas_column_transformer():
+    rng = np.random.default_rng(0)
+    X = pd.DataFrame({"age": rng.normal(40, 10, 50), "income": rng.normal(50, 10, 50)})
+    upstream = ColumnTransformer([("imp", SimpleImputer(), ["age", "income"])]).set_output(transform="pandas")
+    pipe = make_pipeline(upstream, Preprocessor(numerical_method="minmax")).fit(X, rng.normal(size=50))
+    assert cast(Preprocessor, pipe[-1]).get_feature_names_out().tolist() == ["num_imp__age", "num_imp__income"]
+    assert np.asarray(pipe.transform(X)).shape == (50, 2)
diff --git a/tests/integration/test_verbosity.py b/tests/integration/test_verbosity.py
index 3c44e1e..5f51584 100644
--- a/tests/integration/test_verbosity.py
+++ b/tests/integration/test_verbosity.py
@@ -119,6 +119,17 @@ def test_verbose_3_logs_internal_decisions(sample_data, caplog):
     assert "thresholds_" in debug_text or "total_output_dim_" in debug_text
 
 
+def test_verbose_3_logs_internal_decisions_under_the_block_name(caplog):
+    """Issue #71: a label that cannot be part of a step name is logged by its block name."""
+    rng = np.random.default_rng(0)
+    X = pd.DataFrame({"_score": rng.normal(size=60)})
+    caplog.set_level(logging.DEBUG, logger="pretab")
+    Preprocessor(numerical_method="ple", verbose=3).fit(X, rng.normal(size=60))
+    debug_text = "\n".join(r.getMessage() for r in caplog.records if r.levelno == logging.DEBUG)
+    assert "num__score.total_output_dim_" in debug_text
+    assert "num_col0" not in debug_text
+
+
 def test_verbose_true_behaves_like_level_1(sample_data, caplog):
     X, y = sample_data
     caplog.set_level(logging.DEBUG, logger="pretab")