PySpark Jobs skill for Claude Code

A Spark job is done when it uses built-in column expressions instead of Python row functions, reads only the columns and partitions it needs, shuffles once per logical step, has no task running more than a few times longer than the median, and writes files in the 128 MB-1 GB range to a table format that supports atomic commits.

Version notes

• Spark 4.0 turns ANSI SQL mode on by default (spark.sql.ansi.enabled=true): invalid casts and integer overflow raise errors instead of returning NULL. Code ported from 3.x that relied on silent NULLs must use try_cast / try_* functions or handle bad data explicitly.
• Adaptive Query Execution (AQE) is on by default since 3.2: it coalesces shuffle partitions, converts joins to broadcast at runtime, and splits skewed partitions.
• Spark Connect clients (pyspark-connect, Databricks serverless) do not expose sparkContext or RDD APIs; stay on the DataFrame API to remain portable.

Workflow

- [ ] 1. Read the job: inputs, sizes, partitioning, outputs, schedule
- [ ] 2. Check the plan: df.explain("formatted") for scans, filters pushed, exchanges
- [ ] 3. Run on a representative slice; open the Spark UI
- [ ] 4. Find the costly stage: shuffle size, spill, skewed tasks, GC
- [ ] 5. Fix the biggest cost first (join strategy, skew, partitioning, UDFs)
- [ ] 6. Fix output layout (file size, partition columns, table format)
- [ ] 7. Re-run and compare stage metrics and wall time

Code rules

from pyspark.sql import SparkSession, functions as F, types as T, Window

spark = SparkSession.builder.appName("orders_daily").getOrCreate()

schema = T.StructType([
    T.StructField("order_id", T.StringType(), False),
    T.StructField("customer_id", T.LongType(), True),
    T.StructField("amount", T.DecimalType(12, 2), True),
    T.StructField("created_at", T.TimestampType(), True),
])

orders = (
    spark.read.schema(schema).json("s3://raw/orders/dt=2026-10-19/")
    .select("order_id", "customer_id", "amount", "created_at")       # prune early
    .filter(F.col("amount") > 0)
    .withColumn("order_date", F.to_date("created_at"))
)

w = Window.partitionBy("customer_id").orderBy(F.col("created_at").desc())
latest = orders.withColumn("rn", F.row_number().over(w)).filter("rn = 1").drop("rn")

• Supply schemas for JSON/CSV. Schema inference reads the data twice and guesses types differently per batch.
• Use pyspark.sql.functions for everything they cover. A Python UDF serializes every row to a Python worker; it can be 10-100x slower. If unavoidable, use a vectorized @F.pandas_udf or Arrow-optimized UDFs, not a row UDF.
• Never collect() or toPandas() a large DataFrame onto the driver. Use limit() for samples and write results to storage.
• Chain transformations; avoid loops that call withColumn hundreds of times (plan blowup) - build a list of expressions and use one select, or withColumns({...}) (3.3+).
• Do not use count() for control flow on large data (if df.count() > 0); use df.isEmpty() (3.3+) or df.limit(1).count().

Partitions and files

Situation	Action
Shuffle stages with thousands of tiny tasks	let AQE coalesce (spark.sql.adaptive.coalescePartitions.enabled, default on)
Shuffle tasks spilling to disk, multi-GB per task	raise spark.sql.shuffle.partitions (default 200) or spark.sql.adaptive.advisoryPartitionSizeInBytes
Reading many small files	spark.sql.files.maxPartitionBytes (default 128 MB) groups them; better, compact the source
Output has thousands of tiny files	repartition(n, *partition_cols) before write, or table-format compaction (OPTIMIZE)
Reduce partitions without shuffle	coalesce(n) (can create uneven, slow final tasks; prefer repartition when n is small relative to data)

Target output files of 128 MB-1 GB. Partition tables on low-cardinality columns queried in filters (date), never on high-cardinality ids - each value becomes a directory.

Joins

Case	Strategy
One side under ~10 MB (spark.sql.autoBroadcastJoinThreshold) or AQE detects it	broadcast hash join (automatic)
Small side known but above threshold (up to a few hundred MB)	big.join(F.broadcast(small), "key")
Both large	sort-merge join; ensure keys have the same type on both sides
Repeated joins on the same key across jobs	bucketed or clustered tables on that key

