web-dev-qa-db-fra.com

Trouver le nombre maximum de lignes par groupe dans Spark DataFrame

J'essaie d'utiliser Spark dataframes au lieu de RDD car ils semblent être de plus haut niveau que les RDD et ont tendance à produire du code plus lisible.

Dans un cluster Google Dataproc à 14 nœuds, j'ai environ 6 millions de noms traduits en identifiants par deux systèmes différents: sa et sb. Chaque Row contient name, id_sa Et id_sb. Mon objectif est de produire un mappage de id_sa À id_sb De telle sorte que pour chaque id_sa, Le id_sb Correspondant soit l'identifiant le plus fréquent parmi tous les noms associés à id_sa.

Essayons de clarifier avec un exemple. Si j'ai les lignes suivantes:

[Row(name='n1', id_sa='a1', id_sb='b1'),
 Row(name='n2', id_sa='a1', id_sb='b2'),
 Row(name='n3', id_sa='a1', id_sb='b2'),
 Row(name='n4', id_sa='a2', id_sb='b2')]

Mon objectif est de produire un mappage de a1 À b2. En effet, les noms associés à a1 Sont n1, n2 Et n3, Qui correspondent respectivement à b1, b2 et b2, donc b2 est le mappage le plus fréquent dans les noms associés à a1. De la même manière, a2 Sera mappé à b2. Il est normal de supposer qu'il y aura toujours un gagnant: inutile de briser les liens.

J'espérais pouvoir utiliser groupBy(df.id_sa) sur mon cadre de données, mais je ne sais pas quoi faire ensuite. J'espérais une agrégation capable de produire, à la fin, les lignes suivantes:

[Row(id_sa=a1, max_id_sb=b2),
 Row(id_sa=a2, max_id_sb=b2)]

Mais peut-être que j'essaie d'utiliser le mauvais outil et que je devrais simplement revenir à l'utilisation de RDD.

apache-sparkpysparkapache-spark-sql
42
Quentin Pradet

Utilisation de join (il en résultera plus d’une ligne dans le groupe en cas d’égalité):

import pyspark.sql.functions as F
from pyspark.sql.functions import count, col 

cnts = df.groupBy("id_sa", "id_sb").agg(count("*").alias("cnt")).alias("cnts")
maxs = cnts.groupBy("id_sa").agg(F.max("cnt").alias("mx")).alias("maxs")

cnts.join(maxs, 
  (col("cnt") == col("mx")) & (col("cnts.id_sa") == col("maxs.id_sa"))
).select(col("cnts.id_sa"), col("cnts.id_sb"))

Utilisation des fonctions de la fenêtre (laissera tomber les liens):

from pyspark.sql.functions import row_number
from pyspark.sql.window import Window

w = Window().partitionBy("id_sa").orderBy(col("cnt").desc())

(cnts
  .withColumn("rn", row_number().over(w))
  .where(col("rn") == 1)
  .select("id_sa", "id_sb"))

Utilisation de struct ordering:

from pyspark.sql.functions import struct

(cnts
  .groupBy("id_sa")
  .agg(F.max(struct(col("cnt"), col("id_sb"))).alias("max"))
  .select(col("id_sa"), col("max.id_sb")))

Voir aussi Comment sélectionner la première ligne de chaque groupe?

52
zero323

Je pense que vous cherchez peut-être des fonctions de fenêtre: http://spark.Apache.org/docs/latest/api/python/pyspark.sql.html?highlight=window#pyspark.sql.Window

https://databricks.com/blog/2015/07/15/introducing-window-functions-in-spark-sql.html

Voici un exemple dans Scala (Je n’ai pas de Spark Shell avec Hive disponible pour le moment. Par conséquent, je n’ai pas pu tester le code, mais Je pense que ça devrait marcher):

case class MyRow(name: String, id_sa: String, id_sb: String)

val myDF = sc.parallelize(Array(
    MyRow("n1", "a1", "b1"),
    MyRow("n2", "a1", "b2"),
    MyRow("n3", "a1", "b2"),
    MyRow("n1", "a2", "b2")
)).toDF("name", "id_sa", "id_sb")

import org.Apache.spark.sql.expressions.Window

val windowSpec = Window.partitionBy(myDF("id_sa")).orderBy(myDF("id_sb").desc)

myDF.withColumn("max_id_b", first(myDF("id_sb")).over(windowSpec).as("max_id_sb")).filter("id_sb = max_id_sb")

Il existe probablement des moyens plus efficaces d'obtenir les mêmes résultats avec les fonctions Windows, mais j'espère que cela vous oriente dans la bonne direction.

9
alghimo