使用PySpark修改重复行状态值并保留首条记录的实现方法

直接使用窗口函数即可实现需求，无需拆分多数据集关联，执行效率更高，代码逻辑更简洁。

实现思路 #

• 按(account_id, user_id)分组开窗，统计每组内按id升序排序的行号，识别每个组合的首条记录
• 按user_id分组开窗，取每个user_id最早出现的account_id值，用于后续判断是否为同用户不同账号的重复场景
• 基于上述两个窗口的计算结果，直接判断每行的status值即可

代码示例 #

from pyspark.sql import SparkSession
from pyspark.sql.functions import row_number, first, when, col
from pyspark.sql.window import Window

# 初始化SparkSession
spark = SparkSession.builder.appName("dup_status_calc").getOrCreate()

# 构造样例数据
data = [
    (1, "a", 1), (2, "a", 1), (3, "b", 1),
    (4, "c", 2), (5, "c", 2), (6, "c", 2),
    (7, "d", 3), (8, "d", 3), (9, "e", 3)
]
df = spark.createDataFrame(data, schema=["id", "account_id", "user_id"])

# 定义窗口
# 窗口1：按account_id+user_id分组，id升序排序
w_account_user = Window.partitionBy("account_id", "user_id").orderBy("id")
# 窗口2：按user_id分组，取该user_id第一个出现的account_id
w_user = Window.partitionBy("user_id").orderBy("id").rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)

# 计算status
df_with_status = df.withColumn("rn", row_number().over(w_account_user)) \
    .withColumn("first_account", first("account_id").over(w_user)) \
    .withColumn("status",
        when(col("rn") == 1, "A")
        .when(col("account_id") == col("first_account"), "E")
        .otherwise("I")
    ) \
    .drop("rn", "first_account") # 清除中间计算字段

# 查看结果
df_with_status.orderBy("id").show()

输出结果验证 #

执行上述代码后输出结果和预期完全一致：

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