具有多个参数的 PySpark UDF 返回 null

Question

我有一个 PySpark Dataframe，有两列（A, B，其类型为double），其值为0.0或1.0。我正在尝试添加一个新列，它是这两列的总和。我遵循Pyspark 中的示例：Pass multiple columns in UDF

import pyspark.sql.functions as F
from pyspark.sql.types import IntegerType, StringType
sum_cols = F.udf(lambda x: x[0]+x[1], IntegerType())
df_with_sum = df.withColumn('SUM_COL',sum_cols(F.array('A','B')))
df_with_sum.select(['SUM_COL']).toPandas()

这显示了一系列NULLs 而不是我期望的结果。

我尝试了以下任一方法来查看数据类型是否存在问题

sum_cols = F.udf(lambda x: x[0], IntegerType())
sum_cols = F.udf(lambda x: int(x[0]), IntegerType())

仍然得到空值。

我尝试删除数组：

sum_cols = F.udf(lambda x: x, IntegerType())
df_with_sum = df.withColumn('SUM_COL',sum_cols(df.A))

这工作正常并显示0/1

我尝试删除 UDF，但保留数组：

df_with_sum = df.withColumn('SUM_COL', F.array('A','B'))

这工作正常并显示一系列数组[0.0/1.0, 0.0/1.0]

因此，数组工作正常，UDF 工作正常，只是当我尝试将数组传递给 UDF 时，事情就崩溃了。我究竟做错了什么？

Answer 1

问题是你试图在一个应该输出整数的函数中返回一个双精度值，但它不适合，并且 pyspark 默认情况下会NULL在转换失败时默默地求助于：

df_with_doubles = spark.createDataFrame([(1.0,1.0), (2.0,2.0)], ['A', 'B'])
sum_cols = F.udf(lambda x: x[0]+x[1], IntegerType())
df_with_sum = df_with_double.withColumn('SUM_COL',sum_cols(F.array('A','B')))
df_with_sum.select(['SUM_COL']).toPandas()

你得到：

  SUM_COL
0    None
1    None

但是，如果您这样做：

df_with_integers = spark.createDataFrame([(1,1), (2,2)], ['A', 'B'])
sum_cols = F.udf(lambda x: x[0]+x[1], IntegerType())
df_with_sum = df_with_integers.withColumn('SUM_COL',sum_cols(F.array('A','B')))
df_with_sum.select(['SUM_COL']).toPandas()

你得到：

   SUM_COL
0        2
1        4

因此，要么预先将列转换为IntegerType（或将它们转换到 UDF 中），要么将 UDF 的返回类型更改为DoubleType.