如何选择每组的第一行？

我有一个 DataFrame 生成如下:

df.groupBy($"Hour", $"Category")
.agg(sum($"value") as "TotalValue")
.sort($"Hour".asc, $"TotalValue".desc))

结果如下:

+----+--------+----------+
|Hour|Category|TotalValue|
+----+--------+----------+
|   0|   cat26|      30.9|
|   0|   cat13|      22.1|
|   0|   cat95|      19.6|
|   0|  cat105|       1.3|
|   1|   cat67|      28.5|
|   1|    cat4|      26.8|
|   1|   cat13|      12.6|
|   1|   cat23|       5.3|
|   2|   cat56|      39.6|
|   2|   cat40|      29.7|
|   2|  cat187|      27.9|
|   2|   cat68|       9.8|
|   3|    cat8|      35.6|
| ...|    ....|      ....|
+----+--------+----------+

正如您可以看到的，DataFrame 按 Hour的递增顺序排序，然后按 TotalValue的递减顺序排序。

我想选择每组的最上一行，即。

  • 从 Hour = = 0组中选择(0，cat26,30.9)
  • 从 Hour = = 1组中选择(1，cat67,28.5)
  • 从小时 = = 2组中选择(2，cat56,39.6)
  • 诸如此类

因此，理想的输出应该是:

+----+--------+----------+
|Hour|Category|TotalValue|
+----+--------+----------+
|   0|   cat26|      30.9|
|   1|   cat67|      28.5|
|   2|   cat56|      39.6|
|   3|    cat8|      35.6|
| ...|     ...|       ...|
+----+--------+----------+

也可以方便地选择每个组的前 N 行。

非常感谢你的帮助。

162006 次浏览
小开
最佳答案

窗口功能 :

像这样的东西应该可以解决问题:

import org.apache.spark.sql.functions.{row_number, max, broadcast}
import org.apache.spark.sql.expressions.Window


val df = sc.parallelize(Seq(
(0,"cat26",30.9), (0,"cat13",22.1), (0,"cat95",19.6), (0,"cat105",1.3),
(1,"cat67",28.5), (1,"cat4",26.8), (1,"cat13",12.6), (1,"cat23",5.3),
(2,"cat56",39.6), (2,"cat40",29.7), (2,"cat187",27.9), (2,"cat68",9.8),
(3,"cat8",35.6))).toDF("Hour", "Category", "TotalValue")


val w = Window.partitionBy($"hour").orderBy($"TotalValue".desc)


val dfTop = df.withColumn("rn", row_number.over(w)).where($"rn" === 1).drop("rn")


dfTop.show
// +----+--------+----------+
// |Hour|Category|TotalValue|
// +----+--------+----------+
// |   0|   cat26|      30.9|
// |   1|   cat67|      28.5|
// |   2|   cat56|      39.6|
// |   3|    cat8|      35.6|
// +----+--------+----------+

在数据严重倾斜的情况下，这种方法效率不高。这个问题由 星火 -34775跟踪，将来可能会得到解决(火花37099)。

普通 SQL 聚合后跟 join :

或者，你可以加入聚合数据框架:

val dfMax = df.groupBy($"hour".as("max_hour")).agg(max($"TotalValue").as("max_value"))


val dfTopByJoin = df.join(broadcast(dfMax),
($"hour" === $"max_hour") && ($"TotalValue" === $"max_value"))
.drop("max_hour")
.drop("max_value")


dfTopByJoin.show


// +----+--------+----------+
// |Hour|Category|TotalValue|
// +----+--------+----------+
// |   0|   cat26|      30.9|
// |   1|   cat67|      28.5|
// |   2|   cat56|      39.6|
// |   3|    cat8|      35.6|
// +----+--------+----------+

它将保持重复的值(如果每小时有一个以上的类别具有相同的总值)。您可以删除以下内容:

dfTopByJoin
.groupBy($"hour")
.agg(
first("category").alias("category"),
first("TotalValue").alias("TotalValue"))

在 structs上使用订单 :

不需要连接或窗口函数的技巧:

val dfTop = df.select($"Hour", struct($"TotalValue", $"Category").alias("vs"))
.groupBy($"hour")
.agg(max("vs").alias("vs"))
.select($"Hour", $"vs.Category", $"vs.TotalValue")


dfTop.show
// +----+--------+----------+
// |Hour|Category|TotalValue|
// +----+--------+----------+
// |   0|   cat26|      30.9|
// |   1|   cat67|      28.5|
// |   2|   cat56|      39.6|
// |   3|    cat8|      35.6|
// +----+--------+----------+

使用 DataSet API (Spark 1.6 + ，2.0 +) :

