如何将DataFrame列中的NA替换为同列的非NA字符串？

我懂你遇到的困境了——你的DataFrame里code列只有两类值：字符串"NA"和某个不确定的非NA字符串（比如例子里的GSTR），想用UDF但因为摸不准这个目标字符串到底是什么而卡壳了对吧？其实完全不用UDF就能搞定，而且效率更高，下面给你具体操作步骤：

步骤1：先提取列中的唯一非NA值 #

既然题目明确code列只有一种非NA值，那我们先把它提取出来：

# 过滤掉code为"NA"的行，去重后取出唯一的非NA值
non_na_code = df.filter(df.code != "NA").select("code").distinct().collect()[0][0]

这里collect()[0][0]是因为我们拿到的是一个仅一行一列的DataFrame，直接提取对应位置的值就行。

步骤2：批量替换所有"NA"为目标值 #

接下来用Spark内置的when函数完成替换，代码简洁还比UDF性能好：

from pyspark.sql.functions import when, col

# 生成更新后的code列，将"NA"替换为我们提取到的非NA值
updated_df = df.withColumn(
    "code",
    when(col("code") == "NA", non_na_code).otherwise(col("code"))
)

最终效果验证 #

执行完上面的代码后，你的DataFrame就会变成你期望的样子：

id,code
1,GSTR
2,GSTR
3,GSTR
4,GSTR
5,GSTR

为啥不用UDF呢？因为UDF是逐行处理的，没办法直接获取全局的非NA值，而且Spark内置函数的执行效率远高于自定义UDF，这种场景下完全没必要舍近求远~

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