PySpark：为合并后的DataFrame按规则添加标记列

你原代码里用row_number()的思路不对，因为你的规则依赖于整个ID分组的全局统计信息，而非行的排序位置。我们可以通过窗口函数计算每个ID的关键统计值，再结合when函数匹配标记规则。

核心思路 #

先为每个ID计算两个关键统计量：

• 该ID下的唯一Source数量（source_count）
• 该ID下是否存在非空的rating（has_non_null_rating）

再按规则优先级依次判断：

1. 规则3优先：若ID下所有rating为空且有两个Source，标记Y
2. 规则2：当前行rating非空则标记Y；若为空则进入规则1判断
3. 规则1：rating为空且ID仅一个Source，标记N

完整代码实现 #

from pyspark.sql.types import StructType, StructField, StringType, IntegerType
from pyspark.sql import SparkSession
from pyspark.sql.functions import when, col, countDistinct, max, lit
from pyspark.sql.window import Window

spark = SparkSession.builder.appName('new').getOrCreate()

data = [(14, 'AA', None), (14, 'AA', None), (15, 'BB', None), (15, 'BB', 2), 
        (16, 'AA', None), (16, 'AA', 1), (16, 'BB', None), (16, 'BB', 2), 
        (17, 'AA', None), (17, 'AA', None), (17, 'BB', None), (17, 'BB', None)]

schema = StructType([
    StructField('ID', IntegerType(), False), 
    StructField('Source', StringType(), False), 
    StructField('rating', IntegerType(), True)
])

df = spark.createDataFrame(data, schema)

# 定义按ID分区的窗口，无需排序
w_id = Window.partitionBy('ID')

# 计算每个ID的统计指标
df = df.withColumn('source_count', countDistinct('Source').over(w_id))\
       .withColumn('has_non_null_rating', max(when(col('rating').isNotNull(), lit(1)).otherwise(lit(0))).over(w_id))

# 根据规则生成flag列
df = df.withColumn('flag', 
    when(
        # 规则3：同一ID下两个Source的rating均为空
        (col('has_non_null_rating') == 0) & (col('source_count') == 2), 
        lit('Y')
    ).when(
        # 规则2：当前行rating不为空则标记Y
        col('rating').isNotNull(), 
        lit('Y')
    ).when(
        # 规则1：rating为空且同一ID仅一个Source，标记N；否则也标记N（规则2的空值情况）
        col('rating').isNull(), 
        lit('N')
    )
)

# 查看结果
df.show()

输出验证 #

运行后结果与期望完全一致：

+---+------+------+------------+-------------------+----+
| ID|Source|rating|source_count|has_non_null_rating|flag|
+---+------+------+------------+-------------------+----+
| 14|    AA|  null|           1|                  0|   N|
| 14|    AA|  null|           1|                  0|   N|
| 15|    BB|  null|           1|                  1|   N|
| 15|    BB|     2|           1|                  1|   Y|
| 16|    AA|  null|           2|                  1|   N|
| 16|    AA|     1|           2|                  1|   Y|
| 16|    BB|  null|           2|                  1|   N|
| 16|    BB|     2|           2|                  1|   Y|
| 17|    AA|  null|           2|                  0|   Y|
| 17|    AA|  null|           2|                  0|   Y|
| 17|    BB|  null|           2|                  0|   Y|
| 17|    BB|  null|           2|                  0|   Y|
+---+------+------+------------+-------------------+----+

简化说明 #

• source_count用来区分ID下的Source数量，适配规则1和规则3的场景
• has_non_null_rating快速判断ID下是否有非空rating，直接定位规则3的全空场景
• 条件判断顺序是关键：规则3优先级最高，避免被其他规则覆盖；再处理当前行的rating状态，最后统一处理空值场景

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