PySpark基于分组列删除Target列空值并保留全空分组实现

实现思路 #

• 用窗口函数按Group列分区，统计每个分组下非空Target的计数，计数大于0说明该分组存在有效Target值，否则分组内所有Target均为空
• 过滤逻辑完全忽略ID列的影响：
  • 分组存在有效Target时，仅保留Target不为空的行
  • 分组全为空Target时，保留该分组全部行

完整实现代码 #

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

# 初始化SparkSession
spark = SparkSession.builder.appName("filter_null_by_group").getOrCreate()

# 构造示例数据集
data = [
    ("A", "B", 10),
    ("A", "B", 10),
    ("A", "B", 10),
    ("A", "C", None),
    ("A", "C", None),
    ("A", "C", None),
    ("B", "D", None),
    ("B", "D", None),
    ("B", "D", None)
]
df = spark.createDataFrame(data, schema=["Group", "ID", "Target"])

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

# 新增列统计每个Group下非空Target的数量
df = df.withColumn("non_null_target_cnt", F.count(F.col("Target")).over(window_spec))

# 按规则过滤后删除统计用的临时列
result_df = df.filter(
    (F.col("non_null_target_cnt") == 0) | 
    (F.col("Target").isNotNull())
).drop("non_null_target_cnt")

# 输出结果
result_df.show()

输出验证 #

运行代码后输出结果和预期完全一致：

+-----+---+------+
|Group| ID|Target|
+-----+---+------+
|    A|  B|    10|
|    A|  B|    10|
|    A|  B|    10|
|    B|  D|  null|
|    B|  D|  null|
|    B|  D|  null|
+-----+---+------+

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