RDD - Transformation MapPartitions

The mapPartitions() transformation is a high-performance variant of map(). While standard map() applies your function to each individual element of an RDD one-by-one, mapPartitions() applies the function to each individual partition as a whole.

Your function receives a Python generator iterator representing all the elements within a partition, and must return a new iterator/generator of output elements.

The Performance Gain

When doing simple mathematical operations, map and mapPartitions perform similarly. However, if your processing logic requires high-overhead initialization stepssuch as opening database connections, loading machine learning models, initializing thick API clients, or parsing heavy configuration structures:

• map(): Runs the initialization once per row (e.g., millions of database connection open/close calls!).
• mapPartitions(): Runs the initialization once per partition (e.g., if you have 4 partitions, only 4 database connections are opened!).

graph TD
    subgraph MapFlow["map() - Evaluates function for EVERY row"]
        direction TB
        R1["Row 1"] -->|Load Model| E1["Process Row"]
        R2["Row 2"] -->|Load Model| E2["Process Row"]
        R3["Row 3"] -->|Load Model| E3["Process Row"]
    end

    subgraph MapPartitionsFlow["mapPartitions() - Evaluates ONCE per partition block"]
        direction TB
        subgraph Partition1["Partition 1 (Rows 1, 2, 3)"]
            PM["Load Model Once"] --> P1["Process Row 1"]
            PM --> P2["Process Row 2"]
            PM --> P3["Process Row 3"]
        end
    end

    style MapFlow fill:#ffebee,stroke:#c62828,stroke-width:2px;
    style MapPartitionsFlow fill:#e8f5e9,stroke:#2e7d32,stroke-width:2px;

PySpark Code Examples

Setup Spark Session

from pyspark.sql import SparkSession

spark = SparkSession.builder \
    .appName("RDD Transformation MapPartitions") \
    .master("local[*]") \
    .getOrCreate()

sc = spark.sparkContext

Example A: Basic Summation per Partition

Let's see how mapPartitions processes partition datasets as sequences:

# 1. Create an RDD with 6 items split into exactly 3 partitions
numbers = sc.parallelize([1, 2, 3, 4, 5, 6], numSlices=3)
print("Initial partition data layout:", numbers.glom().collect())
# Output: [[1, 2], [3, 4], [5, 6]]

# 2. Define function that sums each partition iterator
def sum_iterator(partition_iterator):
    # Sum the values inside this partition's list
    partition_total = sum(partition_iterator)

    # We must yield or return an iterator/generator
    yield partition_total

# 3. Apply mapPartitions
partition_sums = numbers.mapPartitions(sum_iterator)

print("Sums per partition:", partition_sums.collect())
# Output: Sums per partition: [3, 7, 11]  (Sums of [1,2], [3,4], and [5,6])

Example B: Simulating a Distributed Database Write

Let's write a mock setup showing how to initialize a connection exactly once per partition block instead of row-by-row:

# RDD of user transactions
transactions = sc.parallelize([
    ("Txn101", 100), ("Txn102", 45), ("Txn103", 250), ("Txn104", 15)
], numSlices=2)

def save_to_database(partition_iterator):
    # 1. INITIALIZE CONNECTION ONCE PER PARTITION BLOCK!
    print(">>> Opening database connection for this partition node...")
    db_connection = "ActiveConnMock_3306"

    saved_records = []

    # 2. Iterate through all items in the partition locally
    for txn_id, amount in partition_iterator:
        # Simulate inserting row using our single open connection
        print(f"  [DB Write] Inserting: {txn_id} | Amount: ${amount} using {db_connection}")
        saved_records.append(f"{txn_id}_SUCCESS")

    # 3. CLOSE CONNECTION ONCE PER PARTITION BLOCK!
    print(">>> Closing database connection safely.")

    # Return output iterator
    return iter(saved_records)

# Apply mapPartitions
write_status = transactions.mapPartitions(save_to_database)

print("
Insertion Status Results:")
print(write_status.collect())

# Typical Output:
# >>> Opening database connection for this partition node...
# [DB Write] Inserting: Txn101 | Amount: $100 using ActiveConnMock 3306
# [DB Write] Inserting: Txn102 | Amount: $45 using ActiveConnMock 3306
# >>> Closing database connection safely.
# >>> Opening database connection for this partition node...
# [DB Write] Inserting: Txn103 | Amount: $250 using ActiveConnMock 3306
# [DB Write] Inserting: Txn104 | Amount: $15 using ActiveConnMock 3306
# >>> Closing database connection safely.
#
# Insertion Status Results:
# ['Txn101 SUCCESS', 'Txn102 SUCCESS', 'Txn103 SUCCESS', 'Txn104 SUCCESS']