How to add multiple new columns with when condition in pyspark dataframe?-CodePudding

I need to add two new columns to my existing pyspark dataframe.
Below is my sample data:

Section   Grade     Promotion_grade Section_team
Admin       C       
Account     B       
IT          B   

condition :

If Section = Admin then Promotion_grade = B
If Section = Account then Promotion_grade = A
If Section = IT then
             If Grade = C then Promotion_grade = B & Section_team= team1
             If Grade = D  then Promotion_grade = C & Section_team= team2
             If Grade = A  then Promotion_grade = A  & Section_team= team3

I can add one column for first two conditions. But I don't know for the rest conditions.

def addCols(data):
   data = (data.withColumn('Promotion_grade', F.when(data.Section  =='Admin', 'B')
                                                .when(data.Section  =='Account', 'A')
                                                .otherwise('Not applicable')))
   return data

Please someone can help me in this? May be the way I'm doing is wrong. Thank you

CodePudding user response：

You can nest when conditions to handle nested conditions.

Working Example

from pyspark.sql import functions as F

data = [("Admin", "C", ), 
        ("Account", "B", ), 
        ("IT", "B", ),
        ("IT", "C", ),
        ("IT", "D", ),
        ("IT", "A", ),]

df = spark.createDataFrame(data, ("Section", "Grade", ))

# Define Promotion Grade conditions for IT Section
it_promotion_grade = (F.when(F.col("Grade") == "C", "B")
                       .when(F.col("Grade") == "D", "C")
                       .when(F.col("Grade") == "A", "A ")
                       .otherwise("Not applicable"))

# Define Section Team conditions for IT Section
it_section_team = (F.when(F.col("Grade") == "C", "team1")
                    .when(F.col("Grade") == "D", "team2")
                    .when(F.col("Grade") == "A", "team3")
                    .otherwise("Not applicable"))

(df.withColumn("Promotion_grade", F.when(F.col("Section") == "Admin", "B")
                                  .when(F.col("Section") == "Account", "A")
                                  .when(F.col("Section") == "IT", it_promotion_grade)
                                  .otherwise("Not applicable"))
    .withColumn("Section_team", F.when(F.col("Section") == "IT", it_section_team)
                     .otherwise("Not applicable"))
    .show())

Output

 ------- ----- --------------- -------------- 
|Section|Grade|Promotion_grade|  Section_team|
 ------- ----- --------------- -------------- 
|  Admin|    C|              B|Not applicable|
|Account|    B|              A|Not applicable|
|     IT|    B| Not applicable|Not applicable|
|     IT|    C|              B|         team1|
|     IT|    D|              C|         team2|
|     IT|    A|             A |         team3|
 ------- ----- --------------- --------------