PySpark for Beginners: Building Intermediate-Level Skills

my first two articles in this series, you’ve already covered a lot of ground.

Table of Contents

You know how to create a SparkSession, load data into a DataFrame, define schemas, clean messy columns, join datasets, write Parquet files, and build simple PySpark workflows. I’ll put links to those articles in this series at the end of this one.

That is already enough to do useful work. But once your data grows or your transformations become more complex, a new set of questions starts to appear.

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 is about those next steps.

We are still keeping things friendly. We’re not going deep into cluster tuning, memory internals, Kubernetes, YARN, or obscure PySpark configuration settings. Instead, we’ll focus on practical ideas that help you understand what PySpark does and how to write PySpark code that behaves more predictably.

The overarching theme is this:

Attaining intermediate level PySpark is mostly about understanding data movement.

Once that idea clicks, PySpark starts to feel a lot more predictable. There’s usually a good reason why your PySpark jobs are underperforming; you just need to know where to look.

When PySpark starts to feel slow

When you first learn PySpark, it’s natural to focus on the code.

You write things like:

df.filter(df.city == "London") 
df.groupBy("customer_id").sum("amount") 
df.join(customers, on="customer_id")

But as you move into more advanced PySpark, you need to start thinking about what PySpark has to do behind the scenes.

Some operations are fairly cheap. Filtering rows is a good example. Spark can look at each row, make a yes/no decision, and move on.

df_london = df.filter(df.city == "London")

Other operations are more expensive because PySpark has to move data around. Consider this. Like the filter, it’s just one line of code, but PySpark has to make sure all rows for the same customer end up together.

df_totals = df.groupBy("customer_id").sum("amount")

That’s a much more complicated operation and may need data to move between different parts of the job. That movement is often where the cost is.

So the intermediate level habit is simple:

Before reaching for tuning settings, ask a simpler question: is this line of code forcing Spark to move data around?

This doesn’t mean joins, aggregations, or data sorting are bad. On the contrary, they’re essential, but they are the type of operations you should treat with a little more care.

How PySpark divides the work

A PySpark DataFrame isn’t treated as a single contiguous block of data. PySpark splits it into partitions. You can think of the partitions as chunks of data that PySpark can process separately and in parallel. That is one of the main reasons PySpark is performant and useful for processing large-scale data.

You can check how many partitions a DataFrame has like this:

df.rdd.getNumPartitions()

As output, you’ll get a simple integer number back, for example:

That means PySpark has split the data into eight chunks, and this matters because partitioning affects how much parallel work PySpark can do. If you have too few partitions, PySpark may not be able to use all available CPU cores. If you have too many partitions, PySpark may spend too much time managing lots of tiny tasks and data sets. Neither extreme is ideal.

If either of these situations arises, PySpark offers two useful functions to help.

You can change the number of partitions with repartition()

df_more = df.repartition(16)

And you can reduce the number of partitions with coalesce():

df_fewer = df.coalesce(4)

repartition() can increase or decrease the number of partitions, but it causes PySpark to redistribute the data, which will usually have a performance hit.

coalesce() is used to reduce the number of partitions and needs less data movement.

A practical rule is:

Use repartition() when you need to spread data out more evenly. Use coalesce() when you simply want fewer output files.

For example, say you have a dataset that naturally produces 40 small-ish Parquet files when it’s written out to disk. You might do this instead:

df.coalesce(4).write.mode("overwrite").parquet("output/sales")

That tells PySpark to write the result with four partitions, which typically produces four output files.

You don’t need to obsess over exact partition counts at this stage. The important point is that partitions control how PySpark divides its work. Once you understand these ideas, PySpark becomes much easier to reason about. Jobs that slow down stop feeling like random events and you start to see the likely causes: a shuffle here, a wide join there, too many tiny files somewhere else.

When data has to move

Some PySpark operations can happen inside each partition. Others require data from different partitions to be brought together.

That movement is called a shuffle.

A shuffle happens when PySpark has to redistribute data across partitions. This is one of the first Spark ideas that really changes how you write PySpark. Shuffles aren’t rare, and they aren’t automatically wrong, but they are usually computationally expensive and should be avoided if possible.

Common PySpark operations that may cause shuffles include:

groupBy()
join()
distinct()
orderBy()
repartition()

Take this example:

df.groupBy("customer_id").sum("amount")

Spark can’t calculate the final total for each customer unless all rows for that customer are aggregated. That normally means moving data between partitions.

Sorting is another common example:

df.orderBy("transaction_date")

