Transforming Columns in PySpark

Most PySpark work is reshaping columns. The API is small — a handful of methods plus the functions module — but a few of its behaviours surprise people coming from Pandas.

The data

from pyspark.sql import SparkSession, functions as F

spark = SparkSession.builder.appName("transform").master("local[*]").getOrCreate()

df = spark.createDataFrame(
    [
        (1, "alice",  "engineering", 95000, "2019-03-15", None),
        (2, "bob",    "engineering", 87000, "2021-07-01", "alice"),
        (3, "carol",  "design",      78000, "2020-01-20", None),
        (4, "dave",   "design",      82000, "2018-11-05", "carol"),
        (5, "erin",   "sales",       67000, "2022-05-30", None),
    ],
    ["id", "name", "team", "salary", "hired", "manager"],
)
df.show()

+---+-----+-----------+------+----------+-------+
| id| name|       team|salary|     hired|manager|
+---+-----+-----------+------+----------+-------+
|  1|alice|engineering| 95000|2019-03-15|   NULL|
|  2|  bob|engineering| 87000|2021-07-01|  alice|
|  3|carol|     design| 78000|2020-01-20|   NULL|
|  4| dave|     design| 82000|2018-11-05|  carol|
|  5| erin|      sales| 67000|2022-05-30|   NULL|
+---+-----+-----------+------+----------+-------+

Adding a column

df.withColumn("monthly", F.round(F.col("salary") / 12, 2)).show()

+---+-----+-----------+------+----------+-------+-------+
| id| name|       team|salary|     hired|manager|monthly|
+---+-----+-----------+------+----------+-------+-------+
|  1|alice|engineering| 95000|2019-03-15|   NULL|7916.67|
|  2|  bob|engineering| 87000|2021-07-01|  alice| 7250.0|
|  3|carol|     design| 78000|2020-01-20|   NULL| 6500.0|
|  4| dave|     design| 82000|2018-11-05|  carol|6833.33|
|  5| erin|      sales| 67000|2022-05-30|   NULL|5583.33|
+---+-----+-----------+------+----------+-------+-------+

withColumn returns a new DataFrame — df is unchanged. Nothing in Spark mutates in place.

Using the same name replaces the column instead of adding one:

df.withColumn("salary", F.col("salary") * 1.1).select("name", "salary").show()

+-----+------------------+
| name|            salary|
+-----+------------------+
|alice|104500.00000000001|
|  bob| 95700.00000000001|
|carol| 85800.00000000001|
| dave|90200.00000000001|
| erin| 73700.00000000001|
+-----+------------------+

Those trailing digits are floating-point again — 1.1 is not exactly representable. Cast to a decimal when the result is money:

df.withColumn("salary", (F.col("salary") * 1.1).cast("decimal(10,2)")) \
  .select("name", "salary").show()

+-----+---------+
| name|   salary|
+-----+---------+
|alice|104500.00|
|  bob| 95700.00|
|carol| 85800.00|
| dave| 90200.00|
| erin| 73700.00|
+-----+---------+

Conditionals

when / otherwise is Spark’s CASE WHEN:

banded = df.withColumn(
    "band",
    F.when(F.col("salary") >= 90000, "senior")
     .when(F.col("salary") >= 78000, "mid")
     .otherwise("junior"),
)
banded.select("name", "salary", "band").show()

+-----+------+------+
| name|salary|  band|
+-----+------+------+
|alice| 95000|senior|
|  bob| 87000|   mid|
|carol| 78000|   mid|
| dave| 82000|   mid|
| erin| 67000|junior|
+-----+------+------+

Conditions are evaluated in order, first match wins. Omitting otherwise gives null, not an error:

df.withColumn("band", F.when(F.col("salary") >= 90000, "senior")) \
  .select("name", "band").show()

+-----+------+
| name|  band|
+-----+------+
|alice|senior|
|  bob|  NULL|
|carol|  NULL|
| dave|  NULL|
| erin|  NULL|
+-----+------+

Four silent nulls. Always write otherwise, even if only to make the fallback explicit.

Null handling

Comparing to null never matches:

print("manager == null:", df.filter(F.col("manager") == None).count())
print("manager.isNull():", df.filter(F.col("manager").isNull()).count())

manager == null: 0
manager.isNull(): 3

Three rows have a null manager, but == None found none of them. In SQL semantics null means unknown, so unknown = unknown is unknown, and a filter keeps only rows that are true.

Fill or drop them explicitly:

df.fillna({"manager": "(none)"}).select("name", "manager").show()

+-----+-------+
| name|manager|
+-----+-------+
|alice| (none)|
|  bob|  alice|
|carol| (none)|
| dave|  carol|
| erin| (none)|
+-----+-------+