火花1.6 :

case class Record(Hour: Integer, Category: String, TotalValue: Double)


df.as[Record]
.groupBy($"hour")
.reduce((x, y) => if (x.TotalValue > y.TotalValue) x else y)
.show


// +---+--------------+
// | _1|            _2|
// +---+--------------+
// |[0]|[0,cat26,30.9]|
// |[1]|[1,cat67,28.5]|
// |[2]|[2,cat56,39.6]|
// |[3]| [3,cat8,35.6]|
// +---+--------------+

Spark 2.0或更高版本 :

df.as[Record]
.groupByKey(_.Hour)
.reduceGroups((x, y) => if (x.TotalValue > y.TotalValue) x else y)

最后两个方法可以利用映射侧合并，不需要完全的洗牌，所以大多数时候应该表现出比窗口函数和连接更好的性能。这些甘蔗也用于结构化流在 completed输出模式。

不要使用 :

df.orderBy(...).groupBy(...).agg(first(...), ...)

它可能看起来工作(特别是在 local模式下) ，但是它是不可靠的(参见 SARK-16207，联系 JIRA 相关问题和 火花30335的 扎克 · 佐哈尔学分)。

同样的注意事项也适用于

df.orderBy(...).dropDuplicates(...)

它在内部使用等效的执行计划。

小开

对于按多列分组的 Spark 2.0.2:

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


val w = Window.partitionBy($"col1", $"col2", $"col3").orderBy($"timestamp".desc)


val refined_df = df.withColumn("rn", row_number.over(w)).where($"rn" === 1).drop("rn")
小开

在这里你可以这样做-

   val data = df.groupBy("Hour").agg(first("Hour").as("_1"),first("Category").as("Category"),first("TotalValue").as("TotalValue")).drop("Hour")


data.withColumnRenamed("_1","Hour").show
小开

下面的解决方案只执行一次 groupBy 并一次提取包含 maxValue 的数据框行。不需要进一步的连接，或 Windows。

import org.apache.spark.sql.Row
import org.apache.spark.sql.catalyst.encoders.RowEncoder
import org.apache.spark.sql.DataFrame


//df is the dataframe with Day, Category, TotalValue


implicit val dfEnc = RowEncoder(df.schema)


val res: DataFrame = df
.groupByKey{(r) => r.getInt(0)}
.mapGroups[Row]{(day: Int, rows: Iterator[Row]) => i.maxBy{(r) => r.getDouble(2)}}
小开

这是一个完全相同的 0-323的 回答，但在 SQL 查询方式。

假设创建数据框并将其注册为

df.createOrReplaceTempView("table")
//+----+--------+----------+
//|Hour|Category|TotalValue|
//+----+--------+----------+
//|0   |cat26   |30.9      |
//|0   |cat13   |22.1      |
//|0   |cat95   |19.6      |
//|0   |cat105  |1.3       |
//|1   |cat67   |28.5      |
//|1   |cat4    |26.8      |
//|1   |cat13   |12.6      |
//|1   |cat23   |5.3       |
//|2   |cat56   |39.6      |
//|2   |cat40   |29.7      |
//|2   |cat187  |27.9      |
//|2   |cat68   |9.8       |
//|3   |cat8    |35.6      |
//+----+--------+----------+

窗口功能:

sqlContext.sql("select Hour, Category, TotalValue from (select *, row_number() OVER (PARTITION BY Hour ORDER BY TotalValue DESC) as rn  FROM table) tmp where rn = 1").show(false)
//+----+--------+----------+
//|Hour|Category|TotalValue|
//+----+--------+----------+
//|1   |cat67   |28.5      |
//|3   |cat8    |35.6      |
//|2   |cat56   |39.6      |
//|0   |cat26   |30.9      |
//+----+--------+----------+

普通 SQL 聚合后跟联接:

sqlContext.sql("select Hour, first(Category) as Category, first(TotalValue) as TotalValue from " +
"(select Hour, Category, TotalValue from table tmp1 " +
"join " +
"(select Hour as max_hour, max(TotalValue) as max_value from table group by Hour) tmp2 " +
"on " +
"tmp1.Hour = tmp2.max_hour and tmp1.TotalValue = tmp2.max_value) tmp3 " +
"group by tmp3.Hour")
.show(false)
//+----+--------+----------+
//|Hour|Category|TotalValue|
//+----+--------+----------+
//|1   |cat67   |28.5      |
//|3   |cat8    |35.6      |
//|2   |cat56   |39.6      |
//|0   |cat26   |30.9      |
//+----+--------+----------+

在结构上使用排序:

