DataFrame transformations and actions

⏱ 35 min

Hook

Two transformations look almost identical in your code: one adds a column, the other groups by a key. On the sample file they finish in the same second. In production, the column-adder still finishes in seconds and the grouping runs for twenty minutes. The API gave you no hint that one of them quietly moved every row in the dataset across the network. This lesson is about reading that hint yourself.

Concept

Almost every PySpark pipeline is built from a small set of DataFrame transformations. Learn these five and you can express most batch logic; learn which of them move data and you can predict which will be slow.

select projects columns: it keeps, drops, renames, or computes columns and returns a narrower DataFrame. filter (the same as where) keeps only the rows matching a condition. withColumn adds or replaces a single column, computed from an expression over existing columns. Each of these three works on one row at a time: to produce an output row, Spark needs only the matching input row, and never has to look at any other partition.

groupBy followed by an aggregation (sum, count, avg) collapses many rows sharing a key into one summary row per key. join matches rows from two DataFrames on a key. These two are different in kind: to compute a group's total, Spark must gather every row with that key together, and those rows start out scattered across every partition. The same is true of a join — matching keys have to meet.

That difference has a name. A narrow transformation is one where each output partition depends on exactly one input partition; the work stays local to a partition. A wide transformation is one where each output partition depends on many input partitions, so Spark must move data across the cluster to line up rows by key. select, filter, and withColumn are narrow. groupBy and join are wide.

The precise rule to carry forward: a wide transformation requires a shuffle, because its output partitions depend on multiple input partitions. Narrow transformations can be fused into a single stage and run in one pass; a wide transformation ends that stage, forces a shuffle, and starts a new one. That is the twenty-minute difference from the hook, made predictable.

Transformation	What it does	Kind	Moves data?
select	Keep, rename, or compute columns	Narrow	No
filter / where	Keep matching rows	Narrow	No
withColumn	Add or replace one column	Narrow	No
groupBy + agg	One summary row per key	Wide	Yes — shuffle
join	Match rows of two DataFrames on a key	Wide	Yes — shuffle

Keep the earlier distinction in mind: all of these are transformations, so none of them run until an action — show, count, collect, write — asks for a result. The narrow-versus-wide label tells you what each transformation will cost once that action finally fires.

🧠 Knowledge check

1. Which of these transformations forces Spark to move data across the cluster?

a) withColumn that adds a tax column from an existing amountb) filter that keeps rows where status = "paid"c) groupBy("category") followed by sum("amount")d) select of three existing columns

2. What defines a transformation as narrow rather than wide?

a) It returns fewer columns than it started withb) Each output partition depends on exactly one input partition, so the work stays localc) It runs on the driver instead of the executorsd) It always runs faster than any action

Worked example

Let me build a small pipeline that uses four of these operations, so you can see each one and label it narrow or wide as it goes. The data is a set of ride-share trips; the goal is revenue per city among completed trips.

from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("df-transforms").getOrCreate()

trips = spark.createDataFrame(
    [
        ("t-1", "austin", 12.0, 0.20, "completed"),
        ("t-2", "austin", 8.0, 0.20, "cancelled"),
        ("t-3", "denver", 20.0, 0.15, "completed"),
        ("t-4", "denver", 15.0, 0.15, "completed"),
        ("t-5", "austin", 10.0, 0.20, "completed"),
    ],
    ["trip_id", "city", "fare", "tip_rate", "status"],
)

# filter — narrow: keep only completed trips.
completed = trips.filter(F.col("status") == "completed")

# withColumn — narrow: total = fare plus a computed tip.
with_total = completed.withColumn(
    "total", F.col("fare") + F.col("fare") * F.col("tip_rate")
)

# select — narrow: keep just what the aggregation needs.
trimmed = with_total.select("city", "total")

# groupBy + agg — wide: revenue per city forces a shuffle.
revenue = trimmed.groupBy("city").agg(
    F.round(F.sum("total"), 2).alias("revenue")
)

revenue.show()

+------+-------+
|  city|revenue|
+------+-------+
|austin|  26.40|
|denver|  40.25|
+------+-------+

Read the pipeline as a sequence of costs. The filter, withColumn, and select are all narrow — Spark can fuse them into one stage and run them in a single pass over each partition, no data leaving the executor it started on. Only the final groupBy is wide: to total each city, Spark shuffles the trimmed rows so that every austin row lands together and every denver row lands together, then sums within each group. Notice the deliberate ordering — I filtered and trimmed to two columns before the groupBy, so the shuffle moves the smallest possible rows. That instinct is the whole of Lesson 3.

Hands-on

Your turn, on a different dataset. The exercise below loads retail order records (not trips) and asks you to build a pipeline that filters to shipped orders, adds a net_amount column, and computes the number of orders and total net amount per region. One of your transformations is wide; the exercise asks you to name which.

▶ Spark Forgede202-orders-by-region

Open in Spark Forge ↗

Success criteria: your pipeline returns one row per region with a correct order count and net-amount total, and you correctly label the groupBy as the wide transformation. Spark Forge validates both the output and your classification.

Recap

• You can shape data with the five core transformations: select, filter, withColumn, groupBy, and join.
• You can classify any of them as narrow or wide before running it, using the one-to-one-partition test.
• You can state the rule that a wide transformation requires a shuffle because its output partitions depend on multiple input partitions.
• You can order a pipeline so the narrow, data-shrinking steps come before the wide one.

Next up: what that shuffle actually costs — why wide operations are expensive, how to read the shuffle in a plan, and how caching stops Spark from repeating the whole chain.

🧠 Knowledge check

1. A join between an orders DataFrame and a customers DataFrame on customer_id is which kind of transformation, and why?

a) Narrow, because each row is processed independentlyb) Wide, because matching keys start scattered across partitions and must be brought together, which requires a shufflec) Narrow, because a join only adds columnsd) Neither — a join is an action, not a transformation

2. Why filter and select before a groupBy rather than after it?

a) It changes the aggregation result to be correctb) It shrinks the rows before the wide step, so the shuffle moves less datac) groupBy cannot run unless a filter comes firstd) Spark runs filters on the driver, which is faster

📌 Key takeaways

• Most batch logic is built from five transformations: select, filter, withColumn, groupBy, and join.
• Narrow transformations (select, filter, withColumn) keep work local; wide ones (groupBy, join) move data across the cluster.
• A wide transformation requires a shuffle because its output partitions depend on multiple input partitions.
• Put narrow, data-shrinking steps before the wide step so the shuffle moves less data.