feat: implement DataFrame.checkpoint and DataFrame.localCheckpoint

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

@@ -147,6 +147,22 @@ impl CatalogManager {
     ) -> CatalogResult<Arc<LogicalPlan>> {
         self.tracker.get_tracked_logical_plan(id)
     }
+
+    pub fn track_cached_relation(
+        &self,
+        relation_id: String,
+        plan: Arc<LogicalPlan>,
+    ) -> CatalogResult<()> {
+        self.tracker.track_cached_relation(relation_id, plan)
+    }
+
+    pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult<Arc<LogicalPlan>> {
+        self.tracker.get_cached_relation(relation_id)
+    }
+
+    pub fn remove_cached_relation(&self, relation_id: &str) -> CatalogResult<bool> {
+        self.tracker.remove_cached_relation(relation_id)
+    }
 }
 
 impl CatalogManagerState {

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

@@ -18,6 +18,10 @@ struct CatalogObjectTrackerState {
     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. Populated when the server handles a `CheckpointCommand`,
+    /// and queried when resolving a `CachedRemoteRelation` query node.
+    cached_relations: HashMap<String, Arc<LogicalPlan>>,
 }
 
 /// Tracks in-memory objects (UDFs and logical plans) that cannot be serialized directly,
@@ -75,4 +79,33 @@ 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,
+        plan: Arc<LogicalPlan>,
+    ) -> CatalogResult<()> {
+        let mut state = self.state()?;
+        state.cached_relations.insert(relation_id, plan);
+        Ok(())
+    }
+
+    pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult<Arc<LogicalPlan>> {
+        let state = self.state()?;
+        state
+            .cached_relations
+            .get(relation_id)
+            .cloned()
+            .ok_or_else(|| {
+                CatalogError::NotFound(CatalogObject::LogicalPlan, relation_id.to_string())
+            })
+    }
+
+    pub fn remove_cached_relation(&self, relation_id: &str) -> CatalogResult<bool> {
+        let mut state = self.state()?;
+        Ok(state.cached_relations.remove(relation_id).is_some())
+    }
 }

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

@@ -5,9 +5,12 @@ use datafusion::catalog::MemTable;
 use datafusion_common::{DFSchema, DFSchemaRef, ParamValues};
 use datafusion_expr::{EmptyRelation, Extension, LogicalPlan, UNNAMED_TABLE};
 use log::warn;
+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 +141,24 @@ impl PlanResolver<'_> {
         )
     }
 
