将大型数组列拆分为多个列 - Pyspark

Question

我有：

+---+-------+-------+
| id|   var1|   var2|
+---+-------+-------+
|  a|[1,2,3]|[1,2,3]|
|  b|[2,3,4]|[2,3,4]|
+---+-------+-------+

我想要：

+---+-------+-------+-------+-------+-------+-------+
| id|var1[0]|var1[1]|var1[2]|var2[0]|var2[1]|var2[2]|
+---+-------+-------+-------+-------+-------+-------+
|  a|      1|      2|      3|      1|      2|      3|
|  b|      2|      3|      4|      2|      3|      4|
+---+-------+-------+-------+-------+-------+-------+

由如何在Pyspark中将列表拆分为多个列？提供的解决方案。

df1.select('id', df1.var1[0], df1.var1[1], ...).show()

“工作正常，但我的一些数组非常长（最大为332）。

我该如何编写代码以考虑所有长度的数组？”

- Microsim

---

在您分享的链接中，显示了df2.select('key', 'value.*').show()这一行代码；这有帮助吗？ - Cleb

不幸的是，它只适用于StructType变量，而我想将其保留为数组。 - Microsim

Answer 1

无论您的初始列数和数组大小如何，此解决方案都将适用于您的问题。此外，如果某一列具有不同的数组大小（例如[1,2]，[3,4,5]），则将产生最大数量的带有空值的列以填补差距。

from pyspark.sql import functions as F

df = spark.createDataFrame(sc.parallelize([['a', [1,2,3], [1,2,3]], ['b', [2,3,4], [2,3,4]]]), ["id", "var1", "var2"])

columns = df.drop('id').columns
df_sizes = df.select(*[F.size(col).alias(col) for col in columns])
df_max = df_sizes.agg(*[F.max(col).alias(col) for col in columns])
max_dict = df_max.collect()[0].asDict()

df_result = df.select('id', *[df[col][i] for col in columns for i in range(max_dict[col])])
df_result.show()
>>>
+---+-------+-------+-------+-------+-------+-------+
| id|var1[0]|var1[1]|var1[2]|var2[0]|var2[1]|var2[2]|
+---+-------+-------+-------+-------+-------+-------+
|  a|      1|      2|      3|      1|      2|      3|
|  b|      2|      3|      4|      2|      3|      4|
+---+-------+-------+-------+-------+-------+-------+

- Pierre Gourseaud

---

你可以使用生成器表达式来代替使用列表推导式在内存中构建列表，例如：df.select(*(F.size(col).alias(col) for col in columns))。 - trianta2