fixup! Add case_when API * Used to support conditional assignment operation.

Author: ELHoussineT

pandas/core/case_when.py

@@ -18,6 +18,7 @@ def case_when(*args, default: Any = lib.no_default) -> Callable:
     Create a callable for assignment based on a condition or multiple conditions.
 
     This is useful when you want to assign a column based on multiple conditions.
+    Uses `np.select` to perform the assignment.
 
     Parameters
     ----------
@@ -26,8 +27,9 @@ def case_when(*args, default: Any = lib.no_default) -> Callable:
             `condition0`, `value0`, `condition1`, `value1`, ...
         `condition` can be a 1-D boolean array/series or a callable
         that evaluate to a 1-D boolean array/series.
-    default : Any, default is `None`.
-        The default value to be used if all conditions evaluate False.
+    default : Any, default is `np.nan`.
+        The default value to be used if all conditions evaluate False. This value
+        will be passed to the underlying `np.select` call.
 
     Returns
     -------
@@ -72,7 +74,7 @@ def case_when(*args, default: Any = lib.no_default) -> Callable:
         )
 
     if default is lib.no_default:
-        default = None
+        default = np.nan
 
     def _eval(df: pd.DataFrame) -> np.ndarray:
         booleans = []

pandas/tests/test_case_when.py

@@ -1,3 +1,4 @@
+import numpy as np
 import pytest  # noqa
 
 import pandas as pd
@@ -10,51 +11,51 @@ def test_case_when_multiple_conditions_callable(self):
         result = df.assign(
             new_column=pd.case_when(
                 lambda x: x.a == 1,
-                "first",
+                1,
                 lambda x: (x.a > 1) & (x.b == 5),
-                "second",
+                2,
             )
         )
-        expected = df.assign(new_column=["first", "second", None])
+        expected = df.assign(new_column=[1, 2, np.nan])
         tm.assert_frame_equal(result, expected)
 
     def test_case_when_multiple_conditions_array_series(self):
         df = pd.DataFrame(dict(a=[1, 2, 3], b=[4, 5, 6]))
         result = df.assign(
             new_column=pd.case_when(
                 [True, False, False],
-                "first",
+                1,
                 pd.Series([False, True, False]),
-                "second",
+                2,
             )
         )
-        expected = df.assign(new_column=["first", "second", None])
+        expected = df.assign(new_column=[1, 2, np.nan])
         tm.assert_frame_equal(result, expected)
 
     def test_case_when_multiple_conditions_callable_default(self):
         df = pd.DataFrame(dict(a=[1, 2, 3], b=[4, 5, 6]))
         result = df.assign(
             new_column=pd.case_when(
                 lambda x: x.a == 1,
-                "first",
+                1,
                 lambda x: (x.a > 1) & (x.b == 5),
-                "second",
-                default="default",
+                2,
+                default=-1,
             )
         )
-        expected = df.assign(new_column=["first", "second", "default"])
+        expected = df.assign(new_column=[1, 2, -1])
         tm.assert_frame_equal(result, expected)
 
     def test_case_when_multiple_conditions_callable_default_series(self):
         df = pd.DataFrame(dict(a=[1, 2, 3], b=[4, 5, 6]))
         result = df.assign(
             new_column=pd.case_when(
                 lambda x: x.a == 1,
-                "first",
+                1,
                 lambda x: (x.a > 1) & (x.b == 5),
-                "second",
+                2,
                 default=df.b,
             )
         )
-        expected = df.assign(new_column=["first", "second", "6"])
+        expected = df.assign(new_column=[1, 2, 6])
         tm.assert_frame_equal(result, expected)