Flatten Map Spark Python
Tags: spark, pyspark
Published: 2023-05-15 15:38:20

---
title: Flatten Map Spark Python
published: true
description: 
tags:
  - spark
  - pyspark
canonical_url: https://www.aloneguid.uk/posts/2023/05/flatten-map-spark-python/
---

## Problem

We have a  `DataFrame` with a map property, looking like this:

```
root
 |-- t: timestamp (nullable = true)
 |-- tags: map (nullable = true)
 |    |-- key: string
 |    |-- value: string (valueContainsNull = true)
```

We want to transform it to the following:

```
root
 |-- t: timestamp (nullable = true)
 |-- key1: string
 |-- key2: string
 ...
 |-- keyN: string
```

where `key1...keyN` are all possible combinations of keys in the map.

## Solution

One of the challenges of working with `map` type data is that it does not have a fixed schema. Unlike other types, such as `array` or `struct`, that have a predefined number and name of columns, `map` can have different keys and values for each record. This means that we cannot easily predict how many columns will be generated when we try to flatten a `map` type column into multiple columns. For example, if we have a `map` column that stores the names and ages of different people, one record might have two keys (`name` and `age`), while another record might have three keys (`name`, `age`, and `gender`). This makes it difficult to process and analyze the data in a consistent way.

This approach has a limitation: it depends on the current state of the data. If the data changes in the future, the solution might not work anymore. However, this is not a problem for me because I only need to run it once. For those who use Delta Lake, this is also not an issue because Delta Lake can handle schema changes automatically.

### Get Unique Keys

One of the first steps in data analysis is to identify the unique values or keys in a dataset. This can help us to understand the structure of the data, as well as to create the necessary flat schema.

The pyspark code for finding the keys in a dataframe is as follows:

```python
map_keys = (df
            .select(f.map_keys("tags").alias("tag_keys"))
            .select(f.explode("tag_keys").alias("tag_keys"))
            .distinct())
map_keys = [x["tag_keys"] for x in map_keys.collect()]
```

This code operates on a `DataFrame` named `df` and performs the following operations:

1. The `select` function is used with the `map_keys` transformation from the `pyspark.sql.functions` module. It retrieves the keys from the "tags" map column in the `DataFrame` `df` and renames the resulting column as "tag_keys".
2. The `explode` function is applied to the "tag_keys" column, which transforms each map key into a separate row, effectively "exploding" the map into multiple rows.
3. The `distinct` function is used to remove any duplicate rows from the `DataFrame`.
4. The resulting `DataFrame` is assigned to the variable `map_keys`.
5. The `collect` function is called on the `map_keys` `DataFrame`, which triggers the execution of the Spark job and returns the results as a list of `Row` objects.
6. A list comprehension is used to extract the values of the "tag_keys" column from each Row object in the list, and this list of extracted values is assigned back to the variable `map_keys`.

In summary, this code extracts the distinct keys from a map column named "tags" in the DataFrame `df` and stores them as a list of strings in the `map_keys` variable. You should get a python array similar to this:

```python
['url',
 'b_sys_name',
 'bi',
 'ua',
 'referrer',
 'out',
 'height',
 'b_name',
 'profile',
 'ppn',
 'width',
 'inactive_mins']
```

One step done.

### Generate Flat Schema

To generate flat schema, i'll use python list comprehension again. Honestly, list comprehension is awesome - this is one of the reason I actually started using Python a long time ago. I'm getting of the rails. Ok, the last time:

> "Discover the Magic of Python List Comprehension: Unlocking the Key to Efficiently Generating Flat Schemas. Dive Into the World of Simplified Data Processing in Python and Spark, and Uncover Why List Comprehension Has Been the Secret Ingredient That Made Me Fall in Love With Python Many Years Ago."

Back to the topic. Creating flat schema is as easy as this:

```python
xdf = (df
       .select(*df_all.columns[:-1], *[f.col(f"tags.{x}") for x in map_keys]))
```

This code operates on a `DataFrame` named `df_all` and performs the following operations:

