fix(output): preserve input index in set_output pandas container

pretab/compose/output.py

@@ -114,7 +114,7 @@ def compute_output_report(array, output_format, *, threshold=_SPARSE_AUTO_THRESH
     return fmt, report
 
 
-def to_dataframe_output(array, columns, container):
+def to_dataframe_output(array, columns, container, *, index=None):
     """Wrap a dense stacked array in a pandas or polars DataFrame.
 
     Parameters
@@ -126,6 +126,11 @@ def to_dataframe_output(array, columns, container):
         One name per output column (from ``get_feature_names_out``).
     container : {"pandas", "polars"}
         Target dataframe library.
+    index : pandas.Index, optional
+        Row labels for the pandas container, normally the input's index so the
+        output stays aligned with it (as scikit-learn's own transformers do).
+        ``None`` gives a default ``RangeIndex``. Ignored for polars, which has no
+        index.
 
     Raises
     ------
@@ -138,7 +143,7 @@ def to_dataframe_output(array, columns, container):
     if container == "pandas":
         import pandas as pd
 
-        return pd.DataFrame(array, columns=pd.Index(columns))
+        return pd.DataFrame(array, index=index, columns=pd.Index(columns))
     try:
         import polars as pl  # type: ignore
     except ImportError as exc:  # pragma: no cover - exercised only without polars

pretab/preprocessor.py

@@ -631,7 +631,8 @@ def transform(self, X, embeddings=None, return_array=None):
             "matrix"``); a dict of per-feature blocks when ``output_structure="blocks"``
             (or ``return_array=False``); a SciPy CSR matrix (or CSR blocks) when
             ``output_format`` resolves to ``"sparse"``; or a pandas / polars DataFrame
-            when configured via :meth:`set_output`.
+            when configured via :meth:`set_output` (a pandas DataFrame keeps the index
+            of ``X``).
         """
 
         check_is_fitted(self)
@@ -656,7 +657,9 @@ def transform(self, X, embeddings=None, return_array=None):
         fmt, self.output_report_ = compute_output_report(transformed_X, self.output_format)
 
         if container in ("pandas", "polars"):
-            return to_dataframe_output(transformed_X, self.get_feature_names_out(), container)
+            # Pass X's index explicitly: sklearn's set_output wrapper keeps the index of
+            # a returned DataFrame instead of restoring the input's.
+            return to_dataframe_output(transformed_X, self.get_feature_names_out(), container, index=X.index)
 
         slices = None if resolved_return_array else get_output_slices(self.column_transformer_)
         return format_output(

tests/compose/test_output.py

@@ -1,9 +1,16 @@
 """Unit tests for :mod:`pretab.compose.output`."""
 
 import numpy as np
+import pandas as pd
 import pytest
 
-from pretab.compose.output import attach_embeddings, build_output_dict, format_output, resolve_embedding_dimensions
+from pretab.compose.output import (
+    attach_embeddings,
+    build_output_dict,
+    format_output,
+    resolve_embedding_dimensions,
+    to_dataframe_output,
+)
 from pretab.exceptions import IncompatibleParamsError, PretabDataError
 
 
@@ -109,3 +116,11 @@ def test_format_output_rejects_embeddings_with_array_output():
             embeddings=np.ones((2, 3)),
             embeddings_expected=True,
         )
+
+
+def test_to_dataframe_output_pandas_uses_given_index():
+    """Regression guard for issue #60: the pandas container carries the passed row labels."""
+    index = pd.Index([10, 30, 20])
+    out = to_dataframe_output(np.arange(6).reshape(3, 2), ["a", "b"], "pandas", index=index)
+    pd.testing.assert_index_equal(out.index, index)
+    pd.testing.assert_index_equal(to_dataframe_output(np.zeros((3, 2)), ["a", "b"], "pandas").index, pd.RangeIndex(3))

tests/integration/test_output_format.py

@@ -9,6 +9,9 @@
 import pandas as pd
 import pytest
 from scipy import sparse as sp
+from sklearn.compose import ColumnTransformer
+from sklearn.pipeline import FeatureUnion
+from sklearn.preprocessing import StandardScaler
 
 from pretab import Preprocessor
 from pretab.exceptions import OptionalDependencyError
@@ -20,6 +23,12 @@ def frame():
     return pd.DataFrame({"a": rng.random(30), "b": rng.random(30)})
 
 
+@pytest.fixture
+def indexed_frame(frame):
+    # Non-default, non-monotonic index, like the one train_test_split leaves behind.
+    return frame.set_axis(np.arange(len(frame))[::-1] + 100)
+
+
 @pytest.fixture
 def y():
     rng = np.random.default_rng(1)
@@ -186,6 +195,34 @@ def test_set_output_pandas_fit_transform(frame, y):
     assert out.shape[1] == p.total_output_dim_
 
 
+def test_set_output_pandas_preserves_input_index(indexed_frame, y):
+    """Regression guard for issue #60: the DataFrame keeps X's index, as sklearn's transformers do."""
+    p = _bspline().set_output(transform="pandas")
+    pd.testing.assert_index_equal(p.fit_transform(indexed_frame, y).index, indexed_frame.index)
+    pd.testing.assert_index_equal(p.transform(indexed_frame).index, indexed_frame.index)
+
+
+def test_set_output_pandas_inside_column_transformer(indexed_frame):
+    """Regression guard for issue #60: a reset index made the concatenation raise."""
+    ct = ColumnTransformer(
+        [("pre", Preprocessor(numerical_method="minmax"), ["a"]), ("sc", StandardScaler(), ["b"])]
+    ).set_output(transform="pandas")
+    out = ct.fit_transform(indexed_frame)
+    assert isinstance(out, pd.DataFrame)
+    assert out.shape[0] == len(indexed_frame)
+    pd.testing.assert_index_equal(out.index, indexed_frame.index)
+
+
+def test_set_output_pandas_inside_feature_union(indexed_frame):
+    """Regression guard for issue #60: a reset index made FeatureUnion return 2n NaN-padded rows."""
+    fu = FeatureUnion([("pre", Preprocessor(numerical_method="minmax")), ("sc", StandardScaler())])
+    out = fu.set_output(transform="pandas").fit_transform(indexed_frame)
+    assert isinstance(out, pd.DataFrame)
+    assert out.shape[0] == len(indexed_frame)
+    assert not out.isna().to_numpy().any()
+    pd.testing.assert_index_equal(out.index, indexed_frame.index)
+
+
 def test_set_output_default_still_array(frame, y):
     # set_output(transform="default") only opts out of pandas/polars wrapping; it
     # does not override output_structure, which still resolves to an array.