将 row_number 列添加到分区的 Spark 数据帧

Question

我正在尝试在分区数据框中添加包含 row_num 的列。

最初，我从 Azure blob 读取我的增量数据：

var df = spark.read.format("delta").load(path)

此数据按日期列分区：

df.rdd.getNumPartitions
res28: Int = 5

因此，当我尝试添加 row_num 列时：

df=df.withColumn("id", monotonically_increasing_id()

它生成 5 个不同的序列（每个分区一个），这显然不是我需要的。

我的问题是：无论如何要在分区数据帧上生成适当的行号列？

我正在考虑使用这样的东西：

df=df.coalesce(1).withColumn("id", monotonically_increasing_id() 

但我不知道我的其余代码是否安全，如果这是最佳实践。

谢谢！

Answer 1

您可以将 window 函数与 row_number 一起使用，如下所示

val window = Window.partitionBy("date")

df.withColumn("id", row_number().over(window))