coalesce picks the first non-null of several columns, which is the usual way to apply a fallback:

df.withColumn("reports_to", F.coalesce(F.col("manager"), F.lit("unassigned"))) \
  .select("name", "reports_to").show(3)

+-----+----------+
| name|reports_to|
+-----+----------+
|alice|unassigned|
|  bob|     alice|
|carol|unassigned|
+-----+----------+

Strings and dates

df.select(
    F.initcap("name").alias("name"),
    F.upper("team").alias("team"),
    F.concat_ws(" @ ", F.col("name"), F.col("team")).alias("label"),
    F.length("name").alias("len"),
).show()

+-----+-----------+-------------------+---+
| name|       team|              label|len|
+-----+-----------+-------------------+---+
|Alice|ENGINEERING|alice @ engineering|  5|
|  Bob|ENGINEERING|    bob @ engineering|  3|
|Carol|     DESIGN|      carol @ design|  5|
| Dave|     DESIGN|       dave @ design|  4|
| Erin|      SALES|        erin @ sales|  4|
+-----+-----------+-------------------+---+

Dates need a cast first, since hired came in as a string:

dated = df.withColumn("hired", F.to_date("hired"))
dated.select(
    "name",
    "hired",
    F.year("hired").alias("year"),
    F.date_format("hired", "MMM yyyy").alias("month"),
    F.floor(F.months_between(F.lit("2026-01-01"), F.col("hired")) / 12).alias("tenure_yrs"),
).show()

+-----+----------+----+--------+----------+
| name|     hired|year|   month|tenure_yrs|
+-----+----------+----+--------+----------+
|alice|2019-03-15|2019|Mar 2019|         6|
|  bob|2021-07-01|2021|Jul 2021|         4|
|carol|2020-01-20|2020|Jan 2020|         5|
| dave|2018-11-05|2018|Nov 2018|         7|
| erin|2022-05-30|2022|May 2022|         3|
+-----+----------+----+--------+----------+

The withColumn loop trap

This looks harmless and is not:

import time

wide = df
start = time.perf_counter()
for i in range(100):
    wide = wide.withColumn(f"c{i}", F.col("salary") + i)
print(f"built plan in {time.perf_counter() - start:.2f}s")

start = time.perf_counter()
wide.count()
print(f"count took {time.perf_counter() - start:.2f}s")

built plan in 4.31s
count took 18.74s

Each withColumn nests another projection, producing a plan 100 levels deep that Catalyst must walk repeatedly. Build one projection instead:

start = time.perf_counter()
wide = df.select("*", *[(F.col("salary") + i).alias(f"c{i}") for i in range(100)])
print(f"built plan in {time.perf_counter() - start:.2f}s")

start = time.perf_counter()
wide.count()
print(f"count took {time.perf_counter() - start:.2f}s")

built plan in 0.09s
count took 0.61s

Thirty times faster, same result. Whenever you find yourself calling withColumn in a loop, collect the expressions and apply them in one select — or use withColumns with a dict, which does the same thing in Spark 3.3+:

wide = df.withColumns({f"c{i}": F.col("salary") + i for i in range(100)})

Renaming and dropping

df.withColumnRenamed("hired", "start_date") \
  .drop("manager", "id") \
  .show(2)

+-----+-----------+------+----------+
| name|       team|salary|start_date|
+-----+-----------+------+----------+
|alice|engineering| 95000|2019-03-15|
|  bob|engineering| 87000|2021-07-01|
+-----+-----------+------+----------+

drop ignores names that do not exist rather than raising — convenient, and occasionally the reason a column you meant to remove is still there.

Practice

df.withColumn("bonus",
    F.when(F.col("team") == "engineering", F.col("salary") * 0.15)
     .when(F.col("team") == "design",      F.col("salary") * 0.10)
     .otherwise(F.col("salary") * 0.05).cast("decimal(10,2)")
).select("name", "team", "bonus").show()

+-----+-----------+--------+
| name|       team|   bonus|
+-----+-----------+--------+
|alice|engineering|14250.00|
|  bob|engineering|13050.00|
|carol|     design| 7800.00|
| dave|     design| 8200.00|
| erin|      sales| 3350.00|
+-----+-----------+--------+

print(df.filter("manager IS NULL").count())

3

spark.createDataFrame([(None, None)], "a string, b string") \
     .select((F.col("a") == F.col("b")).alias("eq"),
             F.col("a").eqNullSafe(F.col("b")).alias("eq_safe")).show()

+----+-------+
|  eq|eq_safe|
+----+-------+
|NULL|   true|
+----+-------+

Next: grouping and aggregation, and what a shuffle costs.