How to conditionally replace value in a column based on evaluation of expression based on another column in Pyspark?

Question

import numpy as np  df = spark.createDataFrame(     [(1, 1, None),      (1, 2, float(5)),      (1, 3, np.nan),      (1, 4, None),      (0, 5, float(10)),      (1, 6, float('nan')),      (0, 6, float('nan'))],     ('session', "timestamp1", "id2")) 

+-------+----------+----+ |session|timestamp1| id2| +-------+----------+----+ |      1|         1|null| |      1|         2| 5.0| |      1|         3| NaN| |      1|         4|null| |      0|         5|10.0| |      1|         6| NaN| |      0|         6| NaN| +-------+----------+----+ 

How to replace value of timestamp1 column with value 999 when session==0?

Expected output

+-------+----------+----+ |session|timestamp1| id2| +-------+----------+----+ |      1|         1|null| |      1|         2| 5.0| |      1|         3| NaN| |      1|         4|null| |      0|       999|10.0| |      1|         6| NaN| |      0|       999| NaN| +-------+----------+----+ 

Is it possible to do it using replace() in PySpark?

Answer 1

You should be using the when (with otherwise) function:

from pyspark.sql.functions import when  targetDf = df.withColumn("timestamp1", \               when(df["session"] == 0, 999).otherwise(df["timestamp1"]))