Sorting a whole dataset requires PySpark to compare and arrange data across the entire dataset, not just within each partition. That’s another shuffle operation. The main lesson is:

Shuffles are expensive.

You will need them because real data processing work depends on grouping, joining, and sorting. But if a PySpark job feels slow, shuffles are one of the first things to investigate.

A useful habit to reduce shuffles is to filter and select data before expensive operations that might trigger them.

Instead of this:

df_joined = df_sales.join(df_customers, "customer_id") 
df_london = df_joined.filter(df_joined.city == "London")

Prefer this, where possible:

df_london_customers = df_customers.filter(df_customers.city == "London") 
df_joined = df_sales.join(df_london_customers, "customer_id")

Removing unnecessary data as early as possible in your pipelines means PySpark has less data to move (shuffle) during the join.

Why joins deserve extra care

Joins are where a lot of PySpark jobs start to hit performance issues. Not because joins are bad, but because matching rows across datasets often means moving data.

A simple join looks like this:

df_joined = df_sales.join(df_customers, on="customer_id", how="left")

Spark needs matching customer IDs from both the sales and customers DataFrames to meet in the same place. Depending on the size and layout of the data, that can involve a shuffle. There are three practical things you can do to make joins less painful.

First, try to ensure that your join keys are of the same data type. If that’s not possible, make sure you cast data types appropriately. If one table stores customer IDs as integers and another stores them as strings, PySpark may produce errors or unexpected results.

from pyspark.sql import functions as F  
df_sales = df_sales.withColumn(
    "customer_id",
    F.col("customer_id").cast("string")
)
df_customers = df_customers.withColumn(
    "customer_id",F.col("customer_id").cast("string") 
)

Second, only keep the columns you need by reducing data before joining.

df_customers_small = df_customers.select("customer_id","city","loyalty_level")

And filter early if possible:

df_sales_recent = df_sales.filter(F.col("year") == 2026)

Then join:

df_joined = df_sales_recent.join( df_customers_small,on="customer_id",how="left")

Third, consider broadcasting small lookup tables. If one DataFrame is very small, PySpark can copy it to each worker rather than shuffle both datasets.

from pyspark.sql.functions import broadcast  
df_joined = df_sales.join(broadcast(df_customers_small),on="customer_id,"how="left" )

This is called a broadcast join and it’s especially useful for lookup tables, such as customer details, country codes, product categories, or exchange rates.

A simple rule to remember is this:

If one side of the join is small enough to fit comfortably in memory, a broadcast join may be faster.

Don’t force broadcast joins everywhere. They are useful when the smaller table is genuinely small. Trying to broadcast a large table can create more problems than it solves.

Looking at what PySpark is planning

Because PySpark uses lazy evaluation, it doesn’t run each transformation as soon as it reads it. Instead, it builds up a plan for how the work should be done.

You can take a look at that plan with:

df.explain("formatted")

The first time you view an execution plan, the output can feel overwhelming. You may see terms such as:

Scan 
Filter 
Project 
Exchange 
SortMergeJoin 
BroadcastHashJoin 
HashAggregate

You don’t need to understand every line. At this stage, just start recognising a few useful clues.

Scan means PySpark is reading data.
Filter means PySpark is applying a filter.
Project usually means PySpark is selecting columns or creating derived columns.
Exchange usually means a shuffle.
SortMergeJoin means PySpark is joining data by sorting and matching keys.
HashAggregate means PySpark is grouping rows and calculating aggregate values, such as counts, sums, averages, minimums, or maximums.
BroadcastHashJoin means PySpark is using a broadcast join.

The one I would pay attention to first is:

If you see this, it means PySpark is shuffling data between partitions. Again, that isn’t automatically bad. But it tells you where some of the costs may be. For example:

df_totals = df.groupBy("customer_id").sum("amount") 
df_totals.explain("formatted")

You will likely see an Exchange in the plan because PySpark needs to group matching customer IDs together.

You can also use the PySpark UI to investigate issues, which I talked about briefly in my previous article in this series. When PySpark is running locally, you can usually open the PySpark UI at this URL:

It shows you which jobs have run, how long each stage took, and whether one part of your job is taking much longer than the rest.

You don’t need to understand every TAB in the output of the Spark UI at this point. Just get used to opening it now and again to check that your jobs are progressing as they should.

Using explain() and the PySpark UI is one of the best habits you can develop. You don’t need to become a PySpark internals expert. Just start looking for obvious clues. After a while, you stop trying to read every line and start looking for the handful of clues that matter.

