Fix wrapping metaestimators in Pipeline in cuml.accel

[jcrist]

Our pipeline data transfer optimization didn't work if any of the steps were other compositional metaestimators that wrapped accelerated estimators (since these could then accidentally use the accelerated versions, resulting in a mix of cupy and numpy results).

This PR patches the other two compositional estimators (FeatureUnion and ColumnTransformer) so they always run within a numpy output-type context.

Fixes #8112.
Fixes a few sklearn examples as well (yay!)

[coderabbitai]

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes

  • Enhanced support for scikit-learn's ColumnTransformer and FeatureUnion components when used within cuML-accelerated pipelines, ensuring consistent NumPy output formatting and reducing unnecessary device-to-host memory transfers during pipeline execution.

• Tests

  • Added integration tests to validate that ColumnTransformer and FeatureUnion work correctly with other cuML-accelerated components in complex pipeline workflows.

Walkthrough

This PR adds import-time patching for sklearn.compose.ColumnTransformer and extends sklearn.pipeline.FeatureUnion patching to force NumPy output types. It registers these patches in the core accelerator and validates the integration with new pipeline composition tests, resolving the issue where CuPy outputs from accelerated steps break non-accelerated downstream steps.

Changes

ColumnTransformer and FeatureUnion NumPy output patching

Layer / File(s)	Summary
ColumnTransformer and FeatureUnion patch implementations python/cuml/cuml/accel/_patches/sklearn/compose.py, python/cuml/cuml/accel/_patches/sklearn/pipeline.py	Adds new compose.py module that wraps ColumnTransformer.fit, fit_transform, and transform to execute under using_output_type("numpy"). Expands pipeline.py imports to include FeatureUnion and using_output_type, renames patch_method to patch_pipeline_method, and introduces patch_feature_union_method to wrap FeatureUnion methods with the same NumPy output context. Module exports updated to include both Pipeline and FeatureUnion.
Patch registration in accelerator core python/cuml/cuml/accel/core.py	Registers sklearn.compose in the _PATCHES set, enabling ColumnTransformer patching to be activated during ACCEL.register initialization.
Pipeline integration tests for patched estimators python/cuml/cuml_accel_tests/test_pipeline.py	Adds imports for ColumnTransformer, FeatureUnion, and RobustScaler. Introduces two sklearn≥1.8-gated tests: test_column_transfomer_in_pipeline_works() validates a Pipeline with ColumnTransformer (TruncatedSVD + passthrough) produces NumPy output and chains through RobustScaler and Ridge; test_feature_union_in_pipeline_works() does the same for FeatureUnion (TruncatedSVD + PCA) composition.
Test expectation configuration update python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-examples.yaml	Replaces the prior cuml_accel_bugs xfail entry for implicit CuPy-to-NumPy conversion in ColumnTransformer/Pipeline with a new entry for CUDA OOM in PCA concurrent execution, and updates target tests to decomposition::plot_faces_decomposition.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Possibly related PRs

• rapidsai/cuml#7835: Both PRs modify python/cuml/cuml/accel/_patches/sklearn/pipeline.py to patch sklearn.pipeline.Pipeline (and FeatureUnion) methods—wrapping fit/fit_transform/transform with forced/aware output-type handling for consistent NumPy results.
• rapidsai/cuml#7960: Both PRs touch the same upstream scikit-learn xfail-examples.yaml configuration—main PR updates the ColumnTransformer/Pipeline-related xfail entries to match the new Pipeline/ColumnTransformer/FeatureUnion behavior.

Suggested labels

Cython / Python

Suggested reviewers

• csadorf
• dantegd
• divyegala

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 66.67% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly and specifically describes the main change: fixing the wrapping of metaestimators (ColumnTransformer and FeatureUnion) in Pipeline within cuml.accel.
Description check	✅ Passed	The PR description is directly related to the changeset, explaining the bug fix for pipeline optimization issues with compositional metaestimators and referencing the relevant issue #8112.
Linked Issues check	✅ Passed	All code changes fully address issue #8112: patching ColumnTransformer and FeatureUnion methods [compose.py, pipeline.py], registering sklearn.compose for patching [core.py], and adding test coverage [test_pipeline.py].
Out of Scope Changes check	✅ Passed	All changes are in scope: patches to compositional metaestimators, registration of sklearn.compose module, test coverage for the fix, and removal of an outdated XFAIL entry unrelated to the current bug.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

