akshaychitneni commented on code in PR #2568: URL: https://github.com/apache/datafusion-ballista/pull/2568#discussion_r4242405173
########## ballista/scheduler/src/state/cache_registry.rs: ########## @@ -0,0 +1,485 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Cross-job registry of materialized `DataFrame::cache()` results. +//! +//! A cached dataset is represented in the logical plan by a +//! [`BallistaCacheNode`](ballista_core::extension::BallistaCacheNode). When it +//! is first executed it is materialized as *pinned* shuffle output on the +//! executors (shuffle files that survive normal job cleanup). This registry +//! records, per cache key, the schema and the shuffle [`PartitionLocation`]s of +//! that output so that a later job can read the data back over the existing +//! shuffle path instead of recomputing the subplan. +//! +//! Unlike a job's [`ExecutionGraph`](crate::state::execution_graph::ExecutionGraph), +//! which is garbage-collected shortly after the job completes, entries here +//! outlive the job that produced them — that is the whole point of a cache. +//! +//! # Keying +//! +//! Entries are keyed by [`CacheKey`] = `(session_id, cache_id)`, where +//! `cache_id` is the per-`cache()`-call UUID stamped by +//! `BallistaCacheFactory`. This gives reuse whenever the client holds on to the +//! cached `DataFrame` (its logical plan keeps the same `cache_id`). Canonical +//! plan-based keying for cross-call deduplication is intentionally deferred. +//! +//! # Lifecycle +//! +//! ```text +//! (absent) --begin_materialization--> Pending +//! Pending --complete_materialization--> Materialized +//! Materialized --invalidate / invalidate_executor--> (absent) +//! ``` +//! +//! Invalidation simply removes the entry; the next lookup is therefore a miss +//! and the subplan is re-materialized. This implements the "if an executor +//! holding a cached partition is lost, the cache is invalidated" policy. + +use std::collections::HashSet; + +use ballista_core::serde::scheduler::PartitionLocation; +use dashmap::DashMap; +use datafusion::arrow::datatypes::SchemaRef; + +/// Identifies a cached dataset. +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct CacheKey { + /// The owning session. + pub session_id: String, + /// The per-`cache()`-call identifier carried by `BallistaCacheNode`. + pub cache_id: String, +} + +impl CacheKey { + /// Creates a new cache key. + pub fn new(session_id: impl Into<String>, cache_id: impl Into<String>) -> Self { + Self { + session_id: session_id.into(), + cache_id: cache_id.into(), + } + } +} + +/// The materialized output of a cache: its schema and the shuffle locations of +/// each output partition. +#[derive(Debug, Clone)] +pub struct MaterializedCache { + /// Schema of the cached data. + pub schema: SchemaRef, + /// Shuffle locations indexed by output partition. The inner `Vec` allows a + /// single output partition to be backed by more than one file (e.g. one per + /// map task), mirroring [`ExecutionStage`] shuffle output. + /// + /// [`ExecutionStage`]: crate::state::execution_stage + pub locations: Vec<Vec<PartitionLocation>>, + /// The job that materialized this cache. Used to pin/unpin its shuffle data. + pub materializing_job: String, +} + +/// Internal state of a cache entry. +#[derive(Debug, Clone)] +enum CacheState { + /// A materialization job has been submitted but has not completed; the + /// partition locations are not yet known. + Pending { materializing_job: String }, + /// The cache is fully materialized and readable. + Materialized(MaterializedCache), +} + +/// Cross-job registry of materialized cache entries. +/// +/// Cheap to clone-share via `Arc`; all methods take `&self` and are safe to call +/// concurrently. +#[derive(Debug, Default)] +pub struct CacheRegistry { + entries: DashMap<CacheKey, CacheState>, +} + +/// Outcome of [`CacheRegistry::begin_materialization`]. +#[derive(Debug, PartialEq, Eq)] +pub enum BeginOutcome { + /// This caller claimed the key and should submit the materialization job. + Claimed, + /// A materialization job is already in flight for this key; do not submit. + AlreadyPending, + /// The key is already materialized; the caller should read it back instead. + AlreadyMaterialized, +} + +impl CacheRegistry { + /// Creates an empty registry. + pub fn new() -> Self { + Self::default() + } + + /// Returns the materialized data for `key`, or `None` if the key is absent + /// (a miss) or still pending. + pub fn lookup(&self, key: &CacheKey) -> Option<MaterializedCache> { + match self.entries.get(key)?.value() { + CacheState::Materialized(cache) => Some(cache.clone()), + CacheState::Pending { .. } => None, + } + } + + /// Attempts to claim `key` for materialization by `job_id`. + /// + /// Returns [`BeginOutcome::Claimed`] only when the caller is responsible for + /// submitting the materialization job. If another job already claimed the + /// key, or the key is already materialized, the caller must not submit a + /// duplicate. This races safely: concurrent callers for the same key see + /// exactly one `Claimed`. + pub fn begin_materialization( + &self, + key: CacheKey, + job_id: impl Into<String>, + ) -> BeginOutcome { + use dashmap::mapref::entry::Entry; + match self.entries.entry(key) { + Entry::Occupied(occupied) => match occupied.get() { + CacheState::Pending { .. } => BeginOutcome::AlreadyPending, + CacheState::Materialized(_) => BeginOutcome::AlreadyMaterialized, + }, + Entry::Vacant(vacant) => { + vacant.insert(CacheState::Pending { + materializing_job: job_id.into(), + }); + BeginOutcome::Claimed + } + } + } + + /// Records the materialized output for `key`, transitioning it to + /// [`CacheState::Materialized`]. Overwrites any existing state so that a + /// re-materialization (after invalidation) can refresh in place. `job_id` is + /// the job that produced the shuffle data and is retained for pin/unpin. + pub fn complete_materialization( Review Comment: Totally makes sense. Updated to make sure to only promote pending job if it exists, otherwise returns a new CompleteOutcome::Stale for the caller to handle. Added test case -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
