基于列单调增加ID

时间:2018-06-15 21:02:15

标签: scala apache-spark apache-spark-sql

我正在尝试为我的spark DF添加一个新列。我理解可以使用以下代码:

df.withColumn("row",monotonically_increasing_id)

但我的用例是:

输入DF:

col value
  1
  2
  2
  3
  3
  3

输出DF:

col_value      identifier
  1               1
  2               1
  2               2
  3               1
  3               2
  3               3

有关使用monotonically_increasing或rowWithUniqueIndex获取此内容的任何建议。

1 个答案:

答案 0 :(得分:4)

根据您的要求,一种方法是使用row_number窗口函数:

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

val df = Seq(
  1, 2, 2, 3, 3, 3
).toDF("col_value")

val window = Window.partitionBy("col_value").orderBy("col_value")
df.withColumn("identifier", row_number().over(window)).
  orderBy("col_value").
  show
// +---------+----------+
// |col_value|identifier|
// +---------+----------+
// |        1|         1|
// |        2|         1|
// |        2|         2|
// |        3|         1|
// |        3|         2|
// |        3|         3|
// +---------+----------+