Home / Databricks on Azure

Databricks on AzureHow-toIntermediate

Transformation: single-node pandas jobs to PySpark (or pandas API on Spark)

Fix jobs that run pandas on the driver and die with out-of-memory errors as data grows. Recognise driver-bound code, convert common pandas operations to PySpark, use pandas API on Spark or applyInPandas where it helps, and validate results.

Deepak Bhardwaj2 min read

Applies toAzure Databricks Runtime with PySpark, pandas API on Spark (pyspark.pandas), notebooks and Python jobs

Symptoms of a pandas-on-the-driver job

  • Cluster has 8 workers but only the driver's CPU is busy.
  • toPandas() or pd.read_csv() on large files; driver OOM / "driver is temporarily unavailable".
  • Runtime grows linearly with data and adding workers doesn't help.

1. Find the hot spots

Search code for toPandas(), pd.read_, .apply( loops, iterrows() and collect(). In the Spark UI, if a job shows almost no Spark stages, the work is happening in plain Python on the driver.

2. Convert common operations

pandas PySpark
pd.read_csv(path) spark.read.option("header", True).csv(path)
df[df.amount > 0] df.filter(F.col("amount") > 0)
df["x"] = df.a * 2 df.withColumn("x", F.col("a") * 2)
df.groupby("k").agg({"v": "sum"}) df.groupBy("k").agg(F.sum("v"))
pd.merge(a, b, on="id", how="left") a.join(b, "id", "left")
df.drop_duplicates(["id"]) df.dropDuplicates(["id"])
df.sort_values("ts") df.orderBy("ts") (only if needed — sorting is expensive)
df.apply(func, axis=1) Built-in functions first; else pandas UDF
df.to_csv(...) df.write.saveAsTable(...)

3. Two shortcuts

pandas API on Spark — minimal code changes for pandas-heavy code:

import pyspark.pandas as ps
pdf = ps.read_table("main.sales.orders")
out = pdf[pdf.amount > 0].groupby("store_id").amount.sum()

Not every pandas behaviour is identical (index handling, ordering); test outputs.

Per-group pandas logic (e.g. a model or complex calc per store):

def per_store(pdf):          # plain pandas inside
    pdf["z"] = (pdf.amount - pdf.amount.mean()) / pdf.amount.std()
    return pdf

result = df.groupBy("store_id").applyInPandas(per_store, schema="store_id int, amount double, z double")

Spark runs it in parallel across workers, one group at a time per task.

4. Watch for

  • Row-order assumptions: Spark DataFrames have no inherent order.
  • Index-based logic (.loc, .iloc): rewrite with keys and window functions.
  • Small lookups: broadcast joins (F.broadcast(small_df)) instead of dict lookups in loops.
  • Leftover toPandas() for final small outputs is fine — just after aggregation, not before.

5. Validate

Run old and new on the same input; compare counts and aggregates with a tolerance for floating-point differences. Then right-size the cluster — the driver no longer needs to be huge.

Tip

If the data truly fits in memory and runs fast, a single-node cluster with pandas (or Polars) is fine. Convert the jobs that are growing, not everything.

Checklist

  • Driver-bound jobs identified
  • Operations converted or wrapped with applyInPandas
  • Ordering/index assumptions removed
  • Outputs reconciled
  • Cluster resized

Related fixes