Блог

Допустим, у меня есть этот набор данных:

   val main_df = Seq(("yankees-mets",8,20),("yankees-redsox",4,14),("yankees-mets",6,17),
    ("yankees-redsox",2,10),("yankees-mets",5,17),("yankees-redsox",5,10)).toDF("teams","homeruns","hits")
  

который выглядит так:

Я хочу сделать поворот по столбцам команд, а для всех остальных столбцов вернуть 2 (или N) самых высоких значения для этого столбца. Итак, для yankees-mets и homeruns это вернет это,

Поскольку 2 самых высоких итоговых значения homerun для них были 8 и 6.

Как бы я это сделал в общем случае?

Спасибо

Ваша проблема не очень хорошо подходит для pivot, поскольку pivot означает:

Сводная таблица — это агрегация, в которой один (или несколько в общем случае) столбцов группировки имеют свои отдельные значения, перенесенные в отдельные столбцы.

Вы можете создать дополнительный столбец ранга с помощью функции window, а затем выбрать только строки с рангом 1 или 2:

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

main_df.withColumn(
  "rank", 
  rank()
  .over(
    Window.partitionBy("teams")
    .orderBy($"homeruns".desc)
  )
)
.where($"teams" === "yankees-mets" and ($"rank" === 1 or $"rank" === 2))
.show
  

  ------------ -------- ---- ---- 
|       teams|homeruns|hits|rank|
 ------------ -------- ---- ---- 
|yankees-mets|       8|  20|   1|
|yankees-mets|       6|  17|   2|
 ------------ -------- ---- ---- 

  

Затем, если вам больше не нужен rank столбец, вы можете просто удалить его.