如何在Polars中像Pandas的.loc一样按条件批量设置元素？

我完全理解你想要的那种简洁感——Pandas的loc语法确实直观，Polars虽然不支持原地修改，但其实有更符合它设计风格的简洁写法，不用像你尝试的两种方法那样繁琐。

最简洁的表达式写法：直接用Series.scatter在with_columns中 #

Polars的表达式API其实支持直接在with_columns里调用scatter方法，不需要额外用map_batches或者手动处理索引。你可以这样写：

import polars as pl

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

# 直接在表达式中完成条件筛选和赋值
df_updated = df.with_columns(
    pl.col("A").scatter(
        # 找到所有A < B的行的位置索引
        pl.col("A").lt(pl.col("B")).arg_true(),
        # 对应位置要赋的新值
        [100, 210, 320]
    ).alias("A")
)

print(df_updated)

这个写法和你Pandas的思路完全对应：先定位A < B的行的位置（用arg_true()获取符合条件的索引），再用scatter把指定值覆盖到这些位置上，最后返回新的DataFrame。

为什么这个方法更合适？ #

• 它完全利用了Polars的表达式优化器，性能比map_batches这种自定义函数写法更高（自定义函数会绕过Polars的内部优化）；
• 代码更紧凑，逻辑和Pandas的loc操作一一对应，读起来非常直观；
• 不需要额外处理索引或者创建中间DataFrame，避免了你Attempt 2里的冗余步骤。

补充：另一种条件分支思路 #

如果你更倾向于用条件判断的写法，也可以通过when/then/otherwise结合序列生成来实现：

# 先获取条件掩码，确保新值列表长度匹配
mask = df.select(pl.col("A").lt(pl.col("B"))).to_series()
new_values = pl.Series([100, 210, 320])

df_updated = df.with_columns(
    pl.when(mask)
    .then(new_values.take(pl.arange(0, mask.sum())))
    .otherwise(pl.col("A"))
    .alias("A")
)

不过这种写法比scatter稍显繁琐，更适合需要动态生成新值的场景。

对你之前尝试的小优化 #

如果还是想用你Attempt 2的思路，其实可以不用手动添加index列，但确实不如scatter的写法直接高效。

总结一下，最符合Polars风格且简洁的写法就是直接在with_columns里用scatter表达式，既保留了Pandas那种直观的条件赋值逻辑，又符合Polars不可变DataFrame的设计理念。

备注：内容来源于stack exchange，提问作者HYRY