Implement .where (#453)

README.md

@@ -97,6 +97,7 @@ API Coverage
 - `map`
 - `map_overlap`
 - `map_partitions`
+- `mask`
 - `max`
 - `mean`
 - `memory_usage`
@@ -130,6 +131,7 @@ API Coverage
 - `to_timestamp`
 - `var`
 - `visualize`
+- `where`
 
 
 **`dask_expr.Series`**
@@ -160,6 +162,7 @@ API Coverage
 - `isna`
 - `map`
 - `map_partitions`
+- `mask`
 - `max`
 - `mean`
 - `memory_usage`
@@ -188,6 +191,7 @@ API Coverage
 - `value_counts`
 - `var`
 - `visualize`
+- `where`
 
 
 **`dask_expr.Index`**

dask_expr/_collection.py

@@ -601,6 +601,11 @@ def mask(self, cond, other=np.nan):
     def round(self, decimals=0):
         return new_collection(self.expr.round(decimals))
 
+    def where(self, cond, other=np.nan):
+        cond = cond.expr if isinstance(cond, FrameBase) else cond
+        other = other.expr if isinstance(other, FrameBase) else other
+        return new_collection(self.expr.where(cond, other))
+
     def apply(self, function, *args, **kwargs):
         return new_collection(self.expr.apply(function, *args, **kwargs))
 

dask_expr/_expr.py

@@ -720,6 +720,9 @@ def mask(self, cond, other=np.nan):
     def round(self, decimals=0):
         return Round(self, decimals=decimals)
 
+    def where(self, cond, other=np.nan):
+        return Where(self, cond=cond, other=other)
+
     def apply(self, function, *args, **kwargs):
         return Apply(self, function, args, kwargs)
 
@@ -1766,6 +1769,13 @@ class Round(Elemwise):
     operation = M.round
 
 
+class Where(Elemwise):
+    _projection_passthrough = True
+    _parameters = ["frame", "cond", "other"]
+    _defaults = {"other": np.nan}
+    operation = M.where
+
+
 class Abs(Elemwise):
     _projection_passthrough = True
     _parameters = ["frame"]

dask_expr/tests/test_collection.py

@@ -331,6 +331,12 @@ def test_to_timestamp(pdf, how):
         lambda df: df.x.mask(df.x == 10, 42),
         lambda df: df.abs(),
         lambda df: df.x.abs(),
+        lambda df: df.where(df.x == 10, 42),
+        lambda df: df.where(df.x == 10),
+        lambda df: df.where(lambda df: df.x % 2 == 0, 42),
+        lambda df: df.where(df.x == 10, df + 2),
+        lambda df: df.where(df.x == 10, lambda df: df + 2),
+        lambda df: df.x.where(df.x == 10, 42),
         lambda df: df.rename(columns={"x": "xx"}),
         lambda df: df.rename(columns={"x": "xx"}).xx,
         lambda df: df.rename(columns={"x": "xx"})[["xx"]],