Most engineers never look at this layer, which is why they spend days tuning configurations that don't address the real problem: inefficient transformations that generate bloated plans.

This handbook teaches you to read, interpret, and control those plans, transforming you from someone who writes PySpark code into someone who architects efficient data pipelines with precision and confidence.

Table of Contents

1. Background Information

2. Chapter 1: The Spark Mindset: Why Plans Matter

3. Chapter 2: Understanding the Spark Execution Flow

4. Chapter 3: Reading and Debugging Plans Like a Pro

5. Chapter 4: Writing Efficient Transformations

   • Scenario 1: Rename in One Pass: withColumnRenamed() vs toDF()

   • Scenario 2: Reusing expressions

   • Scenario 3: Batch column ops

   • Scenario 4: Early Filter vs Late Filter

   • Scenario 5: Column Pruning

   • Scenario 6: Filter pushdown vs full scan

   • Scenario 7: De-duplicate right

   • Scenario 8: Count Smarter

   • Scenario 9: Window wisely

   • Scenario 10: Incremental Aggregations with Cache and Persist

   • Scenario 11: Reduce shuffles

   • Scenario 12: Know Your Shuffle Triggers

   • Scenario 13: Tune Parallelism: Shuffle Partitions & AQE

   • Scenario 14: Handle Skew Smartly

   • Scenario 15: Sort Efficiently (orderBy vs sortWithinPartitions)

6. Conclusion

Background Information

What This Handbook is Really About

This is not a tutorial about Spark internals, cluster tuning, or PySpark syntax or APIs.

This is a handbook about writing PySpark code that generates efficient logical plans.

Because when your code produces clean, optimized plans, Spark pushes filters correctly, shuffles reduce instead of multiply, projections stay shallow, and the DAG (Directed Acyclic Graph) becomes predictable, lean, and fast.

When your code produces messy plans, Spark shuffles more than necessary, and projects pile up into deep, expensive stacks. Filters arrive late instead of early, joins explode into wide, slow operations, and the DAG becomes tangled and expensive.

The difference between a fast job and a slow job is not “faster hardware.” It’s the structure of the plan Spark generates from your code. This handbook teaches you to shape that plan deliberately through scenarios.

Who This Handbook Is For

This handbook is written for:

• Data engineers building production ETL pipelines who want to move beyond trial-and-error tuning and understand why jobs perform the way they do

• Analytics engineers working with large datasets in Databricks, EMR, or Glue who need to optimize Spark jobs but don't have time for thousand-page reference manuals

• Data scientists transitioning from pandas to PySpark who find themselves writing code that technically runs but takes forever

• Anyone who has stared at the Spark UI, seen mysterious "Exchange" nodes in the DAG, and wondered, "Why is this shuffling so much data?"

You should already be comfortable writing basic PySpark code , creating DataFrames, applying transformations, running aggregations. This handbookbook won't teach you Spark syntax. Instead, it teaches you how to write transformations that work with the optimizer, not against it.

How This Handbook Is Structured

We’ll start with foundations, then move on to real-world scenarios.

Chapters 1-3 build your mental model. You'll learn what logical plans are, how they connect to physical execution, and how to read the plan output that Spark shows you. These chapters are short and focused – just enough theory to make the practical scenarios meaningful.

Chapter 4 is the heart of the handbook. It contains 15 real-world scenarios, organized by category. Each scenario shows you a common performance problem, explains what's happening in the logical plan, and demonstrates the better approach. You'll see before-and-after code, plan comparisons, and clear explanations of why one approach outperforms another.

What You'll Learn

By the end of this handbook, you'll be able to:

• Read and interpret Spark's logical, optimized, and physical plans

• Identify expensive operations before running your code

• Restructure transformations to minimize shuffles

• Choose the right join strategies for your data

• Avoid common pitfalls that cause memory issues and slow performance

• Debug production issues by examining execution plans

More importantly, you'll develop a Spark mindset, an intuition for how your code translates to cluster operations. You'll stop writing code that "should work" and start writing code that you know will work efficiently.

Technical Prerequisites

I assume that you’re familiar with the following concepts before proceeding:

1. Python fundamentals

2. PySpark basics

   • Creating DataFrames and reading data from files

   • Basic DataFrame operations: select, filter, withColumn, groupBy, join

   • Writing DataFrames back to storage

3. Basic Spark concepts

   • Basic understanding of Spark applications, jobs, stages, and tasks

   • Basic understanding of the difference between transformations and actions

   • Understanding. of partitions and shuffles

4. AWS Glue (Good to have)

Chapter 1: The Spark Mindset: Why Plans Matter

This chapter isn’t about Spark theory or internals. It’s about understanding Spark Plans, and seeing Spark the way the engine sees your code. Once you understand how Spark builds and optimizes a logical plan, optimization stops being trial and error and becomes intentional engineering.

Behind every simple transformation, Spark quietly redraws its internal blueprint. Every transformation you write from "withColumn" to join changes that plan. When the plan is efficient, Spark flies, but when it’s messy, Spark crawls.

The Invisible Layer Behind Every Transformation

When you write PySpark code, it feels like you’re chaining operations step by step. In reality, Spark isn’t executing those lines. It’s quietly building a blueprint, a logical plan describing what to do, not how.

Once this plan is built, the Catalyst Optimizer analyzes it, rearranges operations, eliminates redundancies, and produces an optimized plan. Catalyst is Spark’s query optimization engine.

Every DataFrame or SQL operation we write, such as select, filter, join, groupBy, is first converted into a logical plan. Catalyst then analyzes and transforms this plan using a set of rule-based optimizations, such as predicate pushdown, column pruning, constant folding, and join reordering. The result is an optimized logical plan, which Spark later converts into a physical execution plan. Finally, Spark translates that into a physical plan of what your cluster actually runs. This invisible planning layer decides the job’s performance more than any configuration setting.

From Logical to Optimized to Physical Plans

When you run df.explain(True), Spark actually shows you four stages of reasoning:

1. Logical Plan

The logical plan is the first stage where the initial translation of the code results in a tree structure that shows what operations need to happen, without worrying about how to execute them efficiently. It’s a blueprint of the query’s logic before any optimization or physical planning occurs.

This:

df.filter(col('age') > 25) \
  .select('firstname', 'country') \
  .groupby('country') \
  .count() \
  .explain(True)

results in the following logical plan:

