web-dev-qa-db-fra.com

Agrégation de plusieurs colonnes avec une fonction personnalisée dans Spark

Je me demandais s'il y avait un moyen de spécifier une fonction d'agrégation personnalisée pour spark dataframes sur plusieurs colonnes.

J'ai une table comme celle-ci du type (nom, article, prix):

john | tomato | 1.99
john | carrot | 0.45
bill | Apple  | 0.99
john | banana | 1.29
bill | taco   | 2.59

à:

Je voudrais regrouper l'article et son coût pour chaque personne dans une liste comme celle-ci:

john | (tomato, 1.99), (carrot, 0.45), (banana, 1.29)
bill | (Apple, 0.99), (taco, 2.59)

Est-ce possible dans les trames de données? J'ai récemment entendu parler de collect_list mais il ne semble fonctionner que pour une colonne.

scalaapache-sparkdataframeapache-spark-sqlorc
31
anthonybell

La façon la plus simple de le faire en tant que DataFrame est de collecter d'abord deux listes, puis d'utiliser un UDF pour Zip les deux listes ensemble. Quelque chose comme:

import org.Apache.spark.sql.functions.{collect_list, udf}
import sqlContext.implicits._

val zipper = udf[Seq[(String, Double)], Seq[String], Seq[Double]](_.Zip(_))

val df = Seq(
  ("john", "tomato", 1.99),
  ("john", "carrot", 0.45),
  ("bill", "Apple", 0.99),
  ("john", "banana", 1.29),
  ("bill", "taco", 2.59)
).toDF("name", "food", "price")

val df2 = df.groupBy("name").agg(
  collect_list(col("food")) as "food",
  collect_list(col("price")) as "price" 
).withColumn("food", zipper(col("food"), col("price"))).drop("price")

df2.show(false)
# +----+---------------------------------------------+
# |name|food                                         |
# +----+---------------------------------------------+
# |john|[[tomato,1.99], [carrot,0.45], [banana,1.29]]|
# |bill|[[Apple,0.99], [taco,2.59]]                  |
# +----+---------------------------------------------+
29
David Griffin

Pensez à utiliser la fonction struct pour regrouper les colonnes avant de les collecter sous forme de liste:

import org.Apache.spark.sql.functions.{collect_list, struct}
import sqlContext.implicits._

val df = Seq(
  ("john", "tomato", 1.99),
  ("john", "carrot", 0.45),
  ("bill", "Apple", 0.99),
  ("john", "banana", 1.29),
  ("bill", "taco", 2.59)
).toDF("name", "food", "price")

df.groupBy($"name")
  .agg(collect_list(struct($"food", $"price")).as("foods"))
  .show(false)

Les sorties:

+----+---------------------------------------------+
|name|foods                                        |
+----+---------------------------------------------+
|john|[[tomato,1.99], [carrot,0.45], [banana,1.29]]|
|bill|[[Apple,0.99], [taco,2.59]]                  |
+----+---------------------------------------------+
60
Daniel Siegmann

Peut-être qu'un meilleur moyen que la fonction Zip (puisque UDF et UDAF sont très mauvais pour les performances) est de mettre les deux colonnes dans Struct.

Cela fonctionnerait probablement aussi:

df.select('name, struct('food, 'price).as("Tuple"))
  .groupBy('name)
  .agg(collect_list('Tuple).as("tuples"))
5
Yifan Guo

Voici une option en convertissant la trame de données en un RDD de Map puis en appelant un groupByKey dessus. Le résultat serait une liste de paires clé-valeur où valeur est une liste de tuples.

df.show
+----+------+----+
|  _1|    _2|  _3|
+----+------+----+
|john|tomato|1.99|
|john|carrot|0.45|
|bill| Apple|0.99|
|john|banana|1.29|
|bill|  taco|2.59|
+----+------+----+


val tuples = df.map(row => row(0) -> (row(1), row(2)))
tuples: org.Apache.spark.rdd.RDD[(Any, (Any, Any))] = MapPartitionsRDD[102] at map at <console>:43

tuples.groupByKey().map{ case(x, y) => (x, y.toList) }.collect
res76: Array[(Any, List[(Any, Any)])] = Array((bill,List((Apple,0.99), (taco,2.59))), (john,List((tomato,1.99), (carrot,0.45), (banana,1.29))))
2
Psidom