python - PySpark convert struct field inside array to string

Question

I have a dataframe with schema like this:

|-- order: string (nullable = true)
|-- travel: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- place: struct (nullable = true)
 |    |    |    |-- name: string (nullable = true)
 |    |    |    |-- address: string (nullable = true)
 |    |    |    |-- latitude: double (nullable = true)
 |    |    |    |-- longitude: double (nullable = true)
 |    |    |-- distance_in_kms: float (nullable = true)
 |    |    |-- estimated_time: struct (nullable = true)
 |    |    |    |-- seconds: long (nullable = true)
 |    |    |    |-- nanos: integer (nullable = true)

I want to get the seconds in estimated_time and convert it into a string and concatenate it with s, and then replace estimated_time with the new string value. For example, { "seconds": "988", "nanos": "102" } will be converted to 988s, so the schema will change to

|-- order: string (nullable = true)
|-- travel: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- place: struct (nullable = true)
 |    |    |    |-- name: string (nullable = true)
 |    |    |    |-- address: string (nullable = true)
 |    |    |    |-- latitude: double (nullable = true)
 |    |    |    |-- longitude: double (nullable = true)
 |    |    |-- distance_in_kms: float (nullable = true)
 |    |    |-- estimated_time: string (nullable = true)

How can I do this in PySpark?

More concrete example, I want to transform this DF (visualized in JSON)

{
    "order": "c-331",
    "travel": [
        {
            "place": {
                "name": "A place",
                "address": "The address",
                "latitude": 0.0,
                "longitude": 0.0
            },
            "distance_in_kms": 1.0,
            "estimated_time": {
                "seconds": 988,
                "nanos": 102
            }
        }
    ]
}

into

{
    "order": "c-331",
    "travel": [
        {
            "place": {
                "name": "A place",
                "address": "The address",
                "latitude": 0.0,
                "longitude": 0.0
            },
            "distance_in_kms": 1.0,
            "estimated_time": "988s"
        }
    ]
}

Answer 1

You can do this with the following pyspark functions:

• withColumn lets you create a new column. We will use this to extract "estimated_time"
• concat concatenates string columns
• lit creates a column of a given string

Please have a look at the following example:

from pyspark.sql import functions as F
j = '{"order":"c-331","travel":[{"place":{"name":"A place","address":"The address","latitude":0.0,"longitude":0.0},"distance_in_kms":1.0,"estimated_time":{"seconds":988,"nanos":102}}]}'
df = spark.read.json(sc.parallelize([j]))

#the following command creates a new column called estimated_time2 which contains the values of travel.estimated_time.seconds concatenated with a 's' 
bla = df.withColumn('estimated_time2', F.concat(df.travel.estimated_time.seconds[0].cast("string"), F.lit("s")))

#unfortunately it is currently not possible to use withColumn to add a new member to a struct. Therefore the following command replaces 'travel.estimated_time' with the before created column estimated_time2
bla = bla.select("order"
                , F.array(
                    F.struct(
                        bla.travel.distance_in_kms[0].alias("distance_in_kms")
                        ,bla.travel.place[0].alias("place")
                        , bla.estimated_time2.alias('estimated_time')
                        )).alias("travel"))

bla.show(truncate=False)
bla.printSchema()

And that is the output:

+-----+------------------------------------------+ 
|order|travel                                    | 
+-----+------------------------------------------+ 
|c-331|[[1.0,[The address,0.0,0.0,A place],988s]]| 
+-----+------------------------------------------+


root 
|-- order: string (nullable = true) 
|-- travel: array (nullable = false) 
| |-- element: struct (containsNull = false) 
| | |-- distance_in_kms: double (nullable = true)
| | |-- place: struct (nullable = true) 
| | | |-- address: string (nullable = true) 
| | | |-- latitude: double (nullable = true) 
| | | |-- longitude: double (nullable = true) 
| | | |-- name: string (nullable = true) 
| | |-- estimated_time: string (nullable = true)