PySpark窗口函数处理含NaN数值列时聚合操作异常排查与解决

问题场景 #

给定如下Spark DataFrame：

import pandas as pd
from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()
sdf = spark.createDataFrame(pd.DataFrame([[1, 'name_value', None, 'square', 'polygon'],
                                          [2, 'name_value', 5.5, None, 'polygon']],
                                          columns=['id', 'name', 'area', 'sub_type', 'type']))

需求是按name列去重，为每个name分组选取各字段的非空最优值（优先取非空值）。使用窗口函数实现时，发现字符串列（如sub_type）用first()、collect_list()能正常拿到非空值，但float类型的area列得到的结果是NaN，与预期不符。

异常原因分析 #

核心问题并非float类型本身的问题，而是Spark窗口函数中first()的默认逻辑：

• first()函数在未指定窗口排序规则时，会取分组内物理存储顺序的第一条记录，不会自动优先选择非空值。
• 你的数据中，area列的第一条记录是null，第二条是5.5，所以默认first(area)会取到null；而sub_type列第一条是非空的square，第二条是null，所以first(sub_type)能拿到非空值。两者表现差异是数据顺序导致的，和字段类型无关。

解决方法 #

方法1：使用first()的ignorenulls参数（推荐） #

Spark 2.0及以上版本的first()函数支持ignorenulls=True参数，会自动跳过null值，取分组内第一个非空值：

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

window_spec = Window.partitionBy('name')

result_sdf = sdf.withColumn('prefered_id', F.first('id', ignorenulls=True).over(window_spec))\
                .withColumn('prefered_area', F.first('area', ignorenulls=True).over(window_spec))\
                .withColumn('prefered_sub_type', F.first('sub_type', ignorenulls=True).over(window_spec))\
                .withColumn('prefered_type', F.first('type', ignorenulls=True).over(window_spec))\
                .dropDuplicates(['name'])

方法2：窗口内排序，优先非空值在前 #

通过窗口排序，把非空值排在前面，再用first()取第一条：

window_spec = Window.partitionBy('name').orderBy(F.col('area').isNull(), F.col('sub_type').isNull())

result_sdf = sdf.withColumn('prefered_id', F.first('id').over(window_spec))\
                .withColumn('prefered_area', F.first('area').over(window_spec))\
                .withColumn('prefered_sub_type', F.first('sub_type').over(window_spec))\
                .withColumn('prefered_type', F.first('type').over(window_spec))\
                .dropDuplicates(['name'])

注：F.col('col_name').isNull()会把非空值标记为False（排序时在前），空值标记为True（在后）。

方法3：GroupBy聚合（更简洁） #

直接按name分组，对每个字段用聚合函数取非空值：

result_sdf = sdf.groupBy('name')\
                .agg(F.first('id', ignorenulls=True).alias('prefered_id'),
                     F.first('area', ignorenulls=True).alias('prefered_area'),
                     F.first('sub_type', ignorenulls=True).alias('prefered_sub_type'),
                     F.first('type', ignorenulls=True).alias('prefered_type'))

这种方式不需要窗口函数，逻辑更直接，适合简单的去重取非空需求。

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