Everything is a question
about the shuffle.

The mental model that makes the rest of this document predictable. Get this right and tuning stops being guesswork.

The cast

Narrow vs wide, concretely

# Local / standalone
from pyspark.sql import SparkSession
from pyspark.sql import functions as F, types as T
from pyspark.sql.window import Window

spark = (SparkSession.builder
    .appName("claims_load")
    .config("spark.sql.shuffle.partitions", "200")
    .config("spark.sql.session.timeZone", "UTC")
    .config("spark.sql.execution.arrow.pyspark.enabled", "true")
    .getOrCreate())

spark.sparkContext.setLogLevel("WARN")

# Spark Connect (3.4+) — thin client, no local JVM
spark = SparkSession.builder.remote("sc://host:15002").getOrCreate()

# In Databricks / Glue the session already exists — never build a new one
#   Databricks: `spark` is injected
#   Glue:       glueContext.spark_session

df = (spark.read
    .format("parquet")
    .option("mergeSchema", "false")
    .load("s3://bucket/claims/"))

# CSV — always give it a schema in production
df = (spark.read
    .schema(schema)                       # skip inference: one fewer full scan
    .option("header", True)
    .option("mode", "PERMISSIVE")      # PERMISSIVE | DROPMALFORMED | FAILFAST
    .option("columnNameOfCorruptRecord", "_corrupt")
    .csv("s3://bucket/raw/*.csv"))

# JDBC — parallelise or you get one task and a very long wait
df = (spark.read.format("jdbc")
    .option("url", url).option("dbtable", "claims")
    .option("partitionColumn", "claim_id")
    .option("lowerBound", 1).option("upperBound", 10_000_000)
    .option("numPartitions", 32)
    .option("fetchsize", 10_000)
    .load())

# DDL string — shortest readable form
schema = "claim_id BIGINT, policy_no STRING, paid DECIMAL(18,2), ts TIMESTAMP"

# StructType — when you need nullability or nesting explicit
schema = T.StructType([
    T.StructField("claim_id", T.LongType(), nullable=False),
    T.StructField("paid",     T.DecimalType(18, 2)),
    T.StructField("party",    T.StructType([
        T.StructField("name", T.StringType()),
        T.StructField("tags", T.ArrayType(T.StringType())),
    ])),
])

df.printSchema()
df.schema.json()                      # persist and reuse across jobs
T.StructType.fromJson(json.loads(s))  # round-trip back

# Good: one projection, flat plan
exprs = [F.col(c).cast("decimal(18,2)").alias(c) if c in money else F.col(c)
         for c in df.columns]
df = df.select(*exprs)

# Composable steps
def add_audit(d):
    return d.withColumns({"loaded_at": F.current_timestamp(),
                          "src_file":  F.input_file_name()})

df = raw.transform(add_audit).transform(dedupe).transform(conform)

All narrow. All executed in the JVM with code generation — which is exactly why you reach here before reaching for a UDF (§11).

Conditionals & nulls

Strings

Dates & timestamps

Numbers & misc

agg = (df.groupBy("policy_no", F.date_trunc("month", "ts").alias("period"))
    .agg(
        F.sum("paid").alias("paid_total"),
        F.countDistinct("claim_id").alias("claims"),
        F.max("ts").alias("last_event"),
        F.collect_set("peril").alias("perils"),
        F.sum(F.when(F.col("status") == "OPEN", F.col("paid"))).alias("open_paid"),
    ))

# Column list form — no duplicate key columns in the output
j = fact.join(dim, on=["policy_key"], how="left")

# Expression form — full control, but both key columns survive
j = fact.alias("f").join(dim.alias("d"),
        (F.col("f.policy_key") == F.col("d.policy_key")) &
        (F.col("f.event_ts").between(F.col("d.valid_from"), F.col("d.valid_to"))),
        how="left")                       # SCD2 point-in-time lookup

# Force a broadcast when you know the side is small
j = fact.join(F.broadcast(dim), "policy_key", "left")

Join types

Physical strategies — what the planner actually picks

Hints

w = Window.partitionBy("policy_no").orderBy(F.col("ts").desc())

latest = (df
    .withColumn("rn", F.row_number().over(w))
    .filter("rn = 1").drop("rn"))        # deterministic dedupe

# Running total — explicit frame
wr = (Window.partitionBy("policy_no").orderBy("ts")
        .rowsBetween(Window.unboundedPreceding, Window.currentRow))
df = df.withColumn("cum_paid", F.sum("paid").over(wr))

# SCD2 close-out: next row's start becomes this row's end
ws = Window.partitionBy("nk").orderBy("valid_from")
df = df.withColumn("valid_to",
        F.coalesce(F.lead("valid_from").over(ws) - F.expr("INTERVAL 1 SECOND"),
                   F.lit("9999-12-31").cast("timestamp")))

# Flatten a struct
df.select("party.*")
df.select(F.col("party.name").alias("party_name"))

# Array to rows
df.select("claim_id", F.explode("perils").alias("peril"))
df.select("claim_id", F.explode_outer("perils").alias("peril"))   # keeps empty/null
df.select("claim_id", F.posexplode("perils").alias("pos", "peril"))

# JSON string column to struct
df.withColumn("payload", F.from_json("raw", schema))
df.withColumn("raw", F.to_json("payload"))
spark.read.json(df.rdd.map(lambda r: r.raw))   # infer schema from a sample

# Unpivot: wide to long (3.4+ native)
long = df.unpivot(["claim_id"], ["jan", "feb", "mar"], "month", "amount")
# pre-3.4 equivalent
long = df.select("claim_id", F.expr(
    "stack(3, 'jan', jan, 'feb', feb, 'mar', mar) as (month, amount)"))

