PySparkでArrayTypeの中身を一個づつcolumnにする

どういう時に使うの？

主にPysparkにおいてArrayTypeを保持しているDataFrameをcsvに出力したいときに使うと思います。

イメージ

input

ID	array
a	[1, 2, 3]
b	[4, 5, 6]

上記のようなDataFrameを

output

ID	num1	num2	num3
a	1	2	3
b	4	5	6

こんな感じにArrayの中身を分解するイメージです。

サンプルコード

sample.py

import pyspark.sql.functions as F

hoge = ''
top_N = 10 
# Arrayの中身を左から何個取るか。 1 ≦ top_N ≦ len(array)

for i in range(top_N):
    exec(f'top{i + 1} = F.udf(lambda x : x[{i}])')
    hoge += f'.withColumn("num{i + 1}", top{i + 1}("array"))'
exec(f'df2 = df1{hoge}.drop("array")')

udfでarrayのN番目を取ってきてwithColumnで列を作るだけなのですが、
列の分だけudfの関数を定義し、withColumnを書かなければならなかったのを、execを用いることで自動化しています。

top_Nを変えることで、arrayの中身を左から5個取ってきてTOP5だけ、みたいなことも出来ます。