Creating PySpark UDFs from python method with numpy array input, to calculate and return a single float value

user_15168930 0
Viewed 839

As input I have a csv file with int values in it.

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

I also have some python functions designed for numpy array inputs, that I need to apply on the Spark DataFrame.

The example one:

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

Real function:

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

For the 1. example you can use SQL sum like:

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

But, what I need is a standard solution to transform these functions into Spark UDFs

I tried many ways, like:

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

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

but it always failed with:

TypeError: Invalid argument, not a string or column: Row(values=1114.0) of type <class 'pyspark.sql.types.Row'>. For column literals, use 'lit', 'array', 'struct' or 'create_map' function.

Is it possible to transform the data in a way that works with these functions?

DataFrame sample:

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

Expected output:

A Float value returned from the UDF.

For the Sum function it should be clear.

pyspark-dataframes pyspark apache-spark-sql user-defined-functions numpy
1 Answers

What you want is groupby and use collect_list to get all integer values into an array column then apply your UDF on that column. Also, you need to explicitly return float from 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|
#+--------+
user_13865510
Related