diff --git a/pretab/compose/output.py b/pretab/compose/output.py index 5a0223c..2ab992e 100644 --- a/pretab/compose/output.py +++ b/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 diff --git a/pretab/preprocessor.py b/pretab/preprocessor.py index a8193b8..d70473e 100644 --- a/pretab/preprocessor.py +++ b/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( diff --git a/tests/compose/test_output.py b/tests/compose/test_output.py index 5d583ff..d3ce644 100644 --- a/tests/compose/test_output.py +++ b/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)) diff --git a/tests/integration/test_output_format.py b/tests/integration/test_output_format.py index abe2613..bc6feb5 100644 --- a/tests/integration/test_output_format.py +++ b/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.