This tutorial will explain multiple workarounds to flatten (explode) 2 or more array columns in PySpark.

PySpark: Dataframe Multiple Explode

• Explode function can be used to flatten array in a column in Pyspark.
• Pyspark does not allow 2 or more explode to be present in a single select statement.
• Users will get below error if they will try to use multiple explode in a single select statement.

  Error: pyspark.sql.utils.AnalysisException: Only one generator allowed per select clause but found 2: explode(_2), explode(_3)
  

• This tutorial will explain multiple workarounds to flatten (explode) 2 or more array columns in Pyspark.
  
• Following workarounds will be explained, click on item in the below list and it will take you to the respective section of the page:
  • Using arrays_zip function (Preferred way)
  • Using zip_with function
  • Using join

Sample Data: Following 2 dataset will be used in the below examples.

Dataset 1: 

data_list = [(1, [2,3,4], [5,6,7]), (2,[3,4,5], [6,7,8])]

df= spark.createDataFrame( sc.parallelize(data_list) )

df.show()
+---+---------+---------+
| _1|       _2|       _3|
+---+---------+---------+
|  1|[2, 3, 4]|[5, 6, 7]|
|  2|[3, 4, 5]|[6, 7, 8]|
+---+---------+---------+



Dataset 2: 

data_list_2 = [(1, [2,3,4], [5,6,7]), (2,[], [6,7,8]), (3, [4,5,6], [])]
df_2= spark.createDataFrame( sc.parallelize(data_list_2) )

df_2.show()
+---+---------+---------+
| _1|       _2|       _3|
+---+---------+---------+
|  1|[2, 3, 4]|[5, 6, 7]|
|  2|       []|[6, 7, 8]|
|  3|[4, 5, 6]|       []|
+---+---------+---------+

➠ Using arrays_zip function(): array_zip function can be used along with explode function to flatten multiple columns together. This can work with n numbers of array columns. This is our preferred approach to flatten multiple array columns.

• Syntax: It can take n number of array columns as parameters and returns merged array.

  from pyspark.sql.functions import arrays_zip, explode
  
  arrays_zip(*array_cols)
  

  
  
• Example: Multiple column can be flattened using arrays_zip in 2 steps as shown in this example.
  
  
  • → Step 1: Zipping 2 arrays first and then exploding

    from pyspark.sql.functions import arrays_zip, explode
    
    exp_step1_df = df.select("_1",explode(arrays_zip("_2","_3")))
    
    exp_step1_df.show()
    +---+------+
    | _1|   col|
    +---+------+
    |  1|[2, 5]|
    |  1|[3, 6]|
    |  1|[4, 7]|
    |  2|[3, 6]|
    |  2|[4, 7]|
    |  2|[5, 8]|
    +---+------+
    

    
    
  • → Step 2: Separating zipped column using dot notation.

    
    exp_step2_df = exp_step1_df.select("_1","col._2","col._3")
    
    exp_step2_df.show()
    +---+---+---+
    | _1| _2| _3|
    +---+---+---+
    |  1|  2|  5|
    |  1|  3|  6|
    |  1|  4|  7|
    |  2|  3|  6|
    |  2|  4|  7|
    |  2|  5|  8|
    +---+---+---+
    

➠ Using zip_with function(): zip_with function can be used along with explode function to flatten multiple columns together. This function was added in Spark version 3.1.0. Use coalesce if any array column values are expected to be null else this approach will not give required output.

• Syntax: It will take 2 array columns as parameters and a function as 3rd parameter to merge 2 array columns elementwise using this function.

  from pyspark.sql.functions import zip_with, concat_ws, explode, substring_index
  
  zip_with(array_1, array_2, function)
  

  
  
• Example 1: Multiple column can be flattened using zip_with in 3 steps as shown in this example.
  
  
  • → Step 1: Zipping 2 arrays first using zip_with with concat_ws.

    from pyspark.sql.functions import zip_with, concat_ws, explode, substring_index
    
    exp_step1_df = df.select("_1", zip_with("_2", "_3", lambda x, y: concat_ws("|*|" , x, y) ).alias("join_col"))
    
    exp_step1_df.show(truncate=False)
    
    +---+---------------------+
    |_1 |join_col             |
    +---+---------------------+
    |1  |[2|*|5, 3|*|6, 4|*|7]|
    |2  |[3|*|6, 4|*|7, 5|*|8]|
    +---+---------------------+
    

    
    
  • → Step 2: Flatten zipped data using explode.

    
    exp_step2_df = exp_step1_df.select("_1", explode("join_col"))
    
    exp_step2_df.show()
    +---+-----+
    | _1|  col|
    +---+-----+
    |  1|2|*|5|
    |  1|3|*|6|
    |  1|4|*|7|
    |  2|3|*|6|
    |  2|4|*|7|
    |  2|5|*|8|
    +---+-----+
    

    
    
  • → Step 3: Separating zipped column using substring_index function.

    
    exp_step3_df = exp_step2_df.select("_1", substring_index("col","|*|",1).alias("_2"), substring_index("col","|*|",-1).alias("_3"))
    
    exp_step3_df.show()
    +---+---+---+
    | _1| _2| _3|
    +---+---+---+
    |  1|  2|  5|
    |  1|  3|  6|
    |  1|  4|  7|
    |  2|  3|  6|
    |  2|  4|  7|
    |  2|  5|  8|
    +---+---+---+
    

  
  
