markfairbanks / tidypolars

Tidy interface to polars
http://tidypolars.readthedocs.io
MIT License
321 stars 10 forks source link

Implement `coalesce()` #162

Closed markfairbanks closed 2 years ago

markfairbanks commented 2 years ago

Or should we just recommend using if_else()/case_when()?

markfairbanks commented 2 years ago

Haven't done much testing on this, but this is an initial attempt at it:

import tidypolars as tp
from tidypolars import col

from tidypolars.utils import _args_as_list

def coalesce(*args):
    args = _args_as_list(args)
    expr = tp.if_else(args[0].is_null(), args[1], args[0])
    if len(args) > 2:
        sequence = list(range(2, len(args) + 1, 1))
        for val in sequence:
            expr = tp.if_else(expr.is_null(), args[val], expr)
    return expr

df = tp.Tibble(x = [None, None, 1], y = [2, None, 2], z = [3, 3, 3])

df.mutate(
    coalesce_x = coalesce(col('x'), col('y'), col('z'))
)

┌──────┬──────┬─────┬────────────┐
│ x    ┆ y    ┆ z   ┆ coalesce_x │
│ ---  ┆ ---  ┆ --- ┆ ---        │
│ i64  ┆ i64  ┆ i64 ┆ i64        │
╞══════╪══════╪═════╪════════════╡
│ null ┆ 2    ┆ 3   ┆ 2          │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ null ┆ null ┆ 3   ┆ 3          │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1    ┆ 2    ┆ 3   ┆ 1          │
└──────┴──────┴─────┴────────────┘
markfairbanks commented 2 years ago

FYI @alexandro-ag

import tidypolars as tp
from tidypolars import col

df = tp.Tibble(x = [None, None, 1], y = [2, None, 2], z = [3, 3, 3])

df.mutate(
    coalesce_x = tp.coalesce(col('x'), col('y'), col('z'))
)
┌──────┬──────┬─────┬────────────┐
│ x    ┆ y    ┆ z   ┆ coalesce_x │
│ ---  ┆ ---  ┆ --- ┆ ---        │
│ i64  ┆ i64  ┆ i64 ┆ i64        │
╞══════╪══════╪═════╪════════════╡
│ null ┆ 2    ┆ 3   ┆ 2          │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ null ┆ null ┆ 3   ┆ 3          │
├╌╌╌╌╌╌┼╌╌╌╌╌╌┼╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1    ┆ 2    ┆ 3   ┆ 1          │
└──────┴──────┴─────┴────────────┘