Data Caching

Caching is tempting because it feels like a magic “make this faster” button. Sometimes it does. Other times, it just uses up memory without helping much.

You can mark a DataFrame for caching like this:

df_clean.cache()

But remember: PySpark is lazy. This line doesn’t immediately cache the data. It only tells PySpark:

When this DataFrame is computed, keep it in memory if possible.

To actually materialise the cache, you need an action like:

df_clean.count()

Now PySpark computes df_clean and stores it. Caching is generally only useful when you reuse the same DataFrame multiple times in the same session.

For example:

df_clean.cache() 
df_clean.count()  
df_by_city = df_clean.groupBy("city").sum("amount") 
df_by_level = df_clean.groupBy("loyalty_level").sum("amount") 
df_sample = df_clean.filter(F.col("amount") > 1000)

Here, caching may help because the same cleaned dataset is reused in several later operations, but caching isn’t useful if you only use the DataFrame once.

For example, this cache call is pointless:

df_clean.cache() 
df_final = df_clean.groupBy("city").sum("amount")

If df_clean isn’t reused, caching just adds overhead.

You can remove a DataFrame from cache with:

df_clean.unpersist()

A good rule of thumb is:

Cache only when a DataFrame is expensive to create and reused more than once. Don’t cache everything. PySpark memory is valuable.

Writing data that PySpark can read efficiently

In my previous PySpark article, I introduced Parquet as the preferred file format for reading and writing data. Now we can go one step further.

You can write Parquet data partitioned by one or more columns:

df.write.mode("overwrite") \    
 .partitionBy("year", "month") \     
 .parquet("output/sales")

Depending on the data in the DataFrame, this would create an output folder layout similar to this:

output/sales/   
             -> year=2025/     
                          ->  month=1/   
                          ->  month=5/   
                          ->  month=12/   
             -> year=2026/  
                          ->  month=1/  
                          ->  month=2/

The useful part is that when reading data back in from a folder structure like this, Spark doesn’t always have to read the whole dataset.

If you later read only the 2026 data:

df_2026 = spark.read.parquet("output/sales").filter(F.col("year") == 2026)

Spark can often avoid reading the 2025 folder altogether. This is called partition pruning, and it can make reads much faster when your filters match the partition columns.

Good partition columns are usually fields you commonly filter by, such as:

year 
month 
week
region 
country 
department

But be careful. Don’t partition by a column with too many unique values.

For example, this is usually a bad idea:

.partitionBy("transaction_id")

If every transaction ID is unique, PySpark may create thousands or millions of tiny folders and files. That makes the dataset harder to manage and slower to read.

A good rule to remember is:

Partition by columns you often filter on, but avoid columns with too many distinct values. Date-based partitioning is often a sensible place to start.

Letting PySpark do the work

In PySpark, a UDF stands for a user-defined function. It lets you write your own Python function and apply it to a PySpark column in a Dataframe.

This is genuinely useful in some cases. But Python UDFs can be slower than PySpark’s built-in functions.

Here is a simple built-in version example where we want to convert some text to uppercase:

from pyspark.sql import functions as F  
df = df.withColumn("customer_name_upper",F.upper("customer_name") )

Spark understands F.upper(), and it can optimise it as part of the execution plan.

Now compare that with a Python UDF:

from pyspark.sql.functions import udf 
from pyspark.sql.types import StringType  

def make_upper(value):     
    if value is None:         
        return None    
    return value.upper()  

make_upper_udf = udf(make_upper, StringType())  
df = df.withColumn("customer_name_upper", make_upper_udf("customer_name") )

This works but it’s more verbose and harder to understand and, crucially, PySpark has less ability to optimise it. Data may need to move between PySpark’s engine and Python, which adds overhead.

So the practical intermediate level advice is:

Choose PySpark built-in functions first. Use UDFs sparingly and only when there is no built-in function that does what you need.

Common traps as your data grows

Once you start working with larger datasets, some innocent-looking PySpark code can cause problems. One common trap is using collect() on large data.

rows = df.collect()

This brings all data from PySpark into Python on your driver machine.

For small data, that is fine. For large datasets, this can overwhelm the driver because all rows are loaded into local Python memory.

Instead, use:

df.show()

To see what your data looks like before deciding what to do with it. Viewing the data set might help you formulate a plan to filter it more precisely or process it more optimally.

Another trap is calling count() too often.

df.count()

Like show(), count() is an action. It causes PySpark to scan the entire dataset. It can be useful, but don’t use it everywhere without thinking.

