Post

Apache Spark Transformations Every Data Engineer Should Know

In this article, let us go through the Apache Spark transformations that I use most often while building data pipelines. Knowing the syntax is useful, but it is equally important to know what Spark does with each transformation. A small choice such as using select instead of a Python UDF can make a noticeable difference when the data grows.

We will use PySpark and a small orders data set for the examples. The same concepts apply whether the job runs on Databricks, AWS Glue, Google Cloud Dataproc, or a local Spark installation.

1
2
3
4
5
6
7
8
9
10
11
from pyspark.sql import SparkSession
from pyspark.sql import functions as F

spark = SparkSession.builder.appName("transformations-demo").getOrCreate()

orders = spark.createDataFrame([
    (101, "C001", "AU", 120.50, "COMPLETE"),
    (102, "C002", "AU", 75.00, "PENDING"),
    (103, "C001", "NZ", 220.00, "COMPLETE"),
    (104, "C003", None, 45.50, "CANCELLED")
], ["order_id", "customer_id", "country", "amount", "status"])

One point to remember is that Spark transformations are lazy. Calling filter, select, or groupBy does not immediately process the data. Spark builds a logical plan and starts the work only when we call an action such as show, count, or write.

Quick comparison

TransformationTypical useShuffle likely?
select / withColumnChoose or derive columnsNo
filterRemove unwanted rowsNo
groupByAggregate recordsYes
joinCombine data setsUsually
dropDuplicatesRemove duplicate keysYes
Window functionsRank or compare rows in a groupYes
repartitionRedistribute dataYes

A shuffle means Spark moves records between executors. It is sometimes required, but it normally costs network, disk, and time.

1. Select only the columns needed

select is simple, but it is one of the most useful transformations. Do not carry fifty columns through a pipeline if the output needs only six. Selecting columns early reduces the amount of data that Spark has to serialize and move.

1
2
3
4
5
completed = orders.select(
    "order_id",
    "customer_id",
    F.col("amount").cast("decimal(12,2)").alias("order_amount")
)

We can also derive multiple columns in one select. I prefer this when a stage has several related mappings because the output schema is visible in one place.

2. Add and change columns with withColumn

withColumn is convenient for adding a derived field or replacing an existing one.

1
2
3
4
5
6
7
with_tax = orders.withColumn(
    "amount_with_tax",
    F.round(F.col("amount") * F.lit(1.10), 2)
).withColumn(
    "country",
    F.coalesce(F.col("country"), F.lit("UNKNOWN"))
)

Built-in Spark functions are normally better than Python UDFs. Spark understands built-in expressions and can optimise them. A Python UDF adds serialization between the JVM and Python process and also hides the logic from Spark’s optimiser. I use a UDF only when the operation cannot reasonably be expressed using functions from pyspark.sql.functions.

Be careful about adding hundreds of columns in a loop with repeated withColumn calls. For a large mapping, a single select with a list of expressions is easier to read and can avoid an unnecessarily large query plan.

3. Filter rows as early as possible

Use filter or where to remove records not needed by later stages. Both methods produce the same result.

1
2
3
4
valid_orders = orders.filter(
    (F.col("status") == "COMPLETE") &
    (F.col("amount") > 0)
)

When reading Parquet or Delta files, a filter can be pushed down to the storage layer. Spark may then avoid reading irrelevant row groups or partitions. Filtering on a partition column such as order_date is much better than reading five years of data and filtering after applying a UDF.

Always use parentheses around each condition when combining PySpark column expressions with & or |. Python’s normal and and or operators do not work for Spark columns.

4. Aggregate using groupBy

Most ETL pipelines need totals, counts, or the latest date for a business key.

1
2
3
4
5
6
7
8
9
country_summary = (
    orders.filter(F.col("status") == "COMPLETE")
    .groupBy("country")
    .agg(
        F.countDistinct("order_id").alias("order_count"),
        F.sum("amount").alias("total_amount"),
        F.avg("amount").alias("average_amount")
    )
)

