Introduction

Some of the examples you have seen so far used StructType and StructField. Let's dive deep into StructType and StructFields and their usage.

Spark SQL StructType & StructField classes are used to programmatically specify the schema to the DataFrame and create complex columns like nested struct, array, and map columns. StructType is a collection of StructFields. Using StructField, we can define column name, column data type, a nullable column (a boolean to specify if the field can be nullable or not), and metadata.

Though Spark infers a schema from data, sometimes we may need to define our own column names and data types. 

StructType – Defines the structure of the DataFrame

Spark provides pyspark.sql.types.StructType class to define the structure of the DataFrame and it is a collection or list of StructField objects. By calling Spark DataFrame printSchema() you can print the schema on the console where StructType columns are represented as struct.

StructField – Defines the metadata of the DataFrame column

Spark provides spark.sql.types.StructField class to define the column name (String), column type (DataType), nullable column (Boolean) and metadata (MetaData).

Using StructType & StructField with DataFrame

While creating a Spark DataFrame, we can specify the structure using StructType and StructField classes. As specified in the introduction, StructType is a collection of StructFields which is used to define the column name, data type, and a flag for nullable or not. Using StructField we can also add a nested struct schema, ArrayType for arrays, and MapType for key-value pairs, which we will discuss in detail in later sections.

StructType & StructField case class as follows.

from pyspark.sql.types import StructType, StructField, StringType, IntegerType, Metadata

schema = StructType([
    StructField("firstname", StringType(), True),
    StructField("middlename", StringType(), True),
    StructField("lastname", StringType(), True),
    StructField("id", StringType(), True),
    StructField("gender", StringType(), True),
    StructField("salary", IntegerType(), True)
])

The below example demonstrates a very simple example of how to create a struct using StructType & StructField on DataFrame and its usage with sample data to support it.

from pyspark.sql import SparkSession
from pyspark.sql import Row

spark = SparkSession.builder.master("local[1]").appName("AIStudyLabs.com").getOrCreate()

simpleData = [
    Row("James", "", "Smith", "36636", "M", 3000),
    Row("Michael", "Rose", "", "40288", "M", 4000),
    Row("Robert", "", "Williams", "42114", "M", 4000),
    Row("Maria", "Anne", "Jones", "39192", "F", 4000),
    Row("Jen", "Mary", "Brown", "", "F", -1)
]

simpleSchema = StructType([
    StructField("firstname", StringType(), True),
    StructField("middlename", StringType(), True),
    StructField("lastname", StringType(), True),
    StructField("id", StringType(), True),
    StructField("gender", StringType(), True),
    StructField("salary", IntegerType(), True)
])

df = spark.createDataFrame(simpleData, schema=simpleSchema)

df.printSchema()
df.show()

By running the above snippet, it displays the below outputs:

root
 |-- firstname: string (nullable = true)
 |-- middlename: string (nullable = true)
 |-- lastname: string (nullable = true)
 |-- id: string (nullable = true)
 |-- gender: string (nullable = true)
 |-- salary: integer (nullable = true)

+---------+----------+--------+-----+------+------+
|firstname|middlename|lastname|   id|gender|salary|
+---------+----------+--------+-----+------+------+
|    James|          |   Smith|36636|     M|  3000|
|  Michael|      Rose|        |40288|     M|  4000|
|   Robert|          |Williams|42114|     M|  4000|
|    Maria|      Anne|   Jones|39192|     F|  4000|
|      Jen|      Mary|   Brown|     |     F|    -1|
+---------+----------+--------+-----+------+------+

Defining Nested StructType object struct

We often need to work with the nested struct column, and this can be defined using SQL StructType.

In the below example, I have instantiated StructType and used the add method (instead of StructField) to add column names and datatype. Notice that for column “name” data type is StructType, which is nested.

