diff --git a/ibis/expr/api.py b/ibis/expr/api.py index 1415868cb8aa..2db05618e17a 100644 --- a/ibis/expr/api.py +++ b/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 diff --git a/ibis/tests/expr/test_case.py b/ibis/tests/expr/test_case.py index 4bb6cdcdec2c..1ef15a8fa69e 100644 --- a/ibis/tests/expr/test_case.py +++ b/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