Как объединить два столбца ArrayType (StringType ()) поэлементно в Pyspark?

Question

У меня есть два ArrayType(StringType()) столбца в кадре данных искры, и я хочу объединить два столбца поэлементно:

input :

+-------------+-------------+
|col1         |col2         |
+-------------+-------------+
|['a','b']    |['c','d']    |
|['a','b','c']|['e','f','g']|
+-------------+-------------+

вывод :

+-------------+-------------+----------------+
|col1         |col2         |col3            |
+-------------+-------------+----------------+
|['a','b']    |['c','d']    |['ac', 'bd']    |
|['a','b','c']|['e','f','g']|['ae','bf','cg']|
+-------------+----------- -+----------------+

Спасибо.

Answer 1

Для Spark 2.4+ вы можете использовать функцию transform следующим образом:

col3_expr = "transform(col1, (x, i) -> concat(x, col2[i]))"
df.withColumn("col3", expr(col3_expr)).show()

Функция transform принимает в качестве параметров первый столбец массива col1, выполняет итерации поверх его элементов и применяет лямбда-функцию (x, i) -> concat(x, col2[i]), где x фактический элемент и i его индекс, используемый для получения соответствующего элемента из массива col2.

Дает:

+------+------+--------+
|  col1|  col2|    col3|
+------+------+--------+
|[a, b]|[c, d]|[ac, bd]|
+------+------+--------+

Или еще проще, используя функцию высшего порядка zip_with:

df.withColumn("col3", expr("zip_with(col1, col2, (x, y) -> concat(x, y))"))

Answer 2

Вот альтернативный ответ, который можно использовать для обновленного неоригинального вопроса. Использует array и array_except, чтобы продемонстрировать использование таких методов. Принятый ответ более элегантный.

from pyspark.sql.functions import *
from pyspark.sql.types import *

# Arbitrary max number of elements to apply array over, need not broadcast such a small amount of data afaik.
max_entries = 5 

# Gen in this case numeric data, etc. 3 rows with 2 arrays of varying length,but per row constant length. 
dfA = spark.createDataFrame([   ( list([x,x+1,4, x+100]), 4, list([x+100,x+200,999,x+500])   ) for x in range(3)], ['array1', 'value1', 'array2'] ).withColumn("s",size(col("array1")))    
dfB = spark.createDataFrame([   ( list([x,x+1]), 4, list([x+100,x+200])   ) for x in range(5)], ['array1', 'value1', 'array2'] ).withColumn("s",size(col("array1"))) 
df = dfA.union(dfB)

# concat the array elements which are variable in size up to a max amount.
df2 = df.select(( [concat(col("array1")[i], col("array2")[i]) for i in range(max_entries)]))
df3 = df2.withColumn("res", array(df2.schema.names))

# Get results but strip out null entires from array.
df3.select(array_except(df3.res, array(lit(None)))).show(truncate=False)

Не удалось получить значение s столбца, который должен быть передан в диапазон.

Возвращает:

+------------------------------+
|array_except(res, array(NULL))|
+------------------------------+
|[0100, 1200, 4999, 100500]    |
|[1101, 2201, 4999, 101501]    |
|[2102, 3202, 4999, 102502]    |
|[0100, 1200]                  |
|[1101, 2201]                  |
|[2102, 3202]                  |
|[3103, 4203]                  |
|[4104, 5204]                  |
+------------------------------+

Answer 3

Вот общий ответ c. Просто посмотрите на Res для результата. 2 массива одинакового размера, т.е. n элементов для обоих.

from pyspark.sql.functions import *
from pyspark.sql.types import *

# Gen in this case numeric data, etc. 3 rows with 2 arrays of varying length, but both the same length as in your example
df = spark.createDataFrame([   ( list([x,x+1,4, x+100]), 4, list([x+100,x+200,999,x+500])   ) for x in range(3)], ['array1', 'value1', 'array2'] )    
num_array_elements = len(df.select("array1").first()[0])

# concat
df2 = df.select(([ concat(col("array1")[i], col("array2")[i]) for i in range(num_array_elements)]))
df2.withColumn("res", array(df2.schema.names)).show(truncate=False)

возвращает:

Answer 4

Это не будет масштабироваться, но вы можете получить записи 0th и 1st в каждом массиве, а затем сказать, что col3 - это a[0] + b[0], а затем a[1] + b[1]. Сделайте все 4 записи отдельными значениями и затем выведите их вместе.