work on Polars ifelse and where

Author: JohnMount

data_algebra/polars_model.py

@@ -309,10 +309,10 @@ def _populate_expr_impl_map() -> Dict[int, Dict[str, Callable]]:
         "parse_datetime": lambda x, format : x.cast(str).str.strptime(pl.Datetime, fmt=format, strict=False).cast(pl.Datetime),
     }
     impl_map_3 = {
-        "if_else": lambda a, b, c: pl.when(a).then(b).otherwise(c),
+        "if_else": lambda a, b, c: pl.when(a.is_null()).then(pl.lit(None)).otherwise(pl.when(a).then(b).otherwise(c)),
         "mapv": _mapv,
         "trimstr": lambda a, b, c: a.trimstr(b, c),
-        "where": lambda a, b, c: pl.when(a).then(b).otherwise(c),
+        "where": lambda a, b, c: pl.when(a.is_null()).then(c).otherwise(pl.when(a).then(b).otherwise(c)),
     }
     impl_map = {
         0: impl_map_0,
@@ -436,7 +436,7 @@ def drop_indices(self, df) -> None:
 
     def bad_column_positions(self, x):
         """
-        Return vector indicating which entries are bad (null or nan) (vectorized).
+        Return vector indicating which entries are null (vectorized).
         """
         return x.is_null()
 

tests/test_if_else.py

@@ -139,5 +139,4 @@ def test_if_else_where():
     assert data_algebra.test_util.equivalent_frames(res_sqlite, expect)
     sqlite_handle.close()
     data_algebra.test_util.check_transform(ops=ops, data=d, expect=expect,
-        try_on_Polars=False,  # TODO: turn this on
     )

tests/test_polars.py

@@ -640,10 +640,10 @@ def test_is_inf_polars():
         ops = (
             data_algebra.descr(d=d)
                 .extend({
-                    'is_inf': 'a.is_inf().if_else(1, 0)',
-                    'is_nan': 'a.is_nan().if_else(1, 0)',
-                    'is_bad': 'a.is_bad().if_else(1, 0)',
-                    'is_null': 'a.is_null().if_else(1, 0)',
+                    'is_inf': 'a.is_inf().where(1, 0)',
+                    'is_nan': 'a.is_nan().where(1, 0)',
+                    'is_bad': 'a.is_bad().where(1, 0)',
+                    'is_null': 'a.is_null().where(1, 0)',
                     })
         )
         res_polars = ops.transform(d)