数据框pyspark更新前一行的行

Question

I am using pyspark, and I have a dataframe that looks like that :

CODE  |  POSITION|  COL1 | COL2
A     |  1       |       | 
A     |  2       |       | AAA
A     |  3       |   INF |
A     |  4       |   BIC |
A     |  5       |       |
B     |  1       |       | BBB
B     |  2       |   MIL |
B     |  3       |       |
B     |  4       |       | CCC
B     |  5       |       |
B     |  6       |       |

and I want to have that :

CODE  |  POSITION|  COL1 | COL2
A     |  1       |       | 
A     |  2       |       | AAA
A     |  3       |   INF | AAA
A     |  4       |   BIC | AAA
A     |  5       |       |
B     |  1       |       | BBB
B     |  2       |   MIL | BBB
B     |  3       |       |
B     |  4       |       | CCC
B     |  5       |       |
B     |  6       |       |

I explain, this dataframe is grouped by "CODE" and ordered by "POSITION", I need for a group "CODE" , when I have "COL2" filled (position =2 in this example) to take the value "AAA" and put it in the following positions 3 and 4 (while COL1 is filled)

I know is not that easy(for me!)

Thank you a lot for your help

Answer 1

It can be done using the last function.
F.last returns the last non-null value in an ordered window.

Your dataframe:

from pyspark.sql.functions import col
from pyspark.sql.functions import lag
from pyspark.sql.window import Window
from pyspark.sql import functions as F
import sys

df = sc.parallelize([['A', 1, None, None], ['A', 2, None, 'AAA'], ['A', 3, 'INF', None], ['A', 4, 'BIC', None], ['A', 5, None, None], ['B', 1, None, 'BBB'], ['B', 2, 'MIL', None], ['B', 3, None, None], ['B', 4, None, 'CCC'], ['B', 5, None, None], ['B', 6, None, None]])
df = df.toDF(['code', 'position', 'col1', 'col2'])

w = Window.partitionBy("code").orderBy("position")

df.withColumn("col3", F.last('col2', True).over(w.rowsBetween(-sys.maxsize, 0)))\
    .withColumn("col3", F.when(col("col1").isNull(), col("col2"))
    .otherwise(col("col3")))\
    .drop("col2").withColumnRenamed("col3", "col2")\
    .orderBy("code", "position").show()

Output:

+----+--------+----+----+                                                       
|code|position|col1|col2|
+----+--------+----+----+
|   A|       1|null|null|
|   A|       2|null| AAA|
|   A|       3| INF| AAA|
|   A|       4| BIC| AAA|
|   A|       5|null|null|
|   B|       1|null| BBB|
|   B|       2| MIL| BBB|
|   B|       3|null|null|
|   B|       4|null| CCC|
|   B|       5|null|null|
|   B|       6|null|null|
+----+--------+----+----+

If col1 corresponding to position 6 is filled, it will return CCC in col2 .
It takes the latest non-null value in col2 as it progresses through the window.

+----+--------+----+----+
|   B|       6| XYZ| CCC|
+----+--------+----+----+