Converting pandas data frame logic to pyspark data frame based logic - pandas

Question

Given a data frame with 4 columns group, start_date, available_stock, used_stock. I basically have to figure out how long a stock will last given a group and date. lets say we have a dataframe with the following data
+----------+------------+-----------------+------------+
| group | start_date | available stock | used_stock |
+----------+------------+-----------------+------------+
| group 1 | 01/12/2019 | 100 | 80 |
| group 1 | 08/12/2019 | 60 | 10 |
| group 1 | 15/12/2019 | 60 | 10 |
| group 1 | 22/12/2019 | 150 | 200 |
| group 2 | 15/12/2019 | 80 | 90 |
| group 2 | 22/12/2019 | 150 | 30 |
| group 3 | 22/12/2019 | 50 | 50 |
+----------+------------+-----------------+------------+
Steps:
sort each group by start_date so we get something like the above data set
per group starting from the smallest date we check if the used_stock is greater or equal to the available stock. if it is true the end date is same as start_date
if the above condition is false then add the next dates used_stock to the current used_stock value. continue till the used_stock is greater or equal to available_stock, at which point the end date is same as the start_date of last added used_stock row.
in case no such value is found end date is null
after applying the above steps for every row we should get something like
+----------+------------+-----------------+------------+------------+
| group | start_date | available stock | used_stock | end_date |
+----------+------------+-----------------+------------+------------+
| group 1 | 01/12/2019 | 100 | 80 | 15/12/2019 |
| group 1 | 08/12/2019 | 60 | 10 | 22/12/2019 |
| group 1 | 15/12/2019 | 60 | 10 | 22/12/2019 |
| group 1 | 22/12/2019 | 150 | 200 | 22/12/2019 |
| group 2 | 15/12/2019 | 80 | 90 | 15/12/2019 |
| group 2 | 22/12/2019 | 150 | 30 | null |
| group 3 | 22/12/2019 | 50 | 50 | 22/12/2019 |
+----------+------------+-----------------+------------+------------+
the above logic was prebuilt in pandas and was tweaked and applied in the spark application as a grouped map Pandas UDF. I want to move away from #pandas_udf approach and have a pure spark data frame based approach to check if there will be any performance improvements.Would appreciate any help with this or any improvements on the given logic which would reduce the overall execution time.

Answer 1

With spark 2.4+, you can use SparkSQL builtin function aggregate:
aggregate(array_argument, zero_expression, merge, finish)
and implement the logic in the merge and finish expressions, see below for an example:
from pyspark.sql.functions import collect_list, struct, to_date, expr
from pyspark.sql import Window
w1 = Window.partitionBy('group').orderBy('start_date').rowsBetween(0, Window.unboundedFollowing)
# SQL expression to calculate end_date using aggregate function:
end_date_expr = """
aggregate(
/* argument */
data,
/* zero expression, initialize and specify the aggregator's datatype which is 'struct<end_date:date,total:double>' */
(date(NULL) as end_date, double(0) as total),
/* merge: use acc.total to save accumulated sum of used_stock
* this works similar to Python's reduce function
*/
(acc, y) ->
IF(acc.total >= `available stock`
, (acc.end_date as end_date, acc.total as total)
, (y.start_date as end_date, acc.total + y.used_stock as total)
),
/* finish: post-processing and retrieving only end_date */
z -> IF(z.total >= `available stock`, z.end_date, NULL)
)
"""
df.withColumn('start_date', to_date('start_date', 'dd/MM/yyyy')) \
.withColumn('data', collect_list(struct('start_date','used_stock')).over(w1)) \
.withColumn('end_date', expr(end_date_expr)) \
.select("group", "start_date", "`available stock`", "used_stock", "end_date") \
.show(truncate=False)
+-------+----------+---------------+----------+----------+
|group |start_date|available stock|used_stock|end_date |
+-------+----------+---------------+----------+----------+
|group 1|2019-12-01|100 |80 |2019-12-15|
|group 1|2019-12-08|60 |10 |2019-12-22|
|group 1|2019-12-15|60 |10 |2019-12-22|
|group 1|2019-12-22|150 |200 |2019-12-22|
|group 2|2019-12-15|80 |90 |2019-12-15|
|group 2|2019-12-22|150 |30 |null |
|group 3|2019-12-22|50 |50 |2019-12-22|
+-------+----------+---------------+----------+----------+
Note: this could be less efficient if many of the groups contain a large list of rows(i.e. 1000+ rows), when most of them require to just scan limited rows (i.e. less than 20) to find the first row satisfying the condition. In such case, you might set up two Window specs and do the calculation in two rounds:
from pyspark.sql.functions import collect_list, struct, to_date, col, when, expr
# 1st scan up to the N following rows which can cover majority of end_date satisfying the condition
N = 20
w2 = Window.partitionBy('group').orderBy('start_date').rowsBetween(0, N)
# 2nd scan will cover the full length but only to rows having end_date is NULL
w1 = Window.partitionBy('group').orderBy('start_date').rowsBetween(0, Window.unboundedFollowing)
df.withColumn('start_date', to_date('start_date', 'dd/MM/yyyy')) \
.withColumn('data', collect_list(struct('start_date','used_stock')).over(w2)) \
.withColumn('end_date', expr(end_date_expr)) \
.withColumn('data',
when(col('end_date').isNull(), collect_list(struct('start_date','used_stock')).over(w1))) \
.selectExpr(
"group",
"start_date",
"`available stock`",
"used_stock",
"IF(end_date is NULL, {0}, end_date) AS end_date".format(end_date_expr)
).show(truncate=False)