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

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

@@ -37,6 +37,7 @@
 import cuml
 from cuml.internals.global_settings import _global_settings_data
 from cuml.internals.mixins import DeprecatedGetFeatureNamesMixin
+from cuml.internals.outputs import _is_object_dtype
 from cuml.internals.validation import (
     check_is_fitted,
     check_features,
@@ -300,7 +301,8 @@ def _list_indexing(X, key, key_dtype):
 
 
 def _transform_one(transformer, X, y, weight, **fit_params):
-    with cuml.using_output_type("cupy"):
+    output_type = "numpy" if _is_object_dtype(X) else "cupy"
+    with cuml.using_output_type(output_type):
         res = transformer.transform(X)
 
     # if we have a weight for this transformer, multiply output
@@ -322,7 +324,8 @@ def _fit_transform_one(transformer,
     be multiplied by ``weight``.
     """
     with _print_elapsed_time(message_clsname, message):
-        with cuml.using_output_type("cupy"):
+        output_type = "numpy" if _is_object_dtype(X) else "cupy"
+        with cuml.using_output_type(output_type):
             transformer.accept_sparse = True
             if hasattr(transformer, 'fit_transform'):
                 res = transformer.fit_transform(X, y, **fit_params)
@@ -1014,6 +1017,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

@@ -20,6 +20,7 @@
 
 import cupy as np
 import numpy as cpu_np
+import pandas as pd
 from cupyx.scipy import sparse
 
 import cuml
@@ -29,11 +30,11 @@
     StringInputTagMixin,
     _ensure_transformer_tags,
 )
-from cuml.internals.outputs import mlfunc, ReflectedAttr
+from cuml.internals.outputs import ReflectedAttr, _is_object_dtype, mlfunc
 from cuml.internals.validation import (
-    check_is_fitted,
-    check_inputs,
     check_input_features,
+    check_inputs,
+    check_is_fitted,
 )
 
 from ....thirdparty_adapters import (
@@ -144,15 +145,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__()
@@ -264,10 +274,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):
@@ -288,13 +300,15 @@ def _validate_input(self, X, in_fit):
             ensure_all_finite = "allow-nan"
 
         try:
+            mem_type = "host" if _is_object_dtype(X) else "device"
             X = check_inputs(
                 self,
                 X,
                 accept_sparse="csc",
                 dtype=dtype,
                 ensure_all_finite=ensure_all_finite,
                 copy=self.copy,
+                mem_type=mem_type,
                 reset=in_fit,
             )
         except ValueError as ve:
@@ -305,6 +319,18 @@ def _validate_input(self, X, in_fit):
             else:
                 raise ve
 
+        # Object nulls reach this host path as NaN. Restore pd.NA for its
+        # identity-based path, but reject NaN when None was requested,
+        # matching scikit-learn's validation semantics.
+        if _is_object_dtype(X):
+            if hasattr(X, "flags") and not X.flags.writeable:
+                X = X.copy()
+            nan_mask = _get_mask(X, np.nan)
+            if self.missing_values is pd.NA:
+                X[nan_mask] = pd.NA
+            elif self.missing_values is None and nan_mask.any():
+                raise ValueError("Input contains NaN")
+
         _check_inputs_dtype(X, self.missing_values)
         if X.dtype.kind not in ("i", "u", "f", "O"):
             raise ValueError("SimpleImputer does not support data with dtype "
@@ -427,7 +453,8 @@ 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)
+            xp = np.get_array_module(X)
+            return xp.full(X.shape[1], fill_value, dtype=X.dtype)
 
     @mlfunc
     def transform(self, X):
@@ -443,24 +470,27 @@ def transform(self, X):
         X = self._validate_input(X, in_fit=False)
         X_indicator = super()._transform_indicator(X)
 
-        statistics = self.statistics_
+        # Use the stored value for internal computation. ColumnTransformer
+        # may request cupy output, which cannot reflect object statistics.
+        statistics = type(self).statistics_.get_raw(self)
 
         if X.shape[1] != statistics.shape[0]:
             raise ValueError("X has %d features per sample, expected %d"
-                             % (X.shape[1], self.statistics_.shape[0]))
+                             % (X.shape[1], statistics.shape[0]))
 
         # Delete the invalid columns if strategy is not constant
         if self.strategy == "constant":
             valid_statistics = statistics
         else:
+            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)
@@ -485,9 +515,10 @@ def transform(self, X):
             if self.strategy == "constant":
                 X[mask] = valid_statistics[0]
             else:
-                for i, vi in enumerate(valid_statistics_indexes):
-                    feature_idxs = np.flatnonzero(mask[:, vi])
-                    X[feature_idxs, vi] = valid_statistics[i]
+                xp = np.get_array_module(mask)
+                for i in range(valid_statistics.shape[0]):
+                    feature_idxs = xp.flatnonzero(mask[:, i])
+                    X[feature_idxs, i] = valid_statistics[i]
 
         X = super()._concatenate_indicator(X, X_indicator)
         return X
@@ -508,7 +539,11 @@ def get_feature_names_out(self, input_features=None):
         """
         check_is_fitted(self)
         input_features = check_input_features(self, input_features)
-        non_missing_mask = np.logical_not(_get_mask(self.statistics_, np.nan)).get()
+        statistics = type(self).statistics_.get_raw(self)
+        xp = np.get_array_module(statistics)
+        non_missing_mask = xp.logical_not(_get_mask(statistics, np.nan))
+        if isinstance(non_missing_mask, np.ndarray):
+            non_missing_mask = non_missing_mask.get()
         names = input_features[non_missing_mask]
         if self.add_indicator:
             indicator_names = self.indicator_.get_feature_names_out(input_features)
@@ -657,9 +692,11 @@ def _get_missing_features_info(self, X):
                 imputer_mask = sparse.csc_matrix(imputer_mask)
 
         if self.features == 'all':
-            features_indices = np.arange(X.shape[1])
+            xp = np.get_array_module(imputer_mask)
+            features_indices = xp.arange(X.shape[1])
         else:
-            features_indices = np.flatnonzero(n_missing)
+            xp = np.get_array_module(n_missing)
+            features_indices = xp.flatnonzero(n_missing)
 
         return imputer_mask, features_indices
 
@@ -668,11 +705,13 @@ def _validate_input(self, X, in_fit):
             ensure_all_finite = True
         else:
             ensure_all_finite = "allow-nan"
+        mem_type = "host" if _is_object_dtype(X) else "device"
         X = check_inputs(
             self,
             X,
             accept_sparse=('csc', 'csr'),
             ensure_all_finite=ensure_all_finite,
+            mem_type=mem_type,
             reset=in_fit,
         )
         _check_inputs_dtype(X, self.missing_values)
@@ -802,10 +841,13 @@ def get_feature_names_out(self, input_features=None):
         check_is_fitted(self)
         input_features = check_input_features(self, input_features)
         prefix = self.__class__.__name__.lower()
+        features = type(self).features_.get_raw(self)
+        if isinstance(features, np.ndarray):
+            features = features.get()
         return cpu_np.asarray(
             [
                 f"{prefix}_{feature_name}"
-                for feature_name in input_features[self.features_.get()]
+                for feature_name in input_features[features]
             ],
             dtype=object,
         )

python/cuml/cuml/internals/outputs.py

@@ -284,6 +284,31 @@ def infer_output_type(array, array_like="numpy"):
     return None
 
 
+def _is_object_dtype(res):
+    """Check for object or string dtype on an array/dataframe-like."""
+
+    def is_object_or_string(dtype):
+        try:
+            if np.dtype(dtype).kind == "O":
+                return True
+        except TypeError:
+            pass
+        return pd.api.types.is_string_dtype(dtype)
+
+    dtype = getattr(res, "dtype", None)
+    if dtype is not None:
+        return is_object_or_string(dtype)
+
+    dtypes = getattr(res, "dtypes", None)
+    if dtypes is not None:
+        try:
+            return any(is_object_or_string(dt) for dt in dtypes)
+        except TypeError:
+            return False
+
+    return False
+
+
 class ArrayIndexPair:
     """An array paired with an aligned index.
 
@@ -521,6 +546,40 @@ 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 internal outputs, or wrap them in a
+            # dataframe-like type.
+            if output_type == "cuml":
+                return obj
+            elif output_type == "cupy":
+                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 obj.ndim == 2:
+                if one_col_2d_as_series and obj.shape[1] == 1:
+                    host_df = pd.Series(
+                        obj.flatten(), index=index
+                    ).infer_objects()
+                else:
+                    host_df = pd.DataFrame(obj, index=index).infer_objects()
+            else:
+                host_df = pd.Series(obj, index=index).infer_objects()
+
+            try:
+                return cudf.from_pandas(host_df)
+            except (TypeError, ValueError, NotImplementedError) as exc:
+                raise TypeError(
+                    "Cannot convert an object-dtype output with shape "
+                    f"{obj.shape} to output_type='cudf'. Use "
+                    "output_type='pandas' or output_type='numpy' for this "
+                    "object layout."
+                ) from exc
         else:
             # Other output types use device memory, coerce to cupy and take
             # cupy code path.
@@ -639,6 +698,16 @@ def __reduce__(self):
     def __set_name__(self, owner, name):
         self.name = name
 
+    def get_raw(self, instance):
+        """Return the value so internal work can dispatch on its stored type."""
+        cache = instance.__dict__.get(self.name)
+        if cache is None:
+            raise AttributeError(
+                f"{type(instance).__name__!r} object has no attribute "
+                f"{self.name!r}"
+            )
+        return cache.value
+
     def __get__(self, instance, owner):
         if instance is None:
             return self