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

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').arr.add_element(pl.col('lists').arr.max() 2))

CodePudding user response：

You can do

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

CodePudding user response：

For this you can use the apply method along with a lambda function.

df = pd.DataFrame({'lists': [[1,2,3], [2,3,4], [3,4,5]]})
df['lists'] = df['lists'].apply(lambda x: x   [max(x)   2])

Or you can use a built-in function, such as max like this. It will give the same output as previous example but its more readable.

df = pd.DataFrame({'lists': [[1,2,3], [2,3,4], [3,4,5]]})
df['lists'] = [[*lst, max(lst)   2] for lst in df['lists']]