Scala中基于行填充Spark DataFrame缺失值的最优方案

原始数据 #

预期结果 #

优化实现方式 #

方法一：窗口函数（简洁直观） #

利用每个test_number仅存在一行有效数据的特性，使用first()窗口函数并设置ignorenulls=True，直接提取分组内的有效值：

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

# 定义按test_number分组的窗口
window_spec = Window.partitionBy("test_number")

# 批量填充缺失值与替换0值
filled_df = df.withColumn(
    "test_name",
    F.first("test_name", ignorenulls=True).over(window_spec)
).withColumn(
    "value_1",
    F.first("value_1", ignorenulls=True).over(window_spec)
).withColumn(
    "value_2",
    F.first("value_2", ignorenulls=True).over(window_spec)
)

filled_df.show()

优势：

• 逻辑直接贴合需求，代码简洁易维护
• 无需额外聚合或关联操作，减少数据shuffle次数
• 同时适配字符串类型的test_name和数值类型的value字段

方法二：聚合后关联（适合大数据量场景） #

先按test_number聚合提取有效数据，再与原表关联完成填充，适合超大规模数据集：

# 聚合每个test_number的有效数据
agg_df = df.filter(
    (F.col("test_name").isNotNull()) & (F.col("value_1") != 0)
).select(
    "test_number",
    F.col("test_name").alias("filled_test_name"),
    F.col("value_1").alias("filled_value_1"),
    F.col("value_2").alias("filled_value_2")
)

# 关联原表完成填充
filled_df = df.join(
    agg_df,
    on="test_number",
    how="left"
).drop("test_name", "value_1", "value_2").rename(
    {"filled_test_name": "test_name", "filled_value_1": "value_1", "filled_value_2": "value_2"}
)

filled_df.show()

优势：

• 聚合后数据量远小于原表，关联操作的资源开销更低
• 避免窗口函数对全表的逐行计算，适合数据量极大的场景

对比原思路 #

你提到的max()窗口函数对数值型value可行，但test_name用collect_list需要额外取非空值，不如first()直接。上述两种方法都比原思路更简洁高效，尤其是方法一完全匹配当前场景的设计规则。

内容的提问来源于stack exchange，提问作者Tamás Godányi