web-dev-qa-db-fra.com

Médiane / quantiles au sein du groupe PySparkPar

Je voudrais calculer des quantiles de groupe sur un Spark dataframe (avec PySpark). Un résultat approximatif ou exact conviendrait. Je préfère une solution que je peux utiliser dans le contexte de groupBy/agg, de sorte que je puisse le mélanger à d'autres fonctions d'agrégat PySpark Si cela n'est pas possible pour une raison quelconque, une approche différente conviendrait également.

Cette question est liée mais n'indique pas comment utiliser approxQuantile en tant que fonction d'agrégat.

J'ai aussi accès au percentile_approx Hive UDF mais je ne sais pas comment l'utiliser comme fonction d'agrégat.

Par souci de spécificité, supposons que j'ai le dataframe suivant:

from pyspark import SparkContext
import pyspark.sql.functions as f

sc = SparkContext()    

df = sc.parallelize([
    ['A', 1],
    ['A', 2],
    ['A', 3],
    ['B', 4],
    ['B', 5],
    ['B', 6],
]).toDF(('grp', 'val'))

df_grp = df.groupBy('grp').agg(f.magic_percentile('val', 0.5).alias('med_val'))
df_grp.show()

Le résultat attendu est:

+----+-------+
| grp|med_val|
+----+-------+
|   A|      2|
|   B|      5|
+----+-------+
apache-sparkpysparkapache-spark-sqlpyspark-sql
24
abeboparebop

Je suppose que vous n'en avez plus besoin. Mais le laisserons ici pour les générations futures (c’est-à-dire moi la semaine prochaine quand j’oublierai).

from pyspark.sql import Window
import pyspark.sql.functions as F

grp_window = Window.partitionBy('grp')
magic_percentile = F.expr('percentile_approx(val, 0.5)')

df.withColumn('med_val', magic_percentile.over(grp_window))

Ou pour répondre exactement à votre question, cela fonctionne également:

df.groupBy('gpr').agg(magic_percentile.alias('med_val'))

Et en prime, vous pouvez passer un tableau de centiles:

quantiles = F.expr('percentile_approx(val, array(0.25, 0.5, 0.75))')

Et vous obtiendrez une liste en retour.

28
kael

Depuis que vous avez accès à percentile_approx, une solution simple serait de l’utiliser dans une commande SQL:

from pyspark.sql import SQLContext
sqlContext = SQLContext(sc)

df.registerTempTable("df")
df2 = sqlContext.sql("select grp, percentile_approx(val, 0.5) as med_val from df group by grp")
10
Shaido

Malheureusement, et à ma connaissance, il semble que cela ne soit pas possible avec des commandes "pures" de PySpark (la solution de Shaido fournit une solution de contournement avec SQL), et la raison en est très élémentaire: contrairement à d'autres des fonctions telles que mean, approxQuantile ne renvoient pas un type Column, mais une liste .

Voyons un exemple rapide avec vos exemples de données:

spark.version
# u'2.2.0'

import pyspark.sql.functions as func
from pyspark.sql import DataFrameStatFunctions as statFunc

# aggregate with mean works OK:
df_grp_mean = df.groupBy('grp').agg(func.mean(df['val']).alias('mean_val'))
df_grp_mean.show()
# +---+--------+ 
# |grp|mean_val|
# +---+--------+
# |  B|     5.0|
# |  A|     2.0|
# +---+--------+

# try aggregating by median:
df_grp_med = df.groupBy('grp').agg(statFunc(df).approxQuantile('val', [0.5], 0.1))
# AssertionError: all exprs should be Column

# mean aggregation is a Column, but median is a list:

type(func.mean(df['val']))
# pyspark.sql.column.Column

type(statFunc(df).approxQuantile('val', [0.5], 0.1))
# list

Je doute qu'une approche basée sur une fenêtre puisse faire une différence, car comme je l'ai dit, la raison sous-jacente est très élémentaire.

Voir aussi ma réponse ici pour plus de détails.

8
desertnaut