Case study · Spark Optimization & Capacity Planning Platform

Making a Spark job 44× faster

How I built a toolkit that finds why PySpark jobs are slow, sizes clusters for a deadline and forecasts storage, then used it to speed up a real job.

1. The problem

Spark jobs are often slow because of small coding choices, not because the cluster is too small. Data teams keep asking three questions: why is this job slow, what resources does it need, and when will storage run out? Each usually needs an experienced engineer to answer.

2. My approach

I built one Python toolkit with three tools that share the same core code, available as a command line, a FastAPI service and a browser demo:

• Code analyser: reads PySpark with Python's AST and checks Spark SQL for 25 slow patterns, each with a line number, reason and fix.
• Resource optimizer: applies standard executor sizing rules, sizes shuffle partitions and checks the runtime against an SLA.
• Capacity planner: forecasts data growth and the date storage fills up.
• AI layer: sends the findings to Claude for a plain-language explanation and rewritten code, with a rule-based fallback.

3. The experiment

To test it, I wrote a typical first version of a sales report job. The analyser gave it a health score of 43/100 and found 8 issues.

@udf(returnType=DoubleType())
def amount_with_tax(amount, rate):
    return float(amount) * (1 + float(rate))

enriched = orders.join(countries, "country_id")
enriched = enriched.withColumn("total",
    amount_with_tax("amount", "tax_rate"))
top10 = enriched.orderBy(F.desc("total")) \
    .collect()[:10]

enriched = (orders
    .join(F.broadcast(countries), "country_id")
    .withColumn("total",
        F.col("amount") * (1 + F.col("tax_rate")))
    .cache())

top10 = enriched.orderBy(F.desc("total")) \
    .limit(10).collect()

4. Results

Both versions ran on the same 500,000 orders with PySpark 4.2 in local mode (median of 3 runs). The results were identical.

Most of the gain came from removing the Python UDF and the full collect().

5. Key decisions

• AST instead of text search. Parsing the real code structure means a comment or string that mentions collect() never causes a false alarm.
• Model chosen on held-out data. The capacity planner hides the last 20% of the history and keeps whichever growth model predicts it better, instead of assuming one.
• AI as an extra, not a dependency. With no API key or a failed call, the tool falls back to rule-based explanations, so it always works.
• Honest estimates. The runtime estimate is a simple throughput model, and the docs say it should be calibrated with real job history.

6. Limitations and next steps

The benchmark ran on one machine, not a production cluster, and static analysis cannot see data size or skew. Next, I want to read Spark event logs to find skewed stages and spill, learn throughput numbers from past jobs, and test on a multi-node cluster.

Python · PySpark · NumPy · FastAPI · Pydantic · Claude API · pytest · GitHub Actions · Docker