From d8e38fd10659128f83004b81a9325662a9a2f564 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Wed, 29 Jul 2026 16:28:23 +0200 Subject: [PATCH 1/4] update variable handling module for narwahls --- .../variable_handling/check_all_variables.rst | 46 ++- .../check_categorical_variables.rst | 52 ++- .../check_datetime_variables.rst | 51 +++ .../check_numerical_variables.rst | 29 ++ .../variable_handling/find_all_variables.rst | 66 +++- ...nd_categorical_and_numerical_variables.rst | 48 ++- .../find_categorical_variables.rst | 56 +++ .../find_datetime_variables.rst | 60 ++++ .../find_numerical_variables.rst | 29 ++ .../retain_variables_if_in_df.rst | 29 ++ .../_variable_type_checks.py | 78 +++++ .../variable_handling/check_variables.py | 105 ++++-- .../variable_handling/find_variables.py | 200 ++++++++--- .../variable_handling/retain_variables.py | 13 +- tests/test_variable_handling/conftest.py | 41 +++ .../test_check_variables.py | 136 +++++--- .../test_fe_type_checks.py | 107 ++++++ .../test_find_variables.py | 323 +++++++++--------- .../test_remove_variables.py | 43 ++- 19 files changed, 1219 insertions(+), 293 deletions(-) diff --git a/docs/user_guide/variable_handling/check_all_variables.rst b/docs/user_guide/variable_handling/check_all_variables.rst index d35f7f53d..414002960 100644 --- a/docs/user_guide/variable_handling/check_all_variables.rst +++ b/docs/user_guide/variable_handling/check_all_variables.rst @@ -89,4 +89,48 @@ Below we see the error message: .. code:: python - KeyError: 'Some of the variables are not in the dataframe.' \ No newline at end of file + KeyError: 'Some of the variables are not in the dataframe.' + +With polars +----------- + +:class:`check_all_variables()` works in the same way with a polars dataframe. +Let's create an equivalent toy dataset: + +.. code:: python + + import polars as pl + from datetime import datetime, timedelta + from sklearn.datasets import make_classification + from feature_engine.variable_handling import check_all_variables + + X, y = make_classification( + n_samples=1000, + n_features=4, + n_redundant=1, + n_clusters_per_class=1, + weights=[0.50], + class_sep=2, + random_state=1, + ) + + colnames = [f"num_var_{i+1}" for i in range(4)] + X = pl.DataFrame(X, schema=colnames) + + X = X.with_columns( + pl.lit("Hello").alias("cat_var1"), + pl.lit("Bye").alias("cat_var2"), + pl.Series("date1", [datetime(2020, 2, 24) + timedelta(minutes=i) for i in range(1000)]), + pl.Series("date2", [datetime(2021, 9, 29) + timedelta(hours=i) for i in range(1000)]), + pl.lit("2020-02-24").alias("date3"), + ) + + checked_vars = check_all_variables(X, ["num_var_1", "cat_var1", "date1"]) + + checked_vars + +The output is the list of variable names passed to the function: + +.. code:: python + + ['num_var_1', 'cat_var1', 'date1'] \ No newline at end of file diff --git a/docs/user_guide/variable_handling/check_categorical_variables.rst b/docs/user_guide/variable_handling/check_categorical_variables.rst index d311decd6..90443145b 100644 --- a/docs/user_guide/variable_handling/check_categorical_variables.rst +++ b/docs/user_guide/variable_handling/check_categorical_variables.rst @@ -89,4 +89,54 @@ Below we see the error message: .. code:: python TypeError: Some of the variables are not categorical. Please cast them as object - or categorical before using this transformer. \ No newline at end of file + or categorical before using this transformer. + +With polars +----------- + +:class:`check_categorical_variables()` works in the same way with a polars dataframe. +Let's create an equivalent toy dataset: + +.. code:: python + + import polars as pl + from datetime import datetime, timedelta + from sklearn.datasets import make_classification + from feature_engine.variable_handling import check_categorical_variables + + X, y = make_classification( + n_samples=1000, + n_features=4, + n_redundant=1, + n_clusters_per_class=1, + weights=[0.50], + class_sep=2, + random_state=1, + ) + + colnames = [f"num_var_{i+1}" for i in range(4)] + X = pl.DataFrame(X, schema=colnames) + + X = X.with_columns( + pl.lit("Hello").alias("cat_var1"), + pl.lit("Bye").alias("cat_var2"), + pl.Series("date1", [datetime(2020, 2, 24) + timedelta(minutes=i) for i in range(1000)]), + pl.Series("date2", [datetime(2021, 9, 29) + timedelta(hours=i) for i in range(1000)]), + pl.lit("2020-02-24").alias("date3"), + ) + + var_cat = check_categorical_variables(X, ["cat_var1", "date3"]) + + var_cat + +Both variables are of type string and hence, will be in the resulting list: + +.. code:: python + + ['cat_var1', 'date3'] + +.. note:: + + Polars has no separate "object" dtype the way pandas does - its `String` dtype + is the only way to represent free-form text, so it is accepted here as + categorical, alongside `Categorical` and `Enum` columns. \ No newline at end of file diff --git a/docs/user_guide/variable_handling/check_datetime_variables.rst b/docs/user_guide/variable_handling/check_datetime_variables.rst index e1fb86168..36923b483 100644 --- a/docs/user_guide/variable_handling/check_datetime_variables.rst +++ b/docs/user_guide/variable_handling/check_datetime_variables.rst @@ -92,3 +92,54 @@ Below the error message: .. code:: python TypeError: Some of the variables are not or cannot be parsed as datetime. + +With polars +----------- + +:class:`check_datetime_variables()` works in the same way with a polars dataframe. +Let's create an equivalent toy dataset: + +.. code:: python + + import polars as pl + from datetime import datetime, timedelta + from sklearn.datasets import make_classification + from feature_engine.variable_handling import check_datetime_variables + + X, y = make_classification( + n_samples=1000, + n_features=4, + n_redundant=1, + n_clusters_per_class=1, + weights=[0.50], + class_sep=2, + random_state=1, + ) + + colnames = [f"num_var_{i+1}" for i in range(4)] + X = pl.DataFrame(X, schema=colnames) + + X = X.with_columns( + pl.lit("Hello").alias("cat_var1"), + pl.lit("Bye").alias("cat_var2"), + pl.Series("date1", [datetime(2020, 2, 24) + timedelta(minutes=i) for i in range(1000)]), + pl.Series("date2", [datetime(2021, 9, 29) + timedelta(hours=i) for i in range(1000)]), + pl.lit("2020-02-24").alias("date3"), + ) + + var_date = check_datetime_variables(X, ["date2", "date3"]) + + var_date + +In this case, both variables, if they can be parsed as datetime, will be in the +resulting list: + +.. code:: python + + ['date2', 'date3'] + +.. note:: + + `date3` here is an ISO-8601 string, which polars (via narwhals) can parse. + Pandas' flexible, dateutil-backed date-string guessing recognises additional + formats that polars does not - see the note in :class:`find_datetime_variables()`. diff --git a/docs/user_guide/variable_handling/check_numerical_variables.rst b/docs/user_guide/variable_handling/check_numerical_variables.rst index 795376850..c27f95fd0 100644 --- a/docs/user_guide/variable_handling/check_numerical_variables.rst +++ b/docs/user_guide/variable_handling/check_numerical_variables.rst @@ -67,3 +67,32 @@ Below we see the error message: TypeError: Some of the variables are not numerical. Please cast them as numerical before using this transformer. + +With polars +----------- + +:class:`check_numerical_variables()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from datetime import datetime + from feature_engine.variable_handling import check_numerical_variables + + df = pl.DataFrame({ + "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": [datetime(2020, 2, 24, 0, i) for i in range(4)], + }) + + var_num = check_numerical_variables(df, ['Age', 'Marks']) + + var_num + +If the variables are numerical, the function returns their names in a list: + +.. code:: python + + ['Age', 'Marks'] diff --git a/docs/user_guide/variable_handling/find_all_variables.rst b/docs/user_guide/variable_handling/find_all_variables.rst index 055fac788..8c7835dd0 100644 --- a/docs/user_guide/variable_handling/find_all_variables.rst +++ b/docs/user_guide/variable_handling/find_all_variables.rst @@ -125,4 +125,68 @@ However, this command returns an empty list: X[[ 'date1', 'date2', 'date3']], exclude_datetime=True, return_empty=True, - ) \ No newline at end of file + ) + +With polars +----------- + +:class:`find_all_variables()` works in the same way with a polars dataframe. +Let's create an equivalent toy dataset: + +.. code:: python + + import polars as pl + from datetime import datetime, timedelta + from sklearn.datasets import make_classification + from feature_engine.variable_handling import find_all_variables + + X, y = make_classification( + n_samples=1000, + n_features=4, + n_redundant=1, + n_clusters_per_class=1, + weights=[0.50], + class_sep=2, + random_state=1, + ) + + colnames = [f"num_var_{i+1}" for i in range(4)] + X = pl.DataFrame(X, schema=colnames) + + X = X.with_columns( + pl.lit("Hello").alias("cat_var1"), + pl.lit("Bye").alias("cat_var2"), + pl.Series("date1", [datetime(2020, 2, 24) + timedelta(minutes=i) for i in range(1000)]), + pl.Series("date2", [datetime(2021, 9, 29) + timedelta(hours=i) for i in range(1000)]), + pl.lit("2020-02-24").alias("date3"), + ) + + vars_all = find_all_variables(X) + + vars_all + +We see the variable names in the list below: + +.. code:: python + + ['num_var_1', + 'num_var_2', + 'num_var_3', + 'num_var_4', + 'cat_var1', + 'cat_var2', + 'date1', + 'date2', + 'date3'] + +And, as with pandas, we can exclude the datetime variables: + +.. code:: python + + vars_all = find_all_variables(X, exclude_datetime=True) + + vars_all + +.. code:: python + + ['num_var_1', 'num_var_2', 'num_var_3', 'num_var_4', 'cat_var1', 'cat_var2'] \ No newline at end of file diff --git a/docs/user_guide/variable_handling/find_categorical_and_numerical_variables.rst b/docs/user_guide/variable_handling/find_categorical_and_numerical_variables.rst index 05d5807cc..0b27dc70a 100644 --- a/docs/user_guide/variable_handling/find_categorical_and_numerical_variables.rst +++ b/docs/user_guide/variable_handling/find_categorical_and_numerical_variables.rst @@ -126,4 +126,50 @@ To return empty lists instead, we set `return_empty` to `True`: find_categorical_and_numerical_variables( X[[ 'date1', 'date2', 'date3']], return_empty = True - ) \ No newline at end of file + ) + +With polars +----------- + +:class:`find_categorical_and_numerical_variables()` works in the same way with a +polars dataframe. Let's create an equivalent toy dataset: + +.. code:: python + + import polars as pl + from datetime import datetime, timedelta + from sklearn.datasets import make_classification + from feature_engine.variable_handling import find_categorical_and_numerical_variables + + X, y = make_classification( + n_samples=1000, + n_features=4, + n_redundant=1, + n_clusters_per_class=1, + weights=[0.50], + class_sep=2, + random_state=1, + ) + + colnames = [f"num_var_{i+1}" for i in range(4)] + X = pl.DataFrame(X, schema=colnames) + + X = X.with_columns( + pl.lit("Hello").alias("cat_var1"), + pl.lit("Bye").alias("cat_var2"), + pl.Series("date1", [datetime(2020, 2, 24) + timedelta(minutes=i) for i in range(1000)]), + pl.Series("date2", [datetime(2021, 9, 29) + timedelta(hours=i) for i in range(1000)]), + pl.lit("2020-02-24").alias("date3"), + ) + + var_cat, var_num = find_categorical_and_numerical_variables(X) + + var_cat, var_num + +Below we see the names of the categorical variables, followed by the names of the +numerical variables: + +.. code:: python + + (['cat_var1', 'cat_var2'], + ['num_var_1', 'num_var_2', 'num_var_3', 'num_var_4']) \ No newline at end of file diff --git a/docs/user_guide/variable_handling/find_categorical_variables.rst b/docs/user_guide/variable_handling/find_categorical_variables.rst index a00de72af..7eba698cf 100644 --- a/docs/user_guide/variable_handling/find_categorical_variables.rst +++ b/docs/user_guide/variable_handling/find_categorical_variables.rst @@ -90,3 +90,59 @@ To return an empty list instead of the error we need to set `return_empty` to `T follows: `find_categorical_variables(X[colnames], return_empty=True)`. The previous command returns an empty list: `[]`. + +With polars +----------- + +:class:`find_categorical_variables()` works in the same way with a polars dataframe. +Let's create an equivalent toy dataset: + +.. code:: python + + import polars as pl + from datetime import datetime, timedelta + from sklearn.datasets import make_classification + from feature_engine.variable_handling import find_categorical_variables + + X, y = make_classification( + n_samples=1000, + n_features=4, + n_redundant=1, + n_clusters_per_class=1, + weights=[0.50], + class_sep=2, + random_state=1, + ) + + colnames = [f"num_var_{i+1}" for i in range(4)] + X = pl.DataFrame(X, schema=colnames) + + X = X.with_columns( + pl.lit("Hello").alias("cat_var1"), + pl.lit("Bye").alias("cat_var2"), + pl.Series("date1", [datetime(2020, 2, 24) + timedelta(minutes=i) for i in range(1000)]), + pl.Series("date2", [datetime(2021, 9, 29) + timedelta(hours=i) for i in range(1000)]), + pl.lit("2020-02-24").alias("date3"), + ) + +Now let's find the categorical variables: + +.. code:: python + + var_cat = find_categorical_variables(X) + + var_cat + +We see the variable names in the list below: + +.. code:: python + + ['cat_var1', 'cat_var2'] + +.. note:: + + Unlike for pandas, `date3` above is a plain ISO-8601 string column, which polars + (via narwhals) *can* auto-detect as a datetime candidate, so it is still correctly + excluded here. Only pandas' flexible, dateutil-backed date-string guessing (formats + like "24-Feb-2020" or "02/24/20") has no polars equivalent - see the note in + :class:`find_datetime_variables()`. diff --git a/docs/user_guide/variable_handling/find_datetime_variables.rst b/docs/user_guide/variable_handling/find_datetime_variables.rst index 443a43d4e..0c4fe1920 100644 --- a/docs/user_guide/variable_handling/find_datetime_variables.rst +++ b/docs/user_guide/variable_handling/find_datetime_variables.rst @@ -87,3 +87,63 @@ can be parsed as datetime, it will be captured in the list as well. If there are no datetime variables, :class:`find_datetime_variables()` will raise an error. To return an empty list instead, use the argument `return_empty` to `True`. + +With polars +----------- + +:class:`find_datetime_variables()` works in the same way with a polars dataframe. +Let's create an equivalent toy dataset: + +.. code:: python + + import polars as pl + from datetime import datetime, timedelta + from sklearn.datasets import make_classification + from feature_engine.variable_handling import find_datetime_variables + + X, y = make_classification( + n_samples=1000, + n_features=4, + n_redundant=1, + n_clusters_per_class=1, + weights=[0.50], + class_sep=2, + random_state=1, + ) + + colnames = [f"num_var_{i+1}" for i in range(4)] + X = pl.DataFrame(X, schema=colnames) + + X = X.with_columns( + pl.lit("Hello").alias("cat_var1"), + pl.lit("Bye").alias("cat_var2"), + pl.Series("date1", [datetime(2020, 2, 24) + timedelta(minutes=i) for i in range(1000)]), + pl.Series("date2", [datetime(2021, 9, 29) + timedelta(hours=i) for i in range(1000)]), + pl.lit("2020-02-24").alias("date3"), + ) + +The dataframe has 3 datetime variables: two of them are native polars `Datetime` +columns, and one, `date3`, is an ISO-8601 string. Let's capture all 3: + +.. code:: python + + var_date = find_datetime_variables(X) + + var_date + +Below we see the variable names in the list: + +.. code:: python + + ['date1', 'date2', 'date3'] + +.. note:: + + For pandas dataframes, string columns are parsed with pandas' flexible, + dateutil-backed date guessing, so formats like "24-Feb-2020" or "02/24/20" are + recognised as datetime, in addition to ISO-8601 strings. Polars (and other + non-pandas dataframes, through narwhals) has no equivalent flexible guesser, so + only ISO-8601 strings and native `Date`/`Datetime` columns are recognised + automatically. If your polars dataframe has date strings in another format, + pass `variables` explicitly to the transformer instead of relying on + auto-detection. diff --git a/docs/user_guide/variable_handling/find_numerical_variables.rst b/docs/user_guide/variable_handling/find_numerical_variables.rst index fcb4b40c4..aa315114d 100644 --- a/docs/user_guide/variable_handling/find_numerical_variables.rst +++ b/docs/user_guide/variable_handling/find_numerical_variables.rst @@ -68,3 +68,32 @@ need to set `return_empty` to `True`: find_numerical_variables(df[["Name", "City", "dob"]], return_empty=True) The previous command returns an empty list: `[]`. + +With polars +----------- + +:class:`find_numerical_variables()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from datetime import datetime + from feature_engine.variable_handling import find_numerical_variables + + df = pl.DataFrame({ + "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": [datetime(2020, 2, 24, 0, i) for i in range(4)], + }) + + var_num = find_numerical_variables(df) + + var_num + +We see the names of the numerical variables in the list below: + +.. code:: python + + ['Age', 'Marks'] diff --git a/docs/user_guide/variable_handling/retain_variables_if_in_df.rst b/docs/user_guide/variable_handling/retain_variables_if_in_df.rst index 4f4ab71cd..e95fc36dc 100644 --- a/docs/user_guide/variable_handling/retain_variables_if_in_df.rst +++ b/docs/user_guide/variable_handling/retain_variables_if_in_df.rst @@ -59,6 +59,35 @@ We see the names of the subset of variables that are in the dataframe below: If none of variables in the list are in the dataset, :class:`retain_variables_if_in_df()` will raise an error. +With polars +----------- + +:class:`retain_variables_if_in_df()` works in the same way with a polars dataframe: + +.. code:: python + + import polars as pl + from datetime import datetime + from feature_engine.variable_handling import retain_variables_if_in_df + + df = pl.DataFrame({ + "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": [datetime(2020, 2, 24, 0, i) for i in range(4)], + }) + + vars_in_df = retain_variables_if_in_df(df, variables = ["Name", "City", "Dogs"]) + + vars_in_df + +We see the names of the subset of variables that are in the dataframe below: + +.. code:: python + + ['Name', 'City'] + Uses ---- diff --git a/feature_engine/variable_handling/_variable_type_checks.py b/feature_engine/variable_handling/_variable_type_checks.py index 17eb4e41d..7e720b954 100644 --- a/feature_engine/variable_handling/_variable_type_checks.py +++ b/feature_engine/variable_handling/_variable_type_checks.py @@ -1,8 +1,20 @@ +import narwhals as nw import pandas as pd from pandas.api.types import is_object_dtype, is_string_dtype from pandas.core.dtypes.common import is_datetime64_any_dtype as is_datetime from pandas.core.dtypes.common import is_numeric_dtype as is_numeric +# --------------------------------------------------------------------------- +# pandas-only implementation. +# +# These functions rely on pandas' flexible, dateutil-backed `pd.to_datetime` +# string guessing and on pandas' `object` dtype (which, unlike any narwhals +# dtype, can hold arbitrary non-string Python objects). Neither has a +# polars/narwhals equivalent, so they are kept exactly as they were before the +# narwhals migration and are only ever called on pandas input. See the `_nw_*` +# functions below for the polars/narwhals-backend equivalents. +# --------------------------------------------------------------------------- + def is_object(s) -> bool: return is_object_dtype(s) or is_string_dtype(s) @@ -60,3 +72,69 @@ def _is_categorical_and_is_datetime(column: pd.Series) -> bool: is_dt = False return is_dt + + +# --------------------------------------------------------------------------- +# narwhals implementation, used for every backend other than pandas (polars, +# in practice). +# +# narwhals has no lenient/"try" cast (no `strict=False`, unlike raw polars) +# and its `str.to_datetime()` requires ISO-8601 or an explicit `format=` - it +# cannot reproduce pandas' dateutil-based guessing. So a string column such as +# "01-Jan-2010" or "10/11/12" is not auto-detected as datetime for polars, +# even though it is for pandas. ISO-8601 strings and native Date/Datetime +# columns are detected correctly. Users can always pass `variables` explicitly +# to sidestep this. +# --------------------------------------------------------------------------- + + +def _nw_is_date_or_datetime(dtype) -> bool: + # nw.selectors.datetime() only matches Datetime, not Date, so this needs + # its own explicit check. + return isinstance(dtype, (nw.Date, nw.Datetime)) + + +def _nw_is_convertible_to_num(s: "nw.Series") -> bool: + try: + s.cast(nw.String()).cast(nw.Float64()) + except Exception: + return False + return True + + +def _nw_is_convertible_to_dt(s: "nw.Series") -> bool: + try: + s.cast(nw.String()).str.to_datetime() + except Exception: + return False + return True + + +def _nw_is_categorical_and_is_not_datetime(s: "nw.Series") -> bool: + if isinstance(s.dtype, nw.Enum): + # an explicit, user-defined category set is an unambiguous categorical + # signal, unlike a generic string column, so skip the datetime check + return True + + if isinstance(s.dtype, nw.Categorical): + # polars categorical categories are always string-backed, unlike + # pandas' pd.Categorical, which can have numeric categories + return not _nw_is_convertible_to_dt(s) + + if isinstance(s.dtype, nw.String): + return _nw_is_convertible_to_num(s) or not _nw_is_convertible_to_dt(s) + + return False + + +def _nw_is_categorical_and_is_datetime(s: "nw.Series") -> bool: + if isinstance(s.dtype, nw.Enum): + return False + + if isinstance(s.dtype, nw.Categorical): + return _nw_is_convertible_to_dt(s) + + if isinstance(s.dtype, nw.String): + return not _nw_is_convertible_to_num(s) and _nw_is_convertible_to_dt(s) + + return False diff --git a/feature_engine/variable_handling/check_variables.py b/feature_engine/variable_handling/check_variables.py index 76c4ea7c3..b823063ba 100644 --- a/feature_engine/variable_handling/check_variables.py +++ b/feature_engine/variable_handling/check_variables.py @@ -2,11 +2,15 @@ from typing import List, Union -import pandas as pd -from pandas.api.types import is_numeric_dtype as is_numeric +import narwhals as nw +import narwhals.dependencies as nwd +from narwhals.typing import IntoDataFrame +from pandas.core.dtypes.common import is_numeric_dtype as is_numeric from feature_engine.variable_handling._variable_type_checks import ( _is_categorical_and_is_datetime, + _nw_is_categorical_and_is_datetime, + _nw_is_date_or_datetime, ) from feature_engine.variable_handling.dtypes import DATETIME_TYPES @@ -14,7 +18,7 @@ def check_numerical_variables( - X: pd.DataFrame, variables: Variables + X: IntoDataFrame, variables: Variables ) -> List[Union[str, int]]: """ Checks that the variables in the list are of type numerical. @@ -23,8 +27,9 @@ def check_numerical_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset. + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. variables : List The list with the names of the variables to check. @@ -51,7 +56,15 @@ def check_numerical_variables( if isinstance(variables, (str, int)): variables = [variables] - if len(X[variables].select_dtypes(exclude="number").columns) > 0: + if nwd.is_pandas_dataframe(X): + not_numerical = len(X[variables].select_dtypes(exclude="number").columns) > 0 + else: + sub_X = nw.from_native(X, eager_only=True).select(variables) + not_numerical = any( + not sub_X.schema[column].is_numeric() for column in sub_X.columns + ) + + if not_numerical: raise TypeError( "Some of the variables are not numerical. Please cast them as " "numerical before using this transformer." @@ -61,7 +74,7 @@ def check_numerical_variables( def check_categorical_variables( - X: pd.DataFrame, variables: Variables + X: IntoDataFrame, variables: Variables ) -> List[Union[str, int]]: """ Checks that the variables in the list are of type object or categorical. @@ -70,8 +83,9 @@ def check_categorical_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. variables : list The list with the names of the variables to check. @@ -81,6 +95,13 @@ def check_categorical_variables( variables: List The names of the categorical variables. + Notes + ----- + For polars (and other non-pandas dataframes), plain string columns are + accepted as categorical. Polars has no separate "object" dtype the way + pandas does, so its `String` dtype is the only way to represent free-form + text and is treated as categorical here. + Examples -------- >>> import pandas as pd @@ -98,7 +119,18 @@ def check_categorical_variables( if isinstance(variables, (str, int)): variables = [variables] - if len(X[variables].select_dtypes(exclude=["O", "category"]).columns) > 0: + if nwd.is_pandas_dataframe(X): + not_categorical = ( + len(X[variables].select_dtypes(exclude=["O", "category"]).columns) > 0 + ) + else: + sub_X = nw.from_native(X, eager_only=True).select(variables) + not_categorical = any( + not isinstance(sub_X.schema[column], (nw.Categorical, nw.Enum, nw.String)) + for column in sub_X.columns + ) + + if not_categorical: raise TypeError( "Some of the variables are not categorical. Please cast them as " "object or categorical before using this transformer." @@ -108,7 +140,7 @@ def check_categorical_variables( def check_datetime_variables( - X: pd.DataFrame, + X: IntoDataFrame, variables: Variables, ) -> List[Union[str, int]]: """ @@ -119,8 +151,9 @@ def check_datetime_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. variables : list The list with the names of the variables to check. @@ -130,6 +163,13 @@ def check_datetime_variables( variables: List The names of the datetime variables. + Notes + ----- + For pandas dataframes, string columns are parsed with pandas' flexible, + dateutil-backed date guessing. For polars (and other non-pandas dataframes), + only ISO-8601 strings and native `Date`/`Datetime` columns are recognised - + polars has no equivalent flexible guesser. + Examples -------- >>> import pandas as pd @@ -147,11 +187,27 @@ def check_datetime_variables( if isinstance(variables, (str, int)): variables = [variables] - # find non datetime variables, if any: - non_datetime_vars = [] - for column in X[variables].select_dtypes(exclude=DATETIME_TYPES): - if is_numeric(X[column]) or not _is_categorical_and_is_datetime(X[column]): - non_datetime_vars.append(column) + if nwd.is_pandas_dataframe(X): + # find non datetime variables, if any: + non_datetime_vars = [] + for column in X[variables].select_dtypes(exclude=DATETIME_TYPES): + if is_numeric(X[column]) or not _is_categorical_and_is_datetime( + X[column] + ): + non_datetime_vars.append(column) + else: + sub_X = nw.from_native(X, eager_only=True).select(variables) + candidates = [ + column + for column in sub_X.columns + if not _nw_is_date_or_datetime(sub_X.schema[column]) + ] + non_datetime_vars = [ + column + for column in candidates + if sub_X.schema[column].is_numeric() + or not _nw_is_categorical_and_is_datetime(sub_X[column]) + ] if len(non_datetime_vars) > 0: raise TypeError( @@ -162,7 +218,7 @@ def check_datetime_variables( def check_all_variables( - X: pd.DataFrame, + X: IntoDataFrame, variables: Variables, ) -> List[Union[str, int]]: """ @@ -172,8 +228,9 @@ def check_all_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. variables : list The list with the names of the variables to check. @@ -196,13 +253,15 @@ def check_all_variables( >>> vars_all ['var_num', 'var_cat', 'var_date'] """ + columns = nw.from_native(X, eager_only=True).columns + if isinstance(variables, (str, int)): - if variables not in X.columns.to_list(): + if variables not in columns: raise KeyError(f"The variable {variables} is not in the dataframe.") variables_ = [variables] else: - if not set(variables).issubset(set(X.columns)): + if not set(variables).issubset(set(columns)): raise KeyError("Some of the variables are not in the dataframe.") variables_ = variables diff --git a/feature_engine/variable_handling/find_variables.py b/feature_engine/variable_handling/find_variables.py index 5d072eb56..2244b95b1 100644 --- a/feature_engine/variable_handling/find_variables.py +++ b/feature_engine/variable_handling/find_variables.py @@ -3,19 +3,31 @@ import warnings from typing import List, Tuple, Union -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +from narwhals.typing import IntoDataFrame from pandas.api.types import is_datetime64_any_dtype as is_datetime from pandas.core.dtypes.common import is_numeric_dtype as is_numeric from feature_engine.variable_handling._variable_type_checks import ( _is_categorical_and_is_datetime, _is_categorical_and_is_not_datetime, + _nw_is_categorical_and_is_datetime, + _nw_is_categorical_and_is_not_datetime, + _nw_is_date_or_datetime, ) from feature_engine.variable_handling.dtypes import DATETIME_TYPES +# columns of these narwhals dtypes are candidates for being "categorical" - they +# still need to be run through the datetime-disambiguation helpers, because a +# String/Categorical column may actually hold dates. +_CATEGORICAL_CANDIDATE_SELECTOR = ( + nw.selectors.categorical() | nw.selectors.enum() | nw.selectors.string() +) + def find_numerical_variables( - X: pd.DataFrame, + X: IntoDataFrame, return_empty: bool = False, ) -> List[Union[str, int]]: """ @@ -25,8 +37,9 @@ def find_numerical_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset. + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. return_empty : bool, default=False Whether to return an empty list when no numerical variables are found. @@ -56,7 +69,12 @@ def find_numerical_variables( >>> var_ ['var_num'] """ - variables = list(X.select_dtypes(include="number").columns) + if nwd.is_pandas_dataframe(X): + variables = list(X.select_dtypes(include="number").columns) + else: + nw_X = nw.from_native(X, eager_only=True) + variables = list(nw_X.select(nw.selectors.numeric()).columns) + if len(variables) == 0: if return_empty is False: raise TypeError( @@ -73,7 +91,7 @@ def find_numerical_variables( def find_categorical_variables( - X: pd.DataFrame, + X: IntoDataFrame, return_empty: bool = False, ) -> List[Union[str, int]]: """ @@ -85,8 +103,9 @@ def find_categorical_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset. + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. return_empty : bool, default=False Whether to return an empty list when no categorical variables are found. @@ -116,11 +135,21 @@ def find_categorical_variables( >>> var_ ['var_cat'] """ - variables = [ - column - for column in X.select_dtypes(include=["O", "category", "string"]).columns - if _is_categorical_and_is_not_datetime(X[column]) - ] + if nwd.is_pandas_dataframe(X): + variables = [ + column + for column in X.select_dtypes(include=["O", "category", "string"]).columns + if _is_categorical_and_is_not_datetime(X[column]) + ] + else: + nw_X = nw.from_native(X, eager_only=True) + candidates = nw_X.select(_CATEGORICAL_CANDIDATE_SELECTOR).columns + variables = [ + column + for column in candidates + if _nw_is_categorical_and_is_not_datetime(nw_X[column]) + ] + if len(variables) == 0: if return_empty is False: raise TypeError( @@ -138,7 +167,7 @@ def find_categorical_variables( def find_datetime_variables( - X: pd.DataFrame, + X: IntoDataFrame, return_empty: bool = False, ) -> List[Union[str, int]]: """ @@ -152,8 +181,9 @@ def find_datetime_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset. + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. return_empty : bool, default=False Whether to return an empty list when no datetime variables are found. @@ -170,6 +200,16 @@ def find_datetime_variables( variables: List The names of the datetime variables. + Notes + ----- + For pandas dataframes, string columns are parsed with pandas' flexible, + dateutil-backed date guessing (the same as `pandas.to_datetime`), so formats + like "01-Jan-2010" or "10/11/12" are recognised. For polars (and other + non-pandas dataframes), only ISO-8601 strings and native `Date`/`Datetime` + columns are recognised automatically - polars has no equivalent flexible + guesser. Pass `variables` explicitly to a transformer if your polars date + strings use a different format. + Examples -------- >>> import pandas as pd @@ -183,12 +223,24 @@ def find_datetime_variables( >>> var_date ['var_date'] """ + if nwd.is_pandas_dataframe(X): + variables = [ + column + for column in X.select_dtypes(exclude="number").columns + if is_datetime(X[column]) or _is_categorical_and_is_datetime(X[column]) + ] + else: + nw_X = nw.from_native(X, eager_only=True) + non_numeric = [ + column for column in nw_X.columns if not nw_X.schema[column].is_numeric() + ] + variables = [ + column + for column in non_numeric + if _nw_is_date_or_datetime(nw_X.schema[column]) + or _nw_is_categorical_and_is_datetime(nw_X[column]) + ] - variables = [ - column - for column in X.select_dtypes(exclude="number").columns - if is_datetime(X[column]) or _is_categorical_and_is_datetime(X[column]) - ] if len(variables) == 0: if return_empty is False: raise TypeError( @@ -205,7 +257,7 @@ def find_datetime_variables( def find_all_variables( - X: pd.DataFrame, + X: IntoDataFrame, exclude_datetime: bool = False, return_empty: bool = False, ) -> List[Union[str, int]]: @@ -217,8 +269,9 @@ def find_all_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset. + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. exclude_datetime: bool, default=False Whether to exclude datetime variables. @@ -251,15 +304,32 @@ def find_all_variables( >>> vars_all ['var_num', 'var_cat', 'var_date'] """ - if exclude_datetime is True: - variables = X.select_dtypes(exclude=DATETIME_TYPES).columns.to_list() - variables = [ - var - for var in variables - if is_numeric(X[var]) or not _is_categorical_and_is_datetime(X[var]) - ] + if nwd.is_pandas_dataframe(X): + if exclude_datetime is True: + variables = X.select_dtypes(exclude=DATETIME_TYPES).columns.to_list() + variables = [ + var + for var in variables + if is_numeric(X[var]) or not _is_categorical_and_is_datetime(X[var]) + ] + else: + variables = X.columns.to_list() else: - variables = X.columns.to_list() + nw_X = nw.from_native(X, eager_only=True) + if exclude_datetime is True: + variables = [ + var + for var in nw_X.columns + if not _nw_is_date_or_datetime(nw_X.schema[var]) + ] + variables = [ + var + for var in variables + if nw_X.schema[var].is_numeric() + or not _nw_is_categorical_and_is_datetime(nw_X[var]) + ] + else: + variables = list(nw_X.columns) if len(variables) == 0: if return_empty is False: @@ -276,7 +346,7 @@ def find_all_variables( def find_categorical_and_numerical_variables( - X: pd.DataFrame, + X: IntoDataFrame, variables: Union[None, int, str, List[Union[str, int]]] = None, return_empty: bool = False, ) -> Tuple[List[Union[str, int]], List[Union[str, int]]]: @@ -290,8 +360,9 @@ def find_categorical_and_numerical_variables( Parameters ---------- - X : pandas dataframe of shape = [n_samples, n_features] - The dataset. + X : dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. variables : list, default=None If `None`, the function finds all categorical and numerical variables in X. @@ -329,15 +400,31 @@ def find_categorical_and_numerical_variables( >>> var_cat, var_num (['var_cat'], ['var_num']) """ + is_pandas = nwd.is_pandas_dataframe(X) + # cheap to build even when unused in the pandas branch below - avoids an + # Optional type that mypy can't narrow across the branches + nw_X = nw.from_native(X, eager_only=True) # If the user passes just 1 variable outside a list. if isinstance(variables, (str, int)): - if X[variables].dtype.name == "category" or _is_categorical_and_is_not_datetime( - X[variables] - ): + if is_pandas: + is_cat = X[ + variables + ].dtype.name == "category" or _is_categorical_and_is_not_datetime( + X[variables] + ) + is_num = is_numeric(X[variables]) + else: + s = nw_X[variables] + is_cat = isinstance( + s.dtype, (nw.Categorical, nw.Enum) + ) or _nw_is_categorical_and_is_not_datetime(s) + is_num = s.dtype.is_numeric() + + if is_cat: variables_cat = [variables] variables_num = [] - elif is_numeric(X[variables]): + elif is_num: variables_num = [variables] variables_cat = [] else: @@ -358,12 +445,23 @@ def find_categorical_and_numerical_variables( # If user leaves default None parameter. elif variables is None: - variables_cat = [ - column - for column in X.select_dtypes(include=["O", "category", "string"]).columns - if _is_categorical_and_is_not_datetime(X[column]) - ] - variables_num = list(X.select_dtypes(include="number").columns) + if is_pandas: + variables_cat = [ + column + for column in X.select_dtypes( + include=["O", "category", "string"] + ).columns + if _is_categorical_and_is_not_datetime(X[column]) + ] + variables_num = list(X.select_dtypes(include="number").columns) + else: + candidates = nw_X.select(_CATEGORICAL_CANDIDATE_SELECTOR).columns + variables_cat = [ + column + for column in candidates + if _nw_is_categorical_and_is_not_datetime(nw_X[column]) + ] + variables_num = list(nw_X.select(nw.selectors.numeric()).columns) if len(variables_num) == 0 and len(variables_cat) == 0: if return_empty is False: @@ -398,7 +496,7 @@ def find_categorical_and_numerical_variables( variables_cat = [] variables_num = [] - else: + elif is_pandas: # find categorical variables variables_cat = [ column @@ -410,4 +508,14 @@ def find_categorical_and_numerical_variables( # find numerical variables variables_num = list(X[variables].select_dtypes(include="number").columns) + else: + sub_X = nw_X.select(variables) + candidates = sub_X.select(_CATEGORICAL_CANDIDATE_SELECTOR).columns + variables_cat = [ + column + for column in candidates + if _nw_is_categorical_and_is_not_datetime(sub_X[column]) + ] + variables_num = list(sub_X.select(nw.selectors.numeric()).columns) + return variables_cat, variables_num diff --git a/feature_engine/variable_handling/retain_variables.py b/feature_engine/variable_handling/retain_variables.py index 2a161d066..3413b7c88 100644 --- a/feature_engine/variable_handling/retain_variables.py +++ b/feature_engine/variable_handling/retain_variables.py @@ -2,18 +2,22 @@ from typing import List, Union +import narwhals as nw +from narwhals.typing import IntoDataFrame + Variables = Union[int, str, List[Union[str, int]]] -def retain_variables_if_in_df(X, variables): +def retain_variables_if_in_df(X: IntoDataFrame, variables): """Returns the subset of variables in the list that are present in the dataframe. More details in the :ref:`User Guide `. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features] - The dataset. + X: dataframe of shape = [n_samples, n_features] + The dataset. Can be a pandas, polars, or any other dataframe supported by + narwhals. variables: string, int or list of strings or int. The names of the variables to check. @@ -39,7 +43,8 @@ def retain_variables_if_in_df(X, variables): if isinstance(variables, (str, int)): variables = [variables] - variables_in_df = [var for var in variables if var in X.columns] + columns = nw.from_native(X, eager_only=True).columns + variables_in_df = [var for var in variables if var in columns] # Raise an error if no column is left to work with. if len(variables_in_df) == 0: diff --git a/tests/test_variable_handling/conftest.py b/tests/test_variable_handling/conftest.py index 841656da2..6d029a01a 100644 --- a/tests/test_variable_handling/conftest.py +++ b/tests/test_variable_handling/conftest.py @@ -1,7 +1,48 @@ +from datetime import datetime, timezone + import pandas as pd +import polars as pl import pytest +def cast_categorical(df, columns): + """Cast `columns` to the backend's categorical dtype, whichever backend `df` + (pandas or polars) happens to be. Used to build matched pandas/polars data + for tests parametrized over both libraries. + """ + if isinstance(df, pd.DataFrame): + df = df.copy() + df[columns] = df[columns].astype("category") + return df + return df.with_columns([pl.col(c).cast(pl.Categorical) for c in columns]) + + +# Data shared between the pandas and polars variants of a test. Kept as plain +# dicts/lists (not fixtures) so a test can build both `make_df(BASIC_DATA)` and +# `make_df(BASIC_DATA)` for a different backend without needing to convert +# between frame types. +BASIC_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], +} + +# Datetime formats that both pandas and polars auto-detect: native +# Datetime/Date columns and ISO-8601 strings. Formats that only pandas' +# flexible, dateutil-backed guessing can parse (e.g. "01-Jan-2010", +# "10/11/12", bare time strings) are exercised separately, in pandas-only +# tests, against the `df_datetime` fixture below. +DATETIME_DATA = { + **BASIC_DATA, + "date_range": [datetime(2020, 2, 24, 0, i) for i in range(4)], + "date_obj0": ["2020-02-24", "2020-02-25", "2020-02-26", "2020-02-27"], + "date_range_tz": [ + datetime(2020, 2, 24, 0, i, tzinfo=timezone.utc) for i in range(4) + ], +} + + @pytest.fixture def df(): df = pd.DataFrame( diff --git a/tests/test_variable_handling/test_check_variables.py b/tests/test_variable_handling/test_check_variables.py index 8eba88cb0..5ea8d1485 100644 --- a/tests/test_variable_handling/test_check_variables.py +++ b/tests/test_variable_handling/test_check_variables.py @@ -1,4 +1,5 @@ import pandas as pd +import polars as pl import pytest from feature_engine.variable_handling import ( @@ -7,18 +8,31 @@ check_datetime_variables, check_numerical_variables, ) +from tests.test_variable_handling.conftest import ( + BASIC_DATA, + DATETIME_DATA, + cast_categorical, +) -def test_check_numerical_variables_returns_numerical_variables(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_check_numerical_variables_returns_numerical_variables(make_df): + df = make_df(BASIC_DATA) assert check_numerical_variables(df, ["Age", "Marks"]) == ["Age", "Marks"] assert check_numerical_variables(df, ["Age"]) == ["Age"] assert check_numerical_variables(df, "Age") == ["Age"] + + +def test_check_numerical_variables_returns_numerical_variables_int_names(df_int): + # polars requires string column names, so int-named columns are pandas-only assert check_numerical_variables(df_int, [3, 4]) == [3, 4] assert check_numerical_variables(df_int, [3]) == [3] assert check_numerical_variables(df_int, 4) == [4] -def test_check_numerical_variables_raises_errors_when_not_numerical(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_check_numerical_variables_raises_errors_when_not_numerical(make_df): + df = make_df(BASIC_DATA) msg = ( "Some of the variables are not numerical. Please cast them as " "numerical before using this transformer." @@ -32,15 +46,21 @@ def test_check_numerical_variables_raises_errors_when_not_numerical(df, df_int): assert str(record.value) == msg with pytest.raises(TypeError) as record: - assert check_numerical_variables(df_int, 1) + assert check_numerical_variables(df, ["Name", "Marks"]) assert str(record.value) == msg + +def test_check_numerical_variables_raises_errors_int_names(df_int): + msg = ( + "Some of the variables are not numerical. Please cast them as " + "numerical before using this transformer." + ) with pytest.raises(TypeError) as record: - assert check_numerical_variables(df_int, [1]) + assert check_numerical_variables(df_int, 1) assert str(record.value) == msg with pytest.raises(TypeError) as record: - assert check_numerical_variables(df, ["Name", "Marks"]) + assert check_numerical_variables(df_int, [1]) assert str(record.value) == msg with pytest.raises(TypeError) as record: @@ -48,22 +68,32 @@ def test_check_numerical_variables_raises_errors_when_not_numerical(df, df_int): assert str(record.value) == msg -def test_check_categorical_variables_returns_categorical_variables(df, df_int): - assert check_categorical_variables(df, ["Name", "date_obj0"]) == [ - "Name", - "date_obj0", - ] +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_check_categorical_variables_returns_categorical_variables(make_df): + df = make_df(BASIC_DATA) + assert check_categorical_variables(df, ["Name", "City"]) == ["Name", "City"] assert check_categorical_variables(df, ["Name"]) == ["Name"] - assert check_categorical_variables(df, "date_obj0") == ["date_obj0"] + assert check_categorical_variables(df, "Name") == ["Name"] + + +def test_check_categorical_variables_numeric_categories_pandas_only(): + # pandas allows category dtype with numeric categories (pd.Categorical can + # wrap numbers); polars categoricals are always string-backed, so casting a + # numeric column to Categorical isn't a realistic polars scenario. + df = pd.DataFrame(BASIC_DATA) + df = cast_categorical(df, ["Age", "Marks"]) + assert check_categorical_variables(df, ["Age", "Marks"]) == ["Age", "Marks"] + + +def test_check_categorical_variables_returns_categorical_variables_int_names(df_int): assert check_categorical_variables(df_int, [1, 2]) == [1, 2] assert check_categorical_variables(df_int, [2]) == [2] assert check_categorical_variables(df_int, 2) == [2] - df[["Age", "Marks"]] = df[["Age", "Marks"]].astype(pd.CategoricalDtype) - assert check_categorical_variables(df, ["Age", "Marks"]) == ["Age", "Marks"] - -def test_check_categorical_variables_raises_errors_when_not_categorical(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_check_categorical_variables_raises_errors_when_not_categorical(make_df): + df = make_df(BASIC_DATA) msg = ( "Some of the variables are not categorical. Please cast them as " "object or categorical before using this transformer." @@ -77,15 +107,21 @@ def test_check_categorical_variables_raises_errors_when_not_categorical(df, df_i assert str(record.value) == msg with pytest.raises(TypeError) as record: - assert check_categorical_variables(df_int, 3) + assert check_categorical_variables(df, ["Name", "Marks"]) assert str(record.value) == msg + +def test_check_categorical_variables_raises_errors_int_names(df_int): + msg = ( + "Some of the variables are not categorical. Please cast them as " + "object or categorical before using this transformer." + ) with pytest.raises(TypeError) as record: - assert check_categorical_variables(df_int, [3]) + assert check_categorical_variables(df_int, 3) assert str(record.value) == msg with pytest.raises(TypeError) as record: - assert check_categorical_variables(df, ["Name", "Marks"]) + assert check_categorical_variables(df_int, [3]) assert str(record.value) == msg with pytest.raises(TypeError) as record: @@ -93,17 +129,32 @@ def test_check_categorical_variables_raises_errors_when_not_categorical(df, df_i assert str(record.value) == msg -def test_check_datetime_variables_returns_datetime_variables(df_datetime): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_check_datetime_variables_returns_datetime_variables(make_df): + df = make_df(DATETIME_DATA) var_dt = ["date_range"] var_dt_str = "date_range" + vars_dt = ["date_range", "date_obj0", "date_range_tz"] + tz_time = "date_range_tz" + + assert check_datetime_variables(df, var_dt_str) == [var_dt_str] + assert check_datetime_variables(df, var_dt) == var_dt + assert check_datetime_variables(df, vars_dt) == vars_dt + assert check_datetime_variables(df, tz_time) == [tz_time] + + # only the string column can be cast to categorical - native Datetime + # columns can't be cast to Categorical in polars + df = cast_categorical(df, ["date_obj0"]) + assert check_datetime_variables(df, "date_obj0") == ["date_obj0"] + + +def test_check_datetime_variables_returns_pandas_only_string_formats(df_datetime): + # "01-Jan-2010"-style and "10/11/12"-style strings are only ever recognised + # through pandas' flexible, dateutil-backed guessing - see the note in + # check_datetime_variables' docstring. vars_convertible_to_dt = ["date_range", "date_obj1", "date_obj2", "time_obj"] var_convertible_to_dt = "date_obj1" - tz_time = "time_objTZ" - tz_time_obj = "date_range_tz" - # when variables are specified - assert check_datetime_variables(df_datetime, var_dt_str) == [var_dt_str] - assert check_datetime_variables(df_datetime, var_dt) == var_dt assert check_datetime_variables(df_datetime, var_convertible_to_dt) == [ var_convertible_to_dt ] @@ -111,8 +162,6 @@ def test_check_datetime_variables_returns_datetime_variables(df_datetime): check_datetime_variables(df_datetime, vars_convertible_to_dt) == vars_convertible_to_dt ) - assert check_datetime_variables(df_datetime, tz_time) == [tz_time] - assert check_datetime_variables(df_datetime, tz_time_obj) == [tz_time_obj] df_datetime[vars_convertible_to_dt] = df_datetime[vars_convertible_to_dt].astype( pd.CategoricalDtype @@ -123,54 +172,53 @@ def test_check_datetime_variables_returns_datetime_variables(df_datetime): ) -def test_check_datetime_variables_raises_errors_when_not_datetime(df_datetime): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_check_datetime_variables_raises_errors_when_not_datetime(make_df): + df = make_df(DATETIME_DATA) msg = "Some of the variables are not or cannot be parsed as datetime." with pytest.raises(TypeError) as record: - assert check_datetime_variables(df_datetime, variables="Age") + assert check_datetime_variables(df, variables="Age") assert str(record.value) == msg with pytest.raises(TypeError) as record: - assert check_datetime_variables(df_datetime, variables=["Age", "Name"]) + assert check_datetime_variables(df, variables=["Age", "Name"]) assert str(record.value) == msg with pytest.raises(TypeError): - assert check_datetime_variables(df_datetime, variables=["date_range", "Age"]) + assert check_datetime_variables(df, variables=["date_range", "Age"]) assert str(record.value) == msg +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "input_vars", [ - ["Name", "City", "Age", "Marks", "dob"], - [ - "Name", - "City", - "Age", - "Marks", - ], + ["Name", "City", "Age", "Marks"], + ["Name", "City", "Age"], "Name", ["Age"], ], ) -def test_check_all_variables_returns_all_variables(df_vartypes, input_vars): +def test_check_all_variables_returns_all_variables(make_df, input_vars): + df = make_df(BASIC_DATA) if isinstance(input_vars, list): - assert check_all_variables(df_vartypes, input_vars) == input_vars + assert check_all_variables(df, input_vars) == input_vars else: - assert check_all_variables(df_vartypes, input_vars) == [input_vars] + assert check_all_variables(df, input_vars) == [input_vars] +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) @pytest.mark.parametrize( "input_vars", [["Name", "City", "Absent"], "Absent", ["Absent"]] ) -def test_check_all_variables_raises_errors_when_not_in_dataframe( - df_vartypes, input_vars -): +def test_check_all_variables_raises_errors_when_not_in_dataframe(make_df, input_vars): + df = make_df(BASIC_DATA) msg_ls = "'Some of the variables are not in the dataframe.'" msg_single = "'The variable Absent is not in the dataframe.'" with pytest.raises(KeyError) as record: - assert check_all_variables(df_vartypes, input_vars) + assert check_all_variables(df, input_vars) if isinstance(input_vars, list): assert str(record.value) == msg_ls else: diff --git a/tests/test_variable_handling/test_fe_type_checks.py b/tests/test_variable_handling/test_fe_type_checks.py index de4bc2d38..50a4f6cba 100644 --- a/tests/test_variable_handling/test_fe_type_checks.py +++ b/tests/test_variable_handling/test_fe_type_checks.py @@ -1,4 +1,8 @@ +from datetime import date + +import narwhals as nw import pandas as pd +import polars as pl from feature_engine.variable_handling._variable_type_checks import ( _is_categorical_and_is_datetime, @@ -6,9 +10,21 @@ _is_categories_num, _is_convertible_to_dt, _is_convertible_to_num, + _nw_is_categorical_and_is_datetime, + _nw_is_categorical_and_is_not_datetime, + _nw_is_convertible_to_dt, + _nw_is_convertible_to_num, + _nw_is_date_or_datetime, ) +def nw_series(values, dtype=None): + s = pl.Series("x", values) + if dtype is not None: + s = s.cast(dtype) + return nw.from_native(s, series_only=True) + + def test_is_categories_num(df): assert _is_categories_num(df["Name"]) is False @@ -91,3 +107,94 @@ def test_is_categorical_and_is_not_datetime(df): # Categorical should be True (it hits the 'if' branch) s_cat = pd.Series(["a", "b"], dtype="category") assert _is_categorical_and_is_not_datetime(s_cat) is True + + +# --------------------------------------------------------------------------- +# narwhals (polars) equivalents +# --------------------------------------------------------------------------- + + +def test_nw_is_date_or_datetime(): + assert _nw_is_date_or_datetime(nw_series([date(2020, 1, 1)]).dtype) is True + assert ( + _nw_is_date_or_datetime(nw_series(["2020-01-01"]).str.to_datetime().dtype) + is True + ) + assert _nw_is_date_or_datetime(nw_series(["a", "b"]).dtype) is False + assert _nw_is_date_or_datetime(nw_series([1, 2, 3]).dtype) is False + + +def test_nw_is_convertible_to_num(): + assert _nw_is_convertible_to_num(nw_series(["20", "21", "19"])) is True + assert _nw_is_convertible_to_num(nw_series(["a", "b"])) is False + assert ( + _nw_is_convertible_to_num(nw_series(["20", "21"], dtype=pl.Categorical)) + is True + ) + + +def test_nw_is_convertible_to_dt(): + assert _nw_is_convertible_to_dt(nw_series(["2020-01-01", "2020-01-02"])) is True + assert _nw_is_convertible_to_dt(nw_series(["a", "b"])) is False + assert _nw_is_convertible_to_dt(nw_series(["20", "21"])) is False + + # polars has no dateutil-style guesser, unlike pandas, so non-ISO date + # strings are not recognised + assert _nw_is_convertible_to_dt(nw_series(["01-Jan-2010"])) is False + + +def test_nw_is_categorical_and_is_datetime(): + assert ( + _nw_is_categorical_and_is_datetime( + nw_series(["2020-01-01", "2020-01-02"], dtype=pl.Categorical) + ) + is True + ) + assert ( + _nw_is_categorical_and_is_datetime(nw_series(["a", "b"], dtype=pl.Categorical)) + is False + ) + assert _nw_is_categorical_and_is_datetime(nw_series(["2020-01-01"])) is True + assert _nw_is_categorical_and_is_datetime(nw_series(["20", "21"])) is False + assert _nw_is_categorical_and_is_datetime(nw_series(["a", "b"])) is False + + # an explicit Enum is always treated as categorical, never as datetime + enum_dtype = pl.Enum(["2020-01-01", "2020-01-02"]) + assert ( + _nw_is_categorical_and_is_datetime( + nw_series(["2020-01-01", "2020-01-02"], dtype=enum_dtype) + ) + is False + ) + + # numeric should be False + assert _nw_is_categorical_and_is_datetime(nw_series([1, 2, 3])) is False + + +def test_nw_is_categorical_and_is_not_datetime(): + assert ( + _nw_is_categorical_and_is_not_datetime( + nw_series(["2020-01-01", "2020-01-02"], dtype=pl.Categorical) + ) + is False + ) + assert ( + _nw_is_categorical_and_is_not_datetime( + nw_series(["a", "b"], dtype=pl.Categorical) + ) + is True + ) + assert _nw_is_categorical_and_is_not_datetime(nw_series(["2020-01-01"])) is False + assert _nw_is_categorical_and_is_not_datetime(nw_series(["20", "21"])) is True + assert _nw_is_categorical_and_is_not_datetime(nw_series(["a", "b"])) is True + + # an explicit Enum is always treated as categorical + assert ( + _nw_is_categorical_and_is_not_datetime( + nw_series(["a", "b"], dtype=pl.Enum(["a", "b"])) + ) + is True + ) + + # numeric should be False + assert _nw_is_categorical_and_is_not_datetime(nw_series([1, 2, 3])) is False diff --git a/tests/test_variable_handling/test_find_variables.py b/tests/test_variable_handling/test_find_variables.py index 6ae29384d..8654724b5 100644 --- a/tests/test_variable_handling/test_find_variables.py +++ b/tests/test_variable_handling/test_find_variables.py @@ -1,4 +1,5 @@ import pandas as pd +import polars as pl import pytest from feature_engine.variable_handling import ( @@ -8,89 +9,101 @@ find_datetime_variables, find_numerical_variables, ) +from tests.test_variable_handling.conftest import ( + BASIC_DATA, + DATETIME_DATA, + cast_categorical, +) # --- find_numerical_variables --- # -def test_numerical_variables_finds_variables(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numerical_variables_finds_variables(make_df): + df = make_df(BASIC_DATA) assert find_numerical_variables(df) == ["Age", "Marks"] + + +def test_numerical_variables_finds_variables_with_int_column_names(df_int): + # polars requires string column names, so int-named columns are pandas-only assert find_numerical_variables(df_int) == [3, 4] -def test_numerical_variables_raises_error(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numerical_variables_raises_error(make_df): + df = make_df(BASIC_DATA) msg = "No numerical variables found in this dataframe." with pytest.raises(TypeError, match=msg): - find_numerical_variables(df.drop(["Age", "Marks"], axis=1)) - - with pytest.raises(TypeError, match=msg): - find_numerical_variables(df_int.drop([3, 4], axis=1)) + find_numerical_variables(df[["Name", "City"]]) -def test_numerical_variables_raises_warning(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numerical_variables_raises_warning(make_df): + df = make_df(BASIC_DATA) msg = "No numerical variables found in this dataframe." - - # Test with a regular DataFrame with pytest.warns(UserWarning, match=msg): - find_numerical_variables(df.drop(["Age", "Marks"], axis=1), return_empty=True) + find_numerical_variables(df[["Name", "City"]], return_empty=True) - # Test with integer-only DataFrame - with pytest.warns(UserWarning, match=msg): - find_numerical_variables(df_int.drop([3, 4], axis=1), return_empty=True) - -def test_numerical_variables_returns_empty_list(df, df_int): - assert ( - find_numerical_variables(df.drop(["Age", "Marks"], axis=1), return_empty=True) - == [] - ) - assert ( - find_numerical_variables(df_int.drop([3, 4], axis=1), return_empty=True) == [] - ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numerical_variables_returns_empty_list(make_df): + df = make_df(BASIC_DATA) + assert find_numerical_variables(df[["Name", "City"]], return_empty=True) == [] # --- find_categorical_variables --- # -def test_categorical_variables_finds_variables(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_categorical_variables_finds_variables(make_df): + df = make_df(BASIC_DATA) assert find_categorical_variables(df) == ["Name", "City"] + + +def test_categorical_variables_finds_variables_with_int_column_names(df_int): assert find_categorical_variables(df_int) == [1, 2] -def test_categorical_variables_raises_error(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_categorical_variables_raises_error(make_df): + df = make_df(BASIC_DATA) msg = "No categorical variables found in this dataframe." with pytest.raises(TypeError, match=msg): - find_categorical_variables(df.drop(["Name", "City"], axis=1)) - - with pytest.raises(TypeError, match=msg): - find_categorical_variables(df_int.drop([1, 2], axis=1)) + find_categorical_variables(df[["Age", "Marks"]]) -def test_categorical_variables_raises_warning(df, df_int): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_categorical_variables_raises_warning(make_df): + df = make_df(BASIC_DATA) msg = "No categorical variables found in this dataframe." - - # Test with a regular DataFrame with pytest.warns(UserWarning, match=msg): - find_categorical_variables(df.drop(["Name", "City"], axis=1), return_empty=True) + find_categorical_variables(df[["Age", "Marks"]], return_empty=True) - # Test with integer-only DataFrame - with pytest.warns(UserWarning, match=msg): - find_categorical_variables(df_int.drop([1, 2], axis=1), return_empty=True) - -def test_categorical_variables_returns_empty_list(df, df_int): - assert ( - find_categorical_variables(df.drop(["Name", "City"], axis=1), return_empty=True) - == [] - ) - assert ( - find_categorical_variables(df_int.drop([1, 2], axis=1), return_empty=True) == [] - ) +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_categorical_variables_returns_empty_list(make_df): + df = make_df(BASIC_DATA) + assert find_categorical_variables(df[["Age", "Marks"]], return_empty=True) == [] # --- find_datetime_variables --- # -def test_datetime_variables_finds_variables(df_datetime): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_datetime_variables_finds_variables(make_df): + df = make_df(DATETIME_DATA) + vars_dt = ["date_range", "date_obj0", "date_range_tz"] + assert find_datetime_variables(df) == vars_dt + + assert find_datetime_variables( + df[["date_obj0", "date_range", "date_range_tz"]], + ) == ["date_obj0", "date_range", "date_range_tz"] + + +def test_datetime_variables_finds_pandas_only_string_formats(df_datetime): + # "01-Jan-2010"-style, "10/11/12"-style and bare-time strings are only ever + # recognised through pandas' flexible, dateutil-backed guessing - polars has + # no equivalent, see the note in find_datetime_variables' docstring. vars_dt = [ "date_range", "date_obj0", @@ -100,206 +113,193 @@ def test_datetime_variables_finds_variables(df_datetime): "time_obj", "time_objTZ", ] - assert find_datetime_variables(df_datetime) == vars_dt - assert find_datetime_variables( - df_datetime[vars_dt].reindex(columns=["date_obj1", "date_range", "date_obj2"]), - ) == ["date_obj1", "date_range", "date_obj2"] - -def test_datetime_variables_raises_error(df_datetime): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_datetime_variables_raises_error(make_df): + df = make_df(DATETIME_DATA) msg = "No datetime variables found in this dataframe." - vars_nondt = ["Marks", "Age", "Name"] - with pytest.raises(TypeError, match=msg): - find_datetime_variables(df_datetime.loc[:, vars_nondt]) + find_datetime_variables(df[vars_nondt]) -def test_datetime_variables_raises_warning(df_datetime): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_datetime_variables_raises_warning(make_df): + df = make_df(DATETIME_DATA) msg = "No datetime variables found in this dataframe." vars_nondt = ["Marks", "Age", "Name"] with pytest.warns(UserWarning, match=msg): - find_datetime_variables(df_datetime.loc[:, vars_nondt], return_empty=True) + find_datetime_variables(df[vars_nondt], return_empty=True) -def test_datetime_variables_returns_empty_list(df_datetime): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_datetime_variables_returns_empty_list(make_df): + df = make_df(DATETIME_DATA) vars_nondt = ["Marks", "Age", "Name"] - assert ( - find_datetime_variables(df_datetime.loc[:, vars_nondt], return_empty=True) == [] - ) + assert find_datetime_variables(df[vars_nondt], return_empty=True) == [] # --- find_all_variables --- # -def test_find_all_variables(df): - all_vars = [ - "Name", - "City", - "Age", - "Marks", - "date_range", - "date_obj0", - "date_range_tz", - ] - assert find_all_variables(df, exclude_datetime=False) == all_vars +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_find_all_variables(make_df): + df = make_df(BASIC_DATA) + assert find_all_variables(df, exclude_datetime=False) == list(BASIC_DATA.keys()) -def test_find_all_variables_excludes_dt(df): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_find_all_variables_excludes_dt(make_df): + df = make_df(DATETIME_DATA) all_vars_no_dt = ["Name", "City", "Age", "Marks"] assert find_all_variables(df, exclude_datetime=True) == all_vars_no_dt -def test_find_all_variables_raises_error(df): - dt_vars = [ - "date_range", - "date_obj0", - "date_range_tz", - ] - df = df[dt_vars] +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_find_all_variables_raises_error(make_df): + dt_vars = ["date_range", "date_obj0", "date_range_tz"] + df = make_df(DATETIME_DATA)[dt_vars] msg = "No variables found in this dataframe" with pytest.raises(TypeError, match=msg): find_all_variables(df, exclude_datetime=True) -def test_find_all_variables_raises_warning(df): - dt_vars = [ - "date_range", - "date_obj0", - "date_range_tz", - ] - df = df[dt_vars] +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_find_all_variables_raises_warning(make_df): + dt_vars = ["date_range", "date_obj0", "date_range_tz"] + df = make_df(DATETIME_DATA)[dt_vars] msg = "No variables found in this dataframe" with pytest.warns(UserWarning, match=msg): find_all_variables(df, exclude_datetime=True, return_empty=True) -def test_find_all_variables_returns_empty(df): - dt_vars = [ - "date_range", - "date_obj0", - "date_range_tz", - ] - df = df[dt_vars] +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_find_all_variables_returns_empty(make_df): + dt_vars = ["date_range", "date_obj0", "date_range_tz"] + df = make_df(DATETIME_DATA)[dt_vars] assert find_all_variables(df, exclude_datetime=True, return_empty=True) == [] # --- find_categorical_and_numerical_variables --- # -def test_numcat_user_passes_varlist(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numcat_user_passes_varlist(make_df): + df = make_df(BASIC_DATA) + # Case 1: user passes 1 variable that is categorical - assert find_categorical_and_numerical_variables(df_vartypes, ["Name"]) == ( - ["Name"], - [], - ) - assert find_categorical_and_numerical_variables(df_vartypes, "Name") == ( - ["Name"], - [], - ) + assert find_categorical_and_numerical_variables(df, ["Name"]) == (["Name"], []) + assert find_categorical_and_numerical_variables(df, "Name") == (["Name"], []) # Case 2: user passes 1 variable that is numerical - assert find_categorical_and_numerical_variables(df_vartypes, ["Age"]) == ( - [], - ["Age"], - ) - assert find_categorical_and_numerical_variables(df_vartypes, "Age") == ( - [], - ["Age"], - ) + assert find_categorical_and_numerical_variables(df, ["Age"]) == ([], ["Age"]) + assert find_categorical_and_numerical_variables(df, "Age") == ([], ["Age"]) # Case 3: user passes 1 categorical and 1 numerical variable - assert find_categorical_and_numerical_variables(df_vartypes, ["Age", "Name"]) == ( + assert find_categorical_and_numerical_variables(df, ["Age", "Name"]) == ( ["Name"], ["Age"], ) -def test_numcat_when_var_is_none(df_vartypes): - # Case 4: automatically identify variables - assert find_categorical_and_numerical_variables(df_vartypes, None) == ( +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numcat_when_var_is_none(make_df): + df = make_df(BASIC_DATA) + + assert find_categorical_and_numerical_variables(df, None) == ( ["Name", "City"], ["Age", "Marks"], ) - assert find_categorical_and_numerical_variables( - df_vartypes[["Name", "City"]], None - ) == (["Name", "City"], []) - assert find_categorical_and_numerical_variables( - df_vartypes[["Age", "Marks"]], None - ) == ([], ["Age", "Marks"]) - - -@pytest.fixture(scope="module") -def dfdt(): - X = pd.DataFrame() - X["date1"] = pd.date_range("2020-02-24", periods=1000, freq="min") - X["date2"] = pd.date_range("2021-09-29", periods=1000, freq="h") - X["date3"] = ["2020-02-24"] * 1000 - return X + assert find_categorical_and_numerical_variables(df[["Name", "City"]], None) == ( + ["Name", "City"], + [], + ) + assert find_categorical_and_numerical_variables(df[["Age", "Marks"]], None) == ( + [], + ["Age", "Marks"], + ) -def test_numcat_raises_no_var_error(dfdt): +@pytest.mark.parametrize( + "make_df, assert_error", [(pd.DataFrame, TypeError), (pl.DataFrame, TypeError)] +) +def test_numcat_raises_no_var_error(make_df, assert_error): # Case 5: error when no variable is numerical or categorical + df = make_df( + { + "date1": DATETIME_DATA["date_range"], + "date2": DATETIME_DATA["date_range_tz"], + } + ) msg = "There are no numerical or categorical variables" - with pytest.raises(TypeError, match=msg): - find_categorical_and_numerical_variables(dfdt, None) + with pytest.raises(assert_error, match=msg): + find_categorical_and_numerical_variables(df, None) msg = "The variable entered is neither numerical nor categorical." - with pytest.raises(TypeError, match=msg): - find_categorical_and_numerical_variables(dfdt, "date1") + with pytest.raises(assert_error, match=msg): + find_categorical_and_numerical_variables(df, "date1") -def test_numcat_raises_no_var_warn(dfdt): - # Case 6: warning when no variable is numerical or categorical +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numcat_raises_no_var_warn(make_df): + df = make_df( + { + "date1": DATETIME_DATA["date_range"], + "date2": DATETIME_DATA["date_range_tz"], + } + ) msg = "There are no numerical or categorical variables" with pytest.warns(UserWarning, match=msg): - find_categorical_and_numerical_variables( - dfdt, - None, - return_empty=True, - ) + find_categorical_and_numerical_variables(df, None, return_empty=True) msg = "The variable entered is neither numerical nor" with pytest.warns(UserWarning, match=msg): find_categorical_and_numerical_variables( - dfdt, variables="date1", return_empty=True + df, variables="date1", return_empty=True ) -def test_numcat_returns_empty_lists(dfdt): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numcat_returns_empty_lists(make_df): + df = make_df( + { + "date1": DATETIME_DATA["date_range"], + "date2": DATETIME_DATA["date_range_tz"], + } + ) assert find_categorical_and_numerical_variables( - dfdt, - None, - return_empty=True, + df, None, return_empty=True ) == ([], []) assert find_categorical_and_numerical_variables( - dfdt, - "date1", - return_empty=True, + df, "date1", return_empty=True ) == ([], []) -def test_numcat_on_user_empty_list(df_vartypes): - # Case 7: user passes empty list +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numcat_on_user_empty_list(make_df): + df = make_df(BASIC_DATA) + msg = "The list of variables provided is empty. If this was" with pytest.raises(ValueError, match=msg): - find_categorical_and_numerical_variables(df_vartypes, []) + find_categorical_and_numerical_variables(df, []) msg = "The list of variables provided is empty. Returning " with pytest.warns(UserWarning, match=msg): - find_categorical_and_numerical_variables(df_vartypes, [], return_empty=True) + find_categorical_and_numerical_variables(df, [], return_empty=True) - assert find_categorical_and_numerical_variables( - df_vartypes, [], return_empty=True - ) == ([], []) + assert find_categorical_and_numerical_variables(df, [], return_empty=True) == ( + [], + [], + ) def test_numcat_when_dt_as_object(df_vartypes): - # Case 8: datetime cast as object + # Case 8: datetime cast as object - pandas-only, `df_vartypes["dob"]` is a + # pandas datetime64 column relying on pandas' `.astype("O")`, which has no + # polars equivalent (polars has no generic object dtype to cast into). df = df_vartypes.copy() df["dob"] = df["dob"].astype("O") - # datetime variable is skipped when automatically finding variables, assert find_categorical_and_numerical_variables(df, None) == ( ["Name", "City"], ["Age", "Marks"], @@ -310,10 +310,11 @@ def test_numcat_when_dt_as_object(df_vartypes): ) -def test_numcat_vars_as_category(df_vartypes): +@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) +def test_numcat_vars_as_category(make_df): # Case 9: variables cast as category - df = df_vartypes.copy() - df["City"] = df["City"].astype("category") + df = make_df(BASIC_DATA) + df = cast_categorical(df, ["City"]) assert find_categorical_and_numerical_variables(df, None) == ( ["Name", "City"], ["Age", "Marks"], diff --git a/tests/test_variable_handling/test_remove_variables.py b/tests/test_variable_handling/test_remove_variables.py index 3984d2c45..5f08c96fd 100644 --- a/tests/test_variable_handling/test_remove_variables.py +++ b/tests/test_variable_handling/test_remove_variables.py @@ -1,23 +1,26 @@ import pandas as pd +import polars as pl import pytest from feature_engine.variable_handling.retain_variables import retain_variables_if_in_df + +def make_empty_df(is_pandas: bool, columns): + if is_pandas: + return pd.DataFrame(columns=columns) + return pl.DataFrame(schema=columns) + + test_dict = [ - ( - pd.DataFrame(columns=["A", "B", "C", "D", "E"]), - ["A", "C", "B", "G", "H"], - ["A", "C", "B"], - ["X", "Y"], - ), - (pd.DataFrame(columns=[1, 2, 3, 4, 5]), [1, 2, 4, 6], [1, 2, 4], [6, 7]), - (pd.DataFrame(columns=[1, 2, 3, 4, 5]), 1, [1], 7), - (pd.DataFrame(columns=["A", "B", "C", "D", "E"]), "C", ["C"], "G"), + (["A", "C", "B", "G", "H"], ["A", "C", "B"], ["X", "Y"]), + ("C", ["C"], "G"), ] -@pytest.mark.parametrize("df, variables, overlap, col_not_in_df", test_dict) -def test_retain_variables_if_in_df(df, variables, overlap, col_not_in_df): +@pytest.mark.parametrize("is_pandas", [True, False]) +@pytest.mark.parametrize("variables, overlap, col_not_in_df", test_dict) +def test_retain_variables_if_in_df(is_pandas, variables, overlap, col_not_in_df): + df = make_empty_df(is_pandas, ["A", "B", "C", "D", "E"]) msg = "None of the variables in the list are present in the dataframe." @@ -26,3 +29,21 @@ def test_retain_variables_if_in_df(df, variables, overlap, col_not_in_df): with pytest.raises(ValueError) as record: retain_variables_if_in_df(df, col_not_in_df) assert str(record.value) == msg + + +def test_retain_variables_if_in_df_int_column_names(): + # polars requires string column names, so int-named columns are pandas-only + df = pd.DataFrame(columns=[1, 2, 3, 4, 5]) + + msg = "None of the variables in the list are present in the dataframe." + + assert retain_variables_if_in_df(df, [1, 2, 4, 6]) == [1, 2, 4] + assert retain_variables_if_in_df(df, 1) == [1] + + with pytest.raises(ValueError) as record: + retain_variables_if_in_df(df, [6, 7]) + assert str(record.value) == msg + + with pytest.raises(ValueError) as record: + retain_variables_if_in_df(df, 7) + assert str(record.value) == msg From 8edc831151a132e89de6b82feb0e17c2c217c2c7 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Wed, 29 Jul 2026 17:26:32 +0200 Subject: [PATCH 2/4] creating own datetime parser --- feature_engine/datetime/datetime.py | 38 +++- .../_variable_type_checks.py | 151 ++++++--------- .../variable_handling/check_variables.py | 43 ++--- .../variable_handling/find_variables.py | 181 +++++++----------- pyproject.toml | 1 + .../test_check_variables.py | 7 +- .../test_fe_type_checks.py | 175 ++++++++--------- .../test_find_variables.py | 20 +- 8 files changed, 274 insertions(+), 342 deletions(-) diff --git a/feature_engine/datetime/datetime.py b/feature_engine/datetime/datetime.py index 106d50277..f6f485d08 100644 --- a/feature_engine/datetime/datetime.py +++ b/feature_engine/datetime/datetime.py @@ -5,6 +5,7 @@ import pandas as pd from pandas.api.types import is_datetime64_any_dtype as is_datetime from pandas.api.types import is_numeric_dtype as is_numeric +from pandas.api.types import is_object_dtype, is_string_dtype from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted @@ -38,13 +39,41 @@ FEATURES_SUFFIXES, FEATURES_SUPPORTED, ) -from feature_engine.variable_handling._variable_type_checks import ( - _is_categorical_and_is_datetime, -) from feature_engine.variable_handling.check_variables import check_datetime_variables from feature_engine.variable_handling.find_variables import find_datetime_variables +def _index_is_categorical_and_is_datetime(index: pd.Index) -> bool: + # This file is fully pandas-based (it casts with `pd.to_datetime` during + # `transform()`), and this check only ever runs against a pandas `Index` + # (narwhals has no `Index` concept), so it stays pandas-only rather than + # routing through the narwhals-based `_variable_type_checks` helpers. + is_object = is_object_dtype(index) or is_string_dtype(index) + + if isinstance(index.dtype, pd.CategoricalDtype): + categories_are_numeric = is_numeric(index.categories) + if categories_are_numeric: + return False + try: + return is_datetime(pd.to_datetime(index, utc=True)) + except Exception: + return False + + elif is_object: + try: + is_convertible_to_num = is_numeric(pd.to_numeric(index)) + except (ValueError, TypeError): + is_convertible_to_num = False + if is_convertible_to_num: + return False + try: + return is_datetime(pd.to_datetime(index, utc=True)) + except Exception: + return False + + return False + + @Substitution( return_empty=_return_empty_docstring, feature_names_in_=_feature_names_in_docstring, @@ -265,7 +294,8 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None): if not ( is_datetime(X.index) or ( - not is_numeric(X.index) and _is_categorical_and_is_datetime(X.index) + not is_numeric(X.index) + and _index_is_categorical_and_is_datetime(X.index) ) ): raise TypeError("The dataframe index is not datetime.") diff --git a/feature_engine/variable_handling/_variable_type_checks.py b/feature_engine/variable_handling/_variable_type_checks.py index 7e720b954..885e18583 100644 --- a/feature_engine/variable_handling/_variable_type_checks.py +++ b/feature_engine/variable_handling/_variable_type_checks.py @@ -1,115 +1,79 @@ +from datetime import date, datetime + import narwhals as nw -import pandas as pd -from pandas.api.types import is_object_dtype, is_string_dtype -from pandas.core.dtypes.common import is_datetime64_any_dtype as is_datetime -from pandas.core.dtypes.common import is_numeric_dtype as is_numeric +from dateutil.parser import parse as _dateutil_parse # --------------------------------------------------------------------------- -# pandas-only implementation. +# narwhals implementation, used for every backend (pandas, polars, etc.) # -# These functions rely on pandas' flexible, dateutil-backed `pd.to_datetime` -# string guessing and on pandas' `object` dtype (which, unlike any narwhals -# dtype, can hold arbitrary non-string Python objects). Neither has a -# polars/narwhals equivalent, so they are kept exactly as they were before the -# narwhals migration and are only ever called on pandas input. See the `_nw_*` -# functions below for the polars/narwhals-backend equivalents. +# Flexible date-string recognition (e.g. "01-Jan-2010", "10/11/12", not just +# ISO-8601) is implemented directly on top of `dateutil` - the same library +# pandas.to_datetime delegates to internally for this - so it works +# identically regardless of the underlying dataframe library. # --------------------------------------------------------------------------- -def is_object(s) -> bool: - return is_object_dtype(s) or is_string_dtype(s) - - -def _is_categorical_and_is_not_datetime(column: pd.Series) -> bool: - # check for datetime only if the type of the categories is not numeric - # because pd.to_datetime throws an error when it is an integer - if isinstance(column.dtype, pd.CategoricalDtype): - is_cat = _is_categories_num(column) or not _is_convertible_to_dt(column) - - # check for datetime only if object cannot be cast as numeric because - # if it could pd.to_datetime would convert it to datetime regardless - elif is_object(column): - is_cat = _is_convertible_to_num(column) or not _is_convertible_to_dt(column) - - else: - is_cat = False - - return is_cat +def _nw_is_date_or_datetime(dtype) -> bool: + # nw.selectors.datetime() only matches Datetime, not Date, so this needs + # its own explicit check. + return isinstance(dtype, (nw.Date, nw.Datetime)) -def _is_categories_num(column: pd.Series) -> bool: - return is_numeric(column.dtype.categories) +_DATE_PARSE_DEFAULT_1 = datetime(1, 1, 1, 1, 1, 1) +_DATE_PARSE_DEFAULT_2 = datetime(2, 2, 2, 2, 2, 2) +_DATETIME_FIELDS = ("year", "month", "day", "hour", "minute", "second") -def _is_convertible_to_dt(column: pd.Series) -> bool: +def _looks_like_date_string(value: str) -> bool: + # dateutil.parser.parse() fills in any date/time component that isn't + # present in the string from a `default` datetime, so a bare number like + # "20" "parses" successfully as day=20 - it would wrongly be treated as a + # date. Parsing twice, with two defaults that differ in every field, + # reveals which fields were actually present in the string: those are the + # fields that agree between the two parses. Requiring at least 2 fields to + # be corroborated this way rejects bare numbers while still accepting real + # dates (including non-ISO formats like "01-Jan-2010") and bare times + # (like "21:45:23"). try: - var = pd.to_datetime(column, utc=True) - return is_datetime(var) - except Exception: + first = _dateutil_parse(value, default=_DATE_PARSE_DEFAULT_1) + second = _dateutil_parse(value, default=_DATE_PARSE_DEFAULT_2) + except (ValueError, OverflowError, TypeError): return False - -def _is_convertible_to_num(column: pd.Series) -> bool: - try: - ser = pd.to_numeric(column) - except (ValueError, TypeError): - ser = column - return is_numeric(ser) - - -def _is_categorical_and_is_datetime(column: pd.Series) -> bool: - # check for datetime only if the type of the categories is not numeric - # because pd.to_datetime throws an error when it is an integer - if isinstance(column.dtype, pd.CategoricalDtype): - is_dt = not _is_categories_num(column) and _is_convertible_to_dt(column) - - # check for datetime only if object cannot be cast as numeric because - # if it could pd.to_datetime would convert it to datetime regardless - elif is_object(column): - is_dt = not _is_convertible_to_num(column) and _is_convertible_to_dt(column) - - else: - is_dt = False - - return is_dt - - -# --------------------------------------------------------------------------- -# narwhals implementation, used for every backend other than pandas (polars, -# in practice). -# -# narwhals has no lenient/"try" cast (no `strict=False`, unlike raw polars) -# and its `str.to_datetime()` requires ISO-8601 or an explicit `format=` - it -# cannot reproduce pandas' dateutil-based guessing. So a string column such as -# "01-Jan-2010" or "10/11/12" is not auto-detected as datetime for polars, -# even though it is for pandas. ISO-8601 strings and native Date/Datetime -# columns are detected correctly. Users can always pass `variables` explicitly -# to sidestep this. -# --------------------------------------------------------------------------- - - -def _nw_is_date_or_datetime(dtype) -> bool: - # nw.selectors.datetime() only matches Datetime, not Date, so this needs - # its own explicit check. - return isinstance(dtype, (nw.Date, nw.Datetime)) + corroborated = sum( + 1 for attr in _DATETIME_FIELDS if getattr(first, attr) == getattr(second, attr) + ) + return corroborated >= 2 def _nw_is_convertible_to_num(s: "nw.Series") -> bool: + values = s.drop_nulls().to_list() + if not values: + return False try: - s.cast(nw.String()).cast(nw.Float64()) - except Exception: + for value in values: + float(value) + except (ValueError, TypeError): return False return True def _nw_is_convertible_to_dt(s: "nw.Series") -> bool: - try: - s.cast(nw.String()).str.to_datetime() - except Exception: + values = s.drop_nulls().to_list() + if not values: return False + for value in values: + if isinstance(value, (date, datetime)): + continue + if not _looks_like_date_string(str(value)): + return False return True +def _nw_categories_are_numeric(s: "nw.Series") -> bool: + return s.cat.get_categories().dtype.is_numeric() + + def _nw_is_categorical_and_is_not_datetime(s: "nw.Series") -> bool: if isinstance(s.dtype, nw.Enum): # an explicit, user-defined category set is an unambiguous categorical @@ -117,11 +81,14 @@ def _nw_is_categorical_and_is_not_datetime(s: "nw.Series") -> bool: return True if isinstance(s.dtype, nw.Categorical): - # polars categorical categories are always string-backed, unlike - # pandas' pd.Categorical, which can have numeric categories - return not _nw_is_convertible_to_dt(s) - - if isinstance(s.dtype, nw.String): + # check for datetime only if the categories are not numeric, because + # a numeric-backed categorical (pandas-only - polars categories are + # always string-backed) can never hold dates + return _nw_categories_are_numeric(s) or not _nw_is_convertible_to_dt(s) + + if isinstance(s.dtype, (nw.String, nw.Object)): + # check for datetime only if the column cannot be cast as numeric, + # because if it could, it would be a numeric column, not a date return _nw_is_convertible_to_num(s) or not _nw_is_convertible_to_dt(s) return False @@ -132,9 +99,9 @@ def _nw_is_categorical_and_is_datetime(s: "nw.Series") -> bool: return False if isinstance(s.dtype, nw.Categorical): - return _nw_is_convertible_to_dt(s) + return not _nw_categories_are_numeric(s) and _nw_is_convertible_to_dt(s) - if isinstance(s.dtype, nw.String): + if isinstance(s.dtype, (nw.String, nw.Object)): return not _nw_is_convertible_to_num(s) and _nw_is_convertible_to_dt(s) return False diff --git a/feature_engine/variable_handling/check_variables.py b/feature_engine/variable_handling/check_variables.py index b823063ba..7fb56b530 100644 --- a/feature_engine/variable_handling/check_variables.py +++ b/feature_engine/variable_handling/check_variables.py @@ -5,14 +5,11 @@ import narwhals as nw import narwhals.dependencies as nwd from narwhals.typing import IntoDataFrame -from pandas.core.dtypes.common import is_numeric_dtype as is_numeric from feature_engine.variable_handling._variable_type_checks import ( - _is_categorical_and_is_datetime, _nw_is_categorical_and_is_datetime, _nw_is_date_or_datetime, ) -from feature_engine.variable_handling.dtypes import DATETIME_TYPES Variables = Union[int, str, List[Union[str, int]]] @@ -165,10 +162,9 @@ def check_datetime_variables( Notes ----- - For pandas dataframes, string columns are parsed with pandas' flexible, - dateutil-backed date guessing. For polars (and other non-pandas dataframes), - only ISO-8601 strings and native `Date`/`Datetime` columns are recognised - - polars has no equivalent flexible guesser. + String columns are parsed with flexible, dateutil-backed date guessing, in + addition to ISO-8601 strings and native `Date`/`Datetime` columns, + regardless of the dataframe library backing `X`. Examples -------- @@ -187,27 +183,18 @@ def check_datetime_variables( if isinstance(variables, (str, int)): variables = [variables] - if nwd.is_pandas_dataframe(X): - # find non datetime variables, if any: - non_datetime_vars = [] - for column in X[variables].select_dtypes(exclude=DATETIME_TYPES): - if is_numeric(X[column]) or not _is_categorical_and_is_datetime( - X[column] - ): - non_datetime_vars.append(column) - else: - sub_X = nw.from_native(X, eager_only=True).select(variables) - candidates = [ - column - for column in sub_X.columns - if not _nw_is_date_or_datetime(sub_X.schema[column]) - ] - non_datetime_vars = [ - column - for column in candidates - if sub_X.schema[column].is_numeric() - or not _nw_is_categorical_and_is_datetime(sub_X[column]) - ] + sub_X = nw.from_native(X, eager_only=True).select(variables) + candidates = [ + column + for column in sub_X.columns + if not _nw_is_date_or_datetime(sub_X.schema[column]) + ] + non_datetime_vars = [ + column + for column in candidates + if sub_X.schema[column].is_numeric() + or not _nw_is_categorical_and_is_datetime(sub_X.get_column(column)) + ] if len(non_datetime_vars) > 0: raise TypeError( diff --git a/feature_engine/variable_handling/find_variables.py b/feature_engine/variable_handling/find_variables.py index 2244b95b1..aa8f02b00 100644 --- a/feature_engine/variable_handling/find_variables.py +++ b/feature_engine/variable_handling/find_variables.py @@ -6,26 +6,36 @@ import narwhals as nw import narwhals.dependencies as nwd from narwhals.typing import IntoDataFrame -from pandas.api.types import is_datetime64_any_dtype as is_datetime -from pandas.core.dtypes.common import is_numeric_dtype as is_numeric from feature_engine.variable_handling._variable_type_checks import ( - _is_categorical_and_is_datetime, - _is_categorical_and_is_not_datetime, _nw_is_categorical_and_is_datetime, _nw_is_categorical_and_is_not_datetime, _nw_is_date_or_datetime, ) -from feature_engine.variable_handling.dtypes import DATETIME_TYPES # columns of these narwhals dtypes are candidates for being "categorical" - they # still need to be run through the datetime-disambiguation helpers, because a -# String/Categorical column may actually hold dates. +# String/Categorical column may actually hold dates. `Object` covers pandas +# columns holding a genuine mix of Python objects (e.g. strings and numbers), +# which narwhals cannot classify as String or any other single dtype. _CATEGORICAL_CANDIDATE_SELECTOR = ( - nw.selectors.categorical() | nw.selectors.enum() | nw.selectors.string() + nw.selectors.categorical() + | nw.selectors.enum() + | nw.selectors.string() + | nw.selectors.by_dtype(nw.Object) ) +def _categorical_candidates_in_order(nw_df) -> List[Union[str, int]]: + # nw.selectors' `|` combination returns matches grouped by which + # sub-selector matched (e.g. all categorical() hits, then all string() + # hits) rather than preserving the dataframe's original column order, for + # the pandas backend specifically. Re-sort against the dataframe's own + # column order so results are deterministic and backend-independent. + matched = set(nw_df.select(_CATEGORICAL_CANDIDATE_SELECTOR).columns) + return [column for column in nw_df.columns if column in matched] + + def find_numerical_variables( X: IntoDataFrame, return_empty: bool = False, @@ -135,20 +145,13 @@ def find_categorical_variables( >>> var_ ['var_cat'] """ - if nwd.is_pandas_dataframe(X): - variables = [ - column - for column in X.select_dtypes(include=["O", "category", "string"]).columns - if _is_categorical_and_is_not_datetime(X[column]) - ] - else: - nw_X = nw.from_native(X, eager_only=True) - candidates = nw_X.select(_CATEGORICAL_CANDIDATE_SELECTOR).columns - variables = [ - column - for column in candidates - if _nw_is_categorical_and_is_not_datetime(nw_X[column]) - ] + nw_X = nw.from_native(X, eager_only=True) + candidates = _categorical_candidates_in_order(nw_X) + variables = [ + column + for column in candidates + if _nw_is_categorical_and_is_not_datetime(nw_X.get_column(column)) + ] if len(variables) == 0: if return_empty is False: @@ -202,13 +205,10 @@ def find_datetime_variables( Notes ----- - For pandas dataframes, string columns are parsed with pandas' flexible, - dateutil-backed date guessing (the same as `pandas.to_datetime`), so formats - like "01-Jan-2010" or "10/11/12" are recognised. For polars (and other - non-pandas dataframes), only ISO-8601 strings and native `Date`/`Datetime` - columns are recognised automatically - polars has no equivalent flexible - guesser. Pass `variables` explicitly to a transformer if your polars date - strings use a different format. + String columns are parsed with flexible, dateutil-backed date guessing, so + formats like "01-Jan-2010" or "10/11/12" are recognised, in addition to + ISO-8601 strings and native `Date`/`Datetime` columns, regardless of the + dataframe library backing `X`. Examples -------- @@ -223,23 +223,16 @@ def find_datetime_variables( >>> var_date ['var_date'] """ - if nwd.is_pandas_dataframe(X): - variables = [ - column - for column in X.select_dtypes(exclude="number").columns - if is_datetime(X[column]) or _is_categorical_and_is_datetime(X[column]) - ] - else: - nw_X = nw.from_native(X, eager_only=True) - non_numeric = [ - column for column in nw_X.columns if not nw_X.schema[column].is_numeric() - ] - variables = [ - column - for column in non_numeric - if _nw_is_date_or_datetime(nw_X.schema[column]) - or _nw_is_categorical_and_is_datetime(nw_X[column]) - ] + nw_X = nw.from_native(X, eager_only=True) + non_numeric = [ + column for column in nw_X.columns if not nw_X.schema[column].is_numeric() + ] + variables = [ + column + for column in non_numeric + if _nw_is_date_or_datetime(nw_X.schema[column]) + or _nw_is_categorical_and_is_datetime(nw_X.get_column(column)) + ] if len(variables) == 0: if return_empty is False: @@ -304,32 +297,19 @@ def find_all_variables( >>> vars_all ['var_num', 'var_cat', 'var_date'] """ - if nwd.is_pandas_dataframe(X): - if exclude_datetime is True: - variables = X.select_dtypes(exclude=DATETIME_TYPES).columns.to_list() - variables = [ - var - for var in variables - if is_numeric(X[var]) or not _is_categorical_and_is_datetime(X[var]) - ] - else: - variables = X.columns.to_list() + nw_X = nw.from_native(X, eager_only=True) + if exclude_datetime is True: + variables = [ + var for var in nw_X.columns if not _nw_is_date_or_datetime(nw_X.schema[var]) + ] + variables = [ + var + for var in variables + if nw_X.schema[var].is_numeric() + or not _nw_is_categorical_and_is_datetime(nw_X.get_column(var)) + ] else: - nw_X = nw.from_native(X, eager_only=True) - if exclude_datetime is True: - variables = [ - var - for var in nw_X.columns - if not _nw_is_date_or_datetime(nw_X.schema[var]) - ] - variables = [ - var - for var in variables - if nw_X.schema[var].is_numeric() - or not _nw_is_categorical_and_is_datetime(nw_X[var]) - ] - else: - variables = list(nw_X.columns) + variables = list(nw_X.columns) if len(variables) == 0: if return_empty is False: @@ -400,26 +380,15 @@ def find_categorical_and_numerical_variables( >>> var_cat, var_num (['var_cat'], ['var_num']) """ - is_pandas = nwd.is_pandas_dataframe(X) - # cheap to build even when unused in the pandas branch below - avoids an - # Optional type that mypy can't narrow across the branches nw_X = nw.from_native(X, eager_only=True) # If the user passes just 1 variable outside a list. if isinstance(variables, (str, int)): - if is_pandas: - is_cat = X[ - variables - ].dtype.name == "category" or _is_categorical_and_is_not_datetime( - X[variables] - ) - is_num = is_numeric(X[variables]) - else: - s = nw_X[variables] - is_cat = isinstance( - s.dtype, (nw.Categorical, nw.Enum) - ) or _nw_is_categorical_and_is_not_datetime(s) - is_num = s.dtype.is_numeric() + s = nw_X.get_column(variables) + is_cat = isinstance( + s.dtype, (nw.Categorical, nw.Enum) + ) or _nw_is_categorical_and_is_not_datetime(s) + is_num = s.dtype.is_numeric() if is_cat: variables_cat = [variables] @@ -445,23 +414,13 @@ def find_categorical_and_numerical_variables( # If user leaves default None parameter. elif variables is None: - if is_pandas: - variables_cat = [ - column - for column in X.select_dtypes( - include=["O", "category", "string"] - ).columns - if _is_categorical_and_is_not_datetime(X[column]) - ] - variables_num = list(X.select_dtypes(include="number").columns) - else: - candidates = nw_X.select(_CATEGORICAL_CANDIDATE_SELECTOR).columns - variables_cat = [ - column - for column in candidates - if _nw_is_categorical_and_is_not_datetime(nw_X[column]) - ] - variables_num = list(nw_X.select(nw.selectors.numeric()).columns) + candidates = _categorical_candidates_in_order(nw_X) + variables_cat = [ + column + for column in candidates + if _nw_is_categorical_and_is_not_datetime(nw_X.get_column(column)) + ] + variables_num = list(nw_X.select(nw.selectors.numeric()).columns) if len(variables_num) == 0 and len(variables_cat) == 0: if return_empty is False: @@ -496,25 +455,13 @@ def find_categorical_and_numerical_variables( variables_cat = [] variables_num = [] - elif is_pandas: - # find categorical variables - variables_cat = [ - column - for column in X[variables] - .select_dtypes(include=["O", "category", "string"]) - .columns - if _is_categorical_and_is_not_datetime(X[column]) - ] - # find numerical variables - variables_num = list(X[variables].select_dtypes(include="number").columns) - else: sub_X = nw_X.select(variables) - candidates = sub_X.select(_CATEGORICAL_CANDIDATE_SELECTOR).columns + candidates = _categorical_candidates_in_order(sub_X) variables_cat = [ column for column in candidates - if _nw_is_categorical_and_is_not_datetime(sub_X[column]) + if _nw_is_categorical_and_is_not_datetime(sub_X.get_column(column)) ] variables_num = list(sub_X.select(nw.selectors.numeric()).columns) diff --git a/pyproject.toml b/pyproject.toml index ec9cd9079..64dbb8156 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "scikit-learn>=1.7.0", "scipy>=1.4.1", "narwhals>=2.0.0", + "python-dateutil>=2.8.2", ] classifiers = [ diff --git a/tests/test_variable_handling/test_check_variables.py b/tests/test_variable_handling/test_check_variables.py index 5ea8d1485..e1da482bc 100644 --- a/tests/test_variable_handling/test_check_variables.py +++ b/tests/test_variable_handling/test_check_variables.py @@ -149,9 +149,10 @@ def test_check_datetime_variables_returns_datetime_variables(make_df): def test_check_datetime_variables_returns_pandas_only_string_formats(df_datetime): - # "01-Jan-2010"-style and "10/11/12"-style strings are only ever recognised - # through pandas' flexible, dateutil-backed guessing - see the note in - # check_datetime_variables' docstring. + # "01-Jan-2010"-style and "10/11/12"-style strings are recognised via + # flexible, dateutil-backed guessing - see the note in + # check_datetime_variables' docstring. This fixture is pandas-only, but the + # guessing itself is backend-agnostic (also works for polars). vars_convertible_to_dt = ["date_range", "date_obj1", "date_obj2", "time_obj"] var_convertible_to_dt = "date_obj1" diff --git a/tests/test_variable_handling/test_fe_type_checks.py b/tests/test_variable_handling/test_fe_type_checks.py index 50a4f6cba..8e2992647 100644 --- a/tests/test_variable_handling/test_fe_type_checks.py +++ b/tests/test_variable_handling/test_fe_type_checks.py @@ -5,11 +5,6 @@ import polars as pl from feature_engine.variable_handling._variable_type_checks import ( - _is_categorical_and_is_datetime, - _is_categorical_and_is_not_datetime, - _is_categories_num, - _is_convertible_to_dt, - _is_convertible_to_num, _nw_is_categorical_and_is_datetime, _nw_is_categorical_and_is_not_datetime, _nw_is_convertible_to_dt, @@ -25,93 +20,9 @@ def nw_series(values, dtype=None): return nw.from_native(s, series_only=True) -def test_is_categories_num(df): - assert _is_categories_num(df["Name"]) is False - - df["Age"] = df["Age"].astype("category") - assert _is_categories_num(df["Age"]) is True - - -def test_is_convertible_to_num(df): - assert _is_convertible_to_num(df["Name"]) is False - assert _is_convertible_to_num(df["date_obj0"]) is False - - df["age_str"] = ["20", "21", "19", "18"] - assert _is_convertible_to_num(df["age_str"]) is True - - -def test_is_convertible_to_dt(df): - assert _is_convertible_to_dt(df["date_obj0"]) is True - assert _is_convertible_to_dt(df["date_range"]) is True - assert _is_convertible_to_dt(df["Name"]) is False - - df["age_str"] = ["20", "21", "19", "18"] - assert _is_convertible_to_dt(df["age_str"]) is False - - -def test_is_categorical_and_is_datetime(df, df_datetime): - assert _is_categorical_and_is_datetime(df["date_obj0"]) is True - assert _is_categorical_and_is_datetime(df["Name"]) is False - assert _is_categorical_and_is_datetime(df_datetime["date_obj1"]) is True - - df["age_str"] = ["20", "21", "19", "18"] - assert _is_categorical_and_is_datetime(df["age_str"]) is False - - df = df.copy() - # from pandas 3 onwards, object types that contain strings are not recognised as - # objects any more - df["Age"] = df["Age"].astype("O") - assert _is_categorical_and_is_datetime(df["Age"]) is False - - # Object Datetime - s_obj_dt = pd.Series([pd.Timestamp("2020-01-01")], dtype="object") - assert _is_categorical_and_is_datetime(s_obj_dt) is True - - # StringDtype Datetime (if convertible) - s_str_dt = pd.Series(["2020-01-01", "2020-01-02"], dtype="string") - assert _is_categorical_and_is_datetime(s_str_dt) is True - - # Numeric (should be False for both if and elif branches) - s_num = pd.Series([1, 2, 3]) - assert _is_categorical_and_is_datetime(s_num) is False - - # Categorical (should hit the 'if' branch) - s_cat = pd.Series(["a", "b"], dtype="category") - assert _is_categorical_and_is_datetime(s_cat) is False - - -def test_is_categorical_and_is_not_datetime(df): - assert _is_categorical_and_is_not_datetime(df["date_obj0"]) is False - assert _is_categorical_and_is_not_datetime(df["date_obj0"]) is False - assert _is_categorical_and_is_not_datetime(df["Name"]) is True - - df["age_str"] = ["20", "21", "19", "18"] - assert _is_categorical_and_is_not_datetime(df["age_str"]) is True - - # Object Integer - s_obj_int = pd.Series([1, 2], dtype="object") - assert _is_categorical_and_is_not_datetime(s_obj_int) is True - - # Object Datetime should be False - s_obj_dt = pd.Series([pd.Timestamp("2020-01-01")], dtype="object") - assert _is_categorical_and_is_not_datetime(s_obj_dt) is False - - # StringDtype (not convertible to numeric/datetime) should be True - s_str = pd.Series(["a", "b"], dtype="string") - assert _is_categorical_and_is_not_datetime(s_str) is True - - # Numeric should be False - s_num = pd.Series([1, 2, 3]) - assert _is_categorical_and_is_not_datetime(s_num) is False - - # Categorical should be True (it hits the 'if' branch) - s_cat = pd.Series(["a", "b"], dtype="category") - assert _is_categorical_and_is_not_datetime(s_cat) is True - - -# --------------------------------------------------------------------------- -# narwhals (polars) equivalents -# --------------------------------------------------------------------------- +def nw_pandas_series(values, dtype=None): + s = pd.Series(values, dtype=dtype) + return nw.from_native(s, series_only=True) def test_nw_is_date_or_datetime(): @@ -132,15 +43,35 @@ def test_nw_is_convertible_to_num(): is True ) + # object dtype columns (pandas-only concept - narwhals classifies a plain + # object dtype column of ints as `nw.Object`, not `nw.String`) + assert _nw_is_convertible_to_num(nw_pandas_series([1, 2], dtype="object")) is True + assert ( + _nw_is_convertible_to_num( + nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") + ) + is False + ) + def test_nw_is_convertible_to_dt(): assert _nw_is_convertible_to_dt(nw_series(["2020-01-01", "2020-01-02"])) is True assert _nw_is_convertible_to_dt(nw_series(["a", "b"])) is False assert _nw_is_convertible_to_dt(nw_series(["20", "21"])) is False - # polars has no dateutil-style guesser, unlike pandas, so non-ISO date - # strings are not recognised - assert _nw_is_convertible_to_dt(nw_series(["01-Jan-2010"])) is False + # flexible, dateutil-backed date guessing works for every backend now, not + # just pandas - so non-ISO formats are recognised here too + assert _nw_is_convertible_to_dt(nw_series(["01-Jan-2010"])) is True + assert _nw_is_convertible_to_dt(nw_series(["10/11/12"])) is True + + # an object dtype column holding actual datetime objects (e.g. pandas + # Timestamps) is trivially convertible, without needing to parse anything + assert ( + _nw_is_convertible_to_dt( + nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") + ) + is True + ) def test_nw_is_categorical_and_is_datetime(): @@ -170,6 +101,33 @@ def test_nw_is_categorical_and_is_datetime(): # numeric should be False assert _nw_is_categorical_and_is_datetime(nw_series([1, 2, 3])) is False + # a numeric-backed categorical (pandas-only - polars categories are always + # string-backed) can never be a datetime, regardless of the categories + numeric_cat = nw_pandas_series([20, 21, 19, 18], dtype="category") + assert _nw_is_categorical_and_is_datetime(numeric_cat) is False + + # a string-dtype pandas column with datetime-like values + assert ( + _nw_is_categorical_and_is_datetime( + nw_pandas_series(["2020-01-01", "2020-01-02"], dtype="string") + ) + is True + ) + + # object dtype column holding actual Timestamp objects + assert ( + _nw_is_categorical_and_is_datetime( + nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") + ) + is True + ) + + # object dtype column holding plain ints - not a datetime + assert ( + _nw_is_categorical_and_is_datetime(nw_pandas_series([1, 2], dtype="object")) + is False + ) + def test_nw_is_categorical_and_is_not_datetime(): assert ( @@ -198,3 +156,30 @@ def test_nw_is_categorical_and_is_not_datetime(): # numeric should be False assert _nw_is_categorical_and_is_not_datetime(nw_series([1, 2, 3])) is False + + # a numeric-backed categorical is categorical-and-not-datetime + numeric_cat = nw_pandas_series([20, 21, 19, 18], dtype="category") + assert _nw_is_categorical_and_is_not_datetime(numeric_cat) is True + + # object dtype column of plain ints + assert ( + _nw_is_categorical_and_is_not_datetime(nw_pandas_series([1, 2], dtype="object")) + is True + ) + + # object dtype column holding actual Timestamp objects - is a datetime, so + # not "categorical and not datetime" + assert ( + _nw_is_categorical_and_is_not_datetime( + nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") + ) + is False + ) + + # string-dtype pandas column not convertible to numeric or datetime + assert ( + _nw_is_categorical_and_is_not_datetime( + nw_pandas_series(["a", "b"], dtype="string") + ) + is True + ) diff --git a/tests/test_variable_handling/test_find_variables.py b/tests/test_variable_handling/test_find_variables.py index 8654724b5..e08112eff 100644 --- a/tests/test_variable_handling/test_find_variables.py +++ b/tests/test_variable_handling/test_find_variables.py @@ -101,9 +101,10 @@ def test_datetime_variables_finds_variables(make_df): def test_datetime_variables_finds_pandas_only_string_formats(df_datetime): - # "01-Jan-2010"-style, "10/11/12"-style and bare-time strings are only ever - # recognised through pandas' flexible, dateutil-backed guessing - polars has - # no equivalent, see the note in find_datetime_variables' docstring. + # "01-Jan-2010"-style, "10/11/12"-style and bare-time strings are + # recognised through flexible, dateutil-backed guessing - see the note in + # find_datetime_variables' docstring. This fixture is pandas-only, but the + # guessing itself is backend-agnostic (also works for polars). vars_dt = [ "date_range", "date_obj0", @@ -116,6 +117,19 @@ def test_datetime_variables_finds_pandas_only_string_formats(df_datetime): assert find_datetime_variables(df_datetime) == vars_dt +def test_datetime_variables_finds_flexible_string_formats_in_polars_too(): + # flexible, dateutil-backed date guessing is backend-agnostic, so polars + # now also recognises non-ISO formats it previously could not. + df = pl.DataFrame( + { + "var_num": [1, 2, 3], + "date_obj1": ["01-Jan-2010", "24-Feb-1945", "14-Jun-2100"], + "date_obj2": ["10/11/12", "12/31/09", "06/30/95"], + } + ) + assert find_datetime_variables(df) == ["date_obj1", "date_obj2"] + + @pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame]) def test_datetime_variables_raises_error(make_df): df = make_df(DATETIME_DATA) From c17d9cc4e1325717b8b14d8b44cabe94f56604d1 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Fri, 31 Jul 2026 16:01:11 +0200 Subject: [PATCH 3/4] Improve readability of narwhals date/type-check helpers, add missing tests Replace the double-parse-with-disagreeing-defaults trick in _looks_like_date_string with a direct call to dateutil's parser()._parse(), which exposes which date/time fields were actually found in a string without needing to approximate it - this also drops the now-unneeded sentinel default datetimes and the defensive str() coercion at its call site. Make truthiness checks and compound boolean returns explicit throughout the module, and restore the pre-narwhals function names that PR #978 had prefixed with _nw_ for no continuing reason. Rename test_fe_type_checks.py to test_variable_type_checks.py to match the module it tests, add docstrings, and add coverage for _looks_like_date_string and _is_categories_num, the two functions that previously had no direct tests. Co-Authored-By: Claude Sonnet 5 --- .../_variable_type_checks.py | 77 +++--- .../variable_handling/check_variables.py | 8 +- .../variable_handling/find_variables.py | 22 +- .../test_fe_type_checks.py | 185 ------------- .../test_variable_type_checks.py | 246 ++++++++++++++++++ 5 files changed, 304 insertions(+), 234 deletions(-) delete mode 100644 tests/test_variable_handling/test_fe_type_checks.py create mode 100644 tests/test_variable_handling/test_variable_type_checks.py diff --git a/feature_engine/variable_handling/_variable_type_checks.py b/feature_engine/variable_handling/_variable_type_checks.py index 885e18583..221928464 100644 --- a/feature_engine/variable_handling/_variable_type_checks.py +++ b/feature_engine/variable_handling/_variable_type_checks.py @@ -1,7 +1,7 @@ from datetime import date, datetime import narwhals as nw -from dateutil.parser import parse as _dateutil_parse +from dateutil.parser import parser as _dateutil_parser_cls # --------------------------------------------------------------------------- # narwhals implementation, used for every backend (pandas, polars, etc.) @@ -13,42 +13,43 @@ # --------------------------------------------------------------------------- -def _nw_is_date_or_datetime(dtype) -> bool: +def _is_date_or_datetime(dtype) -> bool: # nw.selectors.datetime() only matches Datetime, not Date, so this needs # its own explicit check. return isinstance(dtype, (nw.Date, nw.Datetime)) -_DATE_PARSE_DEFAULT_1 = datetime(1, 1, 1, 1, 1, 1) -_DATE_PARSE_DEFAULT_2 = datetime(2, 2, 2, 2, 2, 2) -_DATETIME_FIELDS = ("year", "month", "day", "hour", "minute", "second") +# reused across calls, mirroring how pandas keeps its own single DEFAULTPARSER +# instance internally instead of re-instantiating one every time. +_dateutil_parser = _dateutil_parser_cls() -def _looks_like_date_string(value: str) -> bool: - # dateutil.parser.parse() fills in any date/time component that isn't - # present in the string from a `default` datetime, so a bare number like - # "20" "parses" successfully as day=20 - it would wrongly be treated as a - # date. Parsing twice, with two defaults that differ in every field, - # reveals which fields were actually present in the string: those are the - # fields that agree between the two parses. Requiring at least 2 fields to - # be corroborated this way rejects bare numbers while still accepting real - # dates (including non-ISO formats like "01-Jan-2010") and bare times - # (like "21:45:23"). +def _looks_like_date_string(value) -> bool: + # parser().parse() (the public function) fills in any date/time component + # that isn't present in the string from a `default` datetime, so a bare + # number like "20" "parses" successfully as day=20 - it would wrongly be + # treated as a date. parser()._parse() (private) returns the intermediate + # result before that backfilling happens: its fields are None for + # anything not actually found in the string. Requiring at least 2 fields + # to be present rejects bare numbers while still accepting real dates + # (including non-ISO formats like "01-Jan-2010") and bare times (like + # "21:45:23"). try: - first = _dateutil_parse(value, default=_DATE_PARSE_DEFAULT_1) - second = _dateutil_parse(value, default=_DATE_PARSE_DEFAULT_2) - except (ValueError, OverflowError, TypeError): + result, _ = _dateutil_parser._parse(value) + except TypeError: return False - corroborated = sum( - 1 for attr in _DATETIME_FIELDS if getattr(first, attr) == getattr(second, attr) - ) - return corroborated >= 2 + if result is None: + return False + + fields = ("year", "month", "day", "hour", "minute", "second") + found_fields = sum(1 for field in fields if getattr(result, field) is not None) + return found_fields >= 2 -def _nw_is_convertible_to_num(s: "nw.Series") -> bool: +def _is_convertible_to_num(s: "nw.Series") -> bool: values = s.drop_nulls().to_list() - if not values: + if len(values) == 0: return False try: for value in values: @@ -58,23 +59,23 @@ def _nw_is_convertible_to_num(s: "nw.Series") -> bool: return True -def _nw_is_convertible_to_dt(s: "nw.Series") -> bool: +def _is_convertible_to_dt(s: "nw.Series") -> bool: values = s.drop_nulls().to_list() - if not values: + if len(values) == 0: return False for value in values: if isinstance(value, (date, datetime)): continue - if not _looks_like_date_string(str(value)): + if _looks_like_date_string(value) is False: return False return True -def _nw_categories_are_numeric(s: "nw.Series") -> bool: +def _is_categories_num(s: "nw.Series") -> bool: return s.cat.get_categories().dtype.is_numeric() -def _nw_is_categorical_and_is_not_datetime(s: "nw.Series") -> bool: +def _is_categorical_and_is_not_datetime(s: "nw.Series") -> bool: if isinstance(s.dtype, nw.Enum): # an explicit, user-defined category set is an unambiguous categorical # signal, unlike a generic string column, so skip the datetime check @@ -84,24 +85,32 @@ def _nw_is_categorical_and_is_not_datetime(s: "nw.Series") -> bool: # check for datetime only if the categories are not numeric, because # a numeric-backed categorical (pandas-only - polars categories are # always string-backed) can never hold dates - return _nw_categories_are_numeric(s) or not _nw_is_convertible_to_dt(s) + categories_are_numeric = _is_categories_num(s) + is_convertible_to_dt = _is_convertible_to_dt(s) + return categories_are_numeric is True or is_convertible_to_dt is False if isinstance(s.dtype, (nw.String, nw.Object)): # check for datetime only if the column cannot be cast as numeric, # because if it could, it would be a numeric column, not a date - return _nw_is_convertible_to_num(s) or not _nw_is_convertible_to_dt(s) + is_convertible_to_num = _is_convertible_to_num(s) + is_convertible_to_dt = _is_convertible_to_dt(s) + return is_convertible_to_num is True or is_convertible_to_dt is False return False -def _nw_is_categorical_and_is_datetime(s: "nw.Series") -> bool: +def _is_categorical_and_is_datetime(s: "nw.Series") -> bool: if isinstance(s.dtype, nw.Enum): return False if isinstance(s.dtype, nw.Categorical): - return not _nw_categories_are_numeric(s) and _nw_is_convertible_to_dt(s) + categories_are_numeric = _is_categories_num(s) + is_convertible_to_dt = _is_convertible_to_dt(s) + return categories_are_numeric is False and is_convertible_to_dt is True if isinstance(s.dtype, (nw.String, nw.Object)): - return not _nw_is_convertible_to_num(s) and _nw_is_convertible_to_dt(s) + is_convertible_to_num = _is_convertible_to_num(s) + is_convertible_to_dt = _is_convertible_to_dt(s) + return is_convertible_to_num is False and is_convertible_to_dt is True return False diff --git a/feature_engine/variable_handling/check_variables.py b/feature_engine/variable_handling/check_variables.py index 7fb56b530..0d3c6a257 100644 --- a/feature_engine/variable_handling/check_variables.py +++ b/feature_engine/variable_handling/check_variables.py @@ -7,8 +7,8 @@ from narwhals.typing import IntoDataFrame from feature_engine.variable_handling._variable_type_checks import ( - _nw_is_categorical_and_is_datetime, - _nw_is_date_or_datetime, + _is_categorical_and_is_datetime, + _is_date_or_datetime, ) Variables = Union[int, str, List[Union[str, int]]] @@ -187,13 +187,13 @@ def check_datetime_variables( candidates = [ column for column in sub_X.columns - if not _nw_is_date_or_datetime(sub_X.schema[column]) + if not _is_date_or_datetime(sub_X.schema[column]) ] non_datetime_vars = [ column for column in candidates if sub_X.schema[column].is_numeric() - or not _nw_is_categorical_and_is_datetime(sub_X.get_column(column)) + or not _is_categorical_and_is_datetime(sub_X.get_column(column)) ] if len(non_datetime_vars) > 0: diff --git a/feature_engine/variable_handling/find_variables.py b/feature_engine/variable_handling/find_variables.py index aa8f02b00..1cbe9ae45 100644 --- a/feature_engine/variable_handling/find_variables.py +++ b/feature_engine/variable_handling/find_variables.py @@ -8,9 +8,9 @@ from narwhals.typing import IntoDataFrame from feature_engine.variable_handling._variable_type_checks import ( - _nw_is_categorical_and_is_datetime, - _nw_is_categorical_and_is_not_datetime, - _nw_is_date_or_datetime, + _is_categorical_and_is_datetime, + _is_categorical_and_is_not_datetime, + _is_date_or_datetime, ) # columns of these narwhals dtypes are candidates for being "categorical" - they @@ -150,7 +150,7 @@ def find_categorical_variables( variables = [ column for column in candidates - if _nw_is_categorical_and_is_not_datetime(nw_X.get_column(column)) + if _is_categorical_and_is_not_datetime(nw_X.get_column(column)) ] if len(variables) == 0: @@ -230,8 +230,8 @@ def find_datetime_variables( variables = [ column for column in non_numeric - if _nw_is_date_or_datetime(nw_X.schema[column]) - or _nw_is_categorical_and_is_datetime(nw_X.get_column(column)) + if _is_date_or_datetime(nw_X.schema[column]) + or _is_categorical_and_is_datetime(nw_X.get_column(column)) ] if len(variables) == 0: @@ -300,13 +300,13 @@ def find_all_variables( nw_X = nw.from_native(X, eager_only=True) if exclude_datetime is True: variables = [ - var for var in nw_X.columns if not _nw_is_date_or_datetime(nw_X.schema[var]) + var for var in nw_X.columns if not _is_date_or_datetime(nw_X.schema[var]) ] variables = [ var for var in variables if nw_X.schema[var].is_numeric() - or not _nw_is_categorical_and_is_datetime(nw_X.get_column(var)) + or not _is_categorical_and_is_datetime(nw_X.get_column(var)) ] else: variables = list(nw_X.columns) @@ -387,7 +387,7 @@ def find_categorical_and_numerical_variables( s = nw_X.get_column(variables) is_cat = isinstance( s.dtype, (nw.Categorical, nw.Enum) - ) or _nw_is_categorical_and_is_not_datetime(s) + ) or _is_categorical_and_is_not_datetime(s) is_num = s.dtype.is_numeric() if is_cat: @@ -418,7 +418,7 @@ def find_categorical_and_numerical_variables( variables_cat = [ column for column in candidates - if _nw_is_categorical_and_is_not_datetime(nw_X.get_column(column)) + if _is_categorical_and_is_not_datetime(nw_X.get_column(column)) ] variables_num = list(nw_X.select(nw.selectors.numeric()).columns) @@ -461,7 +461,7 @@ def find_categorical_and_numerical_variables( variables_cat = [ column for column in candidates - if _nw_is_categorical_and_is_not_datetime(sub_X.get_column(column)) + if _is_categorical_and_is_not_datetime(sub_X.get_column(column)) ] variables_num = list(sub_X.select(nw.selectors.numeric()).columns) diff --git a/tests/test_variable_handling/test_fe_type_checks.py b/tests/test_variable_handling/test_fe_type_checks.py deleted file mode 100644 index 8e2992647..000000000 --- a/tests/test_variable_handling/test_fe_type_checks.py +++ /dev/null @@ -1,185 +0,0 @@ -from datetime import date - -import narwhals as nw -import pandas as pd -import polars as pl - -from feature_engine.variable_handling._variable_type_checks import ( - _nw_is_categorical_and_is_datetime, - _nw_is_categorical_and_is_not_datetime, - _nw_is_convertible_to_dt, - _nw_is_convertible_to_num, - _nw_is_date_or_datetime, -) - - -def nw_series(values, dtype=None): - s = pl.Series("x", values) - if dtype is not None: - s = s.cast(dtype) - return nw.from_native(s, series_only=True) - - -def nw_pandas_series(values, dtype=None): - s = pd.Series(values, dtype=dtype) - return nw.from_native(s, series_only=True) - - -def test_nw_is_date_or_datetime(): - assert _nw_is_date_or_datetime(nw_series([date(2020, 1, 1)]).dtype) is True - assert ( - _nw_is_date_or_datetime(nw_series(["2020-01-01"]).str.to_datetime().dtype) - is True - ) - assert _nw_is_date_or_datetime(nw_series(["a", "b"]).dtype) is False - assert _nw_is_date_or_datetime(nw_series([1, 2, 3]).dtype) is False - - -def test_nw_is_convertible_to_num(): - assert _nw_is_convertible_to_num(nw_series(["20", "21", "19"])) is True - assert _nw_is_convertible_to_num(nw_series(["a", "b"])) is False - assert ( - _nw_is_convertible_to_num(nw_series(["20", "21"], dtype=pl.Categorical)) - is True - ) - - # object dtype columns (pandas-only concept - narwhals classifies a plain - # object dtype column of ints as `nw.Object`, not `nw.String`) - assert _nw_is_convertible_to_num(nw_pandas_series([1, 2], dtype="object")) is True - assert ( - _nw_is_convertible_to_num( - nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") - ) - is False - ) - - -def test_nw_is_convertible_to_dt(): - assert _nw_is_convertible_to_dt(nw_series(["2020-01-01", "2020-01-02"])) is True - assert _nw_is_convertible_to_dt(nw_series(["a", "b"])) is False - assert _nw_is_convertible_to_dt(nw_series(["20", "21"])) is False - - # flexible, dateutil-backed date guessing works for every backend now, not - # just pandas - so non-ISO formats are recognised here too - assert _nw_is_convertible_to_dt(nw_series(["01-Jan-2010"])) is True - assert _nw_is_convertible_to_dt(nw_series(["10/11/12"])) is True - - # an object dtype column holding actual datetime objects (e.g. pandas - # Timestamps) is trivially convertible, without needing to parse anything - assert ( - _nw_is_convertible_to_dt( - nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") - ) - is True - ) - - -def test_nw_is_categorical_and_is_datetime(): - assert ( - _nw_is_categorical_and_is_datetime( - nw_series(["2020-01-01", "2020-01-02"], dtype=pl.Categorical) - ) - is True - ) - assert ( - _nw_is_categorical_and_is_datetime(nw_series(["a", "b"], dtype=pl.Categorical)) - is False - ) - assert _nw_is_categorical_and_is_datetime(nw_series(["2020-01-01"])) is True - assert _nw_is_categorical_and_is_datetime(nw_series(["20", "21"])) is False - assert _nw_is_categorical_and_is_datetime(nw_series(["a", "b"])) is False - - # an explicit Enum is always treated as categorical, never as datetime - enum_dtype = pl.Enum(["2020-01-01", "2020-01-02"]) - assert ( - _nw_is_categorical_and_is_datetime( - nw_series(["2020-01-01", "2020-01-02"], dtype=enum_dtype) - ) - is False - ) - - # numeric should be False - assert _nw_is_categorical_and_is_datetime(nw_series([1, 2, 3])) is False - - # a numeric-backed categorical (pandas-only - polars categories are always - # string-backed) can never be a datetime, regardless of the categories - numeric_cat = nw_pandas_series([20, 21, 19, 18], dtype="category") - assert _nw_is_categorical_and_is_datetime(numeric_cat) is False - - # a string-dtype pandas column with datetime-like values - assert ( - _nw_is_categorical_and_is_datetime( - nw_pandas_series(["2020-01-01", "2020-01-02"], dtype="string") - ) - is True - ) - - # object dtype column holding actual Timestamp objects - assert ( - _nw_is_categorical_and_is_datetime( - nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") - ) - is True - ) - - # object dtype column holding plain ints - not a datetime - assert ( - _nw_is_categorical_and_is_datetime(nw_pandas_series([1, 2], dtype="object")) - is False - ) - - -def test_nw_is_categorical_and_is_not_datetime(): - assert ( - _nw_is_categorical_and_is_not_datetime( - nw_series(["2020-01-01", "2020-01-02"], dtype=pl.Categorical) - ) - is False - ) - assert ( - _nw_is_categorical_and_is_not_datetime( - nw_series(["a", "b"], dtype=pl.Categorical) - ) - is True - ) - assert _nw_is_categorical_and_is_not_datetime(nw_series(["2020-01-01"])) is False - assert _nw_is_categorical_and_is_not_datetime(nw_series(["20", "21"])) is True - assert _nw_is_categorical_and_is_not_datetime(nw_series(["a", "b"])) is True - - # an explicit Enum is always treated as categorical - assert ( - _nw_is_categorical_and_is_not_datetime( - nw_series(["a", "b"], dtype=pl.Enum(["a", "b"])) - ) - is True - ) - - # numeric should be False - assert _nw_is_categorical_and_is_not_datetime(nw_series([1, 2, 3])) is False - - # a numeric-backed categorical is categorical-and-not-datetime - numeric_cat = nw_pandas_series([20, 21, 19, 18], dtype="category") - assert _nw_is_categorical_and_is_not_datetime(numeric_cat) is True - - # object dtype column of plain ints - assert ( - _nw_is_categorical_and_is_not_datetime(nw_pandas_series([1, 2], dtype="object")) - is True - ) - - # object dtype column holding actual Timestamp objects - is a datetime, so - # not "categorical and not datetime" - assert ( - _nw_is_categorical_and_is_not_datetime( - nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") - ) - is False - ) - - # string-dtype pandas column not convertible to numeric or datetime - assert ( - _nw_is_categorical_and_is_not_datetime( - nw_pandas_series(["a", "b"], dtype="string") - ) - is True - ) diff --git a/tests/test_variable_handling/test_variable_type_checks.py b/tests/test_variable_handling/test_variable_type_checks.py new file mode 100644 index 000000000..e09da4438 --- /dev/null +++ b/tests/test_variable_handling/test_variable_type_checks.py @@ -0,0 +1,246 @@ +from datetime import date + +import narwhals as nw +import pandas as pd +import polars as pl + +from feature_engine.variable_handling._variable_type_checks import ( + _is_categorical_and_is_datetime, + _is_categorical_and_is_not_datetime, + _is_categories_num, + _is_convertible_to_dt, + _is_convertible_to_num, + _is_date_or_datetime, + _looks_like_date_string, +) + + +def nw_series(values, dtype=None): + s = pl.Series("x", values) + if dtype is not None: + s = s.cast(dtype) + return nw.from_native(s, series_only=True) + + +def nw_pandas_series(values, dtype=None): + s = pd.Series(values, dtype=dtype) + return nw.from_native(s, series_only=True) + + +def test_is_date_or_datetime(): + """A dtype is a date or datetime if it is narwhals' Date or Datetime type.""" + assert _is_date_or_datetime(nw_series([date(2020, 1, 1)]).dtype) is True + assert ( + _is_date_or_datetime(nw_series(["2020-01-01"]).str.to_datetime().dtype) + is True + ) + assert _is_date_or_datetime(nw_series(["a", "b"]).dtype) is False + assert _is_date_or_datetime(nw_series([1, 2, 3]).dtype) is False + + +def test_looks_like_date_string(): + """A string looks like a date if dateutil finds at least 2 date/time fields + in it - this rejects bare numbers that dateutil would otherwise happily + "parse" as a single field (e.g. a day), while still accepting real dates + in non-ISO formats and bare times. + """ + # real dates, including non-ISO formats + assert _looks_like_date_string("2020-01-01") is True + assert _looks_like_date_string("01-Jan-2010") is True + assert _looks_like_date_string("10/11/12") is True + + # bare times + assert _looks_like_date_string("21:45:23") is True + assert _looks_like_date_string("08:00") is True + + # partial dates + assert _looks_like_date_string("Jan 2020") is True + + # bare numbers dateutil could misparse as a single date/time field + assert _looks_like_date_string("20") is False + assert _looks_like_date_string("1999") is False + assert _looks_like_date_string("12") is False + + # non-date garbage + assert _looks_like_date_string("hello") is False + assert _looks_like_date_string("") is False + + # non-string values (e.g. from a mixed-type pandas Object column) must not + # raise, they simply aren't date strings + assert _looks_like_date_string(20) is False + assert _looks_like_date_string(1.5) is False + assert _looks_like_date_string(None) is False + + +def test_is_convertible_to_num(): + """A series is convertible to numeric if every non-null value can be cast + to float. + """ + assert _is_convertible_to_num(nw_series(["20", "21", "19"])) is True + assert _is_convertible_to_num(nw_series(["a", "b"])) is False + assert ( + _is_convertible_to_num(nw_series(["20", "21"], dtype=pl.Categorical)) + is True + ) + + # object dtype columns (pandas-only concept - narwhals classifies a plain + # object dtype column of ints as `nw.Object`, not `nw.String`) + assert _is_convertible_to_num(nw_pandas_series([1, 2], dtype="object")) is True + assert ( + _is_convertible_to_num( + nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") + ) + is False + ) + + +def test_is_convertible_to_dt(): + """A series is convertible to datetime if every non-null value is either a + real date/datetime object, or a string that looks like a date. + """ + assert _is_convertible_to_dt(nw_series(["2020-01-01", "2020-01-02"])) is True + assert _is_convertible_to_dt(nw_series(["a", "b"])) is False + assert _is_convertible_to_dt(nw_series(["20", "21"])) is False + + # flexible, dateutil-backed date guessing works for every backend now, not + # just pandas - so non-ISO formats are recognised here too + assert _is_convertible_to_dt(nw_series(["01-Jan-2010"])) is True + assert _is_convertible_to_dt(nw_series(["10/11/12"])) is True + + # an object dtype column holding actual datetime objects (e.g. pandas + # Timestamps) is trivially convertible, without needing to parse anything + assert ( + _is_convertible_to_dt( + nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") + ) + is True + ) + + +def test_is_categories_num(): + """A categorical series' categories are numeric if their dtype is numeric - + only possible for pandas, since polars categories are always string-backed. + """ + non_numeric_cat = nw_series(["a", "b", "c"], dtype=pl.Categorical) + assert _is_categories_num(non_numeric_cat) is False + + numeric_cat = nw_pandas_series([20, 21, 19, 18], dtype="category") + assert _is_categories_num(numeric_cat) is True + + +def test_is_categorical_and_is_datetime(): + """A series is categorical-and-datetime if it is a Categorical/String/Object + column whose values are dates, but not an Enum (an explicit category set is + never treated as a datetime) or a numeric-backed categorical. + """ + assert ( + _is_categorical_and_is_datetime( + nw_series(["2020-01-01", "2020-01-02"], dtype=pl.Categorical) + ) + is True + ) + assert ( + _is_categorical_and_is_datetime(nw_series(["a", "b"], dtype=pl.Categorical)) + is False + ) + assert _is_categorical_and_is_datetime(nw_series(["2020-01-01"])) is True + assert _is_categorical_and_is_datetime(nw_series(["20", "21"])) is False + assert _is_categorical_and_is_datetime(nw_series(["a", "b"])) is False + + # an explicit Enum is always treated as categorical, never as datetime + enum_dtype = pl.Enum(["2020-01-01", "2020-01-02"]) + assert ( + _is_categorical_and_is_datetime( + nw_series(["2020-01-01", "2020-01-02"], dtype=enum_dtype) + ) + is False + ) + + # numeric should be False + assert _is_categorical_and_is_datetime(nw_series([1, 2, 3])) is False + + # a numeric-backed categorical (pandas-only - polars categories are always + # string-backed) can never be a datetime, regardless of the categories + numeric_cat = nw_pandas_series([20, 21, 19, 18], dtype="category") + assert _is_categorical_and_is_datetime(numeric_cat) is False + + # a string-dtype pandas column with datetime-like values + assert ( + _is_categorical_and_is_datetime( + nw_pandas_series(["2020-01-01", "2020-01-02"], dtype="string") + ) + is True + ) + + # object dtype column holding actual Timestamp objects + assert ( + _is_categorical_and_is_datetime( + nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") + ) + is True + ) + + # object dtype column holding plain ints - not a datetime + assert ( + _is_categorical_and_is_datetime(nw_pandas_series([1, 2], dtype="object")) + is False + ) + + +def test_is_categorical_and_is_not_datetime(): + """A series is categorical-and-not-datetime if it is a Categorical/String/ + Object/Enum column whose values are not dates. + """ + assert ( + _is_categorical_and_is_not_datetime( + nw_series(["2020-01-01", "2020-01-02"], dtype=pl.Categorical) + ) + is False + ) + assert ( + _is_categorical_and_is_not_datetime( + nw_series(["a", "b"], dtype=pl.Categorical) + ) + is True + ) + assert _is_categorical_and_is_not_datetime(nw_series(["2020-01-01"])) is False + assert _is_categorical_and_is_not_datetime(nw_series(["20", "21"])) is True + assert _is_categorical_and_is_not_datetime(nw_series(["a", "b"])) is True + + # an explicit Enum is always treated as categorical + assert ( + _is_categorical_and_is_not_datetime( + nw_series(["a", "b"], dtype=pl.Enum(["a", "b"])) + ) + is True + ) + + # numeric should be False + assert _is_categorical_and_is_not_datetime(nw_series([1, 2, 3])) is False + + # a numeric-backed categorical is categorical-and-not-datetime + numeric_cat = nw_pandas_series([20, 21, 19, 18], dtype="category") + assert _is_categorical_and_is_not_datetime(numeric_cat) is True + + # object dtype column of plain ints + assert ( + _is_categorical_and_is_not_datetime(nw_pandas_series([1, 2], dtype="object")) + is True + ) + + # object dtype column holding actual Timestamp objects - is a datetime, so + # not "categorical and not datetime" + assert ( + _is_categorical_and_is_not_datetime( + nw_pandas_series([pd.Timestamp("2020-01-01")], dtype="object") + ) + is False + ) + + # string-dtype pandas column not convertible to numeric or datetime + assert ( + _is_categorical_and_is_not_datetime( + nw_pandas_series(["a", "b"], dtype="string") + ) + is True + ) From 8a0978c3d52e797a6b271eb055bc9f1b072bf5e0 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Fri, 31 Jul 2026 17:35:06 +0200 Subject: [PATCH 4/4] Replace per-column schema access with bulk narwhals selectors for speed nw_X.schema is not cached - every access re-derives the full schema from the underlying native dataframe, so checking dtype-based conditions (is_numeric(), native Date/Datetime, categorical/enum/string) one column at a time inside a loop was quadratic instead of linear. Replace each such loop with a single nw_df.select().columns call converted to a set, then a plain membership test per column - confirmed old vs new give identical results, and measured 8x-120x speedups depending on backend and column count. Also use by_dtype(Date, Datetime) to bulk-detect native datetime columns in one pass, only falling back to the expensive per-value _is_categorical_and_is_datetime check for columns that aren't already known to be numeric or natively datetime. Drop the now-unused _is_date_or_datetime import from both files. Simplify _looks_like_date_string's comment to link directly to the pandas source it mirrors, and instantiate dateutil's parser() per call instead of reusing a module-level instance. Co-Authored-By: Claude Sonnet 5 --- .../_variable_type_checks.py | 29 +------ .../variable_handling/check_variables.py | 26 +++--- .../variable_handling/find_variables.py | 86 +++++++------------ 3 files changed, 51 insertions(+), 90 deletions(-) diff --git a/feature_engine/variable_handling/_variable_type_checks.py b/feature_engine/variable_handling/_variable_type_checks.py index 221928464..b1a0abf29 100644 --- a/feature_engine/variable_handling/_variable_type_checks.py +++ b/feature_engine/variable_handling/_variable_type_checks.py @@ -1,16 +1,7 @@ from datetime import date, datetime import narwhals as nw -from dateutil.parser import parser as _dateutil_parser_cls - -# --------------------------------------------------------------------------- -# narwhals implementation, used for every backend (pandas, polars, etc.) -# -# Flexible date-string recognition (e.g. "01-Jan-2010", "10/11/12", not just -# ISO-8601) is implemented directly on top of `dateutil` - the same library -# pandas.to_datetime delegates to internally for this - so it works -# identically regardless of the underlying dataframe library. -# --------------------------------------------------------------------------- +from dateutil.parser import parser def _is_date_or_datetime(dtype) -> bool: @@ -19,23 +10,11 @@ def _is_date_or_datetime(dtype) -> bool: return isinstance(dtype, (nw.Date, nw.Datetime)) -# reused across calls, mirroring how pandas keeps its own single DEFAULTPARSER -# instance internally instead of re-instantiating one every time. -_dateutil_parser = _dateutil_parser_cls() - - def _looks_like_date_string(value) -> bool: - # parser().parse() (the public function) fills in any date/time component - # that isn't present in the string from a `default` datetime, so a bare - # number like "20" "parses" successfully as day=20 - it would wrongly be - # treated as a date. parser()._parse() (private) returns the intermediate - # result before that backfilling happens: its fields are None for - # anything not actually found in the string. Requiring at least 2 fields - # to be present rejects bare numbers while still accepting real dates - # (including non-ISO formats like "01-Jan-2010") and bare times (like - # "21:45:23"). + # taken from pandas + # https://github.com/pandas-dev/pandas/blob/cbae8aea4a31a4052736ab0d23f284ff1e78aa06/pandas/_libs/tslibs/parsing.pyx#L666 try: - result, _ = _dateutil_parser._parse(value) + result, _ = parser()._parse(value) except TypeError: return False diff --git a/feature_engine/variable_handling/check_variables.py b/feature_engine/variable_handling/check_variables.py index 0d3c6a257..7a77a0c3d 100644 --- a/feature_engine/variable_handling/check_variables.py +++ b/feature_engine/variable_handling/check_variables.py @@ -8,7 +8,6 @@ from feature_engine.variable_handling._variable_type_checks import ( _is_categorical_and_is_datetime, - _is_date_or_datetime, ) Variables = Union[int, str, List[Union[str, int]]] @@ -57,9 +56,8 @@ def check_numerical_variables( not_numerical = len(X[variables].select_dtypes(exclude="number").columns) > 0 else: sub_X = nw.from_native(X, eager_only=True).select(variables) - not_numerical = any( - not sub_X.schema[column].is_numeric() for column in sub_X.columns - ) + numeric_cols = set(sub_X.select(nw.selectors.numeric()).columns) + not_numerical = any(column not in numeric_cols for column in sub_X.columns) if not_numerical: raise TypeError( @@ -122,9 +120,13 @@ def check_categorical_variables( ) else: sub_X = nw.from_native(X, eager_only=True).select(variables) + categorical_cols = set( + sub_X.select( + nw.selectors.categorical() | nw.selectors.enum() | nw.selectors.string() + ).columns + ) not_categorical = any( - not isinstance(sub_X.schema[column], (nw.Categorical, nw.Enum, nw.String)) - for column in sub_X.columns + column not in categorical_cols for column in sub_X.columns ) if not_categorical: @@ -184,15 +186,15 @@ def check_datetime_variables( variables = [variables] sub_X = nw.from_native(X, eager_only=True).select(variables) - candidates = [ - column - for column in sub_X.columns - if not _is_date_or_datetime(sub_X.schema[column]) - ] + datetime_cols = set( + sub_X.select(nw.selectors.by_dtype(nw.Date, nw.Datetime)).columns + ) + numeric_cols = set(sub_X.select(nw.selectors.numeric()).columns) + candidates = [column for column in sub_X.columns if column not in datetime_cols] non_datetime_vars = [ column for column in candidates - if sub_X.schema[column].is_numeric() + if column in numeric_cols or not _is_categorical_and_is_datetime(sub_X.get_column(column)) ] diff --git a/feature_engine/variable_handling/find_variables.py b/feature_engine/variable_handling/find_variables.py index 1cbe9ae45..419bea95e 100644 --- a/feature_engine/variable_handling/find_variables.py +++ b/feature_engine/variable_handling/find_variables.py @@ -4,36 +4,29 @@ from typing import List, Tuple, Union import narwhals as nw -import narwhals.dependencies as nwd from narwhals.typing import IntoDataFrame from feature_engine.variable_handling._variable_type_checks import ( _is_categorical_and_is_datetime, _is_categorical_and_is_not_datetime, - _is_date_or_datetime, -) - -# columns of these narwhals dtypes are candidates for being "categorical" - they -# still need to be run through the datetime-disambiguation helpers, because a -# String/Categorical column may actually hold dates. `Object` covers pandas -# columns holding a genuine mix of Python objects (e.g. strings and numbers), -# which narwhals cannot classify as String or any other single dtype. -_CATEGORICAL_CANDIDATE_SELECTOR = ( - nw.selectors.categorical() - | nw.selectors.enum() - | nw.selectors.string() - | nw.selectors.by_dtype(nw.Object) ) -def _categorical_candidates_in_order(nw_df) -> List[Union[str, int]]: - # nw.selectors' `|` combination returns matches grouped by which - # sub-selector matched (e.g. all categorical() hits, then all string() - # hits) rather than preserving the dataframe's original column order, for - # the pandas backend specifically. Re-sort against the dataframe's own - # column order so results are deterministic and backend-independent. - matched = set(nw_df.select(_CATEGORICAL_CANDIDATE_SELECTOR).columns) - return [column for column in nw_df.columns if column in matched] +def _find_nw_categoricals(nw_df) -> List[Union[str, int]]: + _NW_SELECTOR = ( + nw.selectors.categorical() + | nw.selectors.enum() + | nw.selectors.string() + | nw.selectors.by_dtype(nw.Object) + ) + matched = set(nw_df.select(_NW_SELECTOR).columns) + vars = [column for column in nw_df.columns if column in matched] + variables = [ + column + for column in vars + if _is_categorical_and_is_not_datetime(nw_df.get_column(column)) + ] + return variables def find_numerical_variables( @@ -79,11 +72,8 @@ def find_numerical_variables( >>> var_ ['var_num'] """ - if nwd.is_pandas_dataframe(X): - variables = list(X.select_dtypes(include="number").columns) - else: - nw_X = nw.from_native(X, eager_only=True) - variables = list(nw_X.select(nw.selectors.numeric()).columns) + nw_X = nw.from_native(X, eager_only=True) + variables = list(nw_X.select(nw.selectors.numeric()).columns) if len(variables) == 0: if return_empty is False: @@ -146,12 +136,7 @@ def find_categorical_variables( ['var_cat'] """ nw_X = nw.from_native(X, eager_only=True) - candidates = _categorical_candidates_in_order(nw_X) - variables = [ - column - for column in candidates - if _is_categorical_and_is_not_datetime(nw_X.get_column(column)) - ] + variables = _find_nw_categoricals(nw_X) if len(variables) == 0: if return_empty is False: @@ -224,13 +209,16 @@ def find_datetime_variables( ['var_date'] """ nw_X = nw.from_native(X, eager_only=True) - non_numeric = [ - column for column in nw_X.columns if not nw_X.schema[column].is_numeric() - ] + numeric_cols = set(nw_X.select(nw.selectors.numeric()).columns) + non_numeric = [column for column in nw_X.columns if column not in numeric_cols] + + datetime_cols = set( + nw_X.select(nw.selectors.by_dtype(nw.Date, nw.Datetime)).columns + ) variables = [ column for column in non_numeric - if _is_date_or_datetime(nw_X.schema[column]) + if column in datetime_cols or _is_categorical_and_is_datetime(nw_X.get_column(column)) ] @@ -299,13 +287,15 @@ def find_all_variables( """ nw_X = nw.from_native(X, eager_only=True) if exclude_datetime is True: - variables = [ - var for var in nw_X.columns if not _is_date_or_datetime(nw_X.schema[var]) - ] + datetime_cols = set( + nw_X.select(nw.selectors.by_dtype(nw.Date, nw.Datetime)).columns + ) + numeric_cols = set(nw_X.select(nw.selectors.numeric()).columns) + variables = [var for var in nw_X.columns if var not in datetime_cols] variables = [ var for var in variables - if nw_X.schema[var].is_numeric() + if var in numeric_cols or not _is_categorical_and_is_datetime(nw_X.get_column(var)) ] else: @@ -414,12 +404,7 @@ def find_categorical_and_numerical_variables( # If user leaves default None parameter. elif variables is None: - candidates = _categorical_candidates_in_order(nw_X) - variables_cat = [ - column - for column in candidates - if _is_categorical_and_is_not_datetime(nw_X.get_column(column)) - ] + variables_cat = _find_nw_categoricals(nw_X) variables_num = list(nw_X.select(nw.selectors.numeric()).columns) if len(variables_num) == 0 and len(variables_cat) == 0: @@ -457,12 +442,7 @@ def find_categorical_and_numerical_variables( else: sub_X = nw_X.select(variables) - candidates = _categorical_candidates_in_order(sub_X) - variables_cat = [ - column - for column in candidates - if _is_categorical_and_is_not_datetime(sub_X.get_column(column)) - ] + variables_cat = _find_nw_categoricals(sub_X) variables_num = list(sub_X.select(nw.selectors.numeric()).columns) return variables_cat, variables_num