PySpark无UDF实现：用rangeBetween窗口获取前序最近行product_id

需求说明 #

针对每个product_name，为每行数据新增ref_id列，取值为当前行之前purchase_date最接近的行的product_id。

原方案问题分析 #

你尝试的方法通过rangeBetween和max函数获取最近前置日期，但这种逻辑只能拿到日期值，无法直接关联到该日期对应的product_id，没办法完成从日期到目标ID的映射，所以无法得到正确结果。

正确实现方案 #

直接使用Spark窗口函数中的lag即可解决——按purchase_date排序后，当前行的前一行就是时间上最接近的前置记录，正好匹配需求。

基础版代码 #

from pyspark.sql import Window
from pyspark.sql.functions import col, lag

# 定义窗口：按产品名称分区，按购买日期升序排序
window = Window.partitionBy("product_name").orderBy(col("purchase_date").cast("timestamp"))

# 生成ref_id列：取窗口中前一行的product_id，无前置记录时为null
df = df.withColumn("ref_id", lag("product_id", 1).over(window))

特殊场景处理（同产品同购买日期） #

如果同一product_name下存在多条purchase_date完全相同的记录，可根据业务需求调整：

• 若同日期记录不算"之前"的记录，保持上述代码即可，同日期行不会互相引用
• 若需要在同日期内也获取前置行的product_id，可增加product_id作为辅助排序字段：

window = Window.partitionBy("product_name").orderBy(col("purchase_date").cast("timestamp"), "product_id")
df = df.withColumn("ref_id", lag("product_id", 1).over(window))

方案优势 #

相比你原有的方法，lag函数直接定位到目标行的字段值，无需额外创建临时时间戳列，逻辑更简洁高效，完全满足无UDF的要求。

内容的提问来源于stack exchange，提问作者Neda Ah.