PySpark DataFrame计算用户每月首尾日余额的正确方法

问题描述 #

处理包含name、date、month和balance列的PySpark DataFrame时，需要为每个name+month分组的每行添加balance_on_first_day（当月首日余额）和balance_on_last_day（当月末日余额）列。使用first()和last()窗口函数计算后，部分行的新列出现NULL值，每个name-month组内总有一行的其中一个新列为NULL。

错误原因 #

• 原代码中window_last_day的rowsBetween(Window.unboundedFollowing, Window.unboundedFollowing)仅包含分组排序后的最后一行，因此只有最后一行能获取到last()的值，其余行均为NULL；
• 窗口范围设置不合理，导致部分行无法访问到全组的首尾数据。

正确实现代码 #

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

# 创建SparkSession与示例DataFrame
data = [
    ("John", "2023-01-05", 1, 1000),
    ("Alice", "2023-01-10", 1, 1200),
    ("John", "2023-02-15", 2, 1500),
    ("Alice", "2023-02-20", 2, 1800),
    ("John", "2023-03-01", 3, 2000),
    ("Alice", "2023-03-25", 3, 2200),
    ("John", "2023-03-31", 3, 2100),
]
columns = ["name", "date", "month", "balance"]
spark = SparkSession.builder.appName("BalanceCalculation").getOrCreate()
df = spark.createDataFrame(data, columns)

# 将date列转换为日期类型
df = df.withColumn('date', F.col('date').cast('date'))

# 定义覆盖整个分组的窗口：按name、month分区，按date排序以确定首尾日期
full_group_window = Window.partitionBy('name', 'month').orderBy('date').rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

# 计算当月首日余额：取分组内最早日期对应的balance
df = df.withColumn('balance_on_first_day', F.first('balance', ignorenulls=True).over(full_group_window))
# 计算当月末日余额：取分组内最晚日期对应的balance
df = df.withColumn('balance_on_last_day', F.last('balance', ignorenulls=True).over(full_group_window))

# 查看结果
df.show()

关键说明 #

• full_group_window通过rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)确保窗口覆盖整个分组，让每行都能访问到组内所有数据；
• first(..., ignorenulls=True)和last(..., ignorenulls=True)参数保证即使分组内存在NULL值，也能正确获取到首尾有效余额；
• 如果业务需要处理全组balance均为NULL的场景，可以在函数外嵌套F.coalesce(..., F.lit(0))来替换NULL为0。

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