Polars分组聚合：按Attr1+Attr2分组时取Attr1维度Val1最大值

针对6000万行的大型Polars DataFrame，要实现按Attr1+Attr2分组统计行数，同时为每个分组附加对应Attr1维度下Val1的全局最大值，最优方案是分两步聚合再关联，结合Polars的Lazy API实现低内存、高速度的处理：

核心思路 #

1. 先单独按Attr1分组，计算每个Attr1对应的Val1全局最大值——这一步生成的数据集远小于原始数据，仅包含每个Attr1的唯一值和对应最大值。
2. 再按Attr1+Attr2分组，统计每组的行数。
3. 将两个聚合结果按Attr1关联，得到最终所需数据。

这种方式避免了窗口函数对全量数据的逐行计算，转而在小数据集上做关联操作，大幅降低内存占用并提升执行效率。

代码实现（Lazy API版本） #

import polars as pl

# 以CSV为例，用scan_csv加载为LazyFrame（避免一次性加载全量数据到内存）
df = pl.scan_csv("your_large_data.csv")

# 步骤1：计算每个Attr1的Val1全局最大值
attr1_val1_max = df.group_by("Attr1").agg(
    pl.col("Val1").max().alias("Attr1_Val1_Global_Max")
)

# 步骤2：计算Attr1+Attr2分组的行数统计
group_counts = df.group_by("Attr1", "Attr2").agg(
    pl.count().alias("Group_Count")
)

# 步骤3：关联两个结果集
final_result = group_counts.join(attr1_val1_max, on="Attr1", how="inner")

# 执行计算并获取结果（如需写入文件可直接用write_csv等方法，无需collect）
final_result.collect()

为什么这是最优解？ #

• 内存效率：两次分组聚合的结果行数远低于原始6000万行，关联操作仅在小数据集上进行，内存占用可控。
• 性能优势：Polars的分组聚合基于Rust实现，性能远超纯Python库；Lazy API会自动优化执行计划（比如投影下推、谓词下推），减少不必要的数据加载和计算。
• 对比窗口函数：如果用pl.col("Val1").max().over("Attr1")实现，虽然代码更简洁，但需要对全量6000万行逐行计算窗口值，内存开销和执行时间都会显著增加，不适用于超大型数据集。

注意事项 #

• 始终优先使用Lazy API（scan_*系列方法）加载数据，避免一次性将全量数据读入内存。
• 如果原始数据存在缺失值，可根据需求添加pl.col("Val1").max().drop_nulls()确保最大值计算准确。
• 若需写入结果，直接对final_result调用write_csv等方法，无需先collect，进一步降低内存压力。

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