How to update dataframe column which contains array of structs

user_11740220 0
Viewed 54

Schema of dataframe

  root
    |-- parentColumn: array
    |    |-- element: struct
    |    |    |-- colA: string
    |    |    |-- colB: string
    |    |    |-- colTimestamp: string

value inside dataframe look like this

"parentColumn": [
        {
            "colA": "TestA",
            "colB": "TestB",
            "colTimestamp": "2020-08-17T03:28:44.986000"
        },
        {
            "colA": "TestA",
            "colB": "TestB",
            "colTimestamp": "2020-08-17T03:28:44.986000"
        }
    ]

df.withColumn("parentColumn", ?)

Here I want to format all colTimestamp inside the array to UTC format, I saw many examples of updating values inside array but I'm not able to find a way to Update dict inside an array.

pyspark apache-spark
2 Answers

You can use the transform function to apply a function to each element of the array. Then, since you don't have lots of fields, you can recreate the struct like this:

df.withColumn("parentColumn", transform('parentColumn, x => struct(
    x.getField("colA") as "colA",
    x.getField("colB") as "colB",
    to_utc_timestamp(x.getField("colTimestamp") , "GMT+2") as "colTimestamp"
)))
user_88936860

If you're on spark 3.1+, you can use the transform function with withField within a lambda function.

spark.conf.set('spark.sql.legacy.timeParserPolicy', 'LEGACY')

data_sdf. \
    withColumn('parent_col_new', 
               func.transform('parent_col', 
                              lambda x: x.withField('col_ts', 
                                                    func.to_timestamp(x.col_ts, "yyyy-MM-dd'T'HH:mm:ss")
                                                    )
                              )
               ). \
    show(truncate=False)

# +---+---------------------------------------------------------------------------------------+-------------------------------------------------------------------------+
# |id |parent_col                                                                             |parent_col_new                                                           |
# +---+---------------------------------------------------------------------------------------+-------------------------------------------------------------------------+
# |1  |[{test, testB, 2020-08-17T03:28:44.986000}, {UNREAD, USER, 2020-08-17T03:28:44.986000}]|[{test, testB, 2020-08-17 03:28:44}, {UNREAD, USER, 2020-08-17 03:28:44}]|
# +---+---------------------------------------------------------------------------------------+-------------------------------------------------------------------------+

# root
#  |-- id: integer (nullable = false)
#  |-- parent_col: array (nullable = false)
#  |    |-- element: struct (containsNull = false)
#  |    |    |-- col_a: string (nullable = true)
#  |    |    |-- col_b: string (nullable = true)
#  |    |    |-- col_ts: string (nullable = true)
#  |-- parent_col_new: array (nullable = false)
#  |    |-- element: struct (containsNull = false)
#  |    |    |-- col_a: string (nullable = true)
#  |    |    |-- col_b: string (nullable = true)
#  |    |    |-- col_ts: timestamp (nullable = true)

If withField and/or transform isn't available in your spark version, you can use expr and recreate the struct. It'll result in the same output.

data_sdf. \
    withColumn('parent_col_new',
               func.expr('''
                         transform(parent_col, 
                                   x -> struct(x.col_a as col_a, 
                                               x.col_b as col_b, 
                                               to_timestamp(x.col_ts, "yyyy-MM-dd'T'HH:mm:ss") as col_ts
                                               )
                                   )
                         ''')
               ). \
    show(truncate=False)
user_82795850
Related