ChatGPT解决这个技术问题 Extra ChatGPT

Median / quantiles within PySpark groupBy

I would like to calculate group quantiles on a Spark dataframe (using PySpark). Either an approximate or exact result would be fine. I prefer a solution that I can use within the context of groupBy / agg, so that I can mix it with other PySpark aggregate functions. If this is not possible for some reason, a different approach would be fine as well.

This question is related but does not indicate how to use approxQuantile as an aggregate function.

I also have access to the percentile_approx Hive UDF but I don't know how to use it as an aggregate function.

For the sake of specificity, suppose I have the following dataframe:

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

Expected result is:

+----+-------+
| grp|med_val|
+----+-------+
|   A|      2|
|   B|      5|
+----+-------+
I think you might be able to roll your own in this instance using the underlying rdd and an algorithm for computing distributed quantiles e.g. here and links therein. In fact, the github they link to has some pyspark examples.

C
Chris A.

I guess you don't need it anymore. But will leave it here for future generations (i.e. me next week when I forget).

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

Or to address exactly your question, this also works:

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

And as a bonus, you can pass an array of percentiles:

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

And you'll get a list in return.


Very clean answer. Do you know how can it be done using Pandas UDF (a.k.a. Vectorized UDFs) too?
@CesareIurlaro, I've only wrapped it in a UDF. Never tried with a Pandas one
Would you mind to try? Performace really should shine there: databricks.com/blog/2017/10/30/…
With Spark 3.1.0 it is now possible to use percentile_approx directly in PySpark groupby aggregations: df.groupBy("key").agg(percentile_approx("value", 0.5, lit(1000000)).alias("median")) spark.apache.org/docs/latest/api/python/reference/api/…
@thentangler: the former is an exact percentile, which is not a scalable operation for large datasets, and the latter is approximate but scalable.
d
desertnaut

Since you have access to percentile_approx, one simple solution would be to use it in a SQL command:

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

This works, but I prefer a solution that I can use within groupBy / agg at the PySpark level (so that I can easily mix it with other PySpark aggregate functions).
@abeboparebop I do not beleive it's possible to only use groupBy and agg, however, to use a window-based approach should also work.
I have clarified my ideal solution in the question. Clearly this answer does the job, but it's not quite what I want. I'll leave the question open for some time to see if a cleaner answer comes up.
d
desertnaut

(UPDATE: now it is possible, see accepted answer above)

Unfortunately, and to the best of my knowledge, it seems that it is not possible to do this with "pure" PySpark commands (the solution by Shaido provides a workaround with SQL), and the reason is very elementary: in contrast with other aggregate functions, such as mean, approxQuantile does not return a Column type, but a list.

Let's see a quick example with your sample data:

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

I doubt that a window-based approach will make any difference, since as I said the underlying reason is a very elementary one.

See also my answer here for some more details.


K
Kxrr

The most simple way to do this with pyspark==2.4.5 is:

df \
    .groupby('grp') \
    .agg(expr('percentile(val, array(0.5))')[0].alias('50%')) \
    .show()

output:

|grp|50%|
+---+---+
|  B|5.0|
|  A|2.0|
+---+---+

J
Jan_ewazz

It seems to be completely solved by pyspark >= 3.1.0 using percentile_approx

import pyspark.sql.functions as func    

df.groupBy("grp").agg(func.percentile_approx("val", 0.5).alias("median"))

For further information see: https://spark.apache.org/docs/3.1.1/api/python/reference/api/pyspark.sql.functions.percentile_approx.html


e
eid

problem of "percentile_approx(val, 0.5)": if e.g. range is [1,2,3,4] this function returns 2 (as median) the function below returns 2.5:

import statistics

median_udf = F.udf(lambda x: statistics.median(x) if bool(x) else None, DoubleType())

... .groupBy('something').agg(median_udf(F.collect_list(F.col('value'))).alias('median'))