Skip to content

Migrate StringSimilarityEncoder to narwhals, add polars support #1031

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
solegalli wants to merge 2 commits into
base: narwhals-migration
Choose a base branch
from
Open

Migrate StringSimilarityEncoder to narwhals, add polars support #1031

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8bc929d
Migrate StringSimilarityEncoder to narwhals, add polars support
solegalli Aug 26, 2026
f263a18
Adapt StringSimilarityEncoder to narwhals-returning check_X
solegalli Aug 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions docs/user_guide/encoding/StringSimilarityEncoder.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -299,6 +299,37 @@ Below, we see the resulting dataframe:
393 0.0 0.437500 0.666667 0.666667


With polars
-----------

:class:`StringSimilarityEncoder()` works the same way with polars dataframes:

.. code:: python

import polars as pl
from feature_engine.encoding import StringSimilarityEncoder

df = pl.DataFrame({"words": ["dog", "dig", "cat"]})

encoder = StringSimilarityEncoder()
dft = encoder.fit_transform(df)
dft

We see the same similarity values as with the pandas dataframe:

.. code:: text

shape: (3, 3)
┌───────────┬───────────┬───────────┐
│ words_dog ┆ words_dig ┆ words_cat │
│ --- ┆ --- ┆ --- │
│ f64 ┆ f64 ┆ f64 │
╞═══════════╪═══════════╪═══════════╡
│ 1.0 ┆ 0.666667 ┆ 0.0 │
│ 0.666667 ┆ 1.0 ┆ 0.0 │
│ 0.0 ┆ 0.0 ┆ 1.0 │
└───────────┴───────────┴───────────┘

Additional resources
--------------------

Expand Down
185 changes: 107 additions & 78 deletions feature_engine/encoding/similarity_encoder.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,9 @@
from difflib import SequenceMatcher
from typing import List, Optional, Union

import narwhals as nw
import numpy as np
import pandas as pd
from narwhals.typing import IntoDataFrame, IntoSeries
from sklearn.utils.validation import check_is_fitted