🧹 Nitpick comments (1)

python/cuml/cuml_accel_tests/test_pipeline.py (1)

362-385: ⚡ Quick win

Add explicit NumPy-boundary assertions and cover predict path.

Both new tests currently stop at fit, so they don’t directly verify the patched transform path or that downstream estimator inputs remain NumPy. Please assert input/output array types during both fit and predict (as done in earlier pipeline tests).

Proposed test hardening

 `@requires_sklearn_18`
-def test_column_transfomer_in_pipeline_works():
+def test_column_transfomer_in_pipeline_works(patch_methods):
     """Ensure outputs of steps in `ColumnTransformer` return as numpy"""
+    patch_methods(Ridge, "fit", "predict")
     rng = np.random.default_rng(0)
     X = rng.standard_normal((200, 20)).astype(np.float32)
     y = rng.standard_normal(200).astype(np.float32)
@@
     pipe = Pipeline(
         [
             ("ct", ct),  # Shouldn't be accelerated
             ("scaler", RobustScaler()),  # Not accelerated
             ("ridge", Ridge()),  # Accelerated
         ]
     )

     pipe.fit(X, y)
+    assert isinstance(Ridge.fit.args[0], np.ndarray)
+    out = pipe.predict(X[:10])
+    assert isinstance(Ridge.predict.args[0], np.ndarray)
+    assert isinstance(out, np.ndarray)
@@
 `@requires_sklearn_18`
-def test_feature_union_in_pipeline_works():
+def test_feature_union_in_pipeline_works(patch_methods):
     """Ensure outputs of steps in `FeatureUnion` return as numpy"""
+    patch_methods(Ridge, "fit", "predict")
     rng = np.random.default_rng(0)
     X = rng.standard_normal((200, 20)).astype(np.float32)
     y = rng.standard_normal(200).astype(np.float32)
@@
     pipe = Pipeline(
         [
             ("features", union),  # Shouldn't be accelerated
             ("scaler", RobustScaler()),  # Not accelerated
             ("ridge", Ridge()),  # Accelerated
         ]
     )

     pipe.fit(X, y)
+    assert isinstance(Ridge.fit.args[0], np.ndarray)
+    out = pipe.predict(X[:10])
+    assert isinstance(Ridge.predict.args[0], np.ndarray)
+    assert isinstance(out, np.ndarray)

As per coding guidelines: python/**/test_*.py: “Test files must ... test fit/predict/transform consistency”.

Also applies to: 387-409

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@python/cuml/cuml_accel_tests/test_pipeline.py` around lines 362 - 385, The
test test_column_transfomer_in_pipeline_works stops at pipe.fit and must assert
NumPy-boundaries and cover predict: during fit, assert that data passed into
ColumnTransformer ct and into downstream steps (e.g., "scaler" and "ridge") are
NumPy arrays (use isinstance(..., np.ndarray)) and that ct.transform output is
NumPy; then call pipe.predict(X) and assert the inputs received by downstream
steps and the final prediction output are also NumPy arrays; mirror these same
fit/predict assertions in the sibling test around lines 387-409 to ensure both
transform and predict paths are covered.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@python/cuml/cuml_accel_tests/test_pipeline.py`:
- Around line 362-385: The test test_column_transfomer_in_pipeline_works stops
at pipe.fit and must assert NumPy-boundaries and cover predict: during fit,
assert that data passed into ColumnTransformer ct and into downstream steps
(e.g., "scaler" and "ridge") are NumPy arrays (use isinstance(..., np.ndarray))
and that ct.transform output is NumPy; then call pipe.predict(X) and assert the
inputs received by downstream steps and the final prediction output are also
NumPy arrays; mirror these same fit/predict assertions in the sibling test
around lines 387-409 to ensure both transform and predict paths are covered.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3b5ee11a-f17e-41cb-87fc-74625b64c69c

📥 Commits

Reviewing files that changed from the base of the PR and between ab98e80 and 333e879.

📒 Files selected for processing (5)

• python/cuml/cuml/accel/_patches/sklearn/compose.py
• python/cuml/cuml/accel/_patches/sklearn/pipeline.py
• python/cuml/cuml/accel/core.py
• python/cuml/cuml_accel_tests/test_pipeline.py
• python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-examples.yaml

💤 Files with no reviewable changes (1)

• python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-examples.yaml

[csadorf]

Thanks for the quick turn-around!!

[jcrist]

/merge

[betatim]

Whoop whoop