groupBy causes a shuffle because all records for a country must reach the same partition. A common issue in production is data skew. If one key has millions of records while other keys have only thousands, one task can run much longer. Check the Spark UI before increasing the cluster size. Filtering invalid hot keys, enabling adaptive query execution, or salting a genuinely skewed key may be more useful.

5. Join data sets carefully

Let us enrich the orders using a small customer data set.

1
2
3
4
5
6
7
8
9
10
11
customers = spark.createDataFrame([
    ("C001", "Retail"),
    ("C002", "Business"),
    ("C003", "Retail")
], ["customer_id", "segment"])

enriched = orders.join(
    F.broadcast(customers),
    on="customer_id",
    how="left"
)

Broadcasting sends the small table to every executor and avoids shuffling the large table. It works well for a genuinely small dimension table, but forcing a broadcast for a table that no longer fits in executor memory can fail the job. In production I first check its size and usually let Spark’s automatic broadcast threshold decide unless I have a reason to override it.

Make the join type explicit. An inner join can silently remove orders without a matching customer. A left join keeps them, which is often safer for enrichment, but we should still measure unmatched records and decide how to handle them.

6. Remove duplicates using business keys

dropDuplicates is useful when duplicate records are identical enough that any record can be retained.

1
unique_orders = orders.dropDuplicates(["order_id"])

The important question is what should happen when two rows have the same key but different values. dropDuplicates does not express which row is correct. For incremental ingestion, I normally use a window to retain the record with the latest source timestamp.

1
2
3
4
5
6
7
8
9
10
11
12
13
from pyspark.sql.window import Window

latest_first = Window.partitionBy("order_id").orderBy(
    F.col("updated_at").desc(),
    F.col("ingested_at").desc()
)

deduplicated = (
    incoming_orders
    .withColumn("row_number", F.row_number().over(latest_first))
    .filter(F.col("row_number") == 1)
    .drop("row_number")
)

The second ordering column makes the choice deterministic when two updates have the same business timestamp. This small detail prevents inconsistent results across reruns.

7. Use window functions without losing detail

A groupBy returns one row per group. A window function calculates across related rows while keeping every row. For example, we can calculate each customer’s running spend.

1
2
3
4
5
6
7
8
9
10
customer_window = (
    Window.partitionBy("customer_id")
    .orderBy("order_id")
    .rowsBetween(Window.unboundedPreceding, Window.currentRow)
)

with_running_total = orders.withColumn(
    "running_amount",
    F.sum("amount").over(customer_window)
)

Windows require Spark to partition and sort the records, so they are not free. Keep the partition key meaningful and avoid a window without partitionBy on a large data set, as that can move all records into one logical group.

8. Repartition and coalesce only with a reason

repartition performs a full shuffle and can increase or decrease the number of partitions. It is useful before a large operation when records need better distribution, or before writing partitioned output.

1
2
prepared = orders.repartition("country")
prepared.write.partitionBy("country").mode("overwrite").parquet(output_path)

coalesce normally reduces partitions without a full shuffle. It can help reduce small output files, but coalesce(1) is not a production solution. It pushes the final write through one task and becomes slow as volume increases. A better production approach is to target a sensible file size and compact small files as part of table maintenance.

Check the execution plan

Before tuning a Spark transformation based on guesswork, inspect its plan.

1
enriched.explain("formatted")

Look for unexpected exchanges, sort operations, Cartesian products, and filters that were not pushed down. Then use the Spark UI to check task duration, shuffle read and write, spill, and skew. cache should also be used only when the same expensive DataFrame is reused. Caching every intermediate result can consume memory and make the pipeline slower.

For a simple demo, these transformations can run against in-memory records. In a production pipeline I would also enforce an input schema, validate null and duplicate counts, write rejected rows separately, add data quality metrics, and test the result using representative volumes. I would also make the write idempotent so that retrying a failed job does not duplicate data.

Conclusion

Most Spark pipelines are built from a small set of transformations: select, filter, aggregate, join, deduplicate, and window. The syntax is not the difficult part. The practical work is understanding where data is shuffled, reducing it early, and making business rules such as join and deduplication behaviour explicit. Start with built-in functions, inspect the execution plan, and use the Spark UI before making performance changes.

This post is licensed under CC BY 4.0 by the author.