DEV Community

Karan Pulyala
Karan Pulyala

Posted on

PySpark Caching on Azure Databricks: cache() vs persist() vs checkpoint(), the Disk Cache, and When Caching Hurts

"Should I cache this DataFrame?" is one of those questions that sounds simple in an interview and turns out to have a lot of layers. On Azure Databricks there are actually two different caches in play, and cache() is lazy, can silently serve stale data, and sometimes makes a job slower. This post walks through what each one does, when it helps, and how to prove it from the Spark UI.

Two caches, not one

On Databricks you'll run into:

  1. The Spark cache (df.cache(), df.persist(), CACHE TABLE). It stores the result of a query plan in executor memory and/or local disk, in Spark's in-memory format. You control it explicitly.
  2. The Databricks disk cache (it used to be called the Delta cache). It stores copies of remote Parquet/Delta files on the workers' local SSDs, so repeated reads skip cloud storage. It works automatically when enabled, and you don't call anything per DataFrame.

They solve different problems. The Spark cache saves recomputation (joins, aggregations, UDFs). The disk cache saves I/O from ADLS. Interviewers like asking for this distinction, so it's worth being crisp about it.

The Spark cache: lazy, plan-based, and per cluster

from pyspark.sql import functions as F

orders = spark.read.table("sales.orders")
customers = spark.read.table("sales.customers")

enriched = (
    orders.join(customers, "customer_id")
          .withColumn("order_month", F.date_trunc("month", "order_ts"))
          .filter(F.col("status") == "COMPLETE")
)

enriched.cache()      # nothing happens yet
enriched.count()      # this action materialises the cache

by_month   = enriched.groupBy("order_month").agg(F.sum("amount").alias("revenue"))
by_country = enriched.groupBy("country").agg(F.countDistinct("customer_id").alias("buyers"))

by_month.write.mode("overwrite").saveAsTable("reporting.revenue_by_month")
by_country.write.mode("overwrite").saveAsTable("reporting.buyers_by_country")

enriched.unpersist()
Enter fullscreen mode Exit fullscreen mode

A few things to notice:

  • cache() is a transformation-like marker, not an action. Nothing is stored until an action runs. If you call cache() and then only use the DataFrame once, you paid the cost of writing the cache for nothing.
  • The cache is keyed on the logical plan. Any later query whose plan matches (or contains) the cached plan will read from cache. A DataFrame built with the same code in another notebook attached to the same cluster can hit the cache too, because the cache lives in the cluster's executors, not in your Python variable.
  • count() is a common way to force materialisation, but be aware that for some plans Spark can satisfy count() without computing every column. If you need the full thing materialised, run the action you actually care about, or use the noop writer, which computes every column and writes nothing:
# Forces full computation of every column, writes nothing
enriched.write.format("noop").mode("overwrite").save()
Enter fullscreen mode Exit fullscreen mode

cache() vs persist() and storage levels

cache() is just persist() with the default storage level. For DataFrames that default is memory first, spilling to local disk when memory runs out (MEMORY_AND_DISK style). You can choose explicitly:

from pyspark import StorageLevel

big_lookup = spark.read.table("ref.product_attributes")

big_lookup.persist(StorageLevel.DISK_ONLY)       # keep executor memory free
# big_lookup.persist(StorageLevel.MEMORY_ONLY)   # recompute partitions that don't fit
# big_lookup.persist(StorageLevel.MEMORY_AND_DISK_2)  # replicated on two executors

print(big_lookup.storageLevel)
Enter fullscreen mode Exit fullscreen mode

How I think about the choice:

  • MEMORY_AND_DISK (default): fine for most reuse cases.
  • MEMORY_ONLY: partitions that don't fit are recomputed on demand. Good when recomputation is cheap and memory is tight; bad when the lineage is expensive.
  • DISK_ONLY: useful when the result is large and executor memory is needed for shuffles and joins. Local SSD reads are still far cheaper than recomputing a multi-way join.
  • _2 variants: replicate each partition to a second executor. Rarely worth it unless you're on spot instances and lose executors often.

checkpoint() and localCheckpoint(): cutting the lineage

Caching keeps the lineage. If an executor dies, Spark recomputes the lost partitions from the original plan. In iterative jobs (loops that keep transforming the same DataFrame), the plan itself can get enormous, and planning time starts to dominate. That's where checkpointing helps:

