PySpark DataFrame处理：过滤单关联Idx并计算多记录Result列

需求说明 #

需要保留Idx对应多个Name的记录，同时新增Result列计算Value与Wight的乘积。这类操作属于分组统计过滤+列计算，核心是先识别出分组内记录数大于1的Idx，再完成后续计算。

实现代码 #

from pyspark.sql import SparkSession
from pyspark.sql.window import Window
from pyspark.sql.functions import col, count

# 初始化Spark会话
spark = SparkSession.builder.appName("FilterMultiNameIdx").getOrCreate()

# 构造原始数据
raw_data = [
    ("A", "Peter", 123, 0.1),
    ("B", "Peter", 234, 0.1),
    ("B", "Mary", 287, 0.1),
    ("B", "Jason", 243, 0.1),
    ("C", "Mary", 109, 0.1),
    ("D", "David", 231, 0.1),
    ("D", "Sean", 101, 0.1),
    ("D", "Tim", 238, 0.1),
    ("E", "Ken", 910, 0.1),
    ("E", "Jim", 112, 0.1)
]
df = spark.createDataFrame(raw_data, ["Idx", "Name", "Value", "Wight"])

# 定义窗口：按Idx分组
idx_window = Window.partitionBy("Idx")

# 执行过滤与计算
final_df = df.withColumn("name_count", count("Name").over(idx_window)) \
             .filter(col("name_count") > 1) \
             .withColumn("Result", col("Value") * col("Wight")) \
             .drop("name_count")

# 输出结果
final_df.show()

代码解释 #

1. 窗口函数统计分组数量：通过Window.partitionBy("Idx")将数据按Idx分组，用count("Name").over(idx_window)计算每个Idx对应的Name记录数，生成临时列name_count。
2. 过滤目标记录：筛选出name_count > 1的行，保留对应多个Name的Idx记录。
3. 新增计算列：直接通过col("Value") * col("Wight")计算乘积，生成Result列。
4. 清理临时列：删除用于过滤的name_count列，得到最终结果。

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