Spark将前N个选定的行复制到新数据框中

Question

我有一些具有数百万行的数据框。我需要选择其他数据框中前100个ID的所有行

我知道如何获得顶部的column_ids

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

这将返回列ID及其计数，现在我需要将前100个ID的行过滤到其他数据框中。

下面是我的示例表，由于用户 123 和 678 的行比其他行多，我希望这两行位于单独的数据框中

如何执行此操作？

Answer 1

尝试

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")

Answer 2

您可以使用窗口功能来做到这一点：

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)