spark.sparkContext.setCheckpointDir("abfss://[email protected]/checkpoints/")

df = spark.read.table("graph.nodes")
for i in range(20):
    df = df.withColumn("score", F.col("score") * 0.85 + 0.15).withColumn("iteration", F.lit(i))
    if i % 5 == 4:
        df = df.checkpoint()   # eager by default: writes files and truncates the plan
Enter fullscreen mode Exit fullscreen mode
  • checkpoint() writes the data to reliable storage and returns a DataFrame whose plan is just "read these files". It survives executor loss.
  • localCheckpoint() stores data on executors instead. It's faster, but if an executor is lost the data is gone and the job fails rather than recomputes.
  • Note that checkpoint() returns a new DataFrame. Calling df.checkpoint() without reassigning does the work and throws the result away.

In practice, for a pipeline step that's reused across jobs, writing an intermediate Delta table is often clearer than either: it's inspectable, survives cluster restarts, and you can OPTIMIZE it.

The Databricks disk cache

The disk cache is controlled by a Spark config and is enabled by default on many worker types that have local SSDs. You can check and set it per cluster or per session:

print(spark.conf.get("spark.databricks.io.cache.enabled"))
spark.conf.set("spark.databricks.io.cache.enabled", "true")
Enter fullscreen mode Exit fullscreen mode

It caches the remote data files that queries read, so the second query over the same columns and files is served from local disk. Because it caches files, it's automatically consistent with Delta: when a new version writes new files, those are fetched; old cached files simply stop being referenced.

You can warm it ahead of a dashboard refresh with SQL:

CACHE SELECT customer_id, amount, order_ts
FROM sales.orders
WHERE order_ts >= current_date() - INTERVAL 30 DAYS;
Enter fullscreen mode Exit fullscreen mode

For BI-style workloads that hit the same tables repeatedly, the disk cache usually gives you most of the benefit with none of the bookkeeping of the Spark cache. Also note that serverless compute has its own rules around caching APIs, so check the current limitations for your workspace before relying on cache() there.

When caching makes things worse

This is the part interviewers really want to hear:

  1. Single use. Caching something read once adds a write to memory/disk with no payoff.
  2. Caching too early. Cache after filters and column pruning, not the raw table. spark.read.table("big").cache() stores every column, then your later select can't push down into the cached data.
  3. Memory pressure. Cached blocks compete with execution memory. A large cache can push shuffles and joins into spilling, making the whole job slower. Watch for spill in the stage metrics after adding a cache.
  4. Stale results. The Spark cache is a snapshot. If another job updates the underlying Delta table, a long-running notebook can keep serving the old cached data until you unpersist() or run REFRESH TABLE. The disk cache doesn't have this problem because it caches immutable files.
  5. Forgetting to clean up. On shared all-purpose clusters, leftover caches eat memory for everyone. Use unpersist() when you're done, or spark.catalog.clearCache() at the end of a notebook.

Proving the cache is used

Don't guess; check:

enriched.explain()
Enter fullscreen mode Exit fullscreen mode

Look for InMemoryTableScan (and InMemoryRelation) in the physical plan. If you see the original Scan parquet/FileScan and joins instead, your later query isn't matching the cached plan, often because it was built from a different DataFrame lineage.

In the Spark UI:

  • The Storage tab shows each cached DataFrame, its storage level, the fraction cached, and how much is in memory vs disk. "Fraction cached" below 100% with MEMORY_ONLY means partitions are being recomputed.
  • In the stage details, cached reads show up as fast tasks with no input from cloud storage.
  • For the disk cache, the query profile and the cluster's metrics show cache hit rates and bytes read from cache vs remote.

Interview questions worth practising

  1. What's the difference between the Spark cache and the Databricks disk cache, and when would you use each?
  2. You called df.cache() but the job didn't get faster. Name three possible reasons.
  3. When would you choose checkpoint() over cache()? What about localCheckpoint()?
  4. A notebook shows yesterday's numbers even though the table was updated an hour ago. What's a likely cause?
  5. How do you confirm from the plan that a query is reading from the cache?

If you can answer those with the plan and the Storage tab in front of you, you've got caching covered.


I've put together a free 10-question sample of Azure Databricks and PySpark interview questions with worked answers, if you'd like more practice like this: free sample. The full set is in the Azure Databricks & PySpark Interview Pack.

Top comments (0)