Question

Polars: Replace elements in list of List column

Consider the following example series.

s = pl.Series('s', [[1, 2, 3], [3, 4, 5]])

I'd like to replace all 3s with 10s to obtain the following.

res = pl.Series('s', [[1, 2, 10], [10, 4, 5]])

Is it possible to efficiently replace elements in the lists of a List column in polars?

Note. I've already tried converting to a dataframe and using pl.when().then(), but pl.when() fails for input of type List[bool]. Moreover, I've experimented with pl.Expr.list.eval, but couldn't get much further than the original mask.

 3  58  3
1 Jan 1970

Solution

 3

In the case of a simple replacement, you can also consider using pl.Expr.replace within a pl.Series.list.eval context. This approach is a bit less general, but more concise than a pl.when().then().otherwise() construct.

s.list.eval(pl.element().replace({3: 10}))
shape: (2,)
Series: 's' [list[i64]]
[
    [1, 2, 10]
    [10, 4, 5]
]
2024-07-24
Hericks

Solution

 2

Not sure if I understand your original question, but you could combine list.eval() and when()/ then()/otherwise() together for the replacement.

import polars as pl


s = pl.Series('s', [[1, 2, 3], [3, 4, 5]])

df = s.to_frame().with_columns(
    pl.col('s').list.eval(
        pl.when(pl.element() == 3).then(10).otherwise(pl.element())
    ).alias('s')
)

print(df)

it will output like

shape: (2, 1)
┌─────────────┐
│ s           │
│ ---         │
│ list[i64]   │
├─────────────┤
│ [1, 2, 10]  │
│ [10, 4, 5]  │
└─────────────┘
2024-07-24
linpingta