feat: implement DataFrame.checkpoint and DataFrame.localCheckpoint

crates/sail-catalog/Cargo.toml

@@ -19,3 +19,4 @@ lazy_static = { workspace = true }
 async-trait = { workspace = true }
 thiserror = { workspace = true }
 tokio = { workspace = true }
+url = { workspace = true }

crates/sail-catalog/src/error.rs

@@ -16,6 +16,7 @@ pub enum CatalogObject {
     Function,
     TemporaryView,
     LogicalPlan,
+    CachedRelation,
 }
 
 impl fmt::Display for CatalogObject {
@@ -30,6 +31,7 @@ impl fmt::Display for CatalogObject {
             CatalogObject::Function => "Function",
             CatalogObject::TemporaryView => "Temporary View",
             CatalogObject::LogicalPlan => "Logical Plan",
+            CatalogObject::CachedRelation => "Cached Relation",
         };
         write!(f, "{name}")
     }

crates/sail-catalog/src/manager/mod.rs

@@ -5,7 +5,9 @@ use datafusion_expr::{LogicalPlan, ScalarUDF};
 use sail_common_datafusion::extension::SessionExtension;
 
 use crate::error::{CatalogError, CatalogObject, CatalogResult};
-use crate::manager::tracker::{CatalogFunctionId, CatalogLogicalPlanId, CatalogObjectTracker};
+use crate::manager::tracker::{
+    CatalogCachedRelation, CatalogFunctionId, CatalogLogicalPlanId, CatalogObjectTracker,
+};
 use crate::provider::{CatalogProvider, Namespace};
 use crate::temp_view::TemporaryViewManager;
 
@@ -147,6 +149,29 @@ impl CatalogManager {
     ) -> CatalogResult<Arc<LogicalPlan>> {
         self.tracker.get_tracked_logical_plan(id)
     }
+
+    pub fn track_cached_relation(
+        &self,
+        relation_id: String,
+        relation: CatalogCachedRelation,
+    ) -> CatalogResult<()> {
+        self.tracker.track_cached_relation(relation_id, relation)
+    }
+
+    pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult<CatalogCachedRelation> {
+        self.tracker.get_cached_relation(relation_id)
+    }
+
+    pub fn remove_cached_relation(
+        &self,
+        relation_id: &str,
+    ) -> CatalogResult<Option<CatalogCachedRelation>> {
+        self.tracker.remove_cached_relation(relation_id)
+    }
+
+    pub fn drain_cached_relation_storage_uris(&self) -> CatalogResult<Vec<String>> {
+        self.tracker.drain_cached_relation_storage_uris()
+    }
 }
 
 impl CatalogManagerState {

crates/sail-catalog/src/manager/tracker.rs

@@ -12,12 +12,25 @@ pub struct CatalogFunctionId(u64);
 #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, PartialOrd, Serialize, Deserialize)]
 pub struct CatalogLogicalPlanId(u64);
 
+#[derive(Debug, Clone)]
+pub struct CatalogCachedRelation {
+    pub plan: Arc<LogicalPlan>,
+    pub fields: Vec<String>,
+    pub storage_uri: Option<String>,
+}
+
 #[derive(Default)]
 struct CatalogObjectTrackerState {
     next_function_id: u64,
     next_logical_plan_id: u64,
     functions: HashMap<u64, ScalarUDF>,
     logical_plans: HashMap<u64, Arc<LogicalPlan>>,
+    /// Maps relation IDs from `CachedRemoteRelation` to their checkpointed
+    /// logical plans and user-facing field names. Populated when the server
+    /// handles a `CheckpointCommand`, and queried when resolving a
+    /// `CachedRemoteRelation` query node. Entries are removed when the server
+    /// handles a `RemoveCachedRemoteRelationCommand`.
+    cached_relations: HashMap<String, CatalogCachedRelation>,
 }
 
 /// Tracks in-memory objects (UDFs and logical plans) that cannot be serialized directly,
@@ -75,4 +88,72 @@ impl CatalogObjectTracker {
             .cloned()
             .ok_or_else(|| CatalogError::NotFound(CatalogObject::LogicalPlan, id.0.to_string()))
     }
+
+    /// Stores a logical plan keyed by an opaque relation ID.
+    /// This is used to back checkpointed DataFrames referenced via
+    /// `CachedRemoteRelation` in subsequent Spark Connect requests.
+    pub fn track_cached_relation(
+        &self,
+        relation_id: String,
+        relation: CatalogCachedRelation,
+    ) -> CatalogResult<()> {
+        let mut state = self.state()?;
+        state.cached_relations.insert(relation_id, relation);
+        Ok(())
+    }
+
+    pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult<CatalogCachedRelation> {
+        let state = self.state()?;
+        state
+            .cached_relations
+            .get(relation_id)
+            .cloned()
+            .ok_or_else(|| {
+                CatalogError::NotFound(CatalogObject::CachedRelation, relation_id.to_string())
+            })
+    }
+
+    pub fn remove_cached_relation(
+        &self,
+        relation_id: &str,
+    ) -> CatalogResult<Option<CatalogCachedRelation>> {
+        let mut state = self.state()?;
+        Ok(state.cached_relations.remove(relation_id))
+    }
+
+    pub fn drain_cached_relation_storage_uris(&self) -> CatalogResult<Vec<String>> {
+        let mut state = self.state()?;
+        Ok(state
+            .cached_relations
+            .drain()
+            .filter_map(|(_, relation)| relation.storage_uri)
+            .collect())
+    }
+}
+
+impl Drop for CatalogObjectTracker {
+    /// Best-effort cleanup of any on-disk checkpoint storage for cached
+    /// relations that were never released via `RemoveCachedRemoteRelationCommand`.
+    /// This covers session timeouts, client disconnects, and any other path
+    /// where the catalog manager is dropped without explicit cleanup.
+    fn drop(&mut self) {
+        let Ok(state) = self.state.get_mut() else {
+            return;
+        };
+        for (_, relation) in state.cached_relations.drain() {
+            let Some(uri) = relation.storage_uri else {
+                continue;
+            };
+            if let Ok(url) = url::Url::parse(&uri) {
+                if url.scheme() == "file" {
+                    if let Ok(path) = url.to_file_path() {
+                        // `Drop` runs in a sync context, so we use blocking
+                        // filesystem calls for local checkpoint cleanup.
+                        // Errors are ignored because the caller cannot react.
+                        let _ = std::fs::remove_dir_all(path);
+                    }
+                }
+            }
+        }
+    }
 }

