如何在Spark中按userid和counterparty_iban分组实现类似.shift()的薪资列？

我想在Spark DataFrame里实现类似Python的.shift()函数的效果，目前已有初步尝试，但需要正确的分组逻辑才能得到预期结果。

输入DataFrame #

预期输出 #

需要生成包含previous_salary列的DataFrame，该列显示同一userid和counterparty_iban分组内、按month排序后的上一行amount值，分组内第一行该值为null：

创建输入DataFrame的代码 #

columns = ["month", "userid", 'exactoriginalamount', 'counterparty_iban']
data = [("1", "John Jones", "2400", 'ENG12345'),
        ("4", "John Jones", "1200", 'ENG12345'),
        ("8", "John Jones", "2300", 'ENG12345'),
        ("5", "John Jones", "3000", 'AM12345'),
        ("9", "John Jones", "5000", 'AM12345'),
        ("12", "John Jones", "6000", 'AM12345'),
        ("1", "Joe Jones", "1200", 'AM12345'),
        ("2", "Joe Jones", "2400", 'AM12345'),
        ("3", "Joe Jones", "5000", 'AM12345')]

df = spark.createDataFrame(data=data, schema=columns)

我的尝试（未正确分组） #

w = Window().partitionBy().orderBy(F.col('userid'))
df = df.withColumn('previous_salary', F.lag('exactoriginalamount', 1).over(w))

解决方案 #

你需要调整Window的partitionBy和orderBy逻辑：

1. 分组依据：按userid和counterparty_iban分区，确保同一用户、同一交易对手的记录归为一组；
2. 排序规则：将month转为数值类型后按升序排列，避免字符串排序导致的顺序错误；
3. 使用lag函数获取分组内上一行的exactoriginalamount值，偏移量设为1，分组首行会自动返回null。

完整代码如下：

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

# 将month转为整数类型，修正排序逻辑
df = df.withColumn("month", F.col("month").cast("int"))

# 定义Window规范：按userid和counterparty_iban分区，按month升序排序
w = Window.partitionBy("userid", "counterparty_iban").orderBy("month")

# 添加previous_salary列
df = df.withColumn("previous_salary", F.lag("exactoriginalamount", 1).over(w))

# 查看结果
df.show()

执行后即可得到预期输出，previous_salary会正确显示同一分组内上一条记录的金额值，分组首行显示null。

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