Ordered from cheapest to most expensive. Exhaust the option above before reaching for the one below.

# Vectorised: Series -> Series
from pyspark.sql.functions import pandas_udf

@pandas_udf("double")
def fx(amount: pd.Series, rate: pd.Series) -> pd.Series:
    return amount * rate

df.withColumn("gbp", fx("amount", "rate"))

# Iterator variant — set up an expensive resource once per batch stream
from typing import Iterator

@pandas_udf("string")
def classify(it: Iterator[pd.Series]) -> Iterator[pd.Series]:
    model = load_model()          # once per executor task, not per row
    for s in it:
        yield pd.Series(model.predict(s))

# Grouped map: entire group as a pandas DataFrame
def reconcile(pdf: pd.DataFrame) -> pd.DataFrame:
    pdf["variance"] = pdf["actual"] - pdf["budget"].cumsum()
    return pdf

df.groupBy("cost_centre").applyInPandas(reconcile, schema="cost_centre string, variance double")

# Plain UDF — last resort. Always declare the return type.
from pyspark.sql.functions import udf

@udf(T.StringType())
def normalise(s):
    return None if s is None else s.strip().upper()

df.createOrReplaceTempView("claims")          # session-scoped
df.createOrReplaceGlobalTempView("claims")    # cross-session: global_temp.claims

out = spark.sql("""
    SELECT policy_no, SUM(paid) AS paid
    FROM claims
    WHERE ts >= :cutoff
    GROUP BY policy_no
""", cutoff="2026-01-01")                    # parameterised (3.4+) — no string concat

spark.catalog.listDatabases()
spark.catalog.listTables("analytics")
spark.catalog.tableExists("analytics.dim_policy")
spark.catalog.refreshTable("analytics.fct_claim")   # after external writes
spark.catalog.clearCache()

(df.write
   .format("parquet")
   .mode("overwrite")                      # append | overwrite | ignore | error
   .partitionBy("event_date")
   .option("compression", "snappy")
   .save("s3://bucket/curated/claims/"))

# Replace only the partitions present in this batch, leave the rest alone
spark.conf.set("spark.sql.sources.partitionOverwriteMode", "dynamic")

# Control file count before writing
df.repartition("event_date").write.partitionBy("event_date").save(path)

# Managed / catalogued table
df.writeTo("catalog.schema.fct_claim").partitionedBy("event_date").createOrReplace()

Caching

Diagnose in this order

Skew: salting

# One key holds 60% of the rows; that key's task decides your runtime.
N = 32
fact_s = fact.withColumn("salt", (F.rand() * N).cast("int"))
dim_s  = (dim
    .withColumn("salt", F.explode(F.sequence(F.lit(0), F.lit(N - 1))))) # fan out

joined = (fact_s.join(dim_s, ["policy_key", "salt"], "left")
                .drop("salt"))

Key configuration

Same DataFrame API, incremental execution. Micro-batches by default.

src = (spark.readStream
    .format("cloudFiles")                       # Databricks Auto Loader
    .option("cloudFiles.format", "json")
    .schema(schema)
    .load("s3://bucket/landing/"))

agg = (src
    .withWatermark("event_ts", "10 minutes")      # bounds state; drops late data
    .groupBy(F.window("event_ts", "5 minutes"), "policy_no")
    .agg(F.sum("paid").alias("paid")))

q = (agg.writeStream
    .format("delta")
    .outputMode("append")
    .option("checkpointLocation", "s3://bucket/_ckpt/claims_agg")
    .trigger(availableNow=True)                 # catch up, then stop
    .start("s3://bucket/curated/claims_agg"))

q.awaitTermination()

from delta.tables import DeltaTable

# SCD1 upsert
tgt = DeltaTable.forPath(spark, "s3://bucket/silver/dim_policy")
(tgt.alias("t")
   .merge(updates.alias("s"), "t.policy_key = s.policy_key")
   .whenMatchedUpdateAll(condition="t.hash_key <> s.hash_key")   # skip no-op rows
   .whenNotMatchedInsertAll()
   .whenNotMatchedBySourceUpdate(set={"is_deleted": "true"})       # soft delete
   .execute())

# Time travel
spark.read.format("delta").option("versionAsOf", 42).load(path)
spark.read.format("delta").option("timestampAsOf", "2026-08-01").load(path)

# Change Data Feed — downstream incremental consumption
(spark.read.format("delta")
   .option("readChangeFeed", "true")
   .option("startingVersion", 42).load(path))   # adds _change_type, _commit_version

# Built-in test helpers (3.5+)
from pyspark.testing import assertDataFrameEqual, assertSchemaEqual

assertDataFrameEqual(actual, expected, checkRowOrder=False, rtol=1e-6)
assertSchemaEqual(actual.schema, expected.schema)

# Session fixture — reuse one session for the whole suite
@pytest.fixture(scope="session")
def spark():
    s = (SparkSession.builder.master("local[2]")
         .config("spark.sql.shuffle.partitions", "2")   # 200 is absurd on 10 rows
         .config("spark.sql.session.timeZone", "UTC")
         .getOrCreate())
    yield s
    s.stop()

# Inline contracts on the pipeline itself
df.observe("dq", F.count(F.lit(1)).alias("rows"),
                 F.sum(F.col("paid").isNull().cast("int")).alias("null_paid"))

Errors and what they usually mean

You should rarely need this layer — no Catalyst, no Tungsten, no code generation. Worth recognising for legacy code and for the few things the DataFrame API cannot express.

What is new and worth knowing