PySpark udf returns null when function works in Pandas dataframe

Question

我正在尝试创建一个用户定义的 function ，它采用数组的累积和并将该值与另一列进行比较。 这是一个可重现的示例：

from pyspark.sql.session import SparkSession

# instantiate Spark
spark = SparkSession.builder.getOrCreate()

# make some test data
columns = ['loc', 'id', 'date', 'x', 'y']
vals = [
    ('a', 'b', '2016-07-01', 1, 5),
    ('a', 'b', '2016-07-02', 0, 5),
    ('a', 'b', '2016-07-03', 5, 15),
    ('a', 'b', '2016-07-04', 7, 5),
    ('a', 'b', '2016-07-05', 8, 20),
    ('a', 'b', '2016-07-06', 1, 5)
]

# create DataFrame
temp_sdf = (spark
      .createDataFrame(vals, columns)
      .withColumn('x_ary', collect_list('x').over(Window.partitionBy(['loc','id']).orderBy(desc('date')))))

temp_df = temp_sdf.toPandas()

def test_function(x_ary, y):
  cumsum_array = np.cumsum(x_ary) 
  result = len([x for x in cumsum_array if x <= y])
  return result

test_function_udf = udf(test_function, ArrayType(LongType()))

temp_df['len'] = temp_df.apply(lambda x: test_function(x['x_ary'], x['y']), axis = 1)
display(temp_df)

在 Pandas 中，这是 output：

loc id  date        x   y   x_ary           len
a   b   2016-07-06  1   5   [1]             1
a   b   2016-07-05  8   20  [1,8]           2
a   b   2016-07-04  7   5   [1,8,7]         1
a   b   2016-07-03  5   15  [1,8,7,5]       2
a   b   2016-07-02  0   5   [1,8,7,5,0]     1
a   b   2016-07-01  1   5   [1,8,7,5,0,1]   1

在 Spark 中使用temp_sdf.withColumn('len', test_function_udf('x_ary', 'y')) ，所有len最终都是null 。

有人知道为什么会这样吗？

Also, replacing cumsum_array = np.cumsum(np.flip(x_ary)) fails in pySpark with error AttributeError: module 'numpy' has no attribute 'flip' , but I know it exists as I can run it fine with Pandas dataframe.
这个问题可以解决吗，或者有没有更好的方法用 pySpark 翻转 arrays？

在此先感谢您的帮助。

Answer 1

由于 test_function 返回 integer 而不是列表/数组。 您将获得 null 值，因为您提到了错误的返回类型。 因此，请删除“来自 udf 的 ArrayType”或将返回类型替换为 LongType()，然后它将按以下方式工作。 ：

注意：您可以选择设置 UDF 的返回类型，否则默认返回类型是 StringType。

选项1：

test_function_udf = udf(test_function) # Returns String type

选项2：

test_function_udf = udf(test_function, LongType())  #Returns Long/integer type

temp_sdf = temp_sdf.withColumn('len', 
           test_function_udf('x_ary', 'y'))
temp_sdf.show()