PySpark分组条件应用：优先满足条件1则跳过条件2

原始DataFrame #

data = [
    ("1", "SAP", "A", 1, 51.55, "Team1"),
    ("2", "SAP", "B", 1, 51.55, "Team1"),
    ("3", "SAP", "B", 1, 51.55, "Team1"),
    ("4", "SAP", "A", 1, 55.35, "Team2"),
    ("5", "SAP", "B", 1, 55.35, "Team2"),
    ("6", "SAP", "C", 1, 58.00, "Team3"),
    ("7", "SAP", "D", 1, 47.00, "Team3")
]
df = spark.createDataFrame(data, ["ID", "Source", "Type", "Active", "Weight", "Group"])

需求说明 #

• 按Group字段分组处理main_record标记：
  1. 若组内存在满足条件1的行，仅将满足条件1的行标记为1，组内其他行无论是否符合条件2都标记为0
  2. 仅当组内无满足条件1的行时，才将满足条件2的行标记为1，其余行标记为0
• 条件1：(Type == "A") & (Active == 1) & (Source == "SAP")
• 条件2：(Rank == 1)（Rank为组内按Weight降序排列的排名）

原代码（结果不符合预期） #

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

condition1 = (F.col("Type") == "A") & (F.col("Active") == 1) & (F.col("Source") == "SAP")
condition2 = (F.col("Rank") == 1)

mainRecordWindow = Window.partitionBy("Group").orderBy(F.col("Weight").desc())
df = df.withColumn(
    "Rank", F.rank().over(mainRecordWindow)
)
df = df.withColumn(
    "main_record",
    F.when(condition1, 1).otherwise(F.when(condition2, 1).otherwise(0))
)

df.show()

预期结果 #

+--+------+----+-------+------+-----+----+-----------+
|ID|Source|Type|Active|Weight|Group|Rank|main_record|
+--+------+----+-------+------+-----+----+-----------+
| 1|   SAP|   A|     1| 51.55|Team1|   1|          1|
| 2|   SAP|   B|     1| 51.55|Team1|   1|          0|
| 3|   SAP|   B|     1| 51.55|Team1|   1|          0|
| 4|   SAP|   A|     1| 55.35|Team2|   1|          1|
| 5|   SAP|   B|     1| 55.35|Team2|   1|          0|
| 6|   SAP|   C|     1| 58.00|Team3|   1|          1|
| 7|   SAP|   D|     1| 47.00|Team3|   2|          0|
+--+------+----+-------+------+-----+----+-----------+

修正后的代码 #

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

condition1 = (F.col("Type") == "A") & (F.col("Active") == 1) & (F.col("Source") == "SAP")
condition2 = (F.col("Rank") == 1)

# 定义分组窗口，用于判断组内是否存在满足条件1的行
group_window = Window.partitionBy("Group")
# 定义排名窗口
rank_window = Window.partitionBy("Group").orderBy(F.col("Weight").desc())

df = df.withColumn(
    "has_condition1",
    F.max(F.when(condition1, 1).otherwise(0)).over(group_window)
).withColumn(
    "Rank",
    F.rank().over(rank_window)
).withColumn(
    "main_record",
    F.when(
        F.col("has_condition1") == 1,
        F.when(condition1, 1).otherwise(0)
    ).otherwise(
        F.when(condition2, 1).otherwise(0)
    )
)

df.show()

修正说明 #

原代码未判断整个组内是否存在满足条件1的行，导致组内满足条件2但不满足条件1的行也被错误标记为1。修正逻辑：

1. 新增has_condition1字段：通过组内窗口计算，标记当前组是否存在满足条件1的行（1=存在，0=不存在）
2. 计算main_record时优先依据has_condition1判断：
   • 若组内存在符合条件1的行，仅将满足条件1的行标记为1，其余全为0
   • 若组内无符合条件1的行，再按条件2判断标记

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