polars.when

polars.when(
    *predicates: IntoExprColumn | Iterable[IntoExprColumn] | bool,
    **constraints: Any,
) → When

Начинает выражение when-then-otherwise.

Всегда начинается с pl.when().then()., за которым при необходимости могут следовать одна или несколько инструкций .when().then().

В конце можно добавить необязательный .otherwise(). Если он не указан, используется значение по умолчанию .otherwise(None).

Как и в случае с coalesce(), выбирается значение из первого условия, результатом которого является True.

Если все условия имеют значение False, выбирается значение otherwise.

Параметры:

predicates

Условие или условия, которые должны быть выполнены для применения последующей инструкции. Принимает одно или несколько булевых выражений, которые неявно объединяются с помощью &.

constraints

Условия в виде именованных аргументов col_name = value, которые рассматриваются как проверки на равенство, например x = 123. Как и в параметре predicates, несколько условий неявно объединяются с помощью &.

Предупреждение

Polars вычисляет параллельно все выражения, переданные в when-then-otherwise, а фильтрацию выполняет после этого. Это означает, что каждое выражение должно быть допустимым само по себе, независимо от условий в цепочке when-then-otherwise.

Примечания

• Строковые входные данные, например when("string"), then("string") или otherwise("string"), интерпретируются как имена столбцов. Для создания строковых значений можно использовать lit().
• Имя результата выражения берётся из первой инструкции then. Оно не зависит от predicates и constraints.

Примеры

Ниже мы добавляем столбец со значением 1, если столбец «foo» > 2, и со значением 1 + столбец «bar» в противном случае.

