PySpark Internals — Day 87 of 100 Days of Data Engineering, AI and Azure Challenge

PySpark combines the scalability of Apache Spark (JVM-based) with Python’s simplicity. Below is a detailed breakdown of its internals, including DAGs, lineage, job execution, and the interplay between JVM and Python.

1. PySpark Architecture

PySpark operates on a master-worker model:

Driver Program :

• Runs the Python code (e.g., spark-submit).
• Coordinates the job execution and communicates with the SparkSession.

Cluster Manager (YARN, Kubernetes, or Spark Standalone):

• Allocates resources (CPU, memory) to executors.

Executors :

• Run JVM-based worker processes to process data.
• Execute tasks (units of work) and cache data.

Key Components :

• SparkSession : Entry point to PySpark (replaces older SparkContext, SQLContext).
• RDD/DataFrame/Dataset : Distributed data structures.

2. DAG (Directed Acyclic Graph)

Spark optimizes execution using a DAG of transformations:

• Transformations (lazy operations like map, filter, join) build the DAG.
• Actions (eager operations like count, collect) trigger DAG execution.

DAG Creation :

1. Logical Plan : Generated from transformations (e.g., df.select().filter()).
2. Physical Plan : Optimized by Catalyst (Spark’s query optimizer).
3. Stages : DAG is split into stages at shuffle boundaries (e.g., groupByKey, join).

# Code Example
df = spark.read.csv("data.csv")
df_filtered = df.filter(df.age > 30)
df_grouped = df_filtered.groupBy("city").count()
df_grouped.write.parquet("output")

• Stages :
• Stage 1: Read CSV + Filter (narrow transformation).
• Stage 2: Shuffle + GroupBy + Count (wide transformation).

3. Lineage and Fault Tolerance

• Lineage : RDDs track dependencies (parent RDDs and transformations) to rebuild lost data.
• Fault Tolerance : If a partition is lost, Spark recomputes it using lineage.

rdd = sc.textFile("data.txt")  # Lineage: data.txt → RDD
rdd2 = rdd.map(lambda x: x.upper())  # Lineage: data.txt → map → RDD2
rdd3 = rdd2.filter(lambda x: "ERROR" in x)  # Lineage: data.txt → map → filter → RDD3

If a node fails, rdd3 can recompute lost partitions by reapplying map and filter.

4. Job Execution: Stages and Tasks

Job : Triggered by an action (e.g., count()).

A Job is created when an action (e.g., count(), collect(), write()) is triggered.

What Happens :

• The driver submits the job to the Spark cluster.
• The job is split into stages based on data dependencies (shuffles).

df = spark.read.parquet("data")  # Lazy (no job yet)
df.count()  # Action → Job starts

Stages : Split based on shuffle dependencies (e.g., reduceByKey).

A Stage is a group of tasks that can run in parallel without requiring data shuffling.

Key Points :

• Stages are separated by shuffle boundaries (e.g., groupBy, join).
• Each stage processes data in partitions .
• The DAG Scheduler analyzes dependencies between transformations.
• Narrow dependencies (e.g., filter, map) are grouped into the same stage .
• Wide dependencies (e.g., reduceByKey, join) split the job into new stages .

df.write.parquet("output")  # Action triggers a job

If the data requires a shuffle (e.g., repartitioning), the job splits into:

• Stage 1 : Read data and prepare for shuffle.
• Stage 2 : Write shuffled data to output.

Tasks :

A Task is the smallest unit of work in a stage. Each task processes one partition of data.

What Happens :

• Each stage is split into tasks (one per data partition).
• Tasks are sent to executors for parallel processing.

Key Details :

• Task granularity : Determined by the number of partitions.
• Default partitions : Typically equal to the cluster’s total cores.
• Number of tasks = Number of partitions in the DataFrame/RDD.
• Tasks run in parallel on executor nodes .

Example : If a DataFrame has 10 partitions, a stage will have 10 tasks (one per partition).

• Each stage has tasks (one per partition).
• Tasks are sent to executors for parallel processing.
• A DataFrame with 4 partitions → 4 tasks in a stage.
• Executors process tasks in parallel.

Data Flow

Driver :

• Converts PySpark code into JVM-compatible bytecode (via Py4J).
• Manages DAG creation and task scheduling.

Executors :

• Run JVM tasks for data processing (e.g., filtering, aggregation).
• For Python-specific code (e.g., UDFs), executors launch Python subprocesses (slower due to serialization).

5. JVM and Python Integration

PySpark bridges Python and JVM (Java Virtual Machine) using Py4J :

• Driver : Python code runs in a Python interpreter.
• Workers : Data processing happens in JVM-based executors.

Data Flow :

Serialization :

• Python objects (e.g., lists) are serialized (via Pickle or Arrow) and sent to JVM.

Processing :

• JVM processes data using optimized Scala/Java code.

Deserialization :

• Results are sent back to Python.

Performance Considerations :

• Avoid Python UDFs : They are slower due to serialization overhead.
• Use DataFrames : Catalyst optimizer and Tungsten engine (JVM-based) improve performance.

6. Key Optimizations

• Catalyst Optimizer : Optimizes logical/physical plans (e.g., predicate pushdown).
• Tungsten Engine : Efficient memory management and binary data processing.
• Arrow : Accelerates data transfer between Python and JVM.

7. Common Pitfalls

• Data Skew : Uneven partitions slow down jobs (use repartition or salting).
• Garbage Collection : JVM GC pauses can affect performance (tune memory settings).
• Serialization Overhead : Minimize data transfers between Python and JVM.

8. Monitoring and Debugging

• Spark UI : Track stages, tasks, and DAG visualization at [http://driver:4040.](http://driver:4040.)
• Logs : Check executor/driver logs for errors.

Conclusion

PySpark’s power lies in its ability to leverage JVM-based optimizations while providing Python’s ease of use. Understanding DAGs, lineage, and the JVM/Python bridge helps write efficient, fault-tolerant code. For large-scale data engineering, prefer DataFrame APIs and Catalyst-optimized workflows.

• Python code runs in the driver, but tasks execute in JVM-based executors.
• Use Arrow (spark.sql.execution.arrow.enabled=true) to optimize Python-JVM data transfer.