diff --git a/sklearn/preprocessing/_encoders.py b/sklearn/preprocessing/_encoders.py index a2669bf4665b4..0504396c40f9b 100644 --- a/sklearn/preprocessing/_encoders.py +++ b/sklearn/preprocessing/_encoders.py @@ -16,7 +16,9 @@ _fit_context, ) from sklearn.utils import _align_api_if_sparse, check_array -from sklearn.utils._dataframe import is_df_or_series +from sklearn.utils._dataframe import ( + is_df_or_series, +) from sklearn.utils._encode import _encode, _get_counts, _unique from sklearn.utils._indexing import _safe_indexing from sklearn.utils._mask import _get_mask @@ -40,7 +42,7 @@ class _BaseEncoder(TransformerMixin, BaseEstimator): """ - def _check_X(self, X): + def _check_X(self, X, transform=False): """ Perform custom check_array: - convert list of strings to object dtype @@ -49,10 +51,16 @@ def _check_X(self, X): constructed feature by feature to preserve the data types of pandas DataFrame columns, as otherwise information is lost and cannot be used, e.g. for the `categories_` attribute. + + For small pandas/polars inputs, columns are extracted as plain numpy + arrays instead of going through `narwhals`: at this size, there's no + bulk work to vectorize, so `narwhals`'s own per-column overhead (e.g. + dtype introspection, building a `pandas.Index` for lookups) costs + more than it saves, and dominates `transform` latency (see gh-32368). """ X_columns = [] - if is_df_or_series(X): + if is_df_or_series(X) and (not transform or len(X) > 128): try: X = nw.from_native(X) except TypeError: @@ -176,6 +184,16 @@ def _fit( self.categories_.append(cats) + # Lazily populated by `_transform` (not here) with, per feature, the + # lookup table used to look up values in `categories_[i]`. Building + # it is O(len(categories_[i])), so we avoid paying that cost in + # `fit` for models that are never used for `transform`, and only pay + # it once, on the first `transform` call, rather than on every call. + # This is what makes small-batch `transform` calls fast: `_check_X` + # routes small pandas/polars batches through this same cached path + # (see its docstring), rather than through `narwhals`. + self._transform_cache = [{} for _ in range(n_features)] + output = {"n_samples": n_samples} if return_counts: output["category_counts"] = category_counts @@ -203,7 +221,7 @@ def _transform( warn_on_unknown=False, ignore_category_indices=None, ): - X_list, n_samples, n_features = self._check_X(X) + X_list, n_samples, n_features = self._check_X(X, transform=True) validate_data(self, X=X, reset=False, skip_check_array=True) X_int = np.zeros((n_samples, n_features), dtype=int, order="F") @@ -212,7 +230,9 @@ def _transform( columns_with_unknown = [] for i in range(n_features): Xi = X_list[i] - X_int[:, i] = _encode(Xi, uniques=self.categories_[i]) + X_int[:, i] = _encode( + Xi, uniques=self.categories_[i], cache=self._transform_cache[i] + ) X_mask[:, i] = X_int[:, i] != -1 if not np.all(X_mask[:, i]): diff --git a/sklearn/utils/_encode.py b/sklearn/utils/_encode.py index b247e7c2c5821..ac5bf4d748962 100644 --- a/sklearn/utils/_encode.py +++ b/sklearn/utils/_encode.py @@ -179,13 +179,24 @@ def __missing__(self, key): return -1 -def _map_to_integer(values, uniques): +def _map_to_integer(values, uniques, cache=None): """Map values based on their position in uniques. Values not present in `uniques` are encoded as -1. + + `cache`, if given, is a plain dict that this function may use to avoid + rebuilding the `uniques` lookup table on every call: building it is + O(n_uniques), which otherwise dominates the cost of encoding small + batches against a large set of categories (e.g. repeated single-row + calls to `transform`). """ xp, _ = get_namespace(values, uniques) - table = _nandict({val: i for i, val in enumerate(uniques)}) + if cache is None: + table = _nandict({val: i for i, val in enumerate(uniques)}) + elif "cache" in table: + table = cache["table"] + else: + table = cache["table"] = _nandict({val: i for i, val in enumerate(uniques)}) return xp.asarray([table[v] for v in values], device=array_device(values)) @@ -268,7 +279,7 @@ def _encode_labels(values, *, uniques): return encoded -def _encode(values, *, uniques, return_diff=False): +def _encode(values, *, uniques, return_diff=False, cache=None): """Encode values into [0, n_uniques - 1]. Uses pure python method for object dtype, and numpy method for @@ -290,6 +301,14 @@ def _encode(values, *, uniques, return_diff=False): return_diff : bool, default=False If True, also return the unique values in `values` that are not present in `uniques`. + cache : dict, default=None + Mutable dict used to cache, across repeated calls with the same + `uniques`, the lookup table built from `uniques`. Only used for + object dtype (non-`narwhals.Series`) inputs, where building that + table is O(len(uniques)) and would otherwise dominate the cost of + encoding small batches (e.g. repeated single-row calls to + `transform`). Callers are responsible for invalidating (e.g. + replacing with a fresh dict) the cache whenever `uniques` changes. Returns ------- @@ -303,7 +322,7 @@ def _encode(values, *, uniques, return_diff=False): if isinstance(values, nw.Series): encoded = _encode_series(values, uniques) elif not xp.isdtype(values.dtype, "numeric"): - encoded = _map_to_integer(values, uniques) + encoded = _map_to_integer(values, uniques, cache=cache) else: encoded = xp.searchsorted(uniques, values) if size(uniques):