如何在Polars DataFrame的指定行索引范围更新列？

需求与Pandas实现 #

需要更新Polars DataFrame中指定列的特定行索引范围，对应的Pandas实现代码如下：

import pandas as pd
df = pd.DataFrame({ "foo": [0,0,0,0] })
df["foo"].iloc[0:3] = 1
# 或者另一种写法
df.iloc[0:3, df.columns.get_loc("foo")] = 1

Polars直接切片赋值的问题 #

在Polars中可以通过df[0, "foo"] = 1更新单行，但尝试范围切片赋值会触发错误：

import polars as pl
df = pl.DataFrame({ "foo": [0,0,0,0] })

# 报错：TypeError: cannot use "slice(0, 3, None)" for indexing
df[0:3, "foo"] = 1

# 报错：TypeError: DataFrame object does not support `Series` assignment by index
df[0:3]["foo"] = 1

高效替代方案（避免row_number的性能开销） #

方法1：构造新Series替换目标列 #

直接获取原列数据，切片更新后替换，适合小范围更新场景：

updated_foo = df["foo"].to_list()
updated_foo[0:3] = [1]*3
df = df.with_columns(pl.Series(updated_foo).alias("foo"))

方法2：用pl.int_range替代row_number #

利用pl.int_range直接生成连续行索引序列，避免全表扫描计算行号的开销：

df = df.with_columns(
    pl.when(pl.int_range(0, pl.count()).is_between(0, 3))
    .then(1)
    .otherwise(pl.col("foo"))
    .alias("foo")
)

方法3：切片拆分+拼接（大表最优） #

通过切片拆分DataFrame，更新指定段后重新拼接，Polars的切片和拼接操作基本是零拷贝，性能接近Pandas的iloc赋值：

# 拆分需要更新的部分和剩余部分
top_part = df.slice(0, 3).with_columns(foo=1)
bottom_part = df.slice(3)
# 拼接得到新的DataFrame
df = pl.concat([top_part, bottom_part])

内容的提问来源于stack exchange，提问作者Oreille