使用数组修改数据框列

Question

我有以下数据框：

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

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

而且我不知道该怎么做。

我想我应该用 udf 来做？我不太熟悉它，但是我尝试过。

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

但是，是的，这当然不起作用，因为我需要一个列表......但我想使用数据框的力量......有人给我提示吗？

最佳博恩达尔

Answer 1

你的代码很完美。唯一的错误是你的导入声明应该是

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函数，因为它们需要数据序列化和反序列化，这肯定会降低处理效率，您应该尽可能尝试使用内置函数。

所以简单地说您可以使用array如下col功能来满足您的要求。

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)