Apache Spark -- 将UDF的结果分配给多个数据帧列

11 浏览
0 Comments

Apache Spark -- 将UDF的结果分配给多个数据帧列

我正在使用pyspark和spark-csv将一个大的csv文件加载到一个dataframe中,并且作为预处理步骤,我需要对一个包含json字符串的列中的数据应用多种操作。这些操作将返回X个值,每个值都需要存储在它们自己的单独列中。

这个功能将在一个UDF中实现。然而,我不确定如何从UDF中返回一个值列表,并将它们提供给单独的列。下面是一个简单的例子:

(...)
from pyspark.sql.functions import udf
def udf_test(n):
    return [n/2, n%2]
test_udf=udf(udf_test)
df.select('amount','trans_date').withColumn("test", test_udf("amount")).show(4)

这将产生以下结果:

+------+----------+--------------------+
|amount|trans_date|                test|
+------+----------+--------------------+
|  28.0|2016-02-07|         [14.0, 0.0]|
| 31.01|2016-02-07|[15.5050001144409...|
| 13.41|2016-02-04|[6.70499992370605...|
| 307.7|2015-02-17|[153.850006103515...|
| 22.09|2016-02-05|[11.0450000762939...|
+------+----------+--------------------+
only showing top 5 rows

在这个例子中,将两个值从UDF中返回并存储在单独的列中,最好的方法是什么?现在它们被定义为字符串:

df.select('amount','trans_date').withColumn("test", test_udf("amount")).printSchema()
root
 |-- amount: float (nullable = true)
 |-- trans_date: string (nullable = true)
 |-- test: string (nullable = true)

0