結合PySpark數據幀數組類型的字段爲單數組類型字段

Question

我有一個PySpark數據幀與2名數組類型的字段：

>>>df 
DataFrame[id: string, tokens: array<string>, bigrams: array<string>] 
>>>df.take(1) 
[Row(id='ID1', tokens=['one', 'two', 'two'], bigrams=['one two', 'two two'])] 

我想將它們合併成一個數組類型字段：

>>>df2 
DataFrame[id: string, tokens_bigrams: array<string>] 
>>>df2.take(1) 
[Row(id='ID1', tokens_bigrams=['one', 'two', 'two', 'one two', 'two two'])] 

語法與字符串工作似乎並沒有在這裏工作：

df2 = df.withColumn('tokens_bigrams', df.tokens + df.bigrams) 

謝謝！

Answer 1

遺憾的是並置在一般的情況下array列你需要一個UDF，例如像這樣：

from itertools import chain 
from pyspark.sql.functions import col, udf 
from pyspark.sql.types import * 

def concat(type): 
    def concat_(*args): 
     return list(chain(*args)) 
    return udf(concat_, ArrayType(type)) 


concat_string_arrays = concat(StringType()) 

df.select(concat_string_arrays(col("tokens"), col("bigrams")))