使用数组修改Dataframe列

时间:2018-01-15 15:55:22

标签: apache-spark pyspark apache-spark-sql user-defined-functions

我有以下数据框:

+----------+ 
|col       | 
+----------+ 
|[1, 4, 3] | 
|[1, 5, 11]| 
|[1, 3, 3] | 
|[1, 4, 3] | 
|[1, 6, 3] | 
|[1, 1, 3] | 
+----------+

我想要的是:

+----------+ 
|col_new   | 
+----------+ 
|[3, -1]   | 
|[4, 6]    | 
|[2, 0]    | 
|[3, -1]   | 
|[5, -3]   | 
|[0, 2]    | 
+----------+

=> Diff运算符arr [n + 1] - arr [n]

我不知道该怎么做。

我以为我应该用udf做到这一点?我并不熟悉它,但是我试过了。

from pyspark.sql.functions import col

def diff(a):
    return [a[ii+1]-a[ii] for ii in range(a.__len__()-1)]

function = udf(lambda c: diff(c))
df.withColumn("col_new",function(col("col"))).show(20,False)

但是,因为我需要一个列表,所以当然没有用...但我想使用数据帧的强大功能...... 有人对我有暗示吗?

最佳Boendal

2 个答案:

答案 0 :(得分:3)

您的代码非常完美。唯一的错误在于你的import语句应该是

from pyspark.sql import functions as F

def diff(a):
    return [a[ii+1]-a[ii] for ii in range(a.__len__()-1)]

function = F.udf(lambda c: diff(c))
df.withColumn("col_new",function(F.col("col"))).show(20,False)

你应该都很好

<强>更新

总结一下,我建议你尝试尽可能不使用udf函数,因为它们需要数据序列化和反序列化,这肯定会降低处理效率,你应该总是这样尝试尽可能使用内置函数。

简单地说,您可以使用以下arraycol功能来满足您的要求。

from pyspark.sql import functions as F
df.select(F.array([(F.col("col")[i+1]-F.col("col")[i]) for i in range(2)]).alias("col_new")).show(20,False)

答案 1 :(得分:2)

您可以编写一个UDF,它可以在普通的python中执行您需要的操作:

def diff(array):
     res = []
     for i in range(0, len(array) -1):
             res.append(array[i+1]-array[i])
     return res

import pyspark.sql.functions as fun
f=fun.udf(diff)

这就是您将其应用于数据的方式:

d = sc.parallelize([[[1,4,3]], [[1,5,11]], [[1,3,3]]]).toDF(["col"])
d.show()
+----------+
|       col|
+----------+
| [1, 4, 3]|
|[1, 5, 11]|
| [1, 3, 3]|
+----------+

d.withColumn("new_col", f(d["col"])).drop("col").show()
+-------+
|new_col|
+-------+
|[3, -1]|
| [4, 6]|
| [2, 0]|
+-------+