Создать идентификатор группы над окном в Spark Dataframe

Question

У меня есть датафрейм, в котором я хочу дать идентификаторы в каждом разделе Window. Например, у меня есть

id | col |
1  |  a  |
2  |  a  |
3  |  b  |
4  |  c  |
5  |  c  |

Итак, я хочу (на основе группировки по столбцу col)

id | group |
1  |  1    |
2  |  1    |
3  |  2    |
4  |  3    |
5  |  3    |

Я хочу использовать оконную функцию, но все равно не могу найти, чтобы назначить Id для каждого окна. Мне нужно что-то вроде:

w = Window().partitionBy('col')
df = df.withColumn("group", id().over(w)) 

Есть ли способ достичь чего-то подобного. (Я не могу просто использовать col в качестве идентификатора группы, потому что я заинтересован в создании окна из нескольких столбцов)

Answer 1

Простое использование встроенной функции dense_rank над оконной функцией даст вам желаемый результат как

from pyspark.sql import window as W
import pyspark.sql.functions as f
df.select('id', f.dense_rank().over(W.Window.orderBy('col')).alias('group')).show(truncate=False)

, который должен дать вам

+---+-----+
|id |group|
+---+-----+
|1  |1    |
|2  |1    |
|3  |2    |
|4  |3    |
|5  |3    |
+---+-----+

Answer 2

Вы можете назначить row_number для отдельных col и себя join с исходным кадром данных.

val data = Seq(
  (1, "a"),
  (2, "a"),
  (3, "b"),
  (4, "c"),
  (5, "c")
).toDF("id","col")

val df2 = data.select("col").distinct()
  .withColumn("group", row_number().over(Window.orderBy("col")))


val result = data.join(df2, Seq("col"), "left")
    .drop("col")

Код в scala, но его можно легко изменить на pyspark.

Надеюсь, это поможет