Spark dataframe - 如何用连续的整数值填充空值?

Spark dataframe - how to fill nulls with sequential integer values?

假设我有一个像这样的 pyspark 数据框:

KEY    VALUE
---    -----
623    "cat"
245    "dog"
null   "horse"
null   "pig"
331    "narwhal"
null   "snake"

如何转换此数据框,以便将 KEY 列中的任何 null 值替换为从 1 开始的整数序列?期望的结果如下:

KEY    VALUE
---    -----
623    "cat"
245    "dog"
1      "horse"
2      "pig"
331    "narwhal"
3      "snake"

我知道您要求 Python,但也许 Scala 中的等效项会有所帮助。基本上,您想将 Window 函数 rank 与函数 coalesce 一起使用。首先我们定义一些测试数据:

val df = Seq(
  (Option(623), "cat"),
  (Option(245),"dog"),
  (None, "horse"),
  (None, "pig"),
  (Option(331), "narwhal"),
  (None, "snake")
).toDF("key","value")

然后我们将rank一个key的所有实例,然后我们将使用coalesce来选择原始的key或新的rank ],然后删除我们创建的 rank 列来清理它:

import org.apache.spark.sql.expressions._
import org.apache.spark.sql.functions._

val window = Window.partitionBy(col("key")).orderBy(col("value"))
df.withColumn("rank", rank.over(window))
  .withColumn("key", coalesce(col("key"),col("rank")))
  .drop("rank")