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

Question

У меня есть dataframe:

id | num  
0  |  0  
1  |  1  
2  |  1  
3  |  0  
4  |  1  
5  |  1  

И я хотел бы использовать оконные функции для ранжирования каждой «группы» строк, где группа строк - это все строки с 1 с + первый 0 до следующего нуля. В нашем примере:

    id | num | group_id
    0  |  0  | 0
    1  |  1  | 0
    2  |  1  | 0
    3  |  0  | 1
    4  |  1  | 1
    5  |  1  | 1

Как это можно использовать оконные функции ?

Answer 1

Вы можете рассчитать совокупную сумму для выражения "case num = 0 then 1 else 0":

case class GrpRow(id: Int, num: Int)

val df = spark.sparkContext.parallelize(Seq(
  GrpRow(0, 0),
  GrpRow(1, 1),
  GrpRow(2, 1),
  GrpRow(3, 0),
  GrpRow(4, 1),
  GrpRow(5, 1),
  GrpRow(6, 1),
  GrpRow(7, 0),
  GrpRow(8, 1)
)).toDF()

df
  .withColumn(
    "group_id",
    sum(when(col("num") === 0, 1).otherwise(0))
      .over(Window.rangeBetween(Window.unboundedPreceding, Window.currentRow).orderBy("id"))
      .minus(1) // del this row if you want group_id starts with 1 instead of 0
  )
  .show()