Join hygiene: filter and select before joining, deduplicate the dimension side, and check for null or default keys ("", 0, "unknown") that pile onto one partition.

Skew

Symptom in the Spark UI: one or a few tasks in a stage take far longer (max task time greater than 5-10x the median) and read far more shuffle data.

1. Confirm the hot keys: df.groupBy("key").count().orderBy(F.desc("count")).show(20).
2. If AQE skew join is on (spark.sql.adaptive.skewJoin.enabled, default true), check the plan reports skew handling; tune skewedPartitionFactor / skewedPartitionThresholdInBytes.
3. Null or junk keys: filter them out and process separately.
4. Remaining skew in joins: salt the hot keys.

N = 16
big_salted = big.withColumn("salt", (F.rand(seed=42) * N).cast("int"))
small_salted = small.crossJoin(spark.range(N).withColumnRenamed("id", "salt"))
joined = big_salted.join(small_salted, ["key", "salt"]).drop("salt")

For skewed aggregations, aggregate on (key, salt) first, then on key.

Caching

Cache only a DataFrame reused by two or more actions, and only after the expensive part: df.persist(StorageLevel.MEMORY_AND_DISK), then unpersist() when done. Caching a DataFrame used once wastes memory and adds a materialization step. Check the Storage tab: a “fraction cached” below 100% means it does not fit.

Reading the Spark UI

Where	Look for	Meaning
Jobs / Stages	longest stage, number of tasks	where the time goes
Stage detail: summary metrics	max vs median duration, shuffle read size	skew if max far above median
Stage detail	Spill (memory/disk)	partitions too large; more partitions or more memory per task
Executors	GC time above ~10% of task time	memory pressure; fewer cores per executor or more memory
SQL / DataFrame tab	plan with Exchange nodes, scan “files read”, “number of output rows”	extra shuffles, missing partition pruning

On Databricks, the same metrics are in the Spark UI and query profile; on EMR/Dataproc use the Spark History Server for finished jobs.

Table formats

Write to Delta Lake or Apache Iceberg rather than bare Parquet directories when jobs update data, run concurrently, or need rollback.

# Idempotent daily load: replace only the partition being processed
(daily.write.format("delta").mode("overwrite")
      .option("replaceWhere", "order_date = '2026-10-19'")
      .save("s3://lake/orders"))

# Upsert (Delta)
from delta.tables import DeltaTable
(DeltaTable.forPath(spark, "s3://lake/customers").alias("t")
   .merge(updates.alias("s"), "t.customer_id = s.customer_id")
   .whenMatchedUpdateAll().whenNotMatchedInsertAll().execute())

• With plain Parquet, set spark.sql.sources.partitionOverwriteMode=dynamic before mode("overwrite"), or a partitioned overwrite deletes every other partition.
• Iceberg: df.writeTo("catalog.db.orders").overwritePartitions() and MERGE INTO in SQL.
• Deduplicate the source before MERGE; multiple source rows matching one target row fail the merge.
• Schedule compaction (OPTIMIZE, Iceberg rewrite_data_files) and snapshot expiry (VACUUM, expire_snapshots) - ask before shortening retention, it limits time travel.

Gotchas

• F.rand() without a seed and non-deterministic UDFs make retried tasks produce different data; seed them.
• Timestamps: set spark.sql.session.timeZone explicitly (usually UTC); to_date on a timestamp uses the session time zone.
• dropDuplicates() without columns compares all columns; pass the key columns.
• union matches by position; use unionByName(allowMissingColumns=True).
• Writing then reading the same path in one job (overwrite while reading) can delete the input before it is read; write to a new location or use a table format.
• Local tests: SparkSession.builder.master("local[2]") with small fixtures and spark.sql.shuffle.partitions=2 keeps tests fast.

Verification

• explain("formatted") shows PushedFilters and pruned columns on scans, broadcast where expected, and no unexpected extra Exchange.
• Spark UI after the change: wall time, total shuffle bytes, spill, and max/median task time recorded before and after.
• Output check: row count and a key-level checksum match the previous implementation on the same input; file sizes in the target range.
• Re-running the job for the same date produces the same output (idempotent).