将pySpark将array列展开为pandas单独列

写在前面的话

这一系列主要记录我工作中遇到的一些问题，以及相应的解决办法。希望对大家能有帮助。

pyspark列展开

假设我们的数据中有4列，col_A,col_B,col_C和col_D，其中列A和列B都是数组列，我们需要将这两列进行展开，数组中的每一维都展开为一列。代码如下：

#获取数据中的第一行，主要是为了获取数组长度
first_row = raw_data.first()
#数组列
array_cols = ['col_A', 'col_B']
#存储展开后的列名
columns = []
columns.append("col_C")
columns.append("col_D")
for col in array_cols:
	#对于每一个数组列首先计算该数组的长度
    numAttrs = len(first_row[col])
    print("新增列的个数", numAttrs)
    #生成每一维的名字，attrs表示展开后每一维的名字
    attrs = spark.sparkContext.parallelize([col + str(i) for i in range(numAttrs)]).zipWithIndex().collect()
    #将每一维都展开展开为单独的一维，
    for name, index in attrs:
        raw_data = raw_data.withColumn(name, raw_data[col].getItem(index))
        columns.append(name)
transData = raw_data[columns]

之后我们可以使用toPandas（）,方法直接将spark的dataframe转化为pandas().