Kết hợp các lĩnh vực PySpark DataFrame ArrayType vào lĩnh vực ArrayType đơn

Question

Tôi có một PySpark DataFrame với 2 lĩnh vực ArrayType:

>>>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'])] 

Tôi muốn kết hợp chúng thành một lĩnh vực ArrayType duy nhất:

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

Cú pháp hoạt động với các chuỗi dường như không hoạt động tại đây:

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

Cảm ơn!

Answer 1

Thật không may để nối array cột trong trường hợp chung, bạn sẽ cần một UDF, ví dụ như thế này:

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")))