Как преобразовать два массива каждого столбца в пару для Spark DataFrame?

Question

У меня есть DataFrame, который имеет два столбца значений массива, как показано ниже

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

Думаю, я могу используйте struct, а затем udf. Но я хотел знать, есть ли какой-нибудь встроенный метод более высокого порядка, чтобы сделать это эффективно.

Answer 1

Для 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")))

Это дает мне желаемый результат.

Answer 2

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