SQL Partition By窗口聚合逻辑转等效Pandas Python代码咨询

需求说明 #

将以下带窗口函数的SQL查询转换为等效的Pandas代码，输出结果和SQL执行结果完全一致：

select 
a.*, 
b.vol1 / sum(vol1) over (
  partition by a.sale, a.d_id, 
  a.month, a.p_id
) vol_r, 
a.vol2* b.vol1/ sum(b.vol1) over (
  partition by a.sale, a.d_id, 
  a.month, a.p_id
) vol_t
from 
sales1 a 
left join sales2 b on a.sale = b.sale
and a.d_id = b.d_id 
and a.month = b.month
and a.p_id = b.p_id

实现代码 #

标准实现（可读性优先） #

import pandas as pd

# 1. 执行左连接，完全对齐SQL的LEFT JOIN逻辑
merge_df = pd.merge(
    left=sales1,
    right=sales2,
    on=["sale", "d_id", "month", "p_id"],
    how="left"
)

# 2. 计算每个分区的vol1总和，等效SQL的SUM(vol1) OVER(PARTITION BY ...)
# groupby+transform会将聚合结果广播到分组内的每一行，完全匹配窗口函数行为
partition_sum = merge_df.groupby(["sale", "d_id", "month", "p_id"])["vol1"].transform("sum")

# 3. 计算两个目标字段
merge_df["vol_r"] = merge_df["vol1"] / partition_sum
merge_df["vol_t"] = merge_df["vol2"] * merge_df["vol1"] / partition_sum

# 4. 可选：输出字段和SQL完全对齐（仅保留sales1的所有字段 + 两个计算字段）
final_df = merge_df[sales1.columns.to_list() + ["vol_r", "vol_t"]]

内存优化实现（大数据量场景适用） #

如果处理的数据集较大，可以省略中间列存储，进一步降低内存开销：

import pandas as pd

merge_df = pd.merge(sales1, sales2, on=["sale", "d_id", "month", "p_id"], how="left")
merge_df = merge_df.assign(
    partition_sum = lambda x: x.groupby(["sale", "d_id", "month", "p_id"])["vol1"].transform("sum"),
    vol_r = lambda x: x["vol1"] / x["partition_sum"],
    vol_t = lambda x: x["vol2"] * x["vol1"] / x["partition_sum"]
).drop(columns="partition_sum") # 计算完成后直接删除中间聚合列

注意事项 #

• 该实现和SQL行为完全对齐：左连接匹配失败的行vol1为NaN，后续计算结果也为NaN，和SQL左连接的NULL运算逻辑一致
• 如果需要处理除零场景，可以在计算时使用replace或fillna自定义异常值，比如merge_df["vol_r"] = (merge_df["vol1"] / partition_sum).fillna(0)

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