在 Spark 数据框的列中为每个组添加递增的数字

Question

我有一个包含 2 列“Id”和“category”的数据框。对于每个类别，我想标记编码列“Id”，因此预期的结果将是这样的列“Enc_id”

Id   Category    Enc_id
a1       A         0
a2       A         1
b1       B         0 
c1       C         0
c2       C         1
a3       A         2
b2       B         1
b3       B         2 
b4       B         3 
b4       B         3
b3       B         2

这里的Id可能不是唯一的，所以可能会有重复的行。我想为partitionBy(category) 创建一个窗口，然后在此窗口上应用标签编码（StringIndexer），但它不起作用。请问有什么提示吗？

Answer 1

您可以将window 函数与substring 函数一起使用并计算rank

val window = Window.partitionBy($"Category", substring($"Id", 1,1)).orderBy("Id")

df.withColumn("Enc_id", rank().over(window) - 1) // -1 to start the rank from 0
  .show(false)

输出：

+---+--------+------+
|Id |Category|Enc_id|
+---+--------+------+
|a1 |A       |0     |
|a2 |A       |1     |
|a3 |A       |2     |
|c1 |C       |0     |
|c2 |C       |1     |
|b1 |B       |0     |
|b2 |B       |1     |
|b3 |B       |2     |
|b4 |B       |3     |
+---+--------+------+

更新1： 对于具有重复 id 的更新案例

df1.groupBy("Id", "Category")
  .agg(collect_list("Category") as "list_category")
  .withColumn("Enc_id", rank().over(window) - 1)
  .withColumn("Category", explode($"list_category"))
  .drop("list_category")
  .show(false)

输出：

+---+--------+------+
|Id |Category|Enc_id|
+---+--------+------+
|a1 |A       |0     |
|a2 |A       |1     |
|a3 |A       |2     |
|c1 |C       |0     |
|c2 |C       |1     |
|b1 |B       |0     |
|b2 |B       |1     |
|b3 |B       |2     |
|b3 |B       |2     |
|b4 |B       |3     |
|b4 |B       |3     |
+---+--------+------+