diff --git a/docs/user_guide/creation/MathFeatures.rst b/docs/user_guide/creation/MathFeatures.rst index 6b79af525..0c119d348 100644 --- a/docs/user_guide/creation/MathFeatures.rst +++ b/docs/user_guide/creation/MathFeatures.rst @@ -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. @@ -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 ^^^^^^^^^^^^^^^^^^^ diff --git a/feature_engine/creation/math_features.py b/feature_engine/creation/math_features.py index bea520bb1..edd36bca4 100644 --- a/feature_engine/creation/math_features.py +++ b/feature_engine/creation/math_features.py @@ -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, @@ -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 -> @@ -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 `_. @@ -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__( @@ -237,18 +265,18 @@ 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) @@ -256,41 +284,72 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: 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 diff --git a/tests/test_creation/test_math_features.py b/tests/test_creation/test_math_features.py index 9c4c6b10c..1c695ca88 100644 --- a/tests/test_creation/test_math_features.py +++ b/tests/test_creation/test_math_features.py @@ -1,13 +1,35 @@ import warnings +import narwhals as nw import numpy as np import pandas as pd +import polars as pl import pytest from sklearn.pipeline import Pipeline from feature_engine.creation import MathFeatures -dob_datrange = pd.date_range("2020-02-24", periods=4, freq="min") +DATA = { + "Name": ["tom", "nick", "krish", "jack"], + "City": ["London", "Manchester", "Liverpool", "Bristol"], + "Age": [20, 21, 19, 18], + "Marks": [0.9, 0.8, 0.7, 0.6], +} + + +def _none_to_nan(values): + # Missing values print as None for polars, NaN for pandas float columns + # - both mean "missing" here, so normalize both sides before comparing. + return [np.nan if v is None else v for v in values] + + +def assert_df_equal(X, expected: dict, abs_tol: float = 1e-5) -> None: + result = nw.from_native(X, eager_only=True).to_dict(as_series=False) + assert list(result.keys()) == list(expected.keys()) + for col, values in expected.items(): + assert _none_to_nan(result[col]) == pytest.approx( + _none_to_nan(values), abs=abs_tol, nan_ok=True + ) # test param variables_to_combine @@ -83,138 +105,95 @@ def test_error_new_variable_names_not_permitted(): ) -def test_aggregations_with_strings(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_aggregations_with_strings(make_df): + df = make_df(DATA) transformer = MathFeatures( variables=["Age", "Marks"], func=["sum", "prod", "mean", "std", "max", "min"] ) - X = transformer.fit_transform(df_vartypes) + Xt = transformer.fit_transform(df) - ref = pd.DataFrame.from_dict( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Age": [20, 21, 19, 18], - "Marks": [0.9, 0.8, 0.7, 0.6], - "dob": dob_datrange, - "sum_Age_Marks": [20.9, 21.8, 19.7, 18.6], - "prod_Age_Marks": [18.0, 16.8, 13.299999999999999, 10.799999999999999], - "mean_Age_Marks": [10.45, 10.9, 9.85, 9.3], - "std_Age_Marks": [ - 13.505739520663058, - 14.28355697996826, - 12.94005409571382, - 12.303657992645928, - ], - "max_Age_Marks": [20.0, 21.0, 19.0, 18.0], - "min_Age_Marks": [0.9, 0.8, 0.7, 0.6], - } - ) + expected = dict(DATA) + expected["sum_Age_Marks"] = [20.9, 21.8, 19.7, 18.6] + expected["prod_Age_Marks"] = [18.0, 16.8, 13.3, 10.8] + expected["mean_Age_Marks"] = [10.45, 10.9, 9.85, 9.3] + expected["std_Age_Marks"] = [13.505740, 14.283557, 12.940054, 12.303658] + expected["max_Age_Marks"] = [20.0, 21.0, 19.0, 18.0] + expected["min_Age_Marks"] = [0.9, 0.8, 0.7, 0.6] - # transform params - pd.testing.assert_frame_equal(X, ref) + assert_df_equal(Xt, expected) -def test_aggregations_with_functions(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_aggregations_with_functions(make_df): + df = make_df(DATA) transformer = MathFeatures( variables=["Age", "Marks"], func=[np.sum, np.mean, np.std] ) - X = transformer.fit_transform(df_vartypes) + Xt = transformer.fit_transform(df) - ref = pd.DataFrame.from_dict( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Age": [20, 21, 19, 18], - "Marks": [0.9, 0.8, 0.7, 0.6], - "dob": dob_datrange, - "sum_Age_Marks": [20.9, 21.8, 19.7, 18.6], - "mean_Age_Marks": [10.45, 10.9, 9.85, 9.3], - "std_Age_Marks": [ - 13.505739520663058, - 14.28355697996826, - 12.94005409571382, - 12.303657992645928, - ], - } - ) + expected = dict(DATA) + expected["sum_Age_Marks"] = [20.9, 21.8, 19.7, 18.6] + expected["mean_Age_Marks"] = [10.45, 10.9, 9.85, 9.3] - # TODO: Remove pandas < 3 support when dropping older pandas versions - # In pandas >=3, when the user passes np.std, agg will use numpy. - # In pandas <3, when the user passes np.std, agg will use pd.std. - # Hence the difference in results - if pd.__version__ >= "3": - ref["std_Age_Marks"] = np.std(df_vartypes[["Age", "Marks"]], axis=1) + # np.std uses ddof=0 (population std) everywhere now, except pandas < 3, + # where agg() still routes np.std through pandas' own ddof=1 Series.std(). + # TODO: remove the pandas < 3 branch when dropping older pandas support. + if make_df is pd.DataFrame and int(pd.__version__.split(".")[0]) < 3: + expected["std_Age_Marks"] = [13.505740, 14.283557, 12.940054, 12.303658] + else: + arr = np.array([DATA["Age"], DATA["Marks"]], dtype=float) + expected["std_Age_Marks"] = np.std(arr, axis=0).tolist() - # transform params - pd.testing.assert_frame_equal(X, ref) + assert_df_equal(Xt, expected) -def test_user_enters_two_operations(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_user_enters_two_operations(make_df): + df = make_df(DATA) transformer = MathFeatures(variables=["Age", "Marks"], func=["sum", np.mean]) + Xt = transformer.fit_transform(df) - X = transformer.fit_transform(df_vartypes) + expected = dict(DATA) + expected["sum_Age_Marks"] = [20.9, 21.8, 19.7, 18.6] + expected["mean_Age_Marks"] = [10.45, 10.9, 9.85, 9.3] - ref = pd.DataFrame.from_dict( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Age": [20, 21, 19, 18], - "Marks": [0.9, 0.8, 0.7, 0.6], - "dob": dob_datrange, - "sum_Age_Marks": [20.9, 21.8, 19.7, 18.6], - "mean_Age_Marks": [10.45, 10.9, 9.85, 9.3], - } - ) - - pd.testing.assert_frame_equal(X, ref) + assert_df_equal(Xt, expected) -def test_new_variable_names(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_new_variable_names(make_df): + df = make_df(DATA) transformer = MathFeatures( variables=["Age", "Marks"], func=["sum", "mean"], new_variables_names=["sum_of_two_vars", "mean_of_two_vars"], ) + Xt = transformer.fit_transform(df) - X = transformer.fit_transform(df_vartypes) + expected = dict(DATA) + expected["sum_of_two_vars"] = [20.9, 21.8, 19.7, 18.6] + expected["mean_of_two_vars"] = [10.45, 10.9, 9.85, 9.3] - ref = pd.DataFrame.from_dict( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Age": [20, 21, 19, 18], - "Marks": [0.9, 0.8, 0.7, 0.6], - "dob": dob_datrange, - "sum_of_two_vars": [20.9, 21.8, 19.7, 18.6], - "mean_of_two_vars": [10.45, 10.9, 9.85, 9.3], - } - ) + assert_df_equal(Xt, expected) - pd.testing.assert_frame_equal(X, ref) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_one_mathematical_operation(make_df): + df = make_df(DATA) + expected = dict(DATA) + expected["sum_Age_Marks"] = [20.9, 21.8, 19.7, 18.6] -def test_one_mathematical_operation(df_vartypes): transformer = MathFeatures(variables=["Age", "Marks"], func="sum") - X = transformer.fit_transform(df_vartypes) - - ref = pd.DataFrame.from_dict( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Age": [20, 21, 19, 18], - "Marks": [0.9, 0.8, 0.7, 0.6], - "dob": dob_datrange, - "sum_Age_Marks": [20.9, 21.8, 19.7, 18.6], - } - ) - pd.testing.assert_frame_equal(X, ref) + assert_df_equal(transformer.fit_transform(df), expected) transformer = MathFeatures(variables=["Age", "Marks"], func=["sum"]) - X = transformer.fit_transform(df_vartypes) - pd.testing.assert_frame_equal(X, ref) + assert_df_equal(transformer.fit_transform(df), expected) def test_variable_names_when_df_cols_are_integers(df_numeric_columns): + # polars requires string column names, so int-named columns are + # pandas-only - no polars equivalent to parametrize against here. transformer = MathFeatures( variables=[2, 3], func=["sum", "prod", "mean", "std", "max", "min"] ) @@ -227,7 +206,7 @@ def test_variable_names_when_df_cols_are_integers(df_numeric_columns): 1: ["London", "Manchester", "Liverpool", "Bristol"], 2: [20, 21, 19, 18], 3: [0.9, 0.8, 0.7, 0.6], - 4: dob_datrange, + 4: pd.date_range("2020-02-24", periods=4, freq="min"), "sum_2_3": [20.9, 21.8, 19.7, 18.6], "prod_2_3": [18.0, 16.8, 13.299999999999999, 10.799999999999999], "mean_2_3": [10.45, 10.9, 9.85, 9.3], @@ -245,9 +224,11 @@ def test_variable_names_when_df_cols_are_integers(df_numeric_columns): pd.testing.assert_frame_equal(X, ref) -def test_error_when_null_values_in_variable(df_vartypes): - df_na = df_vartypes.copy() - df_na.loc[1, "Age"] = np.nan +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_error_when_null_values_in_variable(make_df): + data_na = dict(DATA) + data_na["Age"] = [20, None, 19, 18] + df_na = make_df(data_na) math_combinator = MathFeatures( variables=["Age", "Marks"], @@ -258,65 +239,65 @@ def test_error_when_null_values_in_variable(df_vartypes): with pytest.raises(ValueError): math_combinator.fit(df_na) - math_combinator.fit(df_vartypes) + math_combinator.fit(make_df(DATA)) with pytest.raises(ValueError): math_combinator.transform(df_na) -def test_no_error_when_null_values_in_variable(df_vartypes): - df_na = df_vartypes.copy() - df_na.loc[1, "Age"] = np.nan +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_no_error_when_null_values_in_variable(make_df): + data_na = dict(DATA) + data_na["Age"] = [20, None, 19, 18] + df_na = make_df(data_na) transformer = MathFeatures( variables=["Age", "Marks"], func=["sum", "mean"], missing_values="ignore", ) + Xt = transformer.fit_transform(df_na) - X = transformer.fit_transform(df_na) + expected = dict(data_na) + expected["sum_Age_Marks"] = [20.9, 0.8, 19.7, 18.6] + expected["mean_Age_Marks"] = [10.45, 0.8, 9.85, 9.3] - ref = pd.DataFrame.from_dict( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "Age": [20, np.nan, 19, 18], - "Marks": [0.9, 0.8, 0.7, 0.6], - "dob": dob_datrange, - "sum_Age_Marks": [20.9, 0.8, 19.7, 18.6], - "mean_Age_Marks": [10.45, 0.8, 9.85, 9.3], - } - ) - # transform params - pd.testing.assert_frame_equal(X, ref) + assert_df_equal(Xt, expected) -def test_standard_aggregations_match_pandas_with_missing_values(): - X = pd.DataFrame( - { - "a": [1.0, np.nan, np.nan, 4.0], - "b": [3.0, 4.0, np.nan, 6.0], - "c": [5.0, 8.0, np.nan, np.nan], - } - ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_standard_aggregations_match_pandas_with_missing_values(make_df): + data = { + "a": [1.0, np.nan, np.nan, 4.0], + "b": [3.0, 4.0, np.nan, 6.0], + "c": [5.0, 8.0, np.nan, np.nan], + } functions = ["sum", "mean", "std", "var", "min", "max", "prod", "median"] names = [f"result_{function}" for function in functions] + + # pandas' own agg() is the ground truth both backends are checked against. + X_pd = pd.DataFrame(data) with warnings.catch_warnings(): warnings.simplefilter("ignore", RuntimeWarning) - expected = X.agg(functions, axis=1) - expected.columns = names + expected_df = X_pd.agg(functions, axis=1) + expected = {name: expected_df[fn].tolist() for name, fn in zip(names, functions)} + df = make_df(data) transformer = MathFeatures( - variables=list(X.columns), + variables=list(data.keys()), func=functions, new_variables_names=names, missing_values="ignore", ) - result = transformer.fit_transform(X) + result = transformer.fit_transform(df) - pd.testing.assert_frame_equal(result[names], expected) + result_dict = nw.from_native(result, eager_only=True).to_dict(as_series=False) + for name in names: + assert result_dict[name] == pytest.approx(expected[name], nan_ok=True) def test_nullable_dtypes_use_backwards_compatible_aggregation(): + # pandas' nullable "Int64" dtype is pandas-specific - no polars + # equivalent to parametrize against here. X = pd.DataFrame( { "a": pd.Series([1, pd.NA, 3], dtype="Int64"), @@ -339,65 +320,111 @@ def test_nullable_dtypes_use_backwards_compatible_aggregation(): pd.testing.assert_frame_equal(result[names], expected) -def test_custom_function_uses_pandas_aggregation_fallback(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_custom_function_fallback(make_df): + # max()/min()/sum() are built-ins, so they work identically whether + # func receives a pandas Series (pandas' agg(axis=1) fallback) or a + # plain tuple (polars' map_rows fallback) - one callable, one test. def peak_to_peak(row): - return row.max() - row.min() + return max(row) - min(row) - expected = df_vartypes[["Age", "Marks"]].agg(peak_to_peak, axis=1) + df = make_df(DATA) transformer = MathFeatures( variables=["Age", "Marks"], func=peak_to_peak, new_variables_names=["age_marks_range"], ) + Xt = transformer.fit_transform(df) - result = transformer.fit_transform(df_vartypes) + expected = dict(DATA) + expected["age_marks_range"] = [ + a - m for a, m in zip(DATA["Age"], DATA["Marks"]) + ] + assert_df_equal(Xt, expected) - pd.testing.assert_series_equal( - result["age_marks_range"], expected, check_names=False - ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_multiple_custom_functions_fallback(make_df): + def total(row): + return sum(row) -def test_drop_original_variables(df_vartypes): + def spread(row): + return max(row) - min(row) + + df = make_df(DATA) + transformer = MathFeatures( + variables=["Age", "Marks"], + func=[total, spread], + new_variables_names=["total", "spread"], + ) + Xt = transformer.fit_transform(df) + + expected = dict(DATA) + expected["total"] = [a + m for a, m in zip(DATA["Age"], DATA["Marks"])] + expected["spread"] = [a - m for a, m in zip(DATA["Age"], DATA["Marks"])] + assert_df_equal(Xt, expected) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_uncommon_aggregation_string_only_supported_for_pandas(make_df): + # a genuine, documented backend asymmetry, not an oversight: pandas' + # agg() accepts any of its own aggregation strings (even ones outside + # our NumPy-vectorized table), but polars has no way to resolve an + # arbitrary pandas-specific string without pandas itself, so it raises + # instead of silently doing the wrong thing. + df = make_df(DATA) + transformer = MathFeatures(variables=["Age", "Marks"], func="sem") + + if make_df is pd.DataFrame: + Xt = transformer.fit_transform(df) + expected = dict(DATA) + expected["sem_Age_Marks"] = [9.55, 10.10, 9.15, 8.70] + assert_df_equal(Xt, expected) + else: + with pytest.raises(NotImplementedError, match="has no NumPy-vectorized"): + transformer.fit_transform(df) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_drop_original_variables(make_df): + df = make_df(DATA) transformer = MathFeatures( variables=["Age", "Marks"], func=["sum", "mean"], drop_original=True, ) + Xt = transformer.fit_transform(df) - X = transformer.fit_transform(df_vartypes) - - ref = pd.DataFrame.from_dict( - { - "Name": ["tom", "nick", "krish", "jack"], - "City": ["London", "Manchester", "Liverpool", "Bristol"], - "dob": dob_datrange, - "sum_Age_Marks": [20.9, 21.8, 19.7, 18.6], - "mean_Age_Marks": [10.45, 10.9, 9.85, 9.3], - } - ) - - pd.testing.assert_frame_equal(X, ref) + expected = { + "Name": DATA["Name"], + "City": DATA["City"], + "sum_Age_Marks": [20.9, 21.8, 19.7, 18.6], + "mean_Age_Marks": [10.45, 10.9, 9.85, 9.3], + } + assert_df_equal(Xt, expected) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("_varnames", [None, ["var1", "var2"]]) @pytest.mark.parametrize("_drop", [True, False]) -def test_get_feature_names_out(_varnames, _drop, df_vartypes): +def test_get_feature_names_out(make_df, _varnames, _drop): + df = make_df(DATA) tr = MathFeatures( variables=["Age", "Marks"], func=["sum", "mean"], new_variables_names=_varnames, drop_original=_drop, ) - X = tr.fit_transform(df_vartypes) - feat_out = list(X.columns) + Xt = tr.fit_transform(df) + feat_out = list(nw.from_native(Xt, eager_only=True).columns) assert tr.get_feature_names_out(input_features=None) == feat_out - assert tr.get_feature_names_out(input_features=df_vartypes.columns) == feat_out +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("_varnames", [None, ["var1", "var2"]]) @pytest.mark.parametrize("_drop", [True, False]) -def test_get_feature_names_out_from_pipeline(_varnames, _drop, df_vartypes): - # set up transformer +def test_get_feature_names_out_from_pipeline(make_df, _varnames, _drop): + df = make_df(DATA) transformer = MathFeatures( variables=["Age", "Marks"], func=["sum", "mean"], @@ -406,24 +433,21 @@ def test_get_feature_names_out_from_pipeline(_varnames, _drop, df_vartypes): ) pipe = Pipeline([("transformer", transformer)]) + Xt = pipe.fit_transform(df) - # fit transformer - X = pipe.fit_transform(df_vartypes) - - feat_out = list(X.columns) + feat_out = list(nw.from_native(Xt, eager_only=True).columns) assert pipe.get_feature_names_out(input_features=None) == feat_out - assert pipe.get_feature_names_out(input_features=df_vartypes.columns) == feat_out +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("_input_features", ["hola", ["Age", "Marks"]]) -def test_get_feature_names_out_raises_error_when_wrong_param( - _input_features, df_vartypes -): +def test_get_feature_names_out_raises_error_when_wrong_param(make_df, _input_features): + df = make_df(DATA) transformer = MathFeatures( variables=["Age", "Marks"], func=["sum", "mean"], ) - transformer.fit(df_vartypes) + transformer.fit(df) with pytest.raises(ValueError): transformer.get_feature_names_out(input_features=_input_features)