feat(api): support deferred and literal values in ibis.ifelse

ibis/expr/api.py

@@ -1774,40 +1774,47 @@ def watermark(time_col: str, allowed_delay: ir.IntervalScalar) -> Watermark:
 geo_y_min = _geo_deprecated(ir.GeoSpatialValue.y_min)
 geo_unary_union = _geo_deprecated(ir.GeoSpatialColumn.unary_union)
 
-ifelse = _deferred(ir.BooleanValue.ifelse)
-"""Construct a ternary conditional expression.
-
-Parameters
-----------
-condition : ir.BooleanValue
-    A boolean expression
-true_expr : ir.Value
-    Expression to return if `condition` evaluates to `True`
-false_expr : ir.Value
-    Expression to return if `condition` evaluates to `False` or `NULL`
-
-Returns
--------
-Value : ir.Value
-    The value of `true_expr` if `condition` is `True` else `false_expr`
 
-Examples
---------
->>> import ibis
->>> ibis.options.interactive = True
->>> t = ibis.memtable({"is_person": [True, False, True, None]})
->>> ibis.ifelse(t.is_person, "yes", "no")
-┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
-┃ IfElse(is_person, 'yes', 'no') ┃
-┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
-│ string                         │
-├────────────────────────────────┤
-│ yes                            │
-│ no                             │
-│ yes                            │
-│ no                             │
-└────────────────────────────────┘
-"""
+@deferrable
+def ifelse(condition: Any, true_expr: Any, false_expr: Any) -> ir.Value:
+    """Construct a ternary conditional expression.
+
+    Parameters
+    ----------
+    condition
+        A boolean expression
+    true_expr
+        Expression to return if `condition` evaluates to `True`
+    false_expr
+        Expression to return if `condition` evaluates to `False` or `NULL`
+
+    Returns
+    -------
+    Value : ir.Value
+        The value of `true_expr` if `condition` is `True` else `false_expr`
+
+    Examples
+    --------
+    >>> import ibis
+    >>> ibis.options.interactive = True
+    >>> t = ibis.memtable({"condition": [True, False, True, None]})
+    >>> ibis.ifelse(t.condition, "yes", "no")
+    ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
+    ┃ IfElse(condition, 'yes', 'no') ┃
+    ┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┩
+    │ string                         │
+    ├────────────────────────────────┤
+    │ yes                            │
+    │ no                             │
+    │ yes                            │
+    │ no                             │
+    └────────────────────────────────┘
+    """
+    if not isinstance(condition, ir.Value):
+        condition = literal(condition, type="bool")
+    elif not condition.type().is_boolean():
+        condition = condition.cast("bool")
+    return condition.ifelse(true_expr, false_expr)
 
 
 @util.deprecated(instead="use `ibis.ifelse` instead", as_of="7.0")
@@ -1816,11 +1823,11 @@ def where(cond, true_expr, false_expr) -> ir.Value:
 
     Parameters
     ----------
-    cond : ir.Value
+    cond
         Boolean conditional expression
-    true_expr : ir.Value
+    true_expr
         Expression to return if `cond` evaluates to `True`
-    false_expr : ir.Value
+    false_expr
         Expression to return if `cond` evaluates to `False` or `NULL`
 
     Returns

ibis/tests/expr/test_case.py

@@ -4,15 +4,46 @@
 import ibis.expr.datatypes as dt
 import ibis.expr.operations as ops
 import ibis.expr.types as ir
+from ibis import _
 from ibis.tests.util import assert_equal, assert_pickle_roundtrip
 
 
-def test_ifelse(table):
+def test_ifelse_method(table):
     bools = table.g.isnull()
     result = bools.ifelse("foo", "bar")
     assert isinstance(result, ir.StringColumn)
 
 
+def test_ifelse_function_literals():
+    res = ibis.ifelse(True, 1, 2)
+    sol = ibis.literal(True, type="bool").ifelse(1, 2)
+    assert res.equals(sol)
+
+    # condition is explicitly cast to bool
+    res = ibis.ifelse(1, 1, 2)
+    sol = ibis.literal(1, type="bool").ifelse(1, 2)
+    assert res.equals(sol)
+
+
+def test_ifelse_function_exprs(table):
+    res = ibis.ifelse(table.g.isnull(), 1, table.a)
+    sol = table.g.isnull().ifelse(1, table.a)
+    assert res.equals(sol)
+
+    # condition is cast if not already bool
+    res = ibis.ifelse(table.a, 1, table.b)
+    sol = table.a.cast("bool").ifelse(1, table.b)
+    assert res.equals(sol)
+
+
+def test_ifelse_function_deferred(table):
+    expr = ibis.ifelse(_.g.isnull(), _.a, 2)
+    assert repr(expr) == "ifelse(_.g.isnull(), _.a, 2)"
+    res = expr.resolve(table)
+    sol = table.g.isnull().ifelse(table.a, 2)
+    assert res.equals(sol)
+
+
 def test_simple_case_expr(table):
     case1, result1 = "foo", table.a
     case2, result2 = "bar", table.c