[akshaychitneni]

Which issue does this PR close?

Closes #1829.

Rationale for this change

DataFrame::cache() is not supported in Ballista. This PR adds it as steps 2–4
of the design doc's plan (materialize + read-back, pinning, recovery part 1),
on top of the cache registry from step 1. A cached dataset is ordinary shuffle
output, pinned past job cleanup and reused through Ballista's existing shuffle
reader and stage rollback.

What changes are included in this PR?

• Materialize (cold path): a cache miss runs the cached subplan as its own
  build job, whose final shuffle output is pinned. The query job is held as
  Queued, without blocking, and dispatched once the build completes. Builds are
  shared between queries and run in parallel; a failed build fails the queries
  waiting on it. Submission runs on its own task, so a client that drops its
  request cannot leave a claim pending.
• Read-back (warm path): a cache node becomes a TableScan whose provider
  returns a ShuffleReaderExec over the pinned locations, with a projection back
  to the cached names and qualifiers. Any query planner plans a table scan, so
  no session needs anything installed: this works with custom session builders
  and standalone_with_state too. The read is registered as a complete stage
  input under a reserved stage id (EXTERNAL_STAGE_ID_BASE), so Ballista's own
  rollback restores it.
• Pinning and validate-on-read: job-data cleanup keeps a pinned job's data,
  and a location on a dead executor turns a hit into a rebuild. A held query
  re-checks its caches before it is released. An invalidated entry's build data
  is deleted after finished_job_data_clean_up_interval_seconds.
• Recovery (fail cleanly, rebuild on retry): if a running query loses cache
  data, the job fails with "lost DataFrame::cache() data; retry the query",
  under both planners, and the caches it read are dropped so a retry rebuilds
  them. A failed fetch counts as cache data only when it names one of the
  stage's cache locations (executor, build stage, partition), so a failed fetch
  of the job's own output is recovered as usual, even when stage ids collide.
  Rebuild-and-resume is a follow-up PR (see the design doc).
• Shutdown: cache submissions are admitted like any other, and a draining
  scheduler waits for held cache jobs.
