Introduction

Spark map() and mapPartitions() transformations apply a function on each element/record/row of the DataFrame and return a new DataFrame. In this article, I will explain the difference between map() vs mapPartitions() transformations, their syntax, and usages with PySpark examples.

• map() – Spark map() transformation applies a function to each row in a DataFrame and returns the new transformed DataFrame.
• mapPartitions() – This is similar to map(); however, Spark mapPartitions() provides the ability to do heavy initializations (for example, database connection) once for each partition instead of doing it on every DataFrame row. This enhances performance when dealing with heavy-weighted initialization on larger datasets.

Key Points:

• Both transformations return a PySpark DataFrame.
• After applying the transformation function on each row, these transformations return the same number of rows as input, but the schema or number of columns of the result could differ.
• Compared to flatMap(), which can return multiple rows for each input, map() returns only one row for every input row.

Understanding StructType and StructField

Before diving into the examples, it’s essential to understand StructType and StructField. StructType is a collection of StructFields that defines the schema of a DataFrame. Each StructField has a name, data type, and a flag indicating whether null values are allowed.

Let’s see the differences with an example. First, let’s create a Spark DataFrame:

from pyspark.sql import SparkSession
from pyspark.sql import Row
from pyspark.sql.types import StructType, StructField, StringType, IntegerType

spark = SparkSession.builder.appName("example").getOrCreate()

structureData = [
    Row("James", "", "Smith", "36636", "NewYork", 3100),
    Row("Michael", "Rose", "", "40288", "California", 4300),
    Row("Robert", "", "Williams", "42114", "Florida", 1400),
    Row("Maria", "Anne", "Jones", "39192", "Florida", 5500),
    Row("Jen", "Mary", "Brown", "34561", "NewYork", 3000)
]

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

df2 = spark.createDataFrame(structureData, structureSchema)

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

Output:

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

+---------+----------+--------+-----+----------+------+
|firstname|middlename|lastname|id   |location  |salary|
+---------+----------+--------+-----+----------+------+
|James    |          |Smith   |36636|NewYork   |3100  |
|Michael  |Rose      |        |40288|California|4300  |
|Robert   |          |Williams|42114|Florida   |1400  |
|Maria    |Anne      |Jones   |39192|Florida   |5500  |
|Jen      |Mary      |Brown   |34561|NewYork   |3000  |
+---------+----------+--------+-----+----------+------+

To explain map() and mapPartitions() with an example, let’s create a utility function combine(). This function takes three string arguments and combines them with a comma delimiter. In practice, this could represent more complex transformations.

def combine(fname: str, mname: str, lname: str) -> str:
    return f"{fname},{mname},{lname}"

Spark map() transformation

Spark map() transformation applies a function to each row in a DataFrame and returns the new transformed DataFrame. The map() function returns as many rows as the input DataFrame.

For example, if you have 100 rows in a DataFrame, after applying the function with map(), it returns exactly 100 rows. However, the structure or schema of the output could be different.

Usage:

from pyspark.sql.functions import udf
from pyspark.sql.types import StringType, IntegerType

combine_udf = udf(combine, StringType())

df3 = df2.withColumn("fullName", combine_udf(df2.firstname, df2.middlename, df2.lastname)) \
         .select("fullName", "id", "salary")

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

The above transformation initializes and creates the full name for every row of a DataFrame. While this is fine for smaller operations, it could lead to performance issues with costly initializations for every row.

The above example yields below output:

root
 |-- fullName: string (nullable = true)
 |-- id: string (nullable = true)
 |-- salary: integer (nullable = true)

+----------------+-----+------+
|fullName        |id   |salary|
+----------------+-----+------+
|James,,Smith    |36636|3100  |
|Michael,Rose,   |40288|4300  |
|Robert,,Williams|42114|1400  |
|Maria,Anne,Jones|39192|5500  |
|Jen,Mary,Brown  |34561|3000  |
+----------------+-----+------+

As noted, the input DataFrame has 5 rows, so the result of the map also has 5, but the number of columns differs.

Spark mapPartitions() transformation

Spark mapPartitions() provides a facility to do heavy initializations (for example, database connections) once per partition instead of each DataFrame row, enhancing performance for larger datasets.

Usage:

def process_partition(iterator):
    for row in iterator:
        fullName = combine(row.firstname, row.middlename, row.lastname)
        yield (fullName, row.id, row.salary)

df4 = df2.rdd.mapPartitions(process_partition).toDF(["fullName", "id", "salary"])

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

This yields the same output as above, demonstrating the advantage of partition-level processing for heavy initializations.