根据列的当前值更新 pyspark 中的列

Question

假设给定一个数据帧

+-----+-----+-----+
|    x|    y|    z|
+-----|-----+-----+
|    3|    5|    9|
|    2|    4|    6|
+-----+-----+-----+

我想将z列中的所有值乘以y列中的值，其中z列等于 6。

这篇文章显示了我的目标解决方案，使用代码

from pyspark.sql import functions as F
df = df.withColumn('z',
    F.when(df['z']==6, df['z']*df['y']).
    otherwise(df['z']))

问题是df['z']和df['y']被识别为 Column 对象，强制转换它们将不起作用......

我怎样才能正确地做到这一点？

Answer 1

from pyspark.sql import functions as F
from pyspark.sql.types import LongType
df = df.withColumn('new_col', 
            F.when(df.z==6, 
                (df.z.cast(LongType()) * df.y.cast(LongType()))
            ).otherwise(df.z)
     )