使用Spark Scala向每个分区添加一个常量值

我正在尝试使用Spark Scala为每一组日期添加一个id

例如，如果输入为：

date
2019-01-29
2019-01-29
2019-07-31
2019-01-29
2019-07-31

产出将是：

id, date
ABC1, 2019-01-29
ABC1, 2019-01-29
ABC1, 2019-01-29
ABC2, 2019-07-31
ABC2, 2019-07-31

有人能帮我吗

我成功地为每个分区添加了顺序行号，但我希望每个分区都有一个常量值

df.withColumn(lineNumColName, row_number().over(Window.partitionBy(partitionByCol).orderBy(orderByCol))).repartition(1).orderBy(orderByCol, lineNumColName)

选项1（小数据集）：

如果数据集不太大，可以使用Window和dense_rank，如下所示：

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.{concat,lit, dense_rank}

val df = Seq(("2019-01-29"),
("2019-01-29"),
("2019-07-31"),
("2019-01-29"),
("2019-07-31")).toDF("date")

val w = Window.orderBy($"date") 
val d_rank = dense_rank().over(w)
df.withColumn("id",  concat(lit("ABC"), d_rank)).show(false)

输出：

+----------+----+
|date      |id  |
+----------+----+
|2019-01-29|ABC1|
|2019-01-29|ABC1|
|2019-01-29|ABC1|
|2019-07-31|ABC2|
|2019-07-31|ABC2|
+----------+----+

由于我们没有为

partitionBy

部分指定任何值，因此这将只使用一个分区，因此效率非常低

选项2（大数据集）：

更有效的方法是使用

zipWithIndex

函数为大型数据集分配ID：

val df_d = df.distinct.rdd.zipWithIndex().map{ r => (r._1.getString(0), r._2 + 1) }.toDF("date", "id")
df_d.show

// Output:
+----------+---+
|      date| id|
+----------+---+
|2019-01-29|  1|
|2019-07-31|  2|
+----------+---+

首先，我们使用

distinct

获取数据帧的唯一值，然后调用

zipWithIndex

为每个日期记录创建一个唯一的id

最后，我们加入两个数据集：

df.join(df_d, Seq("date"))
.withColumn("id",  concat(lit("ABC"), $"id"))
.show

// Output:
+----------+----+
|      date|  id|
+----------+----+
|2019-01-29|ABC1|
|2019-01-29|ABC1|
|2019-01-29|ABC1|
|2019-07-31|ABC2|
|2019-07-31|ABC2|
+----------+----+