如何将每列的两个数组转换为 Spark 数据帧的一对数组?

Question

我有一个数据帧，它有两列数组值，如下所示

var ds = Seq((Array("a","b"),Array("1","2")),(Array("p","q"),Array("3","4")))
var df = ds.toDF("col1", "col2")
+------+------+
|  col1|  col2|
+------+------+
|[a, b]|[1, 2]|
|[p, q]|[3, 4]|
+------+------+

我想将其转换为如下所示的对数组

+------+------+---------------+
|  col1|  col2|           col3|
+------+------+---------------+
|[a, b]|[1, 2]|[[a, 1],[b, 2]]|
|[p, q]|[3, 4]|[[p, 3],[q, 4]]|
+------+------+---------------+

我想我可以使用结构，然后使用一些 udf。但我想知道是否有任何内置的高阶方法来有效地做到这一点。

Answer 1

从Spark-2.4使用arrays_zip函数。

Example:

df.show()
#+------+------+
#|  col1|  col2|
#+------+------+
#|[a, b]|[1, 2]|
#|[p, q]|[3, 4]|
#+------+------+
from pyspark.sql.functions import *
df.withColumn("col3",arrays_zip(col("col1"),col("col2"))).show()
#+------+------+----------------+
#|  col1|  col2|            col3|
#+------+------+----------------+
#|[a, b]|[1, 2]|[[a, 1], [b, 2]]|
#|[p, q]|[3, 4]|[[p, 3], [q, 4]]|
#+------+------+----------------+

Answer 2

对于Spark-2.3或以下，我发现迭代器 zip 方法对于这个用例非常方便(我在发布问题时没有意识到这一点(。我可以定义一个小的 UDF

val zip = udf((xs: Seq[String], ys: Seq[String]) => xs.zip(ys))

并用作

var out = df.withColumn("col3", zip(df("col1"), df("col2")))

这给了我想要的结果。