DataFrame Cache Design (`DataFrame::cache()`)#

Goal#

df.cache() marks a DataFrame to be computed once and reused by later actions instead of recomputed. In Ballista a cached dataset is shuffle output that is kept (pinned) past normal job cleanup, plus a scheduler-side record of where that output lives. It reuses the native blocking shuffle as its storage and transport — no separate cache store, no new wire format, and the read-back is an ordinary ShuffleReaderExec.

The cache is a hint, not a correctness dependency: a cached node keeps its input subplan (lineage), so lost cache data is simply recomputed — never a wrong or missing answer.

The model#

In the logical plan, a cached dataset is a BallistaCacheNode wrapping the subplan to cache. Materialized, it has two pieces:

Pinned shuffle output — the cached subplan runs as its own job; its final stage’s shuffle files ({work_dir}/{job_id}/{stage_id}/…) survive job cleanup.
A cross-job registry — scheduler state (the CacheState trait, in-memory InMemoryCacheState over CacheRegistry) mapping a cache key to the schema and PartitionLocations of that output, so a different job can read it.

A registry is needed because Ballista materializes the cache in a separate distributed job: the mapping from a cache to its shuffle output must outlive that job and be visible to other jobs, so the scheduler owns it explicitly.

Keying: (session_id, cache_id), where cache_id is a per-cache()-call UUID in the BallistaCacheNode. Reuse happens whenever a submitted plan still contains that node (same cache_id) — the cached DataFrame reused directly or built upon. Dedup across different plans with equal content is deferred.

End-to-end flow#

Three roles: the client holds the plan, the scheduler owns planning and the registry, and the executors run tasks and hold the shuffle files.

df.cache() is lazy: it returns a DataFrame whose logical plan is wrapped in a BallistaCacheNode with a fresh cache_id — nothing executes yet. Reuse is keyed on that cache_id: any plan later submitted that still contains the node — the same cached DataFrame, or a query built on top of it — resolves to the same cache entry (the registry’s record of the materialized output). Calling df.cache() again mints a new cache_id.

