Dynamically fill up rows in a column in pyspark dataframe

Question

I want something like this:

Id  A  B  Flag  COL
 1  5  4   0     0
 1  5  8   1     1
 1  6  4   0     1
 1  4  7   1     2
 2  7  6   0     0
 2  8  9   1     1
 2  3  8   1     2

I have dataframe that has to be partitioned based on id and I have flag based on condition(A<B,Then 1) and I need to get rows in column based on previous row. Logic is if flag is 1, then COL will be previous row value+1 else if flag is 0, the COL will be value of previous row of column itself. P.S- We don't have COL column in df, I am creating it based on above logic. My output should be like above mentioned table.

There is a flaw in your logic. You do not have any ordering rule. If i shuffle your lines, I cannot construct back the dataframe. you need to define a column to order your lines. — Steven
Hi Steven, Logic is whenever I get flag as 1, I need to do +1 to previous row value of "COL" column and whenever I get flag as 0, I just need the same value as previous row value of "COL" column. — Sudip
yes, but what if your rows are shuffled ? Spark does shuffle your data, so line with ID=1 may come one after the other ... on paper, your logic works, but not with dataframe. Please, show your real data, not some example you think is simpler, because it is not. — Steven
Steven, I edited my question . I have used partitioned by and order by 'ID' to get that particular format. Please let me know if it makes any sense now. — Sudip
better but still not working. For ID=1, lines do not have any order, so the column COL can change depending on the order of the lines. It can be 0,0,1,2 or 1,2,2,2 or any other combinaison.We need both this new column and the previous one. — Steven

Steven Steven · Accepted Answer · 2020-09-29T11:50:32

Taking into account the different comments I made, here is my solution based on a valid dataset :

from pyspark.sql import functions as F, Window

df.show()  # Without columns parition and order, it is impossible to compute COL

+---------+-----+---+---+
|partition|order|  A|  B|
+---------+-----+---+---+
|        1|    1|  5|  4|
|        1|    2|  5|  8|
|        1|    3|  6|  4|
|        1|    4|  4|  7|
|        2|    1|  7|  6|
|        2|    2|  8|  9|
|        2|    3|  3|  8|
+---------+-----+---+---+

df.withColumn("flag", F.when(F.col("A") < F.col("B"), 1).otherwise(0)).withColumn(
    "COL",
    F.sum("flag").over(
        Window.partitionBy("partition").orderBy(
            "order"
        )  # Window is the reason why we need these two columns
    ),
).show()

+---------+-----+---+---+----+---+
|partition|order|  A|  B|flag|COL|
+---------+-----+---+---+----+---+
|        1|    1|  5|  4|   0|  0|
|        1|    2|  5|  8|   1|  1|
|        1|    3|  6|  4|   0|  1|
|        1|    4|  4|  7|   1|  2|
|        2|    1|  7|  6|   0|  0|
|        2|    2|  8|  9|   1|  1|
|        2|    3|  3|  8|   1|  2|
+---------+-----+---+---+----+---+

Dynamically fill up rows in a column in pyspark dataframe

1 Answers