如何使用Polars创建列汇总重复文件路径对应的磁盘信息

我一直找不到这个问题的答案，因为很难准确描述需求。我想聚合多个磁盘上的文件列表，其中部分文件重复。希望每个文件只保留一行数据，新增一列记录该文件所在的所有磁盘。

示例DataFrame #

import polars as pl

recordings = pl.DataFrame(
    {
        "disk": ["NT23", "NT24", "NT23", "NT24"],
        "path_on_disk": ["file_a.txt", "file_a.txt", "file_b.txt", "file_b.txt"],
        "other_data": [2.0, 2.0, 3.0, 3.0],
    }
)

示例数据结构 #

┌──────┬──────────────┬────────────┐
│ disk ┆ path_on_disk ┆ other_data │
│ ---  ┆ ---          ┆ ---        │
│ str  ┆ str          ┆ f64        │
╞══════╪══════════════╪════════════╡
│ NT23 ┆ file_a.txt   ┆ 2.0        │
│ NT24 ┆ file_a.txt   ┆ 2.0        │
│ NT23 ┆ file_b.txt   ┆ 3.0        │
│ NT24 ┆ file_b.txt   ┆ 3.0        │
└──────┴──────────────┴────────────┘

实际场景说明 #

• 实际DataFrame包含超过50万行数据
• 实际DataFrame包含更多列，当path_on_disk相同时，除disk列外其余列数据均一致

需求 #

1. 找出所有path_on_disk相同的行
2. 创建新列disks，将对应不同的disk值用, .join()连接起来

期望结果 #

┌──────────────┬────────────┬────────────┐
│ path_on_disk ┆ disks      ┆ other_data │
│ ---          ┆ ---        ┆ ---        │
│ str          ┆ str        ┆ f64        │
╞══════════════╪════════════╪════════════╡
│ file_a.txt   ┆ NT23, NT24 ┆ 2.0        │
│ file_b.txt   ┆ NT23, NT24 ┆ 3.0        │
└──────────────┴────────────┴────────────┘

尝试过的代码及错误 #

我知道可以用recordings.group_by("path_on_disk")完成第一步：

for group_df in recordings.group_by("path_on_disk"):
    if len(group_df) > 1:
        print(group_df)
        break

这段代码会输出第一个path_on_disk重复的分组。

我尝试了以下代码，但出现错误：

def merge_disknames(df: pl.DataFrame): 
    return ", ".join(sorted(df["disk"]))

recordings.group_by("path_on_disk").map_groups(merge_disknames).rename("disks")

错误信息：

PanicException: Could not get DataFrame attribute '_df'. Make sure that you return a DataFrame object.: PyErr { type: <class 'AttributeError'>, value: AttributeError("'str' object has no attribute '_df'"), traceback: None }

解决方案 #

错误原因 #

Polars的map_groups要求处理函数必须返回Polars DataFrame对象，但你的函数返回的是字符串，导致Polars无法识别，从而抛出错误。

方案1：Polars原生聚合（高效适配大数据量） #

对于50万行的数据集，优先使用Polars内置的向量化聚合操作，效率远高于逐组遍历：

import polars as pl

result = recordings.group_by("path_on_disk").agg(
    # 对disk列排序后用", "拼接
    disks=pl.col("disk").sort().str.concat(", "),
    # 其他列取第一个值（因同path下数据一致）
    other_data=pl.col("other_data").first()
)

# 如果有更多列，比如col1、col2，只需扩展agg参数：
# result = recordings.group_by("path_on_disk").agg(
#     disks=pl.col("disk").sort().str.concat(", "),
#     other_data=pl.col("other_data").first(),
#     col1=pl.col("col1").first(),
#     col2=pl.col("col2").first()
# )

方案2：修复map_groups的写法 #

如果坚持要使用map_groups，需确保函数返回Polars DataFrame：

import polars as pl

def merge_disknames(df: pl.DataFrame):
    # 拼接排序后的磁盘名称
    disks_str = ", ".join(sorted(df["disk"].to_list()))
    # 获取分组的path_on_disk值
    path = df["path_on_disk"][0]
    # 获取其他列的第一个值（同path下数据一致）
    other_cols = df.select(pl.exclude(["disk", "path_on_disk"])).row(0)
    # 构造并返回新的DataFrame
    return pl.DataFrame({
        "path_on_disk": [path],
        "disks": [disks_str],
        **dict(zip(df.columns.drop(["disk", "path_on_disk"]), other_cols))
    })

result = recordings.group_by("path_on_disk").map_groups(merge_disknames)

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