如何在 Pyspark UDF 中返回双精度列表？

Posted 2023-04-13

【中文标题】如何在 Pyspark UDF 中返回双精度列表？

【英文标题】：How to return a list of double in a Pyspark UDF?

【发布时间】：2019-11-12 21:27:36

【问题描述】：

from pyspark.sql import functions as func

我有一个名为 df 的 Pyspark 数据框。它具有以下架构：

id: string
item: string
data: double

我对其应用以下操作：

grouped_df = df.groupBy(["id", "item"]).agg(func.collect_list(df.data).alias("dataList"))

另外，我定义了用户定义函数iqrOnList：

@udf
def iqrOnList(accumulatorsList: list):
  import numpy as np 

  Q1 = np.percentile(accumulatorsList, 25)
  Q3 = np.percentile(accumulatorsList, 75) 
  IQR = Q3 - Q1

  lowerFence = Q1 - (1.5 * IQR)
  upperFence = Q3 + (1.5 * IQR)

  return [elem if (elem >= lowerFence and elem <= upperFence) else None for elem in accumulatorsList]

我是这样使用这个UDF的：

grouped_df = grouped_df.withColumn("SecondList", iqrOnList(grouped_df.dataList))

这些操作在输出中返回数据帧grouped_df，如下所示：

id: string
item: string
dataList: array
SecondList: string

问题：

SecondList 具有我所期望的完全正确的值（例如 [1, 2, 3, null, 3, null, 2]），但返回类型错误（string 而不是 array，即使它保留了它的形式）。

问题是我需要将其存储为array，与dataList 完全相同。

问题：

1) 我怎样才能用正确的类型保存它？

2) 这个UDF 在性能方面是昂贵的。 我读到here，Pandas UDF 的性能比普通 UDF 好得多。 Pandas UDF 中这个方法的等价物是什么？

额外问题（优先级较低）： func.collect_list(df.data) 不收集 null 的值，df.data 拥有这些值。我也想收藏，没有replacing all null values with another default value怎么办？

【问题讨论】：

定义udf时必须指定返回类型，否则默认为StringType。你可能想要ArrayType(DoubleType())

【参考方案1】：

你仍然可以使用你当前的语法，只需要在注解声明中提供返回类型

import pyspark.sql.types as Types
@udf(returnType=Types.ArrayType(Types.DoubleType()))

【讨论】：

解决了类型问题。你知道如何将 UDF 写成 Pandas UDF 吗？