PySpark Jobs skill for Codex

Done well means: the job uses built-in column expressions end to end, the plan has no surprise shuffles, joins pick the right strategy, skew is handled, output files land at a sane size, and the job can rerun for the same partition without duplicating data.

Before changing code

1. Read AGENTS.md, the job’s entry point and its submit config (spark-submit args, Databricks job JSON, EMR step, spark-defaults.conf). Note Spark version: spark.version. Spark 4.0 turns ANSI SQL mode on by default: invalid casts and overflow now raise instead of returning null.
2. Get evidence before tuning. From the Spark UI (or history server) for the slow run, record: longest stage, shuffle read/write size, spill (memory and disk), task duration max vs median, and number of tasks. Ask the user for these if Codex cannot reach the UI (it is usually off the sandbox network).
3. Print the plan: df.explain(mode="formatted"). Look for Exchange (shuffle), BroadcastHashJoin vs SortMergeJoin, BatchEvalPython / ArrowEvalPython (Python UDFs), and PushedFilters in the scan.

Job skeleton

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

spark = (
    SparkSession.builder.appName("orders_daily")
    .config("spark.sql.adaptive.enabled", "true")               # default on since 3.2
    .config("spark.sql.sources.partitionOverwriteMode", "dynamic")
    .getOrCreate()
)

def transform(orders, customers, run_date: str):
    latest = Window.partitionBy("order_id").orderBy(F.col("updated_at").desc())
    return (
        orders
        .where(F.col("order_date") == F.lit(run_date))          # partition pruning
        .withColumn("rn", F.row_number().over(latest))
        .where("rn = 1").drop("rn")
        .join(F.broadcast(customers.select("customer_id", "country")), "customer_id", "left")
        .withColumn("tier", F.when(F.col("amount") >= 1000, "high")
                              .when(F.col("amount") >= 100, "mid").otherwise("low"))
    )

if __name__ == "__main__":
    import sys
    run_date = sys.argv[1]
    orders = spark.read.parquet("s3://lake/orders/")
    customers = spark.read.parquet("s3://lake/customers/")
    (transform(orders, customers, run_date)
        .repartition("order_date")
        .write.mode("overwrite")                                 # dynamic: replaces only run_date
        .partitionBy("order_date")
        .parquet("s3://lake/orders_enriched/"))

Keep transform a pure function of DataFrames so it can be unit-tested; I/O stays at the edges.

Joins

Situation	Strategy
One side fits comfortably in executor memory (tens to a few hundred MB)	F.broadcast(small); or raise spark.sql.autoBroadcastJoinThreshold (default 10 MB)
Both sides large	Sort-merge join; make sure join keys have the same type on both sides
Repeated joins on the same key across jobs	Bucketed tables or a pre-partitioned layout (Delta/Iceberg clustering)
Need only matching or non-matching rows	how="left_semi" / "left_anti" instead of join + filter
Key with heavy hot values	See skew below

Null join keys never match; a large null-key population still lands in one partition during a shuffle. Filter or salt nulls explicitly.

Shuffle and partition sizing

• Target 100 to 200 MB per shuffle partition. With AQE on, set spark.sql.adaptive.coalescePartitions.enabled=true and leave spark.sql.shuffle.partitions high enough (for example total shuffle bytes / 128 MB); AQE merges small ones.
• repartition(n) or repartition(col) shuffles; coalesce(n) only merges, cheaply, but can leave the final stage with too few tasks.
• Cache (df.cache() / persist) only a DataFrame reused by two or more actions, and unpersist() when done. Caching everything just fills executor memory.

Skew

Symptom: one or a few tasks take 10x the median and shuffle read far more bytes.

1. Confirm the hot keys: df.groupBy("customer_id").count().orderBy(F.desc("count")).show(20).
2. Let AQE split skewed partitions first: spark.sql.adaptive.skewJoin.enabled=true (default on with AQE). Works for sort-merge joins.
3. Still skewed: broadcast the other side if it fits, or salt the hot keys:

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

4. Skewed aggregations: aggregate on (key, salt) first, then on key.

Python code on the hot path

Pattern	Replace with
@udf doing string, date or arithmetic logic	Built-ins in pyspark.sql.functions
Row-wise Python logic with no built-in equivalent	@pandas_udf (vectorised via Arrow), or useArrow=True on udf
df.collect() / toPandas() on big data	Keep it distributed; collect only aggregates
Loops issuing many small Spark actions	One query with a join or groupBy

Writing output

• Partition by low-cardinality columns that readers filter on (date, region). Partitioning by a high-cardinality ID creates millions of directories.
• Too many small files: repartition by the partition column before write, or set spark.sql.files.maxRecordsPerFile. Aim for 128 MB to 1 GB files.
• Idempotent reruns: dynamic partition overwrite (above) or table formats with MERGE/ replaceWhere. Plain mode("append") duplicates data on retry.
• Overwriting a whole table path with static overwrite mode deletes every partition. Confirm with the user before running any job in overwrite mode against production paths.

Memory errors

Error	Likely cause	First fix
Executor OOM / container killed by YARN or K8s	Skewed partition, huge broadcast, Python worker memory	Fix skew; lower broadcast; raise spark.executor.memoryOverhead for PySpark
Driver OOM	collect, toPandas, huge broadcast built on driver	Remove collect; reduce broadcast size
Heavy spill	Too few shuffle partitions	More partitions or let AQE coalesce from a higher number

Testing

import pytest
from pyspark.sql import Row, SparkSession
from pyspark.testing import assertDataFrameEqual   # Spark 3.5+

@pytest.fixture(scope="session")
def spark():
    return (SparkSession.builder.master("local[2]")
            .config("spark.sql.shuffle.partitions", "2").getOrCreate())

def test_latest_version_wins(spark):
    orders = spark.createDataFrame(
        [("o1", "c1", 50.0, "2025-09-01", "2025-09-01 10:00"),
         ("o1", "c1", 150.0, "2025-09-01", "2025-09-01 12:00")],
        "order_id string, customer_id string, amount double, order_date string, updated_at string")
    customers = spark.createDataFrame([("c1", "DE")], "customer_id string, country string")
    out = transform(orders, customers, "2025-09-01").select("order_id", "amount", "tier")
    assertDataFrameEqual(out, [Row(order_id="o1", amount=150.0, tier="mid")])

Local Spark needs a JVM; check java -version before assuming tests can run in the sandbox.

Verification and report

• Compare row counts and key totals between old and new output for the same partition.
• Re-check the plan and the Spark UI: target stage time, spill, max/median task duration.
• Rerun the same date and confirm the partition is replaced, not doubled.

**Job:** orders_daily (Spark 3.5, 20 x r6i.2xlarge)
**Problem:** stage 7 sort-merge join, 1 task 41 min vs median 40 s (customer c-000 = 18% of rows)
**Changes:** broadcast customers (62 MB), dynamic partition overwrite, removed Python UDF for tier
**Result:** 58 min -> 9 min; shuffle 410 GB -> 12 GB; output 96 files of ~300 MB (was 14,000)
**Parity:** row count and revenue identical for 2025-09-01