diff --git a/docs/user_guide/datetime/DatetimeSubtraction.rst b/docs/user_guide/datetime/DatetimeSubtraction.rst index 662657692..638502472 100644 --- a/docs/user_guide/datetime/DatetimeSubtraction.rst +++ b/docs/user_guide/datetime/DatetimeSubtraction.rst @@ -156,6 +156,42 @@ original variables and also the new variables with the time difference: 4 2019-03-09 2018-04-08 0.917199 +With polars +~~~~~~~~~~~ + +:class:`DatetimeSubtraction()` also works with polars dataframes: + +.. code:: python + + import polars as pl + from feature_engine.datetime import DatetimeSubtraction + + data = pl.DataFrame({ + "date1" : ["2022-09-01", "2022-10-01", "2022-12-01"], + "date2" : ["2022-09-15", "2022-10-15", "2022-12-15"], + "date3" : ["2022-08-01", "2022-09-01", "2022-11-01"], + "date4" : ["2022-08-15", "2022-09-15", "2022-11-15"], + }) + + dtf = DatetimeSubtraction(variables=["date1", "date2"], reference=["date3", "date4"]) + + data = dtf.fit_transform(data) + + print(data) + +.. code:: text + + shape: (3, 8) + ┌────────────┬────────────┬────────────┬────────────┬─────────────────┬─────────────────┬─────────────────┬─────────────────┐ + │ date1 ┆ date2 ┆ date3 ┆ date4 ┆ date1_sub_date3 ┆ date2_sub_date3 ┆ date1_sub_date4 ┆ date2_sub_date4 │ + │ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │ + │ str ┆ str ┆ str ┆ str ┆ f64 ┆ f64 ┆ f64 ┆ f64 │ + ╞════════════╪════════════╪════════════╪════════════╪═════════════════╪═════════════════╪═════════════════╪═════════════════╡ + │ 2022-09-01 ┆ 2022-09-15 ┆ 2022-08-01 ┆ 2022-08-15 ┆ 31.0 ┆ 45.0 ┆ 17.0 ┆ 31.0 │ + │ 2022-10-01 ┆ 2022-10-15 ┆ 2022-09-01 ┆ 2022-09-15 ┆ 30.0 ┆ 44.0 ┆ 16.0 ┆ 30.0 │ + │ 2022-12-01 ┆ 2022-12-15 ┆ 2022-11-01 ┆ 2022-11-15 ┆ 30.0 ┆ 44.0 ┆ 16.0 ┆ 30.0 │ + └────────────┴────────────┴────────────┴────────────┴─────────────────┴─────────────────┴─────────────────┴─────────────────┘ + Drop original variables after computation ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/feature_engine/datetime/datetime_subtraction.py b/feature_engine/datetime/datetime_subtraction.py index 3688253fc..a901b7540 100644 --- a/feature_engine/datetime/datetime_subtraction.py +++ b/feature_engine/datetime/datetime_subtraction.py @@ -1,8 +1,11 @@ -from typing import List, Optional, Union +from datetime import timezone +from typing import Dict, List, Optional, Union +import narwhals as nw +import narwhals.dependencies as nwd import numpy as np -import pandas as pd -from pandas.api.types import is_datetime64_any_dtype as is_datetime +from dateutil.parser import parse as _dateutil_parse +from narwhals.typing import IntoDataFrame, IntoSeries from sklearn.utils.validation import check_is_fitted from feature_engine._check_init_parameters.check_init_input_params import ( @@ -47,6 +50,27 @@ 0 2022-09-18 2022-08-18 31.0 1 2022-10-27 2022-08-27 61.0 2 2022-12-24 2022-06-24 183.0 + + With polars: + + >>> import polars as pl + >>> from feature_engine.datetime import DatetimeSubtraction + >>> X = pl.DataFrame({ + >>> "date1": ["2022-09-18", "2022-10-27", "2022-12-24"], + >>> "date2": ["2022-08-18", "2022-08-27", "2022-06-24"]}) + >>> dtf = DatetimeSubtraction(variables=["date1"], reference=["date2"]) + >>> dtf.fit(X) + >>> dtf.transform(X) + shape: (3, 3) + ┌────────────┬────────────┬─────────────────┐ + │ date1 ┆ date2 ┆ date1_sub_date2 │ + │ --- ┆ --- ┆ --- │ + │ str ┆ str ┆ f64 │ + ╞════════════╪════════════╪═════════════════╡ + │ 2022-09-18 ┆ 2022-08-18 ┆ 31.0 │ + │ 2022-10-27 ┆ 2022-08-27 ┆ 61.0 │ + │ 2022-12-24 ┆ 2022-06-24 ┆ 183.0 │ + └────────────┴────────────┴─────────────────┘ """.rstrip() @@ -219,17 +243,17 @@ def __init__( self.utc = utc self.format = format - def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): + def fit(self, X: IntoDataFrame, y: Optional[IntoSeries] = None): """ This transformer does not learn any parameter. 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 transform. - y: pandas Series, or np.array. Default=None. + y: Series, or np.array. Default=None. It is not needed in this transformer. You can pass y or None. """ # Common checks and attributes @@ -263,29 +287,32 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): # check if dataset contains na if self.missing_values == "raise": - vars = list(set(self.variables_ + self.reference_)) - _check_contains_na(X, vars) + vars_ = list(set(self.variables_ + self.reference_)) + _check_contains_na(X, vars_) # save input features - self.feature_names_in_ = X.columns.tolist() + if nwd.is_pandas_dataframe(X) is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw.from_native(X, eager_only=True).columns # save train set shape self.n_features_in_ = X.shape[1] return self - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """ Add new features. 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 input dataframe plus the new variables. """ @@ -299,43 +326,64 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: _check_X_matches_training_df(X, self.n_features_in_) if self.missing_values == "raise": - vars = list(set(self.variables_ + self.reference_)) - _check_contains_na(X, vars) + vars_ = list(set(self.variables_ + self.reference_)) + _check_contains_na(X, vars_) + + is_pandas = nwd.is_pandas_dataframe(X) # reorder variables to match train set - X = X[self.feature_names_in_] + if is_pandas is True: + X = X[self.feature_names_in_] + else: + X = nw.from_native(X, eager_only=True).select( + self.feature_names_in_ + ).to_native() - X_dt = self._to_datetime(X) + nw_X = nw.from_native(X, eager_only=True) - new_features = self._sub(X_dt) + dt_arrays = self._to_datetime(nw_X, is_pandas) - X = pd.concat([X, new_features], axis=1) + new_series = self._sub(dt_arrays, nw_X.implementation) - if self.drop_original: - X = X.drop( - columns=set(self.variables_ + self.reference_), - ) + nw_X = nw_X.with_columns(*new_series) - return X + if self.drop_original is True: + nw_X = nw_X.drop(list(set(self.variables_ + self.reference_))) - def _to_datetime(self, X: pd.DataFrame): - """convert variables to datetime.""" - # convert datetime variables - datetime_df = pd.concat( - [ - pd.to_datetime( - X[variable], + return nw_X.to_native() + + def _to_datetime( + self, nw_X: nw.DataFrame, is_pandas: bool + ) -> Dict[str, np.ndarray]: + """Convert the variables and reference columns to numpy datetime64 arrays.""" + needed = sorted(set(self.variables_ + self.reference_)) + + # pandas.to_datetime honours dayfirst/yearfirst/utc precisely; grab the + # native namespace once (no `import pandas`) rather than per-column. + if is_pandas is True: + native_ns = nw.get_native_namespace(nw_X) + + arrays = {} + non_dt_columns = [] + for variable in needed: + col = nw_X.get_column(variable) + if is_pandas is True: + parsed_native = native_ns.to_datetime( + col.to_native(), dayfirst=self.dayfirst, yearfirst=self.yearfirst, utc=self.utc, format=self.format, ) - for variable in set(self.variables_ + self.reference_) - ], - axis=1, - ) + parsed = nw.from_native(parsed_native, series_only=True) + else: + parsed = self._parse_non_pandas_column(col) - non_dt_columns = datetime_df.columns[~datetime_df.apply(is_datetime)].tolist() + if not isinstance(parsed.dtype, nw.Datetime): + non_dt_columns.append(variable) + continue + + arrays[variable] = parsed.to_numpy() if non_dt_columns: raise ValueError( @@ -343,23 +391,69 @@ def _to_datetime(self, X: pd.DataFrame): + (len(non_dt_columns) * "{} ").format(*non_dt_columns) + "could not be converted to datetime. Try setting utc=True" ) - return datetime_df - def _sub(self, dt_df: pd.DataFrame): - """make datetime subtraction""" - new_df = pd.DataFrame() - for reference in self.reference_: - new_varnames = [f"{var}_sub_{reference}" for var in self.variables_] - new_df[new_varnames] = ( - dt_df[self.variables_] - .sub(dt_df[reference], axis=0) - .div(np.timedelta64(1, self.output_unit).astype("timedelta64[ns]")) + return arrays + + def _parse_non_pandas_column(self, col: "nw.Series") -> "nw.Series": + """Parse a single non-pandas column to a narwhals Datetime series.""" + if isinstance(col.dtype, nw.Datetime): + return col + if isinstance(col.dtype, nw.Date): + return col.cast(nw.Datetime) + if isinstance(col.dtype, (nw.Categorical, nw.Enum)): + col = col.cast(nw.String) + + try: + return col.str.to_datetime(format=self.format) + except Exception: + if self.format is not None: + raise + # narwhals' vectorized parser needs a single unambiguous format; + # fall back to dateutil per value, same flexible guessing that + # check_datetime_variables already promises across backends. + return self._flexible_parse(col) + + def _flexible_parse(self, col: "nw.Series") -> "nw.Series": + values = [ + None + if value is None + else _dateutil_parse( + value, dayfirst=self.dayfirst, yearfirst=self.yearfirst ) + for value in col.to_list() + ] + if self.utc is True: + values = [ + None + if v is None + else ( + v.astimezone(timezone.utc) + if v.tzinfo is not None + else v.replace(tzinfo=timezone.utc) + ) + for v in values + ] + return nw.new_series(col.name, values, backend=col.implementation) - if self.new_variables_names is not None: - new_df.columns = self.new_variables_names - - return new_df + def _sub(self, dt_arrays: Dict[str, np.ndarray], backend) -> List: + """make datetime subtraction""" + names = self._get_new_features_name() + # "Y"/"M" are non-linear units: numpy can only divide timedeltas by + # them once both sides are cast to a common linear unit (ns), which + # is also what pandas does internally for Timedelta / Timedelta. + unit_td = np.timedelta64(1, self.output_unit).astype("timedelta64[ns]") + + new_series = [] + idx = 0 + for reference in self.reference_: + ref_arr = dt_arrays[reference] + for var in self.variables_: + diff = (dt_arrays[var] - ref_arr).astype("timedelta64[ns]") + result = diff / unit_td + new_series.append(nw.new_series(names[idx], result, backend=backend)) + idx += 1 + + return new_series def _get_new_features_name(self) -> List: """Return names of the created features.""" diff --git a/tests/test_datetime/test_datetime_subtraction.py b/tests/test_datetime/test_datetime_subtraction.py index 4e854d04e..a8c64e614 100644 --- a/tests/test_datetime/test_datetime_subtraction.py +++ b/tests/test_datetime/test_datetime_subtraction.py @@ -1,5 +1,8 @@ -import numpy as np +from datetime import datetime as _datetime + +import narwhals as nw import pandas as pd +import polars as pl import pytest from feature_engine.datetime import DatetimeSubtraction @@ -15,6 +18,38 @@ ) from tests.estimator_checks.non_fitted_error_checks import check_raises_non_fitted_error +DATA_DATETIME = { + "Name": ["tom", "nick", "krish", "jack"], + "Age": [20, 21, 19, 18], + "datetime_range": [ + _datetime(2020, 2, 24), + _datetime(2020, 2, 25), + _datetime(2020, 2, 26), + _datetime(2020, 2, 27), + ], + "date_obj1": ["01-Jan-2010", "24-Feb-1945", "14-Jun-2100", "17-May-1999"], + "date_obj2": ["10/11/12", "12/31/09", "06/30/95", "03/17/04"], + "time_obj": ["21:45:23", "09:15:33", "12:34:59", "03:27:02"], +} + +DATA_NAN = { + "dates_na": ["Feb-2010", None, "Jun-1922", None], + "dates_full": ["Feb-2010", "Mar-2010", "Jun-1922", "Feb-2011"], +} + +DATA_NAN_FILLED = { + "dates_na": ["Feb-2010", "Mar-2010", "Jun-1922", "Mar-2010"], + "dates_full": ["Feb-2010", "Mar-2010", "Jun-1922", "Feb-2011"], +} + + +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 result[col] == pytest.approx(values, abs=abs_tol) + + # ========= init functionality tests @@ -138,24 +173,30 @@ def test_missing_values_raises_error_when_not_valid(param): # ==== fit functionality +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("input_vars", [["Age", "date_obj2"], "Age"]) -def test_raises_error_when_variables_not_datetime(df_datetime, input_vars): +def test_raises_error_when_variables_not_datetime(make_df, input_vars): + df = make_df(DATA_DATETIME) tr = DatetimeSubtraction(variables=input_vars, reference="date_obj1") with pytest.raises(TypeError): - tr.fit(df_datetime) + tr.fit(df) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("input_vars", [["Age", "date_obj2"], "Age"]) -def test_raises_error_when_reference_not_datetime(df_datetime, input_vars): +def test_raises_error_when_reference_not_datetime(make_df, input_vars): + df = make_df(DATA_DATETIME) tr = DatetimeSubtraction(variables=["date_obj1"], reference=input_vars) with pytest.raises(TypeError): - tr.fit(df_datetime) + tr.fit(df) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("input_vars", [["time_obj", "date_obj2"], "date_obj2", None]) -def test_sets_variables_if_datetime(df_datetime, input_vars): +def test_sets_variables_if_datetime(make_df, input_vars): + df = make_df(DATA_DATETIME) tr = DatetimeSubtraction(variables=input_vars, reference=input_vars) - tr.fit(df_datetime) + tr.fit(df) if input_vars is None: dt_vars = ["datetime_range", "date_obj1", "date_obj2", "time_obj"] assert tr.variables_ == dt_vars @@ -168,29 +209,24 @@ def test_sets_variables_if_datetime(df_datetime, input_vars): assert tr.reference_ == ["time_obj", "date_obj2"] +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("new", [["new1", "new2"], ["new1", "new2", "new3"]]) -def test_new_variables_raise_error_if_not_adequate_number(df_datetime, new): +def test_new_variables_raise_error_if_not_adequate_number(make_df, new): + df = make_df(DATA_DATETIME) tr = DatetimeSubtraction( variables="date_obj1", reference="date_obj1", new_variables_names=new ) with pytest.raises(ValueError): - tr.fit(df_datetime) - - -@pytest.fixture -def df_nan(): - df = pd.DataFrame( - { - "dates_na": ["Feb-2010", np.nan, "Jun-1922", np.nan], - "dates_full": ["Feb-2010", "Mar-2010", "Jun-1922", "Feb-2011"], - } - ) - return df + tr.fit(df) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("input_vars_1", ["dates_full", None]) @pytest.mark.parametrize("input_vars_2", ["dates_na", ["dates_full", "dates_na"], None]) -def test_raises_error_when_nan_in_variables_in_fit(df_nan, input_vars_1, input_vars_2): +def test_raises_error_when_nan_in_variables_in_fit( + make_df, input_vars_1, input_vars_2 +): + df_nan = make_df(DATA_NAN) tr = DatetimeSubtraction( variables=input_vars_2, reference=input_vars_1, missing_values="raise" ) @@ -198,9 +234,13 @@ def test_raises_error_when_nan_in_variables_in_fit(df_nan, input_vars_1, input_v tr.fit(df_nan) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("input_vars_1", ["dates_full", None]) @pytest.mark.parametrize("input_vars_2", ["dates_na", ["dates_full", "dates_na"], None]) -def test_raises_error_when_nan_in_reference_in_fit(df_nan, input_vars_1, input_vars_2): +def test_raises_error_when_nan_in_reference_in_fit( + make_df, input_vars_1, input_vars_2 +): + df_nan = make_df(DATA_NAN) tr = DatetimeSubtraction( variables=input_vars_1, reference=input_vars_2, missing_values="raise" ) @@ -209,32 +249,35 @@ def test_raises_error_when_nan_in_reference_in_fit(df_nan, input_vars_1, input_v # transform tests +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("input_vars_1", ["dates_full", None]) @pytest.mark.parametrize("input_vars_2", ["dates_na", ["dates_full", "dates_na"], None]) def test_raises_error_when_nan_in_variables_in_transform( - df_nan, input_vars_1, input_vars_2 + make_df, input_vars_1, input_vars_2 ): tr = DatetimeSubtraction( variables=input_vars_2, reference=input_vars_1, missing_values="raise" ) - tr.fit(df_nan.fillna("Mar-2010")) + tr.fit(make_df(DATA_NAN_FILLED)) with pytest.raises(ValueError): - tr.transform(df_nan) + tr.transform(make_df(DATA_NAN)) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize("input_vars_1", ["dates_full", None]) @pytest.mark.parametrize("input_vars_2", ["dates_na", ["dates_full", "dates_na"], None]) def test_raises_error_when_nan_in_reference_in_transform( - df_nan, input_vars_1, input_vars_2 + make_df, input_vars_1, input_vars_2 ): tr = DatetimeSubtraction( variables=input_vars_1, reference=input_vars_2, missing_values="raise" ) - tr.fit(df_nan.fillna("Mar-2010")) + tr.fit(make_df(DATA_NAN_FILLED)) with pytest.raises(ValueError): - tr.transform(df_nan) + tr.transform(make_df(DATA_NAN)) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "unit, expected", [ @@ -243,91 +286,77 @@ def test_raises_error_when_nan_in_reference_in_transform( ("h", [744.0, 1464.0, 4392.0]), ], ) -def test_subtraction_units(unit, expected): - df_input = pd.DataFrame( - { - "date1": ["2022-09-18", "2022-10-27", "2022-12-24"], - "date2": ["2022-08-18", "2022-08-27", "2022-06-24"], - } - ) - df_expected = pd.DataFrame( - { - "date1": ["2022-09-18", "2022-10-27", "2022-12-24"], - "date2": ["2022-08-18", "2022-08-27", "2022-06-24"], - "date1_sub_date2": expected, - } - ) +def test_subtraction_units(make_df, unit, expected): + data = { + "date1": ["2022-09-18", "2022-10-27", "2022-12-24"], + "date2": ["2022-08-18", "2022-08-27", "2022-06-24"], + } + df_input = make_df(data) dtf = DatetimeSubtraction( variables=["date1"], reference=["date2"], output_unit=unit ) df_output = dtf.fit_transform(df_input) - pd.testing.assert_frame_equal(df_output, df_expected, check_dtype=False) + expected_dict = dict(data) + expected_dict["date1_sub_date2"] = expected + assert_df_equal(df_output, expected_dict) + + +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_multiple_subtractions(make_df): + data = { + "date1": ["2022-09-01", "2022-10-01", "2022-12-01"], + "date2": ["2022-09-15", "2022-10-15", "2022-12-15"], + "date3": ["2022-08-01", "2022-09-01", "2022-11-01"], + "date4": ["2022-08-15", "2022-09-15", "2022-11-15"], + } + df_input = make_df(data) + + expected = dict(data) + expected["date1_sub_date3"] = [31, 30, 30] + expected["date2_sub_date3"] = [45, 44, 44] + expected["date1_sub_date4"] = [17, 16, 16] + expected["date2_sub_date4"] = [31, 30, 30] -def test_multiple_subtractions(): - df_input = pd.DataFrame( - { - "date1": ["2022-09-01", "2022-10-01", "2022-12-01"], - "date2": ["2022-09-15", "2022-10-15", "2022-12-15"], - "date3": ["2022-08-01", "2022-09-01", "2022-11-01"], - "date4": ["2022-08-15", "2022-09-15", "2022-11-15"], - } - ) - df_expected = pd.DataFrame( - { - "date1": ["2022-09-01", "2022-10-01", "2022-12-01"], - "date2": ["2022-09-15", "2022-10-15", "2022-12-15"], - "date3": ["2022-08-01", "2022-09-01", "2022-11-01"], - "date4": ["2022-08-15", "2022-09-15", "2022-11-15"], - "date1_sub_date3": [31, 30, 30], - "date2_sub_date3": [45, 44, 44], - "date1_sub_date4": [17, 16, 16], - "date2_sub_date4": [31, 30, 30], - } - ) dtf = DatetimeSubtraction( variables=["date1", "date2"], reference=["date3", "date4"] ) df_output = dtf.fit_transform(df_input) - pd.testing.assert_frame_equal(df_output, df_expected, check_dtype=False) + assert_df_equal(df_output, expected) -def test_assigns_new_variable_names(): - df_input = pd.DataFrame( - { - "date1": ["2022-09-01", "2022-10-01", "2022-12-01"], - "date2": ["2022-09-15", "2022-10-15", "2022-12-15"], - "date3": ["2022-08-01", "2022-09-01", "2022-11-01"], - "date4": ["2022-08-15", "2022-09-15", "2022-11-15"], - } - ) - df_expected = pd.DataFrame( - { - "date1": ["2022-09-01", "2022-10-01", "2022-12-01"], - "date2": ["2022-09-15", "2022-10-15", "2022-12-15"], - "date3": ["2022-08-01", "2022-09-01", "2022-11-01"], - "date4": ["2022-08-15", "2022-09-15", "2022-11-15"], - "new1": [31, 30, 30], - "new2": [45, 44, 44], - "new3": [17, 16, 16], - "new4": [31, 30, 30], - } - ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_assigns_new_variable_names(make_df): + data = { + "date1": ["2022-09-01", "2022-10-01", "2022-12-01"], + "date2": ["2022-09-15", "2022-10-15", "2022-12-15"], + "date3": ["2022-08-01", "2022-09-01", "2022-11-01"], + "date4": ["2022-08-15", "2022-09-15", "2022-11-15"], + } + df_input = make_df(data) + + expected = dict(data) + expected["new1"] = [31, 30, 30] + expected["new2"] = [45, 44, 44] + expected["new3"] = [17, 16, 16] + expected["new4"] = [31, 30, 30] + dtf = DatetimeSubtraction( variables=["date1", "date2"], reference=["date3", "date4"], new_variables_names=["new1", "new2", "new3", "new4"], ) df_output = dtf.fit_transform(df_input) - pd.testing.assert_frame_equal(df_output, df_expected, check_dtype=False) + assert_df_equal(df_output, expected) # additional methods -def test_get_feature_names_out(): - df = pd.DataFrame( +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_get_feature_names_out(make_df): + df = make_df( { "d1": ["Feb-2010", "Mar-2010", "Jun-1922", "Feb-2011"], "d2": ["Feb-2010", "Mar-2010", "Jun-1922", "Feb-2011"], @@ -335,7 +364,7 @@ def test_get_feature_names_out(): "d4": ["Feb-2010", "Mar-2010", "Jun-1922", "Feb-2011"], } ) - input_vars = df.columns.to_list() + input_vars = list(nw.from_native(df, eager_only=True).columns) tr = DatetimeSubtraction(variables="d1", reference="d2") tr.fit(df)