from feature_engine._docstrings.fit_attributes import (
Expand Down Expand Up @@ -183,6 +184,26 @@ class StringSimilarityEncoder(CategoricalMethodsMixin, CategoricalInitMixin):
1 2 0.666667 1.000000 0.444444 0.4
2 3 0.444444 0.444444 1.000000 0.0
3 4 0.000000 0.400000 0.000000 1.0

With polars

>>> import polars as pl
>>> from feature_engine.encoding import StringSimilarityEncoder
>>> X = pl.DataFrame(dict(x1 = [1,2,3,4], x2 = ["dog", "dig", "dagger", "hi"]))
>>> sse = StringSimilarityEncoder()
>>> sse.fit(X)
>>> sse.transform(X)
shape: (4, 5)
┌─────┬──────────┬──────────┬───────────┬───────┐
│ x1 ┆ x2_dog ┆ x2_dig ┆ x2_dagger ┆ x2_hi │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞═════╪══════════╪══════════╪═══════════╪═══════╡
│ 1 ┆ 1.0 ┆ 0.666667 ┆ 0.444444 ┆ 0.0 │
│ 2 ┆ 0.666667 ┆ 1.0 ┆ 0.444444 ┆ 0.4 │
│ 3 ┆ 0.444444 ┆ 0.444444 ┆ 1.0 ┆ 0.0 │
│ 4 ┆ 0.0 ┆ 0.4 ┆ 0.0 ┆ 1.0 │
└─────┴──────────┴──────────┴───────────┴───────┘
"""

def __init__(
Expand Down Expand Up @@ -217,7 +238,7 @@ def __init__(
self.missing_values = missing_values
self.keywords = keywords

def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None):
"""
Learns the unique categories per variable. If top_categories is indicated,
it will learn the most popular categories. Alternatively, it learns all
Expand All @@ -226,15 +247,15 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
Parameters
----------

X: pandas dataframe of shape = [n_samples, n_features]
X: dataframe of shape = [n_samples, n_features]
The training input samples.
Can be the entire dataframe, not just the variables to encode.

y: pandas series, default=None
y: Series, default=None
Target. It is not needed in this encoder. You can pass y or None.
"""

X = check_X(X)
nw_X = check_X(X)
variables_ = self._check_or_select_variables(X)

if self.keywords and not all(
Expand All @@ -257,107 +278,115 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
else:
cols_to_iterate = variables_

if self.missing_values == "raise":
for var in cols_to_iterate:
self.encoder_dict_[var] = (
X[var]
.astype(str)
.value_counts()
.head(self.top_categories)
.index.tolist()
)
elif self.missing_values == "impute":
for var in cols_to_iterate:
series = X[var]
self.encoder_dict_[var] = (
series.astype(str)
.mask(series.isna(), "")
.value_counts()
.head(self.top_categories)
.index.tolist()
)
elif self.missing_values == "ignore":
for var in cols_to_iterate:
self.encoder_dict_[var] = (
X[var]
.dropna()
.astype(str)
.value_counts(dropna=True)
.head(self.top_categories)
.index.tolist()
# cast(nw.String) preserves nulls as null on both backends (unlike
# pandas' own astype(str), which stringifies NaN to "nan"), so
# "impute" can fill_null("") directly and "ignore" can drop_nulls()
# before casting, with no leftover "nan"/"<NA>" text sentinels to
# special-case downstream.
for var in cols_to_iterate:
col = nw_X.get_column(var)
if self.missing_values == "impute":
col = col.cast(nw.String).fill_null("")
elif self.missing_values == "ignore":
col = col.drop_nulls().cast(nw.String)
elif self.missing_values == "raise":
col = col.cast(nw.String)
else:
# missing_values can be set directly (e.g. via set_params or
# attribute assignment) bypassing the __init__ check above.
raise ValueError(
"Unrecognized value for missing_values. It should be 'raise', "
f"'ignore' or 'impute'. Got {self.missing_values} instead."
)
else:
raise ValueError(
"Unrecognized value for missing_values. It should be 'raise', 'ignore' "
f"or 'impute'. Got {self.missing_values} instead."
)

# sort=True mirrors pandas' own value_counts() default order
# (descending by count, ties broken by first appearance), so
# encoder_dict_ keeps the same category order as before.
counts = col.value_counts(sort=True)
categories = counts.get_column(counts.columns[0]).to_list()
self.encoder_dict_[var] = categories[: self.top_categories]

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

def transform(self, X: pd.DataFrame) -> pd.DataFrame:
def transform(self, X: IntoDataFrame) -> IntoDataFrame:
"""
Replaces the categorical variables with the similarity variables.

Parameters
----------
X: pandas dataframe of shape = [n_samples, n_features]
X: dataframe of shape = [n_samples, n_features]
The data to transform.

Returns
-------
X_new: pandas dataframe.
X_new: dataframe.
The transformed dataframe. The shape of the dataframe will be different from
the original as it includes the similarity variables in place of the
original categorical ones.
"""

check_is_fitted(self)
X = self._check_transform_input_and_state(X)
nw_X = self._check_transform_input_and_state(X)
if self.missing_values == "raise":
_check_contains_na(X, self.variables_, error_msg="optional")

if len(self.variables_) == 0:
return X

new_values = []
return nw_X.to_native()

# String similarity (difflib.SequenceMatcher) has no vectorised
# narwhals/backend equivalent, so it is computed in numpy: the
# similarity matrix is built once per unique value (not per row)
# via broadcasting, then broadcast back to all rows with
# np.unique's inverse index - cheaper than a per-row Python-level
# dict lookup and it is backend agnostic, so a single code path
# covers both pandas and polars. Benchmarked at 10k-100k rows x
# 1-10 columns x 5-50 categories: the difflib computation itself
# dominates wall time by 1-3 orders of magnitude, so a
# pandas-specific fast path for reassembling the output columns
# (as used in DecisionTreeFeatures) would save at most ~10ms out of
# a transform that is already tens of ms to seconds - not worth the
# code duplication here.
new_series = []
for var in self.variables_:
col = nw_X.get_column(var)
categories = self.encoder_dict_[var]

null_mask = None
if self.missing_values == "impute":
series = X[var]
series = series.astype(str).mask(series.isna(), "")
str_col = col.cast(nw.String).fill_null("")
else:
series = X[var].astype(str)

categories = series.unique()
column_encoder_dict = {
x: _gpm_fast_vec(x, self.encoder_dict_[var]) for x in categories
}
# Ensure map result is always an array of the correct size.
# Missing values in categories or unknown categories will map to NaN.
default_nan = np.full(len(self.encoder_dict_[var]), np.nan)
if "nan" not in column_encoder_dict:
column_encoder_dict["nan"] = default_nan
if "<NA>" not in column_encoder_dict:
column_encoder_dict["<NA>"] = default_nan

encoded_series = series.map(column_encoder_dict)

# Robust stacking: replace any float NaNs (from unknown values) with arrays
encoded_list = [
v if isinstance(v, (list, np.ndarray)) else default_nan
for v in encoded_series
]
encoded = np.vstack(encoded_list)
if self.missing_values == "ignore":
encoded[X[var].isna(), :] = np.nan
new_values.append(encoded)

new_features = self._get_new_features_name()
X.loc[:, new_features] = np.hstack(new_values)

return X.drop(self.variables_, axis=1)
str_col = col.cast(nw.String)
if self.missing_values == "ignore":
null_mask = np.array(col.is_null().to_list())

values = np.asarray(str_col.to_list(), dtype=object)
if null_mask is not None:
# placeholder value for null rows: overwritten with NaN
# below, the string itself is never used.
values = np.where(null_mask, "", values)

unique_vals, inverse = np.unique(values, return_inverse=True)
cats_arr = np.asarray(categories, dtype=object)
sim_matrix = _gpm_fast_vec(
unique_vals.reshape(-1, 1), cats_arr.reshape(1, -1)
)
encoded = sim_matrix[inverse]

if null_mask is not None:
encoded[null_mask, :] = np.nan

for j, category in enumerate(categories):
name = f"{var}_nan" if category == "" else f"{var}_{category}"
new_series.append(
nw.new_series(name, encoded[:, j], backend=nw_X.implementation)
)

nw_X = nw_X.with_columns(*new_series).drop(self.variables_)

return nw_X.to_native()

def _get_new_features_name(self) -> List[str]:
"""Return names of the created features."""
Expand All @@ -378,7 +407,7 @@ def _add_new_feature_names(self, feature_names: List[str]) -> List[str]:

return feature_names

def inverse_transform(self, X: pd.DataFrame):
def inverse_transform(self, X: IntoDataFrame):
"""inverse_transform is not implemented for this transformer."""
raise NotImplementedError(
"inverse_transform is not implemented for this transformer."
Expand Down
Loading