Spark Checkpointing,Union, Lineage And Partioning

I have been working on apache spark lately and wanted to consolidated my learnings at one place so that I can refer later and make it available to other developers as well.

While working with spark you will often run into situations where the way you write code normally results in most inefficient results. You can visualise spark plan executions but it is very tedious to navigate and make sense of in terms of what exactly is cause of the problem. In such case having basic knowledge of how spark manages tasks can help arrive at solutions to improve compute efficiency. Below are a set of techniques I have used to improve compute efficiency of spark applications.

SparkContext with a Checkpoint Directory

In Apache Spark, the SparkContext is the entry point to interact with the Spark framework. It manages the execution of jobs, resource allocation, and coordination across clusters.

A checkpoint directory in Spark is used to store checkpointed RDDs or streaming data, ensuring fault tolerance and recovery. It prevents recomputation of lost RDD partitions by persisting intermediate data to a reliable storage location (e.g., HDFS, S3, or local file system).

Setting a Checkpoint Directory

Before checkpointing any RDD, you must set the checkpoint directory in your SparkContext:

from pyspark import SparkContext sc = SparkContext("local", "Checkpoint Example") sc.setCheckpointDir("hdfs://namenode:9000/checkpoints") # Specify a checkpoint directory




Using Checkpointing in RDDs

Checkpointing saves RDD lineage, avoiding recomputation in case of failure:

rdd = sc.parallelize([1, 2, 3, 4, 5])

rdd.checkpoint()  # Mark this RDD for checkpointing

rdd.count()  # Trigger an action to save the checkpoint

Difference Between Cache & Checkpoint

	Cache	Checkpoint
Storage	Stored in memory (optionally disk)	Stored in a reliable storage (e.g., HDFS)
Persistence	Lost if Spark restarts	Available across failures
Lineage	Retains lineage	Truncates lineage after checkpointing

Use Cases for Checkpointing

1. Long Lineage Chains – Avoids performance overhead caused by deep dependency graphs.
2. Fault Tolerance – Ensures data recovery after job failures.
3. Streaming Applications – Required for stateful transformations in Spark Streaming.

While working with spark it is very important to understand spark execution plans, also known as lineage.

Lineage of a Spark RDD

The lineage of a Resilient Distributed Dataset (RDD) in Apache Spark refers to the sequence of transformations that led to the creation of that RDD from its parent(s). It represents the logical execution plan rather than the actual data.

How Lineage Works

• When an RDD is created, Spark does not immediately compute and store its data.
• Instead, it maintains a DAG (Directed Acyclic Graph) of transformations, known as RDD lineage.
• The actual computation happens only when an action (e.g., .count(), .collect()) is triggered.
• If a partition of an RDD is lost (e.g., due to node failure), Spark recomputes only the necessary partitions using the lineage.

---

Example of RDD Lineage

python

from pyspark import SparkContext

sc = SparkContext("local", "RDD Lineage Example")

# Create an initial RDD
rdd1 = sc.parallelize([1, 2, 3, 4, 5])

# Apply transformations
rdd2 = rdd1.map(lambda x: x * 2)   # Transformation 1
rdd3 = rdd2.filter(lambda x: x > 5)  # Transformation 2

# Trigger an action
print(rdd3.collect())

Lineage Graph for This Example:

scss

rdd1 (Parallelize) → rdd2 (Map) → rdd3 (Filter) → Action (collect)

---

Viewing RDD Lineage in Spark

You can check an RDD's lineage using .toDebugString():

python

print(rdd3.toDebugString())

Example output:

scss

(2) PythonRDD[3] at RDD at PythonRDD.scala:53
 |  MapPartitionsRDD[2] at map at <stdin>:1
 |  ParallelCollectionRDD[1] at parallelize at PythonRDD.scala:195

This output shows the dependencies between RDDs and how Spark will recompute data if needed.

---

Why RDD Lineage Matters

1. Fault Tolerance – Spark can recompute lost data using lineage rather than storing everything.
2. Lazy Evaluation – Spark defers execution until an action is triggered, optimizing computation.
3. Optimized Execution – Spark can optimize execution plans by skipping unnecessary computations.

---

How to Break Lineage

Since lineage can grow too long (causing high recomputation costs), you can truncate lineage using:

• Checkpointing: Saves RDDs to reliable storage (e.g., HDFS) and clears lineage.

  python
  
  sc.setCheckpointDir("hdfs://namenode:9000/checkpoints")
  rdd3.checkpoint()
  rdd3.count()  # Trigger checkpointing
  