Sorting huge datasets can also be expensive.

df.orderBy("amount")

Global sorting often requires a shuffle. If you only need a quick look at the largest few rows, you might use this instead:

df.orderBy(F.col("amount").desc()).limit(10)

Spark can often optimise this as a top-k query. It still needs to scan the data, but it may avoid sorting and returning the entire dataset.

Another common problem is joining data before filtering it. Where possible, filter each dataset first, so Spark has fewer rows to move and compare during the join.

Writing hundreds or thousands of small files can also slow down future reads. When the output is small enough, using coalesce() can reduce the number of files Spark writes:

df.coalesce(8).write.mode("overwrite").parquet("output/results")

And finally, avoid, avoid, avoid using Python loops where a DataFrame operation would work better. This pattern can trigger many separate PySpark jobs:

for city in cities:     
    df.filter(F.col("city") == city).count()

This is better:

df.groupBy("city").count()

Spark is happiest when you describe the result you want and let it plan the route.

Putting the pieces together

Let’s pull the ideas we’ve discussed together into a simple workflow. Imagine we have sales data stored as partitioned Parquet, plus a small customer lookup table.

We want to:

read recent sales
keep only useful columns
clean join keys
join to customer data
calculate totals
write partitioned output
inspect the plan

Here is what that might look like in PySpark code:

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")
)

There’s nothing flashy here. But this is a much more mature PySpark workflow than simply loading a CSV and running a join.

What does the execution plan look like?

Since this workflow uses a broadcast join and an aggregation, it’s also a good moment to look at the plan Spark might create. The exact plan depends on your PySpark version, data size, file layout, and whether PySpark decides to honour the broadcast hint. But for this example, the simplified version of the plan might look something like this:

== Physical Plan ==
AdaptiveSparkPlan
+- WriteFiles
   +- Sort
      +- HashAggregate
         +- Exchange
            +- HashAggregate
               +- Project
                  +- BroadcastHashJoin LeftOuter
                     :- Project
                     :  +- Filter
                     :     +- Scan parquet data/sales
                     +- BroadcastExchange
                        +- Project
                           +- Scan parquet data/customers

The key things to notice are:

BroadcastHashJoin - This tells you PySpark is using the broadcast join, which is normally a good thing.

Exchange – This tells you PySpark is doing a shuffle. In this workflow, the shuffle is expected because of:

.groupBy("year", "month", "city")

Filter – And if your sales data is partitioned by year, you may also see PySpark pruning files when you filter:

.filter(F.col("year") == 2026)

That means PySpark can skip irrelevant folders rather than reading the entire dataset.

Summary

Moving to an intermediate level in PySpark is less about memorising settings and more about understanding the ideas that explain how Spark behaves:

Data is divided into partitions.
Some operations run independently, while others require shuffles.
Shuffles and joins can be expensive.
Caching only helps when data is reused.
Parquet file layout affects read and write performance.
Built-in functions are usually faster than Python UDFs.
Execution plans show how Spark intends to process your data.

Once you understand these ideas, slow jobs become easier to diagnose. That’s the real step up.

Your first PySpark milestone was learning how to run code. Next, was learning how to build clean workflows. This third milestone is learning how to make those workflows scale.

You don’t need to become a PySpark performance engineer overnight. But if you can recognise partitions, shuffles, joins, caching, file layout, and execution plans, you are already writing PySpark with a much more professional mindset.

For reference, here are the links to the first two articles in this series.

PySpark for Beginners: Building Intermediate-Level Skills

When PySpark starts to feel slow

How PySpark divides the work

When data has to move

Why joins deserve extra care

Looking at what PySpark is planning

Data Caching

Writing data that PySpark can read efficiently

Letting PySpark do the work

Common traps as your data grows

Putting the pieces together

What does the execution plan look like?

Summary

Next.js Enterprise Applications: Why Teams Are Migrating

React SSR vs CSR: Choosing CSR, SSR, SSG, or ISR

The Real Challenge Limiting AI Models Today

Assemble Each RAG Generation Prompt from a Base Prompt Plus the Rules Each Question Needs

From Hugging Face to Amazon SageMaker Studio in one click

Leave a Reply Cancel reply

When PySpark starts to feel slow

How PySpark divides the work

When data has to move

Why joins deserve extra care

Looking at what PySpark is planning

Data Caching

Writing data that PySpark can read efficiently

Letting PySpark do the work

Common traps as your data grows

Putting the pieces together

What does the execution plan look like?

Summary

Similar Posts

Leave a Reply Cancel reply