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,45 @@ 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())
+    }
 }

crates/sail-common-datafusion/Cargo.toml

@@ -18,6 +18,8 @@ thiserror = { workspace = true }
 either = { workspace = true }
 lexical-core = { workspace = true }
 futures = { workspace = true }
+log = { workspace = true }
+object_store = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 async-trait = { workspace = true }
@@ -27,6 +29,7 @@ pin-project-lite = { workspace = true }
 parquet-variant-compute = { workspace = true }
 parquet-variant-json = { workspace = true }
 tokio = { workspace = true }
+url = { workspace = true }
 serde_arrow = { workspace = true }
 
 [build-dependencies]

crates/sail-common-datafusion/src/checkpoint.rs

@@ -0,0 +1,46 @@
+use datafusion::execution::runtime_env::RuntimeEnv;
+use futures::{StreamExt, TryStreamExt};
+use log::warn;
+use object_store::ObjectStoreScheme;
+
+pub async fn cleanup_checkpoint_storage(runtime_env: &RuntimeEnv, storage_uri: &str) {
+    let Ok(url) = url::Url::parse(storage_uri) else {
+        warn!("failed to parse checkpoint location for cleanup: {storage_uri}");
+        return;
+    };
+    let Ok((_, checkpoint_path)) = ObjectStoreScheme::parse(&url) else {
+        warn!("failed to parse checkpoint object-store path for cleanup: {storage_uri}");
+        return;
+    };
+    let Ok(object_store) = runtime_env.object_store_registry.get_store(&url) else {
+        warn!("failed to resolve checkpoint object store for cleanup: {storage_uri}");
+        return;
+    };
+    let locations = object_store
+        .list(Some(&checkpoint_path))
+        .map_ok(|meta| meta.location)
+        .boxed();
+    if let Err(e) = object_store
+        .delete_stream(locations)
+        .try_collect::<Vec<_>>()
+        .await
+    {
+        warn!("failed to remove checkpoint location {storage_uri}: {e}");
+    }
+    cleanup_file_checkpoint_directory(&url, storage_uri).await;
+}
+
+async fn cleanup_file_checkpoint_directory(url: &url::Url, storage_uri: &str) {
+    if url.scheme() != "file" {
+        return;
+    }
+    let Ok(path) = url.to_file_path() else {
+        warn!("failed to convert checkpoint file URL to path for cleanup: {storage_uri}");
+        return;
+    };
+    if let Err(e) = tokio::fs::remove_dir_all(&path).await {
+        if e.kind() != std::io::ErrorKind::NotFound {
+            warn!("failed to remove checkpoint directory {storage_uri}: {e}");
+        }
+    }
+}

crates/sail-common-datafusion/src/lib.rs

@@ -1,5 +1,6 @@
 pub mod array;
 pub mod catalog;
+pub mod checkpoint;
 pub mod column_features;
 pub mod datasource;
 pub mod display;

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 }