Как создать столбец идентификатора строки в кадре данных Spark для каждого отдельного значения столбца с помощью Scala

Question

У меня есть фрейм данных в scala spark как

категория |оценка |

A |0,2

А |0,3

A |0,3

B |0,9

B |0,8

B |1

Я хотел бы добавить столбец идентификатора строки как

категория |оценка |идентификатор строки

A |0,2 |0

A |0,3 |1

A |0,3 |2

B |0,9 |0

B |0,8 |1

B |1 |2

По сути, я хочу, чтобы идентификатор строки монотонно увеличивался для каждого отдельного значения в категории столбца.У меня уже есть отсортированный фрейм данных, поэтому все строки с той же категорией сгруппированы вместе.Тем не менее, я до сих пор не знаю, как создать row_id, который перезапускается при появлении новой категории.Пожалуйста, помогите!

Answer 1

Это хороший вариант использования для Window функций агрегирования

import org.apache.spark.sql.expressions.Window
import org.apache.spark.sql.functions.row_number
import df.sparkSession.implicits._

val window = Window.partitionBy('category).orderBy('score)
df.withColumn("row-id", row_number.over(window))

Оконные функции работают примерно так же, как groupBy, за исключением того, что вместо каждой группы возвращая одно значение, каждая строка в каждой группе возвращает одно значение.В этом случае значением является позиция строки в пределах группы строк той же категории.Кроме того, если это тот эффект, которого вы пытаетесь достичь, то вам не нужно предварительно сортировать столбец category заранее.