1. The `DataFrame` `df` is used as the starting point for the transformation. This is our original data.
2. The `*df.columns[:-1]` expression uses the unpacking operator `*` to unpack all columns except the last one (`[:-1]`) from the `df` `DataFrame`. This ensures that all columns of `df` except the last one are included in the selection, because the last column `tags` is the one we don't need to include in the result (we want to flatten it!).
3. The expression `[f.col(f"tags.{x}") for x in map_keys]` is a list comprehension that generates a list of column expressions dynamically. For each value `x` in the `map_keys` list, it creates a column expression using `f.col(f"tags.{x}")`. This expression accesses the `tags` map column in the `df` `DataFrame` and retrieves the value associated with the key `x`.
4. The resulting list of column expressions is unpacked using the `*` operator within the `select` function. This includes the dynamically generated column expressions in the selection alongside the columns from `df.columns[:-1]`.
5. The resulting `DataFrame` is assigned to the variable `xdf`.

This will generate a very wide table (in my case) which looks similar to:

```
root
 |-- t: timestamp (nullable = true)
 |-- version: string (nullable = true)
 |-- iid: string (nullable = true)
 |-- ip: string (nullable = true)
 |-- event: string (nullable = true)
 |-- user_name: string (nullable = true)
 |-- windows_version: string (nullable = true)
 |-- domain: string (nullable = true)
...
```

Everything after `t` is taken from the map keys.

And now you can just execute the query and reap the benefits.

## Appendix. Why Flattening can be Beneficial

Flattening nested data structures involves transforming a complex nested structure into a **simple one-dimensional list or array**. This process can be useful for several reasons:

1. **Simplifying data access**: Nested structures can make it challenging to access specific pieces of information. Flattening them makes it easier to **access and manipulate individual data points**.
2. **Easier data processing**: Flattening a nested structure can **simplify data processing and analysis** by eliminating the need for complex nested loops or recursion.
3. **More efficient storage**: Storing nested structures can be **memory-intensive**, and flattening can help **reduce the storage space** required by eliminating unnecessary levels of nesting.
4. **Better data visualization**: Data that is flattened can be more easily **visualized in tables, charts, and graphs**, which are often one-dimensional.
5. **Improved performance**: In some cases, flattening can **improve the performance** of an application by reducing the amount of time and resources required to access and process data.
6. **Simplified data integration**: When working with data from different sources, flattening nested structures can help in **integrating and combining datasets** more easily. It allows for a consistent format and facilitates merging data based on common attributes.
7. **Smoother data transformation**: Flattening nested structures can streamline the process of **data transformation**. It enables the application of operations and transformations uniformly across the entire dataset, without needing to account for complex nested hierarchies.
8. **Enhanced compatibility**: Flattened data structures are often more **compatible with various data analysis and machine learning algorithms**. Many algorithms require a tabular or flat structure, and by flattening the data, it becomes readily usable for such algorithms.
9. **Simplified debugging and troubleshooting**: Working with flattened data structures can make **debugging and troubleshooting** easier. It reduces complexity and enables more straightforward identification and resolution of issues related to data quality, consistency, or specific data points.
10. **Improved data sharing**: Flattening nested structures can facilitate **data sharing** with others. It simplifies the process of exchanging data, collaborating on projects, and sharing insights with stakeholders who may not be familiar with complex nested representations.
11. **Easier data export and integration with external systems**: Flattened data structures are often more compatible and easier to export or integrate with external systems, such as databases, data warehouses, or reporting tools. It simplifies the process of data migration and integration.

By flattening nested structures, data becomes more accessible, easier to process, and compatible with a wider range of applications and tools. It promotes simplicity, efficiency, and effective data utilization across various domains and use cases.

Overall, **flattening nested data structures** can make working with complex data sets more manageable and efficient.

By the way, did you know why a data scientist become a pancake chef?

>  Because they loved **flattening data** almost as much as they loved flattening pancakes! 🥞😄

## References

- [PEP 448 -- Additional Unpacking Generalizations](https://www.python.org/dev/peps/pep-0448/).
- [Python Docs: Unpacking Argument Lists](https://docs.python.org/3/tutorial/controlflow.html#unpacking-argument-lists).
- `PySpark` functions used:
  - [select](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.sql.DataFrame.select.html) - Selects a subset of columns from a DataFrame.
  - [map_keys](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.sql.functions.map_keys.html) - Extracts the keys from a map column.
  - [alias](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.sql.Column.alias.html) - Renames a column.
  - [select](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.sql.DataFrame.select.html) - Selects a subset of columns from a DataFrame.
  - [map_keys](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.sql.functions.map_keys.html) - Extracts the keys from a map column.
  - [alias](https://spark.apache.org/docs/latest/api/python/reference/api/pyspark.sql.Column.alias.html) - Renames a column.