From 56d4b94d47a624ba964e1518edd0ae514f1b8779 Mon Sep 17 00:00:00 2001 From: Arthur Lacote Date: Tue, 21 Jul 2026 16:00:04 +0200 Subject: [PATCH 1/2] fast cached path for small batches --- sklearn/preprocessing/_encoders.py | 30 +++++++++++++++++++++++++----- sklearn/utils/_encode.py | 27 +++++++++++++++++++++++---- 2 files changed, 48 insertions(+), 9 deletions(-) diff --git a/sklearn/preprocessing/_encoders.py b/sklearn/preprocessing/_encoders.py index 2c57691a3fd58..b3493d6b98236 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._mask import _get_mask from sklearn.utils._missing import is_scalar_nan @@ -39,7 +41,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 @@ -48,10 +50,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: @@ -169,6 +177,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 @@ -196,7 +214,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") @@ -205,7 +223,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 2b86bc7c9a20c..d43a705fb9440 100644 --- a/sklearn/utils/_encode.py +++ b/sklearn/utils/_encode.py @@ -175,13 +175,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)}) + else: + table = cache.get("table") + if table is None: + table = cache["table"] = _nandict({val: i for i, val in enumerate(uniques)}) return xp.asarray([table[v] for v in values], device=device(values)) @@ -264,7 +275,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 @@ -286,6 +297,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 ------- @@ -299,7 +318,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): From 9612086e3938b09abda916b9670c9ed2f288ce07 Mon Sep 17 00:00:00 2001 From: Arthur Lacote Date: Thu, 6 Aug 2026 08:56:19 +0200 Subject: [PATCH 2/2] minor improvement --- sklearn/utils/_encode.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sklearn/utils/_encode.py b/sklearn/utils/_encode.py index 67cb97da6d5f7..3e8697859631c 100644 --- a/sklearn/utils/_encode.py +++ b/sklearn/utils/_encode.py @@ -190,10 +190,10 @@ def _map_to_integer(values, uniques, cache=None): xp, _ = get_namespace(values, 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.get("table") - if table is None: - table = cache["table"] = _nandict({val: i for i, val in enumerate(uniques)}) + 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))