PySpark UDF с несколькими аргументами возвращает ноль

Question

У меня есть PySpark Dataframe с двумя столбцами (A, B, тип которых double), значения которых 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()

Это показывает серию NULL s вместо ожидаемых результатов.

Я пыталсяЛюбое из следующего, чтобы увидеть, есть ли проблема с типами данных

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

, все еще получающих Null.

Я попытался удалить массив:

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 по умолчанию беззвучно обращается к 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.