Calculate percentiles ignoring missing values - apache-spark

Question

I have a PySpark dataframe with columns ID and BALANCE.
I am trying to bucket the column balance into 100 percentile (1-100%) buckets and calculate how many IDs fall in each bucket.
I cannot use anything related to RDD, I can only use PySpark syntax. I've tried the code below.
w = Window.orderBy(df.BALANCE)
test = df.withColumn('percentile_col', F.percent_rank().over(w))
I am hoping to get a new column that automatically calculates the percentile of each data point in BALANCE column and ignoring the missing value.

Answer 1

Spark 3.1+ has unionByName which has an optional argument allowMissingColumns. This makes it easier.
Test data:
from pyspark.sql import functions as F, Window as W
df = spark.range(12).withColumn(
'balance',
F.when(~F.col('id').isin([0, 1, 2, 3, 4]), F.col('id') + 500))
df.show()
#+---+-------+
#| id|balance|
#+---+-------+
#| 0| null|
#| 1| null|
#| 2| null|
#| 3| null|
#| 4| null|
#| 5| 505|
#| 6| 506|
#| 7| 507|
#| 8| 508|
#| 9| 509|
#| 10| 510|
#| 11| 511|
#+---+-------+
percent_rank will give you exact percentiles - results may have many numbers after the decimal point. That's why percent_rank alone may not provide what you want.
df = (
df.filter(~F.isnull('balance'))
.withColumn('percentile', F.percent_rank().over(W.orderBy('balance')))
.unionByName(df.filter(F.isnull('balance')), True)
)
df.show()
#+---+-------+-------------------+
#| id|balance| percentile|
#+---+-------+-------------------+
#| 5| 505| 0.0|
#| 6| 506|0.16666666666666666|
#| 7| 507| 0.3333333333333333|
#| 8| 508| 0.5|
#| 9| 509| 0.6666666666666666|
#| 10| 510| 0.8333333333333334|
#| 11| 511| 1.0|
#| 0| null| null|
#| 1| null| null|
#| 2| null| null|
#| 3| null| null|
#| 4| null| null|
#+---+-------+-------------------+
The following should work. Rounding step is added.
pr = F.percent_rank().over(W.orderBy('balance'))
df = (
df.filter(~F.isnull('balance'))
.withColumn('bucket', F.when(pr == 0, 1).otherwise(F.ceil(pr * 100)))
.unionByName(df.filter(F.isnull('balance')), True)
)
df.show()
#+---+-------+------+
#| id|balance|bucket|
#+---+-------+------+
#| 5| 505| 1|
#| 6| 506| 17|
#| 7| 507| 34|
#| 8| 508| 50|
#| 9| 509| 67|
#| 10| 510| 84|
#| 11| 511| 100|
#| 0| null| null|
#| 1| null| null|
#| 2| null| null|
#| 3| null| null|
#| 4| null| null|
#+---+-------+------+
You may also consider ntile. Every value is added to one of n "buckets".
When n=100 (test table has less than 100 items, so only the first "buckets" get values):
df = (
df.filter(~F.isnull('balance'))
.withColumn('ntile_100', F.ntile(100).over(W.orderBy('balance')))
.unionByName(df.filter(F.isnull('balance')), True)
)
df.show()
#+---+-------+---------+
#| id|balance|ntile_100|
#+---+-------+---------+
#| 5| 505| 1|
#| 6| 506| 2|
#| 7| 507| 3|
#| 8| 508| 4|
#| 9| 509| 5|
#| 10| 510| 6|
#| 11| 511| 7|
#| 0| null| null|
#| 1| null| null|
#| 2| null| null|
#| 3| null| null|
#| 4| null| null|
#+---+-------+---------+
When n=4:
df = (
df.filter(~F.isnull('balance'))
.withColumn('ntile_100', F.ntile(4).over(W.orderBy('balance')))
.unionByName(df.filter(F.isnull('balance')), True)
)
df.show()
#+---+-------+---------+
#| id|balance|ntile_100|
#+---+-------+---------+
#| 5| 505| 1|
#| 6| 506| 1|
#| 7| 507| 2|
#| 8| 508| 2|
#| 9| 509| 3|
#| 10| 510| 3|
#| 11| 511| 4|
#| 0| null| null|
#| 1| null| null|
#| 2| null| null|
#| 3| null| null|
#| 4| null| null|
#+---+-------+---------+

Answer 2

Try this.
We are first checking if the df.Balance column has Null values. If it has Null values we are displaying None. Else the percent_rank() function gets applied.
from pyspark.sql import functions as F
w = Window.orderBy(df.BALANCE)
test = df.withColumn('percentile_col',when(df.BALANCE.isNull(), lit(None)).otherwise(F.percent_rank().over(w)))