RDD - Transformation FlatMap

The flatMap() transformation is similar to standard map(), but with one major difference: while map() requires each input element to map to exactly one output element, flatMap() allows each input element to map to zero, one, or more output elements.

Additionally, flatMap() automatically flattens the final output collection. If your user-defined function returns lists or sequences, flatMap() merges those sub-lists into a single, flat, continuous RDD of elements.

Contrast: map vs. flatMap

graph TD
    subgraph MapBehavior["map() - Returns nested lists"]
        direction TB
        M_In["['Hello World']"] -->|split| M_Out["[['Hello', 'World']]"]
    end

    subgraph FlatMapBehavior["flatMap() - Flattens nested lists"]
        direction TB
        F_In["['Hello World']"] -->|split & flatten| F_Out["['Hello', 'World']"]
    end

    style MapBehavior fill:#ffebee,stroke:#c62828,stroke-width:2px;
    style FlatMapBehavior fill:#e8f5e9,stroke:#2e7d32,stroke-width:2px;

PySpark Code Examples

Setup Spark Session

from pyspark.sql import SparkSession

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

sc = spark.sparkContext

Example A: Splitting Sentences into Words (The Word Count Starter)

Let's see how map and flatMap handle splitting a sentence by spaces differently:

# 1. Input RDD with two sentences
sentences = sc.parallelize(["Hello World", "Learn Spark RDD"])

# 2. Map split: Returns an RDD of nested lists
mapped_words = sentences.map(lambda s: s.split(" "))
print("Map Split output:")
print(mapped_words.collect())
# Output: [['Hello', 'World'], ['Learn', 'Spark', 'RDD']]

# 3. FlatMap split: Returns a single flat RDD of words
flat_mapped_words = sentences.flatMap(lambda s: s.split(" "))
print("
FlatMap Split output:")
print(flat_mapped_words.collect())
# Output: ['Hello', 'World', 'Learn', 'Spark', 'RDD']

Example B: Extracting Elements from JSON lists

Assume you have user records containing a list of transaction amounts, and you want to extract every transaction into a single global RDD:

# 1. RDD of users and their transaction lists
user_transactions = sc.parallelize([
    ("User_1", [50, 12, 100]),
    ("User_2", [5, 45]),
    ("User_3", []) # Maps to zero output elements!
])

# 2. Extract transaction list and flatten
all_transactions = user_transactions.flatMap(lambda x: x[1])

print("All Transactions:", all_transactions.collect())
# Output: All Transactions: [50, 12, 100, 5, 45]

Example C: Generating Ranges

Let's map integers to ranges and flatten the output:

# 1. Input numbers
numbers = sc.parallelize([2, 3, 4])

# 2. Generate a list from 1 up to x for each element
ranges = numbers.flatMap(lambda x: list(range(1, x)))

print("Merged Ranges:", ranges.collect())
# Output: Merged Ranges: [1, 1, 2, 1, 2, 3]
# (Explains: range(1,2) is [1], range(1,3) is [1,2], range(1,4) is [1,2,3])