如何在Polars DataFrame中基于过滤与分组操作添加新列（Pandas迁移场景）

首先，你的核心需求是把每个sid对应的wm_mask的voxel_count提取出来，作为新列wm_mask__count添加到原DataFrame中。Polars可以通过筛选+左连接的方式轻松实现，比Pandas的索引操作更简洁直观。

完整实现代码 #

import polars as pl

# 原数据集
df = pl.from_repr("""
┌────────────┬─────────┬───────────┬───────┬───────┬─────────────┐
│ sid        ┆ roi     ┆ endpoint  ┆ value ┆ std   ┆ voxel_count │
│ ---        ┆ ---     ┆ ---       ┆ ---   ┆ ---   ┆ ---         │
│ str        ┆ str     ┆ str       ┆ f64   ┆ f64   ┆ i64         │
╞════════════╪═════════╪═══════════╪═══════╪═══════╪═════════════╡
│ 4213-a3_bl ┆ AF_L    ┆ afd_along ┆ 0.4   ┆ 0.21  ┆ 57334       │
│ 4213-a3_bl ┆ AF_L    ┆ radfODF   ┆ 0.08  ┆ 0.045 ┆ 57334       │
│ 4213-a3_bl ┆ AF_R    ┆ afd_along ┆ 0.42  ┆ 0.22  ┆ 53916       │
│ 4213-a3_bl ┆ AF_R    ┆ radfODF   ┆ 0.08  ┆ 0.04  ┆ 53916       │
│ 4213-a3_bl ┆ CC_1    ┆ afd_along ┆ null  ┆ null  ┆ null        │
│ 4213-a3_bl ┆ CC_2a   ┆ afd_along ┆ 0.54  ┆ 0.3   ┆ 3264        │
│ 4213-a3_bl ┆ wm_mask ┆ null      ┆ null  ┆ null  ┆ 602620      │
│ 4225-a3_bl ┆ CC_2a   ┆ radfODF   ┆ 0.06  ┆ 0.04  ┆ 3264        │
│ 4225-a3_bl ┆ CC_2b   ┆ afd_along ┆ 0.47  ┆ 0.24  ┆ 18833       │
│ 4225-a3_bl ┆ wm_mask ┆ null      ┆ null  ┆ null  ┆ 718758      │
└────────────┴─────────┴───────────┴───────┴───────┴─────────────┘
""")

# 步骤1：提取每个sid对应的wm_mask的voxel_count，并重命名列
wm_counts = df.filter(pl.col("roi") == "wm_mask").select(
    "sid",
    pl.col("voxel_count").alias("wm_mask__count")
)

# 步骤2：将提取的结果左连接到原DataFrame，按sid匹配
result_df = df.join(wm_counts, on="sid", how="left")

# 查看结果
print(result_df)

代码解释 #

1. 筛选提取：用filter选出roi等于wm_mask的行，然后通过select只保留sid和重命名后的wm_mask__count列——因为每个sid只有一条wm_mask记录，不需要额外分组取第一行（如果有多行的话，可以加.group_by("sid").agg(pl.col("voxel_count").first().alias("wm_mask__count"))来确保取第一行）。
2. 左连接：用join操作将提取的wm_counts与原DataFrame按sid关联，how="left"保证原DataFrame的所有行都被保留，同时匹配到对应的wm_mask__count值。

输出结果 #

运行后会得到带新列的完整DataFrame，前几行示例如下：

shape: (10, 7)
┌────────────┬─────────┬───────────┬───────┬───────┬─────────────┬────────────────┐
│ sid        ┆ roi     ┆ endpoint  ┆ value ┆ std   ┆ voxel_count ┆ wm_mask__count │
│ ---        ┆ ---     ┆ ---       ┆ ---   ┆ ---   ┆ ---         ┆ ---            │
│ str        ┆ str     ┆ str       ┆ f64   ┆ f64   ┆ i64         ┆ i64            │
╞════════════╪═════════╪═══════════╪═══════╪═══════╪═════════════╪════════════════╡
│ 4213-a3_bl ┆ AF_L    ┆ afd_along ┆ 0.4   ┆ 0.21  ┆ 57334       ┆ 602620         │
│ 4213-a3_bl ┆ AF_L    ┆ radfODF   ┆ 0.08  ┆ 0.045 ┆ 57334       ┆ 602620         │
│ 4213-a3_bl ┆ AF_R    ┆ afd_along ┆ 0.42  ┆ 0.22  ┆ 53916       ┆ 602620         │
...

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