[Python] How to optimize field calculation for large dataset in Azure Synapse Spark notebooks...

I have a dataset with 3M+ records and several columns. Here is a sample of my dataset:

item item_base date quantity_1 quantity_2
1 20 202410 600 7493
1 20 202411 17000 16431

Each item-item_base-date makes a unique key. I need to calculate a new column 'actual_value' based on the following logic:

Let's assume we have a rank operation on the table for item-item_base and order by date, then,
For rank = 1, actual_value = quantity_1,
For rank = 2, actual_value = quantity_1 - quantity_2.
For rank > 2, actual_value = quantity_1 - sum(all prev quantity_1) - sum(all prev quantity_2) - sum(all prev actual_value)

Here is my approach of solving this:

First, I create 2 additional columns cumulative_1 and cumulative_2 which are basically sum of quantity_1 and quantity_2 using sql windows.
SUM(quantity_1) OVER(PARTITION BY item, item_base ORDER BY date ROWS BETWEEN UNBOUNDED PRECEIDING AND 1 PRECEDING) as cumulative_1 and so on. Also, I am creating a rank column as a row_no identifier.

Spark does not support recursive CTEs so implementing sum(all prev actual_value) is tedious. I had to switch to pandas dataframe to complete the calculations. Here is my code:

my_df = df.toPandas()
my_df['actual_value'] = 0.0

for i in range(len(my_df)):
if my_df.at[i, 'rank'] == 1:
my_df.at[i, 'actual_value'] = my_df.at[i, 'quantity_1']
elif my_df.at[i, 'rank'] == 2:
my_df.at[i, 'actual_value'] = my_df.at[i, 'quantity_1'] - my_df.at[i, 'quantity_2']
else:
previous_actual_values = my_df.loc[(my_df['item'] == my_df.at[i, 'item']) &
(my_df['item_base'] == my_df.at[i, 'item_base']) &
(my_df['date'] < my_df.at[i, 'date']), 'actual_value'].sum()

my_df.at[i, 'actual_value'] = my_df.at[i, 'quantity_1'] - my_df.at[i, 'cumulative_2'] - my_df.at[i, 'cumulative_1'] - previous_actual_values

if my_df.at[i, 'actual_value'] < 0:
my_df.at[i, 'actual_value'] = 0



The code does the job and gives me the proper output.

item | item_base| date | quantity_1 | quantity_2 | cumulative_1 | cumulative_2 | rank | actual_value
------------|----------|---------|------------|------------|--------------|--------------|------|--------------
1 | 20 | 202410 | 600 | 7493 | | | 1 | 600
1 | 20 | 202411 | 17000 | 16431 | 600 | 7493 | 2 | 569
1 | 20 | 202412 | 785 | 24456 | 17600 | 23924 | 3 | 0
1 | 20 | 202501 | 0 | 25775 | 18385 | 48380 | 4 | 0
1 | 20 | 202502 | | 26131 | 18385 | 74155 | 5 |
1 | 20 | 202503 | 0 | 39452 | 18385 | 100286 | 6 | 0
1 | 20 | 202504 | | 38087 | 18385 | 139738 | 7 |
1 | 20 | 202505 | 2856 | 28916 | 18385 | 177825 | 8 | 0
1 | 20 | 202506 | 500000 | 42254 | 21241 | 206741 | 9 | 270849
1 | 20 | 202507 | | 36776 | 521241 | 248995 | 10 |
1 | 20 | 202508 | 660 | 23523 | 521241 | 285771 | 11 | 0
1 | 20 | 202509 | 1316000 | 25543 | 521901 | 309294 | 12 | 212787
1 | 20 | 202510 | 265220 | 30589 | 1837901 | 334837 | 13 | 0
1 | 20 | 202511 | 47580 | | 1864421 | 365426 | 14 | 0


Now, the problem. Because I have to use pandas, the code takes forever to work for larger datasets. I need to either find a way to do this in Spark itself or improve to efficiency of the above code. I have considered vectorizing the calculations, but I’m struggling to find an efficient way to calculate the cumulative actual_value for rows where rank > 2.

EDIT: I am unable to fix the format of the output table, here is a screenshot of the output:

Continue reading...