== Parsed Logical Plan ==
'Aggregate ['country], ['country, 'count(1) AS count#108]
+- Project [firstname#95, country#97]
   +- Filter (age#96L > cast(25 as bigint))
      +- LogicalRDD [firstname#95, age#96L, country#97], false

2. Analyzed Logical Plan

The analyzed logical plan is the second stage in Spark’s query optimization. In this stage, Spark validates the query by checking if tables and columns actually exist in the Catalog and resolving all references. It converts all the unresolved logical plans into a resolved one with correct data types and column bindings before optimization.

3. Optimized Logical Plan

The optimized logical plan is where Spark's Catalyst optimizer improves the logical plan by applying smart rules like filtering data early, removing unnecessary columns, and combining operations to reduce computation. It's the smarter, more efficient version of your original plan that will execute faster and use fewer resources.

Let’s understand using a simple code example:

df.select('firstname', 'country') \
  .groupby('country') \
  .count() \
  .filter(col('country') == 'USA') \
  .explain(True)

Here’s the parsed logical plan:

== Parsed Logical Plan ==
'Filter '`=`('country, USA)
+- Aggregate [country#97], [country#97, count(1) AS count#122L]
   +- Project [firstname#95, country#97]
      +- LogicalRDD [firstname#95, age#96L, country#97], false

What this means:

• Spark first projects firstname and country

• Then aggregates by country

• Then applies the filter country = 'USA' after aggregation

(because that’s how you wrote it).

Here’s the optimized logical plan:

== Optimized Logical Plan ==
Aggregate [country#97], [country#97, count(1) AS count#122L]
+- Project [country#97]
   +- Filter (isnotnull(country#97) AND (country#97 = USA))
      +- LogicalRDD [firstname#95, age#96L, country#97], false

Key improvements Catalyst applied:

• Filter pushdown: The filter country = 'USA' is pushed below the aggregation, so Spark only groups U.S. rows.

• Column pruning: “firstname” is automatically removed because it’s never used in the final output.

• Cleaner projection: Intermediate columns are dropped early, reducing I/O and in-memory footprint.

4. Physical Plan

The physical plan is Spark's final execution blueprint that shows exactly how the query will run: which specific algorithms to use, how to distribute work across machines, and the order of low-level operations. It's the concrete, executable version of the optimized logical plan, translated into actual Spark operations like “ShuffleExchange”, “HashAggregate”, and “FileScan” that will run on your cluster.

Catalyst may, for example:

• Fold constants (col("x") * 1 → col("x"))

• Push filters closer to the data source

• Replace a regular join with a broadcast join when data fits in memory

Once the physical plan is finalized, Spark’s scheduler converts it into a DAG of stages and tasks that run across the cluster. Understanding that lineage, from your code → plan → DAG, is what separates fast jobs from slow ones.

How to Read a Logical Plan

A logical plan prints as a tree: the bottom is your data source, and each higher node represents a transformation.

Each node represents a transformation. Execution flows from the bottom up. Let's understand with a basic example:

from pyspark.sql import SparkSession
from pyspark.sql.functions import *
from pyspark.sql.types import *

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

employees_data = [
    (1, "John", "Doe", "Engineering", 80000, 28, "2020-01-15", "USA"),
    (2, "Jane", "Smith", "Engineering", 85000, 32, "2019-03-20", "USA"),
    (3, "Alice", "Johnson", "Sales", 60000, 25, "2021-06-10", "UK"),
    (4, "Bob", "Brown", "Engineering", 90000, 35, "2018-07-01", "USA"),
    (5, "Charlie", "Wilson", "Sales", 65000, 29, "2020-11-05", "UK"),
    (6, "David", "Lee", "HR", 55000, 27, "2021-01-20", "USA"),
    (7, "Eve", "Davis", "Engineering", 95000, 40, "2017-04-12", "Canada"),
    (8, "Frank", "Miller", "Sales", 70000, 33, "2019-09-25", "UK"),
    (9, "Grace", "Taylor", "HR", 58000, 26, "2021-08-15", "Canada"),
    (10, "Henry", "Anderson", "Engineering", 88000, 31, "2020-02-28", "USA")
]

df = spark.createDataFrame(employees_data,  
    ["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"])

Version A: withColumn → filter

In this version, we’re using a derived column "withColumn" and then applying a filter to the dataset. This ordering is logically correct and produces the expected result: it shows how introducing derived columns early affects the logical plan. This example shows what happens when Spark is asked to compute a new column before any rows are eliminated.

df_filtered = df \
.withColumn('bonus', col('salary') * 82) \
.filter(col('age') > 35) \
.explain(True)

Parsed Logical Plan (Simplified)

Filter (age > 35)
└─ Project [*, (salary * 82) AS bonus]
   └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

So what’s going on here? Execution flows from the bottom up.

• Spark first reads the LogicalRDD.

• Then applies the Project node, keeping all columns and adding bonus.

• Finally, the Filter removes rows where age ≤ 35.

This means Spark computes the bonus for every employee, even those who are later filtered out. It's harmless here, but costly on millions of rows, more computation, more I/O, more shuffle volume.

Version B: Filter → Project

In this version, we apply the filter before introducing the derived column. The idea is to show how pushing row-reducing operations earlier allows Catalyst to produce a leaner logical plan. Compared to Version A, this example demonstrates that the same logic, written in a different order, can significantly reduce the amount of work Spark needs to perform.

df_filtered = df \
.filter(col('age') > 35) \
.withColumn('bonus', col('salary') * 82) \
.explain(True)

Parsed Logical Plan (Simplified)

Project [*, (salary * 82) AS bonus]

└─ Filter (age > 35)

└─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

So what’s going on here?

• Spark starts from the LogicalRDD.

• It immediately applies the Filter, reducing the dataset to only employees with age > 35.

• Then the Project node adds the derived column bonus for this smaller subset.

Now the Filter sits below the Project in the plan, cutting data movement and minimizing computation. Spark prunes data first, then derives new columns. This order reduces both the volume of data processed and the amount transferred, leading to a lighter and faster plan.

Why You Should Look at the Plan Every Time by running df.explain(True)

This is the quickest way to spot performance issues before they hit production. It shows:

• Whether filters sit in the right place.

• How many Project nodes exist (each adds overhead).

• Where Exchange nodes appear (these are shuffle boundaries).

• If Catalyst pushed filters or rewrote joins as expected.

A quick explain() takes seconds, while debugging a bad shuffle in production takes hours. Run explain() whenever you add or reorder transformations. The plan never lies.

What Spark Does Under the Hood

Catalyst can sometimes reorder simple filters automatically, but once you use UDFs, nested logic, or joins, it often can’t. That’s why the best habit is to write transformations in a way that already makes sense to the optimizer. Filter early, avoid redundant projections, and keep plans as shallow as possible.

Optimizing Spark isn’t about tuning cluster configs – it’s about writing code that yields efficient plans. If your plan shows late filters, too many projections, or multiple Exchange nodes, it’s already explaining why your job will run slow.

Chapter 2: Understanding the Spark Execution Flow

In Chapter 1, you learned how Spark interprets your transformations into logical plans – blueprints of what the job intends to do.

But Spark doesn't stop there. It must translate those plans into distributed actions across a cluster of executors, coordinate data movement, and handle any failures that may occur.

This chapter reveals what happens when that plan leaves the driver: how Spark breaks your job into stages, tasks, and a directed acyclic graph (DAG) that actually runs.

By the end, you’ll understand why some operations shuffle terabytes while others fly, and how to predict it before execution begins.

From Plans to Stages to Tasks

A Spark job evolves through three conceptual layers:

The Execution Trigger: Actions vs Transformations

Here's the critical distinction that determines when execution actually begins:

df1 = spark.paraquet("data.paraquet")
df2 = spark.filter(col("age") > 25)
df3 = spark.groupby("city").count()

Nothing executes yet! Spark just builds up the logical plan, adding each transformation as a node in the plan tree. No data is read, no filters run, no shuffles happen.

Actions Trigger Execution

Spark transformations are lazy. When a sequence of DataFrame operations is defined, a logical plan is created, but no computation takes place. It’s only when Spark encounters an action, an operation that needs a result to be returned to the driver or written out, that execution takes place.

For example:

result = df3.collect()

At this stage, Spark materializes the logical plan, applies optimizations, creates a physical plan, and executes the job. Until Spark is asked to act, such as collect(), count(), or write(), it’s just describing what it needs to do – but it’s not actually doing it.

The Complete Execution Flow

Spark execution is initiated after the execution of an operation such as collect(). The driver then sends the optimized physical plan to the SparkContext, which is then forwarded to the DAG Scheduler. The physical plan is analyzed to determine shuffle boundaries created by wide operations such as groupBy or orderBy.

The plan is then divided into stages that contain narrow operations. These stages are sent to the Task Scheduler as a TaskSet. Each stage has a single task per partition.

The tasks are then assigned to the cores of the executor based on data locality. The execution of the tasks is then initiated. The execution of the stages is initiated after the completion of the previous stage. The final stage is initiated after the completion of the previous stage. The results of the final stage are then returned to the driver or stored.

What Triggers a Shuffle

A shuffle occurs when Spark needs to redistribute data across partitions, typically because the operation requires grouping, joining, or repartitioning data in a way that can’t be done locally within existing partitions.

Common shuffle triggers:

df.groupBy("user_id”) \
  .agg(sum("amount"))

In Stage 1 (Map), each task performs a partial aggregation on its partition and writes a shuffle file to disk. During the shuffle, each executor retrieves these files across the network such that all records with the same hash(user_id) % numPartitions are colocated.

In Stage 2 (Reduce), each task performs a final aggregation on its partitioned data and writes back to disk. Because Spark has tracked this process as a DAG, a failed task can re-read only the affected shuffle files instead of re-computing the entire DAG.

In practice, a healthy job has 2-6 stages. Seeing 20+ stages for such simple logic usually means unnecessary shuffles or bad partitioning.

Why Shuffles Create Stage Boundaries

Shuffles force data to move across the network between executors. Spark cannot continue processing until:

• All tasks in the current stage write their shuffle output to disk

• The shuffle data is available for the next stage to read over the network

This dependency creates a natural boundary – so a new stage begins after every shuffle. The DAG Scheduler uses these boundaries to determine where stages must wait for previous stages to complete.

Common Performance Bottlenecks

Let’s ground this with a small but realistic pattern using the same employees DataFrame. Goal: average salary per department and country, only for employees older than 30.

Naïve approach:

from pyspark.sql.functions import col, when, avg

df_dept_country = df.select("department", "country").distinct()

df_result = (
    df.withColumn(
        "age_group",
        when(col("age") < 30, "junior")
        .when(col("age") < 40, "mid")
        .otherwise("senior")
    )
    .join(df_dept_country, ["department"], "inner")
    .groupBy("department", "country")
    .agg(avg("salary").alias("avg_salary"))

This looks harmless, but:

• The join on "department" introduces a wide dependency → shuffle #1.

• The groupBy("department", "country") introduces another wide dependency → shuffle #2.

So we have two shuffles for what should be a simple aggregation. If you run explain on the df_result, you’ll see two exchange nodes, each marking a shuffle and stage boundary.

Optimized Approach

We can do better by filtering early, broadcasting the small dimension (df_dept_country), and keeping only one global shuffle for aggregation.

from pyspark.sql.functions import broadcast

df_dept_country = df.select("department", "country").distinct()

df_result_optimized = (
    df.filter(col("age") > 30)
        .join(broadcast(df_dept_country), ["department"], "inner")
        .groupBy("department", "country")
        .agg(avg("salary").alias("avg_salary"))
)

What changed:

• filter(col("age") > 30) is narrow and runs before any shuffle.

• broadcast(df_dept_country) avoids a shuffle for the join.

• Only the groupBy("department", "country") causes a single shuffle.

Now explain shows just one Exchange.

Chapter 3: Reading and Debugging Plans Like a Pro

As explained in Chapter 1, Spark executes transformations based on three levels: the logical plan, the optimized logical plan (Catalyst), and the physical plan. This chapter will expand on this explanation and concentrate on the impact of the logical plan on shuffle and execution performance.

By now, you understand how Spark builds and executes plans. But reading those plans and instantly spotting inefficiencies is the real superpower of a performance-focused data engineer.

Spark’s explain() output isn’t random jargon. It’s a precise log of Spark’s thought process. Once you learn to read it, every optimization becomes obvious.

Three Layers in Spark

As we talked about above, every Spark plan has three key views, printed when you call df.explain(True). Let’s review them now:

1. Parsed Logical Plan: The raw intent Spark inferred from your code. It may include unresolved column names or expressions.

2. Analyzed / Optimized Logical Plan: After Spark applies Catalyst optimizations: constant folding, predicate pushdown, column pruning, and plan rearrangements.

3. Physical Plan: What your executors actually run: joins, shuffles, exchanges, scans, and code-generated operators.

Each stage narrows the gap between what you asked Spark to do and what Spark decides to do.

df_avg = df.filter(col("age") > 30)
        .groupBy("department")
        .agg(avg("salary").alias("avg_salary"))

df_avg.explain(True)

1. Parsed Logical Plan

== Parsed Logical Plan ==
'Aggregate ['department], ['department, 'avg('salary) AS avg_salary#8]
+- Filter (age#5L > cast(30 as bigint))
   +- LogicalRDD [id#0L, firstname#1, lastname#2, department#3, salary#4L, age#5L, hire_date#6, country#7], false

How to read this

• Bottom → data source (LogicalRDD).

• Middle → Filter: Spark hasn’t yet optimized column references.

• Top → Aggregate: high-level grouping intent.

At this stage, the plan may include unresolved symbols (like 'department or 'avg('salary)), meaning Spark hasn’t yet validated column existence or data types.

2. Optimized Logical Plan

== Optimized Logical Plan ==
Aggregate [department#3], [department#3, avg(salary#4L) AS avg_salary#8]
+- Project [department#3, salary#4L]
   +- Filter (isnotnull(age#5L) AND (age#5L > 30))
      +- LogicalRDD [id#0L, firstname#1, lastname#2, department#3, salary#4L, age#5L, hire_date#6, country#7], false

Here, Catalyst has done its job:

• Column IDs (#11, #12L) are resolved.

• Unused columns are pruned – no need to carry them forward.

• The plan now accurately reflects Spark’s optimized logical intent.

If you ever wonder whether Spark pruned columns or pushed filters, this is the section to check.

3. Physical Plan

== Physical Plan ==
AdaptiveSparkPlan isFinalPlan=false
+- HashAggregate(keys=[department#3], functions=[avg(salary#4L)], output=[department#3, avg_salary#8])
   +- Exchange hashpartitioning(department#3, 200), ENSURE_REQUIREMENTS, [plan_id=19]
      +- HashAggregate(keys=[department#3], functions=[partial_avg(salary#4L)], output=[department#3, sum#20, count#21L])
         +- Project [department#3, salary#4L]
            +- Filter (isnotnull(age#5L) AND (age#5L > 30))
               +- Scan ExistingRDD[id#0L,firstname#1,lastname#2,department#3,salary#4L,age#5L,hire_date#6,country#7]

Breakdown

• Scan ExistingRDD → Spark reading from the in-memory DataFrame.

• Filter → narrow transformation, no shuffle.

• HashAggregate → partial aggregation per partition.

• Exchange → wide dependency: data is shuffled by department.

• Top HashAggregate → final aggregation after shuffle.

This structure – partial agg → shuffle → final agg – is Spark’s default two-phase aggregation pattern.

Recognizing Common Nodes

Repeated invocations of withColumn stack multiple Project nodes, which increases the plan depth. Instead, combine these invocations using select.

Multiple Exchange nodes imply repeated data shuffles. You can eliminate these by broadcasting the data or filtering.

Multiple scans of the same table within a single operation imply that some caching of strategic intermediates is lacking.

And frequent SortMergeJoin operations imply that Spark is unnecessarily sorting and shuffling the data. You can eliminate these by broadcasting the smaller dataframe or bucketing.

Debugging Strategy: Read Plans from Top to Bottom

Remember: Spark executes plans from bottom up (from data source to final result). But when you're debugging, you read from the top down (from the output schema back to the root cause). This reversal is intentional: you start with what's wrong at the output level, then trace backward through the plan to find where the inefficiency was introduced.

When debugging a slow job:

• Start at the top: Identify output schema and major operators (HashAggregate, Join, and so on).

• Scroll for Exchanges: Count them. Each = stage boundary. Ask “Why do I need this shuffle?”

• Trace backward: See if filters or projections appear below or above joins.

• Look for duplication: Same scan twice? Missing cache? Re-derived columns?

• Check join strategy: If it’s SortMergeJoin but one table is small, force a broadcast().

• Re-run explain after optimization: You should literally see the extra nodes disappear.

Catalyst Optimizer in Action

Catalyst applies dozens of rules automatically. Knowing a few helps you interpret what changed:

Putting it all together: every plan tells a story:

As you progress through the practical scenarios in Chapter 4, read every plan before and after. Your goal isn't memorization – it's intuition.

Chapter 4: Writing Efficient Transformations

Every Spark job tells a story, not in code, but in plans. By now, you've seen how Spark interprets transformations (Chapter 1), how it executes them through stages and tasks (Chapter 2), and how to read plans like a detective (Chapter 3). Now comes the part where you apply that knowledge: writing transformations that yield efficient logical plans.

This chapter is the heart of the handbook. It's where we move from understanding Spark's mind to writing code that speaks its language fluently.

Why Transformations Matter

In PySpark, most performance issues don’t start in clusters or configurations. They start in transformations: the way we chain, filter, rename, or join data. Every transformation reshapes the logical plan, influencing how Spark optimizes, when it shuffles, and whether the final DAG is streamlined or tangled.

A good transformation sequence:

• Keeps plans shallow, not nested.

• Applies filters early, not after computation.

• Reduces data movement, not just data size.

• Let’s Catalyst and AQE optimize freely, without user-induced constraints.

A bad one can double runtime, and you won't see it in your code, only in your plan.

The Goal of this Chapter

We’ll explore a series of real-world optimization scenarios, drawn from production ETL and analytical pipelines, each showing how a small change in code can completely reshape the logical plan and execution behavior.

Each scenario is practical and short, following a consistent structure. By the end of this chapter, you’ll be able to see optimization opportunities the moment you write code, because you’ll know exactly how they alter the logical plan beneath.

Before You Dive In:

Open a Spark shell or notebook. Load your familiar employees DataFrame. Run every example, and compare the explain("formatted") output before and after the fix. Because in this chapter, performance isn’t about more theory, it’s about seeing the difference in the plan and feeling the difference in execution time.

Scenario 1: Rename in One Pass: withColumnRenamed() vs toDF()

If you’ve worked with PySpark DataFrames, you’ve probably had to rename columns, either by calling withColumnRenamed() repeatedly or by using toDF() in one shot.

At first glance, both approaches produce identical results: the columns have the new names you wanted. But beneath the surface, Spark treats them very differently – and that difference shows up directly in your logical plan.

df_renamed = (df.withColumnRenamed("id", "emp_id")
    .withColumnRenamed("firstname", "first_name")
    .withColumnRenamed("lastname", "last_name")
    .withColumnRenamed("department", "dept")
    .withColumnRenamed("salary", "base_salary")
    .withColumnRenamed("age", "age_years")
    .withColumnRenamed("hire_date", "hired_on")
    .withColumnRenamed("country", "country_code")
)

This is simple and readable. But Spark builds the plan step by step, adding one Project node for every rename. Each Project node copies all existing columns, plus the newly renamed one. In large schemas (hundreds of columns), this silently bloats the plan.

Logical Plan Impact:

Project [emp_id, first_name, last_name, dept, base_salary, age_years, hired_on, country_code]

└─ Project [id, first_name, last_name, dept, base_salary, age_years, hired_on, country_code]

└─ Project [id, firstname, last_name, dept, base_salary, age_years, hired_on, country_code]

└─ Project [id, firstname, lastname, dept, base_salary, age_years, hire_date, country_code]

└─ Project [id, firstname, lastname, department, base_salary, age_years, hire_date, country]

└─ Project [id, firstname, lastname, department, salary, age_years, hire_date, country]

└─ Project [id, firstname, lastname, department, salary, age, hire_date, country]

└─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

Each rename adds a new Project layer, deepening the DAG. Spark now has to materialize intermediate projections before applying the next one. You can see this by running: df.explain(True).

The Better Approach: Rename Once with toDF()

Instead of chaining multiple renames, rename all columns in a single pass:

new_cols = ["id", "first_name", "last_name", "department",
            "salary", "age", "hired_on", "country"]

df_renamed = df.toDF(*new_cols)

Logical Plan Impact:

Project [id, first_name, last_name, department, salary, age, hired_on, country]

└─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

Now there’s just one Project node, which means one projection over the source data. This gives us a flatter, more efficient plan.

Under the Hood: What Spark Actually Does

Every time you call withColumnRenamed(), Spark rewrites the entire projection list. Catalyst treats the rename as a full column re-selection from the previous node, not as a light-weight alias update. When you chain several renames, Catalyst duplicates internal column metadata for each intermediate step.

By contrast, toDF() rebases the schema in a single action. Catalyst interprets it as a single schema rebinding, so no redundant metadata trees are created.

Real-World Timing: Glue Job Benchmark

To see if chained withColumnRenamed calls add real overhead, here's a simple timing test performed on a Glue job using a DataFrame with 1M rows. First using withColumnRenamed:

import time
from pyspark.sql import SparkSession

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

employees_data = [
    (1, "John", "Doe", "Engineering", 80000, 28, "2020-01-15", "USA"),
    (2, "Jane", "Smith", "Engineering", 85000, 32, "2019-03-20", "USA"),
    (3, "Alice", "Johnson", "Sales", 60000, 25, "2021-06-10", "UK"),
    (4, "Bob", "Brown", "Engineering", 90000, 35, "2018-07-01", "USA"),
    (5, "Charlie", "Wilson", "Sales", 65000, 29, "2020-11-05", "UK"),
    (6, "David", "Lee", "HR", 55000, 27, "2021-01-20", "USA"),
    (7, "Eve", "Davis", "Engineering", 95000, 40, "2017-04-12", "Canada"),
    (8, "Frank", "Miller", "Sales", 70000, 33, "2019-09-25", "UK"),
    (9, "Grace", "Taylor", "HR", 58000, 26, "2021-08-15", "Canada"),
    (10, "Henry", "Anderson", "Engineering", 88000, 31, "2020-02-28", "USA")
]

multiplied_data = [(i, f"firstname_{i}", f"lastname_{i}",
                    employees_data[i % 10][3],  # department
                    employees_data[i % 10][4],  # salary
                    employees_data[i % 10][5],  # age
                    employees_data[i % 10][6],  # hire_date
                    employees_data[i % 10][7])  # country
                   for i in range(1, 1_000_001)]

df = spark.createDataFrame(multiplied_data,
                           ["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"])

start = time.time()
df1 = (df
       .withColumnRenamed("firstname", "first_name")
       .withColumnRenamed("lastname", "last_name")
       .withColumnRenamed("department", "dept_name")
       .withColumnRenamed("salary", "annual_salary")
       .withColumnRenamed("age", "emp_age")
       .withColumnRenamed("hire_date", "hired_on")
       .withColumnRenamed("country", "work_country"))

print("withColumnRenamed Count:", df1.count())
print("withColumnRenamed time:", round(time.time() - start, 2), "seconds")

Using toDF:

start = time.time()
df2 = df.toDF("id", "first_name", "last_name", "dept_name", "annual_salary", "emp_age", "hired_on", "work_country")
print("toDF Count:", df2.count())
print("toDF time:", round(time.time() - start, 2), "seconds")

spark.stop()

The difference becomes important at larger sizes or in complex pipelines, especially on managed runtimes such as AWS Glue (where planning overhead becomes important), or when tens of millions of rows are involved, where each additional Project increases column resolution, metadata work, and DAG height. And since Spark can’t collapse chained projections when column names are changed, renaming all columns in one go with toDF() results in a flatter logical and physical plan: one rename, one projection, and faster execution.

Scenario 2: Reusing Expressions

Sometimes Spark jobs run slower, not because of shuffles or joins, but because the same computation is performed repeatedly within the logical plan. Every time you repeat an expression, say, col("salary") * 0.1 in multiple places, Spark treats it as a new derived column, expanding the logical plan and forcing redundant work.

The Problem: Repeated Expressions

Let’s say we’re calculating bonus and total compensation for employees:

df_expr = (
    df.withColumn("bonus", col("salary") * 0.10)
      .withColumn("total_comp", col("salary") + (col("salary") * 0.10))
)

At first glance, it’s simple enough. But Spark’s optimizer doesn’t automatically know that the (col("salary") * 0.10) in the second column is identical to the one computed in the first. Both get evaluated separately in the logical plan.

Simplified Logical Plan:

Project [id, firstname, lastname, department,

salary, age, hire_date, country,

(salary * 0.10) AS bonus,

(salary + (salary * 0.10)) AS total_comp]

└─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

While this looks compact, Spark must compute (salary * 0.10) twice, once for bonus, again inside total_comp. For a large dataset (say 100 M rows), that’s two full column evaluations. The waste compounds when your expression is complex, imagine parsing JSON, applying UDFs, or running date arithmetic multiple times.

The Better Approach: Compute Once, Reuse Everywhere

Compute the expression once, store it as a column, and reference it later:

df_expr = (
    df.withColumn("bonus", col("salary") * 0.10)
      .withColumn("total_comp", col("salary") + col("bonus"))
)

Simplified Logical Plan:

Project [id, firstname, lastname, department,

salary, age, hire_date, country,

(salary * 0.10) AS bonus,

(salary + bonus) AS total_comp]

└─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

Now Spark calculates (salary * 0.10) once, stores it in the bonus column, and reuses that column when computing total_comp. This single change cuts CPU cost and memory usage.

Under the Hood: Why Repetition Hurts

Spark’s Catalyst optimizer doesn’t automatically factor out repeated expressions across different columns. Each withColumn() creates a new Project node with its own expression tree. If multiple nodes reuse the same arithmetic or function, Catalyst re-evaluates them independently.

On small DataFrames, this cost is invisible. On wide, computation-heavy jobs (think feature engineering pipelines), it can add hundreds of milliseconds per task.

Each redundant expression increases:

• Catalyst’s internal expression resolution time

• The size of generated Java code in WholeStageCodegen

• CPU cycles per row, since Spark cannot share intermediate results between columns in the same node

Real-World Benchmark: AWS Glue

We tested this pattern on AWS Glue (Spark 3.3) with 10 million rows and a simulated expensive computation on the similar dataset we used in Scenario 1.

df = spark.createDataFrame(multiplied_data,
                           ["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"])

expr = sqrt(exp(log(col("salary") + 1)))

start = time.time()

df_repeated = (
    df.withColumn("metric_a", expr)
      .withColumn("metric_b", expr * 2)
      .withColumn("metric_c", expr / 10)
)

df_repeated.count()
time_repeated = round(time.time() - start, 2)

start = time.time()

df_reused = (
    df.withColumn("metric", expr)
      .withColumn("metric_a", col("metric"))
      .withColumn("metric_b", col("metric") * 2)
      .withColumn("metric_c", col("metric") / 10)
)

df_reused.count()

print("Repeated expr time:", time_repeated, "seconds")
print("Reused expr time:", round(time.time() - start, 2), "seconds")

spark.stop()

The performance gap widens further with genuinely expensive expressions (like regex extraction, JSON parsing, or UDFs).

Physical Plan Implication

In the physical plan, repeated expressions expand into multiple Java blocks within the same WholeStageCodegen node:

*(1) Project [sqrt(exp(log(salary + 1))) AS metric_a,

(sqrt(exp(log(salary + 1))) * 2) AS metric_b,

(sqrt(exp(log(salary + 1))) / 10) AS metric_c, ...]

Spark literally embeds three copies of the same logic.

Each is JIT-compiled separately, leading to:

• Larger generated Java classes

• Higher CPU utilization

• Longer code-generation time before tasks even start

When reusing a column, Spark generates one expression and references it by name, dramatically shrinking the codegen footprint. If you have complex transformations (nested when, UDFs, regex extractions, and so on), compute them once and reuse them with col("alias"). For even heavier expressions that appear across multiple pipelines, consider persisting the intermediate.

DataFrame:

df_features = df.withColumn("complex_feature", complex_logic)

df_features.cache()

That cache can save multiple recomputations across downstream steps.

Scenario 3: Batch Column Ops

Most PySpark pipelines don’t die because of one big, obvious mistake. They slow down from a thousand tiny cuts: one extra withColumn() here, another there, until the logical plan turns into a tall stack of projections.

On its own, withColumn() is fine. The problem is how we use it:

• 10–30 chained calls in a row

• Re-deriving similar expressions

• Spreading logic across many tiny steps

This scenario shows how batching column operations into a single select() produces a flatter, cleaner logical plan that scales better and is easier to reason about.

The Problem: Chaining withColumn() Forever

from pyspark.sql.functions import col, concat_ws, when, lit

df_transformed = (
    df.withColumn("full_name", concat_ws(" ", col("firstname"), col("lastname")))
      .withColumn("is_senior", when(col("age") >= 35, lit(1)).otherwise(lit(0)))
      .withColumn("salary_k", col("salary") / 1000.0)
      .withColumn("experience_band",
                  when(col("age") < 30, "junior")
                  .when((col("age") >= 30) & (col("age") < 40), "mid")
                  .otherwise("senior"))
      .withColumn("country_upper", col("country").upper())
)

It reads nicely, it runs, and everyone moves on. But under the hood, Spark builds this as multiple Project nodes, one per withColumn() call.

Simplified Logical Plan (Chained): Conceptually

Project [..., country_upper]

└─ Project [..., experience_band]

   └─ Project [..., salary_k]

      └─ Project [..., is_senior]

         └─ Project [..., full_name]

            └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

Each layer re-selects all existing columns, adds one more derived column, and deepens the plan.

The Better Approach: Batch with select()

Instead of incrementally patching the schema, build it once.

df_transformed = df.select(
    col("id"),
    col("firstname"),
    col("lastname"),
    col("department"),
    col("salary"),
    col("age"),
    col("hire_date"),
    col("country"),
    concat_ws(" ", col("firstname"), col("lastname")).alias("full_name"),
    when(col("age") >= 35, lit(1)).otherwise(lit(0)).alias("is_senior"),
    (col("salary") / 1000.0).alias("salary_k"),
    when(col("age") < 30, "junior")
        .when((col("age") >= 30) & (col("age") < 40), "mid")
        .otherwise("senior").alias("experience_band"),
    col("country").upper().alias("country_upper")
)

Simplified Logical Plan (Batched):

Project [id, firstname, lastname, department, salary, age, hire_date, country,

         full_name, is_senior, salary_k, experience_band, country_upper]

└─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

One Project. All derived columns are defined together. Flatter DAG. Cleaner plan.

Under the Hood: Why This Matters

Each withColumn() is syntactic sugar for: “Take the previous plan, and create a new Project on top of it.” So 10 withColumn() calls = 10 projections wrapped on top of each other.

Catalyst can sometimes collapse adjacent Project nodes, but:

• Not always (especially when aliases shadow each other).

• Not when expressions become complex or interdependent.

• Not when UDFs or analysis barriers appear.

Batching with select():

• Gives Catalyst a single, complete view of all expressions.

• Enables more aggressive optimizations (constant folding, expression reuse, pruning).

• Keeps expression trees shallower and codegen output smaller.

Think of it as the difference between editing a sentence 10 times in a row and writing the final sentence once, cleanly.

Real-World Example: Using the Employees DF at Scale:

Chained version (many withColumn()):

from pyspark.sql.functions import col, concat_ws, when, lit, upper
import time

start = time.time()
df_chain = (
    df.withColumn("full_name", concat_ws(" ", col("firstname"), col("lastname")))
      .withColumn("is_senior", when(col("age") >= 35, 1).otherwise(0))
      .withColumn("salary_k", col("salary") / 1000.0)
      .withColumn("high_earner", when(col("salary") >= 90000, 1).otherwise(0))
      .withColumn("experience_band",
                  when(col("age") < 30, "junior")
                  .when((col("age") >= 30) & (col("age") < 40), "mid")
                  .otherwise("senior"))
      .withColumn("country_upper", upper(col("country")))
)

df_chain.count()
time_chain = round(time.time() - start, 2)

Batched version (single select()):

start = time.time()
df_batch = df.select(
    "id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country",
    concat_ws(" ", col("firstname"), col("lastname")).alias("full_name"),
    when(col("age") >= 35, 1).otherwise(0).alias("is_senior"),
    (col("salary") / 1000.0).alias("salary_k"),
    when(col("salary") >= 90000, 1).otherwise(0).alias("high_earner"),
    when(col("age") < 30, "junior")
        .when((col("age") >= 30) & (col("age") < 40), "mid")
        .otherwise("senior").alias("experience_band"),
    upper(col("country")).alias("country_upper")
)

df_batch.count()
time_batch = round(time.time() - start, 2)

The distinction is most evident when there are more derived columns, more complex expressions (UDFs, window functions), or when executing on managed runtimes such as AWS Glue.

In the chained cases, there are more Project nodes, code generation is fragmented, and expression evaluation is less amenable to global optimization.

In the batched cases, Spark generates a single Project node, more work is consolidated into a single WholeStageCodegen pipeline, code generation is reduced, the JVM is less stressed, and the plan is flatter and more amenable to optimization. This is not only cleaner, but it’s also faster, more reliable, and friendlier to Spark’s optimizer.

Scenario 4: Early Filter vs Late Filter

Many pipelines apply transformations first, adding columns, joining datasets, or calculating derived metrics, before filtering records. That order looks harmless in code but can double or triple the workload at execution.

Problem: Late Filtering

df_late = (
    df.withColumn("bonus", col("salary") * 0.1)
      .withColumn("salary_k", col("salary") / 1000)
      .filter(col("age") > 35)
)

This means Spark first computes all columns for every employee, then discards most rows.

Simplified Logical Plan:

Filter (age > 35)

└─ Project [id, firstname, lastname, department, salary, age, hire_date, country,

            (salary * 0.1) AS bonus,

            (salary / 1000) AS salary_k]

   └─ LogicalRDD [...]

Catalyst can sometimes reorder this automatically, but when it can't (due to UDFs or complex logic), you're doing unnecessary work on data that's thrown away.

Better Approach: Early Filtering

df_early = (
    df.filter(col("age") > 35)
      .withColumn("bonus", col("salary") * 0.1)
      .withColumn("salary_k", col("salary") / 1000)
)

Simplified Logical Plan:

Project [id, firstname, lastname, department, salary, age, hire_date, country,

         (salary * 0.1) AS bonus,

         (salary / 1000) AS salary_k]

└─ Filter (age > 35)

   └─ LogicalRDD [...]

Now Spark prunes the dataset first, then applies transformations. The result: smaller intermediate data, less codegen, shorter logical plan, shorter DAG, and smaller shuffle footprint.

Real-World Benchmark: AWS Glue

Late Filtering:

df = spark.createDataFrame(
    multiplied_data,
    ["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"]
)

start_late = time.time()

df_late = (
    df.withColumn("bonus", col("salary") * 0.1)
      .withColumn("salary_k", col("salary") / 1000)
      .filter(col("age") > 35)   
)

df_late.count()
time_late = round(time.time() - start_late, 2)

Early Filtering:

start_early = time.time()

df_early = (
    df.filter(col("age") > 35)    
      .withColumn("bonus", col("salary") * 0.1)
      .withColumn("salary_k", col("salary") / 1000)
)

df_early.count()
time_early = round(time.time() - start_early, 2)

print("Late Filter Time:", time_late, "seconds")
print("Early Filter Time:", time_early, "seconds")

spark.stop()

The early filter approach processes significantly less data before the projection, leading to faster execution and less memory pressure.

Always filter as early as possible, before joins, aggregations, expensive transformations (such as UDFs or window functions), and even during file reads via Parquet/ORC pushdown, since filtering at the source touches fewer partitions and leads to faster jobs.

Scenario 5: Column Pruning

When working with Spark DataFrames, convenience often wins over correctness and nothing feels more convenient than select("*"). It’s quick, flexible, and perfect for exploration.

But in production pipelines, that little star silently costs CPU, memory, network bandwidth, and runtime efficiency. Every time you write select("*"), Spark expands it into every column from your schema, even if you’re using just one or two later.

Those extra attributes flow through every stage of the plan, from filters and joins to aggregations and shuffles. The result: inflated logical plans, bigger shuffle files, and slower queries.

The Problem: “The Lazy Star”

df_star = (
    df.select("*")
      .filter(col("department") == "Engineering")
      .groupBy("country")
      .agg(avg("salary").alias("avg_salary"))
)

At first glance, this seems harmless. But the problem is: only two columns (country and salary) are needed for the aggregation, but Spark carries all eight (id, firstname, lastname, department, salary, age, hire_date, country) through every transformation.

Simplified Logical Plan:

Aggregate [country], [avg(salary) AS avg_salary]

└─ Filter (department = Engineering)

   └─ Project [id, firstname, lastname, department, salary, age, hire_date, country]

      └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

Every node in this tree carries all columns. Catalyst can’t prune them because you explicitly asked for "*". The excess attributes are serialized, shuffled, and deserialized across the cluster, even though they serve no purpose in the final result.

The Fix: Select Only What You Need

Be deliberate with your projections. Select the minimal schema required for the task.

df_pruned = (
    df.select("department", "salary", "country")
      .filter(col("department") == "Engineering")
      .groupBy("country")
      .agg(avg("salary").alias("avg_salary"))
)

Simplified Logical Plan:

Aggregate [country], [avg(salary) AS avg_salary]

└─ Filter (department = Engineering)

   └─ Project [department, salary, country]

      └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

Now Spark reads and processes only the three required columns: department, salary, and country. The plan is narrower, the DAG simpler, and execution faster.

Real-World Benchmark: AWS Glue

Wide Projection:

df = spark.createDataFrame(multiplied_data,
                           ["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"])

start = time.time()
df_star = (
    df.select("*")
      .filter(col("department") == "Engineering")
      .groupBy("country")
      .agg(avg("salary").alias("avg_salary"))
)

df_star.count()
time_star = round(time.time() - start, 2)

Pruned Projection:

start = time.time()

df_pruned = (
    df.select("department", "salary", "country")
      .filter(col("department") == "Engineering")
      .groupBy("country")
      .agg(avg("salary").alias("avg_salary"))
)

df_pruned.count()
time_pruned = round(time.time() - start, 2)

print(f"select('*') time: {time_star}s")
print(f"pruned columns time: {time_pruned}s")

spark.stop()

Under the Hood: How Catalyst Handles Columns

When you call select("*"), Catalyst resolves every attribute into the logical plan. Each subsequent transformation inherits that full attribute list, increasing plan depth and overhead.

Catalyst includes a rule called ColumnPruning, which removes unused attributes but it only works when Spark can see which columns are necessary. If you use "*" or dynamically reference df.columns, Catalyst loses visibility.

Works:

df \
    .select("salary", "country") \
    .groupBy("country") \
    .agg(avg("salary"))

Doesn’t Work:

cols = df.columns

df.select(cols) \
  .groupBy("country") \
  .agg(avg("salary"))

In the second case, Catalyst can’t prune anything because cols might include everything.

Physical Plan Differences

Wide Projection (select("*")):

*(1) HashAggregate(keys=[country], functions=[avg(salary)])

+- *(1) Project [id, firstname, lastname, department, salary, age, hire_date, country]

   +- *(1) Filter (department = Engineering)

      +- *(1) Scan parquet ...

Pruned Projection:

*(1) HashAggregate(keys=[country], functions=[avg(salary)])

+- *(1) Project [department, salary, country]

   +- *(1) Filter (department = Engineering)

      +- *(1) Scan parquet [department, salary, country]

Notice the last line: Spark physically scans only the three referenced columns from Parquet. That’s genuine I/O reduction, not just logical simplification. Using select(*) increases shuffle file sizes, memory usage during serialization, Catalyst planning time, and I/O and network traffic, and the solution requires no more than specifying the necessary columns.

But in managed environments like AWS Glue or Databricks, this simple practice can greatly reduce ETL time, particularly for Parquet or Delta files, due to effective column pruning during explicit projection. It’s one of the easiest and highest-impact Spark optimization techniques, starting with typing fewer asterisks.

Scenario 6: Filter Pushdown vs Full Scan

When a Spark job feels slow right from the start, even before joins or aggregations, the culprit is often hidden at the data-read layer. Spark spends seconds (or minutes) scanning every record, even though most rows are useless for the query.

That’s where filter pushdown comes in. It tells Spark to push your filter logic down to the file reader so that Parquet / ORC / Delta formats return only the relevant rows from disk. Done right, this optimization can reduce scan size significantly. Done wrong, Spark performs a full scan, reading everything before filtering in memory.

The Problem: Late Filters and Full Scans

employees_df = spark.read.parquet("s3://data/employee_data/")

df_full = (
    employees_df
        .select("*")  # reads all columns
        .filter(col("country") == "Canada")
)

Looks fine, right? But Spark can’t push this filter to the Parquet reader because it’s applied after the select("*") projection step. Catalyst sees the filter as operating on a projected DataFrame, not the raw scan, so the pushdown boundary is lost.

Simplified Logical Plan:

Filter (country = Canada)

└─ Project [id, firstname, lastname, department, salary, age, hire_date, country]

   └─ Scan parquet employee_data [id, firstname, lastname, department, salary, age, hire_date, country]

Every record from every Parquet file is read into memory before the filter executes. In large tables, this means scanning terabytes when you only need megabytes.

The Fix: Filter Early and Project Light

Move filters as close as possible to the data source and limit columns before Spark reads them:

df_pushdown = (
    spark.read.parquet("s3://data/employee_data/")
        .select("id", "firstname", "department", "salary", "country")
        .filter(col("country") == "Canada")
)

Simplified Logical Plan:

Project [id, firstname, department, salary, country]

└─ Scan parquet employee_data [id, firstname, department, salary, country]

PushedFilters: [country = Canada]

Notice the difference: PushedFilters appears in the plan. That means the Parquet reader handles the predicate, returning only matching blocks and rows.

Under the Hood: What Actually Happens

When Spark performs filter pushdown, it leverages the Parquet metadata (min/max statistics and row-group indexes) stored in file footers.

• Spark inspects file-level metadata for the predicate column (country).

• It skips any row group whose values don’t match (country ≠ Canada).

• It reads only the necessary row groups and columns from disk.

• Those records enter the DAG directly – no in-memory filtering required.

This optimization happens entirely before Spark begins executing stages, reducing both I/O and network transfer.

Real-World Benchmark: AWS Glue

import time
from pyspark.sql import SparkSession
from pyspark.sql.functions import col

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

start = time.time()
df_full = (
    spark.read.parquet("s3://data/employee_data/")
        .select("*")                         # all columns
        .filter(col("country") == "Canada")  
)
df_full.count()
time_full = round(time.time() - start, 2)

start = time.time()
df_pushdown = (
    spark.read.parquet("s3://data/employee_data/")
        .select("id", "firstname", "department", "salary", "country")
        .filter(col("country") == "Canada")  
)
df_pushdown.count()
time_push = round(time.time() - start, 2)

print("Full Scan Time:", time_full, "sec")
print("Filter Pushdown Time:", time_push, "sec")

spark.stop()

Physical Plan Comparison

Full Scan:

*(1) Filter (country = Canada)

+- *(1) ColumnarToRow

   +- *(1) FileScan parquet [id, firstname, lastname, department, salary, age, hire_date, country]

      Batched: true, DataFilters: [], PushedFilters: []

Pushdown:

*(1) ColumnarToRow

+- *(1) FileScan parquet [id, firstname, department, salary, country]

   Batched: true, DataFilters: [isnotnull(country)], PushedFilters: [country = Canada]

The difference is clear: PushedFilters confirms that Spark applied predicate pushdown, skipping unnecessary row groups at the scan stage.

Reflection: Why Pushdown Matters

Pushdown isn’t a micro-optimization. It’s actually often the single biggest performance lever in Spark ETL. In data lakes with hundreds of files, full scans waste hours and inflate AWS S3 I/O costs. By filtering and projecting early, Spark prunes both rows and columns before execution even begins.

Apply filters as early as possible in the read pipeline, combine filter pushdown with column pruning, verify PushedFilters in explain("formatted"), avoid UDFs and select("*") at read time, and let pushdown turn “read everything and discard most” into “read only what you need.”

Scenario 7: De-duplicate Right

The Problem: “All-Row Deduplication” and Why It Hurts

When we use this:

df.dropDuplicates()

Spark removes identical rows across all columns. It sounds simple, but this operation forces Spark to treat every column as part of the deduplication key.

Internally, it means:

• Every attribute is serialized and hashed.

• Every unique combination of all columns is shuffled across the cluster to ensure global uniqueness.

• Even small changes in a non-essential field (like hire_date) cause new keys and destroy aggregation locality.

In wide tables, this is one of the heaviest shuffle operations Spark can perform: df.dropDuplicates()

Simplified Logical Plan:

Aggregate [id, firstname, lastname, department, salary, age, hire_date, country], [first(id) AS id, ...]

└─ Exchange hashpartitioning(id, firstname, lastname, department, salary, age, hire_date, country, 200)

   └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

Notice the Exchange: that’s a full shuffle across all columns. Spark must send every record to the partition responsible for its unique combination of all fields. This is slow, memory-intensive, and scales poorly as columns grow.

The Better Approach: Key-Based Deduplication

In most real datasets, duplicates are determined by a primary or business key, not all attributes. For example, if id uniquely identifies an employee, we only need to keep one record per id.

df.dropDuplicates(["id"])

Now Spark deduplicates based only on the id column.

Aggregate [id], [first(id) AS id, first(firstname) AS firstname, ...]

└─ Exchange hashpartitioning(id, 200)

   └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

The shuffle is dramatically narrower. Instead of hashing across all columns, Spark redistributes data only by id. Fewer bytes, smaller shuffle files, faster reduce stage

Real-World Benchmark: AWS Glue

import time
from pyspark.sql.functions import exp, log, sqrt, col, concat_ws, when, upper, avg
from pyspark.sql import SparkSession

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

employees_data = [
    (1, "John", "Doe", "Engineering", 80000, 28, "2020-01-15", "USA"),
    (2, "Jane", "Smith", "Engineering", 85000, 32, "2019-03-20", "USA"),
    (3, "Alice", "Johnson", "Sales", 60000, 25, "2021-06-10", "UK"),
    (4, "Bob", "Brown", "Engineering", 90000, 35, "2018-07-01", "USA"),
    (5, "Charlie", "Wilson", "Sales", 65000, 29, "2020-11-05", "UK"),
    (6, "David", "Lee", "HR", 55000, 27, "2021-01-20", "USA"),
    (7, "Eve", "Davis", "Engineering", 95000, 40, "2017-04-12", "Canada"),
    (8, "Frank", "Miller", "Sales", 70000, 33, "2019-09-25", "UK"),
    (9, "Grace", "Taylor", "HR", 58000, 26, "2021-08-15", "Canada"),
    (10, "Henry", "Anderson", "Engineering", 88000, 31, "2020-02-28", "USA")
]

multiplied_data = [(i, f"firstname_{i}", f"lastname_{i}",
                    employees_data[i % 10][3],   # department
                    employees_data[i % 10][4],   # salary
                    employees_data[i % 10][5],   # age
                    employees_data[i % 10][6],   # hire_date
                    employees_data[i % 10][7]    # country
                    )
                   for i in range(1, 1_000_001)]

df = spark.createDataFrame(
    multiplied_data,
    ["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"]
)

start = time.time()
dedup_full = df.dropDuplicates()
dedup_full.count()
time_full = round(time.time() - start, 2)

start = time.time()
dedup_key = df.dropDuplicates(["id"])
dedup_key.count()
time_key = round(time.time() - start, 2)

print(f"Full-row dedup time: {time_full}s")
print(f"Key-based dedup time: {time_key}s")

spark.stop()

Under the Hood: What Catalyst Does

When you specify a key list, Catalyst rewrites dropDuplicates(keys) into a partial + final aggregate plan, just like a groupBy:

HashAggregate(keys=[id], functions=[first(...)])

This allows Spark to:

• Perform map-side partial aggregation on each partition (before shuffle).

• Exchange only the grouping key (id).

• Perform a final aggregation on the reduced data.

The all-column version can’t do that optimization because every column participates in uniqueness Spark must ensure complete data redistribution.

Best Practices for Deduplication

Combining Deduplication & Aggregation

You can merge deduplication with aggregation for even better results:

df_dedup_agg = (
    df.dropDuplicates(["id"])
        .groupBy("department")
        .agg(avg("salary").alias("avg_salary"))
)

Spark now reuses the same shuffle partitioning for both operations, one shuffle instead of two. The plan will show:

HashAggregate(keys=[department], functions=[avg(salary)])

└─ HashAggregate(keys=[id], functions=[first(...), first(department)])

   └─ Exchange hashpartitioning(id, 200)

Prefer dropDuplicates(["key_col"]) over dropDuplicates() to deduplicate by business or surrogate keys rather than the entire schema. Combine deduplication with projection to reduce I/O, and remember that one narrow shuffle is always better than a wide shuffle. Deduplication isn’t just cleanup – it’s an optimization strategy. Choose your keys wisely, and Spark will reward you with faster jobs and lighter DAGs.

Scenario 8: Count Smarter

In production, one of the most common performance pitfalls is the simplest line of code:

if df.count() > 0:

At first glance, this seems harmless. You just want to know whether the DataFrame has any data before writing, joining, or aggregating. But in Spark, count() is not metadata lookup, it’s a full cluster-wide job.

What Really Happens with count()
When you call df.count(), Spark executes a complete action:

• It scans every partition.

• Deserializes every row.

• Counts records locally on each executor.

• Reduces the counts to the driver.

That means your “empty check” runs a full distributed computation, even when the dataset has billions of rows or lives in S3.

df.count()

Simplified Logical Plan:

*(1) HashAggregate(keys=[], functions=[count(1)])

+- *(1) ColumnarToRow

   +- *(1) FileScan parquet [id, firstname, lastname, department, salary, age, hire_date, country]

Every record is read, aggregated, and returned just to produce a single integer.

Now imagine this runs in the middle of your Glue job, before a write, before a filter, or inside a loop. You’ve just added a full-table scan to your DAG for no reason.

The Smarter Way: limit(1) or head(1)

If all you need to know is whether data exists, you don’t need to count every record. You just need to know if there’s at least one.

Two efficient alternatives

df.head(1)
#or
df.limit(1).collect()

Both execute a lazy scan that stops as soon as one record is found.

Simplified Logical Plan:

TakeOrderedAndProject(limit=1)

└─ *(1) FileScan parquet [id, firstname, lastname, department, salary, age, hire_date, country]

• No global aggregation.

• No shuffle.

• No full scan.

Real-World Benchmark: AWS Glue

import time
from pyspark.sql.functions import exp, log, sqrt, col, concat_ws, when, upper, avg
from pyspark.sql import SparkSession

# Initialize Spark session
spark = SparkSession.builder.appName("MillionRowsRenameTest").getOrCreate()

# Base dataset (10 sample employees)
employees_data = [
    (1, "John", "Doe", "Engineering", 80000, 28, "2020-01-15", "USA"),
    (2, "Jane", "Smith", "Engineering", 85000, 32, "2019-03-20", "USA"),
    (3, "Alice", "Johnson", "Sales", 60000, 25, "2021-06-10", "UK"),
    (4, "Bob", "Brown", "Engineering", 90000, 35, "2018-07-01", "USA"),
    (5, "Charlie", "Wilson", "Sales", 65000, 29, "2020-11-05", "UK"),
    (6, "David", "Lee", "HR", 55000, 27, "2021-01-20", "USA"),
    (7, "Eve", "Davis", "Engineering", 95000, 40, "2017-04-12", "Canada"),
    (8, "Frank", "Miller", "Sales", 70000, 33, "2019-09-25", "UK"),
    (9, "Grace", "Taylor", "HR", 58000, 26, "2021-08-15", "Canada"),
    (10, "Henry", "Anderson", "Engineering", 88000, 31, "2020-02-28", "USA")
]

# Create 1 million rows
multiplied_data = [
    (i, f"firstname_{i}", f"lastname_{i}",
     employees_data[i % 10][3],
     employees_data[i % 10][4],
     employees_data[i % 10][5],
     employees_data[i % 10][6],
     employees_data[i % 10][7])
    for i in range(1, 1_000_001)
]

df = spark.createDataFrame(
    multiplied_data,
    ["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"]
)
# Create DataFrame
df = spark.createDataFrame(
    multiplied_data,
    ["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"]
)

start = time.time()
df.count()
count_time = round(time.time() - start, 2)

start = time.time()
df.limit(1).collect()
limit_time = round(time.time() - start, 2)

start = time.time()
df.head(1)
head_time = round(time.time() - start, 2)

spark.stop()

The difference is significant for the same logical check.

So why does this difference exist? Spark’s execution model treats every action as a trigger for computation. count() is an aggregation action, requiring global communication, and limit(1) and head(1) are sampling actions, short-circuiting the job after fetching the first record. Catalyst generates a TakeOrderedAndProject node instead of HashAggregate, and the scheduler terminates once one task finishes.

Plan comparison:

Avoid using count() to check emptiness since it triggers a full scan. Use limit(1) or head(1) for lightweight existence checks. And reserve count() only when the total is required, because Spark will always process all data unless explicitly told to stop. Other alternatives

Scenario 9: Window Wisely

Window functions (rank(), dense_rank(), lag(), avg() with over(), and so on) are essential in analytics. They let you calculate running totals, rankings, or time-based metrics.

But in Spark, they’re not cheap, because they rely on shuffles and ordering.

Each window operation:

• Requires all rows for the same partition key to be co-located on the same node.

• Requires sorting those rows by the orderBy() clause within each partition.

If you omit partitionBy() (or use it with too broad a key), Spark treats the entire dataset as one partition, triggering a massive shuffle and global sort.

Global Window: The Wrong Way

Let’s compute employee rankings by salary without partitioning:

from pyspark.sql.window import Window
from pyspark.sql.functions import rank, col

window_spec = Window.orderBy(col("salary").desc())

df_ranked = df.withColumn("salary_rank", rank().over(window_spec))

Simplified Logical Plan:

Window [rank() windowspecdefinition(orderBy=[salary DESC]) AS salary_rank]

└─ Sort [salary DESC], true

   └─ Exchange rangepartitioning(salary DESC, 200)

      └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

Spark must shuffle and sort the entire dataset globally, a full sort across all rows. Every executor gets a slice of this single global range, and all data must move through the network.

Partition by a Selective Key: The Better Way

Most analytics don’t need a global ranking. You likely want rankings within a department or group, not across the entire company.

window_spec = Window.partitionBy("department").orderBy(col("salary").desc())

df_ranked = df.withColumn("salary_rank", rank().over(window_spec))

Now Spark builds separate windows per department. Each partition’s data stays local, dramatically reducing shuffle size.

Simplified Logical Plan:

Window [rank() windowspecdefinition(partitionBy=[department], orderBy=[salary DESC]) AS salary_rank]

└─ Sort [department ASC, salary DESC], false

   └─ Exchange hashpartitioning(department, 200)

      └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

The Exchange now partitions data only by department. The shuffle boundary is narrower, fewer bytes transferred, fewer sort comparisons, and smaller spill risk.

Real-World Benchmark: AWS Glue

We can execute the windows function on the same 1 million row dataset:

df = spark.createDataFrame(multiplied_data,
["id", "firstname", "lastname", "department", "salary", "age",
 "hire_date", "country"])

start = time.time()
window_global = Window.orderBy(col("salary").desc())
df_global = df.withColumn("salary_rank", rank().over(window_global))
df_global.count()
global_time = round(time.time() - start, 2)
print(f'global_time:{global_time}')

start = time.time()
window_local = Window.partitionBy("department").orderBy(col("salary").desc())
df_local = df.withColumn("salary_rank", rank().over(window_local))
df_local.count()
local_time = round(time.time() - start, 2)
print(f'local_time:{local_time}')

spark.stop()

Partitioning the window reduces shuffle data volume significantly and runtime as well. The difference grows exponentially as data scales.

Under the Hood: What Spark Actually Does

Each Window transformation adds a physical plan node like:

WindowExec [rank() windowspecdefinition(...)], frame=RangeFrame

This node is non-pipelined – it materializes input partitions before computing window metrics. Catalyst optimizer can’t push filters or projections inside WindowExec, which means:

• If you rank before filtering, Spark computes ranks for all rows.

• If you order globally, Spark must sort everything before starting.

That’s why window placement in your code matters almost as much as partition keys.

Common Anti-Patterns:

Partition on selective keys to reduce shuffle volume, and avoid global windows that force full sorts and shuffles. Prefer bounded frames to keep state in memory and limit disk spill, and filter early while combining metrics to minimize unnecessary data flowing through WindowExec. Windows are powerful, but unbounded ones can silently crush performance. In Spark, partitioning isn’t optional. It’s the line between analytics and overhead.

Scenario 10: Incremental Aggregations with Cache and Persist

When multiple actions depend on the same expensive base computation, don’t recompute it every time. Materialize it once with cache() or persist(), then reuse it. Most Spark teams get this wrong in two ways:

• They never cache, so Spark recomputes long lineages (filters, joins, window ops) for every action.

• They cache everything, blowing executor memory and making things worse.

This scenario shows how to do it intelligently.

The Problem: Recomputing the Same Work for Every Metric

from pyspark.sql.functions import col, avg, max as max_, count

base = (
    df.filter(col("department") == "Engineering")
      .filter(col("country") == "USA")
      .filter(col("salary") > 70000)
)

avg_salary = base.groupBy("department").agg(avg("salary").alias("avg_salary"))
max_salary = base.groupBy("department").agg(max_("salary").alias("max_salary"))
cnt_salary = base.groupBy("department").agg(count("*").alias("cnt"))

Looks totally fine at a glance. But remember: Spark is lazy.
Every time you trigger an action:

avg_salary.show()
max_salary.show()
cnt_salary.show()

Spark walks back to the same base definition and re-runs all filters and shuffles for each metric – unless you persist.

So instead of 1 filtered + shuffled dataset reused 3 times, you effectively get:

• 3 jobs

• 3 scans / filter chains

• 3 groupBy shuffles

for the same input slice.

Simplified Logical Plan Shape (Without Cache):

HashAggregate [department], [avg/max/count]

└─ Exchange hashpartitioning(department)

   └─ Filter (department = 'Engineering' AND country = 'USA' AND salary > 70000)

      └─ Scan ...

And Spark builds this three times. Even though the filter logic is identical, each action triggers a new job with:

• new stages,

• new shuffles, and

• new scans.

On large datasets (hundreds of GBs), this is brutal.

The Better Approach: Cache the Shared Base

from pyspark.sql import StorageLevel

base = (
    df.filter(col("department") == "Engineering")
      .filter(col("country") == "USA")
      .filter(col("salary") > 70000)
)

base = base.persist(StorageLevel.MEMORY_AND_DISK)

base.count()

avg_salary = base.groupBy("department").agg(avg("salary").alias("avg_salary"))
max_salary = base.groupBy("department").agg(max_("salary").alias("max_salary"))
cnt_salary = base.groupBy("department").agg(count("*").alias("cnt"))

avg_salary.show()
max_salary.show()
cnt_salary.show()

base.unpersist()

Now, the filters and initial scan run once, the results are cached, and all subsequent aggregates read from cached data instead of recomputing upstream logic.

Logical Plan Shape (With Cache):

Before materialization (base.count()), the plan still shows the lineage. Afterward, subsequent actions operate off the cached node.

InMemoryRelation [department, salary, country, ...]

   └─ * Cached from:

      Filter (department = 'Engineering' AND country = 'USA' AND salary > 70000)

      └─ Scan parquet employees_large ...

Then:

HashAggregate [department], [avg/max/count]

└─ InMemoryRelation [...]

One heavy pipeline, many cheap reads. The DAG becomes flatter:

• Expensive scan & filter & shuffle: once.

• Cheap aggregations: N times from memory/disk.

Real-World Benchmark: AWS Glue

df = spark.createDataFrame(multiplied_data,
["id", "firstname", "lastname", "department", "salary", "age",
"hire_date", "country"])

base = (
    df.filter(col("department") == "Engineering")
      .filter(col("country") == "USA")
      .filter(col("salary") > 85000)
)


start = time.time()

avg_salary = base.groupBy("department").agg(avg("salary").alias("avg_salary"))
max_salary = base.groupBy("department").agg(max_("salary").alias("max_salary"))
cnt = base.groupBy("department").agg(count("*").alias("emp_count"))

print("---- Without Cache ----")
avg_salary.show()
max_salary.show()
cnt.show()

no_cache_time = round(time.time() - start, 2)
print(f"Total time without cache: {no_cache_time} seconds")


from pyspark.sql import DataFrame

base_cached = base.persist(StorageLevel.MEMORY_AND_DISK)
base_cached.count()  # materialize cache

start = time.time()

avg_salary_c = base_cached.groupBy("department").agg(avg("salary").alias("avg_salary"))
max_salary_c = base_cached.groupBy("department").agg(max_("salary").alias("max_salary"))
cnt_c = base_cached.groupBy("department").agg(count("*").alias("emp_count"))

print("---- With Cache ----")
avg_salary_c.show()
max_salary_c.show()
cnt_c.show()

cache_time = round(time.time() - start, 2)
print(f"Total time with cache: {cache_time} seconds")

# Cleanup
base_cached.unpersist()

print("\n==== Summary ====")
print(f"Without cache: {no_cache_time}s | With cache: {cache_time}s")
print("=================")

spark.stop()

Under the Hood: Why This Works

Using cache() or persist() in Spark inserts an InMemoryRelation / InMemoryTableScanExec node so that expensive intermediate results are stored in executor memory (or memory+disk). This allows future jobs to reuse cached blocks instead of re-scanning sources or re-computing shuffles. This shortens downstream logical plans, reduces repeated shuffles, and lowers load on systems like S3, HDFS, or JDBC.

Without caching, every action replays the full lineage and Spark recomputes the data unless another operator or AQE optimization has already materialized part of it. But caching should not become “cache everything”. Rather, you should avoid caching very large DataFrames used only once, wide raw inputs instead of filtered/aggregated subsets, or long-lived caches that are never unpersisted.

A good rule of thumb is to cache only when the DataFrame is expensive to recompute (joins, filters, windows, UDFs), is used at least twice, and is reasonably sized after filtering so it can fit in memory or work with MEMORY_AND_DISK. Otherwise, allow Spark to recompute.

Conceptually, caching converts a tall, repetitive DAG such as repeated “HashAggregate → Exchange → Filter → Scan” sequences into a hub-and-spoke design where one heavy cached hub feeds multiple lightweight downstream aggregates.

When multiple actions depend on the same expensive computation, cache or persist the shared base to flatten the DAG, eliminate repeated scans and shuffles, and improve end-to-end performance. All this while being intentional by caching only when reuse is real, the data size is safe, and always calling unpersist() when done.

Don’t make Spark re-solve the same puzzle three times. Let it solve it once, remember the answer, and move on.

Scenario 11: Reduce Shuffles

Shuffles are Spark’s invisible tax collectors. Every time your data crosses executors, you pay in CPU, disk I/O, and network bandwidth.

Two of the most common yet misunderstood transformations that trigger or avoid shuffles are coalesce() and repartition(). Both change partition counts, but they do it in fundamentally different ways.

The Problem

Writing df_result = df.repartition(10) and thinking “I’m just changing partitions so Spark won’t move data unnecessarily.” But that assumption is wrong. repartition() always performs a full shuffle, even when:

• You are reducing partitions (from 200 → 10), or

• You are increasing partitions (from 10 → 200).

In both cases, Spark redistributes every row across the cluster according to a new hash partitioning scheme. So even if your data is already partitioned optimally, repartition() will still reshuffle it, adding a stage boundary.

Logical Plan:

Exchange hashpartitioning(...)

└─ LogicalRDD [...]

That Exchange node signals a wide dependency: Spark spills intermediate data to disk, transfers it over the network, and reloads it before the next stage. In short: repartition() = "new shuffle, no matter what."

The Better Approach: coalesce()

If your goal is to reduce the number of partitions, for example, before writing results to S3 or Snowflake – use coalesce() instead.

df_result = df.coalesce(10)

coalesce() merges existing partitions locally within each executor, avoiding the costly reshuffle step. It uses a narrow dependency, meaning each output partition depends on one or more existing partitions from the same node.

Coalesce

└─ LogicalRDD [...]

• No Exchange.

• No network shuffle.

• Just local merges – fast and cheap.

Real-World Benchmark: AWS Glue

df = spark.createDataFrame(multiplied_data,
["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"])

start = time.time()
df_repart = df.repartition(10)
df_repart.count()
print("Repartition time:", round(time.time() - start, 2), "sec")

start = time.time()
df_coalesced = df.coalesce(10)
df_coalesced.count()
print("Coalesce time:", round(time.time() - start, 2), "sec")

spark.stop()

Even though both ended with 10 partitions, repartition() took significantly longer all because of the unnecessary shuffle.

Why This Matters

Each Exchange node in your logical plan creates a new stage in your DAG, meaning:

• Extra disk I/O

• Extra serialization

• Extra network transfer

That’s why avoiding just one shuffle in a Glue ETL pipeline can save seconds to minutes per run, especially on wide datasets.

When to use which:

AQE and Auto Coalescing

You can enable Adaptive Query Execution (AQE) in AWS Glue 3.0+ to let Spark merge small shuffle partitions automatically:

spark.conf.set("spark.sql.adaptive.enabled", "true")

spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")

With AQE, Spark dynamically combines small partitions after shuffle to balance performance and I/O.

repartition() always triggers a shuffle, while coalesce() avoids shuffles and is ideal for local merges before writes. You should always inspect Exchange nodes to identify shuffle points. Note that in AWS Glue, avoiding even one shuffle can yield ~7× runtime improvement at the 1M-row scale. Finally, use AQE to enable dynamic partition coalescing in larger workflows.

Scenario 12: Know Your Shuffle Triggers

Much of Spark's performance comes from invisible data movement. Every shuffle boundary adds a new stage, a new write–read cycle, and sometimes minutes of extra execution time.

In Spark, any operation that requires rearranging data between partitions introduces a wide dependency, represented in the logical plan as an Exchange node.

Common shuffle triggers:

Each Exchange means a shuffle stage: Spark writes partition data to disk, transfers it over the network, and reads it back into memory on the next stage. That’s your hidden performance cliff.

df_result = (
    df.groupBy("department")
      .agg(sum("salary").alias("total_salary"))
      .join(df.select("department", "country")
            .distinct(), "department")
      .orderBy("total_salary", ascending=False)
)

df_result.explain("formatted")

Logical Plan Simplified:

Sort [total_salary DESC]

└─ Exchange (global sort)

   └─ SortMergeJoin [department]

      ├─ Exchange (groupBy shuffle)

      │   └─ HashAggregate (sum salary)

      └─ Exchange (distinct shuffle)

          └─ Aggregate (department, country)

We can see three Exchange nodes, one for the aggregation, one for the distinct join, and one for the global sort. That’s three separate shuffles, three full dataset transfers.

Better Approach

Whenever possible, combine wide transformations into a single stage before an action. For instance, you can compute aggregates and join results in one consistent shuffle domain:

agg_df = df.groupBy("department") \
    .agg(sum("salary") \
    .alias("total_salary"))

country_df = df.select("department", "country").distinct()

df_result = (
    agg_df.join(country_df, "department")
          .sortWithinPartitions("total_salary", ascending=False)
)

Logical Plan Simplified:

SortWithinPartitions [total_salary DESC]

└─ SortMergeJoin [department]

   ├─ Exchange (shared shuffle for join)

   └─ Exchange (shared shuffle for distinct)

Now Spark reuses shuffle partitions across compatible operations – only one shuffle boundary remains. The rest execute as narrow transformations.

Real-World Benchmark: AWS Glue (1M)

df = spark.createDataFrame(multiplied_data,
["id", "firstname", "lastname", "department", "salary", "age", "hire_date", "country"]).repartition(20)

from pyspark.sql.functions import sum as sum_

start = time.time()

dept_salary = (
    df.groupBy("department")
      .agg(sum_("salary").alias("total_salary"))
)

dept_country = (
    df.select("department", "country")
      .distinct()
)

naive_result = (
    dept_salary.join(dept_country, "department", "inner")
               .orderBy(col("total_salary").desc())
)

naive_count = naive_result.count()
naive_time = round(time.time() - start, 2)


start = time.time()

dept_country_once = (
    df.select("department", "country")
      .distinct()
)

optimized = (
    df.groupBy("department")
      .agg(sum_("salary").alias("total_salary"))
      .join(dept_country_once, "department", "inner")
      .sortWithinPartitions(col("total_salary").desc())
      # local ordering, avoids extra global shuffle
)

opt_count = optimized.count()
opt_time = round(time.time() - start, 2)

print("Optimized result count:", opt_count)
print("Optimized pipeline time:", opt_time, "sec")

print("\nOptimized plan:")
optimized.explain("formatted")

spark.stop()

By merging compatible stages and using sortWithinPartitions() instead of global orderBy(), the job ran significantly faster on the same dataset, with fewer Exchange nodes and shorter lineage. Run df.explain and search for Exchange. Each one signals a full shuffle. You can also check Spark UI → SQL tab → Exchange Read/Write Size to see exactly how much data moved.

Every Exchange represents a shuffle, adding serialization, network I/O, and stage overhead, so avoid chaining wide operations back-to-back by combining them under a consistent partition key. Prefer sortWithinPartitions() over global orderBy() when ordering is local, monitor plan depth to catch consecutive wide dependencies, and note that in AWS Glue eliminating even one shuffle in a 1M-row job can significantly reduce runtime.

Scenario 13: Tune Parallelism: Shuffle Partitions & AQE

Most Spark jobs are either over-parallelized (thousands of tiny tasks doing almost nothing, flooding the driver and filesystem) or under-parallelized (a handful of huge tasks doing all the work, causing slow stages and skew-like behavior). Both waste resources. We can control this behavior using spark.sql.shuffle.partitions and Adaptive Query Execution (AQE).

By default (in many environments), the default value spark.conf.get("spark.sql.shuffle.partitions") is 200, meaning that every shuffle produces approximately 200 shuffle partitions, regardless of data size. That means every shuffle (groupBy, join, distinct, and so on) creates ~200 shuffle partitions. Whether this default is reasonable depends entirely on the workload:

• If you’re processing 2 GB, 200 partitions might be great.

• If you’re processing 5 MB, 200 partitions is comedy – 200 tiny tasks, overhead > work.

• If you’re processing 2 TB, 200 partitions might be too few – tasks become huge and slow.

Example A: The Default Plan (Too Many Tiny Tasks)

from pyspark.sql import SparkSession
from pyspark.sql.functions import sum as sum_

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

spark.conf.get("spark.sql.shuffle.partitions")  # '200'

data = [
    (1, "John", "Engineering", 90000),
    (2, "Alice", "Engineering", 85000),
    (3, "Bob", "Sales", 75000),
    (4, "Eve", "Sales", 72000),
    (5, "Grace", "HR", 65000),
]

df = spark.createDataFrame(data, ["id", "name", "department", "salary"])

agg_df = df.groupBy("department").agg(sum_("salary").alias("total_salary"))
agg_df.explain("formatted")

Even though there are only 3 departments, Spark will still create 200 shuffle partitions – meaning 200 tasks for 3 groups of data.

Effect: Each task has almost nothing to do. Spark spends more time planning and scheduling than actually computing.

Example B: Tuned Plan (Balanced Parallelism)

spark.conf.set("spark.sql.shuffle.partitions", "8")
agg_df = df.groupBy("department").agg(sum_("salary").alias("total_salary"))
agg_df.explain("formatted")

Now Spark launches only 8 partitions still parallelized, but not wasteful. Even in this small example, you can visually feel the difference: one logical change, but a completely leaner physical plan.

The Real Problem: Static Tuning Doesn’t Scale

In production, job sizes vary:

• Today: 10 GB

• Tomorrow: 500 GB

• Next week: 200 MB (sampling run)

Manually changing shuffle partitions for each run is neither practical nor reliable. That’s where Adaptive Query Execution (AQE) steps in.

Adaptive Query Execution (AQE): Smarter, Dynamic Parallelism

AQE doesn’t guess. It measures actual shuffle statistics at runtime and rewrites the plan while the job is running.

spark.conf.set("spark.sql.adaptive.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.enabled", "true")
spark.conf.set("spark.sql.adaptive.coalescePartitions.minPartitionSize", "64m")
spark.conf.set("spark.sql.adaptive.coalescePartitions.maxPartitionSize", "256m")

AQE merges tiny shuffle partitions, or splits huge ones, based on real-time data metrics, not pre-set assumptions.

df = spark.createDataFrame(multiplied_data,
    ["id", "firstname", "lastname", "department", "salary", "age",
     "hire_date", "country"])

start = time.time()
agg_df = df.groupBy("department").agg(sum_("salary").alias("total_salary"))
agg_df.count()

print(f'Num Partitions df: {df.rdd.getNumPartitions()}')
print(f'Num Partitions aggdf: {agg_df.rdd.getNumPartitions()}')
print("Execution time:", round(time.time() - start, 2), "sec")

spark.stop()

Understanding the Plan

Before AQE (static):

Exchange hashpartitioning(department, 200)

With AQE: AdaptiveSparkPlan (coalesced)

HashAggregate(keys=[department], functions=[sum(salary)])

Exchange hashpartitioning(department, 200) # runtime coalesced to 12

The logical plan remains the same, but the physical execution plan is rewritten during runtime. Spark intelligently reduces or merges shuffle partitions based on data volume.

Spark’s default 200 shuffle partitions often misfit real workloads. Static tuning may work for predictable pipelines, but fails with variable data. On the other hand, AQE uses shuffle statistics to dynamically coalesce partitions at runtime, use it with sensible ceilings (for example, 400 partitions) and always verify in the Spark UI to catch over-partitioning (many tasks reading KBs) or under-partitioning (few tasks reading GBs).

Scenario 14: Handle Skew Smartly

In an ideal Spark world, all partitions contain roughly equal amounts of data. But real datasets are rarely that kind. If one key (say "USA", "2024", or "customer_123") holds millions of rows while others have only a few, Spark ends up with one or two massive partitions. Those partitions take disproportionately longer to process, leaving other executors idle. That’s data skew: the silent killer of parallelism.

You’ll often spot it in Spark UI:

• 198 tasks finish quickly.

• 2 tasks take 10× longer.

• Stage stays stuck at 98% for minutes.

Example A: The Skew Problem

from pyspark.sql import SparkSession, functions as F

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

# Create skewed dataset
df = spark.range(0, 10000).toDF("id") \
    .withColumn("department",
        F.when(F.col("id") < 8000, "Engineering")  # 80% of data
         .when(F.col("id") < 9000, "Sales")
         .otherwise("HR")) \
    .withColumn("salary", (F.rand() * 100000).cast("int"))

df.groupBy("department").count().show()

Spark will hash “Engineering” into just one reducer partition, making it heavier than others. That single task becomes a bottleneck, the shuffle has technically completed, but the stage waits for that one lagging task.

Example B: The Solution: Salting Hot Keys

To handle skew, we the hot key (Engineering) into multiple pseudo-keys using a random salt. This redistributes that large partition across multiple reducers.

from pyspark.sql.functions import rand, concat, lit, floor

salt_buckets = 10

df_salted = (
    df.withColumn(
        "department_salted",
        F.when(F.col("department") == "Engineering",
            F.concat(F.col("department"), lit("_"),
                     (F.floor(rand() * salt_buckets))))
         .otherwise(F.col("department"))
    )
)

df_salted.groupBy("department_salted").agg(F.avg("salary"))

Now “Engineering” isn’t one hot key – it’s 10 smaller keys like Engineering_0, Engineering_1, ..., Engineering_9. Each one goes to a separate reducer partition, enabling parallel processing.

Example C: Post-Aggregation Desalting

After aggregating, recombine salted keys to get the original department names:

df_final = (
    df_salted.groupBy("department_salted")
        .agg(F.avg("salary").alias("avg_salary"))
        .withColumn("department", F.split(F.col("department_salted"), "_")
            .getItem(0))
        .groupBy("department")
        .agg(F.avg("avg_salary").alias("final_avg_salary"))
)

When to Use Salting

Use salting when:

• You observe stage skew (one or few long tasks).

• Shuffle read sizes vary drastically between tasks.

• The skew originates from a few dominant key values.

Avoid it when:

• The dataset is small (< 1 GB).

• You already use partitioning or bucketing keys with uniform distribution.

Alternative approaches:

Glue-Specific Tip

AWS Glue 3.0+ includes Spark 3.x, meaning you can also enable AQE’s built-in skew optimization:

spark.conf.set("spark.sql.adaptive.skewJoin.enabled", "true")

spark.conf.set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "128m")

Spark will automatically detect large shuffle partitions and split them, effectively auto-salting hot keys at runtime. Data skew causes uneven shuffle sizes across tasks and can be detected in the Spark UI or via shuffle read/write metrics. Mitigate heavy-key skew with manual salting (recombined later) or rely on AQE skew join optimization for mild cases, and always validate improvements in the Spark UI SQL tab by checking “Shuffle Read Size.”

Scenario 15: Sort Efficiently (orderBy vs sortWithinPartitions)

Most Spark jobs need sorted data at some point – for window functions, for writing ordered files, or for downstream processing. The instinct is to reach for orderBy(). But those instincts cost you a full shuffle every single time.

The Problem: Global Sort When You Don't Need It

Let's say you want to write employee data partitioned by department, sorted by salary within each department:

from pyspark.sql.functions import col

# Naive approach: global sort
df_sorted = df.orderBy(col("department"), col("salary").desc())

df_sorted.write.partitionBy("department").parquet("s3://output/employees/")

This looks reasonable. You're sorting by department and salary, then writing partitioned files. Clean and simple. But here's what Spark actually does:

Simplified Logical Plan:

Sort [department ASC, salary DESC], true

└─ Exchange rangepartitioning(department ASC, salary DESC, 200)

   └─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

That Exchange rangepartitioning is a full shuffle. So Spark:

• Samples the data to determine range boundaries

• Redistributes every row across 200 partitions based on sort keys

• Sorts each partition locally

• Produces globally ordered output

You just shuffled 1 million rows across the cluster to achieve global ordering – even though you're immediately partitioning by department on write, which destroys that global order anyway.

Why This Hurts

Range partitioning for global sort is one of the most expensive shuffles Spark performs:

• Sampling overhead: Spark must scan data twice (once to sample, once to process)

• Network transfer: Every row moves to a new executor based on range boundaries

• Disk I/O: Shuffle files written and read from disk

• Wasted work: Global ordering across departments is meaningless when you partition by department

For 1M rows, this adds 8-12 seconds of pure shuffle overhead.

The Better Approach: Sort Locally Within Partitions

If you only need ordering within each department (or within each output partition), use sortWithinPartitions():

# Optimized approach: local sort only
df_sorted = df.sortWithinPartitions(col("department"), col("salary").desc())
df_sorted.write.partitionBy("department").parquet("s3://output/employees/")

Simplified Logical Plan:

Sort [department ASC, salary DESC], false

└─ LogicalRDD [id, firstname, lastname, department, salary, age, hire_date, country]

• No Exchange.

• No shuffle.

• Just local sorting within existing partitions.

Spark sorts each partition in-place, without moving data across the network. The false flag in the Sort node indicates this is a local sort, not a global one.

Real-World Benchmark: AWS Glue

Let's measure the difference on 1 million employee records: First, will start with Global Sort with orderBy:

print("\n--- Testing orderBy() (global sort) ---")

start = time.time()

df_global = df.orderBy(col("department"), col("salary").desc())
df_global.write.mode("overwrite").parquet("/tmp/global_sort_output")

global_time = round(time.time() - start, 2)
print(f"orderBy() time: {global_time}s")

Local Sort:

print("\n--- Testing sortWithinPartitions() (local sort) ---")

start = time.time()

df_local = df.sortWithinPartitions(col("department"), col("salary").desc())
df_local.write.mode("overwrite").parquet("/tmp/local_sort_output")

local_time = round(time.time() - start, 2)
print(f"sortWithinPartitions() time: {local_time}s")

Physical Plan Differences:

orderBy() Physical Plan:

*(2) Sort [department ASC NULLS FIRST, salary DESC NULLS LAST], true, 0

+- Exchange rangepartitioning(department ASC NULLS FIRST, salary DESC NULLS LAST, 200)

   +- *(1) Project [id, firstname, lastname, department, salary, age, hire_date, country]

      +- *(1) Scan ExistingRDD[id, firstname, lastname, department, salary, age, hire_date, country]

The Exchange rangepartitioning node marks the shuffle boundary. Spark must:

• Sample data to determine range splits

• Redistribute all rows across executors

• Sort within each range partition

sortWithinPartitions() Physical Plan:

*(1) Sort [department ASC NULLS FIRST, salary DESC NULLS LAST], false, 0

+- *(1) Project [id, firstname, lastname, department, salary, age, hire_date, country]

   +- *(1) Scan ExistingRDD[id, firstname, lastname, department, salary, age, hire_date, country]

No Exchange. The false flag in Sort indicates local sorting only. Each partition is sorted independently, in parallel, without any data movement.

When to Use Which:

Common Anti-Pattern

df.orderBy("department", "salary") \
  .write.partitionBy("department") \
  .parquet("output/")

Problem: You're globally sorting by department, then immediately partitioning by department. The global order is destroyed during partitioning.

Here’s the fix:

df.sortWithinPartitions("department", "salary") \
  .write.partitionBy("department") \
  .parquet("output/")

Or even better, if you're partitioning by department anyway:

# Best: let partitioning handle distribution
df.write.partitionBy("department") \
    .sortBy("salary") \
    .parquet("output/")

orderBy() triggers an expensive full shuffle using range partitioning, while sortWithinPartitions() sorts data locally without a shuffle and is often 4–5× faster. Use it when writing partitioned files, computing window functions with partitionBy(), or when order is needed only within groups, and reserve orderBy() strictly for true global ordering, because in most production ETL, the best sort is the one that doesn’t shuffle.

Conclusion

You began this handbook likely wondering why your Spark application was slow, and now you see that the answer was both clear and not so clear: your problem was never your Spark application, your configuration, or your version of Spark. It was your plan all along.

You now understand that Spark runs plans, not code, that transformation order affects logical plans, that shuffles generate stages and are key to runtime performance, and that examining your physical plans allows you to directly link your application performance issues back to your problematic line of code.

And you’ve seen this pattern repeat across many scenarios: problem, plan, solution, improved plan, and so forth, until optimization feels less like a dark art and more like a certainty.

This is the Spark optimization mindset: read plans before you write code, and challenge every single Exchange. Engineers who write high-performance Spark jobs minimize shuffles, filter early, project narrowly, deal with skew carefully, and validate everything via explain() and the Spark UI. Once you learn to read the plan, Spark performance becomes mechanical.

Java Pattern Matching allows you to write more concise and readable code when working with complex data structures. It simplifies extracting data from data structures and performing operations on them.

Let's dive in and learn more about pattern matching and sealed classes.

What is Pattern Matching in Java?

Pattern Matching matches a value against a pattern that includes variables and conditions. If the value matches the pattern, the corresponding parts of the value are bound to the variables in the pattern. This allows for more readable and intuitive code.

There are two types of pattern matching: traditional and modern. Let's see the differences between the two in the following sections.

Traditional Pattern Matching

In traditional pattern matching, the switch statement is extended to support pattern matching by adding the case keyword with a pattern argument. The switch statement can match against a primitive type, wrappers, enums, and strings.

For example:

private static void printGreetingBasedOnInput(String input){
        switch (input){
            case "hello":
                System.out.println("Hi There");
                break;
            case "goodbye":
                System.out.println("See you Later!");
                break;
            case "thank you":
                System.out.println("You are welcome");
                break;
            default:
                System.out.println("I don't understand");
                break;
        }
    }

The Java method printGreetingBasedOnInput takes an input string and prints a corresponding greeting message based on its value using a switch-case statement. It covers cases for "hello," "goodbye," and "thank you," providing appropriate responses, and defaults to "I don't understand" for any other inputs.

Modern Pattern Matching

In modern pattern matching, the switch statement can match against a variety of patterns like any type of objects, enums, or primitives. The case keyword is used to specify the pattern to match against.

private static void printGreetingBasedOnInput(String input){
        switch (input){
            case "hello" -> System.out.println("Hi There");
            case "goodbye" -> System.out.println("See you Later!");
            case "thank you" -> System.out.println("You are welcome");
            default -> System.out.println("I don't understand");
        }
    }

This code snippet uses a more concise syntax. It simplifies the code by directly specifying the action to perform for each case label.

Before Java 16, we needed to check the object's type and then explicitly cast it to a variable. The enhanced instanceof operator introduced in Java 16 can both verify the type and perform an implicit cast to a variable, like in the example below:

private static void printType(Object input){
        switch (input) {
            case Integer i -> System.out.println("Integer");
            case String s -> System.out.println("String!");
            default -> System.out.println("I don't understand");
        }
    }

The enhancement of instanceof becomes particularly valuable when working with pattern guards. Pattern guards are a way to make case statements in Java pattern matching more specific by including boolean expressions.

This allows for more fine-grained control over how patterns are matched and can make code more readable and expressive.

private static void printType(Object input){                                                 
    switch (input) {                                                                         
        case Integer i && i > 10 -> System.out.println("Integer is greater than 10");        
        case String s && !s.isEmpty()-> System.out.println("String!");                       
        default -> System.out.println("Invalid Input");                                      
    }                                                                                        
}

Based on the examples shown above, you can hopefully see that Java pattern matching offers various benefits:

• It improves code readability by allowing for efficient matching of values against patterns and extracting data.

• It reduces code duplication by handling different cases with a single piece of code.

• It enhances type safety by enabling the matching of values against specific types.

• Pattern guards can be used within cases to further enhance code readability and maintainability.

What Are Sealed Classes in Java?

Sealed classes allow developers to restrict the set of classes that can extend or implement a given class or interface.

Sealed classes provide a way to create a hierarchy of classes or interfaces that can be extended or implemented by only a specified set of classes.

For example:

public sealed class Result permits Success, Failure {
    protected String response;

    public String message(){
        return response;
    }
}

public final class Success extends Result {

    @Override
    public String message() {
        return "Success!";
    }
}

public final class Failure extends Result {

    @Override
    public String message() {
        return "Failure!";
    }
}

In this example, we've defined a sealed class called Result which can be extended by either Success or Failure classes.

Any other class that attempts to extend Result will result in a compilation error.

This provides a way to restrict the set of classes that can be used to extend Result, making the code more maintainable and extensible.

A few important points to keep in mind:

• If a subclass wants to be a permitted subclass of a sealed class in Java, it must be defined in the same package as the sealed class. If the subclass is not defined in the same package, a compilation error will occur.

• If a subclass is to be permitted to extend a sealed class in Java, it must have one of three modifiers: final, sealed, or non-sealed.

• A sealed subclass must define the same or more restrictive set of permitted subclasses as its sealed superclass. Sealed subclasses must be either final or sealed. Non-sealed subclasses are not allowed as permitted subclasses of a sealed superclass and all permitted subclasses must belong to the same package as the sealed superclass.

How to Combine Java Pattern Matching and Sealed Classes

You can use sealed classes and their permitted subclasses in switch statements with pattern matching.

This can make the code more concise and easier to read. Here's an example:

private static String checkResult(Result result){                                     
    return switch (result) {                                                          
        case Success s -> s.message();                                                
        case Failure f -> f.message();                                                
        default -> throw new IllegalArgumentException("Unexpected Input: " + result); 
    };                                                                                
}

In the case of sealed classes, the compiler requires a default branch in pattern matching to ensure that all possible cases are covered.

Since sealed classes have a fixed set of permitted subclasses, it is possible to cover all cases with a finite number of case statements.

If a default branch is not included, it's possible to add a new subclass to the hierarchy in the future, which would not be covered by the existing case statements. This would result in a runtime error, which could be difficult to debug.

By requiring a default branch, the compiler ensures that the code is complete and covers all possible cases, even if new subclasses are added to the sealed class hierarchy in the future.

This helps to prevent runtime errors and makes the code more robust and maintainable.

If we modify the Result class to include a new subclass called Pending and we have not included that in our pattern matching, it will be covered by the default branch.

What is a Sealed Interface in Java?

When working with a Sealed Interface in Java, the compiler will not require a default branch in pattern matching if all cases are covered.

In case a branch is missing, the compiler will require a default branch to ensure that all possible cases are handled.

We are always required to include the default branch when working with Sealed Classes. Here's a code example:

public sealed interface OtherResult permits Pending, Timeout {
    void message();
}

public final class Pending implements OtherResult{
    @Override
    public void message() {
        System.out.println("Pending!");
    }
}

public final class Timeout implements OtherResult{
    @Override
    public void message() {
        System.out.println("Timeout!");
    }
}

private static void checkResult(OtherResult result){
        switch (result) {
            case Pending p -> p.message();
            case Timeout t -> t.message();
        };
    }

Conclusion

Here are some key takeaways of using Pattern Matching and Sealed Classes in Java code:

• Improved readability: Pattern Matching and Sealed Classes can make code more expressive and easier to read, since they allow for more concise and intuitive syntax.

• Greater control over class hierarchies: Sealed Classes provide a way to control class hierarchies and ensure that only permitted subclasses can be used. This can improve code safety and maintainability.

• Implicit type safety: Pattern Matching and Sealed Classes provide implicit type safety, which can reduce the risk of runtime errors and make code easier to maintain.

• Reduced code duplication: Pattern Matching and Sealed Classes can reduce code duplication by allowing different cases to be handled in a single piece of code.

• Better code organization: Sealed Classes can help to organize code and reduce the complexity of class hierarchies by grouping related classes together.

• Improved maintainability: Pattern Matching and Sealed Classes can improve the maintainability of code by making it easier to understand and update, which can save time and effort in the long run.

Thank you so much for reading.

This article explores how to implement FP concepts in Java, including viewing functions as first-class citizens, chaining, and composing them to create function pipelines.

We'll also discuss the technique of currying, which allows a function that takes multiple arguments to be transformed into a chain of functions that each take a single argument. This can simplify the use of complex functions and make them more reusable.

In this article, I'll show you examples of how to implement these concepts in Java using modern language features, like “java.util.function.Function”, “java.util.function.BiFunction”, and a user-defined TriFunction. We'll then see how to overcome its limitation using currying.

The java.util.function.Function interface

The java.util.function.Function interface is a key component of the Java 8 functional programming API. The java.util.function.Function is a functional interface in Java that takes input of type 'T' and produces an output of type 'R'.

In functional programming, functions are first-class citizens, meaning that they can be passed around as values, stored in variables or data structures, and used as arguments or return values of other functions.

The function interface provides a way to define and manipulate functions in Java code.

The function interface represents a function that takes one input and produces one output. It has a single abstract method, apply(), which takes an argument of a specified type and returns a result of a specified type.

You can use the function interface to define new functions, as well as to manipulate and compose existing functions. For example, you can use it to convert a list of objects of one type to a list of objects of another type, or to apply a series of transformations to a stream of data.

One of the main benefits of the function interface is that it allows you to write code that is more concise and expressive. By defining functions as values and passing them around as arguments or return values, developers can create more modular and reusable code. Also, by using lambdas to define functions, Java code can be more expressive and easier to read.

You can think about java.util.function.Function like this:

Function<A, B>

The java.util.function.BiFunction is also a functional interface in Java that takes two inputs 'T' and 'U' and produces an output of type 'R'. In short, BiFunction takes 2 inputs and returns an output:

BiFunction<Integer, Integer, Integer> sum = (a, b) -> a + b;

You can think about java.util.function.BiFunction like this:

BiFunction<A, B, C>

Function Composition and Chaining

Function chaining is a technique in functional programming that involves composing multiple functions into a single pipeline or chain.

In Java FP, function chaining is often used to transform data in a series of steps, where each step applies a specific transformation to the data and passes it on to the next step in the chain.

The function interface in Java provides a powerful tool for function chaining. Each function in the chain is defined as a separate instance of the Function interface. The output of each function becomes the input to the next function in the chain. This allows for a series of transformations to be applied to the data, one after another, until the final result is produced.

How to use the andThen method

The "andThen()" method is a default method provided by the function interface in Java. This method takes a sequence of two functions and applies them in succession, using the output of the first function as the input to the second function. This chaining of the functions results in a new function that combines the behavior of both functions in a single transformation.

The syntax of andThen looks like this:

default <V> Function<T, V> andThen(Function<? super R, ? extends V> after)

andThen

We can break the example down into series of steps. Initially the "multiply" function will be executed and its output passed as input to the "add" function. Then the "logOutput" function is used to log the resulting output.

public class FunctionChainingExample {

    static Logger logger = Logger.getLogger(FunctionChainingExample.class.getName());

    private static Function<Integer, Integer> multiply = x -> x * 2;

    private static Function<Integer, Integer> add = x -> x + 2;

    private static Function<Integer, Unit> logOutput = x -> {
        logger.info("Data:" + x);
        return Unit.unit();
    };

    public static Unit execute(Integer input) {
        Function<Integer, Unit> pipeline = multiply
                                               .andThen(add)
                                               .andThen(logOutput);
        return pipeline.apply(input);
    }

    public static void main(String[] args) {
        execute(10);
    }

}

The the above code snippet showcases the creation of three functions, namely multiply, add, and logOutput.

The function multiply accepts an integer input and produces an integer output, just like the add function. But, the logOutput function accepts an integer input and returns nothing (as represented by the Unit object, which implies the absence of a value).

The execute function chains the three functions together, where the output of multiply is utilized as the input for the add function, and the resulting output of add is passed to the logOutput function for logging purposes. The above example showcases function chaining using the andThen default method.

How to use the compose method

Unlike the andThen method, the compose method is another default method provided by the function interface in Java. It applies the first function to the output of the second function.

This means that the second function is applied to the input first, and then the first function is applied to the output of the second function. As a result, a chain of functions is created where the output of the second function becomes the input of the first function.

The compose function looks like this:

default  Function<V, R> compose(Function<? super V, ? extends T> before)

compose

When to use andThen vs compose depends on the order in which you want the functions to be applied.

If you want to apply the functions in the order they are defined, from left to right, then you should use the andThen method. This method applies the first function to the input, and then applies the second function to the output of the first function.

This is useful when you want to chain together functions that process the input data in a specific order.

On the other hand, if you want to apply the functions in the opposite order, from right to left, then you should use the compose method. This method applies the second function to the input, and then applies the first function to the output of the second function.

This is useful when you want to chain together functions that need to be applied in the opposite order of the way they are defined.

In general, you should use andThen when you want to apply the functions in the order they are defined, and use compose when you want to apply the functions in the opposite order. But the choice between the two methods ultimately depends on the specific requirements of your use case.

 public static Unit compose(Integer input) {
        Function<Integer, Unit> pipeline = logOutput
                                            .compose(add)
                                            .compose(multiply);
        return pipeline.apply(input);
    }

In the compose example, the functions are executed in a right-to-left order. Firstly, multiply is executed and its output is passed to the add function. Then, the resulting output of the add function is passed to logOutput for logging purposes.

How to use the apply() function

The apply() method takes an input and returns a result. It is used to apply a function to an argument and compute a result.

The apply() function is a method of functional interfaces, such as the function interface, that takes an argument of a specified type and returns a result of a specified type. It is the single abstract method of these interfaces, which is required for them to be used as functional interfaces.

The apply() function defines the behavior of the functional interface. When an instance of a functional interface is created, the apply() function is implemented to define what the functional interface does when it is called with an argument.

Function<Integer, Integer> multiply = x -> x * x;
        int result = multiply.apply(5);

In this code, we define a function called multiply that takes an integer argument and returns the square of that integer. We then call the apply() function on this function, passing in the integer value of 5. The apply() function executes the lambda expression defined in the square Function, and returns the result of the computation, which in this case is 25.

Let's understand by looking at a real-world example, FileReaderPipeline.

public class FileReaderPipeline {

    static Function<String, String> trim = String::trim;
    static Function<String, String> toUpperCase = String::toUpperCase;
    static Function<String, String> replaceSpecialCharacters = 
        str -> str.replaceAll("[^\\p{Alpha}]+", "");

    static Function<String, String> pipeline = 
                                trim
                                  .andThen(toUpperCase)
                                  .andThen(replaceSpecialCharacters);

 public static void main(String[] args) throws IOException {
        try (BufferedReader reader = new BufferedReader(new     FileReader("example.txt"))) {
            String line;
            while ((line = reader.readLine()) != null) {
                String result = pipeline.apply(line);
                System.out.println(result);
            }
        }
    }
}

The above code defines a FileReaderPipeline class, which reads lines from a file, processes them through a pipeline of functions, and prints the resulting output.

First, three functions trim, toUpperCase, and replaceSpecialCharacters are defined using method references and lambda expressions. trim removes leading and trailing whitespace from a String, toUpperCase converts the string to uppercase, and replaceSpecialCharacters removes any non-alphabetic characters from the string.

Next, a function pipeline is created using the andThen method, which chains the three functions together. The input to the pipeline function is a String, and the output is also a String. When a String is passed to the pipeline function, it first removes any leading or trailing whitespace, then converts the string to uppercase, and finally removes any non-alphabetic characters.

Finally, in the main method, the program reads lines from a file (in this example, "example.txt") using a BufferedReader. Each line is processed through the pipeline function using the apply method, which applies each function in the pipeline in sequence, and returns the resulting output. The resulting output is then printed to the console using System.out.println.

Function Currying

The Java util package contains two functional interfaces known as "Function<A, B>" and "BiFunction<A, B, C>". The Function interface takes a single input and produces an output, whereas the BiFunction interface takes two inputs and produces an output. Here is an illustration:

 Function<String, String> toUpper = str -> str.toUpperCase();
 BiFunction<Integer, Integer, Integer> sum = (a, b) -> a + b;

This is a clear constraint as there are only two functions, with one accepting a single input and the other accepting two inputs. Consequently, if we need a function that takes three inputs, we must construct our own customized function.

To this end, let's design a function that accepts three inputs and name it TriFunction. Here is an example implementation:

@FunctionalInterface
public interface TriFunction<A, B, C, O> {
    O apply(A a, B b, C c);

    default <R> TriFunction<A, B, C, O> andThen(TriFunction<A, B, C, O> after) {
        return (A a, B b, C c) -> after.apply(a,b,c);
    }
}

TriFunction<Integer, Integer, Integer, Integer> sum = (a, b, c) -> a + b + c;                                                                
int result = sum.apply(1, 2, 3);

In situations where we require more parameters than what the available functions can accept, we can leverage currying to overcome this limitation.

We can refactor the previous function using currying, which involves decomposing a function that takes multiple arguments into a sequence of functions, each of which accepts only one argument. This is in line with the definition of currying, which is a technique employed in Functional Programming.

Function<Integer, Function<Integer, Function<Integer, Integer>>>     sumWithThreeParams = (a) -> (b) -> (c) -> a + b + c;
                                                                       Function<Integer, Function<Integer, Function<Integer, Function<Integer, Integer>>>> sumWithFourParams = (a)->(b)->(c)->(d) -> a + b + c + d;

The above declares a new Function named sumWithThreeParams that takes an Integer input. It returns a new Function that takes an Integer input and returns yet another new Function that takes an Integer input and returns an Integer output. Similarly a new Function named sumWithFourParams is created.

Note that the curried function sumWithThreeParams and sumWithFourParams allows us to partially apply arguments to the function, which can be useful in situations where we have some of the inputs available but not all of them

Function<Integer, Function<Integer, Integer>> partialSumWithTwoParams                                     = sumWithThreeParams.apply(5);
                                                                        Function<Integer, Integer> partialSumWithOneParams 
                                 = partialSumWithTwoParams.apply(10);

                                                                        int c = 15;
int result = partialSumWithOneParams.apply(c);
System.out.println(result); //30

Conclusion

In this article, we covered the implementation of functional programming (FP) concepts in Java, such as treating functions as first-class citizens, composing them into pipelines, and utilizing currying to simplify complex functions.

We have provided examples using modern language features like java.util.function.Function and java.util.function.BiFunction, as well as a user-defined TriFunction with currying to overcome its limitations.

By applying these techniques, developers can write more concise, reusable, and maintainable code in Java.

With the growing popularity of FP in the industry, understanding and implementing these concepts is becoming increasingly important for Java developers.

After reading this article, if you have found it useful and would like to explore more practical examples, please visit the repository and consider giving it a star.

SpringBoot has a number of in-built actuators, and it also allows us to create our own Actuator Endpoint.

For frameworks written in Python like Flask or FastAPI, we can incorporate actuators by integrating a library called Pyctuator.

In this article I am going to explain how to monitor applications written in FastAPI using the Pyctuator library. I'll also show you how to manage the actuator endpoints using the SpringBootAdmin server.

What are Actuators?

We use actuators for monitoring and managing application usage in production. This usage information gets exposed to us via REST endpoints.

For example, we can access the application logs in production, environment details, and HTTP traces. And if something has gone wrong within the application, we can even access the applications “threaddump” for debugging purposes.

Here are some examples of few important actuator endpoints:

Actuators

What is Pyctuator?

Actuators become popular because of SpringBoot, but you can implement them in frameworks like FastAPI or Flask by integrating a module called Pyctuator.

Pyctuator is a Python Module, which is a partial implementation of SpringBoot Actuators. Pyctuator is managed by SolarEdge.

Some of the actuators supported by Pyctuators are:

• /health: This endpoint in Pyctuator has built-in monitoring for Redis and MySQL

• /env

• /metrics

• /logfile

• /threaddump

• /httptrace

• /loggers

What is SpringBootAdmin?

Imagine a service having all these actuators for checking metrics, httptrace, threaddump and so on. It would be pretty tedious to invoke each one of them individually to check what’s going on within the service. And if we have many services and each one of them has its own actuator endpoints, this makes monitoring even more difficult.

That’s where you can use SpringBootAdmin to manage and monitor applications.

In a nutshell, SpringBootAdmin provides a nice dashboard for all the actuator endpoints in one place.

Admin Dashboard

Use-Case for Pyctuator

The use-case is straightforward: we are going to develop a RESTful service using FastAPI framework and configure the actuators in the service using the Pyctuator module.

The service has 3 endpoints as shown in the API-Docs (Swagger) below

API-Docs

• GET /users: Return all the users that exists in the system.

• POST /users: Create user

• GET /users/{id}: Return a user with a given id

You can find the code here.

In the User-Service we are going to enable the actuators using Pyctuator and monitor them using SpringBootAdmin dashboard. We are also going to explore how we can enhance the /health actuator for monitoring Redis.

For configuring the Actuators, first we need to install “pyctuator”. You can do that using the command “pip install pyctuator”.

After installation, simply instantiating the Pyctuator object is the entry point for seeing in-built actuators within a web-framework.

Pyctuator Constructor

Before instantiating Pyctuator, if you access the /pyctuator endpoint you will get the “Not Found” message:

Without Pyctuator Configuration

After instantiation, on accessing the /pyctuator endpoint you will see all the actuators enabled by default. This is because we have defined "pyactuator_endpoint_url" within Pyctuator.

After Pyctuator Configuration

I strongly recommend going through the Pyctuator object as it explains what the mandatory and optional arguments are that we need to provide.

Understanding Constructor parameters

The mandatory parameters are:

• app – instance of FastAPI or Flask

• the application name – displayed in the info section in SpringBootAdmin

• “pyctuator_endpoint_url” – what we have seen which returns all the actuator endpoints

• “registration_url” – you will understand this one shortly.

How to Enhance the /health Endpoint

You can enhance the /health endpoint in Pyctuator to monitor Redis or MySQL databases. Say you are using Redis in your application – then we need to use RedisHealthProvider and pass the redis instance to it.

Redis Health

How to Start the SpringBootAdmin Server

To run the SpringBootAdmin server on local, we have two options: first, we can do it by creating SpringBootAdmin manually by going to start.spring.io and adding libraries.

Spring Web & Spring Boot Admin (Server):

Creating SpringBootAdmin

The second option is to run a Docker image:

docker run --rm --name spring-boot-admin -p 8080:8080 michayaak/spring-boot-admin:2.2.3-1

Once the admin server is up, we need to provide the “registration_url” to the Pyctuator Constructor as discussed earlier.

URL registration in SpringBootAdmin

The Admin Server is running on localhost:8080 and this should register our application to SpringBootAdmin. We can access all the actuator endpoints in one place:

Dashboard

All Configured Actuator Endpoints

I executed the /users endpoint few times and now HTTP Traces on the Admin side showcases all the Request-Response exchange details.

HTTP Traces

Wrapping Up

Actuators are extremely helpful in monitoring and debugging applications in production. By accessing endpoints we can get details on thread dumps, heap dumps, HTTP Traces and so on.

Pyctuator simplifies having actuators in Python APIs to a great extent. By simply importing the library and defining an object, all the actuators are ready for us within our application.

These Docker containers are lightweight, and once the tests are finished the containers get destroyed.

In the article we are going to understand what the TestContainers is and how it helps you write more reliable Tests.

We are also going to understand the important components (Annotations and Methods) of the library which help you write the Tests.

Finally, we will also learn to write a proper Integration Test in SpringBoot using the TestContainers library and its components.

Limitations of Testing with an H2 In-Memory Database

The most common approach to Integration Testing today is to use an H2 in-memory database. But there are certain limitations to this method.

First of all, say we are using version 8.0 of MySQL in production, but our integration tests are using H2. We can never execute our tests for the database version running on Production.

SpringBoot app with MySQL DB and H2

‌Secondly, the test cases are less reliable because in production we are using an altogether different database and the tests are pointing to H2. The application may run into issues in production, but the integration tests may succeed.

I was trying to access my RESTful service on local and faced this error:

“Caused by: org.postgresql.util.PSQLException: FATAL: database "example_db" does not exist”.

It happened because of a permission issue, but the tests on local worked fine.

And finally, as documented here, H2 is compatible with other databases only up to a certain point. There are few areas where H2 is incompatible. If you need to use “nativeQueries” in a SpringBoot application, for example, then using H2 may cause problems.

Enter the TestContainers Library

By using TestContainers we can overcome the limitations of H2.

• Integration tests will point to the same version of the database as it’s in production. So we can tie our TestContainer Database Image to the same version running on production.

• Integration tests are lot more reliable because both application and tests are using the same database type and version and there won't be any compatibility issues in Testcases.

What is TestContainers?

The TestContainers library is a wrapper API over Docker. When we write code to create a container behind the scenes it may be translated to some Docker command, for example‌:

MySQLContainer Creation

This code may be translated to something like the following:

docker run -d --env MYSQL_DATABASE=example_db --env MYSQL_USER=test --env MYSQL_PASSWORD=test ‘mysql:latest’

TestContainers has a method name “withCommand”. You use it to set the command that should be run inside the Docker container which confirms that TestContainers is a wrapper API over Docker.

TestContainers downloads the MySQL, Postgres, Kafka, Redis images and runs in a container. The MySQLContainer will run a MySQL Database in a container and the Testcases can connect to it on the local machine. Once the execution is over the Database will be gone – it just deletes it from the machine. In the Testcases we can start as many container images as we want.

TestContainers supports JUnit 4, JUnit 5 and Spock. If you go to the TestContainers.org website, just visit the QuickStart section that explains how to use it:

TestContainers.org how to start with Test Framework

TestContainers supports almost every Database from MySQL and Postgres to CockroachDB. You can find more info about this on the TestContainers.org website under the Modules section:‌

TestContainers support for Database Modules

‌TestContainers also supports Cloud Modules like GCloud Module and Azure Module as well. If your application is running on Google Cloud, then TestContainers has support for Cloud Spanner, Firestore, Datastore and so on.‌

TestContainers support for GCloud Module

In the article so far, we have discussed only about Databases, but TestContainers supports various other components like Kafka, SOLR, Redis, and more.

How to Use the TestContainers Library

In this article we are going to explore TestContainers with JUnit 5. To implement TestContainers we need to understand a few important TestContainers annotations, methods, and the libraries that we need to implement in our project.

TestContainers libraries

Annotations in TestContainers

Two important annotations are required in our Tests for TestContainers to work: @TestContainers and @Container.

@TestContainer is JUnit-Jupiter extension which automatically starts and stops the containers that are used in the tests. This annotation finds the fields that are marked with @Container and calls the specific Container life-cycle methods. Here, MySQLContainer life-cycle methods will be invoked.‌

MySQLContainer

The MySQLContainer is declared as static because if we declare Container as static then a single container is started and it will be shared across all the test methods.

If it’s an instance variable, then a new container is created for each test method.

TestContainers Library Methods

There are few important methods in TestContainers library that you'll use in the tests. They're good to know before using the library.

• withInitScript: Using ‘withInitScript’ we can execute the .SQL to define the schema, tables, and plus add the data into the database. In short, this method is used to run the .SQL to populate the database.

• withReuse (true): Using “withReuse” method we can enable the reuse of containers. This method works well in conjunction with enabling the “testcontainers.reuse.enable:true” property in the “.testcontainers.properties” file.

• start: we use this to start the container.

• withClasspathResourceMapping: This maps a resource (file or directory) on the classpath to a path inside the container. This will only work if you are running your tests outside a Docker container.

• withCommand: Set the command that should be run inside the Docker container.

• withExposedPorts: Used to set the port that the container listens on.

• withFileSystemBind: Used to map a file / directory from the local filesystem into the container.

TestContainers Use Case

In the example we'll look at now, the application will communicate only with the database and write the integration tests for it using TestContainers. Then we'll extend the use-case by implementing Redis in between.

If the data exists in the Redis cache it will be returned, otherwise it'll dip into the database for saving and retrieval based on the Key.

Use-Case

The service is simple. It has 2 endpoints – the first one is to create a user and the second one is to find a user by email. If the user is found it is returned, otherwise we get a 404. The service class code looks something like this:‌

Service Component

We are going to write tests for this class. You can find the entire codebase here:‌

Test Class

The test class is marked with @TestContainers annotation which starts/stops the container. We use the @Container annotation to call the specific container's life-cycle methods.

Also, the “MySQLContainer” is declared as static because then a single container is started. Then it gets shared across all the test methods (we have already discussed the importance of these annotations).

BeforeAll

Next we need to write a setup method marked with @BeforeAll, where we have enabled the “withReuse” method. This helps us reuse the existing containers. We are using the “withInitScript” method to execute the “.sql” file and then starting the container.‌

Overwriting Properties

‌@DynamicPropertySource helps us override the properties declared in the properties file. We write this method to allow TestContainers to create the URL, username, and password on its own – otherwise we may face errors.

For example on removing username and password we may face an ‘Access denied’ error which may confuse us. So it’s better to allow TestContainer to assign these properties dynamically on its own.

That’s it – we are ready to run the Testcases:

Test Cases

Execute @AfterAll to stop the container, otherwise it may keep running on your local machine if you don't explicitly stop it. ‌

How to Use GenericContainer

‌‌GenericContainer is the most flexible container. It makes it easy to run any container images within GenericContainer.

Now we have Redis in place, all we need to do in our Testcase is to spin up a GenericContainer with the Redis image.

GenericContainer for Redis

Then we start the Generic Redis container in @BeforeAll and stop it with the @AfterAll tear down method.

Starting Containers

Stopping Containers

Wrapping Up

It's extremely easy to use TestContainers in our application to write better tests. The learning curve is not too steep and it has support for various different modules from a variety of databases like Kafka, Redis and others.

Writing tests using TestContainers makes our tests lot more reliable. The only flip side is that the tests are slow compared to H2. This is because H2 is in memory and TestContainers takes time to download the image, run the container, and execute the entire setup we have discussed in this article.

A workflow as a sequence of operations, from start to finish. The workflows in Airflow are authored as Directed Acyclic Graphs (DAG) using standard Python programming.

You can configure when a DAG should start execution and when it should finish. You can also set up workflow monitoring through the very intuitive Airflow UI.

You can be up and running on Airflow in no time – it’s easy to use and you only need some basic Python knowledge. It's also completely open source.

Apache Airflow also has a helpful collection of operators that work easily with the Google Cloud, Azure, and AWS platforms.

In this article we are going to cover

• What are Directed Acyclic Graphs (DAGs)?

• What are Operators?

• How to Create your First DAG

• A Use-Case for DAGs

• How to Set Up Cloud Composer

• How to Run the Pipeline on Composer

What are Directed Acyclic Graphs, or DAGs?

DAGs, or Directed Acyclic Graphs, have nodes and edges. DAGs should not contain any loops and their edges should always be directed.

In short, a DAG is a data pipeline and each node in a DAG is a task. Some examples of nodes are downloading a file from GCS (Google Cloud Storage) to Local, applying business logic on a file using Pandas, querying the database, making a rest call, or uploading a file again to a GCS bucket.

Visualizing DAGs

Correct DAG with no loops

Incorrect DAG with Loop

You can schedule DAGs in Airflow using the schedule_interval attribute. By default it’s "None" which means that the DAG can be run only using the Airflow UI.

You can schedule the DAG to run once every hour, every day, once a week, monthly, yearly or whatever you wish using the cron presets options (@hour, @daily, @weekly, @hourly, @monthly, @yearly).

If you need to run the DAG every 5 mins, every 10 mins, every day at 14:00, or once on a specific day like every Thursday at 10:00am, then you should use these cron-based expressions.

*/5 * * * * = Every 5 minutes

0 14 * * * = Every day at 14:00

What are Operators?

A DAG consists of multiple tasks. You can create tasks in a DAG using operators which are nodes in the graph.

There are various ready to use operators available in Airflow, such as:

• LocalFilesystemToGCSOperator – use it to upload a file from Local to GCS bucket.

• PythonOperator – use it to execute Python callables.

• functionEmailOperator – use it to send email.

• SimpleHTTPOperator – use it to make an HTTP Request.

How to Create Your First DAG

The example DAG we are going to create consists of only one operator (the Python operator) which executes a Python function.

from airflow import DAG
from datetime import datetime
from airflow.operators.python_operator import PythonOperator

def message():
    print("First DAG executed Successfully!!")

with DAG(dag_id="FirstDAG", start_date=datetime(2022,1,23), schedule_interval="@hourly",
         catchup=False) as dag:

    task = PythonOperator(
        task_id="task",
        python_callable=message)

task

The first step is to import the necessary modules required for DAG development. The line with DAG is the DAG which is a data pipeline that has basic parameters like dag_id, start_date, and schedule_interval.

The schedule_interval is configured as @hourly which indicates that the DAG will run every hour.

The task in the DAG is to print a message in the logs. We have used the PythonOperator here. This operator is used to execute any callable Python function.

Once the execution is complete, we should see the message “First DAG executed Successfully” in the logs. We are going to execute all our DAGs on GCP Cloud Composer.

Airflow UI

After successful execution, the message is printed on the logs:

Logs

A Use-Case for DAGs

The use-case we are going to cover in this article involves a three-step process.

In step one, we will upload a .csv file in some input GCS bucket. This file should be processed by PythonOperator in the DAG. The function which will be executed by the PythonOperator consists of Pandas code, which represents how users can use Pandas code for transforming the data in the Airflow Data Pipeline.

In step two, we'll upload the transformed .csv file to another GCS bucket. This task will be handled by the GCSToGCSOperator.

Step three is to send the status email indicating the that the pipeline execution is completed which will be handled by the EmailOperator.

In this use-case we will also cover how to notify the team via email in case any step of the execution failed.

How to Install Cloud Composer

In GCP, Cloud Composer is a managed service built on Apache Airflow. Cloud Composer has default integration with other GCP Services such as GCS, BigQuery, Cloud Dataflow and so on.

First, we need to create the Cloud Composer Environment. So search for Cloud Composer on the search bar and click on "Create Environment" as shown below:

Create Environment

In the Environments option, I am selecting the "Composer 1" option as we don’t need auto-scaling.

Once we select the type of composer we need, we'll need to do some basic configuration just like in any GCP managed service ("Instance Name", "Location", and so on).

The node count here should always be 3 as GCP will setup the 3 services needed for Airflow.

Once we're done with that, it'll set up an Airflow instance for us. To upload a DAG, we need to open the DAGs folder shown in ‘DAGs folder’ section.

Airflow Instance

If you go to the "Kubernetes Engine" section on GCP, we can see 3 services up and running:

Kubernetes Engine

All DAGs will reside in a bucket created by Airflow.

Airflow Instance bucket for DAGs

How to Create and Run the Pipeline on Composer

In the Pipeline, we have two buckets. input_csv will contain the csv which requires some transformation, and the transformed_csv bucket will be the location where the file will be uploaded once the transformation is done.

The entire pipeline code is the following:

from airflow import DAG
from datetime import datetime
import pandas as pd

from airflow.utils.email import send_email
from airflow.operators.python_operator import PythonOperator
from airflow.operators.email_operator import EmailOperator
from airflow.providers.google.cloud.transfers.gcs_to_gcs import GCSToGCSOperator


def transformation():
    trainDetailsDF = pd.read_csv('gs://input_csv/Event_File_03_16_2022.csv')
    print(trainDetailsDF.head())


with DAG(
        dag_id="pipeline_demo",
        schedule_interval="@hourly",
        start_date=datetime(2022, 1, 23),
        catchup=False
) as dag:
    buisness_logic_task = PythonOperator(
        task_id='ApplyBusinessLogic',
        python_callable=transformation,
        dag=dag)

    upload_task = GCSToGCSOperator(
        task_id='upload_task',
        source_bucket='input_csv',
        destination_bucket='transformed_csv',
        source_object='Event_File_03_16_2022.csv',
        move_object=True,
        dag=dag
    )

    email_task = EmailOperator(
        task_id="SendStatusEmail",
        depends_on_past=True,
        to='youremail',
        subject='Pipeline Status!',
        html_content='<p>Hi Everyone, Process completed Successfully! <p>',
        dag=dag)

    buisness_logic_task >> upload_task >> email_task

In the first task, all we are doing is creating a DataFrame from the input file and printing the head elements. In the logs it looks like this:

DataFrame Head

In the second task, GCSToGCSOperator, we have used the attribute move_object=True which will delete the file from the Source bucket.

Once we upload the file to the bucket, we can see that the DAG is being scheduled. The name of the DAG is "pipeline_demo".

DAGs

Note that in case if you encounter any "import errors" after uploading or executing a DAG, something like this:

You can upload these missing packages through the "PYPI Packages" option in GCP. This will update the environment after few minutes.

Updating environment with missing Packages

To open an Airflow UI, Click on the "Airflow" link under Airflow webserver.

Airflow Instance, click Airflow link to Open UI

The Airflow UI looks like this:

Upon successful execution of Pipeline, here's what you should see:

In order to send email if a task fails, you can use the on_failure_callback like this:

def notify_email(contextDict, **kwargs):
    title = "Airflow alert: {task_name} Failed".format(**contextDict)
    body = """
    Task Name :{task_name} Failed.<br>
    """.format(**contextDict)
    send_email('youremail’, title, body)



buisness_logic_task = PythonOperator(
    task_id='ApplyBusinessLogic',
    python_callable=transformation,
    on_failure_callback=notify_email,
    dag=dag)

We're doing the notification email configuration on composer through Sendgrid. Also, once you are done with Cloud Composer, don't forget to delete the instance as it cannot be stopped.

Conclusion

Apache Airflow is a fairly easy-to-use tool. There's also a lot of help now available on the internet and the community is growing.

GCP simplified working with Airflow a lot by creating a separate managed service for it.

Dataproc is a Google Cloud Platform managed service for Spark and Hadoop which helps you with Big Data Processing, ETL, and Machine Learning. It provides a Hadoop cluster and supports Hadoop ecosystems tools like Flink, Hive, Presto, Pig, and Spark.

Dataproc is an auto-scaling cluster which manages logging, monitoring, cluster creation of your choice and job orchestration. You'll need to manually provision the cluster, but once the cluster is provisioned you can submit jobs to Spark, Flink, Presto, and Hadoop.

Dataproc has implicit integration with other GCP products like Compute Engine, Cloud Storage, Bigtable, BigQuery, Cloud Monitoring, and so on. The jobs supported by Dataproc are MapReduce, Spark, PySpark, SparkSQL, SparkR, Hive and Pig.

Apart from that, Dataproc allows native integration with Jupyter Notebooks as well, which we'll cover later in this article.

In the article, we are going to cover:

1. Dataproc cluster types and how to set Dataproc up

2. How to submit a PySpark job to Dataproc

3. How to create a Notebook instance and execute PySpark jobs through Jupyter Notebook.

How to Create a Dataproc Cluster

Dataproc has three cluster types:

1. Standard

2. Single Node

3. High Availability

The Standard cluster can consist of 1 master and N worker nodes. The Single Node has only 1 master and 0 worker nodes. For production purposes, you should use the High Availability cluster which has 3 master and N worker nodes.

For our learning purposes, a single node cluster is sufficient which has only 1 master Node.

Creating Dataproc clusters in GCP is straightforward. First, we'll need to enable Dataproc, and then we'll be able to create the cluster.

Start Dataproc cluster creation

When you click "Create Cluster", GCP gives you the option to select Cluster Type, Name of Cluster, Location, Auto-Scaling Options, and more.

Parameters required for Cluster

Since we've selected the Single Node Cluster option, this means that auto-scaling is disabled as the cluster consists of only 1 master node.

The Configure Nodes option allows us to select the type of machine family like Compute Optimized, GPU and General-Purpose.

In this tutorial, we'll be using the General-Purpose machine option. Through this, you can select Machine Type, Primary Disk Size, and Disk-Type options.

The Machine Type we're going to select is n1-standard-2 which has 2 CPU’s and 7.5 GB of memory. The Primary Disk size is 100GB which is sufficient for our demo purposes here.

Master Node Configuration

We've selected the cluster type of Single Node, which is why the configuration consists only of a master node. If you select any other Cluster Type, then you'll also need to configure the master node and worker nodes.

From the Customise Cluster option, select the default network configuration:

Use the option "Scheduled Deletion" in case no cluster is required at a specified future time (or say after a few hours, days, or minutes).

Schedule Deleting Setting

Here, we've set "Timeout" to be 2 hours, so the cluster will be automatically deleted after 2 hours.

We'll use the default security option which is a Google-managed encryption key. When you click "Create", it'll start creating the cluster.

You can also create the cluster using the ‘gcloud’ command which you'll find on the ‘EQUIVALENT COMMAND LINE’ option as shown in image below.

And you can create a cluster using a POST request which you'll find in the ‘Equivalent REST’ option.

gcloud and REST option for Cluster creation

After few minutes the cluster with 1 master node will be ready for use.

Cluster Up and Running

You can find details about the VM instances if you click on "Cluster Name":

How to Submit a PySpark Job

Let’s briefly understand how a PySpark Job works before submitting one to Dataproc. It’s a simple job of identifying the distinct elements from the list containing duplicate elements.

#! /usr/bin/python

import pyspark

#Create List
numbers = [1,2,1,2,3,4,4,6]

#SparkContext
sc = pyspark.SparkContext()

# Creating RDD using parallelize method of SparkContext
rdd = sc.parallelize(numbers)

#Returning distinct elements from RDD
distinct_numbers = rdd.distinct().collect()

#Print
print('Distinct Numbers:', distinct_numbers)

Upload the .py file to the GCS bucket, and we'll need its reference while configuring the PySpark Job.

Job GCS Location

Submitting jobs in Dataproc is straightforward. You just need to select “Submit Job” option:

Job Submission

For submitting a Job, you'll need to provide the Job ID which is the name of the job, the region, the cluster name (which is going to be the name of cluster, "first-data-proc-cluster"), and the job type which is going to be PySpark.

Parameters required for Job Submission

You can get the Python file location from the GCS bucket where the Python file is uploaded – you'll find it at gsutil URI.

No other additional parameters are required, and we can now submit the job:

After execution, you should be able to find the distinct numbers in the logs:

Logs

How to Create a Jupyter Notebook Instance

You can associate a notebook instance with Dataproc Hub. To do that, GCP provisions a cluster for each Notebook Instance. We can execute PySpark and SparkR types of jobs from the notebook.

To create a notebook, use the "Workbench" option like below:

Make sure you go through the usual configurations like Notebook Name, Region, Environment (Dataproc Hub), and Machine Configuration (we're using 2 vCPUs with 7.5 GB RAM). We're using the default Network settings, and in the Permission section, select the "Service account" option.

Parameters required for Notebook Cluster Creation

Click Create:

Notebook Cluster Up & Running

The "OPEN JUPYTYERLAB" option allows users to specify the cluster options and zone for their notebook.

Once the provisioning is completed, the Notebook gives you a few kernel options:

Click on PySpark which will allow you to execute jobs through the Notebook.

A SparkContext instance will already be available, so you don't need to explicitly create SparkContext. Apart from that, the program remains the same.

Code snapshot on Notebook

Conclusion

Working on Spark and Hadoop becomes much easier when you're using GCP Dataproc. The best part is that you can create a notebook cluster which makes development simpler.