How to update rows with many conditions in Dataframe, Pyspark

user_7092711 0
Viewed 243

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

pyspark dataframe python
1 Answers
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|
+---------+---------+------+---------+------+------+
user_116985620
Related