[Python] How to set multiple elements conditionally in Polars similar to Pandas?

I am trying to set multiple elements in a Polars DataFrame based on a condition, similar to how it is done in Pandas. Here’s an example in Pandas:

import pandas as pd

df = pd.DataFrame(dict(
A=[1, 2, 3, 4, 5],
B=[0, 5, 9, 2, 10],
))

df.loc[df['A'] < df['B'], 'A'] = [100, 210, 320]
print(df)


This updates column A where A < B with [100, 210, 320].

In Polars, I know that updating a DataFrame in place is not possible, and it is fine to return a new DataFrame with the updated elements. I have tried the following methods:

Attempt 1: Using Series.scatter with map_batches


import polars as pl

df = pl.DataFrame(dict(
A=[1, 2, 3, 4, 5],
B=[0, 5, 9, 2, 10],
))

def set_elements(cols):
a, b = cols
return a.scatter((a < b).arg_true(), [100, 210, 320])

df = df.with_columns(
pl.map_batches(['A', 'B'], set_elements)
)

Attempt 2: Creating an update DataFrame and using update()


df = df.with_row_index()
df_update = df.filter(pl.col('A') < pl.col('B')).select(
'index',
pl.Series('A', [100, 210, 320])
)
df = df.update(df_update, on='index').drop('index')


Both approaches work, but they feel cumbersome compared to the straightforward Pandas syntax.

Question:
Is there a simpler or more idiomatic way in Polars to set multiple elements conditionally in a column, similar to the Pandas loc syntax?

Continue reading...