PySpark提取指定月份value全为空值的账户数据问题

问题分析 #

现有代码仅筛选了mn=1/2/3且单条记录value为Null的行，但未校验该账户所有1-3月的value是否全部为Null，导致错误包含了accnt=2的mn=3记录（该账户mn=1/2的value不为Null）。同时原代码使用isnan()判断Null存在逻辑问题：isnan()仅针对数值型NaN，而原始数据中的Null是缺失值，应使用isNull()判断。

解决方案（窗口函数实现） #

通过窗口函数按账户分组，统计每个账户在1-3月中value不为Null的记录数，筛选出统计数为0的账户（即1-3月value全为Null），再提取这些账户的1-3月数据。

from pyspark.sql import Window
from pyspark.sql.functions import col, count, when

# 原始数据初始化
df_data = {'accnt': [1,1,1,1,2,2,2,2,3,3,3,3],
           'mn': [1,2,3,4,1,2,3,4,1,2,3,4], 
           'value': [None,None,None,5,3,2,None,4, None, None, None, None],
           'val2': [5,5,5,5,4.5,4.5,4.5,4.5, None, None, None, None],
          }

import pandas as pd
df_pandas = pd.DataFrame.from_dict(df_data)
df = spark_session.createDataFrame(df_pandas)

# 定义窗口：按账户分组，确保每个账户的统计独立
window_spec = Window.partitionBy("accnt")

# 计算每个账户在1-3月中value不为Null的记录数
df_with_flag = df.withColumn(
    "non_null_count",
    count(when((col("mn").isin(1,2,3)) & col("value").isNotNull(), 1)).over(window_spec)
)

# 筛选符合条件的记录：1-3月value全为Null（non_null_count=0）且mn在1-3范围内
df_final = df_with_flag.where(
    (col("non_null_count") == 0) & (col("mn").isin(1,2,3))
).select("accnt", "mn", "value", "val2")

# 查看结果
df_final.show()

代码说明 #

1. 窗口定义：Window.partitionBy("accnt")确保统计逻辑按账户独立计算。
2. 非Null计数：count(when(...))仅统计1-3月中value不为Null的记录数，若结果为0，说明该账户这三个月的value全为Null。
3. 最终筛选：同时满足non_null_count=0和mn在1-3的记录，即为目标数据。

运行结果 #

+-----+---+-----+----+
|accnt| mn|value|val2|
+-----+---+-----+----+
|    1|  1| null|   5|
|    1|  2| null|   5|
|    1|  3| null|   5|
|    3|  1| null|null|
|    3|  2| null|null|
|    3|  3| null|null|
+-----+---+-----+----+

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