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|
+-------+----------+----+

Session == 0でtimestamp1列の値を値999に置き換える方法は？

期待される出力

+-------+----------+----+
|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|
+-------+----------+----+

PySparkでreplace（）を使用してそれを行うことは可能ですか？

Answer 1

when（otherwiseを使用）関数を使用する必要があります。

from pyspark.sql.functions import when

targetDf = df.withColumn("timestamp1", \
              when(df["session"] == 0, 999).otherwise(df["timestamp1"]))