• Example 2: This example will use coalesce function since Array column is expected to have null values.
  
  
  • → Step 1: Zipping 2 arrays first using zip_with with concat_ws and coalesce.

    from pyspark.sql.functions import zip_with, concat_ws, explode, substring_index, coalesce, lit
    
    exp_step1_df = df_2.select("_1", zip_with("_2", "_3", lambda x, y: concat_ws("|*|" , coalesce(x,lit("")), coalesce(y,lit(""))) ).alias("join_col"))
    
    exp_step1_df.show(truncate=False)
    
    +---+---------------------+
    |_1 |join_col             |
    +---+---------------------+
    |1  |[2|*|5, 3|*|6, 4|*|7]|
    |2  |[|*|6, |*|7, |*|8]   |
    |3  |[4|*|, 5|*|, 6|*|]   |
    +---+---------------------+
    

    
    
  • → Step 2: Flatten zipped data using explode.

    
    exp_step2_df = exp_step1_df.select("_1", explode("join_col"))
    
    exp_step2_df.show()
    +---+-----+
    | _1|  col|
    +---+-----+
    |  1|2|*|5|
    |  1|3|*|6|
    |  1|4|*|7|
    |  2| |*|6|
    |  2| |*|7|
    |  2| |*|8|
    |  3| 4|*||
    |  3| 5|*||
    |  3| 6|*||
    +---+-----+
    

    
    
  • → Step 3: Separating zipped column using substring_index function.

    
    exp_step3_df = exp_step2_df.select("_1", substring_index("col","|*|",1).alias("_2"), substring_index("col","|*|",-1).alias("_3"))
    
    exp_step3_df.show()
    +---+---+---+
    | _1| _2| _3|
    +---+---+---+
    |  1|  2|  5|
    |  1|  3|  6|
    |  1|  4|  7|
    |  2|   |  6|
    |  2|   |  7|
    |  2|   |  8|
    |  3|  4|   |
    |  3|  5|   |
    |  3|  6|   |
    +---+---+---+
    

➠ Using join: Multiple array columns can be flatten individually and then joined to achieve required result. This method will not give required output if there is null values in any array columns.

• Example: Multiple column can be flattened individually and then joined again in 4 steps as shown in this example.
  
  
  • → Step 1: Flatten 1st array column using posexplode.

    from pyspark.sql.functions import posexplode
    
    exp_step1_df = df.select("_1", posexplode("_2"))
    
    exp_step1_df.show(truncate=False)
    
    +---+---+---+
    |_1 |pos|col|
    +---+---+---+
    |1  |0  |2  |
    |1  |1  |3  |
    |1  |2  |4  |
    |2  |0  |3  |
    |2  |1  |4  |
    |2  |2  |5  |
    +---+---+---+
    

    
    
  • → Step 2: Flatten 2nd array column using posexplode.

    
    exp_step2_df = df.select("_1", posexplode("_3"))
    
    exp_step2_df.show()
    +---+---+---+
    | _1|pos|col|
    +---+---+---+
    |  1|  0|  5|
    |  1|  1|  6|
    |  1|  2|  7|
    |  2|  0|  6|
    |  2|  1|  7|
    |  2|  2|  8|
    +---+---+---+
    

    
    
  • → Step 3: Join individually flatter columns using position and non array column.

    
    exp_step3_df = exp_step1_df.join( exp_step2_df,( ( exp_step1_df._1 == exp_step2_df._1 ) & ( exp_step1_df.pos == exp_step2_df.pos)))
    
    exp_step3_df.show()
    +---+---+---+---+---+---+
    | _1|pos|col| _1|pos|col|
    +---+---+---+---+---+---+
    |  1|  2|  4|  1|  2|  7|
    |  2|  2|  5|  2|  2|  8|
    |  1|  1|  3|  1|  1|  6|
    |  2|  1|  4|  2|  1|  7|
    |  1|  0|  2|  1|  0|  5|
    |  2|  0|  3|  2|  0|  6|
    +---+---+---+---+---+---+
    

    
    
  • → Step 4: Selecting the required columns from the joined dataframe.

    
    exp_step4_df = exp_step3_df.select(exp_step1_df._1, exp_step1_df.col.alias("_2"), exp_step2_df.col.alias("_3"))
    
    exp_step4_df.show()
    +---+---+---+
    | _1| _2| _3|
    +---+---+---+
    |  1|  4|  7|
    |  2|  5|  8|
    |  1|  3|  6|
    |  2|  4|  7|
    |  1|  2|  5|
    |  2|  3|  6|
    +---+---+---+