Polars中如何用单次group_by高效实现两级Max聚合的均值计算？

针对你需要的两级聚合逻辑——先按NE+DT分组取UT的最大值，再按NE分组对这些最大值取整体最大值和平均值——完全可以通过单次group_by结合窗口函数实现，避免两次分组带来的额外数据扫描，尤其适合数百万行的大表场景。

实现代码 #

import polars as pl
df = pl.DataFrame(
    {
        'NE' : ["P1", "P1", "P1", "P1", "P2", "P2", "P2", "P2", "P3", "P3", "P3", "P3" ],
        'DT' : ["D1", "D1", "D2", "D2", "D1", "D1", "D2", "D2", "D1", "D1", "D2", "D2" ],
        'TM' : ["H1", "H2", "H1", "H2", "H1", "H2", "H1", "H2", "H1", "H2", "H1", "H2" ],
        'UT' : [ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 ]
    }
)

# 单次group_by实现两级聚合
result = df.group_by('NE').agg(
    UT_max=pl.col('UT').max().over('DT').max(),
    UT_avg=pl.col('UT').max().over('DT').mean()
)

print(result)

运行结果 #

shape: (3, 3)
┌─────┬────────┬────────┐
│ NE  ┆ UT_max ┆ UT_avg │
│ --- ┆ ---    ┆ ---    │
│ str ┆ i64    ┆ f64    │
╞═════╪════════╪════════╡
│ P2  ┆ 8      ┆ 7.0    │
│ P3  ┆ 12     ┆ 11.0   │
│ P1  ┆ 4      ┆ 3.0    │
└─────┴────────┴────────┘

原理说明 #

1. 窗口函数over('DT')：在每个NE分组内部，先计算出每个DT对应的UT最大值（等价于原逻辑中第一级group_by('NE','DT').agg(pl.max('UT'))的结果）。
2. 二次聚合：直接对窗口函数得到的DT级最大值，计算max（即UT_max）和mean（即UT_avg），完成第二级聚合。

这种方式只需要对数据进行一次扫描，相比两次group_by的方案，在大表场景下能显著减少计算开销。

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