listtransformationpython-polars

divide every nth element of a list with a number in polars DataFrame column of lists

I have the following DataFrame

import polars as pl

data = {
    'col': [[11, 21, 31, 41, 51], [12, 22, 32, 42, 52], [13, 23, 33, 43, 53]]
}
df = pl.DataFrame(data)

shape: (3, 1)
┌────────────────┐
│ col            │
│ ---            │
│ list[i64]      │
╞════════════════╡
│ [11, 21, … 51] │
│ [12, 22, … 52] │
│ [13, 23, … 53] │
└────────────────┘

Starting from the first element of each list, I want to divide every two elements of the list with a number, and then starting from the second element of the list, divide again every two elements with another number. For example, if these two numbers are 5 and 10 respectively, the first list will be transformed like this

[11/5, 21/10, 31/5, 41/10, 51/5]

resulting in

[2.2, 2.1, 6.2, 4.1, 10.2]

I want to do the same transformation for all the lists of the column. How can I do that using the polars API?

Solution

  • Use list.eval and when/then/otherwise

    import polars as pl

data = {"col": [[11, 21, 31, 41, 51], [12, 22, 32, 42, 52], [13, 23, 33, 43, 53]]}
df = pl.DataFrame(data)


def func(x):
    return pl.when(x.cum_count() % 2 == 0).then(x / 5).otherwise(x / 10)


df.select(result=pl.col("col").list.eval(func(pl.element())))

    shape: (3, 1)
┌────────────────────┐
│ result             │
│ ---                │
│ list[f64]          │
╞════════════════════╡
│ [2.2, 2.1, … 10.2] │
│ [2.4, 2.2, … 10.4] │
│ [2.6, 2.3, … 10.6] │
└────────────────────┘