pyspark approxQuantile function

Question

I have dataframe with these columns id, price, timestamp.

I would like to find median value grouped by id.

I am using this code to find it but it's giving me this error.

from pyspark.sql import DataFrameStatFunctions as statFunc
windowSpec = Window.partitionBy("id")
median = statFunc.approxQuantile("price",
                                 [0.5],
                                 0) \
                 .over(windowSpec)

return df.withColumn("Median", median)

Is it not possible to use DataFrameStatFunctions to fill values in new column?

TypeError: unbound method approxQuantile() must be called with DataFrameStatFunctions instance as first argument (got str instance instead)

Answer 1

Since PySpark 3.1.0 the percentile_approx function has been introduced that solves this problem.

The function percentile_approx returns a list, thus you need to slice the first element.

As in:

windowSpec = Window.partitionBy("id")
df.withColumn("Median", F.percentile_approx("price", [0.5]).over(windowSpec)[0])

Answer 2

If you are fine with aggregation instead of the window function, there is also the option to use a pandas_udf. They are not as fast as pure Spark though. Here is an adapted example from the docs:

from pyspark.sql.functions import pandas_udf, PandasUDFType

df = spark.createDataFrame(
    [(1, 1.0), (1, 2.0), (2, 3.0), (2, 5.0), (2, 10.0)], ("id", "price")
)

@pandas_udf("double", PandasUDFType.GROUPED_AGG)
def median_udf(v):
    return v.median()

df.groupby("id").agg(median_udf(df["price"])).show()