Miss and hit are not separate stories — they converge: a miss materializes and then reads the pinned output, while a hit does the same read and skips the materialize.

                     df.cache()                 client; lazy — wraps the logical
                         │                       plan in BallistaCacheNode
                         ▼                       {cache_id, session_id, input}
            submit a plan carrying the node      client ──► scheduler
                         │
                         ▼
             lookup((session_id, cache_id))      scheduler
                         │
            ┌────────────┴─────────────┐
          MISS                        HIT
            │                          │
            ▼                          ▼
   begin_materialization        validate locations
   run `input` as its own       against live executors
   job (two-job split, #1993)   (any dead ⇒ re-MISS)
   → executors write the                │
     pinned shuffle files               │
            │                           │
            ▼                           │
   complete_materialization             │
   → Materialized (fresh locs)          │
            │                           │
            └────────────┬──────────────┘
                         ▼
        lower BallistaCacheNode ──► physical ShuffleReaderExec,
        reading the pinned output served by the executors
                         │
                         ▼
                   results ──► client

Cold path — first action (miss → materialize):

The client submits a plan containing BallistaCacheNode{cache_id, session_id, input}.
The scheduler forms key = (session_id, cache_id) and calls lookup(key) → miss.
begin_materialization(key, job) claims the key; the scheduler runs input as its own job (a logical-plan rewrite that splits the query into two jobs, see #1993), writing pinned shuffle output on the executors.
On that job’s JobFinished, complete_materialization(key, job, schema, locations) records the output → the cache entry becomes Materialized.
The outer query reads that output and returns results — so the first action both materializes and reads back.

Warm path — later action (hit → read-back):

The client submits a plan that still carries the same cache_id.
lookup(key) → hit. Its locations are validated against the live-executor set; a dead location downgrades the hit to a miss (see Invalidation).
The scheduler lowers the BallistaCacheNode to a physical ShuffleReaderExec over the recorded locations — the same mechanism stages use to read each other — so input is never planned or executed.
Results return with no materialization job.

Ordering & concurrency. Within one action the read is sequenced after the create: the read’s ShuffleReaderExec needs the recorded locations, so the scheduler resolves it only on the create job’s JobFinished — event-driven, no polling. A concurrent action that finds the key Pending (begin_materialization → AlreadyPending) must either wait for that materialization or fall back to computing the subplan inline; the exact policy — wait-and-notify vs. poll vs. inline, and releasing waiters if the job aborts — is an open question.

Registry lifecycle#

  (absent) --begin--> Pending{job} --complete(job)--> Materialized
                      Pending{job} --abort(job)-----> (absent)
                      Materialized --invalidate------> (absent)

The contract is safe against the out-of-order events a distributed scheduler produces (a completion callback racing invalidation, teardown, or a re-claim):

begin_materialization atomically gives one job the claim.
complete_materialization promotes only a Pending entry owned by that job (duplicate completions are idempotent); anything else is Stale, so a late completion cannot resurrect a removed entry or overwrite a newer claim.
abort_materialization clears a failed job’s pending claim.
invalidate / invalidate_executor / remove_session remove entries and return them so the caller can unpin the shuffle files (Ballista deletes explicitly — there is no reclaimer).
pinned_stages() lists the (job_id, stage_id) stages to protect from cleanup.

Pinning#

Shuffle output is normally reclaimed after a job finishes. A cached stage must be exempted, so job-data cleanup consults pinned_stages() and skips those stage dirs. Pins are per-(job_id, stage_id) because shuffle data is stored and deleted per stage, and a cached stage is often an intermediate stage (reaped immediately on job success) that a job-level pin could not protect.

Invalidation#

An entry is removed when its data is no longer trustworthy or needed:

Executor loss (invalidate_executor) — any materialized entry with a partition on the lost executor is dropped.
Session teardown (remove_session) — all of a session’s entries are dropped.
Explicit / eviction (invalidate) — a single entry is dropped (used by the future eviction sweeper).

Invalidation is whole-entry: the entire cache is dropped and re-materialized in full on the next access — we never serve a cache with missing partitions (per-partition rebuild is a future optimization). Because a cached node keeps its lineage, invalidation always degrades to a recompute, never a wrong answer.

Each removal returns the removed entries so the caller can unpin and delete their shuffle files — Ballista has no background reclaimer, so cleanup is explicit. A removed entry’s next lookup is a miss, which re-triggers materialization.

Executor loss during materialization (executor-loss follow-up). If an executor is lost after the cached stage finishes but before its job completes, complete_materialization can record a now-dead location (the in-graph recovery only rebuilds output an unfinished stage still reads). The fix is validate-on-read: on a hit the planner checks locations against the live-executor set and, if any is dead, treats it as a miss and invalidates it — enforcing “never read from a dead executor” under all timings.

Durability & reconciliation#

CacheState is a pluggable backend; the in-memory one loses entries on scheduler restart (caches become orphaned files, reclaimed by cleanup), and a durable backend can persist and rehydrate them — an optional efficiency extension. A stale or missing location is caught by validate-on-read (executor gone) or the ShuffleReader fetch-failure path (file gone) and degrades to re-materialization.

Future work#

Eviction: entries today leave only via invalidation/executor-loss/session teardown, so a long-running scheduler accumulates cache data. Planned: a TTL / size-cap sweeper (via invalidate) and/or idle-session expiry (via remove_session), using the removal-returns-entries contract above.
Locality: place read-back tasks on the executor holding each partition (TaskDistributionPolicy::Custom), Flight fetch as fallback (see #2319).

Implementation plan#

Registry (done) — CacheRegistry + CacheState/InMemoryCacheState.
Materialize-on-miss + read-back — wire lookup/begin/complete/abort into planning and the job handlers; builds on the two-job rewrite in #1993.
Pinning — teach job-data cleanup to honor pinned_stages().
Locality — place read-back tasks where partitions live (#2319).
Executor-loss invalidation + session teardown — wire invalidate_executor / remove_session and the unpin path, add validate-on-read on the hit path, and enable an end-to-end no-recompute test.
Eviction — TTL / size cap / idle-session expiry.

DataFrame Cache Design (DataFrame::cache())#