Python Polars - add element to columns of lists which has value equal to a function of the list

Question

if I have a columns of lists, for example:

df = pl.DataFrame({'lists': [[1,2,3], [2,3,4], [3,4,5]]})

And I want to add max(list)+2 to each list as an additional element, to get:

expected_df = pl.DataFrame({'lists': [[1,2,3,5], [2,3,4,6], [3,4,5,7]]})

How do I do this in polars?

I'd have thought it would be something like:

expected_df.with_columns(pl.col('lists').list.add_element(pl.col('lists').list.max()+2))

Answer 1

You can do:

>>> df.with_columns(pl.col('lists').list.concat(pl.col('lists').list.max() + 2))
shape: (3, 1)
┌──────────────┐
│ lists        │
│ ---          │
│ list[i64]    │
╞══════════════╡
│ [1, 2, 3, 5] │
│ [2, 3, 4, 6] │
│ [3, 4, 5, 7] │
└──────────────┘

Let's make sure that the expected result is obtained:

>>> expected_df = pl.DataFrame({'lists': [[1,2,3,5], [2,3,4,6], [3,4,5,7]]})
>>> df.with_columns(pl.col('lists').list.concat(pl.col('lists').list.max() + 2)).equals(expected_df)
True