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
| Transformation | Typical use | Shuffle likely? |
|---|---|---|
select / withColumn | Choose or derive columns | No |
filter | Remove unwanted rows | No |
groupBy | Aggregate records | Yes |
join | Combine data sets | Usually |
dropDuplicates | Remove duplicate keys | Yes |
| Window functions | Rank or compare rows in a group | Yes |
repartition | Redistribute data | Yes |
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.
