database - 用 Spark Dataframe 中另一个分类列的平均值替换列的空值

Question

我有这样的数据集

id    category     value
1     A            NaN
2     B            NaN
3     A            10.5
5     A            2.0
6     B            1.0

我想用各自类别的平均值填充 NAN 值。如下图

id    category     value
1     A            4.16
2     B            0.5
3     A            10.5
5     A            2.0
6     B            1.0

我尝试使用分组依据计算每个类别的第一个平均值

val df2 = dataFrame.groupBy(category).agg(mean(value)).rdd.map{
      case r:Row => (r.getAs[String](category),r.get(1))
    }.collect().toMap
    println(df2)

我得到了每个类别的 map 及其各自的平均值。输出: map (A ->4.16,B->0.5) 现在我尝试在 Sparksql 中更新查询来填充列，但似乎 spqrkSql dosnt 支持更新查询。我试图在数据框中填充空值，但没有这样做。 我能做什么？我们可以在 pandas 中做同样的事情，如 Pandas: How to fill null values with mean of a groupby? 所示。 但是我该怎么做才能使用 spark dataframe

Answer 1

最简单的解决方案是使用 groupby 和 join:

 val df2 = df.filter(!(isnan($"value"))).groupBy("category").agg(avg($"value").as("avg"))
 df.join(df2, "category").withColumn("value", when(col("value").isNaN, $"avg").otherwise($"value")).drop("avg")

请注意，如果有一个类别全部为 NaN，它将从结果中删除