Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Split large array columns into multiple columns - Pyspark

Tags:

pyspark

I have:

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

I want:

+---+-------+-------+-------+-------+-------+-------+
| 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|
+---+-------+-------+-------+-------+-------+-------+

The solution provided by How to split a list to multiple columns in Pyspark?

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

works, but some of my arrays are very long (max 332).

How can I write this so that it takes account of all length arrays?

like image 726
Microsim Avatar asked Aug 02 '18 07:08

Microsim


People also ask

How do you split columns in PySpark?

The PySpark SQL provides the split() function to convert delimiter separated String to an Array (StringType to ArrayType) column on DataFrame It can be done by splitting the string column on the delimiter like space, comma, pipe, etc.

How do you explode an array in PySpark?

To split multiple array column data into rows pyspark provides a function called explode(). Using explode, we will get a new row for each element in the array.

How do I split a single column into multiple columns in Python?

We can use str. split() to split one column to multiple columns by specifying expand=True option. We can use str. extract() to exract multiple columns using regex expression in which multiple capturing groups are defined.


1 Answers

This solution will work for your problem, no matter the number of initial columns and the size of your arrays. Moreover, if a column has different array sizes (eg [1,2], [3,4,5]), it will result in the maximum number of columns with null values filling the gap.

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|
+---+-------+-------+-------+-------+-------+-------+
like image 104
Pierre Gourseaud Avatar answered Sep 28 '22 03:09

Pierre Gourseaud