Spark技术问题：将DataFrame数组列中的Nan替换为0

假设你的DataFrame名为df，数组列名为array_column，以下是两种可行的实现方式：

方法一：使用Spark内置函数（推荐，无需自定义UDF） #

Spark 3.0及以上版本支持transform函数，可遍历数组内每个元素，结合coalesce将NaN替换为0，性能更优：

// Scala 版本
import org.apache.spark.sql.functions.{transform, coalesce, lit}

val resultDf = df.withColumn(
  "array_column",
  transform($"array_column", x => coalesce(x, lit(0)))
)

# PySpark 版本
from pyspark.sql.functions import transform, coalesce, lit

result_df = df.withColumn(
    "array_column",
    transform("array_column", lambda x: coalesce(x, lit(0)))
)

transform会逐个处理数组元素，coalesce返回第一个非NaN的值，若元素为NaN则返回lit(0)生成的0值。

方法二：自定义UDF（兼容低版本Spark） #

若你的Spark版本低于3.0，可通过自定义UDF实现需求：

// Scala 版本
import org.apache.spark.sql.functions.udf
import scala.collection.mutable.WrappedArray

val replaceNaNWithZero = udf((arr: WrappedArray[Double]) => {
  arr.map(x => if (x.isNaN) 0.0 else x)
})

val resultDf = df.withColumn("array_column", replaceNaNWithZero($"array_column"))

# PySpark 版本
from pyspark.sql.functions import udf
from pyspark.sql.types import ArrayType, DoubleType
import numpy as np

def replace_nan_with_zero(arr):
    return [0.0 if np.isnan(x) else x for x in arr]

replace_nan_udf = udf(replace_nan_with_zero, ArrayType(DoubleType()))
result_df = df.withColumn("array_column", replace_nan_udf("array_column"))

验证结果 #

执行代码后，可通过show()查看处理后的数据：

// Scala
resultDf.show(false)

# PySpark
result_df.show(truncate=False)

输出结果与预期一致：

+---+-----------+
|Id |array_column|
+---+-----------+
|1  |[1.0,2.0,3.0]|
|2  |[0.0,4.0,0.0]|
+---+-----------+

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