akshaychitneni commented on code in PR #2568:
URL:
https://github.com/apache/datafusion-ballista/pull/2568#discussion_r4242517014
##########
ballista/scheduler/src/cluster/mod.rs:
##########
@@ -363,6 +378,62 @@ pub trait JobState: Send + Sync {
fn produce_config(&self) -> SessionConfig;
}
+/// Scheduler-side metadata store for materialized `DataFrame::cache()`
results —
+/// the mapping from a cache key to the shuffle partition locations of its
output.
+///
+/// This is the third piece of scheduler state alongside [`ClusterState`] and
+/// [`JobState`], and follows the same shape: an async trait with an in-memory
+/// implementation today
([`InMemoryCacheState`](crate::cluster::memory::InMemoryCacheState))
+/// and an [`init`](CacheState::init) hook so a future durable backend (e.g.
one
+/// backed by an object store) can rehydrate entries on startup. The methods
are
+/// `async` even though the in-memory backend does no I/O, so a durable backend
+/// slots in without changing any signatures.
+///
+/// A cache *miss* is simply [`lookup`](CacheState::lookup) returning `None`
(the
+/// key is absent, or materialization is still in flight); a *hit* is `Some`.
+#[async_trait::async_trait]
+pub trait CacheState: Send + Sync + 'static {
+ /// Initializes the backend. A no-op for the in-memory backend; durable
+ /// backends load previously persisted entries here.
+ async fn init(&self) -> Result<()> {
+ Ok(())
+ }
+
+ /// Returns the materialized data for `key`, or `None` on a miss (absent or
+ /// still materializing).
+ async fn lookup(&self, key: &CacheKey) ->
Result<Option<MaterializedCache>>;
+
+ /// Attempts to claim `key` for materialization by `job_id`, deduplicating
+ /// concurrent first-time misses to a single [`BeginOutcome::Claimed`].
+ async fn begin_materialization(
+ &self,
+ key: CacheKey,
+ job_id: String,
+ ) -> Result<BeginOutcome>;
+
+ /// Records the materialized shuffle locations for `key`.
+ async fn complete_materialization(
+ &self,
+ key: CacheKey,
+ job_id: String,
+ schema: SchemaRef,
+ locations: Vec<Vec<PartitionLocation>>,
+ ) -> Result<()>;
+
+ /// Invalidates a single entry; the next [`lookup`](CacheState::lookup) is
a miss.
+ async fn invalidate(&self, key: &CacheKey) -> Result<()>;
+
+ /// Invalidates every entry holding a partition on `executor_id` (the
+ /// executor-loss policy) and returns the invalidated keys.
+ async fn invalidate_executor(&self, executor_id: &str) ->
Result<Vec<CacheKey>>;
Review Comment:
Updated api to return entries for callers to unpin and remove data
--
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]
