如何在pandas中用partition by模拟windows功能？

Question

我在original_eur列中得到了带有空值的数据。

	事件_id	类别	rounds_bot_date	original_eur
0	43ee085d-40b4-44f7-bdf5-c9e649284527	类别1	2024-03-25 00:00:00	200
1	bd9a68b1-c62d-4481-a9bd-01709ccdf237	类别1	2024-03-25 00:00:00	南
2	f905ad71-d199-435a-879f-7520604b698d	类别 2	2024-03-25 00:00:00	南
3	f905ad71-d199-435a-879f-7520604b698d	类别 2	2024-03-25 00:00:00	150
4	b1435b90-2fa6-4ec4-8b6c-29c630a59812	类别 2	2024-03-25 00:00:00	150
5	6ce6fbfb-7ddf-4023-9f96-442180559480	类别1	2024-03-25 00:00:00	南
6	bb4b995c-a478-46c7-acd6-9f39e1e48d2a	第 3 类	2024-03-25 00:00:00	南
7	39863eb6-6abd-492e-9563-860b58d6c352	类别 2	2024-03-25 00:00:00	150
8	8df861d4-4d8e-4da8-ba3f-35f1f25f9b26	第 3 类	2024-03-25 00:00:00	60
9	d54df0e4-f072-4099-ab64-65af7865a1e0	类别 2	2024-03-25 00:00:00	150

我需要将列中的每个空值替换为适当的 event_id、category、rounds_bot_date 的中值。

通过 SQL，我可以使用 case + 中值窗口函数：

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

对于熊猫，我用中位数制作表格：

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 中模拟中值窗函数？

P。 S. 我用相同的逻辑 make iterrows，但它不如 SQL 函数快。

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

注意： 使用 Pandas、Numpy、Polars 和类似的 Python 库时，请记住：

1. 我们99.99%的问题以前都被其他人遇到过并解决过，所以先查看API参考，你通常可以在其中找到现成的答案；
2. 自制函数的性能总是比内置函数差很多，所以请避免使用它们。如果必须的话，至少要避免其中出现循环。
3. 如果你熟悉 SQL，我推荐你使用 Polars，而不是 Pandas。相信你很快就会习惯的。