如何在PySpark中按条件聚合行？附具体数据场景

可以用窗口函数一步搞定，思路直接明了：先算出每个客户所有location为空行的总购买量，再把这个值加到该客户所有非空location行的购买量上，最后过滤掉空行即可。

示例数据 #

假设原始DataFrame结构如下：

期望输出结果：

实现代码 #

from pyspark.sql import functions as F
from pyspark.sql.window import Window

# 定义窗口：按customer分组
window_spec = Window.partitionBy("customer")

result_df = (
    # 原始DataFrame
    df
    # 计算每个客户所有location为空行的总购买量
    .withColumn(
        "null_total",
        F.sum(F.when(F.col("location").isNull(), F.col("total_purchased")).otherwise(0)).over(window_spec)
    )
    # 将空行的总购买量加到对应客户的非空行上
    .withColumn(
        "total_purchased",
        F.col("total_purchased") + F.col("null_total")
    )
    # 过滤掉location为空的行
    .filter(F.col("location").isNotNull())
    # 删除临时计算列
    .drop("null_total")
)

关键逻辑说明 #

1. 窗口分组计算：通过partitionBy("customer")确保只在同一客户范围内计算空行的总购买量；
2. 条件求和：用when判断location是否为空，仅对空行的total_purchased求和；
3. 更新购买量：将原购买量与空行总购买量相加，得到最终的聚合后数值；
4. 过滤空行：移除所有location为空的记录，保留处理后的有效行。

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