如何在Polars中像Pandas一样有条件地设置多个元素？

Question

我试图根据条件在 Polars DataFrame 中设置多个元素，类似于 Pandas 中的操作方式。这是 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)

此更新列

A

，其中

A < B

为

[100, 210, 320]

。

在 Polars 中，我知道不可能就地更新 DataFrame，并且可以返回包含更新元素的新 DataFrame。我尝试过以下方法：

尝试 1：将

Series.scatter

与

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)
)

尝试 2：创建更新 DataFrame 并使用

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')

这两种方法都有效，但与简单的 Pandas 语法相比，它们感觉很麻烦。

问题：
Polars 中是否有更简单或更惯用的方法来有条件地在列中设置多个元素，类似于 Pandas

loc

语法？

Answer 1

pl.when(condition).then(value).otherwise(value) - is used for conditional logic to a column.

alias() is used to assign a name to the resulting column.

pl.col() to refer columns in the expressions.

您可以在 Polars 中尝试这种方法。 例如

df = pl.DataFrame({"foo": [1, 3, 4], "bar": [3, 4, 0]})    
df.with_columns(pl.when(pl.col("foo") > 2).then(1).otherwise(-1).alias("val"))

希望这对你有帮助