如何为Python Polars的map_groups、map_elements添加进度条？

问题背景 #

你想要为Polars中使用自定义函数的map_groups循环添加进度条，示例代码如下：

import polars as pl

df = pl.DataFrame(
    {
        "team": ["A", "A", "A", "B", "B", "C"],
        "conference": ["East", "East", "East", "West", "West", "East"],
        "points": [11, 8, 10, 6, 6, 5],
        "rebounds": [7, 7, 6, 9, 12, 8]
    }
)

df.group_by("team").map_groups(lambda x: x.select(pl.col("points").mean()))

你通过工具函数实现了进度条逻辑，但遇到两个问题：一是控制台中进度条直接跳至100%，二是函数本身有效但在PyCharm中无法正常显示进度条。

解决方案 #

1. 优化进度条实现代码

调整代码，移除不必要的全局变量，优化进度条刷新逻辑：

import polars as pl
from rich.progress import Progress

def pl_progress_applier(func, task_id, progress, **kwargs):
    progress.update(task_id, advance=1)
    return func(**kwargs)

def pl_groupby_progress_apply(data, group_by, func, drop_cols=[], **kwargs):
    with Progress() as progress:
        num_groups = len(data.select(group_by).unique())
        task_id = progress.add_task('处理分组中...', total=num_groups)
        return (
            data
                .group_by(group_by)
                .map_groups(lambda x: pl_progress_applier(
                    x=x.drop(drop_cols), func=func, task_id=task_id, progress=progress, **kwargs)
                )
        )

def custom_func(x):
    return x.select(pl.col('points').mean())

# 调用示例
df = pl.DataFrame(
    {
        "team": ["A", "A", "A", "B", "B", "C"],
        "conference": ["East", "East", "East", "West", "West", "East"],
        "points": [11, 8, 10, 6, 6, 5],
        "rebounds": [7, 7, 6, 9, 12, 8]
    }
)

result = pl_groupby_progress_apply(
    data=df,
    group_by='team',
    func=custom_func
)

print(result)

• 移除global progress：with块内的progress可被lambda直接捕获，避免全局变量带来的耦合问题。
• 去掉refresh=True：rich的Progress组件会自动处理刷新，强制刷新可能降低性能。
• 修改进度条提示文本为中文，提升可读性。

2. 解决PyCharm中进度条不显示的问题

问题根源是PyCharm默认控制台不支持交互式终端的ANSI转义序列，而rich动态进度条依赖该特性。开启终端模拟即可解决：

1. 点击PyCharm右上角运行配置下拉框，选择Edit Configurations...
2. 在配置窗口中找到当前的Python运行配置
3. 切换到Execution标签页，勾选Emulate terminal in output console
4. 点击Apply和OK，重新运行代码

此时PyCharm会模拟终端环境，进度条就能正常动态显示。

3. 额外说明

如果在系统原生终端（Windows CMD/PowerShell、Linux/macOS Terminal）中运行代码，优化后的代码无需额外配置就能正常显示进度条，问题仅存在于PyCharm的默认控制台环境。

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