如何在Pandas中模拟带Partition By的Median窗口函数？

DevOps 小满

2026-6-27

问题描述

我有一份数据集，其中original_eur列存在空值（NaN），数据如下：

	event_id	category	rounds_bot_date	original_eur
0	1	Category 1	2024-03-25 00:00:00	200
1	1	Category 1	2024-03-25 00:00:00	nan
2	2	Category 2	2024-03-25 00:00:00	nan
3	2	Category 2	2024-03-25 00:00:00	150
4	2	Category 2	2024-03-25 00:00:00	150
5	2	Category 1	2024-03-25 00:00:00	nan
6	3	Category 3	2024-03-25 00:00:00	nan
7	3	Category 2	2024-03-25 00:00:00	150
8	3	Category 3	2024-03-25 00:00:00	60
9	3	Category 2	2024-03-25 00:00:00	150

需要将original_eur列中的每个空值替换为对应event_id、category、rounds_bot_date分组的中位数。

在SQL中，可以用CASE语句结合MEDIAN窗口函数实现：

case
    when original_eur = NaN
    then median(original_eur) over(partition by event_id, category, rounds_bot_date)
    else original_eur
end as original_eur

我自己在Pandas中先生成中位数表：

median_table = (
    dataset
    .groupby(['event_id', 'category', 'rounds_bot_date'])
    .agg(original_eur_median = ('original_eur', 'median'))
    .reset_index()
)

再通过自定义函数填充空值：

import math

def fill_na(value, event_id, category, rounds_bot_date, median_table):
    if math.isnan(value):
        value = (
            median_table[
                (median_table['event_id'] == event_id) & 
                (median_table['category'] == category) & 
                (median_table['rounds_bot_date'] == rounds_bot_date)]['original_eur_median'].values[0]
        )
        return value
    else:
        return value

dataset['original_eur'] = (
    dataset
    .apply(
        lambda x: fill_na(x['original_eur'], x['event_id'], x['category'], x['rounds_bot_date'], median_table),
        axis = 1)
)

但这种方式效率不高，用iterrows实现逻辑速度也不如SQL，请问有没有办法优化代码，在Pandas中模拟MEDIAN窗口函数的效果？

解决方案

可以利用Pandas的transform方法直接生成分组中位数列，再填充空值，代码如下：

# 添加分组中位数列
dataset['original_eur_median'] = (
    dataset
    .groupby(['event_id', 'category', 'rounds_bot_date'])['original_eur']
    .transform('median')
)
# 用中位数填充空值
dataset['original_eur'] = dataset['original_eur'].fillna(dataset['original_eur_median'])

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