structureData = [
    Row(Row("James", "", "Smith"), "36636", "M", 3100),
    Row(Row("Michael", "Rose", ""), "40288", "M", 4300),
    Row(Row("Robert", "", "Williams"), "42114", "M", 1400),
    Row(Row("Maria", "Anne", "Jones"), "39192", "F", 5500),
    Row(Row("Jen", "Mary", "Brown"), "", "F", -1)
]

structureSchema = StructType([
    StructField("name", StructType([
        StructField("firstname", StringType(), True),
        StructField("middlename", StringType(), True),
        StructField("lastname", StringType(), True)
    ])),
    StructField("id", StringType(), True),
    StructField("gender", StringType(), True),
    StructField("salary", IntegerType(), True)
])

df2 = spark.createDataFrame(structureData, schema=structureSchema)

df2.printSchema()
df2.show()

Outputs below the schema and the DataFrame:

root
 |-- name: struct (nullable = true)
 |    |-- firstname: string (nullable = true)
 |    |-- middlename: string (nullable = true)
 |    |-- lastname: string (nullable = true)
 |-- id: string (nullable = true)
 |-- gender: string (nullable = true)
 |-- salary: integer (nullable = true)

+--------------------+-----+------+------+ 
|                name|   id|gender|salary| 
+--------------------+-----+------+------+ 
|   [James , , Smith]|36636|     M|  3100| 
|  [Michael , Rose, ]|40288|     M|  4300| 
|[Robert , , Willi...|42114|     M|  1400| 
|[Maria , Anne, Jo...|39192|     F|  5500| 
|  [Jen, Mary, Brown]|     |     F|    -1| 
+--------------------+-----+------+------+ 

Creating StructType object struct from JSON file

If you have too many columns and the structure of the DataFrame changes now and then, it’s a good practice to load the SQL StructType schema from a JSON file. Note the definition in JSON uses a different layout and you can get this by using schema.json().

import json
from pyspark.sql.types import DataType

json_schema = """
{
  "type": "struct",
  "fields": [
    {
      "name": "name",
      "type": {
        "type": "struct",
        "fields": [
          {"name": "firstname", "type": "string", "nullable": true, "metadata": {}},
          {"name": "middlename", "type": "string", "nullable": true, "metadata": {}},
          {"name": "lastname", "type": "string", "nullable": true, "metadata": {}}
        ]
      },
      "nullable": true,
      "metadata": {}
    },
    {"name": "dob", "type": "string", "nullable": true, "metadata": {}},
    {"name": "gender", "type": "string", "nullable": true, "metadata": {}},
    {"name": "salary", "type": "integer", "nullable": true, "metadata": {}}
  ]
}
"""

schemaFromJson = StructType.fromJson(json.loads(json_schema))

df3 = spark.createDataFrame(structureData, schema=schemaFromJson)

df3.printSchema()

This prints the same output as the previous section. You can also have a name, type, and flag for nullable in a comma-separated file, and we can use these to create a StructType programmatically, which is left for you to explore.

Adding & Changing struct of the DataFrame

Using Spark SQL function struct(), we can change the struct of the existing DataFrame and add a new StructType to it. The below example demonstrates how to copy the columns from one structure to another and add a new column.

from pyspark.sql.functions import col, struct, when

updatedDF = df2.withColumn("OtherInfo",
    struct(
        col("id").alias("identifier"),
        col("gender"),
        col("salary"),
        when(col("salary").cast(IntegerType()) < 2000, "Low")
            .when(col("salary").cast(IntegerType()) < 4000, "Medium")
            .otherwise("High").alias("Salary_Grade")
    )
).drop("id", "gender", "salary")

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

Here, it copies “gender“, “salary” and “id” to the new struct “OtherInfo” and adds a new column “Salary_Grade“.

root
 |-- name: struct (nullable = true)
 |    |-- firstname: string (nullable = true)
 |    |-- middlename: string (nullable = true)
 |    |-- lastname: string (nullable = true)
 |-- OtherInfo: struct (nullable = false)
 |    |-- identifier: string (nullable = true)
 |    |-- gender: string (nullable = true)
 |    |-- salary: integer (nullable = true)
 |    |-- Salary_Grade: string (nullable = false)

Using SQL ArrayType and MapType

SQL StructType also supports ArrayType and MapType to define the DataFrame columns for array and map collections respectively. In the below example, column “hobbies” is defined as ArrayType(StringType) and “properties” is defined as MapType(StringType,StringType), meaning both key and value as String.

from pyspark.sql.types import ArrayType, MapType

arrayStructureData = [
    Row(Row("James", "", "Smith"), ["Cricket", "Movies"], {"hair": "black", "eye": "brown"}),
    Row(Row("Michael", "Rose", ""), ["Tennis"], {"hair": "brown", "eye": "black"}),
    Row(Row("Robert", "", "Williams"), ["Cooking", "Football"], {"hair": "red", "eye": "gray"}),
    Row(Row("Maria", "Anne", "Jones"), None, {"hair": "blond", "eye": "red"}),
    Row(Row("Jen", "Mary", "Brown"), ["Blogging"], {"white": "black", "eye": "black"})
]

arrayStructureSchema = StructType([
    StructField("name", StructType([
        StructField("firstname", StringType(), True),
        StructField("middlename", StringType(), True),
        StructField("lastname", StringType(), True)
    ])),
    StructField("hobbies", ArrayType(StringType()), True),
    StructField("properties", MapType(StringType(), StringType()), True)
])

df5 = spark.createDataFrame(arrayStructureData, schema=arrayStructureSchema)

df5.printSchema()
df5.show()

Outputs below the schema and the DataFrame data. Note that field “hobbies” is array type and “properties” is map type.

root
 |-- name: struct (nullable = true)
 |    |-- firstname: string (nullable = true)
 |    |-- middlename: string (nullable = true)
 |    |-- lastname: string (nullable = true)
 |-- hobbies: array (nullable = true)
 |    |-- element: string (containsNull = true)
 |-- properties: map (nullable = true)
 |    |-- key: string
 |    |-- value: string (valueContainsNull = true)

+---------------------+-------------------+------------------------------+
|name                 |hobbies            |properties                    |
+---------------------+-------------------+------------------------------+
|[James , , Smith]    |[Cricket, Movies]  |[hair -> black, eye -> brown] |
|[Michael , Rose, ]   |[Tennis]           |[hair -> brown, eye -> black] |
|[Robert , , Williams]|[Cooking, Football]|[hair -> red, eye -> gray]    |
|[Maria , Anne, Jones]|null               |[hair -> blond, eye -> red]   |
|[Jen, Mary, Brown]   |[Blogging]         |[white -> black, eye -> black]|
+---------------------+-------------------+------------------------------+

Convert case class to Spark StructType

While PySpark does not directly use case classes like Scala, it translates to defining a schema or DataSet in Python. However, you can convert a PySpark DataFrame schema to a StructType.

Creating StructType object struct from DDL String

Like loading structure from JSON string, we can also create it from DDL using fromDDL() method available in Spark 2.4 and above. You can also generate DDL from a schema using toDDL(). printTreeString() on the struct object prints the schema similar to printSchema function returns.

ddlSchemaStr = "`fullName` STRUCT<`first`: STRING, `last`: STRING, `middle`: STRING>,`age` INT,`gender` STRING"
ddlSchema = StructType.fromDDL(ddlSchemaStr)

ddlSchema.printTreeString()

Checking if a field exists in a DataFrame

If you want to perform some checks on metadata of the DataFrame, for example, if a column or field exists in a DataFrame or the data type of the column; we can easily do this using several functions on SQL StructType and StructField.

print("firstname" in df.schema.fieldNames())
print(StructField("firstname", StringType(), True) in df.schema)

This example returns “true” for both scenarios. And for the second one, if you have IntegerType instead of StringType, it returns false as the data type for the first name column is String, as it checks every property in the field. Similarly, you can also check if two schemas are equal and more.