具有多个参数的 PySpark UDF 返回空值

Question

我有一个PySpark Dataframe，它有两列（A，B，其类型为双），其值为0.0或1.0。我正在尝试添加一个新列，这是这两个列的总和。我遵循Pyspark中的示例：在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()

这显示了一系列的< code>NULL,而不是我期望的结果。

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

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

问题在于，您试图在一个应该输出整数的函数中返回一个double值，而这个函数并不合适，默认情况下，当转换失败时，pyspark会自动返回< code>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。