Apache spark 如果满足条件，Pyspark将函数应用于列值

给定pyspark数据帧，例如：

ls = [
    ['1', 2],
    ['2', 7],
    ['1', 3],
    ['2',-6],
    ['1', 3],
    ['1', 5],
    ['1', 4],
    ['2', 7]
]
df = spark.createDataFrame(pd.DataFrame(ls, columns=['col1', 'col2']))
df.show()

+----+-----+
|col1| col2|
+----+-----+
|   1|    2|
|   2|    7|
|   1|    3|
|   2|   -6|
|   1|    3|
|   1|    5|
|   1|    4|
|   2|    7|
+----+-----+

如何将函数应用于col1=='1'的col2值并将结果存储在新列中？ 例如，功能是：

f = x**2

结果应该如下所示：

+----+-----+-----+
|col1| col2|    y|
+----+-----+-----+
|   1|    2|    4|
|   2|    7| null|
|   1|    3|    9|
|   2|   -6| null|
|   1|    3|    9|
|   1|    5|   25|
|   1|    4|   16|
|   2|    7| null|
+----+-----+-----+

我尝试定义一个单独的函数，并使用df.withColumn（y）.when（条件，函数），但它不起作用

那么，怎样才能做到这一点呢？

我希望这有助于：

def myFun(x):
  return (x**2).cast(IntegerType())

df2 = df.withColumn("y", when(df.col1 == 1, myFun(df.col2)).otherwise(None))

df2.show()

+----+----+----+
|col1|col2|   y|
+----+----+----+
|   1|   2|   4|
|   2|   7|null|
|   1|   3|   9|
|   2|  -6|null|
|   1|   3|   9|
|   1|   5|  25|
|   1|   4|  16|
|   2|   7|null|
+----+----+----+