How to create columns of ArrayType and MapType in PySpark using Azure Databricks?

Are you looking to find out how to add a new ArrayType column with a constant value in Azure Databricks cloud or maybe you are looking for a solution, to add a MapType column with literal value on PySpark’s DataFrame in PySpark Databricks using the lit() function? If you are looking for any of these problem solutions, you have landed on the correct page. I will also show you how to use PySpark to add columns on dataframe using the lit() function in Azure Databricks. I will explain it by taking a practical example. So please don’t waste time let’s start with a step-by-step guide to understanding how to add new array and map type columns in PySpark DataFrame using the lit() function.

In this blog, I will teach you the following with practical examples:

• Syntax of lit() function
• Adding a new column of ArrayType using lit()
• Adding a new column of MapType using lit()

The PySpark’s lit() function is a function used to add new columns of DataFrame in PySpark Azure Databricks. Lit takes a literal or constant value and returns a new Column.

Syntax:

lit(constant_value)

What is the syntax of the lit() function in PySpark Azure Databricks?

The syntax is as follows:

lit(constant_value)

Apache Spark Official documentation link: lit()

Create a simple DataFrame

Let’s understand the use of the lit() function with a variety of examples. Let’s start by creating a DataFrame.

Gentle reminder:

In Databricks,

• sparkSession made available as spark
• sparkContext made available as sc

In case, you want to create it manually, use the below code.

from pyspark.sql.session import SparkSession

spark = SparkSession.builder 
    .master("local[*]") 
    .appName("azurelib.com") 
    .getOrCreate()

sc = spark.sparkContext

a) Create manual PySpark DataFrame

data = [
    ("Mamie","Treharne"),
    ("Erv","Colam"),
    ("Daren","Salliss"),
    ("Vania","Laundon"),
    ("Jay","Kees"),
]

df = spark.createDataFrame(data, schema=["f_name","l_name"])
df.printSchema()
df.show(5, truncate=False)

"""
root
 |-- f_name: string (nullable = true)
 |-- l_name: string (nullable = true)

+------+--------+
|f_name|l_name  |
+------+--------+
|Mamie |Treharne|
|Erv   |Colam   |
|Daren |Salliss |
|Vania |Laundon |
|Jay   |Kees    |
+------+--------+
"""

b) Creating a DataFrame by reading files

Download and use the below source file.

# replace the file_path with the source file location which you have downloaded.

df_2 = spark.read.format("csv").option("header", True).load(file_path)
df_2.printSchema()

"""
root
 |-- f_name: string (nullable = true)
 |-- l_name: string (nullable = true)
"""

Note: Here, I will be using the manually created DataFrame.

How to create a new column of ArrayType on PySpark DataFrame using the lit() function?

The typedlit() function is a function that can handle collection types that are not available. In PySpark, we don’t have this functionality as of now; therefore, I have created an ArrayType column using the array() function.

Example:

from pyspark.sql.functions import lit, array

arr_df = df.select("*", array([lit(1), lit(2)]).alias("array_col"))
arr_df.printSchema()
arr_df.show()

"""
Output:

root
 |-- f_name: string (nullable = true)
 |-- l_name: string (nullable = true)
 |-- array_col: array (nullable = false)
 |    |-- element: integer (containsNull = false)

+------+--------+---------+
|f_name|  l_name|array_col|
+------+--------+---------+
| Mamie|Treharne|   [1, 2]|
|   Erv|   Colam|   [1, 2]|
| Daren| Salliss|   [1, 2]|
| Vania| Laundon|   [1, 2]|
|   Jay|    Kees|   [1, 2]|
+------+--------+---------+

"""

How to create a new column of MapType on PySpark DataFrame using the lit() function?

As typedlit() is a function that can handle collection types that are which is not available. In PySpark, we don’t have this functionality as of now, therefore I have created an ArrayType column using the from_json() function.

Example:

from pyspark.sql.types import StringType, MapType
from pyspark.sql.functions import from_json

map_str_df = df.select("*", lit("{'key': 'value'}").alias("map_str_col"))
map_df = map_str_df.withColumn("map_col", 
                               from_json("map_str_col", MapType(StringType(), StringType())))
map_df.printSchema()
map_df.show()

"""
Output:

root
 |-- f_name: string (nullable = true)
 |-- l_name: string (nullable = true)
 |-- map_str_col: string (nullable = false)
 |-- map_col: map (nullable = true)
 |    |-- key: string
 |    |-- value: string (valueContainsNull = true)

+------+--------+----------------+--------------+
|f_name|  l_name|     map_str_col|       map_col|
+------+--------+----------------+--------------+
| Mamie|Treharne|{'key': 'value'}|{key -> value}|
|   Erv|   Colam|{'key': 'value'}|{key -> value}|
| Daren| Salliss|{'key': 'value'}|{key -> value}|
| Vania| Laundon|{'key': 'value'}|{key -> value}|
|   Jay|    Kees|{'key': 'value'}|{key -> value}|
+------+--------+----------------+--------------+

"""

Note: In PySpark string value inside a column is considered a column name.

I have attached the complete code used in this blog in a notebook format in this GitHub link. You can download and import this notebook in databricks, jupyter notebook, etc.

When should you use the PySpark lit() in Azure Databricks?

These could be the possible reasons:

1. For creating new columns using a literal value
2. For creating columns of different DataTypes

Real World Use Case Scenarios for PySpark DataFrame lit() in Azure Databricks?

• Assume you want to create new columns using a constant value, for example, if you have been asked to create a country column, you can add it by passing the country value into the lit() function.
• Assume you were asked to create a column of ArrayType or a MapType. For example, if you have been asked to create a column with a key and value pair, you can use lit() function to achieve this.

What are the alternatives to the lit() function in PySpark Azure Databricks?

The Spark provides typedlit() function to create or replace existing values with array or map values. But unfortunately, we don’t have this function in PySpark. Therefore we have formulated our own way of creating it by using the lit() function. These alternative methods are discussed above with examples.

Final Thoughts

In this article, we have learned about the PySpark lit() method to add new columns on DataFrame in Azure Databricks along with the examples explained clearly. I have also covered different scenarios with practical examples that could be possible. I hope the information that was provided helped in gaining knowledge.

Please share your comments and suggestions in the comment section below and I will try to answer all your queries as time permits.

• For Azure Study material Join Telegram group : Telegram group link:
• Azure Jobs and other updates Follow me on LinkedIn: Azure Updates on LinkedIn
• Azure Tutorial Videos: Videos Link

• Azure Databricks Lesson 1
• Azure Databricks Lesson 2
• Azure Databricks Lesson 3
• Azure Databricks Lesson 4
• Azure Databricks Lesson 5
• Azure Databricks Lesson 6
• Azure Databricks Lesson 7

Arud Seka Berne S

As a big data engineer, I design and build scalable data processing systems and integrate them with various data sources and databases. I have a strong background in Python and am proficient in big data technologies such as Hadoop, Hive, Spark, Databricks, and Azure. My interest lies in working with large datasets and deriving actionable insights to support informed business decisions.