pysparkデータフレームのcollect_setをフラットマップ

Question

私は2つのデータフレームを持っており、groupbyを使用した後、aggでcollect_set（）を使用しています。集計後に結果の配列をフラット・マップする最良の方法は何ですか？

schema = ['col1', 'col2', 'col3', 'col4'] 

a = [[1, [23, 32], [11, 22], [9989]]] 

df1 = spark.createDataFrame(a, schema=schema) 

b = [[1, [34], [43, 22], [888, 777]]] 

df2 = spark.createDataFrame(b, schema=schema) 

df = df1.union(
     df2 
    ).groupby(
     'col1' 
    ).agg(
     collect_set('col2').alias('col2'), 
     collect_set('col3').alias('col3'), 
     collect_set('col4').alias('col4') 
    ) 

df.collect() 

私は出力

[Row(col1=1, col2=[[34], [23, 32]], col3=[[11, 22], [43, 22]], col4=[[9989], [888, 777]])] 

としてこれを取得しています。しかし、私はあなたがudfを使用することができ、出力

[Row(col1=1, col2=[23, 32, 34], col3=[11, 22, 43], col4=[9989, 888, 777])] 

Answer 1

としてこれをしたい：

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

flatten = udf(lambda x: list(chain.from_iterable(x)), ArrayType(IntegerType())) 

df.withColumn('col2_flat', flatten('col2'))