时间戳

首页 > 解决方案 > 使用 numpy 数组输入从 python 方法创建 PySpark UDF，以计算并返回单个浮点值

numpy - 使用 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 DataFrame 上应用它们。

示例一：

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 UDF 的标准解决方案

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

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

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

但它总是失败：

TypeError：参数无效，不是字符串或列：<class 'pyspark.sql.types.Row'> 类型的行（值 = 1114.0）。对于列文字，使用“lit”、“array”、“struct”或“create_map”函数。

是否可以以与这些功能一起使用的方式转换数据？

数据框示例：

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

标签： numpypysparkapache-spark-sqluser-defined-functionspyspark-dataframes

解决方案

您想要的是 groupby 并用于collect_list将所有整数值放入数组列中，然后在该列上应用您的 UDF。此外，您需要从以下位置显式返回浮点数calc_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|
#+--------+

推荐阅读