能否用Polars优化CPU密集型Pandas分组聚合代码？

原Pandas实现 #

这段代码用于按A、B、C列分组，合并重复行并对其他列取最后一个非空值：

import numpy as np
import pandas as pd

def last_non_null(s):
    return s.dropna().iloc[-1] if not s.dropna().empty else np.nan

def merge_rows_of_final_df(df_final):
    # Group by columns A, B, and C
    cols = ['A', 'B', 'C']

    # Apply function to each column not in the subset
    # NOTE: the other columns have np.float32 as dtype
    agg_dict = {col: last_non_null for col in df_final.columns.difference(cols)}

    # Group and aggregate
    return df_final.groupby(cols).agg(agg_dict).reset_index()

# 示例调用
df_merged = merge_rows_of_final_df(df_final)

示例输入输出 #

示例输入

df = pd.DataFrame({
    "A":[1,1,1,1,4,4,4],
    "B":[2,2,2,2,5,5,5],
    "C":[3,3,3,3,6,6,6],
    "D":[4, np.nan, np.nan, np.nan, 1, np.nan, np.nan],
    "E":[np.nan, np.nan, np.nan, 7, np.nan, np.nan, 2],
    "F":[np.nan, np.nan, np.nan, np.nan, np.nan, np.nan, np.nan],
})

示例输出

A   B   C   D   E   F
0   1   2   3   4.0 7.0 NaN
1   4   5   6   1.0 2.0 NaN

Polars优化实现 #

Polars可以通过内置的last_valid函数高效实现该需求，无需自定义UDF，大幅降低CPU开销：

import polars as pl

def merge_rows_of_final_df_pl(df_final_pl):
    cols = ['A', 'B', 'C']
    # 对非分组列应用last_valid聚合
    agg_exprs = [
        pl.col(col).last_valid().alias(col)
        for col in df_final_pl.columns if col not in cols
    ]
    # 分组聚合后保持原顺序
    return df_final_pl.group_by(cols, maintain_order=True).agg(agg_exprs)

# 示例调用：若输入是Pandas DataFrame，先转Polars格式
df_pl = pl.from_pandas(df)
df_merged_pl = merge_rows_of_final_df_pl(df_pl)

# 如需转回Pandas格式
df_merged_pl.to_pandas()

优化说明 #

• 无自定义UDF开销：Polars的last_valid是内置矢量化函数，比Pandas自定义UDF的执行效率高得多，大数据集场景下优势更明显
• 内存与并行优势：Polars采用列式存储与延迟执行机制，内存占用更低，天然支持并行处理
• 类型自动保持：自动保留原列的数据类型（包括float32），无需额外处理

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