diff --git a/tests/test_processing.py b/tests/test_processing.py index fbe94ab..77b2aaa 100644 --- a/tests/test_processing.py +++ b/tests/test_processing.py @@ -93,6 +93,20 @@ def test_is_nan(): is_nan(pl.Series([np.nan, 1.0, None])).to_numpy(), np.array([True, False, None]), ) + for series in ( + pl.Series(["a", None, "b"]), + pl.Series([True, None, False]), + pl.Series([1, None, 2]), + pl.Series(["a", None, "b"], dtype=pl.Categorical), + ): + np.testing.assert_equal( + is_nan(series).to_numpy(), + np.array([False, False, False]), + ) + np.testing.assert_equal( + is_nan_or_none(series).to_numpy(), + np.array([False, True, False]), + ) def test_is_none(): diff --git a/utilsforecast/processing.py b/utilsforecast/processing.py index c6bf34d..691fcfd 100644 --- a/utilsforecast/processing.py +++ b/utilsforecast/processing.py @@ -233,8 +233,11 @@ def filter_with_mask( def is_nan(s: Series) -> Series: if isinstance(s, pd.Series): out = s.isna() - else: + elif s.dtype.is_float(): out = s.is_nan() + else: + # polars only supports is_nan on numeric columns and only floats can hold NaN + out = pl.repeat(False, s.len(), eager=True).alias(s.name) return out