Skip to content

Migrate MathFeatures to narwhals, add polars support #994

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

Merged
solegalli merged 3 commits into from
Aug 24, 2026
Merged

Migrate MathFeatures to narwhals, add polars support #994

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
03dcba7
Migrate MathFeatures to narwhals, add polars support
solegalli Aug 24, 2026
88bbe10
Rewrite MathFeatures tests to run the same test against both backends
solegalli Aug 24, 2026
f6e1c1e
fix: correct _pandas_version() return type hint from bool to int
solegalli Aug 24, 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
63 changes: 56 additions & 7 deletions docs/user_guide/creation/MathFeatures.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,11 +143,11 @@ We obtain the following dataframe:
2 krish Liverpool 19 0.7 2020-02-24 00:02:00 19.7
3 jack Bristol 18 0.6 2020-02-24 00:03:00 18.6

prod_Age_Marks amin_Age_Marks amax_Age_Marks std_Age_Marks
0 18.0 0.9 20.0 13.505740
1 16.8 0.8 21.0 14.283557
2 13.3 0.7 19.0 12.940054
3 10.8 0.6 18.0 12.303658
prod_Age_Marks min_Age_Marks max_Age_Marks std_Age_Marks
0 18.0 0.9 20.0 9.55
1 16.8 0.8 21.0 10.10
2 13.3 0.7 19.0 9.15
3 10.8 0.6 18.0 8.70

We have the option to set the parameter `drop_original` to True to drop the variables
after performing the calculations.
Expand All @@ -169,11 +169,60 @@ Which will return the names of all the variables in the transformed data:
'dob',
'sum_Age_Marks',
'prod_Age_Marks',
'amin_Age_Marks',
'amax_Age_Marks',
'min_Age_Marks',
'max_Age_Marks',
'std_Age_Marks']


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

:class:`MathFeatures()` works in the same way with a polars dataframe:

.. code:: python

import polars as pl
from feature_engine.creation import MathFeatures

df = pl.DataFrame({
"Age": [20, 21, 19, 18],
"Marks": [0.9, 0.8, 0.7, 0.6],
})

transformer = MathFeatures(
variables=["Age", "Marks"],
func=["sum", "prod", "min", "max", "std"],
)

print(transformer.fit_transform(df))

The resulting values match those found with pandas:

.. code:: text

shape: (4, 7)
┌─────┬───────┬───────────────┬────────────────┬───────────────┬───────────────┬───────────────┐
│ Age ┆ Marks ┆ sum_Age_Marks ┆ prod_Age_Marks ┆ min_Age_Marks ┆ max_Age_Marks ┆ std_Age_Marks │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 ┆ f64 │
╞═════╪═══════╪═══════════════╪════════════════╪═══════════════╪═══════════════╪═══════════════╡
│ 20 ┆ 0.9 ┆ 20.9 ┆ 18.0 ┆ 0.9 ┆ 20.0 ┆ 13.50574 │
│ 21 ┆ 0.8 ┆ 21.8 ┆ 16.8 ┆ 0.8 ┆ 21.0 ┆ 14.283557 │
│ 19 ┆ 0.7 ┆ 19.7 ┆ 13.3 ┆ 0.7 ┆ 19.0 ┆ 12.940054 │
│ 18 ┆ 0.6 ┆ 18.6 ┆ 10.8 ┆ 0.6 ┆ 18.0 ┆ 12.303658 │
└─────┴───────┴───────────────┴────────────────┴───────────────┴───────────────┴───────────────┘

`new_variables_names`, `drop_original`, and `get_feature_names_out()` work
identically to the pandas examples above.

If you pass a custom Python callable as `func` (instead of a string or one
of the common aggregations above, which are always NumPy-vectorized), note
that the callable receives a **plain tuple** of values for polars input,
not a pandas `Series` — so `lambda row: max(row) - min(row)` works on both
backends, but `lambda row: row.max() - row.min()` (which relies on `Series`
methods) only works with pandas.


New variables names
^^^^^^^^^^^^^^^^^^^

Expand Down
109 changes: 84 additions & 25 deletions feature_engine/creation/math_features.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,8 +1,10 @@
import warnings
from typing import Any, List, Optional, Union

import narwhals as nw
import narwhals.dependencies as nwd
import numpy as np
import pandas as pd
from narwhals.typing import IntoDataFrame