• Persisting or Caching: Stores the RDD in memory or disk to avoid recomputation.

        python

     rdd3.persist() # or rdd3.cache()

Spark union() – Combining RDDs or DataFrames
In Apache Spark, the union() operation is used to combine two or more RDDs or DataFrames by concatenating their elements. It does not remove duplicates, unlike distinct().
1. union() with DataFrames
When working with Spark DataFrames, union() requires that both DataFrames have the same schema.
Syntax:
python

df3 = df1.union(df2)
Example:
python

from pyspark.sql import SparkSession

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

# Creating DataFrames
data1 = [(1, "Alice"), (2, "Bob")]
data2 = [(3, "Charlie"), (4, "David")]

columns = ["id", "name"]

df1 = spark.createDataFrame(data1, columns)
df2 = spark.createDataFrame(data2, columns)

# Applying union
df3 = df1.union(df2)

# Showing the result
df3.show()
Output:
pgsql

+---+-------+
| id|  name |
+---+-------+
|  1| Alice |
|  2|   Bob |
|  3|Charlie|
|  4| David |
+---+-------+
2. Handling Duplicates with distinct()
If you want to remove duplicates after a union(), use distinct():
For DataFrames:
python

df_union_distinct = df1.union(df2).distinct()
df_union_distinct.show()
3. When to Use union()
Merging multiple datasets without worrying about duplicates.
Appending data from one source to another (e.g., logs, incremental data).
Preprocessing before removing duplicates using distinct().
The union operation in Apache Spark can be computationally expensive for several reasons:
No Deduplication, but Potentially Large Data Volume
The union operation combines two datasets without removing duplicates. If both datasets are large, the result can be even larger, leading to higher memory and disk usage.
Shuffling Costs (if Followed by Actions Like Aggregations or Joins)
If a union is followed by operations like groupBy, distinct, or join, Spark may need to shuffle data across nodes, which is an expensive operation in terms of network I/O.
Different Lineages and Recomputations
Each dataset in Spark has its own lineage (DAG of transformations). When a union is performed, Spark maintains both lineages, meaning if an action (like .count() or .collect()) is triggered, Spark may need to recompute data from both original datasets.
Increased Memory Pressure
If the unioned dataset is large and does not fit in memory, Spark may have to spill to disk, increasing execution time due to disk I/O.
Schema Enforcement in DataFrames
In Spark DataFrames (unlike RDDs), schemas must be consistent across datasets in a union. If schemas differ, Spark has to perform additional transformations to align them, adding computational overhead.
Partitioning and Parallelism Considerations
If the datasets being unioned have different numbers of partitions, Spark may need to repartition them for efficient parallel execution, adding further overhead.
To optimize union operations:
Ensure datasets have the same number of partitions using .coalesce() or .repartition().
If duplicates are unnecessary, use .distinct() after union() to reduce redundant data.
Use cache() or persist() if the unioned dataset is used multiple times.
If working with DataFrames, ensure schemas are aligned before the union to avoid costly transformations.
When I was trying to optimise my spark code my issue surrounded around Partitioning and Parallelism Considerations, so lets expand on that:
Partitioning and Parallelism Considerations in Spark Union
Apache Spark processes data in parallel by dividing it into partitions. When performing a union operation, the way the partitions are distributed across nodes in the cluster can impact performance.
Why Does Partitioning Matter in union?
If two DataFrames or RDDs being unioned have different numbers of partitions, Spark may struggle to efficiently distribute work across the cluster. This can lead to:
Data Skew – Some partitions may become much larger than others, causing certain tasks to take longer while others finish quickly.
Uneven Parallelism – If one dataset has significantly fewer partitions than the other, Spark may not fully utilize available resources, leading to inefficient execution.
Repartitioning Overhead – If Spark detects an imbalance, it may repartition the data, which is an expensive operation involving data movement and shuffling.
Example Scenario
Let’s say:
df1 has 200 partitions.
df2 has 50 partitions.
When you do df1.union(df2), Spark will retain the partitioning of df1, but the partitions of df2 may not align with df1. This can lead to inefficiencies if later operations (like groupBy or join) require even partitioning.
How to Optimize?
Match the Number of Partitions Before Union
If datasets have mismatched partitions, align them using .repartition(n) or .coalesce(n).
python

df1 = df1.repartition(100)
df2 = df2.repartition(100)
df_union = df1.union(df2)
Use coalesce(n) instead of repartition(n) if reducing partitions, as it avoids full shuffle.
Ensure Partitioning Strategy Matches Downstream Operations
If you're going to aggregate (groupBy), join, or filter the data later, use consistent partitioning to reduce shuffle costs.