如何在Python Polars中实现DataFrame子集赋值（对标Pandas操作）

Polars的DataFrame默认是不可变结构，无法像Pandas那样直接给筛选后的结果赋值，需要通过**条件表达式结合with_columns**或update方法实现相同逻辑。以下是对应原Pandas代码的Polars完整实现：

步骤1：创建Polars DataFrame #

import polars as pl

data = pl.DataFrame(
    {   
        "era": ["01", "01", "02", "02", "03", "10"],
        "pred1": [1, 2, 3, 4, 5,6],
        "pred2": [2,4,5,6,7,8],
        "pred3": [3,5,6,8,9,1],
        "something_else": [5,4,3,67,5,4],
    }
)
pred_cols = ["pred1", "pred2", "pred3"]
ERA_COL = "era"
DOWNSAMPLE_CROSS_VAL = 10

test_split = ['01', '02', '10']

步骤2：核心条件赋值实现 #

要替换era属于test_split的行的pred1列，推荐用when/then/otherwise表达式配合with_columns，这是Polars的标准 immutable 操作方式：

# 假设somefunction()返回Polars DataFrame，若返回Pandas可先转成Polars：pl.from_pandas(somefunction())
new_preds = somefunction()["another_column"]

# 生成新的DataFrame，仅修改符合条件的pred1值
data = data.with_columns(
    pl.when(pl.col(ERA_COL).is_in(test_split))
    .then(new_preds)
    .otherwise(pl.col("pred1"))
    .alias("pred1")
)

如果需要直接修改原DataFrame（Polars 0.19+支持），也可以用update结合掩码实现：

# 生成筛选掩码
test_split_mask = data.select(pl.col(ERA_COL).is_in(test_split)).to_series()
# 获取待替换的值
new_preds = somefunction()["another_column"]
# 直接更新原DataFrame的指定列
data.update(pl.col("pred1").set(new_preds, mask=test_split_mask))

补充：下采样训练集的实现 #

原Pandas代码中train_split_index未定义，推测为训练集（era不在test_split中的行），对应Polars下采样逻辑：

# 筛选训练集行
train_split_mask = ~pl.col(ERA_COL).is_in(test_split)
# 对训练集每10行取1行做下采样
downsampled_train = data.filter(train_split_mask).slice(0, None, DOWNSAMPLE_CROSS_VAL)
# 若需要获取下采样行的索引
downsampled_train_indices = downsampled_train.get_row_indices()

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