diff --git a/py-polars/polars/functions/whenthen.py b/py-polars/polars/functions/whenthen.py index 9cab73bcb4da..39ad2958a8ae 100644 --- a/py-polars/polars/functions/whenthen.py +++ b/py-polars/polars/functions/whenthen.py @@ -36,8 +36,7 @@ def then( See Also -------- - when : Start another when, then, otherwise layer. - otherwise : Values to return in case of the predicate being `False`. + pl.when : Documentation for `when, then, otherwise` """ expr_ = expr_to_lit_or_expr(expr) @@ -52,8 +51,7 @@ def otherwise( See Also -------- - when : Start another when, then, otherwise layer. - then : Values to return in case of the predicate being `True`. + pl.when : Documentation for `when, then, otherwise` """ expr = expr_to_lit_or_expr(expr) @@ -91,8 +89,7 @@ def otherwise( See Also -------- - when : Start another when, then, otherwise layer. - then : Values to return in case of the predicate being `True`. + pl.when : Documentation for `when, then, otherwise` """ expr = expr_to_lit_or_expr(expr) @@ -125,8 +122,7 @@ def then( See Also -------- - when : Start another when, then, otherwise layer. - otherwise : Values to return in case of the predicate being `False`. + pl.when : Documentation for `when, then, otherwise` """ expr = expr_to_lit_or_expr(expr) @@ -138,6 +134,13 @@ def when(expr: Expr | bool | Series) -> When: """ Start a "when, then, otherwise" expression. + Expression similar to an `if-else` statement in Python. Always initiated by a + `pl.when().then()`. Optionally followed by chaining + one or more `.when().then()` statements. If none of the conditions + are `True`, an optional `.otherwise()` can be + appended at the end. If not appended, and none of the conditions are `True`, `None` + will be returned. + Examples -------- Below we add a column with the value 1, where column "foo" > 2 and the value -1 @@ -145,18 +148,21 @@ def when(expr: Expr | bool | Series) -> When: >>> df = pl.DataFrame({"foo": [1, 3, 4], "bar": [3, 4, 0]}) >>> df.with_columns( - ... pl.when(pl.col("foo") > 2).then(pl.lit(1)).otherwise(pl.lit(-1)) + ... pl.when(pl.col("foo") > 2) + ... .then(pl.lit(1)) + ... .otherwise(pl.lit(-1)) + ... .alias("val") ... ) shape: (3, 3) - ┌─────┬─────┬─────────┐ - │ foo ┆ bar ┆ literal │ - │ --- ┆ --- ┆ --- │ - │ i64 ┆ i64 ┆ i32 │ - ╞═════╪═════╪═════════╡ - │ 1 ┆ 3 ┆ -1 │ - │ 3 ┆ 4 ┆ 1 │ - │ 4 ┆ 0 ┆ 1 │ - └─────┴─────┴─────────┘ + ┌─────┬─────┬─────┐ + │ foo ┆ bar ┆ val │ + │ --- ┆ --- ┆ --- │ + │ i64 ┆ i64 ┆ i32 │ + ╞═════╪═════╪═════╡ + │ 1 ┆ 3 ┆ -1 │ + │ 3 ┆ 4 ┆ 1 │ + │ 4 ┆ 0 ┆ 1 │ + └─────┴─────┴─────┘ Or with multiple `when, thens` chained: @@ -166,22 +172,40 @@ def when(expr: Expr | bool | Series) -> When: ... .when(pl.col("bar") > 2) ... .then(4) ... .otherwise(-1) + ... .alias("val") ... ) shape: (3, 3) - ┌─────┬─────┬─────────┐ - │ foo ┆ bar ┆ literal │ - │ --- ┆ --- ┆ --- │ - │ i64 ┆ i64 ┆ i32 │ - ╞═════╪═════╪═════════╡ - │ 1 ┆ 3 ┆ 4 │ - │ 3 ┆ 4 ┆ 1 │ - │ 4 ┆ 0 ┆ 1 │ - └─────┴─────┴─────────┘ - - See Also - -------- - then : Values to return in case of the predicate being `True`. - otherwise : Values to return in case of the predicate being `False`. + ┌─────┬─────┬─────┐ + │ foo ┆ bar ┆ val │ + │ --- ┆ --- ┆ --- │ + │ i64 ┆ i64 ┆ i32 │ + ╞═════╪═════╪═════╡ + │ 1 ┆ 3 ┆ 4 │ + │ 3 ┆ 4 ┆ 1 │ + │ 4 ┆ 0 ┆ 1 │ + └─────┴─────┴─────┘ + + Chained `when, thens` should be read as `if, elif, ... elif`, not + as `if, if, ... if`, i.e. the first condition that evaluates to True will + be picked. Note how in the example above for the second row in the dataframe, + where `foo=3` and `bar=4`, the first `when` evaluates to `True`, and therefore + the second `when`, which is also `True`, is not evaluated. + + The `otherwise` at the end is optional. If left out, any rows where none + of the `when` expressions evaluate to True, are set to `null`: + + >>> df.with_columns(pl.when(pl.col("foo") > 2).then(pl.lit(1)).alias("val")) + shape: (3, 3) + ┌─────┬─────┬──────┐ + │ foo ┆ bar ┆ val │ + │ --- ┆ --- ┆ --- │ + │ i64 ┆ i64 ┆ i32 │ + ╞═════╪═════╪══════╡ + │ 1 ┆ 3 ┆ null │ + │ 3 ┆ 4 ┆ 1 │ + │ 4 ┆ 0 ┆ 1 │ + └─────┴─────┴──────┘ + """ expr = expr_to_lit_or_expr(expr)