apache-spark - 在 Pyspark 中减去两个数组以获得一个新数组

Question

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

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

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

df.withColumn('差异', 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)

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

数组看起来像这样(我只是输入一些整数):

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

所以第 1 行的结果数组应该是: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...|
# +---------------+---------------+--------------------+