Как выбрать первую строку каждой группы?

Question

У меня есть 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)
• из группы 1 час == выберите (1, cat67,28.5)
• из группы Hour == 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 строк каждой группы.

Любая помощь очень ценится.

Answer 1

Оконные функции:

Что-то, как это должно сделать трюк:

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

Этот метод будет неэффективным в случае значительного перекоса данных. агрегация Plain SQL

следует join:

В качестве альтернативы вы можете присоединиться к агрегированному кадру данных:

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

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 (Искры 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]| 
// +---+--------------+ 

Спарк 2.0 или более поздняя версия:

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

Последние два метода можно использовать карты со стороны комбината и не требуют полной перетасовать поэтому большую часть времени должна демонстрировать более высокую производительность по сравнению с оконными функциями и присоединяется.

Не используйте:

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

Это может показаться на работу (особенно в режиме local), но это ненадежно (SPARK-16207). Кредиты к Tzach Zohar за linking relevant JIRA issue.

То же замечание относится и к

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

, который внутренне использует эквивалентный план выполнения.

Answer 2

Если dataframe должен быть сгруппирован по нескольким столбцам, это может помочь

val keys = List("Hour", "Category"); 
val selectFirstValueOfNoneGroupedColumns = 
df.columns 
    .filterNot(keys.toSet) 
    .map(_ -> "first").toMap 
val grouped = 
df.groupBy(keys.head, keys.tail: _*) 
    .agg(selectFirstValueOfNoneGroupedColumns) 

Надеется, что это помогает кто-то с подобной проблемой

Answer 3

Для Спарка 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") 

Answer 4

Мы можем использовать функцию ранга() (где вы бы выбрали ранг = 1) Ранг просто добавляет число для каждой строки группы (в этом случай будет час)

вот пример. (От https://github.com/jaceklaskowski/mastering-apache-spark-book/blob/master/spark-sql-functions.adoc#rank)

val dataset = spark.range(9).withColumn("bucket", 'id % 3) 

import org.apache.spark.sql.expressions.Window 
val byBucket = Window.partitionBy('bucket).orderBy('id) 

scala> dataset.withColumn("rank", rank over byBucket).show 
+---+------+----+ 
| id|bucket|rank| 
+---+------+----+ 
| 0|  0| 1| 
| 3|  0| 2| 
| 6|  0| 3| 
| 1|  1| 1| 
| 4|  1| 2| 
| 7|  1| 3| 
| 2|  2| 1| 
| 5|  2| 2| 
| 8|  2| 3| 
+---+------+----+ 

Answer 5

Для искры> 2.0 мы можем просто сделать:
groupBy($"Hour").agg(df_op.columns.map((_, "first")).toMap)

установки: OP в деталях:

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

df_op.groupBy($"Hour").agg(df_op.columns.map((_, "first")).toMap) 

Это использует agg метод RelationalGroupedDataset до Compute aggregates by specifying a map from column name to aggregate methods.. first - это функция агрегации sql.

Answer 6

Здесь вы можете сделать, как это -

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