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

作为输入，我有一个包含int值的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))

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

实函数：

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))

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

对于1。示例：您可以使用SQL sum，如：

但是，我需要的是将这些函数转换为Spark UDF的标准解决方案

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

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

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

但它总是在以下方面失败：

TypeError:参数无效，不是字符串或列： 类型的行（值=1114.0）。专栏 文字，使用“lit”、“数组”、“结构”或“创建映射”函数

是否可以使用这些功能转换数据？

数据帧示例：

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))

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

预期输出：

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))

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函数，应该很清楚。

您需要的是groupby，并使用

collect\u list

将所有整数值放入数组列，然后在该列上应用您的自定义项。此外，还需要从

calc\u rms

显式返回浮点值：

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|
#+--------+