Expanding Window combined with Percentile operation on PySpark?

OmniscientCM 0
Viewed 27

I have a PySpark DataFrame with a "SPEED" column per vehicle trip and I want to calculate the 0.95 expanding(1) quantile in a new column.

That is, in Python, starting with:

   TRIP_ID  SPEED
0        1     90
1        1     75
2        1     60
3        2     65
4        2     80
5        2     95
6        3    100
7        3    110
8        3    115

and applying

df.groupby("TRIP_ID")["SPEED"].expanding().quantile(0.95)

gives (after some manipulation to the indexes and columns of the new df):

   TRIP_ID  SPEED  95_PERC
0        1     90    90.00
1        1     75    89.25
2        1     60    88.50
3        2     65    65.00
4        2     80    79.25
5        2     95    93.50
6        3    100   100.00
7        3    110   109.50
8        3    115   114.50

The last column represents the .95 quantile of column "SPEED", calculated in an expanding window fashion: a new value is calculated as a new row is added to the window.

Is there a way to do this in PySpark without having to revert to pandas DF?

quantile rolling-computation pyspark python
1 Answers

I have actually solved this quite easily by:

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

df.withColumn(
    "95_PERC",
    F.expr("percentile(SPEED, array(0.95))").over(Window.partitionBy("TRIP_ID").orderBy().rowsBetween(-sys.maxsize, 0))
)
OmniscientCM0
Related