← All Scenarios

Looping & Iterative Processing.

Understanding when to use driver-side loops versus distributed partition-level iteration.

1. Driver-Side Looping

Looping through a DataFrame using Python for loops is generally discouraged unless you are iterating over metadata (like a list of table names) to trigger separate Spark jobs.

# Iterating over a list of paths to process multiple files
file_list = ["sales_jan.csv", "sales_feb.csv", "sales_mar.csv"]

for file in file_list:
    df = spark.read.csv(f"s3://raw/{file}")
    # Process and save
    df.write.mode("append").save("s3://silver/yearly_sales")
    

2. Advanced: mapPartitions

When you need to perform an expensive setup (like opening a database connection or loading a machine learning model), mapPartitions is the gold standard. It allows you to execute logic once per partition rather than once per row.

def process_batch(partition_iterator):
    # Setup: Initialize a heavy resource once per partition
    model = load_my_custom_model() 
    
    results = []
    for row in partition_iterator:
        # Process each row using the resource
        prediction = model.predict(row.features)
        results.append((row.id, prediction))
    return iter(results)

# Applying the partition-level logic
rdd_results = df.rdd.mapPartitions(process_batch)
    

3. flatMap for One-to-Many

While map returns exactly one element for every input, flatMap allows you to return a collection of elements that Spark then flattens into a single RDD/DataFrame. This is commonly used for tokenizing text.

# Splitting sentences into individual words
words_rdd = text_df.rdd.flatMap(lambda row: row.sentence.split(" "))
    

Interview Q&A