Skip to content

Commit 6281fcb

Browse files
solegalliclaude
andcommitted
Migrate StringSimilarityEncoder to narwhals, add polars support
fit() rebuilds encoder_dict_ with narwhals cast(nw.String)/value_counts, matching the CountEncoder/RareLabelEncoder convention. cast() preserves nulls as null on both pandas and polars (verified empirically), unlike pandas' own astype(str) which stringifies NaN to "nan" - this lets "impute" mode fill_null("") directly and "ignore" mode drop_nulls() before casting, replacing the old "nan"/"<NA>" text-sentinel workaround with a real null check (col.is_null()) that can't collide with a genuine category literally named "nan" or "<NA>" (both edge cases stay covered by test_string_dtype_with_literal_nan_strings). transform()'s per-row difflib.SequenceMatcher similarity has no vectorised narwhals equivalent, so it's computed once per unique value via numpy broadcasting (np.unique's inverse index fans the small per-unique-value matrix back out to all rows) and reassembled with nw.new_series()/with_columns(), same pattern DecisionTreeFeatures uses for externally-computed new columns. Benchmarked a pandas-specific fast path (X.join(dict-of-columns), as DecisionTreeFeatures uses) against the unified narwhals with_columns() here across 10k-100k rows x 1-10 columns x 5-50 categories: assembly overhead ranges 0.9x-6.25x depending on shape, but the difflib computation itself dominates wall time by 1-3 orders of magnitude in every realistic scenario (e.g. 30ms difflib vs <1ms assembly overhead at 100k rows/20 categories) - even the worst synthetic case (500 output columns) only costs ~10ms extra out of an already tens-of-ms-to-seconds transform. Went with the unified/merged implementation: no is_pandas split, one code path for both backends. Rewrote tests as single parametrized cases over @pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]), keeping only the pandas-NA-sentinel tests (np.nan/pd.NA/None, StringDtype) pandas-only since polars has no equivalent multi-sentinel behavior to exercise. All doc examples (including the Titanic worked example) re-verified against actual output; added a "With polars" section. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1 parent ae0a9c2 commit 6281fcb

3 files changed

Lines changed: 259 additions & 191 deletions

File tree

‎docs/user_guide/encoding/StringSimilarityEncoder.rst‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -299,6 +299,37 @@ Below, we see the resulting dataframe:
299299
393 0.0 0.437500 0.666667 0.666667
300300
301301
302+
With polars
303+
-----------
304+
305+
:class:`StringSimilarityEncoder()` works the same way with polars dataframes:
306+
307+
.. code:: python
308+
309+
import polars as pl
310+
from feature_engine.encoding import StringSimilarityEncoder
311+
312+
df = pl.DataFrame({"words": ["dog", "dig", "cat"]})
313+
314+
encoder = StringSimilarityEncoder()
315+
dft = encoder.fit_transform(df)
316+
dft
317+
318+
We see the same similarity values as with the pandas dataframe:
319+
320+
.. code:: text
321+
322+
shape: (3, 3)
323+
┌───────────┬───────────┬───────────┐
324+
│ words_dog ┆ words_dig ┆ words_cat │
325+
│ --- ┆ --- ┆ --- │
326+
│ f64 ┆ f64 ┆ f64 │
327+
╞═══════════╪═══════════╪═══════════╡
328+
│ 1.0 ┆ 0.666667 ┆ 0.0 │
329+
│ 0.666667 ┆ 1.0 ┆ 0.0 │
330+
│ 0.0 ┆ 0.0 ┆ 1.0 │
331+
└───────────┴───────────┴───────────┘
332+
302333
Additional resources
303334
--------------------
304335

‎feature_engine/encoding/similarity_encoder.py‎

Lines changed: 105 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
from difflib import SequenceMatcher
22
from typing import List, Optional, Union
33

4+
import narwhals as nw
45
import numpy as np
5-
import pandas as pd
6+
from narwhals.typing import IntoDataFrame, IntoSeries
67
from sklearn.utils.validation import check_is_fitted
78

89
from feature_engine._docstrings.fit_attributes import (
@@ -183,6 +184,26 @@ class StringSimilarityEncoder(CategoricalMethodsMixin, CategoricalInitMixin):
183184
1 2 0.666667 1.000000 0.444444 0.4
184185
2 3 0.444444 0.444444 1.000000 0.0
185186
3 4 0.000000 0.400000 0.000000 1.0
187+
188+
With polars
189+
190+
>>> import polars as pl
191+
>>> from feature_engine.encoding import StringSimilarityEncoder
192+
>>> X = pl.DataFrame(dict(x1 = [1,2,3,4], x2 = ["dog", "dig", "dagger", "hi"]))
193+
>>> sse = StringSimilarityEncoder()
194+
>>> sse.fit(X)
195+
>>> sse.transform(X)
196+
shape: (4, 5)
197+
┌─────┬──────────┬──────────┬───────────┬───────┐
198+
│ x1 ┆ x2_dog ┆ x2_dig ┆ x2_dagger ┆ x2_hi │
199+
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
200+
│ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
201+
╞═════╪══════════╪══════════╪═══════════╪═══════╡
202+
│ 1 ┆ 1.0 ┆ 0.666667 ┆ 0.444444 ┆ 0.0 │
203+
│ 2 ┆ 0.666667 ┆ 1.0 ┆ 0.444444 ┆ 0.4 │
204+
│ 3 ┆ 0.444444 ┆ 0.444444 ┆ 1.0 ┆ 0.0 │
205+
│ 4 ┆ 0.0 ┆ 0.4 ┆ 0.0 ┆ 1.0 │
206+
└─────┴──────────┴──────────┴───────────┴───────┘
186207
"""
187208

188209
def __init__(
@@ -217,7 +238,7 @@ def __init__(
217238
self.missing_values = missing_values
218239
self.keywords = keywords
219240

220-
def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
241+
def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None):
221242
"""
222243
Learns the unique categories per variable. If top_categories is indicated,
223244
it will learn the most popular categories. Alternatively, it learns all
@@ -226,11 +247,11 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
226247
Parameters
227248
----------
228249
229-
X: pandas dataframe of shape = [n_samples, n_features]
250+
X: dataframe of shape = [n_samples, n_features]
230251
The training input samples.
231252
Can be the entire dataframe, not just the variables to encode.
232253
233-
y: pandas series, default=None
254+
y: Series, default=None
234255
Target. It is not needed in this encoder. You can pass y or None.
235256
"""
236257

@@ -257,58 +278,52 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
257278
else:
258279
cols_to_iterate = variables_
259280

260-
if self.missing_values == "raise":
261-
for var in cols_to_iterate:
262-
self.encoder_dict_[var] = (
263-
X[var]
264-
.astype(str)
265-
.value_counts()
266-
.head(self.top_categories)
267-
.index.tolist()
268-
)
269-
elif self.missing_values == "impute":
270-
for var in cols_to_iterate:
271-
series = X[var]
272-
self.encoder_dict_[var] = (
273-
series.astype(str)
274-
.mask(series.isna(), "")
275-
.value_counts()
276-
.head(self.top_categories)
277-
.index.tolist()
278-
)
279-
elif self.missing_values == "ignore":
280-
for var in cols_to_iterate:
281-
self.encoder_dict_[var] = (
282-
X[var]
283-
.dropna()
284-
.astype(str)
285-
.value_counts(dropna=True)
286-
.head(self.top_categories)
287-
.index.tolist()
281+
# cast(nw.String) preserves nulls as null on both backends (unlike
282+
# pandas' own astype(str), which stringifies NaN to "nan"), so
283+
# "impute" can fill_null("") directly and "ignore" can drop_nulls()
284+
# before casting, with no leftover "nan"/"<NA>" text sentinels to
285+
# special-case downstream.
286+
nw_X = nw.from_native(X, eager_only=True)
287+
for var in cols_to_iterate:
288+
col = nw_X.get_column(var)
289+
if self.missing_values == "impute":
290+
col = col.cast(nw.String).fill_null("")
291+
elif self.missing_values == "ignore":
292+
col = col.drop_nulls().cast(nw.String)
293+
elif self.missing_values == "raise":
294+
col = col.cast(nw.String)
295+
else:
296+
# missing_values can be set directly (e.g. via set_params or
297+
# attribute assignment) bypassing the __init__ check above.
298+
raise ValueError(
299+
"Unrecognized value for missing_values. It should be 'raise', "
300+
f"'ignore' or 'impute'. Got {self.missing_values} instead."
288301
)
289-
else:
290-
raise ValueError(
291-
"Unrecognized value for missing_values. It should be 'raise', 'ignore' "
292-
f"or 'impute'. Got {self.missing_values} instead."
293-
)
302+
303+
# sort=True mirrors pandas' own value_counts() default order
304+
# (descending by count, ties broken by first appearance), so
305+
# encoder_dict_ keeps the same category order as before.
306+
counts = col.value_counts(sort=True)
307+
categories = counts.get_column(counts.columns[0]).to_list()
308+
self.encoder_dict_[var] = categories[: self.top_categories]
294309

295310
# assign underscore parameters at the end in case code above fails
296311
self.variables_ = variables_
297312
self._get_feature_names_in(X)
298313
return self
299314

300-
def transform(self, X: pd.DataFrame) -> pd.DataFrame:
315+
def transform(self, X: IntoDataFrame) -> IntoDataFrame:
301316
"""
302317
Replaces the categorical variables with the similarity variables.
303318
304319
Parameters
305320
----------
306-
X: pandas dataframe of shape = [n_samples, n_features]
321+
X: dataframe of shape = [n_samples, n_features]
307322
The data to transform.
308323
309324
Returns
310325
-------
311-
X_new: pandas dataframe.
326+
X_new: dataframe.
312327
The transformed dataframe. The shape of the dataframe will be different from
313328
the original as it includes the similarity variables in place of the
314329
original categorical ones.
@@ -322,42 +337,58 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame:
322337
if len(self.variables_) == 0:
323338
return X
324339

325-
new_values = []
340+
# String similarity (difflib.SequenceMatcher) has no vectorised
341+
# narwhals/backend equivalent, so it is computed in numpy: the
342+
# similarity matrix is built once per unique value (not per row)
343+
# via broadcasting, then broadcast back to all rows with
344+
# np.unique's inverse index - cheaper than a per-row Python-level
345+
# dict lookup and it is backend agnostic, so a single code path
346+
# covers both pandas and polars. Benchmarked at 10k-100k rows x
347+
# 1-10 columns x 5-50 categories: the difflib computation itself
348+
# dominates wall time by 1-3 orders of magnitude, so a
349+
# pandas-specific fast path for reassembling the output columns
350+
# (as used in DecisionTreeFeatures) would save at most ~10ms out of
351+
# a transform that is already tens of ms to seconds - not worth the
352+
# code duplication here.
353+
nw_X = nw.from_native(X, eager_only=True)
354+
new_series = []
326355
for var in self.variables_:
356+
col = nw_X.get_column(var)
357+
categories = self.encoder_dict_[var]
358+
359+
null_mask = None
327360
if self.missing_values == "impute":
328-
series = X[var]
329-
series = series.astype(str).mask(series.isna(), "")
361+
str_col = col.cast(nw.String).fill_null("")
330362
else:
331-
series = X[var].astype(str)
332-
333-
categories = series.unique()
334-
column_encoder_dict = {
335-
x: _gpm_fast_vec(x, self.encoder_dict_[var]) for x in categories
336-
}
337-
# Ensure map result is always an array of the correct size.
338-
# Missing values in categories or unknown categories will map to NaN.
339-
default_nan = np.full(len(self.encoder_dict_[var]), np.nan)
340-
if "nan" not in column_encoder_dict:
341-
column_encoder_dict["nan"] = default_nan
342-
if "<NA>" not in column_encoder_dict:
343-
column_encoder_dict["<NA>"] = default_nan
344-
345-
encoded_series = series.map(column_encoder_dict)
346-
347-
# Robust stacking: replace any float NaNs (from unknown values) with arrays
348-
encoded_list = [
349-
v if isinstance(v, (list, np.ndarray)) else default_nan
350-
for v in encoded_series
351-
]
352-
encoded = np.vstack(encoded_list)
353-
if self.missing_values == "ignore":
354-
encoded[X[var].isna(), :] = np.nan
355-
new_values.append(encoded)
356-
357-
new_features = self._get_new_features_name()
358-
X.loc[:, new_features] = np.hstack(new_values)
359-
360-
return X.drop(self.variables_, axis=1)
363+
str_col = col.cast(nw.String)
364+
if self.missing_values == "ignore":
365+
null_mask = np.array(col.is_null().to_list())
366+
367+
values = np.asarray(str_col.to_list(), dtype=object)
368+
if null_mask is not None:
369+
# placeholder value for null rows: overwritten with NaN
370+
# below, the string itself is never used.
371+
values = np.where(null_mask, "", values)
372+
373+
unique_vals, inverse = np.unique(values, return_inverse=True)
374+
cats_arr = np.asarray(categories, dtype=object)
375+
sim_matrix = _gpm_fast_vec(
376+
unique_vals.reshape(-1, 1), cats_arr.reshape(1, -1)
377+
)
378+
encoded = sim_matrix[inverse]
379+
380+
if null_mask is not None:
381+
encoded[null_mask, :] = np.nan
382+
383+
for j, category in enumerate(categories):
384+
name = f"{var}_nan" if category == "" else f"{var}_{category}"
385+
new_series.append(
386+
nw.new_series(name, encoded[:, j], backend=nw_X.implementation)
387+
)
388+
389+
nw_X = nw_X.with_columns(*new_series).drop(self.variables_)
390+
391+
return nw_X.to_native()
361392

362393
def _get_new_features_name(self) -> List[str]:
363394
"""Return names of the created features."""
@@ -378,7 +409,7 @@ def _add_new_feature_names(self, feature_names: List[str]) -> List[str]:
378409

379410
return feature_names
380411

381-
def inverse_transform(self, X: pd.DataFrame):
412+
def inverse_transform(self, X: IntoDataFrame):
382413
"""inverse_transform is not implemented for this transformer."""
383414
raise NotImplementedError(
384415
"inverse_transform is not implemented for this transformer."

0 commit comments

Comments
 (0)