Spark скопировать верхние N выбранных строк в новый фрейм данных

Question

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

Я знаю, как получить top column_ids

 df.groupBy("some_column_id").count()

Это вернет идентификаторы столбцов и их количество, теперь мне нужно отфильтровать строки этих верхних 100 идентификаторов в другом фрейме данных.

Моя примерная таблица ниже, так как у пользователя 123 и 678 больше строк, чем у других, я хочу, чтобы эти два были в отдельном фрейме данных

Как это сделать?

Answer 1

Вы можете сделать это с помощью оконных функций:

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

val newDF = df
  .withColumn("count", count("*").over(Window.partitionBy("student_id")))
  .withColumn("rank", rank().over(Window.orderBy(col("count").desc)))
  .where(col("rank") <= 100)

Answer 2

Попробуйте

val c = df.groupBy("student id").count.orderBy(desc("count")).limit(100)
val a = df.join(c,df.col("student id") === c.col("student id"), "leftsemi")