Convert an Array column to Array of Structs in PySpark dataframe

Question

I have a Dataframe containing 3 columns

| str1      | array_of_str1        | array_of_str2  |
+-----------+----------------------+----------------+
| John      | [Size, Color]        | [M, Black]     |
| Tom       | [Size, Color]        | [L, White]     |
| Matteo    | [Size, Color]        | [M, Red]       |

I want to add the Array column that contains the 3 columns in a struct type

| str1      | array_of_str1        | array_of_str2  | concat_result                                 |
+-----------+----------------------+----------------+-----------------------------------------------+
| John      | [Size, Color]        | [M, Black]     | [[[John, Size , M], [John, Color, Black]]]    |
| Tom       | [Size, Color]        | [L, White]     | [[[Tom, Size , L], [Tom, Color, White]]]      |
| Matteo    | [Size, Color]        | [M, Red]       | [[[Matteo, Size , M], [Matteo, Color, Red]]]  |

|-- AdditionalAttribute: array (nullable = true) | | |-- element: struct (containsNull = true) | | | |-- str1: string (nullable = true) | | | |-- array_of_str1: string (nullable = true) | | | |-- array_of_str2: string (nullable = true) — lazycoder
– lazycoder, Commented Jan 6, 2020 at 9:31
@lazycoder, what is AdditionalAttribute, is this a field-name of the column concat_result? is your spark 2.4+ or below? — jxc
– jxc, Commented Jan 8, 2020 at 3:09
@jxc AdditionalAttribute is the array name. I'm using Spark 2.4.3 — lazycoder
– lazycoder, Commented Jan 8, 2020 at 8:40
@lazycoder, so AdditionalAttribute is your desired column name, not concat_result shown in your post? and the new column has a schema of array of structs with 3 string fields? — jxc
– jxc, Commented Jan 8, 2020 at 12:14

Marioanzas · Accepted Answer · 2021-01-01 07:01:26Z

9

If the number of elements in the arrays in fixed, it is quite straightforward using the array and struct functions. Here is a bit of code in scala.

val result = df
    .withColumn("concat_result", array((0 to 1).map(i => struct(
                     col("str1"),
                     col("array_of_str1").getItem(i),
                     col("array_of_str2").getItem(i)
    )) : _*))

And in python, since you were asking about pyspark:

import pyspark.sql.functions as F

df.withColumn("concat_result", F.array(*[ F.struct(
                  F.col("str1"),
                  F.col("array_of_str1").getItem(i),
                  F.col("array_of_str2").getItem(i))
              for i in range(2)]))

And you get the following schema:

root
 |-- str1: string (nullable = true)
 |-- array_of_str1: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- array_of_str2: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- concat_result: array (nullable = false)
 |    |-- element: struct (containsNull = false)
 |    |    |-- str1: string (nullable = true)
 |    |    |-- col2: string (nullable = true)
 |    |    |-- col3: string (nullable = true)

edited Jan 1, 2021 at 7:01

Marioanzas

1,9452 gold badges21 silver badges38 bronze badges

answered Jan 7, 2020 at 14:43

Oli

10.5k5 gold badges31 silver badges51 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

lazycoder Over a year ago

what if the size of the array_of_str1 and array_of_str2 is variable. How can i dynamically specify the range here

Kafels · Accepted Answer · 2020-12-31 03:06:30Z

Spark >= 2.4.x

For dynamically values you can use high-order functions:

import pyspark.sql.functions as f

expr = "TRANSFORM(arrays_zip(array_of_str1, array_of_str2), x -> struct(str1, concat(x.array_of_str1), concat(x.array_of_str2)))"
df = df.withColumn('concat_result', f.expr(expr))

df.show(truncate=False)

Schema and output:

root
 |-- array_of_str1: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- array_of_str2: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- str1: string (nullable = true)
 |-- concat_result: array (nullable = true)
 |    |-- element: struct (containsNull = false)
 |    |    |-- str1: string (nullable = true)
 |    |    |-- col2: string (nullable = true)
 |    |    |-- col3: string (nullable = true)

+-------------+-------------+------+-----------------------------------------+
|array_of_str1|array_of_str2|str1  |concat_result                            |
+-------------+-------------+------+-----------------------------------------+
|[Size, Color]|[M, Black]   |John  |[[John, Size, M], [John, Color, Black]]  |
|[Size, Color]|[L, White]   |Tom   |[[Tom, Size, L], [Tom, Color, White]]    |
|[Size, Color]|[M, Red]     |Matteo|[[Matteo, Size, M], [Matteo, Color, Red]]|
+-------------+-------------+------+-----------------------------------------+

Collectives™ on Stack Overflow

Convert an Array column to Array of Structs in PySpark dataframe

2 Answers 2

1 Comment

Spark >= 2.4.x

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

Spark >= 2.4.x

Comments

Your Answer

Sign up or log in

Post as a guest

Related