crates/sail-common/src/config/application.rs

@@ -509,6 +509,7 @@ pub enum CatalogType {
 pub struct SparkConfig {
     pub session_timeout_secs: u64,
     pub execution_heartbeat_interval_secs: u64,
+    pub checkpoint_dir: String,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]

crates/sail-common/src/config/application.yaml

@@ -737,6 +737,20 @@
     the session active.
   experimental: true
 
+- key: spark.checkpoint_dir
+  type: string
+  default: ""
+  description: |
+    The directory where reliable (disk) checkpoints created via
+    `DataFrame.checkpoint()` and `DataFrame.localCheckpoint()` are written.
+    When empty, a `sail-checkpoints/` subdirectory of the system temporary
+    directory is used. In cluster deployments, this should be set to a
+    location that is shared and readable by all workers, such as an
+    object-store URI or a shared filesystem path, since the default location
+    is node-local. Checkpoint files are best-effort cleaned up when the client
+    releases the cached relation or when the session ends.
+  experimental: true
+
 - key: flight.session_timeout_secs
   type: number
   default: "3600"

crates/sail-plan/src/resolver/query/misc.rs

@@ -5,9 +5,13 @@ use datafusion::catalog::MemTable;
 use datafusion_common::{DFSchema, DFSchemaRef, ParamValues};
 use datafusion_expr::{EmptyRelation, Extension, LogicalPlan, UNNAMED_TABLE};
 use log::warn;
+use sail_catalog::error::{CatalogError, CatalogObject};
+use sail_catalog::manager::CatalogManager;
 use sail_common::spec;
 use sail_common_datafusion::array::record_batch::{cast_record_batch, read_record_batches};
+use sail_common_datafusion::extension::SessionExtensionAccessor;
 use sail_common_datafusion::literal::LiteralEvaluator;
+use sail_common_datafusion::rename::logical_plan::rename_logical_plan;
 use sail_logical_plan::range::RangeNode;
 
 use crate::error::{PlanError, PlanResult};
@@ -138,6 +142,33 @@ impl PlanResolver<'_> {
         )
     }
 
+    /// Resolves a query that refers to a cached remote relation by looking up
+    /// the previously checkpointed logical plan from the catalog manager. The
+    /// cached plan was stored with internal field IDs and its user-facing field
+    /// names, so we register fresh field IDs in the current resolver state and
+    /// rename the plan accordingly. This preserves duplicate user-facing names
+    /// without relying on them as `Expr::Column` references inside the cached
+    /// plan.
+    pub(super) async fn resolve_query_cached_remote_relation(
+        &self,
+        relation_id: String,
+        state: &mut PlanResolverState,
+    ) -> PlanResult<LogicalPlan> {
+        let manager = self.ctx.extension::<CatalogManager>()?;
+        let relation = match manager.get_cached_relation(&relation_id) {
+            Ok(relation) => relation,
+            Err(CatalogError::NotFound(CatalogObject::CachedRelation, _)) => {
+                return Err(PlanError::AnalysisError(format!(
+                    "cached remote relation not found: {relation_id}"
+                )));
+            }
+            Err(error) => return Err(error.into()),
+        };
+        let names = state.register_field_names(relation.fields.clone());
+        let plan = rename_logical_plan((*relation.plan).clone(), &names)?;
+        Ok(plan)
+    }
+
     pub(super) async fn resolve_query_hint(
         &self,
         input: spec::QueryPlan,

crates/sail-plan/src/resolver/query/mod.rs

@@ -227,8 +227,9 @@ impl PlanResolver<'_> {
             QueryNode::CachedLocalRelation { .. } => {
                 return Err(PlanError::todo("cached local relation"));
             }
-            QueryNode::CachedRemoteRelation { .. } => {
-                return Err(PlanError::todo("cached remote relation"));
+            QueryNode::CachedRemoteRelation { relation_id } => {
+                self.resolve_query_cached_remote_relation(relation_id, state)
+                    .await?
             }
             QueryNode::CommonInlineUserDefinedTableFunction(udtf) => {
                 self.resolve_query_common_inline_udtf(udtf, state).await?

crates/sail-session/Cargo.toml

@@ -42,9 +42,11 @@ datafusion = { workspace = true }
 datafusion-common = { workspace = true }
 datafusion-expr = { workspace = true }
 datafusion-physical-expr = { workspace = true }
+object_store = { workspace = true }
 secrecy = { workspace = true }
 tonic = { workspace = true }
 tokio = { workspace = true }
 chrono = { workspace = true }
 indexmap = { workspace = true }
 futures = { workspace = true }
+url = { workspace = true }