Converting for-loop to windowed function

Question

We have a large (spark) dataframe, and we need to compute a new column. Each row is calculated from the value in the previous row in the same column (the first row in the new column is simply 1). This is trivial in a for-loop, but due to the extremely large number of rows we wish to do this in a window function. Because the input to the current row is the previous row, it is not obvious to us how we can achieve this, if it's possible.

We have a large dataframe with a column containing one of three values: A, B and C. Each of these 3 options represents a formula to compute a new value in a new column in the same row.

If it is A, then the new value is 1.
If it is B, then the new value is the same as the previous row.
if it is C, then the new value is the same as the previous row + 1.

For example:

A
B
B
C
B
A
C
B
C
A

Should become:

A 1
B 1
B 1
C 2
B 2
A 1
C 2
B 2
C 3
A 1

We can achieve this behavior as follows using a for loop (pseudocode):

for index in range(my_df):
    if index == 0:
        my_df[new_column][index] = 1

    elseif my_df[letter_column][index] == 'A':
        my_df[new_column][index] = 1

    elseif my_df[letter_column][index] == 'B':
        my_df[new_column][index] = my_df[new_column][index-1]

    elseif my_df[letter_column][index] == 'C':
        my_df[new_column][index] = my_df[new_column][index-1] + 1

We wish to replace the for loop with a window function. We tried using the 'lag' keyword, but the previous row's value depends on previous calculations. Is there a way to do this or is it fundamentally impossible to do this with a window (or map) function? And if it's impossible, is there an alternative that would be faster than the for loop? (A reduce function would have similar performance?)

Again, due to the extremely large number of rows, this is about performance. We should have enough RAM to hold everything in memory, but we wish the processing to be as quick as possible (and to learn how to solve analogues of this problem more generally: applying window functions that require data calculated in previous rows of that window function). Any help would be much appreciated!!

Kind regards, Mick

Converting for-loop to windowed function

0 Answers