From ea0d5bc84e3a5c05e9b1a53fb55a7ec6bd16f269 Mon Sep 17 00:00:00 2001
From: ChrisW09 <50968720+ChrisW09@users.noreply.github.com>
Date: Mon, 27 Jul 2026 22:14:26 +0200
Subject: [PATCH] perf(preprocessor): slice transform blocks from
 output_indices_

The dict form of ``transform`` -- the default -- measured each block's width by
calling ``transformer.transform(X[columns])`` again, so every per-feature
pipeline ran twice per call. On 20k rows x 20 features with PLE that was 43.5 ms
against 21.9 ms for ``return_array=True``, for identical numbers.

``column_transformer_.output_indices_`` already records the exact output slice
for each transformer, and ``output_dims_`` on this same class already reads it.
Use it here too. Also more precise than what it replaces: the ``width = 1``
fallback for a non-estimator entry was a guess.

    return_array=True :  21.9 ms -> 21.9 ms
    return_array=False:  43.5 ms -> 21.7 ms

Closes #20

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 pretab/preprocessor.py     | 19 +++++-----
 tests/test_preprocessor.py | 72 ++++++++++++++++++++++++++++++++++++--
 2 files changed, 80 insertions(+), 11 deletions(-)

diff --git a/pretab/preprocessor.py b/pretab/preprocessor.py
index fcf228a..b53bfae 100644
--- a/pretab/preprocessor.py
+++ b/pretab/preprocessor.py
@@ -457,17 +457,18 @@ def transform(self, X, embeddings=None, return_array=False):
         if return_array:
             return transformed_X
 
+        # ``output_indices_`` already records the exact output slice each
+        # transformer occupies, so the blocks can be sliced straight out of the
+        # stacked array. Re-running ``transformer.transform`` here just to read
+        # ``.shape[1]`` transformed the whole frame a second time, roughly
+        # doubling the cost of the default (dict) path.
+        output_indices = self.column_transformer_.output_indices_
         transformed_dict = {}
-        start = 0
-        for name, transformer, columns in self.column_transformer_.transformers_:
-            if transformer == "drop":
+        for name, _transformer, _columns in self.column_transformer_.transformers_:
+            span = output_indices[name]
+            if span.stop - span.start == 0:
                 continue
-            if hasattr(transformer, "transform"):
-                width = transformer.transform(X[columns]).shape[1]
-            else:
-                width = 1
-            transformed_dict[name] = transformed_X[:, start : start + width]
-            start += width
+            transformed_dict[name] = transformed_X[:, span]
 
         if embeddings is not None:
             if not self.embeddings_:
diff --git a/tests/test_preprocessor.py b/tests/test_preprocessor.py
index 7222f82..c44dd77 100644
--- a/tests/test_preprocessor.py
+++ b/tests/test_preprocessor.py
@@ -1,10 +1,13 @@
-import pytest
+from typing import cast
+
 import numpy as np
 import pandas as pd
+import pytest
 from sklearn.base import clone
 from sklearn.exceptions import NotFittedError
 from sklearn.utils.validation import check_is_fitted
-from pretab.preprocessor import Preprocessor  # Adjust the import as needed
+
+from pretab.preprocessor import Preprocessor
 
 
 @pytest.fixture
@@ -254,3 +257,68 @@ def test_output_dims_and_total_before_fit_raise():
     with pytest.raises(NotFittedError):
         _ = Preprocessor().total_output_dim_
 
+
+
+# --------------------------------------------------------------------------- #
+# The dict form of ``transform`` must not re-run the pipelines.
+#
+# Widths used to be measured by calling ``transformer.transform(X[columns])``
+# again for every block, transforming the whole frame a second time.
+# ``column_transformer_.output_indices_`` already records the exact slices.
+# --------------------------------------------------------------------------- #
+def _fitted_preprocessor():
+    rng = np.random.default_rng(0)
+    frame = pd.DataFrame(
+        {
+            "a": rng.normal(size=200),
+            "b": rng.normal(size=200),
+            "c": rng.choice(["x", "y", "z"], size=200),
+        }
+    )
+    y = rng.normal(size=200)
+    return Preprocessor(numerical_method="rbf", categorical_method="one-hot").fit(frame, y), frame
+
+
+def test_dict_transform_calls_each_pipeline_once():
+    pre, frame = _fitted_preprocessor()
+    calls: dict[str, int] = {}
+
+    for name, pipe, _cols in pre.column_transformer_.transformers_:
+        original = pipe.transform
+
+        def counting(Z, _name=name, _original=original):
+            calls[_name] = calls.get(_name, 0) + 1
+            return _original(Z)
+
+        pipe.transform = counting
+
+    pre.transform(frame)
+
+    assert calls and set(calls.values()) == {1}
+
+
+def test_dict_blocks_reassemble_the_stacked_array():
+    pre, frame = _fitted_preprocessor()
+
+    blocks = cast("dict[str, np.ndarray]", pre.transform(frame))
+    stacked = cast("np.ndarray", pre.transform(frame, return_array=True))
+
+    np.testing.assert_allclose(
+        np.hstack([blocks[name] for name in blocks]).astype(float), stacked.astype(float)
+    )
+
+
+def test_dict_block_widths_match_output_dims():
+    pre, frame = _fitted_preprocessor()
+
+    blocks = cast("dict[str, np.ndarray]", pre.transform(frame))
+    widths = {name.split("_", 1)[1]: block.shape[1] for name, block in blocks.items()}
+
+    assert widths == pre.output_dims_
+    assert sum(widths.values()) == len(pre.get_feature_names_out())
+
+
+def test_dict_keys_cover_every_input_feature():
+    pre, frame = _fitted_preprocessor()
+
+    assert sorted(pre.transform(frame)) == ["cat_c", "num_a", "num_b"]