Resolved: Pyspark Conditional statement

+-----------+---------+------------------+----------------------+-----------+
|  DATE    |  ID      |sal               |   vat                |     flag  |
+-----------+---------+------------------+----------------------+------------
|10-may-2022|     1   |             1000.0|                  12.0       1   |
|12-may-2022|     2   |              50.0|                   6.0|       1   |
+-----------+---------+------------------+----------------------+------------

df = srcdf.withColumn("sum",col("sal")*2)
display(df)

df = srcdf.withColumn("sum",col("sal")*4)
display(df)

flag = srcdf.select(col("flag"))

if flag == 1 :

df = srcdf.withColumn("sum",col("sal")*2)
display(df)

else:
df = srcdf.withColumn("sum",col("sal")*4)
display(df)

from pyspark.sql.functions import when, col
df = srcdf\
   .withColumn("sum", when(col("flag") == 1, col("sal") * 2)\
                     .when(col("flag") == 2, col("sal") * 4)
   )

from pyspark.sql.functions import when, col
df = srcdf\
   .withColumn("sum", when(col("flag") == 1, col("sal") * 2)\
                     .otherwise(col("sal") * 4)
   )