Spark DataFrame extract value from array with where

Question

I have a dataframe with the following schema:

root
 |-- id: long (nullable = true)
 |-- raw_data: struct (nullable = true)
 |    |-- address_components: array (nullable = true)
 |    |    |-- element: struct (containsNull = true)
 |    |    |    |-- long_name: string (nullable = true)
 |    |    |    |-- short_name: string (nullable = true)
 |    |    |    |-- types: array (nullable = true)
 |    |    |    |    |-- element: string (containsNull = true)

Example of address_components:

{
   "address_components":[
      {
         "long_name":"Portugal",
         "short_name":"PT",
         "types":[
            "country",
            "political"
         ]
      },
      {
         "long_name":"8200-591",
         "short_name":"8200-591",
         "types":[
            "postal_code"
         ]
      }
   ]
}

I want to create a new root level attribute: Country: string that should contain PT. However, the selection should be based on array_contains(col("types"), "country")

I figured part of it out like this:

df = df.withColumn("country", expr("filter(raw_data.address_components, c -> array_contains(c.types, 'country'))"))
       .withColumn("country", col("country").getItem(0).getItem("long_name"))

is there a smarter/shorter way to do this?

Answer 1

I fixed it using expressions in combination with withColumn:

df = df.withColumn("country", expr("filter(raw_data.address_components, c -> array_contains(c.types, 'country'))[0].short_name"))