Python Polars：如何高效聚合函数返回的DataFrame

核心方案：使用groupby.map_groups #

Polars的groupby.map_groups是处理这类分组生成DataFrame场景的最优选择，它直接对每个分组应用函数并自动合并结果，全程保持矢量化操作，避免了行迭代的性能损耗。

改进后的代码 #

import polars as pl

def dummy_dataframe(val):
    return pl.DataFrame(
        {
            "x": [val + 1, 0, 1],
            "y": [val + 2, 8, 9],
            "z": [val + 3, 5, 6],
        }
    )

df = pl.DataFrame(
    {
        "a": ["a", "b", "a", "b"],
    }
)

# 使用map_groups处理每个分组
df_results = df.groupby('a').map_groups(
    lambda group: dummy_dataframe(group['a'].unique()[0]).with_columns(group_name=group['a'].first())
)

print(df_results)

为什么这比原方法高效？ #

• 原代码中iter_rows()会将列存储的数据转换为行式结构，完全违背Polars的设计优势，带来额外的性能开销。
• map_groups是Polars原生的分组处理API，内部基于列存储进行优化，不需要手动迭代和拼接结果，全程保持数据的矢量化处理。

补充说明 #

如果你的dummy_dataframe函数可以接受批量输入（比如直接处理整个分组的列数据），还可以进一步优化，比如直接利用Polars的矢量化计算生成结果，避免在函数内创建小DataFrame的开销。但针对当前示例的场景，map_groups已经是最优的解决方案。

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