将函数应用于数组列 pyspark 中的所有值

Question

我想将 pyspark 数据框中的数组列中的所有值设为负数，而不会爆炸（！）。我尝试了这个 udf 但它不起作用：

negative = func.udf(lambda x: x * -1, T.ArrayType(T.FloatType()))
cast_contracts = cast_contracts \
    .withColumn('forecast_values', negative('forecast_values'))

有人可以帮忙吗？

数据框示例：

df = sc..parallelize(
   [Row(name='Joe', forecast_values=[1.0,2.0,3.0]),
    Row(name='Mary', forecast_values=[4.0,7.1])]).toDF()
>>> df.show()
    +----+---------------+
    |name|forecast_values|
    +----+---------------+
    | Joe|[1.0, 2.0, 3.0]|
    |Mary|     [4.0, 7.1]|
    +----+---------------+

谢谢

Answer 1

我知道这是一年前的帖子，所以我即将提供的解决方案以前可能不是一个选项（它是 Spark 3 的新功能）。如果您在 PySpark API 中使用 Spark 3.0 及更高版本，则应考虑使用spark.sql.function.transforminside pyspark.sql.functions.expr。请不要spark.sql.function.transform与 PySpark 的transform()链接混淆。无论如何，解决方案如下：

df.withColumn("negative", F.expr("transform(forecast_values, x -> x * -1)"))

您唯一需要确保的是将值转换为 int 或 float。突出显示的方法比分解数组或使用 UDF 更有效。

Answer 2

只是您没有循环列表值以将它们乘以 -1

import pyspark.sql.functions as F
import pyspark.sql.types as T

negative = F.udf(lambda x: [i * -1 for i in x], T.ArrayType(T.FloatType()))
cast_contracts = df \
    .withColumn('forecast_values', negative('forecast_values'))

你无法逃避，udf但却是最好的方法。如果你有很大的列表，效果会更好：

import numpy as np

negative = F.udf(lambda x: np.negative(x).tolist(), T.ArrayType(T.FloatType()))
cast_contracts = abdf \
    .withColumn('forecast_values', negative('forecast_values'))
cast_contracts.show()
+------------------+----+
|   forecast_values|name|
+------------------+----+
|[-1.0, -2.0, -3.0]| Joe|
|            [-3.0]|Mary|
|      [-4.0, -7.1]|Mary|
+------------------+----+