Spark Scala DataFrame合并相同主键行并聚合is_deleted字段

我有如下DataFrame：

| notification_id|       el1|       el2|is_deleted| 
+---------------+----------+----------+----------+ 
|notificationId1|element1_1|element1_2|     false| 
|notificationId2|element2_1|element2_2|     false| 
|notificationId3|element3_1|element3_2|     false| 
|notificationId1|      null|      null|      true| 
|notificationId4|      null|      null|      true| 
+---------------+----------+----------+----------+

主键为notification_id，且满足以下规则：

• 当is_deleted = true时，除主键外其他列均为null值；
• 当is_deleted = false时，主键具有唯一性。

希望合并具有相同主键的行，得到合并is_deleted字段后的目标DataFrame：

| notification_id|       el1|       el2|is_deleted| 
+---------------+----------+----------+----------+ 
|notificationId1|element1_1|element1_2|      true| 
|notificationId2|element2_1|element2_2|     false| 
|notificationId3|element3_1|element3_2|     false| 
|notificationId4|      null|      null|      true| 
+---------------+----------+----------+----------+

解决方案 #

Pandas 实现 #

基于给定规则，每个主键组内最多存在一条非空业务数据（is_deleted=false），通过分组聚合提取非空字段，并判断组内是否存在is_deleted=true的记录即可：

import pandas as pd

# 构造原始DataFrame
df = pd.DataFrame({
    'notification_id': ['notificationId1', 'notificationId2', 'notificationId3', 'notificationId1', 'notificationId4'],
    'el1': ['element1_1', 'element2_1', 'element3_1', None, None],
    'el2': ['element1_2', 'element2_2', 'element3_2', None, None],
    'is_deleted': [False, False, False, True, True]
})

# 分组合并：取业务字段第一个非空值，is_deleted只要组内有True则返回True
result_df = df.groupby('notification_id').agg(
    el1=('el1', lambda x: x.dropna().iloc[0] if not x.dropna().empty else None),
    el2=('el2', lambda x: x.dropna().iloc[0] if not x.dropna().empty else None),
    is_deleted=('is_deleted', 'any')
).reset_index()

print(result_df)

PySpark 实现 #

利用Spark的分组聚合函数，指定忽略null值提取业务字段，通过max函数判断is_deleted状态（布尔值中True大于False，只要组内存在True就会返回True）：

from pyspark.sql import SparkSession
from pyspark.sql.functions import first, max

spark = SparkSession.builder.appName("merge_deleted_records").getOrCreate()

# 构造原始DataFrame
data = [
    ("notificationId1", "element1_1", "element1_2", False),
    ("notificationId2", "element2_1", "element2_2", False),
    ("notificationId3", "element3_1", "element3_2", False),
    ("notificationId1", None, None, True),
    ("notificationId4", None, None, True)
]
df = spark.createDataFrame(data, ["notification_id", "el1", "el2", "is_deleted"])

# 分组合并
result_df = df.groupBy("notification_id").agg(
    first("el1", ignoreNulls=True).alias("el1"),
    first("el2", ignoreNulls=True).alias("el2"),
    max("is_deleted").alias("is_deleted")
)

result_df.show()

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