>>> df = pl.DataFrame({"foo": [1, 3, 4], "bar": [3, 4, 0]})
>>> df.with_columns(
...     pl.when(pl.col.foo > 2).then(1).otherwise(1 + pl.col.bar).alias("val")
... )
shape: (3, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ val │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 3   ┆ 4   │
│ 3   ┆ 4   ┆ 1   │
│ 4   ┆ 0   ┆ 1   │
└─────┴─────┴─────┘

Обратите внимание, что when-then всегда выполняет все выражения.

Результаты объединяются слева направо: выбирается значение then из первого условия when, которое имеет значение True.

Если ни одно условие when не имеет значения True, выбирается значение otherwise.

>>> df.with_columns(
...     when = pl.col.foo > 2,
...     then = 1,
...     otherwise = 1 + pl.col.bar
... ).with_columns(
...     pl.when("when").then("then").otherwise("otherwise").alias("val")
... )
shape: (3, 6)
┌─────┬─────┬───────┬──────┬───────────┬─────┐
│ foo ┆ bar ┆ when  ┆ then ┆ otherwise ┆ val │
│ --- ┆ --- ┆ ---   ┆ ---  ┆ ---       ┆ --- │
│ i64 ┆ i64 ┆ bool  ┆ i32  ┆ i64       ┆ i64 │
╞═════╪═════╪═══════╪══════╪═══════════╪═════╡
│ 1   ┆ 3   ┆ false ┆ 1    ┆ 4         ┆ 4   │
│ 3   ┆ 4   ┆ true  ┆ 1    ┆ 5         ┆ 1   │
│ 4   ┆ 0   ┆ true  ┆ 1    ┆ 1         ┆ 1   │
└─────┴─────┴───────┴──────┴───────────┴─────┘

Обратите внимание: при обычном использовании Polars одиночная строка интерпретируется как имя столбца.

>>> df.with_columns(
...     when = pl.col.foo > 2,
...     then = "foo",
...     otherwise = "bar"
... )
shape: (3, 5)
┌─────┬─────┬───────┬──────┬───────────┐
│ foo ┆ bar ┆ when  ┆ then ┆ otherwise │
│ --- ┆ --- ┆ ---   ┆ ---  ┆ ---       │
│ i64 ┆ i64 ┆ bool  ┆ i64  ┆ i64       │
╞═════╪═════╪═══════╪══════╪═══════════╡
│ 1   ┆ 3   ┆ false ┆ 1    ┆ 3         │
│ 3   ┆ 4   ┆ true  ┆ 3    ┆ 4         │
│ 4   ┆ 0   ┆ true  ┆ 4    ┆ 0         │
└─────┴─────┴───────┴──────┴───────────┘

Для единообразия when-then работает так же.

>>> df.with_columns(
...     pl.when(pl.col.foo > 2).then("foo").otherwise("bar").alias("val")
... )
shape: (3, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ val │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 1   ┆ 3   ┆ 3   │
│ 3   ┆ 4   ┆ 3   │
│ 4   ┆ 0   ┆ 4   │
└─────┴─────┴─────┘

Для создания строковых значений можно использовать lit().

>>> df.with_columns(
...     pl.when(pl.col.foo > 2)
...     .then(pl.lit("foo"))
...     .otherwise(pl.lit("bar"))
...     .alias("val")
... )
shape: (3, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ val │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ str │
╞═════╪═════╪═════╡
│ 1   ┆ 3   ┆ bar │
│ 3   ┆ 4   ┆ foo │
│ 4   ┆ 0   ┆ foo │
└─────┴─────┴─────┘

Можно объединять в цепочку несколько инструкций when-then.

>>> df.with_columns(
...     pl.when(pl.col.foo > 2)
...     .then(1)
...     .when(pl.col.bar > 2)
...     .then(4)
...     .otherwise(-1)
...     .alias("val")
... )
shape: (3, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ val │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i32 │
╞═════╪═════╪═════╡
│ 1   ┆ 3   ┆ 4   │
│ 3   ┆ 4   ┆ 1   │
│ 4   ┆ 0   ┆ 1   │
└─────┴─────┴─────┘

В случае с foo=3 и bar=4 оба условия имеют значение True, но выбирается первое значение (то есть 1).

>>> df.with_columns(
...     when1 = pl.col.foo > 2,
...     then1 = 1,
...     when2 = pl.col.bar > 2,
...     then2 = 4,
...     otherwise = -1
... )
shape: (3, 7)
┌─────┬─────┬───────┬───────┬───────┬───────┬───────────┐
│ foo ┆ bar ┆ when1 ┆ then1 ┆ when2 ┆ then2 ┆ otherwise │
│ --- ┆ --- ┆ ---   ┆ ---   ┆ ---   ┆ ---   ┆ ---       │
│ i64 ┆ i64 ┆ bool  ┆ i32   ┆ bool  ┆ i32   ┆ i32       │
╞═════╪═════╪═══════╪═══════╪═══════╪═══════╪═══════════╡
│ 1   ┆ 3   ┆ false ┆ 1     ┆ true  ┆ 4     ┆ -1        │
│ 3   ┆ 4   ┆ true  ┆ 1     ┆ true  ┆ 4     ┆ -1        │
│ 4   ┆ 0   ┆ true  ┆ 1     ┆ false ┆ 4     ┆ -1        │
└─────┴─────┴───────┴───────┴───────┴───────┴───────────┘

Инструкция otherwise необязательна; если она не указана, используется .otherwise(None).

Этот приём часто используют, чтобы заменять значения на null.

>>> df.with_columns(pl.when(pl.col.foo == 3).then("bar"))
shape: (3, 2)
┌─────┬──────┐
│ foo ┆ bar  │
│ --- ┆ ---  │
│ i64 ┆ i64  │
╞═════╪══════╡
│ 1   ┆ null │
│ 3   ┆ 4    │
│ 4   ┆ null │
└─────┴──────┘

when принимает именованные аргументы как сокращённую запись условий равенства.

>>> df.with_columns(pl.when(foo=3).then("bar"))
shape: (3, 2)
┌─────┬──────┐
│ foo ┆ bar  │
│ --- ┆ ---  │
│ i64 ┆ i64  │
╞═════╪══════╡
│ 1   ┆ null │
│ 3   ┆ 4    │
│ 4   ┆ null │
└─────┴──────┘

Несколько предикатов, переданных в when, объединяются с помощью &

>>> df.with_columns(
...     pl.when(pl.col.foo > 2, pl.col.bar < 3) # when((pred1) & (pred2))
...     .then(pl.lit("Yes"))
...     .otherwise(pl.lit("No"))
...     .alias("val")
... )
shape: (3, 3)
┌─────┬─────┬─────┐
│ foo ┆ bar ┆ val │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ str │
╞═════╪═════╪═════╡
│ 1   ┆ 3   ┆ No  │
│ 3   ┆ 4   ┆ No  │
│ 4   ┆ 0   ┆ Yes │
└─────┴─────┴─────┘

Это также можно представить как неявное присутствие all_horizontal().

>>> df.with_columns(
...     when = pl.all_horizontal(pl.col.foo > 2, pl.col.bar < 3)
... )
shape: (3, 3)
┌─────┬─────┬───────┐
│ foo ┆ bar ┆ when  │
│ --- ┆ --- ┆ ---   │
│ i64 ┆ i64 ┆ bool  │
╞═════╪═════╪═══════╡
│ 1   ┆ 3   ┆ false │
│ 3   ┆ 4   ┆ false │
│ 4   ┆ 0   ┆ true  │
└─────┴─────┴───────┘

Структуры можно использовать для возврата нескольких значений.

Здесь мы меняем местами значения «foo» и «bar», если «foo» больше 2.

>>> df.with_columns(
...     pl.when(pl.col.foo > 2)
...     .then(pl.struct(foo="bar", bar="foo"))
...     .otherwise(pl.struct("foo", "bar"))
...     .struct.unnest()
... )
shape: (3, 2)
┌─────┬─────┐
│ foo ┆ bar │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 1   ┆ 3   │
│ 4   ┆ 3   │
│ 0   ┆ 4   │
└─────┴─────┘

Полям структуры присваиваются те же имена, что и целевым столбцам, после чего структура разворачивается.

>>> df.with_columns(
...     when = pl.col.foo > 2,
...     then = pl.struct(foo="bar", bar="foo"),
...     otherwise = pl.struct("foo", "bar")
... )
shape: (3, 5)
┌─────┬─────┬───────┬───────────┬───────────┐
│ foo ┆ bar ┆ when  ┆ then      ┆ otherwise │
│ --- ┆ --- ┆ ---   ┆ ---       ┆ ---       │
│ i64 ┆ i64 ┆ bool  ┆ struct[2] ┆ struct[2] │
╞═════╪═════╪═══════╪═══════════╪═══════════╡
│ 1   ┆ 3   ┆ false ┆ {3,1}     ┆ {1,3}     │
│ 3   ┆ 4   ┆ true  ┆ {4,3}     ┆ {3,4}     │
│ 4   ┆ 0   ┆ true  ┆ {0,4}     ┆ {4,0}     │
└─────┴─────┴───────┴───────────┴───────────┘

Имя результата выражения when-then берётся из первой ветви then.

Здесь мы пытаемся присвоить всем столбцам значение 0, если хотя бы один столбец содержит значение меньше 2.

>>> df.with_columns( 
...    pl.when(pl.any_horizontal(pl.all() < 2))
...    .then(0)
...    .otherwise(pl.all())
... )
# ComputeError: the name 'literal' passed to `LazyFrame.with_columns` is duplicate

Чтобы отдать предпочтение выражению столбца, можно использовать name.keep().

>>> df.with_columns(
...    pl.when(pl.any_horizontal(pl.all() < 2))
...    .then(0)
...    .otherwise(pl.all())
...    .name.keep()
... )
shape: (3, 2)
┌─────┬─────┐
│ foo ┆ bar │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 0   ┆ 0   │
│ 3   ┆ 4   │
│ 0   ┆ 0   │
└─────┴─────┘

Логику также можно изменить, переместив выражение столбца внутрь then.

>>> df.with_columns(
...     pl.when(pl.any_horizontal(pl.all() < 2).not_())
...     .then(pl.all())
...     .otherwise(0)
... )
shape: (3, 2)
┌─────┬─────┐
│ foo ┆ bar │
│ --- ┆ --- │
│ i64 ┆ i64 │
╞═════╪═════╡
│ 0   ┆ 0   │
│ 3   ┆ 4   │
│ 0   ┆ 0   │
└─────┴─────┘