如何更高效地为Polars DataFrame按组填充首元素至新列？

给定字典：

group_dict = {'doc_7': 1, 'doc_8': 2, 'doc_9': 2, 'doc_10': 2, 'doc_11': 5, 'doc_12': 1}

转换为Polars DataFrame：

doc_group_df = pl.from_records(list(group_dict.items()), orient='row', schema=['doc', 'group'])

得到的DataFrame结构：

shape: (6, 2)
┌────────┬───────┐
│ doc    ┆ group │
│ ---    ┆ ---   │
│ str    ┆ i64   │
╞════════╪═══════╡
│ doc_7  ┆ 1     │
│ doc_8  ┆ 2     │
│ doc_9  ┆ 2     │
│ doc_10 ┆ 2     │
│ doc_11 ┆ 5     │
│ doc_12 ┆ 1     │
└────────┴───────┘

需要添加original_doc列，实现效果：

shape: (6, 3)
┌────────┬───────┬──────────────┐
│ doc    ┆ group ┆ original_doc │
│ ---    ┆ ---   ┆ ---          │
│ str    ┆ i64   ┆ str          │
╞════════╪═══════╪══════════════╡
│ doc_7  ┆ 1     ┆              │
│ doc_8  ┆ 2     ┆              │
│ doc_9  ┆ 2     ┆ doc_8        │
│ doc_10 ┆ 2     ┆ doc_8        │
│ doc_11 ┆ 5     ┆              │
│ doc_12 ┆ 1     ┆ doc_7        │
└────────┴───────┴──────────────┘

当前实现代码：

def original_doc(row):
    original_doc = sorted(doc_group_df.filter(pl.col('group')==row['group'])['doc'].to_list())[0]
    return original_doc if original_doc != row['doc'] else ''

doc_group_df.with_columns(pl.struct(pl.all()).map_elements(original_doc).alias('original_doc'))

询问是否有更优实现方式。

你当前用map_elements逐行处理的方式效率较低，尤其是数据量大时，因为它会触发Python层面的循环。可以用Polars原生的分组窗口函数实现，性能更优：

doc_group_df.with_columns(
    # 按group分组，取每组排序后的第一个doc作为基准
    pl.col("doc").sort().first().over("group").alias("original_doc")
).with_columns(
    # 当前doc等于基准值时设为空字符串，否则保留基准值
    pl.when(pl.col("doc") == pl.col("original_doc"))
      .then(pl.lit(""))
      .otherwise(pl.col("original_doc"))
      .alias("original_doc")
)

代码说明 #

1. 第一步用over("group")窗口函数，对每个group分组后排序doc并取第一个值，得到每组的基准原始文档。
2. 第二步用when/then/otherwise条件判断，将当前doc等于基准值的行设为空字符串，其他行保留基准值，完全匹配需求。

这种方式利用Polars的向量化操作，避免了Python循环，数据规模越大，性能优势越明显。

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