sqlContext.sql("select Hour, vs.Category, vs.TotalValue from (select Hour, max(struct(TotalValue, Category)) as vs from table group by Hour)").show(false)
//+----+--------+----------+
//|Hour|Category|TotalValue|
//+----+--------+----------+
//|1   |cat67   |28.5      |
//|3   |cat8    |35.6      |
//|2   |cat56   |39.6      |
//|0   |cat26   |30.9      |
//+----+--------+----------+

数据集方式 和 别这样与原答案相同

小开

使用 dataframe api 的一个很好的方法是像这样使用 argmax 逻辑

  val df = Seq(
(0,"cat26",30.9), (0,"cat13",22.1), (0,"cat95",19.6), (0,"cat105",1.3),
(1,"cat67",28.5), (1,"cat4",26.8), (1,"cat13",12.6), (1,"cat23",5.3),
(2,"cat56",39.6), (2,"cat40",29.7), (2,"cat187",27.9), (2,"cat68",9.8),
(3,"cat8",35.6)).toDF("Hour", "Category", "TotalValue")


df.groupBy($"Hour")
.agg(max(struct($"TotalValue", $"Category")).as("argmax"))
.select($"Hour", $"argmax.*").show


+----+----------+--------+
|Hour|TotalValue|Category|
+----+----------+--------+
|   1|      28.5|   cat67|
|   3|      35.6|    cat8|
|   2|      39.6|   cat56|
|   0|      30.9|   cat26|
+----+----------+--------+
小开

模式是 对每个组进行处理，例如 reduce = > 返回到 dataframe

在这种情况下，我认为数据框架的抽象有点麻烦，所以我使用了 RDD 功能

 val rdd: RDD[Row] = originalDf
.rdd
.groupBy(row => row.getAs[String]("grouping_row"))
.map(iterableTuple => {
iterableTuple._2.reduce(reduceFunction)
})


val productDf = sqlContext.createDataFrame(rdd, originalDf.schema)
小开

使用 Apache DataFu可以很容易地做到这一点(实现类似于 安东宁的回答)。

import datafu.spark.DataFrameOps._


val df = sc.parallelize(Seq(
(0,"cat26",30.9), (0,"cat13",22.1), (0,"cat95",19.6), (0,"cat105",1.3),
(1,"cat67",28.5), (1,"cat4",26.8), (1,"cat13",12.6), (1,"cat23",5.3),
(2,"cat56",39.6), (2,"cat40",29.7), (2,"cat187",27.9), (2,"cat68",9.8),
(3,"cat8",35.6))).toDF("Hour", "Category", "TotalValue")


df.dedupWithOrder($"Hour", $"TotalValue".desc).show

这将导致

+----+--------+----------+
|Hour|Category|TotalValue|
+----+--------+----------+
|   0|   cat26|      30.9|
|   3|    cat8|      35.6|
|   1|   cat67|      28.5|
|   2|   cat56|      39.6|
+----+--------+----------+

(是的，结果不会由 一小时排序，但是如果重要的话，您可以稍后再排序)

还有一个 API-DeupTopN-用于获取最上面的 N行。以及另一个 API-DeupWithCombiner-当您期望每个分组有大量行时。

(完全公开-我是 DataFu 项目的一部分)

小开

您可以从 Spark 3.0中使用 max_by()函数！

Https://spark.apache.org/docs/3.0.0-preview/api/sql/index.html#max_by

val df = sc.parallelize(Seq(
(0,"cat26",30.9), (0,"cat13",22.1), (0,"cat95",19.6), (0,"cat105",1.3),
(1,"cat67",28.5), (1,"cat4",26.8), (1,"cat13",12.6), (1,"cat23",5.3),
(2,"cat56",39.6), (2,"cat40",29.7), (2,"cat187",27.9), (2,"cat68",9.8),
(3,"cat8",35.6))).toDF("Hour", "Category", "TotalValue")


// Register the DataFrame as a SQL temporary view
df.createOrReplaceTempView("table")


// Using SQL
val result = spark.sql("select Hour, max_by(Category, TotalValue) AS Category, max(TotalValue) as TotalValue FROM table group by Hour order by Hour")


// or Using DataFrame API
val result = df.groupBy("Hour").
agg(expr("max_by(Category, TotalValue)").as("Category"), max("TotalValue").as("TotalValue")).
sort("Hour")


+----+--------+----------+
|Hour|Category|TotalValue|
+----+--------+----------+
|   0|   cat26|      30.9|
|   1|   cat67|      28.5|
|   2|   cat56|      39.6|
|   3|    cat8|      35.6|
+----+--------+----------+

提交答案