milenkovicm commented on code in PR #2568: URL: https://github.com/apache/datafusion-ballista/pull/2568#discussion_r4250189513
########## ballista/scheduler/src/state/cache_registry.rs: ########## @@ -0,0 +1,748 @@ +// 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`](ballista_core::serde::scheduler::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`](crate::state::cache_registry::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 Review Comment: what would happen if executors holding cache partition is missing? do we handle that case? do we rebuild whole cache or we can rebuild just that one partition ? ########## ballista/scheduler/src/state/cache_registry.rs: ########## @@ -0,0 +1,748 @@ +// 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`](ballista_core::serde::scheduler::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`](crate::state::cache_registry::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{job} +//! Pending{job} --complete_materialization(job) [match]----> Materialized +//! Pending{job} --abort_materialization(job) [match]----> (absent) +//! Pending{job} --complete_materialization(other|absent)---> Stale (no change) +//! Materialized --invalidate / invalidate_executor / remove_session--> (absent) +//! ``` +//! +//! Transitions out of `Pending` are **job-id matched**: a late completion from a +//! job that no longer owns the claim (because the entry was invalidated and +//! re-claimed, or the session was removed) is reported +//! [`CompleteOutcome::Stale`](crate::state::cache_registry::CompleteOutcome::Stale) +//! and changes nothing, so stale output cannot resurrect a removed entry or +//! overwrite a newer claim. 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::JobId; +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 EntryState { + /// 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. Kept `pub(crate)`: the public surface is the +/// [`CacheState`](crate::cluster::CacheState) trait, so this type can evolve +/// freely while the feature grows. +#[derive(Debug, Default)] +pub(crate) struct CacheRegistry { + entries: DashMap<CacheKey, EntryState>, +} + +/// Outcome of claiming a key for materialization +/// ([`CacheState::begin_materialization`](crate::cluster::CacheState::begin_materialization)). +#[derive(Debug, PartialEq, Eq)] +#[non_exhaustive] +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, +} + +/// Outcome of recording a materialization +/// ([`CacheState::complete_materialization`](crate::cluster::CacheState::complete_materialization)). +#[derive(Debug, PartialEq, Eq)] +#[non_exhaustive] +pub enum CompleteOutcome { Review Comment: i guess we can have cache where partitions are missing (due to executor failure) -- 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]