from feature_engine._docstrings.fit_attributes import (
_feature_names_in_docstring,
Expand All @@ -21,7 +23,10 @@
from feature_engine._docstrings.substitute import Substitution
from feature_engine.creation.base_creation import BaseCreation

_PANDAS_LT_3 = int(pd.__version__.split(".")[0]) < 3

def _pandas_version() -> int:
return int(nwd.get_pandas().__version__.split(".")[0])


# In pandas < 3, agg() maps these callables to the pandas methods and warns that
# this will change; the string alias keeps that behaviour (e.g., np.std ->
Expand Down Expand Up @@ -83,7 +88,11 @@ class MathFeatures(BaseCreation):
"""
MathFeatures() applies functions across multiple features returning one or more
additional features as a result. Common reductions use vectorized NumPy
operations. Other functions fall back to `pandas.agg()` with `axis=1`.
operations. Other functions fall back to `pandas.agg()` with `axis=1` for
pandas input, or to polars' native `map_rows()` for polars input — in that
case, the callable receives each row as a plain tuple, not a `Series`, so
it must not rely on `Series` methods (e.g. use `max(row)` instead of
`row.max()`) to work on both backends.

For supported aggregation functions, see `pandas documentation
<https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.agg.html>`_.
Expand Down Expand Up @@ -174,11 +183,30 @@ class MathFeatures(BaseCreation):

>>> mf = MathFeatures(variables = ["x1","x2"], func = "mean")
>>> mf.fit(X)
>>> mf.transform(X))
>>> mf.transform(X)
x1 x2 mean_x1_x2
0 1 4 2.5
1 2 5 3.5
2 3 6 4.5

With polars:

>>> import polars as pl
>>> from feature_engine.creation import MathFeatures
>>> X = pl.DataFrame({"x1": [1, 2, 3], "x2": [4, 5, 6]})
>>> mf = MathFeatures(variables=["x1", "x2"], func="sum")
>>> mf.fit(X)
>>> mf.transform(X)
shape: (3, 3)
┌─────┬─────┬───────────┐
│ x1 ┆ x2 ┆ sum_x1_x2 │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═══════════╡
│ 1 ┆ 4 ┆ 5 │
│ 2 ┆ 5 ┆ 7 │
│ 3 ┆ 6 ┆ 9 │
└─────┴─────┴───────────┘
"""

def __init__(
Expand Down Expand Up @@ -237,60 +265,91 @@ def __init__(
self.func = func
self.new_variables_names = new_variables_names

def transform(self, X: pd.DataFrame) -> pd.DataFrame:
def transform(self, X: IntoDataFrame) -> IntoDataFrame:
"""
Create and add new 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, shape = [n_samples, n_features + n_operations]
X_new: dataframe, shape = [n_samples, n_features + n_operations]
The input dataframe plus the new variables.
"""
X = self._check_transform_input_and_state(X)

new_variable_names = self._get_new_features_name()

func = self.func
if _PANDAS_LT_3:
is_pandas = nwd.is_pandas_dataframe(X)
if is_pandas is True and _pandas_version() < 3:
if isinstance(func, list):
func = [_FUNC_TO_STRING_ALIAS.get(fun, fun) for fun in func]
else:
func = _FUNC_TO_STRING_ALIAS.get(func, func)

variables = X[self.variables]
functions = func if isinstance(func, list) else [func]
reducers = [_get_numpy_reducer(fun) for fun in functions]
values = variables.to_numpy()

nw_X = nw.from_native(X, eager_only=True)
if is_pandas is True:
values = X[self.variables].to_numpy()
else:
values = nw_X.select(self.variables).to_numpy()

# Nullable extension dtypes produce object arrays. Keep those, custom
# callables, and less common pandas aggregations on the exact legacy path.
# callables, and less common aggregations on the fallback path below.
if reducers and values.dtype.kind in "biuf" and all(reducers):
results = []
for reducer, kwargs in reducers:
new_series = []
for (reducer, kwargs), name in zip(reducers, new_variable_names):
# pandas' named reductions do not warn for empty/all-missing rows.
# NumPy returns the same values but emits RuntimeWarning for some
# reducers, so silence only those warnings on this equivalent path.
with warnings.catch_warnings():
warnings.simplefilter("ignore", RuntimeWarning)
result = reducer(values, axis=1, **kwargs)
results.append(pd.Series(result, index=X.index))

result = results[0] if len(results) == 1 else pd.concat(results, axis=1)
else:
result = variables.agg(func, axis=1)

if len(new_variable_names) == 1:
X[new_variable_names[0]] = result
new_series.append(
nw.new_series(name, result, backend=nw_X.implementation)
)
nw_X = nw_X.with_columns(*new_series)
if self.drop_original is True:
nw_X = nw_X.drop(self.variables)
X = nw_X.to_native()
elif is_pandas is True:
result = X[self.variables].agg(func, axis=1)
if len(new_variable_names) == 1:
X[new_variable_names[0]] = result
else:
X[new_variable_names] = result
if self.drop_original is True:
X = X.drop(columns=self.variables)
else:
X[new_variable_names] = result

if self.drop_original:
X.drop(columns=self.variables, inplace=True)
# polars has no equivalent to pandas' agg(func, axis=1): apply each
# function natively via map_rows, one call per function. map_rows
# passes each row as a plain tuple, not a Series, so callables that
# rely on Series methods (e.g. `row.max()`) need `max(row)` instead.
sub_native = nw_X.select(self.variables).to_native()
new_series = []
for fun, name in zip(functions, new_variable_names):
if not callable(fun):
raise NotImplementedError(
f"'{fun}' has no NumPy-vectorized implementation, and "
"non-callable aggregation names are not supported for "
"polars input. Pass a Python callable instead."
)
result_df = sub_native.map_rows(fun)
new_series.append(
nw.new_series(
name, result_df.to_series(0), backend=nw_X.implementation
)
)
nw_X = nw_X.with_columns(*new_series)
if self.drop_original is True:
nw_X = nw_X.drop(self.variables)
X = nw_X.to_native()

return X

Expand Down
Loading