Convert array to struct in dataframe

Question

In my dataframe, I need to convert an array data type column to struct. I can manually do that with a sample of data (by modifying in editor) and it is the data that I need. I need to do it in PySpark.

Input dataframe schema:

root
 |-- id: string (nullable = true)
 |-- description: string (nullable = true)
 |-- documents: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- id: string (nullable = true)
 |    |    |-- doc_name: string (nullable = true)
 |    |    |-- obligations: struct (containsNull = true)
 |-- contacts: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- id: string (nullable = true)
 |    |    |-- contact_first_name: string (nullable = true)
 |    |    |-- contact_last_name: string (nullable = true)

Data:

{
   "id":"123",
   "description": "agreement",
   "documents":[
     {
       "id":"doc_id_1",
       "doc_name":"doc_name_1",
       "obligations":{}
     }
   ],
   "contacts":[
    {
      "id":"contact_id_1",
      "contact_first_name":"John",
      "contact_last_name":"Doe"
    }
  ]
}

Schema that I need:

root
 |-- id: string (nullable = true)
 |-- description: string (nullable = true)
 |-- documents: struct (containsNull = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- id: string (nullable = true)
 |    |    |-- doc_name: string (nullable = true)
 |    |    |-- obligations: struct (containsNull = true)
 |-- contacts: struct (containsNull = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- id: string (nullable = true)
 |    |    |-- contact_first_name: string (nullable = true)
 |    |    |-- contact_last_name: string (nullable = true)

Data that I need:

{
   "id":"123",
   "description": "agreement",
   "documents":{
     {
       "id":"doc_id_1",
       "doc_name":"doc_name_1",
       "obligations":{}
     }
   },
   "contacts":{
    {
      "id":"contact_id_1",
      "contact_first_name":"John",
      "contact_last_name":"Doe"
    }
  }
}

ZygD · Accepted Answer · 2022-10-19 20:51:51Z

Arrays differ from structs in a way that arrays can hold many items. In your current setup, you have an array of structs - that array may potentially hold many structs. Only if you are sure that your array holds just one struct, you can safely just extract the first element in the array and put it one level higher (removing the array) like this:

df = df.withColumn('contacts', F.col('contacts')[0])

Full example:

from pyspark.sql import functions as F
df = spark.createDataFrame(
    [("123", "agreement", [("doc_id_1", "doc_name_1",())], [("contact_id_1", "John", "Doe")],)],
    "id string, description string, documents array<struct<id:string,doc_name:string,obligations:struct<>>>, contacts array<struct<id:string,contact_first_name:string,contact_last_name:string>>")
df.printSchema()
# root
#  |-- id: string (nullable = true)
#  |-- description: string (nullable = true)
#  |-- documents: array (nullable = true)
#  |    |-- element: struct (containsNull = true)
#  |    |    |-- id: string (nullable = true)
#  |    |    |-- doc_name: string (nullable = true)
#  |    |    |-- obligations: struct (nullable = true)
#  |-- contacts: array (nullable = true)
#  |    |-- element: struct (containsNull = true)
#  |    |    |-- id: string (nullable = true)
#  |    |    |-- contact_first_name: string (nullable = true)
#  |    |    |-- contact_last_name: string (nullable = true)

df = df.withColumn('documents', F.col('documents')[0])
df = df.withColumn('contacts', F.col('contacts')[0])

df.printSchema()
# root
#  |-- id: string (nullable = true)
#  |-- description: string (nullable = true)
#  |-- documents: struct (nullable = true)
#  |    |-- id: string (nullable = true)
#  |    |-- doc_name: string (nullable = true)
#  |    |-- obligations: struct (nullable = true)
#  |-- contacts: struct (nullable = true)
#  |    |-- id: string (nullable = true)
#  |    |-- contact_first_name: string (nullable = true)
#  |    |-- contact_last_name: string (nullable = true)

Collectives™ on Stack Overflow

Convert array to struct in dataframe

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related