PySpark DataFrame分组内多行列值比较并赋值status字段

原始DataFrame #

+--------+--------+----------------+--------------+
| company|      id|        first_nm|        sur_nm|
+--------+--------+----------------+--------------+
|SYNTHE01|SYNTHE02|           JAMES|        FOWLER|
|SYNTHE01|SYNTHE03|          MONICA|        FOWLER|
|SYNTHE01|SYNTHE04|          GEORGE|        FOWLER|
|SYNTHE08|SYNTHE05|           JAMES|        FIWLER|
|SYNTHE08|SYNTHE06|           JAMES|        FUWLER|
|SYNTHE08|SYNTHE07|           JAMES|        FAWLER|
|SYNTHE08|SYNTHE08|           JAMES|        FEWLER|
|SYNTHE11|SYNTHE12|           JAMES|        FOWLER|
|SYNTHE11|SYNTHE11|           JAMES|        FOWLER|
|SYNTHE09|SYNTHE0X|            Null|          Null|
|SYNTHE09|SYNTHE0Y|            Null|          Null|
|SYNTHE09|SYNTHE0Z|            Null|          Null|
+--------+--------+----------------+--------------+

Status赋值规则 #

• 若某company组内所有行的first_nm与sur_nm均匹配（组内两列各自只有唯一非Null值），status为0；
• 若仅组内所有行的first_nm匹配（组内first_nm唯一非Null，sur_nm有多个不同值），status为1；
• 若仅组内所有行的sur_nm匹配（组内sur_nm唯一非Null，first_nm有多个不同值），status为2；
• 若无匹配或列值为Null，status为99。

预期输出DataFrame #

+--------+--------+----------------+--------------+-------+
| company|      id|        first_nm|        sur_nm| status|
+--------+--------+----------------+--------------+-------+
|SYNTHE01|SYNTHE02|           JAMES|        FOWLER|      2|
|SYNTHE01|SYNTHE03|          MONICA|        FOWLER|      2|
|SYNTHE01|SYNTHE04|          GEORGE|        FOWLER|      2|
|SYNTHE08|SYNTHE05|           JAMES|        FIWLER|      1|
|SYNTHE08|SYNTHE06|           JAMES|        FUWLER|      1|
|SYNTHE08|SYNTHE07|           JAMES|        FAWLER|      1|
|SYNTHE08|SYNTHE08|           JAMES|        FEWLER|      1|
|SYNTHE11|SYNTHE12|           JAMES|        FOWLER|      0|
|SYNTHE11|SYNTHE11|           JAMES|        FOWLER|      0|
|SYNTHE09|SYNTHE0X|            Null|          Null|     99|
|SYNTHE09|SYNTHE0Y|            Null|          Null|     99|
|SYNTHE09|SYNTHE0Z|            Null|          Null|     99|
+--------+--------+----------------+--------------+-------+

实现代码与说明 #

步骤1：导入依赖并创建测试DataFrame #

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, countDistinct, when, sum

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

# 创建测试数据
data = [
    ("SYNTHE01", "SYNTHE02", "JAMES", "FOWLER"),
    ("SYNTHE01", "SYNTHE03", "MONICA", "FOWLER"),
    ("SYNTHE01", "SYNTHE04", "GEORGE", "FOWLER"),
    ("SYNTHE08", "SYNTHE05", "JAMES", "FIWLER"),
    ("SYNTHE08", "SYNTHE06", "JAMES", "FUWLER"),
    ("SYNTHE08", "SYNTHE07", "JAMES", "FAWLER"),
    ("SYNTHE08", "SYNTHE08", "JAMES", "FEWLER"),
    ("SYNTHE11", "SYNTHE12", "JAMES", "FOWLER"),
    ("SYNTHE11", "SYNTHE11", "JAMES", "FOWLER"),
    ("SYNTHE09", "SYNTHE0X", None, None),
    ("SYNTHE09", "SYNTHE0Y", None, None),
    ("SYNTHE09", "SYNTHE0Z", None, None)
]

df = spark.createDataFrame(data, ["company", "id", "first_nm", "sur_nm"])

步骤2：计算分组统计指标 #

按company分组，计算以下核心指标：

• first_distinct_count：组内非Null的first_nm去重数量
• sur_distinct_count：组内非Null的sur_nm去重数量
• has_null：组内是否存在first_nm或sur_nm为Null的行

group_stats = df.groupBy("company")\
    .agg(
        countDistinct(when(col("first_nm").isNotNull(), col("first_nm"))).alias("first_distinct_count"),
        countDistinct(when(col("sur_nm").isNotNull(), col("sur_nm"))).alias("sur_distinct_count"),
        (sum(when(col("first_nm").isNull() | col("sur_nm").isNull(), 1).otherwise(0)) > 0).alias("has_null")
    )

步骤3：关联统计指标并赋值status #

将分组统计结果关联回原DataFrame，通过条件判断为每行赋值status：

result_df = df.join(group_stats, on="company", how="left")\
    .withColumn("status", 
        when(col("has_null"), 99)
        .when((col("first_distinct_count") == 1) & (col("sur_distinct_count") == 1), 0)
        .when(col("first_distinct_count") == 1, 1)
        .when(col("sur_distinct_count") == 1, 2)
        .otherwise(99)
    )\
    .drop("first_distinct_count", "sur_distinct_count", "has_null")

# 查看最终结果
result_df.show()

逻辑说明 #

1. 优先判断组内是否存在Null值，若存在直接赋值99；
2. 若组内first_nm和sur_nm都只有唯一非Null值，赋值0；
3. 若仅first_nm唯一非Null，赋值1；
4. 若仅sur_nm唯一非Null，赋值2；
5. 其他情况（如两列都有多个不同值）赋值99。

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