如何在Polars中快速存储groupBy分组数据？单机器提速方案

针对你4900多万行的DataFrame分组写入慢的问题，给出以下可落地的优化点：

1. 提前批量创建目录，减少IO判断开销 #

原代码每次循环都判断目录是否存在，频繁的文件系统检查会拖慢速度。可以先提取所有唯一的codeID，批量创建不存在的目录：

import os
from pathlib import Path

# 获取所有唯一codeID
unique_codes = df["codeID"].unique().to_list()
# 批量创建目录（exist_ok=True自动跳过已存在的目录）
for code in unique_codes:
    Path(code).mkdir(exist_ok=True)

2. 使用Polars Lazy API流式处理分组，降低内存占用 #

原循环会把整个分组数据加载到内存后再写入，改用Lazy模式可以流式处理分组，避免一次性占用大量内存，同时Polars会自动优化执行流程：

(
    df.lazy()
    .group_by("codeID")
    .apply(lambda group: group.write_ipc(f"{group['codeID'][0]}/{trade_date}.feather"))
    .collect()
)

3. 优化Feather/IPC写入参数，减少CPU与转换开销 #

Polars的write_ipc支持几个参数可以提升写入速度：

• 关闭压缩：如果不需要压缩，设置compression=None，减少CPU运算开销
• 使用Polars原生写入：设置use_pyarrow=False，避免PyArrow的中间转换损耗
  修改后的写入代码：

data.write_ipc(f"{code}/{trade_date}.feather", compression=None, use_pyarrow=False)

4. 用Pathlib替代字符串路径拼接 #

字符串拼接路径不仅容易出错，效率也不如pathlib的操作系统原生路径操作：

from pathlib import Path

file_path = Path(code) / f"{trade_date}.feather"
data.write_ipc(file_path, compression=None, use_pyarrow=False)

5. 禁用并行分组避免内存过载 #

如果之前尝试并行导致内存不足，可以强制Polars用单线程分组，避免内存暴涨：

import polars as pl

pl.set_option("dataframe.group_by.parallel", False)

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