如何用PySpark实现分组后仅保留单个指定Name的通用函数？

核心思路 #

• 先定位出分组内存在多个不同姓名的组
• 将这些组与原表关联，筛选出这些组的所有记录
• 最后仅保留这些组里指定姓名的所有记录

可复用函数代码 #

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

def keep_specified_name_in_groups(df: DataFrame, group_cols: list, name_col: str, target_name: str) -> DataFrame:
    # 筛选出分组内不同姓名数量超过1的组
    problematic_groups = df.groupBy(group_cols)\
                           .agg(F.countDistinct(name_col).alias("distinct_name_count"))\
                           .filter(F.col("distinct_name_count") > 1)\
                           .select(group_cols)
    
    # 关联原表，只保留目标分组里指定姓名的记录
    result = df.join(problematic_groups, on=group_cols, how="inner")\
               .filter(F.col(name_col) == target_name)
    
    return result

使用示例 #

假设你的原始数据存在df里，分组列是job、age、gender，姓名列是name，要保留的是Mike的记录，直接调用函数即可：

filtered_df = keep_specified_name_in_groups(df, ['job','age','gender'], 'name', 'Mike')
filtered_df.show()

参数说明 #

• df: 要处理的PySpark DataFrame
• group_cols: 分组用的列名列表，比如['job','age','gender']
• name_col: 存储姓名的列名，比如'name'
• target_name: 需要保留的目标姓名，比如'Mike'

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