如何在pandas中模拟带有partition by的窗口函数？

Question

我在原始数据的original_eur这一列里发现了一些空值（Null）。

	事件ID	类别	日期	原始欧元
0	1	类别 1	2024-03-25 00:00:00	200
1	1	类别 1	2024-03-25 00:00:00	nan
2	2	类别 2	2024-03-25 00:00:00	nan
3	2	类别 2	2024-03-25 00:00:00	150
4	2	类别 2	2024-03-25 00:00:00	150
5	2	类别 1	2024-03-25 00:00:00	nan
6	3	类别 3	2024-03-25 00:00:00	nan
7	3	类别 2	2024-03-25 00:00:00	150
8	3	类别 3	2024-03-25 00:00:00	60
9	3	类别 2	2024-03-25 00:00:00	150

我需要用每个事件ID、类别和日期对应的中位数值来替换这些空值。

在SQL中，我可以使用条件语句加上中位数窗口函数来实现这个。

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()
)

然后把这个函数应用到数据集中：

def fill_na(value, event_id, category, rounds_bot_date, median_table: pd.DataFrame):
    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)
)

有没有办法优化这个代码，并在Pandas中模拟中位数窗口函数呢？

附注：我可以用相同的逻辑做迭代，但速度没有SQL函数快。

解决方案：

# add new column with median values
dataset['original_eur_median'] = (
    dataset
    .groupby(['event_id', 'category', 'rounds_bot_date'])['original_eur']
    .transform('median')
)
# fill NaN with median values.
dataset['original_eur'] = dataset['original_eur'].fillna(dataset['original_eur_median'])

Answer 1

因为你的代码有点复杂，而且没有显示理想的输出，我觉得你可以用 pandas.DataFrame.groupby().transform() 来完成这个任务。试试这个：

df['fixed_eur'] = df.groupby(['event_id', 'category', 'rounds_bot_date'])['original_eur'].transform(lambda x: x.fillna(x.median()))

注意：

1. 我们99.99%的问题都是别人遇到过并解决过的，所以首先查查API文档，通常可以找到现成的答案；
2. 自己写的函数通常比内置的慢很多，所以尽量避免使用。如果必须用，至少要避免在里面使用循环。
3. 如果你对SQL比较熟悉，我推荐你使用Polars，而不是Pandas。我相信你会很快适应的。