UDF function taking much time

Question

I have a DataFrame like that:

SCORE = spark.createDataFrame(
    [
      ('a', "Joe", 1),
      ('b', "Doe", 2),
      ('c', "Carl", 3),
      ('d', "CJ", 4),
      ('e', "Tom", 5),
    ], 
      StructType(
        [
            StructField("id", StringType(), False),
            StructField("user", StringType(), False),
            StructField("score", IntegerType(), False),
        ]
    )
)

id	user	score
a	Joe	1
b	Doe	2
c	Carl	3
d	CJ	4
e	Tom	5

I have written a UDF to calculate a percentile_score, and it's based on the whole score column. And it's working, is generating a new column called percentile_score:

from pyspark.sql.functions import udf, collect_list


def calculate_percentile(user_score, score_list):
    data_prs_score.sort()
    scores_count = len(score_list)

    cumulative_frequency = 0
    frequency = 0
    for score in score_list:
        if score == user_score:
            frequency += 1
        elif score > user_score:
            break

        cumulative_frequency += 1

    return (cumulative_frequency - (0.5 * frequency)) / scores_count


def make_score_list(score_list):
     return udf(lambda user_score: calculate_percentile(user_score, score_list), FloatType())
SCORE.withColumn('percentile_score', make_prs_score_list(SCORE.select(collect_list('score')).collect()[0][0])(col('score'))).show()

My problem is, this function is taking 1 hour to run.

I think the reason that is taking so long is that I'm using the collect() on the UDF. However, I can't see another way to build that.

So I was wondering what kind of optimization I can do here.

user_2129801 · Answer

You can use percent_rank to calculate the percentile of each row:

from pyspark.sql import functions as F

SCORE.withColumn("percentile_score", 
  F.percent_rank().over(Window.orderBy("score"))) \
  .show()

prints

+---+----+-----+----------------+                                               
| id|user|score|percentile_score|
+---+----+-----+----------------+
|  a| Joe|    1|             0.0|
|  b| Doe|    2|            0.25|
|  c|Carl|    3|             0.5|
|  d|  CJ|    4|            0.75|
|  e| Tom|    5|             1.0|
+---+----+-----+----------------+

However the following warning will appear:

WARN WindowExec: No Partition Defined for Window operation! Moving all data to a single partition, this can cause serious performance degradation.

Like the warning says all data will be colleted into a single partition and so you lose (temporarily) the parallelism of Spark. But chances are that this approach is still faster than the UDF.

UDF function taking much time

1 Answers