Repository: scikit-learn/scikit-learn
File: sklearn/tests/test_pipeline.py (Python)

Code before:
    pipe = make_pipeline(StandardScaler(), LogisticRegression())
    pipe.set_output(transform="pandas")
    pipe.fit(X, y)

    feature_names_in_ = pipe[:-1].get_feature_names_out()
    log_reg_feature_names = pipe[-1].feature_names_in_

    assert_array_equal(feature_names_in_, log_reg_feature_names)


def test_feature_union_set_output():
    """Test feature union with set_output API."""
    pd = pytest.importorskip("pandas")

    X, _ = load_iris(as_frame=True, return_X_y=True)
    X_train, X_test = train_test_split(X, random_state=0)
    union = FeatureUnion([("scalar", StandardScaler()), ("pca", PCA())])
    union.set_output(transform="pandas")
    union.fit(X_train)

    X_trans = union.transform(X_test)
    assert isinstance(X_trans, pd.DataFrame)
    assert_array_equal(X_trans.columns, union.get_feature_names_out())
    assert_array_equal(X_trans.index, X_test.index)

    # Test with polars
    pl = pytest.importorskip("polars")
    union.set_output(transform="polars")
    X_trans_pl = union.transform(X_test)
    assert isinstance(X_trans_pl, pl.DataFrame)
    assert_array_equal(X_trans_pl.columns, union.get_feature_names_out())


def test_feature_union_getitem():
    """Check FeatureUnion.__getitem__ returns expected results."""
    scalar = StandardScaler()
    pca = PCA()
    union = FeatureUnion(
        [
            ("scalar", scalar),
            ("pca", pca),
            ("pass", "passthrough"),
            ("drop_me", "drop"),
        ]
    )
    assert union["scalar"] is scalar
    assert union["pca"] is pca
    assert union["pass"] == "passthrough"
    assert union["drop_me"] == "drop"



Review comment (suggestion):
I would rather use `@pytest.mark.parametrize(df_library, ["pandas", "polars"])` so that the full test is not marked as skipped if only one of the two libs is installed.

Code after:
    Non-regresion test for #21349.
    """
    X, y = iris.data, iris.target

    class AddPrefixStandardScalar(StandardScaler):
        def get_feature_names_out(self, input_features=None):
            names = super().get_feature_names_out(input_features=input_features)
            return np.asarray([f"my_prefix_{name}" for name in names], dtype=object)

    pipe = make_pipeline(AddPrefixStandardScalar(), StandardScaler())
    pipe.fit(X, y)

    input_names = iris.feature_names
    feature_names_out = pipe.get_feature_names_out(input_names)

    assert_array_equal(feature_names_out, [f"my_prefix_{name}" for name in input_names])


def test_pipeline_set_output_integration():
    """Test pipeline's set_output with feature names."""
    pytest.importorskip("pandas")

    X, y = load_iris(as_frame=True, return_X_y=True)

    pipe = make_pipeline(StandardScaler(), LogisticRegression())
    pipe.set_output(transform="pandas")
    pipe.fit(X, y)

    feature_names_in_ = pipe[:-1].get_feature_names_out()
    log_reg_feature_names = pipe[-1].feature_names_in_

    assert_array_equal(feature_names_in_, log_reg_feature_names)


@pytest.mark.parametrize("df_library", ["pandas", "polars"])
def test_feature_union_set_output(df_library):
    """Test feature union with set_output API."""
    lib = pytest.importorskip(df_library)

    X, _ = load_iris(as_frame=True, return_X_y=True)
    X_train, X_test = train_test_split(X, random_state=0)
    union = FeatureUnion([("scalar", StandardScaler()), ("pca", PCA())])
    union.set_output(transform=df_library)
    union.fit(X_train)

    X_trans = union.transform(X_test)
    assert isinstance(X_trans, lib.DataFrame)
    assert_array_equal(X_trans.columns, union.get_feature_names_out())
    if df_library == "pandas":
        assert_array_equal(X_trans.index, X_test.index)