PySpark：按唯一ID分组，对满足PurchasePrice>0的首次出现行赋值1

核心思路 #

通过窗口函数定位每个ID下首次满足PurchasePrice>0的记录位置，直接匹配生成目标列即可，无需额外做分组关联，代码简洁且执行性能更好。

具体实现代码 #

首先导入依赖：

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

执行列计算：

# 定义按ID分区、按时间升序排序的窗口
window_spec = Window.partitionBy("ID").orderBy("MonYer")

df = df.withColumn("row_num", F.row_number().over(window_spec)) \
       # 计算每个ID下首次PurchasePrice>0的记录对应的行号
       .withColumn("first_valid_row", F.min(F.when(F.col("PurchasePrice") > 0, F.col("row_num"))).over(Window.partitionBy("ID"))) \
       # 匹配生成目标列
       .withColumn("Wanted Column", F.when(F.col("row_num") == F.col("first_valid_row"), 1).otherwise(0)) \
       # 清理临时计算列
       .drop("row_num", "first_valid_row")

逻辑说明 #

• 先给每个ID下的所有记录按MonYer升序分配唯一的行号row_num
• 用min + when组合，筛选出每个ID里PurchasePrice>0的记录对应的最小行号，也就是首次出现符合条件记录的位置
• 只有行号和该最小行号相等的记录赋值为1，其余所有记录（包括首次符合条件记录之前、之后的所有行）均赋值为0
• 若某个ID下从来没有出现过PurchasePrice>0的记录，first_valid_row为null，该ID下所有记录的Wanted Column均为0，符合常规业务预期。

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