Pyspark中两个数组相减得到一个新数组

Question

我是 Spark 的新手。我可以在 python Pandas&Numpy 中对数组进行求和、减法或乘法。但是我在 Spark (python) 中做类似的事情有困难。我在 Databricks 上。

例如，这种方法给出了一个巨大的错误消息，我不想在这里复制粘贴：

differencer=udf(lambda x,y: x-y, ArrayType(FloatType()))

df.withColumn('difference', differencer('Array1', 'Array2'))

架构如下所示：

root
 |-- col1: integer (nullable = true)
 |-- time: timestamp (nullable = true)
 |-- num: integer (nullable = true)
 |-- part: integer (nullable = true)
 |-- result: integer (nullable = true)
 |-- Array1: array (nullable = true)
 |    |-- element: float (containsNull = true)
 |-- Array2: array (nullable = false)
 |    |-- element: float (containsNull = true)

我只想创建一个减去这两个数组列的新列。其实，我会得到RMSE他们之间的。但是我认为一旦我学会了如何获得这种差异，我就可以处理它。

数组看起来像这样（我只是输入一些整数）：

Array1_row1[5, 4, 2, 4, 3] Array2_row1[4, 3, 1, 2, 1]

所以 row1 的结果数组应该是： DiffCol_row1[1, 1, 1, 2, 2]

感谢您的建议或指导。谢谢你。

Answer 1

你可以zip_arrays和transform

from pyspark.sql.functions import expr

df = spark.createDataFrame(
    [([5, 4, 2, 4, 3], [4, 3, 1, 2, 1])], ("array1", "array2")
) 

df.withColumn(
    "array3", 
    expr("transform(arrays_zip(array1, array2), x -> x.array1 - x.array2)")
).show()                                                                         
# +---------------+---------------+---------------+       
# |         array1|         array2|         array3|
# +---------------+---------------+---------------+
# |[5, 4, 2, 4, 3]|[4, 3, 1, 2, 1]|[1, 1, 1, 2, 2]|
# +---------------+---------------+---------------+

有效udf将需要等效的逻辑，即

from pyspark.sql.functions import udf

@udf("array<double>")
def differencer(xs, ys):
    if xs and ys:
        return [float(x - y) for x, y in zip(xs, ys)]

df.withColumn("array3", differencer("array1", "array2")).show()
# +---------------+---------------+--------------------+
# |         array1|         array2|              array3|
# +---------------+---------------+--------------------+
# |[5, 4, 2, 4, 3]|[4, 3, 1, 2, 1]|[1.0, 1.0, 1.0, 2...|
# +---------------+---------------+--------------------+