使用麻木数组输入从 python 方法创建 PySpark UDF，以计算并返回单个浮点值

Question

作为输入，我有一个包含整数值的csv文件。

spark_df = spark.read.option("header", "false").csv("../int_values.csv")

df = spark_df.selectExpr("_c0 as something")

_df = df.withColumn("values", df.something.cast(FloatType())).select("values")

我还有一些为numpy数组输入设计的python函数，我需要在Spark数据帧上应用这些函数。

示例一：

def calc_sum(float_array):
    return np.sum(float_array)

实函数：

def calc_rms(float_array):
    return np.sqrt(np.mean(np.diff(float_array)**2))

对于1。示例您可以像这样使用SQL sum:

_df.groupBy().sum().collect()

但是，我需要的是一个将这些函数转换成Spark UDFs的标准解决方案

我尝试了很多方法，比如：

udf_sum = udf(lambda x : calc_sum(x), FloatType())

_df.rdd.flatMap(udf_sum).collect()

但是它总是失败:

TypeError:参数无效，不是类型的字符串或列:行(值=1114.0)

有没有可能用这些函数来转换数据？

数据帧示例:

In [6]: spark_df.show()
+----+
| _c0|
+----+
|1114|
|1113|
|1066|
|1119|
|1062|
|1089|
|1093|
| 975|
|1099|
|1062|
|1062|
|1162|
|1057|
|1123|
|1141|
|1089|
|1172|
|1096|
|1164|
|1146|
+----+
only showing top 20 rows

预期输出：

从 UDF 返回的浮点值。

对于 Sum 函数，它应该很清楚。

Answer 1

您需要的是groupby，并使用＜code＞collect_list＜/code＞将所有整数值获取到数组列中，然后在该列上应用您的UDF。此外，您还需要显式地从＜code＞calc_rms＜code＞返回float：

from pyspark.sql import functions as F
from pyspark.sql.types import FloatType


def calc_rms(float_array):
    return float(np.sqrt(np.mean(np.diff(float_array) ** 2)))


calc_rms_udf = F.udf(calc_rms, FloatType())


df.groupby().agg(F.collect_list("_c0").alias("_c0")) \
    .select(calc_rms_udf(F.col("_c0")).alias("rms")) \
    .show()

#+--------+
#|     rms|
#+--------+
#|67.16202|
#+--------+