[解決済み] 各グループの最初の行を選択するにはどうすればよいですか？

質問

以下のような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のグループから select (0,cat26,30.9)
• 時間のグループから==1 select (1,cat67,28.5)
• 時間グループから==2 select (2,cat56,39.6)
• などなど

つまり、望ましい出力は次のようになる。

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

各グループの上位N行も選択できると便利かもしれませんね。

よろしくお願いします。

解決方法は？

ウィンドウの機能 :

こんな感じでいいんじゃないでしょうか。

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|
// +----+--------+----------+

この方法は、データの傾きが大きい場合、効率が悪くなる。この問題を追跡するのが SPARK-34775 であり、将来的に解決されるかもしれない( SPARK-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)

最後の2つの方法は、マップ側の結合を利用することができ、完全なシャッフルを必要としないので、ほとんどの場合、ウィンドウ関数や結合と比較してより良い性能を示すはずです。また、これらは completed 出力モードです。

を使用しないでください。 :

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

一見うまくいっているように見えるかもしれませんが、（特に local モード)ですが、信頼性に欠けます( SPARK-16207 にクレジットされています。 ザック・ゾハル について 関連する JIRA 課題へのリンク と SPARK-30335 ).

にも同様の注意事項があります。

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

で、内部的には同等の実行計画を使用します。