• Tests and docs: scheduler unit and flow tests, execution-graph tests,
  client end-to-end tests (including a scheduler built from the user's state), a
  chaos scenario (kill the executor holding a cache), the design doc
  docs/source/contributors-guide/dataframe-cache.md, and a "Caching
  DataFrames" section in the tuning guide.

No proto changes and no protocol version bump. ballista-core only changes the
ballista.cache.noop description. CacheState (unreleased, from #2568) gains
pending_job and invalidate_if, and ExecutionGraph gains a defaulted
method. TaskManager::executor_lost now also returns the jobs that lost cache
data, so a caller has to fail them. Existing tests pass unchanged.

Are there any user-facing changes?

Yes, opt-in:

• Setting ballista.cache.noop = false makes df.cache() materialize and reuse
  results. The setting is on by default, so cache() stays a no-op unless the
  user changes it.
• New scheduler option --disable-dataframe-cache (operator switch).
• EXPLAIN of a cached frame never materializes it.
• A query that loses cache data mid-run fails with a clear error; retrying it
  rebuilds the cache.

[akshaychitneni]

cc @andygrove @milenkovicm please take a look

[comphead]

Thanks for this. Reusing pinned shuffle output and the stage rollback path is a good fit. Inline comments below. The main one is a possible hang in the static graph when a cache's build-stage id collides with an own input id (traced by reading, not run). The rest are mostly simplification and per-job overhead.

Two questions:

• After a scheduler restart, does a read job held as Queued stay Queued, since its pending state lives only in the process-local CacheLinker?
• Pinned caches are only removed by executor TTL cleanup. Is that acceptable until eviction lands?

[milenkovicm]

Thanks @akshaychitneni will have a look, 5k line change will take a bit to review

[akshaychitneni]

Thanks for this. Reusing pinned shuffle output and the stage rollback path is a good fit. Inline comments below. The main one is a possible hang in the static graph when a cache's build-stage id collides with an own input id (traced by reading, not run). The rest are mostly simplification and per-job overhead.

Thanks for the thorough review. I have addressed comments. Fixed the correctness bug you reported and will add more chaos test cases in the follow up pr to keep this short. Also removed cachestorage api for now. It was added to work on in-memory cache in future but will work on it later.

• After a scheduler restart, does a read job held as Queued stay Queued, since its pending state lives only in the process-local CacheLinker?

Currently, JobState backend is in-mem only, a restart drops job. Only a durable backend could leave job stuck queued. I will look into it as a followup

Pinned caches are only removed by executor TTL cleanup. Is that acceptable until eviction lands?

Yes and next read rebuilds it. I plan to work on cleanup on remove_session and also ttl based eviction tasks next after this pr.

[akshaychitneni]

Thanks @akshaychitneni will have a look, 5k line change will take a bit to review

Thanks! I kept materialize and cache reads together in this PR. The recovery path adds a fair amount of code, so if it would make review easier, I'm happy to move it to a separate PR. The design doc has been updated to describe the approach.

[akshaychitneni]

@milenkovicm @andygrove @comphead I've updated this PR to cover only the materialize and read paths. If a running query loses cache data while reading it, the query now fails cleanly, and a retry rebuilds the cache. To keep this PR small, I'll handle recovery (rebuild and resume) in a separate PR, together with its chaos tests. The design doc still describes that approach. Please take a look when you can. Thanks!

[andygrove]

Thanks for splitting this down, it's much easier to follow now! The fix for comphead's stage id collision looks good, and the test for it is great. I also like that nothing changes when ballista.cache.noop is left at its default, and the extra cost for jobs that don't use a cache is tiny.

I ran into three things that I think need fixing before this lands, and I reproduced each of them with a small test (attached to the inline comments, feel free to use them as regression tests). There's also one suggestion about cleanup that I'm fine with handling in the eviction follow-up.

A few smaller things, not blocking:

• read_schema in scheduler/src/cache.rs falls back to the logical schema when the field counts differ or the qualified schema can't be built. That would fail later at run time with a schema mismatch, so an error at planning time might be friendlier.
• TaskManager::executor_lost is still pub but has no callers now, and it drops the lost-cache result, so a caller could leave jobs hanging. Maybe remove it or have it fail those jobs?
• The ballista.cache.noop description in config.rs still says "Disable default cache node extension". It'd be nice to explain what turning it off does, and a short user-guide section would help people find the feature.
• The EXTERNAL_STAGE_ID_BASE doc says both planners assert ids stay below it, but those are debug_assert!s. And the design doc says lost cache data is "never a wrong or missing answer", while in this PR the query fails and asks for a retry.

The new CacheState methods don't need the api-change label since the trait hasn't been released yet.

[andygrove]

I think this check is a bit too broad. It matches any external input with data on the failing executor, so if a stage's own shuffle fetch from E1 fails for a transient reason while E1 also holds part of a cache, the job fails with "lost DataFrame::cache() data" instead of rerunning the input stage. on_job_terminal then drops every cache the job read too. On small clusters the cache and the job's own output almost always share executors, so this could hit most cache reads.

I reproduced it with the test below in execution_graph.rs. The cache is on executor1 under build stage 7, and the fetch failure is for the job's own stage 1 on executor1. The job fails with Stage 2 lost DataFrame::cache() data on executor executor-id1 and stage 1 is never rerun.

Could it match (executor_id, map_stage_id, map_partition_id) against the external locations instead, and fall through to the normal path when only an own input matches? FetchPartitionError has no job id, so when both match it's ambiguous, but treating that case as cache loss seems fine.

Repro test

    #[tokio::test]
    async fn own_fetch_failure_on_cache_executor_reruns_input() -> Result<()> {
        let executor1 = mock_executor("executor-id1".to_string());
        let mut agg_graph = test_aggregation_plan(4).await;
        revive_graph_and_complete_next_stage(&mut agg_graph)?;
        agg_graph.revive();
        let Some(ExecutionStage::Running(stage)) = agg_graph.stages.get_mut(&2) else {
            panic!("stage 2 should be running");
        };
        // The cache sits on executor1 too, under build stage 7.
        let mut cache_loc = external_location(3, &executor1.id);
        cache_loc.partition_id = PartitionId::new(&JobId::new("cache-job"), 7, 3);
        let mut output = StageOutput::new();
        output.partition_locations.insert(0, vec![cache_loc]);
        output.complete = true;
        stage.inputs.insert(EXTERNAL_STAGE_ID_BASE, output);

        // Fetching the job's own stage 1 output from executor1 fails.
        let task = agg_graph.pop_next_task(&executor1.id)?.unwrap();
        let failed =
            mock_failed_task(task, wrapped_fetch_failed_task(&executor1.id, 1, 0));
        let events = agg_graph.update_task_status(&executor1, vec![failed], 4, 4)?;

        assert!(
            !events.iter().any(|event| matches!(
                event,
                QueryStageSchedulerEvent::JobRunningFailed { .. }
            )),
            "a fetch failure on the job's own output must not fail the job: {events:?}"
        );
        Ok(())
    }

[akshaychitneni]

Fixed: a failed fetch now counts as lost cache data only when it matches a cache location's executor, build stage and partition (an own-input match too is treated as cache loss, as you suggested). Your repro is a regression test in a_failed_fetch_fails_the_job_only_for_cache_data

[andygrove]

This runs inside the execute_query handler, so if the client drops the request after prepare_caches claims the build but before submit_released posts it, the claim stays Pending forever. Every later read of that cache then waits on a build that never runs. The awaits in that window are job_id_generator.next_id() and post_event when the event channel is full.

I reproduced it with the test below in this file. A job id generator stalls on the read job's id and the submit times out. After that the registry still says Pending(job-0) and the next query over the cache stays Queued. The stall = false case finishes with the same number of ticks.

Could everything after the first claim run in a spawned task, or roll the claims back with a drop guard?

Repro test

    /// Hands out the build job's id, then stalls on the read job's id.
    struct StallSecondId {
        calls: std::sync::atomic::AtomicUsize,
        stall: std::sync::atomic::AtomicBool,
    }

    #[async_trait::async_trait]
    impl crate::scheduler_server::JobIdGenerator for StallSecondId {
        async fn next_id(&self) -> String {
            let n = self.calls.fetch_add(1, std::sync::atomic::Ordering::SeqCst);
            if n == 1 && self.stall.load(std::sync::atomic::Ordering::SeqCst) {
                std::future::pending::<()>().await;
            }
            format!("job-{n}")
        }
    }

    #[rstest::rstest]
    #[tokio::test]
    async fn dropped_submit_does_not_strand_the_cache(#[values(true, false)] stall: bool) {
        let generator = Arc::new(StallSecondId {
            calls: Default::default(),
            stall: std::sync::atomic::AtomicBool::new(stall),
        });
        let config = SchedulerConfig::default().with_job_id_generator(generator.clone());
        let (mut test, _metrics) = scheduler_test_with(config, None).await;
        let plan = cache_over("c1", scan());

        // The client gives up while the read job's id is being generated.
        let _ = tokio::time::timeout(Duration::from_millis(300), test.submit("dropped", &plan))
            .await;
        generator.stall.store(false, std::sync::atomic::Ordering::SeqCst);

        let again = test.submit("again", &plan).await.unwrap();
        for _ in 0..30 {
            let _ = tokio::time::timeout(Duration::from_millis(100), test.tick()).await;
        }
        let status = test.job_status(&again).await.unwrap();
        assert!(status.as_ref().is_some_and(is_terminal), "{status:?}");
    }

[andygrove]

I think this means standalone_with_state and remote_with_state can't read caches at all. new_standalone_scheduler_from_state gives the scheduler the user's built state, and a built SessionState always has a query planner, so CacheReadPlanner never gets added. The same goes for an override_session_builder that builds its own state. The build job runs and gets pinned, then the read fails, and because that error isn't treated as lost data every retry fails the same way and the pinned data stays around.

I reproduced it with the test below in context_cache.rs. Both cases fail on the first collect and again on the retry with this scheduler's session cannot plan DataFrame::cache() reads: add CacheReadPlanner to its query planner.

Could the scheduler lower CacheReadNode itself while resolving the cache reads, so it doesn't depend on the session's planner? If that's too much for this PR, checking that the planner can handle it before claiming the build would at least avoid the wasted build and the leftover data. The design doc's note about states that "have one [planner]" would need an update either way.

Repro test

    #[rstest]
    #[case::standalone_with_state(crate::common::standalone_context_with_state())]
    #[case::remote_with_state(crate::common::remote_context_with_state())]
    #[tokio::test]
    async fn cache_with_state(
        #[future(awt)]
        #[case]
        ctx: SessionContext,
        test_data: String,
    ) {
        enable_cache(&ctx).await;
        let cached = alltypes(&ctx, &test_data)
            .await
            .select(vec![col("id")])
            .unwrap()
            .cache()
            .await
            .unwrap();
        assert_eq!(cached.clone().count().await.unwrap(), 8);
        assert_eq!(cached.count().await.unwrap(), 8);
    }

[akshaychitneni]

Fixed by making the read planner-independent: a cache is now read through a plain TableScan, so any session plans it. Your cache_with_state test is in, over a self-join to cover restored qualifiers.

[andygrove]

Suggestion, fine for the eviction follow-up: a pinned job skips clean_up_job_data_delayed entirely, and the entry that invalidate_if removes is dropped without scheduling any cleanup. So every rebuild (after CACHE_DATA_LOST or a dead-executor hit) leaves the old build's data on disk unless the executor janitor is turned on, and it's off by default. Registry entries also stay for the life of the scheduler. Could invalidation at least schedule cleanup for the entry it removes? That would keep rebuilds from piling up data until proper eviction lands.

[akshaychitneni]

Done

[akshaychitneni]

@andygrove Thanks for the thorough review and the repros. All three issues are fixed. I also updated design doc to match the approach on read path using tablescan instead of planner extension.