在PySpark中实现SAS的retain功能 分区同步rank=1的Rate字段值

你需要的按分区保留首行指定字段值的效果，可直接通过窗口聚合函数实现，两种可行实现方式如下：

方案1：基于已计算好的Rank字段实现（适配你当前的代码进度，逻辑更稳妥） #

不需要修改你已有的Rank字段计算逻辑，直接通过条件窗口聚合提取每个分区Rank=1对应的Rate值覆盖全分区即可：

from pyspark.sql import Window
from pyspark.sql.functions import col, when, max

# 定义仅按Prod_Code分区的窗口
window_part = Window.partitionBy("Prod_Code")

# 提取每个分区Rank=1的Rate值填充到全分区所有行
df_output = df.withColumn("Rate", max(when(col("Rank") == 1, col("Rate"))).over(window_part))

方案2：合并Rank计算和Rate填充逻辑（可减少一次窗口计算） #

如果不需要保留原始的Rate值，也可以将Rank计算和Rate填充合并，一步完成：

from pyspark.sql import Window
from pyspark.sql.functions import row_number, first

# 定义分区窗口，排序规则和你计算Rank时保持一致（示例中按Rate降序排列）
window_spec = Window.partitionBy("Prod_Code").orderBy(col("Rate").desc())

# 同时生成Rank字段和填充后的Rate字段
df_output = df.withColumn("Rank", row_number().over(window_spec))\
              .withColumn("Rate", first("Rate").over(window_spec))

两种方案执行后输出结果均和你给出的期望结果一致。

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