fix(compose): support object dtype in ColumnTransformer + SimpleImputer on cuDF

python/cuml/cuml/_thirdparty/sklearn/preprocessing/_column_transformer.py

@@ -25,6 +25,7 @@
 import cudf
 import cupy as np
 import numba
+import numpy as cpu_np
 import pandas as pd
 import scipy.sparse as sp_sparse
 from cupyx.scipy import sparse as cu_sparse
@@ -35,6 +36,7 @@
 
 import cuml
 from cuml.internals.global_settings import _global_settings_data
+from cuml.internals.outputs import _is_object_dtype
 from cuml.internals.validation import check_is_fitted, check_features, check_array
 
 from ..preprocessing._function_transformer import FunctionTransformer
@@ -982,6 +984,12 @@ def _hstack(self, Xs):
             return cu_sparse.hstack(converted_Xs).tocsr()
         else:
             Xs = [f.toarray() if issparse(f) else f for f in Xs]
+            if any(_is_object_dtype(X) for X in Xs):
+                # Object dtype (e.g. string columns from a categorical
+                # SimpleImputer) has no device representation - cupy has no
+                # way to store it. Stack on host instead.
+                Xs = [X.get() if isinstance(X, np.ndarray) else X for X in Xs]
+                return cpu_np.hstack(Xs)
             return np.hstack(Xs)
 
 

python/cuml/cuml/_thirdparty/sklearn/preprocessing/_imputation.py

@@ -29,7 +29,7 @@
     StringInputTagMixin,
     _ensure_transformer_tags,
 )
