Pandas .transform传入自定义函数的执行顺序异常问题排查

问题背景 #

为确认pandas.core.groupby.GroupBy.transform传入的自定义函数接收的对象类型，编写测试代码后发现：函数执行顺序不符合预期，且缺少最后一条EXIT日志。

实验代码 #

import pandas as pd
import numpy as np
import threading

speeds = pd.DataFrame(
    [
        ("bird", "Falconiformes", 389.0),
        ("bird", "Psittaciformes", 24.0),
        ("mammal", "Carnivora", 80.2),
        ("mammal", "Primates", np.nan),
        ("mammal", "Carnivora", 58),
    ],
    index=["falcon", "parrot", "lion", "monkey", "leopard"],
    columns=("class", "order", "max_speed"),
)
enter_count = 0
exit_count = 0
def f(x):
    global enter_count, exit_count
    print(f"--- ENTER f {enter_count} --- {threading.get_ident()}")
    enter_count += 1
    print("RESULT:", type(x), x.dtype)
    print(f"--- EXIT f {exit_count} ---")
    exit_count += 1
    return x
speeds.groupby("class").transform(f)

实际输出 #

--- ENTER f 0 --- 6776
RESULT: <class 'pandas.core.series.Series'> object
--- EXIT f 0 ---
--- ENTER f 1 --- 6776
RESULT: <class 'pandas.core.series.Series'> float64
--- EXIT f 1 ---
--- ENTER f 2 --- 6776
--- ENTER f 3 --- 6776
RESULT: <class 'pandas.core.series.Series'> object
--- EXIT f 2 ---
--- ENTER f 4 --- 6776
RESULT: <class 'pandas.core.series.Series'> float64
--- EXIT f 3 ---

预期输出 #

--- ENTER f 0 --- 6776
...
--- EXIT f 0 --- 6776
--- ENTER f 1 --- 6776
...
--- EXIT f 1 --- 6776
--- ENTER f 2 --- 6776
...
--- EXIT f 2 --- 6776
--- ENTER f 3 --- 6776
...
--- EXIT f 3 --- 6776

原因分析 #

1. 输出缓冲导致顺序混乱 #

Python的print默认采用行缓冲机制，当Pandas内部执行逻辑与标准输出缓冲交互时，会出现打印语句显示顺序与代码实际执行顺序不一致的情况。单线程环境下函数不可能嵌套执行，看似“ENTER f2后直接进入ENTER f3”的现象，实际是前一次函数的RESULT和EXIT打印被延迟缓冲，后一次函数的ENTER先被输出。

2. 缺少最后一条EXIT日志 #

最后一次函数调用的EXIT打印内容仍留在输出缓冲区中，脚本执行结束时未被强制刷新。在print语句中添加flush=True参数即可解决，修改后所有日志会按执行顺序即时输出。

3. 函数调用次数多于预期 #

你的DataFrame有2个分组（bird、mammal），每个分组包含2列（order、max_speed），理论上应调用4次函数，但实际出现了5次调用。这是因为Pandas在内部处理transform时，针对含NaN的列或返回原数据的场景，会触发额外的内部验证或复制操作，导致多一次函数调用。

验证方案 #

修改函数中的print语句，添加flush=True强制刷新缓冲区：

def f(x):
    global enter_count, exit_count
    print(f"--- ENTER f {enter_count} --- {threading.get_ident()}", flush=True)
    enter_count += 1
    print("RESULT:", type(x), x.dtype, flush=True)
    print(f"--- EXIT f {exit_count} ---", flush=True)
    exit_count += 1
    return x

修改后输出顺序会与实际执行顺序一致，且所有EXIT日志都会正常显示。

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