**Scalable Data Workflows: Navigating PySpark Performance and Execution**
In the previous articles of this series, we introduced foundational PySpark concepts, including how to create a `SparkSession`, load and clean data, define schemas, join datasets, write Parquet files, and structure basic workflows. Those skills are enough to accomplish a lot — but only up to a point.
Once your data grows or your transformations become more complex, new questions arise:
– Why did this job suddenly become slow?
– Why does a simple `groupBy()` take longer than expected?
– Why does PySpark write so many output files?
– When should you cache a DataFrame?
– And what on earth is a shuffle?
This article explores those questions, focusing on practical understanding rather than deep cluster tuning or obscure configurations. The goal is to help you write PySpark code that behaves more predictably as workloads increase.
### Attaining Intermediate-Level PySpark Means Understanding Data Movement
The central idea of this stage is simple but powerful:
> **Attaining intermediate-level PySpark is mostly about understanding data movement.**
Once this concept clicks, many performance issues become explainable, and PySpark starts to feel predictable. Jobs that seem slow usually are — and there is usually a reason grounded in how data is being shuffled, transformed, or stored.
—
### How PySpark Divides the Work
A PySpark DataFrame is not a single block of data. Instead, it is split into **partitions**, which are chunks processed independently and in parallel.
– You can check partition count with:
“`python
df.rdd.getNumPartitions()
“`
– Too few partitions limit parallelism; too many create scheduling and management overhead.
– Use `repartition()` to increase or redistribute partitions, and `coalesce()` to reduce them with less data movement.
A practical rule:
> Use `repartition()` when you need even data distribution; use `coalesce()` mainly to reduce output file counts.
—
### When Data Has to Move: Understanding Shuffles
Some operations happen within a partition; others require data to move across partitions. This movement is called a **shuffle**, and it is one of the most important performance concepts in Spark.
Common operations that cause shuffles:
– `groupBy()`
– `join()`
– `distinct()`
– `orderBy()`
– `repartition()`
Example:
“`python
df.groupBy(“customer_id”).sum(“amount”)
“`
To aggregate by customer, Spark must move data so all rows for a given key end up together — and that is costly.
**Key takeaway:**
> Shuffles are expensive and should be minimized, but they are often necessary. Design pipelines to reduce unnecessary data movement.
One effective habit is filtering before expensive operations:
“`python
# Prefer this
df_london = df.filter(df.city == “London”)
df_joined = df_sales.join(df_london, “customer_id”)
# Rather than this
df_joined = df_sales.join(df_customers, “customer_id”)
df_london = df_joined.filter(df_joined.city == “London”)
“`
—
### Why Joins Deserve Extra Care
Joins are a major source of performance issues because matching rows across datasets usually requires shuffling.
Best practices for joins:
1. Ensure join keys share the same data type and consider casting explicitly.
2. Reduce columns and filter rows before joining.
3. For small lookup tables, consider broadcast joins:
“`python
from pyspark.sql.functions import broadcast
df_joined = df_sales.join(broadcast(df_customers_small), “customer_id”, “left”)
“`
> If one side of a join can fit comfortably in memory, a broadcast join is often faster.
—
### Inspecting What PySpark Is Planning
PySpark uses lazy evaluation and builds an execution plan before running anything. You can inspect that plan with:
“`python
df.explain(“formatted”)
“`
Key terms to recognize:
– **Scan**: reading data
– **Filter / Project**: filtering or selecting columns
– **Exchange**: shuffle in progress
– **SortMergeJoin / BroadcastHashJoin**: join strategy
– **HashAggregate**: grouping and aggregation
Watching for **Exchange** in plans is especially useful — it tells you where shuffles are happening and where costs may arise.
You can also use the PySpark UI (usually at `http://localhost:4040` when running locally) to inspect job durations and stages.
—
### Caching: Useful, but Not Free
Caching can feel like a magic performance button, but it is not always helpful:
“`python
df_clean.cache()
df_clean.count() # triggers caching
“`
Rules for caching:
> Cache only when a DataFrame is expensive to create **and** reused multiple times. Avoid caching unnecessarily — memory in PySpark is valuable.
Use `df.unpersist()` to remove cached data when no longer needed.
—
### Writing Data Efficiently
Prefer **Parquet** as your file format. You can also write partitioned Parquet:
“`python
df.write.mode(“overwrite”).partitionBy(“year”, “month”).parquet(“output/sales”)
“`
Partition pruning helps Spark skip irrelevant data when reading later:
“`python
df_2026 = spark.read.parquet(“output/sales”).filter(“year = 2026”)
“`
Good partition columns are low-cardinality fields used often in filters (e.g., `year`, `month`, `region`). Avoid high-cardinality columns like unique IDs.
—
### Prefer Built-in Functions Over Python UDFs
PySpark can optimize built-in functions far better than Python UDFs:
“`python
# Fast and optimizable
df.withColumn(“name_upper”, F.upper(“name”))
# Slower, less optimizable
from pyspark.sql.functions import udf, StringType
to_upper = udf(lambda v: v.upper() if v else None, StringType())
df.withColumn(“name_upper”, to_upper(“name”))
“`
> Use built-ins first; use UDFs only when necessary.
—
### Common Traps as Data Grows
– **`collect()`** brings all data to the driver — dangerous for large datasets.
– **Excessive `count()`** triggers full scans — use sparingly.
– **Global sorts** are expensive; prefer `limit()` for top-k queries.
– **Joining before filtering** moves more data than necessary.
– **Writing thousands of small files** slows future reads; use `coalesce()` to reduce file count.
– **Python loops over data** can spawn many small jobs; prefer vectorized DataFrame operations.
Instead of:
“`python
for city in cities:
df.filter(F.col(“city”) == city).count()
“`
Prefer:
“`python
df.groupBy(“city”).count()
“`
—
### Putting It All Together: A Practical Workflow
Imagine you have partitioned sales data and a small customer lookup table. You want to:
1. Read recent sales
2. Select and clean key columns
3. Cast join keys consistently
4. Join with customer data (broadcast if small)
5. Aggregate totals
6. Write partitioned output
7. Inspect the execution plan
Example code:
“`python
from pyspark.sql import functions as F
from pyspark.sql.functions import broadcast
sales = spark.read.parquet(“data/sales”)
sales_recent = (
sales
.filter(F.col(“year”) == 2026)
.select(“transaction_id”, “customer_id”, “amount”, “year”, “month”)
.dropna(subset=[“customer_id”, “amount”])
.withColumn(“customer_id”, F.col(“customer_id”).cast(“string”))
)
customers = (
spark.read.parquet(“data/customers”)
.select(“customer_id”, “city”, “loyalty_level”)
.withColumn(“customer_id”, F.col(“customer_id”).cast(“string”))
)
enriched = sales_recent.join(
broadcast(customers),
on=”customer_id”,
how=”left”,
)
monthly_totals = (
enriched
.groupBy(“year”, “month”, “city”)
.agg(
F.sum(“amount”).alias(“total_amount”),
F.count(“*”).alias(“transaction_count”),
)
)
monthly_totals.explain(“formatted”)
(
monthly_totals.write
.mode(“overwrite”)
.partitionBy(“year”, “month”)
.parquet(“output/monthly_totals”)
)
“`
The execution plan will likely show:
– `BroadcastHashJoin`
– `Exchange` (expected due to aggregation)
– Partition pruning on `year`
—
### Summary
Advancing to intermediate PySpark is less about memorizing settings and more about understanding core ideas:
– Data lives in partitions
– Some operations cause shuffles
– Shuffles and joins are expensive
– Caching only helps with reuse
– File layout affects performance
– Built-in functions outperform UDFs
– Execution plans reveal how Spark will process your job
Once you recognize these patterns, slow jobs become easier to diagnose and fix. You move from running code to building clean workflows — and finally to building workflows that scale.
**References**
– *Previous articles in this PySpark series* (available in the original post)
– PySpark documentation on [DataFrame transformations and actions](https://spark.apache.org/docs/latest/api/python/)
– Best practices on partitioning, caching, and join strategies
—
**Original article source:**
The content above is based on the post **”From Zero to Hero: Understanding Data Movement and Performance in PySpark – Part 3″**, available at the original author’s site (referenced in the opening links of the source post).