-from cuml.internals.outputs import mlfunc, ReflectedAttr
+from cuml.internals.outputs import _is_object_dtype, mlfunc, ReflectedAttr
 from cuml.internals.validation import check_is_fitted, check_inputs
 
 from ....thirdparty_adapters import (
@@ -140,15 +140,24 @@ def _concatenate_indicator(self, X_imputed, X_indicator):
         if not self.add_indicator:
             return X_imputed
 
-        hstack = sparse.hstack if sparse.issparse(X_imputed) else np.hstack
         if X_indicator is None:
             raise ValueError(
                 "Data from the missing indicator are not provided. Call "
                 "_fit_indicator and _transform_indicator in the imputer "
                 "implementation."
             )
 
-        return hstack((X_imputed, X_indicator))
+        if sparse.issparse(X_imputed):
+            return sparse.hstack((X_imputed, X_indicator))
+
+        if _is_object_dtype(X_imputed) or _is_object_dtype(X_indicator):
+            arrays = [
+                array.get() if isinstance(array, np.ndarray) else array
+                for array in (X_imputed, X_indicator)
+            ]
+            return cpu_np.hstack(arrays)
+
+        return np.hstack((X_imputed, X_indicator))
 
     def __sklearn_tags__(self):
         tags = super().__sklearn_tags__()
@@ -260,10 +269,12 @@ def __init__(self, *, missing_values=np.nan, strategy="mean",
     @classmethod
     def _get_param_names(cls):
         return super()._get_param_names() + [
+            "missing_values",
             "strategy",
             "fill_value",
             "verbose",
-            "copy"
+            "copy",
+            "add_indicator",
         ]
 
     def _validate_input(self, X, in_fit):
@@ -423,7 +434,11 @@ def _dense_fit(self, X, strategy, missing_values, fill_value):
 
         # Constant
         elif strategy == "constant":
-            return np.full(X.shape[1], fill_value, dtype=X.dtype)
+            # Object-dtype X (e.g. strings) lives on host, since cupy has no
+            # way to represent it on device. Dispatch to X's own module so
+            # this works for both device (numeric) and host (object) input.
+            xp = np.get_array_module(X)
+            return xp.full(X.shape[1], fill_value, dtype=X.dtype)
 
     @mlfunc
     def transform(self, X):
@@ -449,14 +464,18 @@ def transform(self, X):
         if self.strategy == "constant":
             valid_statistics = statistics
         else:
+            # Object-dtype statistics (e.g. strings) live on host, since cupy
+            # has no way to represent them on device. Dispatch to their own
+            # module so this works for both device and host input.
+            xp = np.get_array_module(statistics)
             # same as np.isnan but also works for object dtypes
             invalid_mask = _get_mask(statistics, np.nan)
-            valid_mask = np.logical_not(invalid_mask)
+            valid_mask = xp.logical_not(invalid_mask)
             valid_statistics = statistics[valid_mask]
-            valid_statistics_indexes = np.flatnonzero(valid_mask)
+            valid_statistics_indexes = xp.flatnonzero(valid_mask)
 
             if invalid_mask.any():
-                missing = np.arange(X.shape[1])[invalid_mask]
+                missing = xp.arange(X.shape[1])[invalid_mask]
                 if self.verbose:
                     warnings.warn("Deleting features without "
                                   "observed values: %s" % missing)
@@ -481,8 +500,9 @@ def transform(self, X):
             if self.strategy == "constant":
                 X[mask] = valid_statistics[0]
             else:
+                xp = np.get_array_module(mask)
                 for i, vi in enumerate(valid_statistics_indexes):
-                    feature_idxs = np.flatnonzero(mask[:, vi])
+                    feature_idxs = xp.flatnonzero(mask[:, vi])
                     X[feature_idxs, vi] = valid_statistics[i]
 
         X = super()._concatenate_indicator(X, X_indicator)

python/cuml/cuml/internals/outputs.py

@@ -326,6 +326,28 @@ def infer_output_type(array, array_like="numpy"):
     return None
 
 
+def _is_object_dtype(res):
+    """Check for NumPy object dtype on an array-like or dataframe-like."""
+    dtype = getattr(res, "dtype", None)
+    if dtype is not None:
+        # Array-like or Series: a single (possibly extension) dtype. Only a
+        # plain numpy object dtype has no device representation; extension
+        # dtypes are a separate, unsupported case handled elsewhere.
+        return isinstance(dtype, np.dtype) and dtype.kind == "O"
+
+    dtypes = getattr(res, "dtypes", None)
+    if dtypes is not None:
+        # DataFrame-like: dtypes is per-column.
+        try:
+            return any(
+                isinstance(dt, np.dtype) and dt.kind == "O" for dt in dtypes
+            )
+        except TypeError:
+            return False
+
+    return False
+
+
 class ArrayIndexPair:
     """An array paired with an aligned index.
 
@@ -534,6 +556,57 @@ def convert_arrays(
                     return pd.Series(obj.flatten(), index=index)
                 return pd.DataFrame(obj, index=index)
             return pd.Series(obj, index=index)
+        elif _is_object_dtype(obj):
+            # NumPy object arrays have no device representation. Preserve
+            # host arrays for "array" and internal "cuml" outputs, or wrap
+            # them in a dataframe-like type.
+            if output_type in ("array", "cuml"):
+                return obj
+            elif output_type in ("cupy", "numba"):
+                raise TypeError(
+                    f"{output_type=!r} doesn't support outputs of dtype "
+                    f"object and shape {obj.shape}"
+                )
+
+            if hasattr(index, "to_pandas"):
+                index = index.to_pandas()
+            if output_type == "series":
+                if obj.ndim == 2:
+                    if obj.shape[1] == 1:
+                        obj = obj.flatten()
+                    else:
+                        raise ValueError(
+                            "Only single dimensional arrays can be transformed to"
+                            " Series."
+                        )
+                elif obj.ndim == 0:
+                    obj = obj[None]
+            elif output_type == "dataframe":
+                if obj.ndim == 1:
+                    obj = obj[:, None]
+                elif obj.ndim == 0:
+                    obj = obj[None, None]
+
+            if obj.ndim == 2:
+                if (
+                    one_col_2d_as_series
+                    and obj.shape[1] == 1
+                    and output_type != "dataframe"
+                ):
+                    host_df = pd.Series(obj.flatten(), index=index)
+                else:
+                    host_df = pd.DataFrame(obj, index=index)
+            else:
+                host_df = pd.Series(obj, index=index)
+
+            try:
+                return cudf.from_pandas(host_df)
+            except (TypeError, ValueError, NotImplementedError):
+                if isinstance(host_df, pd.DataFrame):
+                    # cudf cannot represent every mixed object DataFrame.
+                    # Preserve the host DataFrame instead of losing its data.
+                    return host_df
+                raise
         else:
             # Other output types use device memory, coerce to cupy and take
             # cupy code path.

python/cuml/cuml/internals/validation.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 import numbers
@@ -697,6 +697,15 @@ def check_array(
         # No original dtype, use first provided dtype
         dtype = dtype[0]
 
+    # Object dtype (e.g. strings) has no device-memory representation - cupy
+    # has no way to store arbitrary Python objects on the GPU. Force host
+    # output in that case regardless of the requested `mem_type`, rather than
+    # crashing when attempting the device conversion below. This is required
+    # for estimators (e.g. SimpleImputer) that explicitly support categorical
+    # / string data with an object dtype.
+    if isinstance(dtype, np.dtype) and dtype.kind == "O":
+        mem_type = "host"
+
     # Coerce `array` to numpy/cupy/scipy.sparse/cupyx.scipy.sparse values as
     # requested. For dataframe-like inputs also extract the index for later use.
     index = None

python/cuml/cuml/thirdparty_adapters/adapters.py

@@ -1,15 +1,35 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 import cupy as cp
 import numpy as np
+import pandas as pd
+
+
+def _is_na_sentinel(value):
+    """Return whether a scalar missing-value sentinel is NA-like."""
+    if isinstance(value, str):
+        return value == "NaN"
+    return pd.isna(value)
 
 
 def _get_mask(X, value_to_mask):
     """Compute the boolean mask X == missing_values."""
-    if value_to_mask == "NaN" or cp.isnan(value_to_mask):
-        return cp.isnan(X)
+    if isinstance(value_to_mask, str) and value_to_mask == "NaN":
+        if isinstance(X, cp.ndarray):
+            return cp.isnan(X)
+        return pd.isna(X)
+    # NaN-like sentinels (np.nan, None, pd.NA, pd.NaT, ...) require an
+    # NA-aware comparison: a plain `==` against e.g. pd.NA propagates
+    # instead of returning a boolean mask, and `cp.isnan` doesn't accept
+    # non-numeric scalars like pd.NA in the first place.
+    if _is_na_sentinel(value_to_mask):
+        if isinstance(X, cp.ndarray):
+            return cp.isnan(X)
+        # Host (e.g. object dtype) arrays can't use isnan - fall back to an
+        # NA-aware elementwise check that also recognizes None/pd.NA.
+        return pd.isna(X)
     else:
         return X == value_to_mask
 
@@ -20,7 +40,7 @@ def _masked_column_median(arr, masked_value):
     mask = _get_mask(arr, masked_value)
     if arr.size == 0:
         return cp.full(arr.shape[1], cp.nan)
-    if not cp.isnan(masked_value):
+    if not _is_na_sentinel(masked_value):
         arr_sorted = arr.copy()
         # If nan is not the missing value, any column with nans should
         # have a median of nan
@@ -59,7 +79,7 @@ def _masked_column_mean(arr, masked_value):
     count_missing_values = mask.sum(axis=0)
     n_elems = arr.shape[0] - count_missing_values
     mean = cp.nansum(arr, axis=0)
-    if not cp.isnan(masked_value):
+    if not _is_na_sentinel(masked_value):
         mean -= count_missing_values * masked_value
     mean /= n_elems
     return mean
@@ -68,18 +88,22 @@ def _masked_column_mean(arr, masked_value):
 def _masked_column_mode(arr, masked_value):
     """Determine the most frequently appearing element in each column in the 2D
     array arr, ignoring any instances of masked_value"""
+    # Object-dtype arrays (e.g. strings) live on host, since cupy has no way
+    # to represent them on device. Dispatch to the array's own module so this
+    # works for both device (numeric) and host (object dtype) input.
+    xp = cp.get_array_module(arr)
     mask = _get_mask(arr, masked_value)
     n_features = arr.shape[1]
     most_frequent = np.empty(n_features, dtype=arr.dtype)
     for i in range(n_features):
-        feature_mask_idxs = cp.where(~mask[:, i])[0]
-        values, counts = cp.unique(
+        feature_mask_idxs = xp.where(~mask[:, i])[0]
+        values, counts = xp.unique(
             arr[feature_mask_idxs, i], return_counts=True
         )
         count_max = counts.max()
         if count_max > 0:
             value = values[counts == count_max].min()
         else:
-            value = cp.nan
+            value = xp.nan
         most_frequent[i] = value
-    return cp.array(most_frequent)
+    return xp.array(most_frequent)

python/cuml/tests/test_adapters.py

@@ -1,5 +1,5 @@
 #
-# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 #
 
@@ -8,6 +8,7 @@
 import cupy as cp
 import cupyx as cpx
 import numpy as np
+import pandas as pd
 import pytest
 from scipy import stats
 from sklearn.utils._mask import _get_mask as sk_get_mask
@@ -110,6 +111,30 @@ def test_get_mask(failure_logger, mask_dataset):
     assert_allclose(cu_mask, sk_mask)
 
 
+def test_get_mask_nan_string_on_host_object_array():
+    X = np.array([1.0, np.nan, "present"], dtype=object)
+
+    np.testing.assert_array_equal(
+        cu_get_mask(X, value_to_mask="NaN"), [False, True, False]
+    )
+
+
+@pytest.mark.parametrize(
+    ("function", "expected"),
+    [
+        (_masked_column_mean, [3.0, 5.0]),
+        (_masked_column_median, [3.0, 5.0]),
+    ],
+)
+@pytest.mark.parametrize("missing_value", [pd.NA, "NaN"])
+def test_masked_column_numeric_na_sentinel(function, expected, missing_value):
+    X = cp.array([[1.0, cp.nan], [3.0, 4.0], [5.0, 6.0]])
+
+    result = function(X, missing_value)
+
+    np.testing.assert_allclose(result.get(), expected)
+
+
 def test_masked_column_median(failure_logger, mask_dataset):
     mask_value, X_np, X = mask_dataset
     median = _masked_column_median(X, mask_value).get()