+    /// Resolves a [`spec::QueryNode::CachedRemoteRelation`] by looking up the
+    /// previously checkpointed logical plan from the catalog manager. The cached
+    /// plan was stored with user-facing field names, so we register fresh field
+    /// IDs in the current resolver state and rename the plan accordingly.
+    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 plan = manager.get_cached_relation(&relation_id).map_err(|_| {
+            PlanError::invalid(format!("cached remote relation not found: {relation_id}"))
+        })?;
+        let names = state.register_fields(plan.schema().inner().fields());
+        let plan = rename_logical_plan((*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-spark-connect/Cargo.toml

@@ -8,6 +8,7 @@ workspace = true
 
 [dependencies]
 sail-cache = { path = "../sail-cache" }
+sail-catalog = { path = "../sail-catalog" }
 sail-common = { path = "../sail-common" }
 sail-common-datafusion = { path = "../sail-common-datafusion" }
 sail-data-source = { path = "../sail-data-source" }

crates/sail-spark-connect/src/server.rs

@@ -100,8 +100,9 @@ async fn handle_command(
         CommandType::CheckpointCommand(checkpoint) => {
             service::handle_execute_checkpoint_command(ctx, checkpoint, metadata).await
         }
-        CommandType::RemoveCachedRemoteRelationCommand(_) => {
-            Err(SparkError::todo("remove cached remote relation command"))
+        CommandType::RemoveCachedRemoteRelationCommand(command) => {
+            service::handle_execute_remove_cached_remote_relation_command(ctx, command, metadata)
+                .await
         }
         CommandType::MergeIntoTableCommand(command) => {
             service::handle_execute_merge_into_table_command(ctx, command, metadata).await

crates/sail-spark-connect/src/service/plan_executor.rs

@@ -1,17 +1,25 @@
 use std::pin::Pin;
+use std::sync::Arc;
 use std::task::{Context, Poll};
 
 use datafusion::arrow::compute::concat_batches;
+use datafusion::catalog::MemTable;
+use datafusion::datasource::provider_as_source;
+use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, UNNAMED_TABLE};
 use datafusion::prelude::SessionContext;
 use fastrace::collector::SpanContext;
 use fastrace::future::FutureExt;
 use fastrace::Span;
 use futures::stream;
 use log::{debug, warn};
+use sail_catalog::manager::CatalogManager;
 use sail_common::spec;
 use sail_common_datafusion::extension::SessionExtensionAccessor;
+use sail_common_datafusion::rename::logical_plan::rename_logical_plan;
 use sail_common_datafusion::session::job::JobService;
 use sail_plan::resolve_and_execute_plan;
+use sail_plan::resolver::plan::NamedPlan;
+use sail_plan::resolver::PlanResolver;
 use tonic::codegen::tokio_stream::wrappers::ReceiverStream;
 use tonic::codegen::tokio_stream::Stream;
 use tonic::Status;
@@ -26,10 +34,11 @@ use crate::spark::connect::execute_plan_response::{
     ResponseType, ResultComplete, SqlCommandResult,
 };
 use crate::spark::connect::{
-    relation, CheckpointCommand, CheckpointCommandResult, CommonInlineUserDefinedDataSource,
-    CommonInlineUserDefinedFunction, CommonInlineUserDefinedTableFunction,
-    CreateDataFrameViewCommand, ExecutePlanResponse, GetResourcesCommand, LocalRelation,
-    MergeIntoTableCommand, Relation, SqlCommand, StreamingQueryCommand,
+    relation, CachedRemoteRelation, CheckpointCommand, CheckpointCommandResult,
+    CommonInlineUserDefinedDataSource, CommonInlineUserDefinedFunction,
+    CommonInlineUserDefinedTableFunction, CreateDataFrameViewCommand, ExecutePlanResponse,
+    GetResourcesCommand, LocalRelation, MergeIntoTableCommand, Relation,
+    RemoveCachedRemoteRelationCommand, SqlCommand, StreamingQueryCommand,
     StreamingQueryCommandResult, StreamingQueryListenerBusCommand, StreamingQueryManagerCommand,
     StreamingQueryManagerCommandResult, WriteOperation, WriteOperationV2,
     WriteStreamOperationStart, WriteStreamOperationStartResult,
@@ -515,13 +524,47 @@ pub(crate) async fn handle_execute_streaming_query_listener_bus_command(
 
 pub(crate) async fn handle_execute_checkpoint_command(
     ctx: &SessionContext,
-    _checkpoint: CheckpointCommand,
+    checkpoint: CheckpointCommand,
     metadata: ExecutorMetadata,
 ) -> SparkResult<ExecutePlanResponseStream> {
-    // TODO: Implement
-    warn!("Checkpoint operation is not yet supported and is a no-op");
+    let CheckpointCommand {
+        relation,
+        local: _,
+        eager,
+        // Sail uses an in-memory store for both local and reliable checkpoints,
+        // so the storage level is currently informational only.
+        storage_level: _,
+    } = checkpoint;
+    let relation = relation.required("checkpoint relation")?;
+    let plan = spec::Plan::Query((relation).try_into()?);
+
     let spark = ctx.extension::<SparkSession>()?;
-    let result = CheckpointCommandResult { relation: None };
+    let resolver = PlanResolver::new(ctx, spark.plan_config()?);
+    let NamedPlan {
+        plan: logical_plan,
+        fields,
+    } = resolver.resolve_named_plan(plan).await?;
+    let logical_plan = if let Some(fields) = fields {
+        rename_logical_plan(logical_plan, &fields)?
+    } else {
+        logical_plan
+    };
+
+    let cached_plan = if eager {
+        materialize_logical_plan(ctx, logical_plan).await?
+    } else {
+        logical_plan
+    };
+
+    let relation_id = uuid::Uuid::new_v4().to_string();
+    let manager = ctx.extension::<CatalogManager>()?;
+    manager
+        .track_cached_relation(relation_id.clone(), Arc::new(cached_plan))
+        .map_err(|e| SparkError::internal(format!("failed to track cached relation: {e}")))?;
+
+    let result = CheckpointCommandResult {
+        relation: Some(CachedRemoteRelation { relation_id }),
+    };
     let mut output = vec![ExecutorOutput::new(ExecutorBatch::CheckpointCommandResult(
         Box::new(result),
     ))];
@@ -535,6 +578,51 @@ pub(crate) async fn handle_execute_checkpoint_command(
     ))
 }
 
+/// Eagerly executes the resolved logical plan, collects all batches, and returns
+/// a new logical plan that scans the materialized data from an in-memory table.
+async fn materialize_logical_plan(
+    ctx: &SessionContext,
+    plan: LogicalPlan,
+) -> SparkResult<LogicalPlan> {
+    let arrow_schema = Arc::new(plan.schema().inner().as_ref().clone());
+    let df = ctx.execute_logical_plan(plan).await?;
+    let batches = df.collect().await?;
+    let table = Arc::new(MemTable::try_new(arrow_schema, vec![batches])?);
+    let scan = LogicalPlanBuilder::scan(
+        datafusion::common::TableReference::bare(UNNAMED_TABLE),
+        provider_as_source(table),
+        None,
+    )?
+    .build()?;
+    Ok(scan)
+}
+
+pub(crate) async fn handle_execute_remove_cached_remote_relation_command(
+    ctx: &SessionContext,
+    command: RemoveCachedRemoteRelationCommand,
+    metadata: ExecutorMetadata,
+) -> SparkResult<ExecutePlanResponseStream> {
+    let RemoveCachedRemoteRelationCommand { relation } = command;
+    let relation = relation.required("remove cached remote relation")?;
+    let manager = ctx.extension::<CatalogManager>()?;
+    // Removing a relation that does not exist is treated as a no-op to match
+    // Spark Connect, which tolerates duplicate or out-of-order cleanup calls.
+    let _existed = manager
+        .remove_cached_relation(&relation.relation_id)
+        .map_err(|e| SparkError::internal(format!("failed to remove cached relation: {e}")))?;
+
+    let spark = ctx.extension::<SparkSession>()?;
+    let mut output = Vec::new();
+    if metadata.reattachable {
+        output.push(ExecutorOutput::complete());
+    }
+    Ok(ExecutePlanResponseStream::new(
+        spark.session_id().to_string(),
+        metadata.operation_id,
+        Box::pin(stream::iter(output)),
+    ))
+}
+
 pub(crate) async fn handle_interrupt_all(ctx: &SessionContext) -> SparkResult<Vec<String>> {
     let spark = ctx.extension::<SparkSession>()?;
     let mut results = vec![];