Como selecionar a primeira linha de cada grupo?

143

Eu tenho um DataFrame gerado da seguinte maneira:

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

Os resultados são parecidos com:

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

Como você pode ver, o DataFrame é ordenado por Houruma ordem crescente e depois por TotalValueuma ordem decrescente.

Gostaria de selecionar a linha superior de cada grupo, ou seja,

do grupo de horas == 0, selecione (0, cat26,30,9)
do grupo de horas == 1, selecione (1, cat67,28,5)
do grupo de horas == 2, selecione (2, cat56,39,6)
e assim por diante

Portanto, a saída desejada seria:

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

Pode ser útil também selecionar as N linhas principais de cada grupo.

Qualquer ajuda é muito apreciada.

sql scala apache-spark dataframe apache-spark-sql Rami
fonte

231

Funções da janela :

Algo assim deve fazer o truque:

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

Este método será ineficiente em caso de distorção significativa dos dados.

Agregação SQL simples seguida dejoin :

Como alternativa, você pode ingressar no quadro de dados agregados:

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

Ele manterá valores duplicados (se houver mais de uma categoria por hora com o mesmo valor total). Você pode removê-los da seguinte maneira:

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

Usando o pedido acimastructs :

Puro, embora não muito bem testado, truque que não requer junções ou funções de janela:

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

Com a API DataSet (Spark 1.6+, 2.0+):

Spark 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 ou posterior :

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

Os dois últimos métodos podem alavancar a combinação lateral do mapa e não exigem a reprodução aleatória completa; portanto, a maior parte do tempo deve exibir um melhor desempenho em comparação com as funções e junções da janela. Essas canas também podem ser usadas com o Streaming Estruturado nocompleted modo de saída.

Não use :

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

Pode parecer funcionar (especialmente no localmodo), mas não é confiável (consulte SPARK-16207 , créditos para Tzach Zohar por vincular o problema relevante do JIRA e SPARK-30335 ).

A mesma nota se aplica a

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

que usa internamente plano de execução equivalente.

zero323
fonte

3

Parece que, desde o Spark 1.6, é row_number () em vez de rowNumber #

Adam Szałucha

Sobre o Não use df.orderBy (...). GropBy (...). Em que circunstâncias podemos confiar em orderBy (...)? ou se não tivermos certeza se orderBy () fornecerá o resultado correto, que alternativas temos?

Ignacio Alorre 27/09

Eu posso estar ignorando alguma coisa, mas, em geral, é recomendável evitar groupByKey , em vez disso, use reduzida. Além disso, você estará salvando uma linha.

Thomas

3

O @Thomas evitar groupBy / groupByKey é justamente quando se lida com RDDs, você notará que a API do conjunto de dados nem sequer possui uma função reduzirByKey.

soote

@Thomas DataFrame / grupo de conjuntos de dadosPor comportamento / otimização

Alper t. Turker 06/06

16

Para o Spark 2.0.2 com agrupamento por várias colunas:

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

Antonín Hoskovec
fonte

8

Esta é uma exata mesma de zero323 's resposta , mas na forma de consulta SQL.

Supondo que o quadro de dados seja criado e registrado como

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

Função de janela:

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

Agregação SQL simples seguida de junção:

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

Usando a ordenação sobre estruturas:

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

DataSets maneira e não fazer s são os mesmos que na resposta original

Ramesh Maharjan
fonte

2

O padrão é agrupar por chaves => fazer algo para cada grupo, por exemplo, reduzir => retornar ao quadro de dados

Eu pensei que a abstração do Dataframe é um pouco complicada nesse caso, então usei a funcionalidade 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)

Pato de borracha
fonte

1

A solução abaixo faz apenas um groupBy e extrai as linhas do seu quadro de dados que contêm o maxValue em uma captura. Não há necessidade de mais junções ou 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)}}

elghoto
fonte

Mas embaralha tudo primeiro. Não é uma melhoria (talvez não seja pior do que as funções da janela, dependendo dos dados).

Alper t. Turker

você tem um primeiro lugar em grupo, que desencadeia uma reprodução aleatória. Não é pior que a função de janela, porque em uma função de janela avaliará a janela para cada linha única no quadro de dados.

Elghoto #

1

Uma boa maneira de fazer isso com a API de dataframe é usar a lógica argmax da seguinte maneira

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

randal25
fonte

0

Aqui você pode fazer assim -

   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

Shubham Agrawal
fonte

-2

Podemos usar a função da janela rank () (onde você escolheria o rank = 1) rank apenas adiciona um número para cada linha de um grupo (nesse caso, seria a hora)

aqui está um exemplo (em 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|
+---+------+----+

Vasile Surdu
fonte

Como selecionar a primeira linha de cada grupo?

Respostas: