如何在 Dataframe、Pyspark 中更新具有多个条件的行

Question

How to update rows in DataFrame(Pyspark, not scala) where the update should happen on certain conditions? We dont know how many conditions will there be nor what they are during design time, so the conditions and the update values are to be applied at runtime. Sample DataFrame.

Table T1:

| Emp_LName | Emp_FName     |  Sal          |   Sal_Grade  | 
| -------- | -------------- |-------------- |------------- |
| Smith    | Bob            |100000         |B             |
| Barnes   | Jim            |90000          |B             |
| Rogers   | Eric           |120000         |A             |
| Carson   | Ben            |45000          |C             |

df_source = spark.sql("Select * from T1)

lst_Conditions =[(Sal= 45000,Sal_Grade=E),(E_Name='Bob',E_Name='Robert),(One more candition),...].

(Basically the conditions are to be interpreted as 'Where Sal=45000, Then set Sal_Grade=E).

How do I write code for each condition to update the dataframe? Any help is appreciated.

Thank you R

Answer 1

df_source = spark.createDataFrame(
  [
     ('Smith','Bob',100000,'B')
    ,('Barnes','Jim',90000,'B')
    ,('Rogers','Eric',120000,'A')
    ,('Carson','Ben',45000,'C')
  ], ['Emp_LName','Emp_FName','Sal','Sal_Grade']
)
                                           
lst_Conditions = [
    ('Cond_1', 'CASE WHEN Sal = 45000 THEN "E" END'),
    ('Cond_2', 'CASE WHEN Emp_FName = "Bob" THEN "V" END')
]

from pyspark.sql import functions as F

for c in lst_Conditions:
  df_source = df_source.withColumn(c[0], F.expr(c[1]))
  
df_source.show()

+---------+---------+------+---------+------+------+
|Emp_LName|Emp_FName|   Sal|Sal_Grade|Cond_1|Cond_2|
+---------+---------+------+---------+------+------+
|    Smith|      Bob|100000|        B|  null|     V|
|   Barnes|      Jim| 90000|        B|  null|  null|
|   Rogers|     Eric|120000|        A|  null|  null|
|   Carson|      Ben| 45000|        C|     E|  null|
+---------+---------+------+---------+------+------+