Создайте другой столбец, используя значение другого столбца.

Question

У меня есть фрейм данных, в который мне нужно добавить еще один столбец на основе логики группировки c.

Dataframe

 id|x_id|y_id|val_id|
  1| 2  |  3 |  4   |  
 10| 2  |  3 |  40  | 
  1| 12 | 13 | 14   | 

Мне нужно добавить другие столбцы parent_id, которые будут основаны на этом правиле:

over x_id и y_id выберите максимальное значение в столбце val_id и используйте соответствующее ему id значение

Окончательный кадр будет выглядеть следующим образом

 id|x_id|y_id|val_id| parent_id
 91| 2  |  3 |  4   |  10 (coming from row 2)
 10| 2  |  3 |  40  |  10 (coming from row 2)
  1| 12 | 13 | 14   |  14

Я пытался использовать withColumn, но Я могу только установить строку над этой группой, чтобы ее значение было родительским.

Объяснение : Здесь parent_id равно 10, поскольку оно исходит от col id. Строка 2 была выбрана, поскольку она имеет максимальное значение val_id для группы x_id и y_id

Я использую scala

Answer 1

Используйте Window, чтобы разделить идентификаторы и рассчитать максимум по окну, отсортировав для каждого раздела по val_id.

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

val w = Window.partitionBy('x_id, 'y_id).orderBy('val_id.desc)

df.withColumn("parent_id", first('id).over(w))
  .show(false)

Результат:

+---+----+----+------+---------+
|id |x_id|y_id|val_id|parent_id|
+---+----+----+------+---------+
|10 |2   |3   |40    |10       |
|1  |2   |3   |4     |10       |
|1  |12  |13  |14    |1        |
+---+----+----+------+---------+