Nesting Columns like a Pro: A Guide to Mastering Nested Structs in PySpark

Task

Are you tired of struggling with adding nested columns to already nested structs in PySpark DataFrames? Well, you're in luck because in this blog post, we'll be diving into the solution to this common problem. With the help of PySpark's withColumn and struct functions, we'll be showing you how to easily add nested columns to your DataFrames, but only if the column doesn't already exist. So, whether you're a beginner or an experienced PySpark developer, this post is for you. Keep reading to learn how to master nested columns in PySpark DataFrames!

Problem

In PySpark, adding columns to a DataFrame is not as simple as selecting them like you would with a SELECT statement. Instead, when you want to add a new column that is nested within other structs, you have to rebuild all of the structs that are part of that column. This can be a bit tricky because you also have to make sure to include all of the other fields that were already in those structs, otherwise they will be lost. It's like a puzzle, you have to make sure to keep all the pieces in the right place.

Solution

from typing import List
from pyspark.sql import functions, Column, DataFrame

def add_column_to_struct(dataTypes, path: List[str], current_struct_name: str, value: any) -> Column:
    current_column_level = path[0]
    next_column_level = path[1]

    should_place_column = len(path) == 2
    if should_place_column:
        new_column = []
        if next_column_level not in dataTypes.names:
            new_column = [functions.lit(value).alias(next_column_level)]
        return functions.struct(
            *([functions.col(current_column_level)[c].alias(c) for c in dataTypes.names] + new_column)
        ).alias(current_struct_name)

    current_struct_fields = []
    for index, column in enumerate(dataTypes.names):

        is_column_in_next_level = column == next_column_level
        if is_column_in_next_level:
            grouped_nested_path = [f"{current_column_level}.{next_column_level}", *path[2:]]

                        current_struct = add_column_to_struct(
                dataTypes.fields[index].dataType,
                grouped_nested_path,
                next_column_level,
                value
            )
        else:
            current_struct = functions.col(current_column_level)[column].alias(column)

        current_struct_fields.append(current_struct)

    return functions.struct(*current_struct_fields).alias(current_struct_name)

def create_column_if_not_exist(df: DataFrame, path: List[str], value: any) -> DataFrame:

    current_column_level = path[0]
    not_nested = len(path) == 1
    if not_nested:
        if current_column_level in df.columns:
            return df
        return df.withColumn(current_column_level, functions.lit(value))

    return df.withColumn(
        current_column_level,
        add_column_to_struct(
            df.schema[current_column_level].dataType,
            path,
            current_column_level,
            value
        )
    )

The previous code defines two functions create_column_if_not_exist and add_column_to_struct that allow adding a new column to a nested struct column in a PySpark DataFrame only if the column doesn't already exist.

The create_column_if_not_exist function takes in three parameters:

• df: a PySpark DataFrame where the column will be added
• path: a list of strings representing the path to the nested struct column where the new column will be added. For example, if the struct column is called address and is nested within another struct column called user, the path would be ["user", "address"]
• value: the value that will be assigned to the new column

The function first checks if the path has only one element, meaning that the column is not nested within any other structs. In this case, it checks if the column already exists in the DataFrame, if it does the function returns the DataFrame unmodified, otherwise, it adds the new column using the withColumn method and assigns it the value passed in.

If the path has more than one element, it means that the column is nested within other structs, the function then calls the add_column_to_struct function, passing the struct's dataType and the other parameters to it.

The add_column_to_struct function takes four parameters:

• dataTypes: the dataType of the struct column where the new column will be added
• path: the path to the nested struct column where the new column will be added
• current_struct_name: the name of the struct column where the new column will be added
• value: the value that will be assigned to the new column

It starts by extracting the first and second element of the path, these will be used to determine if we should add the new column in this level of nesting or if it should be added in a deeper level.

If the path has only two elements, the function checks if the new column already exists in the struct, if it does not, it creates the new column and assigns it the value passed in. It then creates a new struct column containing the existing columns and the new column.

If the path has more than two elements, the function iterates through all the columns of the struct and checks if the next level of nesting is reached. If it is, it calls the add_column_to_struct function again, passing the dataType of the next level of struct and updating the path and struct name accordingly. Otherwise, it simply uses the existing column.

Once all the columns have been processed, the function creates a new struct column containing all the columns and returns it.

Conclusion

In conclusion, working with nested structs in PySpark DataFrames can be like a game of Tetris - it takes a bit of skill to manoeuvre the columns in the right place, but with the right tools, it's totally doable. And let's be real, there's nothing more satisfying than fitting that last piece into the perfect spot. The provided code is like the cheat code to this game, it gives you hopefully a clear and clean way to add new columns to nested structs without losing any existing data. So, whether you're a PySpark pro or a beginner, this solution will have you adding nested columns like a boss in no time.