From 035ad9a5ec103bb60bab3070e40804a18d81ca83 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 01:08:51 +0000 Subject: [PATCH 01/22] feat: implement DataFrame.checkpoint and DataFrame.localCheckpoint Agent-Logs-Url: https://github.com/lakehq/sail/sessions/5f3f46a3-3a38-4586-ac3b-5ff83088166b Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- Cargo.lock | 1 + crates/sail-catalog/src/manager/mod.rs | 16 +++ crates/sail-catalog/src/manager/tracker.rs | 30 +++++ crates/sail-plan/src/resolver/query/misc.rs | 21 ++++ crates/sail-plan/src/resolver/query/mod.rs | 5 +- crates/sail-spark-connect/Cargo.toml | 1 + crates/sail-spark-connect/src/server.rs | 5 +- .../src/service/plan_executor.rs | 102 +++++++++++++-- python/pysail/tests/spark/test_checkpoint.py | 119 ++++++++++++++++++ 9 files changed, 288 insertions(+), 12 deletions(-) create mode 100644 python/pysail/tests/spark/test_checkpoint.py diff --git a/Cargo.lock b/Cargo.lock index 8b2b6bb3c3..66ce79d452 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7666,6 +7666,7 @@ dependencies = [ "quote", "regex", "sail-cache", + "sail-catalog", "sail-common", "sail-common-datafusion", "sail-data-source", diff --git a/crates/sail-catalog/src/manager/mod.rs b/crates/sail-catalog/src/manager/mod.rs index db94464a22..23a372ce67 100644 --- a/crates/sail-catalog/src/manager/mod.rs +++ b/crates/sail-catalog/src/manager/mod.rs @@ -147,6 +147,22 @@ impl CatalogManager { ) -> CatalogResult> { self.tracker.get_tracked_logical_plan(id) } + + pub fn track_cached_relation( + &self, + relation_id: String, + plan: Arc, + ) -> CatalogResult<()> { + self.tracker.track_cached_relation(relation_id, plan) + } + + pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult> { + self.tracker.get_cached_relation(relation_id) + } + + pub fn remove_cached_relation(&self, relation_id: &str) -> CatalogResult { + self.tracker.remove_cached_relation(relation_id) + } } impl CatalogManagerState { diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index b7549055c3..0ac3f60f45 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -18,6 +18,7 @@ struct CatalogObjectTrackerState { next_logical_plan_id: u64, functions: HashMap, logical_plans: HashMap>, + cached_relations: HashMap>, } /// Tracks in-memory objects (UDFs and logical plans) that cannot be serialized directly, @@ -75,4 +76,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, + ) -> CatalogResult<()> { + let mut state = self.state()?; + state.cached_relations.insert(relation_id, plan); + Ok(()) + } + + pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult> { + 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 { + let mut state = self.state()?; + Ok(state.cached_relations.remove(relation_id).is_some()) + } } diff --git a/crates/sail-plan/src/resolver/query/misc.rs b/crates/sail-plan/src/resolver/query/misc.rs index a49c825d8c..8c3fd95763 100644 --- a/crates/sail-plan/src/resolver/query/misc.rs +++ b/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 { + let manager = self.ctx.extension::()?; + 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, diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index 13dbb0a01a..8177c4c1d9 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/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? diff --git a/crates/sail-spark-connect/Cargo.toml b/crates/sail-spark-connect/Cargo.toml index 68045d157a..5730cff7bc 100644 --- a/crates/sail-spark-connect/Cargo.toml +++ b/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" } diff --git a/crates/sail-spark-connect/src/server.rs b/crates/sail-spark-connect/src/server.rs index 9c47c96bfd..c92981f0e8 100644 --- a/crates/sail-spark-connect/src/server.rs +++ b/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 diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index bf2a30e617..a14ad0ed88 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/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 { - // 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::()?; - 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::()?; + 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,49 @@ 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 { + 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 { + let RemoveCachedRemoteRelationCommand { relation } = command; + let relation = relation.required("remove cached remote relation")?; + let manager = ctx.extension::()?; + let _ = manager + .remove_cached_relation(&relation.relation_id) + .map_err(|e| SparkError::internal(format!("failed to remove cached relation: {e}")))?; + + let spark = ctx.extension::()?; + 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> { let spark = ctx.extension::()?; let mut results = vec![]; diff --git a/python/pysail/tests/spark/test_checkpoint.py b/python/pysail/tests/spark/test_checkpoint.py new file mode 100644 index 0000000000..98cdb8e5e4 --- /dev/null +++ b/python/pysail/tests/spark/test_checkpoint.py @@ -0,0 +1,119 @@ +import pandas as pd +import pytest +from pandas.testing import assert_frame_equal +from pyspark import StorageLevel +from pyspark.sql import Row +from pyspark.sql.functions import col, lit + + +def _roundtrip(df): + return df.sort(*df.columns).toPandas() + + +@pytest.mark.parametrize("eager", [True, False]) +def test_dataframe_local_checkpoint(spark, eager): + df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + checkpointed = df.localCheckpoint(eager) + + # The checkpointed DataFrame preserves the schema of the source DataFrame. + assert checkpointed.schema == df.schema + assert checkpointed.columns == ["age", "name"] + assert repr(checkpointed) == "DataFrame[age: bigint, name: string]" + + # The checkpointed DataFrame returns the same rows as the source DataFrame. + assert_frame_equal( + _roundtrip(checkpointed), + pd.DataFrame({"age": [14, 16, 23], "name": ["Tom", "Bob", "Alice"]}).sort_values( + ["age", "name"], ignore_index=True + ), + ) + + +@pytest.mark.parametrize("eager", [True, False]) +def test_dataframe_checkpoint(spark, eager): + df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["id", "value"]) + checkpointed = df.checkpoint(eager) + + assert checkpointed.schema == df.schema + assert repr(checkpointed) == "DataFrame[id: bigint, value: string]" + assert_frame_equal( + _roundtrip(checkpointed), + pd.DataFrame({"id": [1, 2, 3], "value": ["a", "b", "c"]}), + ) + + +def test_dataframe_local_checkpoint_default_eager(spark): + df = spark.createDataFrame([(1,)], ["x"]) + # Default `eager=True` should also work. + checkpointed = df.localCheckpoint() + assert checkpointed.columns == ["x"] + assert_frame_equal(checkpointed.toPandas(), pd.DataFrame({"x": [1]})) + + +def test_dataframe_local_checkpoint_with_storage_level(spark): + df = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "name"]) + # The storage level argument is accepted and the result is a usable DataFrame. + checkpointed = df.localCheckpoint(eager=True, storageLevel=StorageLevel.MEMORY_AND_DISK) + assert checkpointed.columns == ["id", "name"] + assert_frame_equal( + _roundtrip(checkpointed), + pd.DataFrame({"id": [1, 2], "name": ["a", "b"]}), + ) + + +def test_dataframe_checkpoint_after_transformation(spark): + df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) + transformed = df.filter(col("age") > lit(15)).withColumn("country", lit("US")) + checkpointed = transformed.localCheckpoint() + + assert checkpointed.columns == ["age", "name", "country"] + assert_frame_equal( + _roundtrip(checkpointed), + pd.DataFrame({"age": [16, 23], "name": ["Bob", "Alice"], "country": ["US", "US"]}).sort_values( + ["age", "country", "name"], ignore_index=True + ), + ) + + # The checkpointed DataFrame can be used in further operations. + aggregated = checkpointed.groupBy("country").count().toPandas() + assert_frame_equal( + aggregated.sort_values("country", ignore_index=True), + pd.DataFrame({"country": ["US"], "count": [2]}), + ) + + +def test_dataframe_local_checkpoint_with_nulls(spark): + df = spark.createDataFrame([Row(a=1, b="x"), Row(a=None, b="y"), Row(a=3, b=None)]) + checkpointed = df.localCheckpoint() + pdf = checkpointed.sort("b").toPandas() + # Null values are preserved through the checkpoint. + assert pdf["a"].tolist()[:2] == [1, None] or pdf["a"].isna().sum() == 1 + assert pdf["b"].isna().sum() == 1 + + +def test_dataframe_local_checkpoint_join(spark): + left = spark.createDataFrame([(1, "a"), (2, "b")], ["id", "name"]) + right = spark.createDataFrame([(1, 100), (2, 200)], ["id", "value"]) + checkpointed = left.localCheckpoint() + joined = checkpointed.join(right, "id").toPandas() + assert_frame_equal( + joined.sort_values("id", ignore_index=True), + pd.DataFrame({"id": [1, 2], "name": ["a", "b"], "value": [100, 200]}), + ) + + +def test_dataframe_local_checkpoint_empty(spark): + df = spark.createDataFrame([], "id long, name string") + checkpointed = df.localCheckpoint() + assert checkpointed.columns == ["id", "name"] + assert checkpointed.count() == 0 + + +def test_dataframe_local_checkpoint_chained(spark): + df = spark.createDataFrame([(1,), (2,), (3,)], ["x"]) + cp1 = df.localCheckpoint() + cp2 = cp1.filter(col("x") > 1).localCheckpoint() + assert_frame_equal( + cp2.sort("x").toPandas(), + pd.DataFrame({"x": [2, 3]}), + ) From 0ad35c2bed2a3c438d32a860722d662ddb82b044 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 01:14:44 +0000 Subject: [PATCH 02/22] test: clarify null preservation assertions; doc cached_relations Agent-Logs-Url: https://github.com/lakehq/sail/sessions/5f3f46a3-3a38-4586-ac3b-5ff83088166b Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-catalog/src/manager/tracker.rs | 3 +++ crates/sail-spark-connect/src/service/plan_executor.rs | 4 +++- python/pysail/tests/spark/test_checkpoint.py | 5 +++-- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index 0ac3f60f45..dbe1910567 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -18,6 +18,9 @@ struct CatalogObjectTrackerState { next_logical_plan_id: u64, functions: HashMap, logical_plans: HashMap>, + /// 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>, } diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index a14ad0ed88..0cf367381e 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -605,7 +605,9 @@ pub(crate) async fn handle_execute_remove_cached_remote_relation_command( let RemoveCachedRemoteRelationCommand { relation } = command; let relation = relation.required("remove cached remote relation")?; let manager = ctx.extension::()?; - let _ = manager + // 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}")))?; diff --git a/python/pysail/tests/spark/test_checkpoint.py b/python/pysail/tests/spark/test_checkpoint.py index 98cdb8e5e4..39aad36b31 100644 --- a/python/pysail/tests/spark/test_checkpoint.py +++ b/python/pysail/tests/spark/test_checkpoint.py @@ -85,9 +85,10 @@ def test_dataframe_checkpoint_after_transformation(spark): def test_dataframe_local_checkpoint_with_nulls(spark): df = spark.createDataFrame([Row(a=1, b="x"), Row(a=None, b="y"), Row(a=3, b=None)]) checkpointed = df.localCheckpoint() - pdf = checkpointed.sort("b").toPandas() + pdf = checkpointed.toPandas() # Null values are preserved through the checkpoint. - assert pdf["a"].tolist()[:2] == [1, None] or pdf["a"].isna().sum() == 1 + assert len(pdf) == 3 + assert pdf["a"].isna().sum() == 1 assert pdf["b"].isna().sum() == 1 From 262e7f79b0de2886643a2e0ad29765fc05858cc2 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 01:42:38 +0000 Subject: [PATCH 03/22] fix: preserve duplicate columns in checkpointed relations Agent-Logs-Url: https://github.com/lakehq/sail/sessions/8133bd14-cb94-4c56-ac57-5629cfed18d9 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-catalog/src/manager/mod.rs | 10 ++++---- crates/sail-catalog/src/manager/tracker.rs | 19 ++++++++++----- crates/sail-plan/src/resolver/query/misc.rs | 23 +++++++++++++------ .../src/service/plan_executor.rs | 21 ++++++++++++----- python/pysail/tests/spark/test_checkpoint.py | 15 +++++++++++- 5 files changed, 64 insertions(+), 24 deletions(-) diff --git a/crates/sail-catalog/src/manager/mod.rs b/crates/sail-catalog/src/manager/mod.rs index 23a372ce67..ec3cc10e69 100644 --- a/crates/sail-catalog/src/manager/mod.rs +++ b/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; @@ -151,12 +153,12 @@ impl CatalogManager { pub fn track_cached_relation( &self, relation_id: String, - plan: Arc, + relation: CatalogCachedRelation, ) -> CatalogResult<()> { - self.tracker.track_cached_relation(relation_id, plan) + self.tracker.track_cached_relation(relation_id, relation) } - pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult> { + pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult { self.tracker.get_cached_relation(relation_id) } diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index dbe1910567..9bd5cc43cf 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -12,6 +12,12 @@ 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, + pub fields: Vec, +} + #[derive(Default)] struct CatalogObjectTrackerState { next_function_id: u64, @@ -19,9 +25,10 @@ struct CatalogObjectTrackerState { functions: HashMap, logical_plans: HashMap>, /// 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>, + /// logical plans and user-facing field names. Populated when the server + /// handles a `CheckpointCommand`, and queried when resolving a + /// `CachedRemoteRelation` query node. + cached_relations: HashMap, } /// Tracks in-memory objects (UDFs and logical plans) that cannot be serialized directly, @@ -86,14 +93,14 @@ impl CatalogObjectTracker { pub fn track_cached_relation( &self, relation_id: String, - plan: Arc, + relation: CatalogCachedRelation, ) -> CatalogResult<()> { let mut state = self.state()?; - state.cached_relations.insert(relation_id, plan); + state.cached_relations.insert(relation_id, relation); Ok(()) } - pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult> { + pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult { let state = self.state()?; state .cached_relations diff --git a/crates/sail-plan/src/resolver/query/misc.rs b/crates/sail-plan/src/resolver/query/misc.rs index 8c3fd95763..3097cef819 100644 --- a/crates/sail-plan/src/resolver/query/misc.rs +++ b/crates/sail-plan/src/resolver/query/misc.rs @@ -5,6 +5,7 @@ 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}; @@ -143,19 +144,27 @@ 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. + /// 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 { let manager = self.ctx.extension::()?; - 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)?; + let relation = match manager.get_cached_relation(&relation_id) { + Ok(relation) => relation, + Err(CatalogError::NotFound(CatalogObject::LogicalPlan, _)) => { + 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) } diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 0cf367381e..1033930adc 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -12,10 +12,10 @@ use fastrace::future::FutureExt; use fastrace::Span; use futures::stream; use log::{debug, warn}; +use sail_catalog::manager::tracker::CatalogCachedRelation; 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; @@ -544,11 +544,14 @@ pub(crate) async fn handle_execute_checkpoint_command( 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 { + let fields = fields.unwrap_or_else(|| { logical_plan - }; + .schema() + .fields() + .iter() + .map(|field| field.name().to_string()) + .collect() + }); let cached_plan = if eager { materialize_logical_plan(ctx, logical_plan).await? @@ -559,7 +562,13 @@ pub(crate) async fn handle_execute_checkpoint_command( let relation_id = uuid::Uuid::new_v4().to_string(); let manager = ctx.extension::()?; manager - .track_cached_relation(relation_id.clone(), Arc::new(cached_plan)) + .track_cached_relation( + relation_id.clone(), + CatalogCachedRelation { + plan: Arc::new(cached_plan), + fields, + }, + ) .map_err(|e| SparkError::internal(format!("failed to track cached relation: {e}")))?; let result = CheckpointCommandResult { diff --git a/python/pysail/tests/spark/test_checkpoint.py b/python/pysail/tests/spark/test_checkpoint.py index 39aad36b31..c620ca68a3 100644 --- a/python/pysail/tests/spark/test_checkpoint.py +++ b/python/pysail/tests/spark/test_checkpoint.py @@ -61,6 +61,18 @@ def test_dataframe_local_checkpoint_with_storage_level(spark): ) +@pytest.mark.parametrize("method", ["checkpoint", "localCheckpoint"]) +@pytest.mark.parametrize("eager", [True, False]) +def test_dataframe_checkpoint_with_duplicate_columns(spark, method, eager): + df = spark.createDataFrame([(1, 2), (3, 4)], ["left", "right"]).selectExpr("left as id", "right as id") + checkpointed = getattr(df, method)(eager) + pdf = checkpointed.toPandas() + + assert checkpointed.columns == ["id", "id"] + assert pdf.columns.tolist() == ["id", "id"] + assert pdf.values.tolist() == [[1, 2], [3, 4]] + + def test_dataframe_checkpoint_after_transformation(spark): df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) transformed = df.filter(col("age") > lit(15)).withColumn("country", lit("US")) @@ -83,11 +95,12 @@ def test_dataframe_checkpoint_after_transformation(spark): def test_dataframe_local_checkpoint_with_nulls(spark): + expected_rows = 3 df = spark.createDataFrame([Row(a=1, b="x"), Row(a=None, b="y"), Row(a=3, b=None)]) checkpointed = df.localCheckpoint() pdf = checkpointed.toPandas() # Null values are preserved through the checkpoint. - assert len(pdf) == 3 + assert len(pdf) == expected_rows assert pdf["a"].isna().sum() == 1 assert pdf["b"].isna().sum() == 1 From b4fa7b48b14cf22af4dfa76434b76201cee49902 Mon Sep 17 00:00:00 2001 From: Shehab Amin <11789402+shehabgamin@users.noreply.github.com> Date: Sun, 10 May 2026 02:34:52 +0000 Subject: [PATCH 04/22] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- crates/sail-spark-connect/src/service/plan_executor.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 1033930adc..24e53368e8 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -593,7 +593,7 @@ async fn materialize_logical_plan( ctx: &SessionContext, plan: LogicalPlan, ) -> SparkResult { - let arrow_schema = Arc::new(plan.schema().inner().as_ref().clone()); + let arrow_schema = plan.schema().inner().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])?); From 0fecd128dc0642635d8dc2a33f0ca2237a546adb Mon Sep 17 00:00:00 2001 From: Shehab Amin <11789402+shehabgamin@users.noreply.github.com> Date: Sun, 10 May 2026 02:55:34 +0000 Subject: [PATCH 05/22] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- crates/sail-catalog/src/manager/tracker.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index 9bd5cc43cf..dcf1874016 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -107,7 +107,10 @@ impl CatalogObjectTracker { .get(relation_id) .cloned() .ok_or_else(|| { - CatalogError::NotFound(CatalogObject::LogicalPlan, relation_id.to_string()) + CatalogError::Internal(format!( + "Cached relation not found: {}", + relation_id + )) }) } From 5148efd410b07b04ff29c4b7c951415cee96eeb9 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 04:57:52 +0000 Subject: [PATCH 06/22] fix: checkpoint storage parity and formatting Agent-Logs-Url: https://github.com/lakehq/sail/sessions/aa2d0845-e4a5-4df6-96f7-42b8415124c8 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-catalog/src/manager/tracker.rs | 5 +- .../src/service/plan_executor.rs | 49 ++++++++++++++++--- python/pysail/tests/spark/test_checkpoint.py | 7 +++ 3 files changed, 49 insertions(+), 12 deletions(-) diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index dcf1874016..9bd5cc43cf 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -107,10 +107,7 @@ impl CatalogObjectTracker { .get(relation_id) .cloned() .ok_or_else(|| { - CatalogError::Internal(format!( - "Cached relation not found: {}", - relation_id - )) + CatalogError::NotFound(CatalogObject::LogicalPlan, relation_id.to_string()) }) } diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 24e53368e8..826d12a61e 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -4,9 +4,10 @@ use std::task::{Context, Poll}; use datafusion::arrow::compute::concat_batches; use datafusion::catalog::MemTable; +use datafusion::dataframe::DataFrameWriteOptions; use datafusion::datasource::provider_as_source; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, UNNAMED_TABLE}; -use datafusion::prelude::SessionContext; +use datafusion::prelude::{ParquetReadOptions, SessionContext}; use fastrace::collector::SpanContext; use fastrace::future::FutureExt; use fastrace::Span; @@ -529,11 +530,9 @@ pub(crate) async fn handle_execute_checkpoint_command( ) -> SparkResult { let CheckpointCommand { relation, - local: _, + local, eager, - // Sail uses an in-memory store for both local and reliable checkpoints, - // so the storage level is currently informational only. - storage_level: _, + storage_level, } = checkpoint; let relation = relation.required("checkpoint relation")?; let plan = spec::Plan::Query((relation).try_into()?); @@ -553,8 +552,11 @@ pub(crate) async fn handle_execute_checkpoint_command( .collect() }); - let cached_plan = if eager { - materialize_logical_plan(ctx, logical_plan).await? + let use_disk = checkpoint_uses_disk(local, storage_level.as_ref()); + let cached_plan = if eager && use_disk { + materialize_logical_plan_to_disk(ctx, logical_plan).await? + } else if eager { + materialize_logical_plan_to_memory(ctx, logical_plan).await? } else { logical_plan }; @@ -587,9 +589,21 @@ pub(crate) async fn handle_execute_checkpoint_command( )) } +fn checkpoint_uses_disk( + local: bool, + storage_level: Option<&crate::spark::connect::StorageLevel>, +) -> bool { + if !local { + return true; + } + storage_level + .map(|level| level.use_disk && !level.use_memory) + .unwrap_or(false) +} + /// 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( +async fn materialize_logical_plan_to_memory( ctx: &SessionContext, plan: LogicalPlan, ) -> SparkResult { @@ -606,6 +620,25 @@ async fn materialize_logical_plan( Ok(scan) } +/// Eagerly executes the resolved logical plan, writes the result to local +/// checkpoint files, and returns a new logical plan that scans those files. +async fn materialize_logical_plan_to_disk( + ctx: &SessionContext, + plan: LogicalPlan, +) -> SparkResult { + let checkpoint_path = std::env::temp_dir() + .join("sail-checkpoints") + .join(uuid::Uuid::new_v4().to_string()); + let checkpoint_path = checkpoint_path.to_string_lossy().into_owned(); + let df = ctx.execute_logical_plan(plan).await?; + df.write_parquet(&checkpoint_path, DataFrameWriteOptions::new(), None) + .await?; + let df = ctx + .read_parquet(&checkpoint_path, ParquetReadOptions::default()) + .await?; + Ok(df.into_unoptimized_plan()) +} + pub(crate) async fn handle_execute_remove_cached_remote_relation_command( ctx: &SessionContext, command: RemoveCachedRemoteRelationCommand, diff --git a/python/pysail/tests/spark/test_checkpoint.py b/python/pysail/tests/spark/test_checkpoint.py index c620ca68a3..ecab5a9256 100644 --- a/python/pysail/tests/spark/test_checkpoint.py +++ b/python/pysail/tests/spark/test_checkpoint.py @@ -5,6 +5,13 @@ from pyspark.sql import Row from pyspark.sql.functions import col, lit +from pysail.testing.spark.utils.common import pyspark_version + +pytestmark = pytest.mark.skipif( + pyspark_version() < (4,), + reason="DataFrame checkpoint APIs require PySpark 4+ in Spark Connect", +) + def _roundtrip(df): return df.sort(*df.columns).toPandas() From 3b1bdc630ba46a31e85e9d44ee6d2d97413ccad8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 05:12:37 +0000 Subject: [PATCH 07/22] fix: create checkpoint storage directory Agent-Logs-Url: https://github.com/lakehq/sail/sessions/aa2d0845-e4a5-4df6-96f7-42b8415124c8 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-plan/src/resolver/query/misc.rs | 13 +++++++------ .../sail-spark-connect/src/service/plan_executor.rs | 7 ++++--- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/crates/sail-plan/src/resolver/query/misc.rs b/crates/sail-plan/src/resolver/query/misc.rs index 3097cef819..7c39254e14 100644 --- a/crates/sail-plan/src/resolver/query/misc.rs +++ b/crates/sail-plan/src/resolver/query/misc.rs @@ -142,12 +142,13 @@ 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 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. + /// 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, diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 826d12a61e..06f5683494 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -626,9 +626,10 @@ async fn materialize_logical_plan_to_disk( ctx: &SessionContext, plan: LogicalPlan, ) -> SparkResult { - let checkpoint_path = std::env::temp_dir() - .join("sail-checkpoints") - .join(uuid::Uuid::new_v4().to_string()); + let checkpoint_root = std::env::temp_dir().join("sail-checkpoints"); + std::fs::create_dir_all(&checkpoint_root) + .map_err(|e| SparkError::internal(format!("failed to create checkpoint directory: {e}")))?; + let checkpoint_path = checkpoint_root.join(uuid::Uuid::new_v4().to_string()); let checkpoint_path = checkpoint_path.to_string_lossy().into_owned(); let df = ctx.execute_logical_plan(plan).await?; df.write_parquet(&checkpoint_path, DataFrameWriteOptions::new(), None) From 6246edc93341d5fff9f535b81eb7f6ee9047cb66 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 06:22:06 +0000 Subject: [PATCH 08/22] fix: clean up checkpoint storage Agent-Logs-Url: https://github.com/lakehq/sail/sessions/8a18d45f-499a-43ed-9b73-c1940913d938 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-catalog/src/error.rs | 2 + crates/sail-catalog/src/manager/mod.rs | 5 +- crates/sail-catalog/src/manager/tracker.rs | 11 ++- crates/sail-plan/src/resolver/query/misc.rs | 2 +- .../src/service/plan_executor.rs | 87 +++++++++++++------ 5 files changed, 77 insertions(+), 30 deletions(-) diff --git a/crates/sail-catalog/src/error.rs b/crates/sail-catalog/src/error.rs index 6669a40411..130f58029f 100644 --- a/crates/sail-catalog/src/error.rs +++ b/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}") } diff --git a/crates/sail-catalog/src/manager/mod.rs b/crates/sail-catalog/src/manager/mod.rs index ec3cc10e69..9f1405b2df 100644 --- a/crates/sail-catalog/src/manager/mod.rs +++ b/crates/sail-catalog/src/manager/mod.rs @@ -162,7 +162,10 @@ impl CatalogManager { self.tracker.get_cached_relation(relation_id) } - pub fn remove_cached_relation(&self, relation_id: &str) -> CatalogResult { + pub fn remove_cached_relation( + &self, + relation_id: &str, + ) -> CatalogResult> { self.tracker.remove_cached_relation(relation_id) } } diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index 9bd5cc43cf..862217a9f5 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -1,4 +1,5 @@ use std::collections::HashMap; +use std::path::PathBuf; use std::sync::{Arc, Mutex, MutexGuard}; use datafusion_expr::{LogicalPlan, ScalarUDF}; @@ -16,6 +17,7 @@ pub struct CatalogLogicalPlanId(u64); pub struct CatalogCachedRelation { pub plan: Arc, pub fields: Vec, + pub storage_path: Option, } #[derive(Default)] @@ -107,12 +109,15 @@ impl CatalogObjectTracker { .get(relation_id) .cloned() .ok_or_else(|| { - CatalogError::NotFound(CatalogObject::LogicalPlan, relation_id.to_string()) + CatalogError::NotFound(CatalogObject::CachedRelation, relation_id.to_string()) }) } - pub fn remove_cached_relation(&self, relation_id: &str) -> CatalogResult { + pub fn remove_cached_relation( + &self, + relation_id: &str, + ) -> CatalogResult> { let mut state = self.state()?; - Ok(state.cached_relations.remove(relation_id).is_some()) + Ok(state.cached_relations.remove(relation_id)) } } diff --git a/crates/sail-plan/src/resolver/query/misc.rs b/crates/sail-plan/src/resolver/query/misc.rs index 7c39254e14..e65cdbc940 100644 --- a/crates/sail-plan/src/resolver/query/misc.rs +++ b/crates/sail-plan/src/resolver/query/misc.rs @@ -157,7 +157,7 @@ impl PlanResolver<'_> { let manager = self.ctx.extension::()?; let relation = match manager.get_cached_relation(&relation_id) { Ok(relation) => relation, - Err(CatalogError::NotFound(CatalogObject::LogicalPlan, _)) => { + Err(CatalogError::NotFound(CatalogObject::CachedRelation, _)) => { return Err(PlanError::AnalysisError(format!( "cached remote relation not found: {relation_id}" ))); diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 06f5683494..b3989e83a3 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -1,3 +1,4 @@ +use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -553,25 +554,34 @@ pub(crate) async fn handle_execute_checkpoint_command( }); let use_disk = checkpoint_uses_disk(local, storage_level.as_ref()); - let cached_plan = if eager && use_disk { + let (cached_plan, storage_path) = if eager && use_disk { materialize_logical_plan_to_disk(ctx, logical_plan).await? } else if eager { - materialize_logical_plan_to_memory(ctx, logical_plan).await? + ( + materialize_logical_plan_to_memory(ctx, logical_plan).await?, + None, + ) } else { - logical_plan + (logical_plan, None) }; let relation_id = uuid::Uuid::new_v4().to_string(); let manager = ctx.extension::()?; - manager - .track_cached_relation( - relation_id.clone(), - CatalogCachedRelation { - plan: Arc::new(cached_plan), - fields, - }, - ) - .map_err(|e| SparkError::internal(format!("failed to track cached relation: {e}")))?; + if let Err(error) = manager.track_cached_relation( + relation_id.clone(), + CatalogCachedRelation { + plan: Arc::new(cached_plan), + fields, + storage_path: storage_path.clone(), + }, + ) { + if let Some(path) = storage_path { + cleanup_checkpoint_path(&path).await; + } + return Err(SparkError::internal(format!( + "failed to track cached relation: {error}" + ))); + } let result = CheckpointCommandResult { relation: Some(CachedRemoteRelation { relation_id }), @@ -596,9 +606,7 @@ fn checkpoint_uses_disk( if !local { return true; } - storage_level - .map(|level| level.use_disk && !level.use_memory) - .unwrap_or(false) + storage_level.map(|level| level.use_disk).unwrap_or(false) } /// Eagerly executes the resolved logical plan, collects all batches, and returns @@ -625,19 +633,45 @@ async fn materialize_logical_plan_to_memory( async fn materialize_logical_plan_to_disk( ctx: &SessionContext, plan: LogicalPlan, -) -> SparkResult { +) -> SparkResult<(LogicalPlan, Option)> { let checkpoint_root = std::env::temp_dir().join("sail-checkpoints"); - std::fs::create_dir_all(&checkpoint_root) + tokio::fs::create_dir_all(&checkpoint_root) + .await .map_err(|e| SparkError::internal(format!("failed to create checkpoint directory: {e}")))?; let checkpoint_path = checkpoint_root.join(uuid::Uuid::new_v4().to_string()); - let checkpoint_path = checkpoint_path.to_string_lossy().into_owned(); + let checkpoint_path_string = checkpoint_path.to_string_lossy().into_owned(); let df = ctx.execute_logical_plan(plan).await?; - df.write_parquet(&checkpoint_path, DataFrameWriteOptions::new(), None) - .await?; - let df = ctx - .read_parquet(&checkpoint_path, ParquetReadOptions::default()) - .await?; - Ok(df.into_unoptimized_plan()) + if let Err(error) = df + .write_parquet(&checkpoint_path_string, DataFrameWriteOptions::new(), None) + .await + { + cleanup_checkpoint_path(&checkpoint_path).await; + return Err(error.into()); + } + let df = match ctx + .read_parquet(&checkpoint_path_string, ParquetReadOptions::default()) + .await + { + Ok(df) => df, + Err(error) => { + cleanup_checkpoint_path(&checkpoint_path).await; + return Err(error.into()); + } + }; + Ok((df.into_unoptimized_plan(), Some(checkpoint_path))) +} + +async fn cleanup_checkpoint_path(path: &Path) { + match tokio::fs::remove_dir_all(path).await { + Ok(()) => {} + Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} + Err(e) => { + warn!( + "failed to remove checkpoint directory {}: {e}", + path.display() + ); + } + } } pub(crate) async fn handle_execute_remove_cached_remote_relation_command( @@ -650,9 +684,12 @@ pub(crate) async fn handle_execute_remove_cached_remote_relation_command( let manager = ctx.extension::()?; // 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 + let removed = manager .remove_cached_relation(&relation.relation_id) .map_err(|e| SparkError::internal(format!("failed to remove cached relation: {e}")))?; + if let Some(path) = removed.and_then(|relation| relation.storage_path) { + cleanup_checkpoint_path(&path).await; + } let spark = ctx.extension::()?; let mut output = Vec::new(); From 8f808189e7707a29f09f4ec757e2c9f7b10616aa Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 06:35:43 +0000 Subject: [PATCH 09/22] docs: clarify checkpoint storage lifecycle Agent-Logs-Url: https://github.com/lakehq/sail/sessions/8a18d45f-499a-43ed-9b73-c1940913d938 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-catalog/src/manager/tracker.rs | 3 ++- .../src/service/plan_executor.rs | 22 ++++++++++++------- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index 862217a9f5..2e5b70444e 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -29,7 +29,8 @@ struct CatalogObjectTrackerState { /// 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. + /// `CachedRemoteRelation` query node. Entries are removed when the server + /// handles a `RemoveCachedRemoteRelationCommand`. cached_relations: HashMap, } diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index b3989e83a3..eabbc7b34e 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -544,14 +544,7 @@ pub(crate) async fn handle_execute_checkpoint_command( plan: logical_plan, fields, } = resolver.resolve_named_plan(plan).await?; - let fields = fields.unwrap_or_else(|| { - logical_plan - .schema() - .fields() - .iter() - .map(|field| field.name().to_string()) - .collect() - }); + let fields = named_plan_fields(&logical_plan, fields); let use_disk = checkpoint_uses_disk(local, storage_level.as_ref()); let (cached_plan, storage_path) = if eager && use_disk { @@ -599,6 +592,16 @@ pub(crate) async fn handle_execute_checkpoint_command( )) } +fn named_plan_fields(plan: &LogicalPlan, fields: Option>) -> Vec { + fields.unwrap_or_else(|| { + plan.schema() + .fields() + .iter() + .map(|field| field.name().to_string()) + .collect() + }) +} + fn checkpoint_uses_disk( local: bool, storage_level: Option<&crate::spark::connect::StorageLevel>, @@ -634,6 +637,9 @@ async fn materialize_logical_plan_to_disk( ctx: &SessionContext, plan: LogicalPlan, ) -> SparkResult<(LogicalPlan, Option)> { + // Spark Connect does not send a user checkpoint directory with this command, + // so Sail stores reliable checkpoint files in a per-server temporary area + // and deletes them when the client releases the cached remote relation. let checkpoint_root = std::env::temp_dir().join("sail-checkpoints"); tokio::fs::create_dir_all(&checkpoint_root) .await From 72e3002022764b4a1304246798a55aad37869ee6 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 07:46:29 +0000 Subject: [PATCH 10/22] fix: clean up checkpoint storage on session drop Agent-Logs-Url: https://github.com/lakehq/sail/sessions/fc16dcc5-b33f-49a6-82dc-d97f81833140 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-catalog/src/manager/tracker.rs | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index 2e5b70444e..00835cbdd4 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -122,3 +122,23 @@ impl CatalogObjectTracker { Ok(state.cached_relations.remove(relation_id)) } } + +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(path) = relation.storage_path else { + continue; + }; + // `Drop` runs in a sync context, so we use blocking filesystem calls. + // Errors are ignored because the caller cannot react to them. + let _ = std::fs::remove_dir_all(&path); + } + } +} From 43ad6635cb357b7037184d5320b7caedb72588cb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 08:51:54 +0000 Subject: [PATCH 11/22] fix: default local checkpoint to disk storage to avoid OOM Agent-Logs-Url: https://github.com/lakehq/sail/sessions/93e14e98-56b8-44c1-9124-b434406e9f21 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-spark-connect/src/service/plan_executor.rs | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index eabbc7b34e..17f043732e 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -602,6 +602,14 @@ fn named_plan_fields(plan: &LogicalPlan, fields: Option>) -> Vec, @@ -609,7 +617,7 @@ fn checkpoint_uses_disk( if !local { return true; } - storage_level.map(|level| level.use_disk).unwrap_or(false) + storage_level.map(|level| level.use_disk).unwrap_or(true) } /// Eagerly executes the resolved logical plan, collects all batches, and returns From 0fb4ff50fed73cfaf7423fb89201ca74382810e8 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 09:13:46 +0000 Subject: [PATCH 12/22] fix: handle empty DataFrame in disk checkpoint materialization Agent-Logs-Url: https://github.com/lakehq/sail/sessions/17107c00-6af1-463c-9c90-b4edbd2d452a Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- .../src/service/plan_executor.rs | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 17f043732e..86a63e30f5 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -654,6 +654,7 @@ async fn materialize_logical_plan_to_disk( .map_err(|e| SparkError::internal(format!("failed to create checkpoint directory: {e}")))?; let checkpoint_path = checkpoint_root.join(uuid::Uuid::new_v4().to_string()); let checkpoint_path_string = checkpoint_path.to_string_lossy().into_owned(); + let arrow_schema = plan.schema().inner().clone(); let df = ctx.execute_logical_plan(plan).await?; if let Err(error) = df .write_parquet(&checkpoint_path_string, DataFrameWriteOptions::new(), None) @@ -662,6 +663,23 @@ async fn materialize_logical_plan_to_disk( cleanup_checkpoint_path(&checkpoint_path).await; return Err(error.into()); } + // DataFusion's `write_parquet` may not create the checkpoint directory when + // the DataFrame is empty (zero rows). In that case we fall back to an + // in-memory table so that downstream scans see the correct (empty) schema + // instead of failing with a "path does not exist" error. + let path_exists = tokio::fs::try_exists(&checkpoint_path) + .await + .unwrap_or(false); + if !path_exists { + let table = Arc::new(MemTable::try_new(arrow_schema, vec![vec![]])?); + let scan = LogicalPlanBuilder::scan( + datafusion::common::TableReference::bare(UNNAMED_TABLE), + provider_as_source(table), + None, + )? + .build()?; + return Ok((scan, None)); + } let df = match ctx .read_parquet(&checkpoint_path_string, ParquetReadOptions::default()) .await From 1b69b9640fb555d61c57ecccd58c1d3e7698c8f0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sun, 10 May 2026 09:17:06 +0000 Subject: [PATCH 13/22] fix: improve error handling for checkpoint path existence check Agent-Logs-Url: https://github.com/lakehq/sail/sessions/17107c00-6af1-463c-9c90-b4edbd2d452a Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- .../sail-spark-connect/src/service/plan_executor.rs | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 86a63e30f5..af7ab7c9b8 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -667,9 +667,16 @@ async fn materialize_logical_plan_to_disk( // the DataFrame is empty (zero rows). In that case we fall back to an // in-memory table so that downstream scans see the correct (empty) schema // instead of failing with a "path does not exist" error. - let path_exists = tokio::fs::try_exists(&checkpoint_path) - .await - .unwrap_or(false); + let path_exists = match tokio::fs::try_exists(&checkpoint_path).await { + Ok(exists) => exists, + Err(e) => { + warn!( + "failed to check checkpoint path {}: {e}; assuming it exists", + checkpoint_path.display() + ); + true + } + }; if !path_exists { let table = Arc::new(MemTable::try_new(arrow_schema, vec![vec![]])?); let scan = LogicalPlanBuilder::scan( From f5341dae2964d454238a7032b88e97353ee42628 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 12 May 2026 21:41:15 +0000 Subject: [PATCH 14/22] fix: skip checkpoint doctests for PySpark 3.5.x and handle empty DataFrame disk checkpoint Agent-Logs-Url: https://github.com/lakehq/sail/sessions/044174a0-bf6d-4f01-9f74-dd0b6c6703aa Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- .../src/service/plan_executor.rs | 40 +++++++++---------- scripts/spark-tests/plugins/spark.py | 23 +++++++++++ 2 files changed, 41 insertions(+), 22 deletions(-) diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index af7ab7c9b8..22d45035fc 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -656,28 +656,13 @@ async fn materialize_logical_plan_to_disk( let checkpoint_path_string = checkpoint_path.to_string_lossy().into_owned(); let arrow_schema = plan.schema().inner().clone(); let df = ctx.execute_logical_plan(plan).await?; - if let Err(error) = df - .write_parquet(&checkpoint_path_string, DataFrameWriteOptions::new(), None) - .await - { - cleanup_checkpoint_path(&checkpoint_path).await; - return Err(error.into()); - } - // DataFusion's `write_parquet` may not create the checkpoint directory when - // the DataFrame is empty (zero rows). In that case we fall back to an - // in-memory table so that downstream scans see the correct (empty) schema - // instead of failing with a "path does not exist" error. - let path_exists = match tokio::fs::try_exists(&checkpoint_path).await { - Ok(exists) => exists, - Err(e) => { - warn!( - "failed to check checkpoint path {}: {e}; assuming it exists", - checkpoint_path.display() - ); - true - } - }; - if !path_exists { + let batches = df.collect().await?; + + // For empty DataFrames, skip writing to disk entirely. + // DataFusion's `write_parquet` creates an empty directory with no parquet files, + // and `read_parquet` then fails because the directory has no `.parquet` files. + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + if total_rows == 0 { let table = Arc::new(MemTable::try_new(arrow_schema, vec![vec![]])?); let scan = LogicalPlanBuilder::scan( datafusion::common::TableReference::bare(UNNAMED_TABLE), @@ -687,6 +672,17 @@ async fn materialize_logical_plan_to_disk( .build()?; return Ok((scan, None)); } + + // Re-create a DataFrame from the collected batches and write to parquet. + let table = Arc::new(MemTable::try_new(arrow_schema, vec![batches])?); + let df = ctx.read_table(table)?; + if let Err(error) = df + .write_parquet(&checkpoint_path_string, DataFrameWriteOptions::new(), None) + .await + { + cleanup_checkpoint_path(&checkpoint_path).await; + return Err(error.into()); + } let df = match ctx .read_parquet(&checkpoint_path_string, ParquetReadOptions::default()) .await diff --git a/scripts/spark-tests/plugins/spark.py b/scripts/spark-tests/plugins/spark.py index 70a7ed995f..9886715a64 100644 --- a/scripts/spark-tests/plugins/spark.py +++ b/scripts/spark-tests/plugins/spark.py @@ -233,6 +233,10 @@ def assertEqual(self, first, second, msg=None): # noqa: N802 class TestMarker: keywords: list[str] reason: str + pyspark_versions: list[str] | None = None + """Optional list of PySpark version prefixes (e.g. ``["3.5."]``) to restrict + when this marker applies. When ``None`` (default) the marker applies to all + versions.""" SKIPPED_SPARK_TESTS = [ @@ -319,6 +323,20 @@ class TestMarker: keywords=["test_parity_job_cancellation.py"], reason="Slow test not working yet", ), + # The PySpark 3.5.x Spark Connect client does not implement checkpoint + # or localCheckpoint on the DataFrame, so the doctests fail with + # PySparkNotImplementedError. PySpark 4.x supports these APIs natively. + # Sail's own checkpoint tests (in PySail) cover the functionality. + TestMarker( + keywords=["pyspark.sql.dataframe.DataFrame.checkpoint"], + reason="PySpark 3.5.x Spark Connect does not support checkpoint; ported to PySail test suite", + pyspark_versions=["3.4.", "3.5."], + ), + TestMarker( + keywords=["pyspark.sql.dataframe.DataFrame.localCheckpoint"], + reason="PySpark 3.5.x Spark Connect does not support localCheckpoint; ported to PySail test suite", + pyspark_versions=["3.4.", "3.5."], + ), # The following tests rely on direct JVM access (RDD API, Java gateway), # which is not supported by Sail. TestMarker( @@ -377,8 +395,13 @@ class TestMarker: def add_pyspark_test_markers(items: list[pytest.Item]): + import pyspark + + version = pyspark.__version__ for item in items: for test in SKIPPED_SPARK_TESTS: + if test.pyspark_versions is not None and not any(version.startswith(v) for v in test.pyspark_versions): + continue if all(k in item.keywords for k in test.keywords): item.add_marker(pytest.mark.skip(reason=test.reason)) From b3b58c087482b2d4995399c2df336209de3dbccc Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 12 May 2026 22:06:05 +0000 Subject: [PATCH 15/22] refactor: address checkpoint review feedback (runner, config, eager semantics) Agent-Logs-Url: https://github.com/lakehq/sail/sessions/47e7ff4d-c51a-4857-b2bb-975ce3858297 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- crates/sail-common/src/config/application.rs | 1 + .../sail-common/src/config/application.yaml | 14 +++ .../src/service/plan_executor.rs | 85 +++++++++++++++---- crates/sail-spark-connect/src/session.rs | 5 ++ .../sail-spark-connect/src/session_manager.rs | 1 + 5 files changed, 88 insertions(+), 18 deletions(-) diff --git a/crates/sail-common/src/config/application.rs b/crates/sail-common/src/config/application.rs index d9028778a1..0472a6bbfc 100644 --- a/crates/sail-common/src/config/application.rs +++ b/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)] diff --git a/crates/sail-common/src/config/application.yaml b/crates/sail-common/src/config/application.yaml index 4cd21832a7..0979964620 100644 --- a/crates/sail-common/src/config/application.yaml +++ b/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 path + that is shared and readable by all workers (e.g. a shared filesystem or + object-store URI), 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" diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 22d45035fc..d21377fd95 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -532,7 +532,14 @@ pub(crate) async fn handle_execute_checkpoint_command( let CheckpointCommand { relation, local, - eager, + // The `eager` flag is a Spark hint requesting that materialization be + // deferred until the cached relation is first used. Sail does not + // implement deferred materialization (which would require intercepting + // `CachedRemoteRelation` reads and updating the cached entry on first + // use), so we always materialize eagerly. This is a correctness- + // preserving choice: lineage is truncated and the checkpoint exists + // immediately, matching Spark's observable semantics for `eager=true`. + eager: _, storage_level, } = checkpoint; let relation = relation.required("checkpoint relation")?; @@ -547,15 +554,13 @@ pub(crate) async fn handle_execute_checkpoint_command( let fields = named_plan_fields(&logical_plan, fields); let use_disk = checkpoint_uses_disk(local, storage_level.as_ref()); - let (cached_plan, storage_path) = if eager && use_disk { + let (cached_plan, storage_path) = if use_disk { materialize_logical_plan_to_disk(ctx, logical_plan).await? - } else if eager { + } else { ( materialize_logical_plan_to_memory(ctx, logical_plan).await?, None, ) - } else { - (logical_plan, None) }; let relation_id = uuid::Uuid::new_v4().to_string(); @@ -620,15 +625,38 @@ fn checkpoint_uses_disk( storage_level.map(|level| level.use_disk).unwrap_or(true) } -/// 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_to_memory( +/// Eagerly executes a logical plan via the [`JobService`] runner used by +/// regular query execution. This ensures checkpoint materialization runs on +/// the same execution path as user queries (including in cluster mode) and +/// participates in job tracking and cancellation. +async fn execute_plan_via_runner( ctx: &SessionContext, plan: LogicalPlan, -) -> SparkResult { +) -> SparkResult<( + datafusion::arrow::datatypes::SchemaRef, + Vec, +)> { let arrow_schema = plan.schema().inner().clone(); let df = ctx.execute_logical_plan(plan).await?; - let batches = df.collect().await?; + let (session_state, plan) = df.into_parts(); + let plan = session_state.optimize(&plan)?; + let physical_plan = session_state + .query_planner() + .create_physical_plan(&plan, &session_state) + .await?; + let service = ctx.extension::()?; + let stream = service.runner().execute(ctx, physical_plan).await?; + let batches = read_stream(stream).await?; + Ok((arrow_schema, batches)) +} + +/// Eagerly executes the resolved logical plan and returns a new logical plan +/// that scans the materialized data from an in-memory table. +async fn materialize_logical_plan_to_memory( + ctx: &SessionContext, + plan: LogicalPlan, +) -> SparkResult { + let (arrow_schema, batches) = execute_plan_via_runner(ctx, plan).await?; let table = Arc::new(MemTable::try_new(arrow_schema, vec![batches])?); let scan = LogicalPlanBuilder::scan( datafusion::common::TableReference::bare(UNNAMED_TABLE), @@ -639,24 +667,45 @@ async fn materialize_logical_plan_to_memory( Ok(scan) } -/// Eagerly executes the resolved logical plan, writes the result to local -/// checkpoint files, and returns a new logical plan that scans those files. +/// Returns the configured checkpoint root directory. When the configuration +/// option is empty, falls back to a `sail-checkpoints/` subdirectory of the +/// system temporary directory. +/// +/// In cluster deployments the operator should configure `spark.checkpoint_dir` +/// to a path that is shared and readable by all workers (for example, a +/// distributed filesystem mount or an object-store URI). The default +/// node-local path is suitable for single-node deployments and tests but +/// will not work correctly for distributed execution since remote workers +/// cannot read from the driver's local temporary directory. +fn resolve_checkpoint_root(spark: &SparkSession) -> PathBuf { + let configured = spark.options().checkpoint_dir.trim(); + if configured.is_empty() { + std::env::temp_dir().join("sail-checkpoints") + } else { + PathBuf::from(configured) + } +} + +/// Eagerly executes the resolved logical plan, writes the result to checkpoint +/// files, and returns a new logical plan that scans those files. async fn materialize_logical_plan_to_disk( ctx: &SessionContext, plan: LogicalPlan, ) -> SparkResult<(LogicalPlan, Option)> { // Spark Connect does not send a user checkpoint directory with this command, - // so Sail stores reliable checkpoint files in a per-server temporary area - // and deletes them when the client releases the cached remote relation. - let checkpoint_root = std::env::temp_dir().join("sail-checkpoints"); + // so Sail uses the directory configured by `spark.checkpoint_dir` (or a + // server-local temporary area if unset). Files are best-effort cleaned + // up when the client releases the cached remote relation or when the + // session ends. + let spark = ctx.extension::()?; + let checkpoint_root = resolve_checkpoint_root(&spark); tokio::fs::create_dir_all(&checkpoint_root) .await .map_err(|e| SparkError::internal(format!("failed to create checkpoint directory: {e}")))?; let checkpoint_path = checkpoint_root.join(uuid::Uuid::new_v4().to_string()); let checkpoint_path_string = checkpoint_path.to_string_lossy().into_owned(); - let arrow_schema = plan.schema().inner().clone(); - let df = ctx.execute_logical_plan(plan).await?; - let batches = df.collect().await?; + + let (arrow_schema, batches) = execute_plan_via_runner(ctx, plan).await?; // For empty DataFrames, skip writing to disk entirely. // DataFusion's `write_parquet` creates an empty directory with no parquet files, diff --git a/crates/sail-spark-connect/src/session.rs b/crates/sail-spark-connect/src/session.rs index ef70593f6f..bd0aff35c5 100644 --- a/crates/sail-spark-connect/src/session.rs +++ b/crates/sail-spark-connect/src/session.rs @@ -21,6 +21,11 @@ use crate::streaming::{ #[derive(Debug, Clone)] pub(crate) struct SparkSessionOptions { pub execution_heartbeat_interval: Duration, + /// Directory where reliable (disk) checkpoints are written. When empty, + /// a `sail-checkpoints/` subdirectory of the system temporary directory + /// is used. In cluster mode this should be set to a path that is shared + /// and readable by all workers. + pub checkpoint_dir: String, } /// A Spark session extension to the DataFusion [`SessionContext`]. diff --git a/crates/sail-spark-connect/src/session_manager.rs b/crates/sail-spark-connect/src/session_manager.rs index b94582dbb4..298fc29ef0 100644 --- a/crates/sail-spark-connect/src/session_manager.rs +++ b/crates/sail-spark-connect/src/session_manager.rs @@ -41,6 +41,7 @@ impl ServerSessionMutator for SparkSessionMutator { execution_heartbeat_interval: Duration::from_secs( self.config.spark.execution_heartbeat_interval_secs, ), + checkpoint_dir: self.config.spark.checkpoint_dir.clone(), }, ) .map_err(|e| internal_datafusion_err!("{e}"))?; From 241c8e33111e3376c67d3ebdc4fe87974dd41e20 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 12 May 2026 22:50:25 +0000 Subject: [PATCH 16/22] refactor: stream checkpoint output directly to parquet and clarify checkpoint_dir docs Agent-Logs-Url: https://github.com/lakehq/sail/sessions/2feca938-4dd6-4acf-add9-39cda1195a9f Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- .../sail-common/src/config/application.yaml | 12 +-- .../src/service/plan_executor.rs | 89 ++++++++++++------- scripts/spark-tests/plugins/spark.py | 6 +- 3 files changed, 66 insertions(+), 41 deletions(-) diff --git a/crates/sail-common/src/config/application.yaml b/crates/sail-common/src/config/application.yaml index 0979964620..f027bb4bfc 100644 --- a/crates/sail-common/src/config/application.yaml +++ b/crates/sail-common/src/config/application.yaml @@ -744,11 +744,13 @@ 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 path - that is shared and readable by all workers (e.g. a shared filesystem or - object-store URI), 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. + directory is used. In cluster deployments, this should be set to a + filesystem path that is shared and readable by all workers (for example, + a network mount such as NFS), since the default location is node-local. + Only local or shared filesystem paths are currently supported; remote + object-store URIs (such as `s3://`) are not. 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 diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index d21377fd95..9d14f04b40 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -5,14 +5,14 @@ use std::task::{Context, Poll}; use datafusion::arrow::compute::concat_batches; use datafusion::catalog::MemTable; -use datafusion::dataframe::DataFrameWriteOptions; use datafusion::datasource::provider_as_source; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, UNNAMED_TABLE}; +use datafusion::parquet::arrow::AsyncArrowWriter; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use fastrace::collector::SpanContext; use fastrace::future::FutureExt; use fastrace::Span; -use futures::stream; +use futures::{stream, TryStreamExt}; use log::{debug, warn}; use sail_catalog::manager::tracker::CatalogCachedRelation; use sail_catalog::manager::CatalogManager; @@ -672,11 +672,14 @@ async fn materialize_logical_plan_to_memory( /// system temporary directory. /// /// In cluster deployments the operator should configure `spark.checkpoint_dir` -/// to a path that is shared and readable by all workers (for example, a -/// distributed filesystem mount or an object-store URI). The default -/// node-local path is suitable for single-node deployments and tests but -/// will not work correctly for distributed execution since remote workers -/// cannot read from the driver's local temporary directory. +/// to a filesystem path that is shared and readable by all workers (for +/// example, a network mount such as NFS). The default node-local path is +/// suitable for single-node deployments and tests but will not work +/// correctly for distributed execution since remote workers cannot read +/// from the driver's local temporary directory. Only local or shared +/// filesystem paths are currently supported; remote object-store URIs +/// (such as `s3://`) are not and will surface as filesystem I/O errors +/// when the checkpoint directory is created. fn resolve_checkpoint_root(spark: &SparkSession) -> PathBuf { let configured = spark.options().checkpoint_dir.trim(); if configured.is_empty() { @@ -686,8 +689,15 @@ fn resolve_checkpoint_root(spark: &SparkSession) -> PathBuf { } } -/// Eagerly executes the resolved logical plan, writes the result to checkpoint -/// files, and returns a new logical plan that scans those files. +/// Eagerly executes the resolved logical plan, streaming the result directly +/// to a checkpoint Parquet file on disk, and returns a new logical plan that +/// scans those files. +/// +/// Streaming the physical plan output through an `AsyncArrowWriter` avoids +/// buffering the entire materialized DataFrame in driver memory, which is +/// the whole point of using the on-disk checkpoint path. Execution still +/// goes through the [`JobService`] runner so that the work participates in +/// distributed execution, job tracking, and cancellation. async fn materialize_logical_plan_to_disk( ctx: &SessionContext, plan: LogicalPlan, @@ -703,35 +713,48 @@ async fn materialize_logical_plan_to_disk( .await .map_err(|e| SparkError::internal(format!("failed to create checkpoint directory: {e}")))?; let checkpoint_path = checkpoint_root.join(uuid::Uuid::new_v4().to_string()); - let checkpoint_path_string = checkpoint_path.to_string_lossy().into_owned(); - let (arrow_schema, batches) = execute_plan_via_runner(ctx, plan).await?; - - // For empty DataFrames, skip writing to disk entirely. - // DataFusion's `write_parquet` creates an empty directory with no parquet files, - // and `read_parquet` then fails because the directory has no `.parquet` files. - let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); - if total_rows == 0 { - let table = Arc::new(MemTable::try_new(arrow_schema, vec![vec![]])?); - let scan = LogicalPlanBuilder::scan( - datafusion::common::TableReference::bare(UNNAMED_TABLE), - provider_as_source(table), - None, - )? - .build()?; - return Ok((scan, None)); + // Plan and execute via the `JobService` runner so checkpoint + // materialization goes through the same execution path as user queries. + let arrow_schema = plan.schema().inner().clone(); + let df = ctx.execute_logical_plan(plan).await?; + let (session_state, plan) = df.into_parts(); + let plan = session_state.optimize(&plan)?; + let physical_plan = session_state + .query_planner() + .create_physical_plan(&plan, &session_state) + .await?; + let service = ctx.extension::()?; + let mut stream = service.runner().execute(ctx, physical_plan).await?; + + // Stream batches directly to a single Parquet file inside the checkpoint + // directory without collecting them all into memory first. + let write_result: SparkResult<()> = async { + tokio::fs::create_dir_all(&checkpoint_path).await?; + let parquet_path = checkpoint_path.join("part-0.parquet"); + let file = tokio::fs::File::create(&parquet_path).await?; + let mut writer = AsyncArrowWriter::try_new(file, arrow_schema, None) + .map_err(|e| SparkError::internal(format!("failed to create parquet writer: {e}")))?; + while let Some(batch) = stream.try_next().await? { + writer + .write(&batch) + .await + .map_err(|e| SparkError::internal(format!("failed to write parquet batch: {e}")))?; + } + writer + .close() + .await + .map_err(|e| SparkError::internal(format!("failed to close parquet writer: {e}")))?; + Ok(()) } + .await; - // Re-create a DataFrame from the collected batches and write to parquet. - let table = Arc::new(MemTable::try_new(arrow_schema, vec![batches])?); - let df = ctx.read_table(table)?; - if let Err(error) = df - .write_parquet(&checkpoint_path_string, DataFrameWriteOptions::new(), None) - .await - { + if let Err(error) = write_result { cleanup_checkpoint_path(&checkpoint_path).await; - return Err(error.into()); + return Err(error); } + + let checkpoint_path_string = checkpoint_path.to_string_lossy().into_owned(); let df = match ctx .read_parquet(&checkpoint_path_string, ParquetReadOptions::default()) .await diff --git a/scripts/spark-tests/plugins/spark.py b/scripts/spark-tests/plugins/spark.py index 9886715a64..735202a7d8 100644 --- a/scripts/spark-tests/plugins/spark.py +++ b/scripts/spark-tests/plugins/spark.py @@ -233,10 +233,10 @@ def assertEqual(self, first, second, msg=None): # noqa: N802 class TestMarker: keywords: list[str] reason: str + # Optional list of PySpark version prefixes (e.g. ``["3.5."]``) to restrict + # when this marker applies. When ``None`` (default) the marker applies to + # all versions. pyspark_versions: list[str] | None = None - """Optional list of PySpark version prefixes (e.g. ``["3.5."]``) to restrict - when this marker applies. When ``None`` (default) the marker applies to all - versions.""" SKIPPED_SPARK_TESTS = [ From ac15f6bd430d7596a57b55bb6eaf9f1b4ad553cb Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 12 May 2026 23:41:50 +0000 Subject: [PATCH 17/22] feat: support object-store checkpoint locations Agent-Logs-Url: https://github.com/lakehq/sail/sessions/31be1e08-c59a-4597-86fd-f617c6186791 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- Cargo.lock | 3 + crates/sail-catalog/Cargo.toml | 1 + crates/sail-catalog/src/manager/tracker.rs | 18 ++- .../sail-common/src/config/application.yaml | 10 +- crates/sail-spark-connect/Cargo.toml | 2 + .../src/service/plan_executor.rs | 136 ++++++++++++------ 6 files changed, 113 insertions(+), 57 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 66ce79d452..4b5f0ff659 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6975,6 +6975,7 @@ dependencies = [ "serde_arrow", "thiserror 2.0.18", "tokio", + "url", ] [[package]] @@ -7654,6 +7655,7 @@ dependencies = [ "lazy_static", "log", "monostate", + "object_store", "pbjson", "pbjson-build", "pbjson-types", @@ -7685,6 +7687,7 @@ dependencies = [ "tonic-prost", "tonic-prost-build", "tonic-types", + "url", "uuid", ] diff --git a/crates/sail-catalog/Cargo.toml b/crates/sail-catalog/Cargo.toml index 06c6a8cf6e..2c2273aa5f 100644 --- a/crates/sail-catalog/Cargo.toml +++ b/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 } diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index 00835cbdd4..bd89562e7a 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -1,5 +1,4 @@ use std::collections::HashMap; -use std::path::PathBuf; use std::sync::{Arc, Mutex, MutexGuard}; use datafusion_expr::{LogicalPlan, ScalarUDF}; @@ -17,7 +16,7 @@ pub struct CatalogLogicalPlanId(u64); pub struct CatalogCachedRelation { pub plan: Arc, pub fields: Vec, - pub storage_path: Option, + pub storage_uri: Option, } #[derive(Default)] @@ -133,12 +132,19 @@ impl Drop for CatalogObjectTracker { return; }; for (_, relation) in state.cached_relations.drain() { - let Some(path) = relation.storage_path else { + let Some(uri) = relation.storage_uri else { continue; }; - // `Drop` runs in a sync context, so we use blocking filesystem calls. - // Errors are ignored because the caller cannot react to them. - let _ = std::fs::remove_dir_all(&path); + 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); + } + } + } } } } diff --git a/crates/sail-common/src/config/application.yaml b/crates/sail-common/src/config/application.yaml index f027bb4bfc..bccfd58a29 100644 --- a/crates/sail-common/src/config/application.yaml +++ b/crates/sail-common/src/config/application.yaml @@ -745,12 +745,10 @@ `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 - filesystem path that is shared and readable by all workers (for example, - a network mount such as NFS), since the default location is node-local. - Only local or shared filesystem paths are currently supported; remote - object-store URIs (such as `s3://`) are not. Checkpoint files are - best-effort cleaned up when the client releases the cached relation or - when the session ends. + 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 diff --git a/crates/sail-spark-connect/Cargo.toml b/crates/sail-spark-connect/Cargo.toml index 5730cff7bc..60dd09c998 100644 --- a/crates/sail-spark-connect/Cargo.toml +++ b/crates/sail-spark-connect/Cargo.toml @@ -41,6 +41,8 @@ log = { workspace = true } futures = { workspace = true } fastrace = { workspace = true } pyo3 = { workspace = true } +object_store = { workspace = true } +url = { workspace = true } [build-dependencies] prost-build = { workspace = true } diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 9d14f04b40..d443927412 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -1,4 +1,4 @@ -use std::path::{Path, PathBuf}; +use std::path::PathBuf; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -7,13 +7,16 @@ 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::parquet::arrow::async_writer::ParquetObjectWriter; use datafusion::parquet::arrow::AsyncArrowWriter; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use fastrace::collector::SpanContext; use fastrace::future::FutureExt; use fastrace::Span; -use futures::{stream, TryStreamExt}; +use futures::{stream, StreamExt, TryStreamExt}; use log::{debug, warn}; +use object_store::path::Path as ObjectStorePath; +use object_store::ObjectStoreScheme; use sail_catalog::manager::tracker::CatalogCachedRelation; use sail_catalog::manager::CatalogManager; use sail_common::spec; @@ -554,7 +557,7 @@ pub(crate) async fn handle_execute_checkpoint_command( let fields = named_plan_fields(&logical_plan, fields); let use_disk = checkpoint_uses_disk(local, storage_level.as_ref()); - let (cached_plan, storage_path) = if use_disk { + let (cached_plan, storage_uri) = if use_disk { materialize_logical_plan_to_disk(ctx, logical_plan).await? } else { ( @@ -570,11 +573,11 @@ pub(crate) async fn handle_execute_checkpoint_command( CatalogCachedRelation { plan: Arc::new(cached_plan), fields, - storage_path: storage_path.clone(), + storage_uri: storage_uri.clone(), }, ) { - if let Some(path) = storage_path { - cleanup_checkpoint_path(&path).await; + if let Some(uri) = storage_uri { + cleanup_checkpoint_path(ctx, &uri).await; } return Err(SparkError::internal(format!( "failed to track cached relation: {error}" @@ -667,28 +670,62 @@ async fn materialize_logical_plan_to_memory( Ok(scan) } -/// Returns the configured checkpoint root directory. When the configuration +fn directory_file_url(path: PathBuf) -> SparkResult { + let path = if path.is_absolute() { + path + } else { + std::env::current_dir()?.join(path) + }; + url::Url::from_directory_path(&path).map_err(|_| { + SparkError::invalid(format!( + "checkpoint directory is not a valid file path: {}", + path.display() + )) + }) +} + +fn ensure_directory_url(mut url: url::Url) -> SparkResult { + if !url.path().ends_with('/') { + let path = format!("{}/", url.path()); + url.set_path(&path); + } + Ok(url) +} + +/// Returns the configured checkpoint root location. When the configuration /// option is empty, falls back to a `sail-checkpoints/` subdirectory of the /// system temporary directory. /// /// In cluster deployments the operator should configure `spark.checkpoint_dir` -/// to a filesystem path that is shared and readable by all workers (for -/// example, a network mount such as NFS). The default node-local path is -/// suitable for single-node deployments and tests but will not work -/// correctly for distributed execution since remote workers cannot read -/// from the driver's local temporary directory. Only local or shared -/// filesystem paths are currently supported; remote object-store URIs -/// (such as `s3://`) are not and will surface as filesystem I/O errors -/// when the checkpoint directory is created. -fn resolve_checkpoint_root(spark: &SparkSession) -> PathBuf { +/// to a location that is shared and readable by all workers, such as an +/// object-store URI or a shared filesystem path. The default node-local path +/// is suitable for single-node deployments and tests but will not work +/// correctly for distributed execution since remote workers cannot read from +/// the driver's local temporary directory. +fn resolve_checkpoint_root(spark: &SparkSession) -> SparkResult { let configured = spark.options().checkpoint_dir.trim(); if configured.is_empty() { - std::env::temp_dir().join("sail-checkpoints") + directory_file_url(std::env::temp_dir().join("sail-checkpoints")) + } else if let Ok(url) = url::Url::parse(configured) { + ensure_directory_url(url) } else { - PathBuf::from(configured) + directory_file_url(PathBuf::from(configured)) } } +fn checkpoint_store_and_path( + ctx: &SessionContext, + checkpoint_url: &url::Url, +) -> SparkResult<(Arc, ObjectStorePath)> { + let (_, checkpoint_path) = ObjectStoreScheme::parse(checkpoint_url) + .map_err(|e| SparkError::invalid(format!("invalid checkpoint location: {e}")))?; + let object_store = ctx + .runtime_env() + .object_store_registry + .get_store(checkpoint_url)?; + Ok((object_store, checkpoint_path)) +} + /// Eagerly executes the resolved logical plan, streaming the result directly /// to a checkpoint Parquet file on disk, and returns a new logical plan that /// scans those files. @@ -701,18 +738,19 @@ fn resolve_checkpoint_root(spark: &SparkSession) -> PathBuf { async fn materialize_logical_plan_to_disk( ctx: &SessionContext, plan: LogicalPlan, -) -> SparkResult<(LogicalPlan, Option)> { +) -> SparkResult<(LogicalPlan, Option)> { // Spark Connect does not send a user checkpoint directory with this command, // so Sail uses the directory configured by `spark.checkpoint_dir` (or a // server-local temporary area if unset). Files are best-effort cleaned // up when the client releases the cached remote relation or when the // session ends. let spark = ctx.extension::()?; - let checkpoint_root = resolve_checkpoint_root(&spark); - tokio::fs::create_dir_all(&checkpoint_root) - .await - .map_err(|e| SparkError::internal(format!("failed to create checkpoint directory: {e}")))?; - let checkpoint_path = checkpoint_root.join(uuid::Uuid::new_v4().to_string()); + let checkpoint_root = resolve_checkpoint_root(&spark)?; + let checkpoint_url = checkpoint_root + .join(&format!("{}/", uuid::Uuid::new_v4())) + .map_err(|e| SparkError::invalid(format!("invalid checkpoint location: {e}")))?; + let checkpoint_uri = checkpoint_url.to_string(); + let (object_store, checkpoint_path) = checkpoint_store_and_path(ctx, &checkpoint_url)?; // Plan and execute via the `JobService` runner so checkpoint // materialization goes through the same execution path as user queries. @@ -730,10 +768,9 @@ async fn materialize_logical_plan_to_disk( // Stream batches directly to a single Parquet file inside the checkpoint // directory without collecting them all into memory first. let write_result: SparkResult<()> = async { - tokio::fs::create_dir_all(&checkpoint_path).await?; - let parquet_path = checkpoint_path.join("part-0.parquet"); - let file = tokio::fs::File::create(&parquet_path).await?; - let mut writer = AsyncArrowWriter::try_new(file, arrow_schema, None) + let parquet_path = ObjectStorePath::from(format!("{checkpoint_path}/part-0.parquet")); + let object_writer = ParquetObjectWriter::new(object_store.clone(), parquet_path); + let mut writer = AsyncArrowWriter::try_new(object_writer, arrow_schema, None) .map_err(|e| SparkError::internal(format!("failed to create parquet writer: {e}")))?; while let Some(batch) = stream.try_next().await? { writer @@ -750,34 +787,43 @@ async fn materialize_logical_plan_to_disk( .await; if let Err(error) = write_result { - cleanup_checkpoint_path(&checkpoint_path).await; + cleanup_checkpoint_path(ctx, &checkpoint_uri).await; return Err(error); } - let checkpoint_path_string = checkpoint_path.to_string_lossy().into_owned(); let df = match ctx - .read_parquet(&checkpoint_path_string, ParquetReadOptions::default()) + .read_parquet(&checkpoint_uri, ParquetReadOptions::default()) .await { Ok(df) => df, Err(error) => { - cleanup_checkpoint_path(&checkpoint_path).await; + cleanup_checkpoint_path(ctx, &checkpoint_uri).await; return Err(error.into()); } }; - Ok((df.into_unoptimized_plan(), Some(checkpoint_path))) + Ok((df.into_unoptimized_plan(), Some(checkpoint_uri))) } -async fn cleanup_checkpoint_path(path: &Path) { - match tokio::fs::remove_dir_all(path).await { - Ok(()) => {} - Err(e) if e.kind() == std::io::ErrorKind::NotFound => {} - Err(e) => { - warn!( - "failed to remove checkpoint directory {}: {e}", - path.display() - ); - } +async fn cleanup_checkpoint_path(ctx: &SessionContext, checkpoint_uri: &str) { + let Ok(checkpoint_url) = url::Url::parse(checkpoint_uri) else { + warn!("failed to parse checkpoint location for cleanup: {checkpoint_uri}"); + return; + }; + let Ok((object_store, checkpoint_path)) = checkpoint_store_and_path(ctx, &checkpoint_url) + else { + warn!("failed to resolve checkpoint object store for cleanup: {checkpoint_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::>() + .await + { + warn!("failed to remove checkpoint location {checkpoint_uri}: {e}"); } } @@ -794,8 +840,8 @@ pub(crate) async fn handle_execute_remove_cached_remote_relation_command( let removed = manager .remove_cached_relation(&relation.relation_id) .map_err(|e| SparkError::internal(format!("failed to remove cached relation: {e}")))?; - if let Some(path) = removed.and_then(|relation| relation.storage_path) { - cleanup_checkpoint_path(&path).await; + if let Some(uri) = removed.and_then(|relation| relation.storage_uri) { + cleanup_checkpoint_path(ctx, &uri).await; } let spark = ctx.extension::()?; From b5831b7342fc042f5c3238e72705ad9407251026 Mon Sep 17 00:00:00 2001 From: Shehab Amin <11789402+shehabgamin@users.noreply.github.com> Date: Tue, 12 May 2026 23:53:11 +0000 Subject: [PATCH 18/22] Apply suggestions from code review Co-authored-by: Shehab Amin <11789402+shehabgamin@users.noreply.github.com> --- scripts/spark-tests/plugins/spark.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/spark-tests/plugins/spark.py b/scripts/spark-tests/plugins/spark.py index 735202a7d8..f6c3296dc1 100644 --- a/scripts/spark-tests/plugins/spark.py +++ b/scripts/spark-tests/plugins/spark.py @@ -330,12 +330,12 @@ class TestMarker: TestMarker( keywords=["pyspark.sql.dataframe.DataFrame.checkpoint"], reason="PySpark 3.5.x Spark Connect does not support checkpoint; ported to PySail test suite", - pyspark_versions=["3.4.", "3.5."], + pyspark_versions=["3.5."], ), TestMarker( keywords=["pyspark.sql.dataframe.DataFrame.localCheckpoint"], reason="PySpark 3.5.x Spark Connect does not support localCheckpoint; ported to PySail test suite", - pyspark_versions=["3.4.", "3.5."], + pyspark_versions=["3.5."], ), # The following tests rely on direct JVM access (RDD API, Java gateway), # which is not supported by Sail. From d1e4dc77c7889379361d9ccaf70813cff0121f95 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 13 May 2026 00:07:16 +0000 Subject: [PATCH 19/22] fix: address checkpoint review follow-ups Agent-Logs-Url: https://github.com/lakehq/sail/sessions/52f051c1-f32c-4857-a6d0-3bcf42c09452 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- Cargo.lock | 2 + crates/sail-catalog/src/manager/mod.rs | 4 ++ crates/sail-catalog/src/manager/tracker.rs | 9 ++++ crates/sail-session/Cargo.toml | 2 + .../src/session_manager/actor/handler.rs | 54 ++++++++++++++++++- .../src/service/plan_executor.rs | 52 +++++++++++------- 6 files changed, 102 insertions(+), 21 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 4b5f0ff659..a7b40026a9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7612,6 +7612,7 @@ dependencies = [ "futures", "indexmap 2.14.0", "log", + "object_store", "readonly", "sail-cache", "sail-catalog", @@ -7642,6 +7643,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tonic", + "url", ] [[package]] diff --git a/crates/sail-catalog/src/manager/mod.rs b/crates/sail-catalog/src/manager/mod.rs index 9f1405b2df..86f2f0c87f 100644 --- a/crates/sail-catalog/src/manager/mod.rs +++ b/crates/sail-catalog/src/manager/mod.rs @@ -168,6 +168,10 @@ impl CatalogManager { ) -> CatalogResult> { self.tracker.remove_cached_relation(relation_id) } + + pub fn drain_cached_relation_storage_uris(&self) -> CatalogResult> { + self.tracker.drain_cached_relation_storage_uris() + } } impl CatalogManagerState { diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index bd89562e7a..1f107e732b 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -120,6 +120,15 @@ impl CatalogObjectTracker { let mut state = self.state()?; Ok(state.cached_relations.remove(relation_id)) } + + pub fn drain_cached_relation_storage_uris(&self) -> CatalogResult> { + let mut state = self.state()?; + Ok(state + .cached_relations + .drain() + .filter_map(|(_, relation)| relation.storage_uri) + .collect()) + } } impl Drop for CatalogObjectTracker { diff --git a/crates/sail-session/Cargo.toml b/crates/sail-session/Cargo.toml index 9d5bb34c7a..b5f3b1d65d 100644 --- a/crates/sail-session/Cargo.toml +++ b/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 } diff --git a/crates/sail-session/src/session_manager/actor/handler.rs b/crates/sail-session/src/session_manager/actor/handler.rs index 900bff95d1..c4122b7015 100644 --- a/crates/sail-session/src/session_manager/actor/handler.rs +++ b/crates/sail-session/src/session_manager/actor/handler.rs @@ -1,10 +1,14 @@ use std::sync::Arc; use chrono::Utc; +use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::prelude::SessionContext; use fastrace::collector::SpanContext; use fastrace::Span; +use futures::{StreamExt, TryStreamExt}; use log::{info, warn}; +use object_store::ObjectStoreScheme; +use sail_catalog::manager::CatalogManager; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::activity::ActivityTracker; use sail_common_datafusion::session::job::JobService; @@ -326,11 +330,27 @@ impl SessionManagerActor { warn!("job service not found for session {session_id}"); return; }; + let checkpoint_storage_uris = match context.extension::() { + Ok(manager) => match manager.drain_cached_relation_storage_uris() { + Ok(uris) => uris, + Err(e) => { + warn!("failed to drain cached relation checkpoint locations for session {session_id}: {e}"); + vec![] + } + }, + Err(e) => { + warn!("catalog manager not found for session {session_id}: {e}"); + vec![] + } + }; + let runtime_env = context.runtime_env(); let handle = ctx.handle().clone(); let (tx, rx) = oneshot::channel(); ctx.spawn(async move { service.runner().stop(tx).await; - let message = match rx.await { + let history = rx.await; + cleanup_cached_relation_storage(runtime_env, checkpoint_storage_uris).await; + let message = match history { Ok(x) => SessionManagerEvent::SetSessionHistory { session_id, history: SessionHistory { job_runner: x }, @@ -341,3 +361,35 @@ impl SessionManagerActor { }); } } + +async fn cleanup_cached_relation_storage(runtime_env: Arc, storage_uris: Vec) { + for uri in storage_uris { + cleanup_checkpoint_storage(&runtime_env, &uri).await; + } +} + +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::>() + .await + { + warn!("failed to remove checkpoint location {storage_uri}: {e}"); + } +} diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index d443927412..3c10076e7f 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -1,11 +1,13 @@ -use std::path::PathBuf; +use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; use datafusion::arrow::compute::concat_batches; +use datafusion::arrow::datatypes::SchemaRef; use datafusion::catalog::MemTable; use datafusion::datasource::provider_as_source; +use datafusion::execution::SendableRecordBatchStream; use datafusion::logical_expr::{LogicalPlan, LogicalPlanBuilder, UNNAMED_TABLE}; use datafusion::parquet::arrow::async_writer::ParquetObjectWriter; use datafusion::parquet::arrow::AsyncArrowWriter; @@ -632,13 +634,10 @@ fn checkpoint_uses_disk( /// regular query execution. This ensures checkpoint materialization runs on /// the same execution path as user queries (including in cluster mode) and /// participates in job tracking and cancellation. -async fn execute_plan_via_runner( +async fn execute_plan_via_runner_stream( ctx: &SessionContext, plan: LogicalPlan, -) -> SparkResult<( - datafusion::arrow::datatypes::SchemaRef, - Vec, -)> { +) -> SparkResult<(SchemaRef, SendableRecordBatchStream)> { let arrow_schema = plan.schema().inner().clone(); let df = ctx.execute_logical_plan(plan).await?; let (session_state, plan) = df.into_parts(); @@ -649,6 +648,14 @@ async fn execute_plan_via_runner( .await?; let service = ctx.extension::()?; let stream = service.runner().execute(ctx, physical_plan).await?; + Ok((arrow_schema, stream)) +} + +async fn collect_plan_via_runner( + ctx: &SessionContext, + plan: LogicalPlan, +) -> SparkResult<(SchemaRef, Vec)> { + let (arrow_schema, stream) = execute_plan_via_runner_stream(ctx, plan).await?; let batches = read_stream(stream).await?; Ok((arrow_schema, batches)) } @@ -659,7 +666,7 @@ async fn materialize_logical_plan_to_memory( ctx: &SessionContext, plan: LogicalPlan, ) -> SparkResult { - let (arrow_schema, batches) = execute_plan_via_runner(ctx, plan).await?; + let (arrow_schema, batches) = collect_plan_via_runner(ctx, plan).await?; let table = Arc::new(MemTable::try_new(arrow_schema, vec![batches])?); let scan = LogicalPlanBuilder::scan( datafusion::common::TableReference::bare(UNNAMED_TABLE), @@ -692,6 +699,18 @@ fn ensure_directory_url(mut url: url::Url) -> SparkResult { Ok(url) } +fn is_windows_drive_path(path: &str) -> bool { + let bytes = path.as_bytes(); + bytes.len() >= 3 + && bytes[0].is_ascii_alphabetic() + && bytes[1] == b':' + && matches!(bytes[2], b'/' | b'\\') +} + +fn is_uri_like(path: &str) -> bool { + path.contains("://") +} + /// Returns the configured checkpoint root location. When the configuration /// option is empty, falls back to a `sail-checkpoints/` subdirectory of the /// system temporary directory. @@ -706,7 +725,11 @@ fn resolve_checkpoint_root(spark: &SparkSession) -> SparkResult { let configured = spark.options().checkpoint_dir.trim(); if configured.is_empty() { directory_file_url(std::env::temp_dir().join("sail-checkpoints")) - } else if let Ok(url) = url::Url::parse(configured) { + } else if Path::new(configured).is_absolute() || is_windows_drive_path(configured) { + directory_file_url(PathBuf::from(configured)) + } else if is_uri_like(configured) { + let url = url::Url::parse(configured) + .map_err(|e| SparkError::invalid(format!("invalid checkpoint URI: {e}")))?; ensure_directory_url(url) } else { directory_file_url(PathBuf::from(configured)) @@ -752,18 +775,7 @@ async fn materialize_logical_plan_to_disk( let checkpoint_uri = checkpoint_url.to_string(); let (object_store, checkpoint_path) = checkpoint_store_and_path(ctx, &checkpoint_url)?; - // Plan and execute via the `JobService` runner so checkpoint - // materialization goes through the same execution path as user queries. - let arrow_schema = plan.schema().inner().clone(); - let df = ctx.execute_logical_plan(plan).await?; - let (session_state, plan) = df.into_parts(); - let plan = session_state.optimize(&plan)?; - let physical_plan = session_state - .query_planner() - .create_physical_plan(&plan, &session_state) - .await?; - let service = ctx.extension::()?; - let mut stream = service.runner().execute(ctx, physical_plan).await?; + let (arrow_schema, mut stream) = execute_plan_via_runner_stream(ctx, plan).await?; // Stream batches directly to a single Parquet file inside the checkpoint // directory without collecting them all into memory first. From aa44e2b8e6bc43f418795108b987a9f3a216b5a6 Mon Sep 17 00:00:00 2001 From: Shehab Amin <11789402+shehabgamin@users.noreply.github.com> Date: Wed, 13 May 2026 01:29:46 +0000 Subject: [PATCH 20/22] Potential fix for pull request finding Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- python/pysail/tests/spark/test_checkpoint.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/pysail/tests/spark/test_checkpoint.py b/python/pysail/tests/spark/test_checkpoint.py index ecab5a9256..2fc0ca29c1 100644 --- a/python/pysail/tests/spark/test_checkpoint.py +++ b/python/pysail/tests/spark/test_checkpoint.py @@ -74,10 +74,11 @@ def test_dataframe_checkpoint_with_duplicate_columns(spark, method, eager): df = spark.createDataFrame([(1, 2), (3, 4)], ["left", "right"]).selectExpr("left as id", "right as id") checkpointed = getattr(df, method)(eager) pdf = checkpointed.toPandas() + rows = sorted(pdf.values.tolist()) assert checkpointed.columns == ["id", "id"] assert pdf.columns.tolist() == ["id", "id"] - assert pdf.values.tolist() == [[1, 2], [3, 4]] + assert rows == [[1, 2], [3, 4]] def test_dataframe_checkpoint_after_transformation(spark): From f771809dab09ac6f1a8e74149beadb697cceb7f1 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 13 May 2026 01:35:57 +0000 Subject: [PATCH 21/22] fix: remove local checkpoint directories Agent-Logs-Url: https://github.com/lakehq/sail/sessions/d72a8b0a-1114-4626-abaa-4566793a6975 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- .../src/session_manager/actor/handler.rs | 16 ++++++++++++++++ .../src/service/plan_executor.rs | 16 ++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/crates/sail-session/src/session_manager/actor/handler.rs b/crates/sail-session/src/session_manager/actor/handler.rs index c4122b7015..675f437607 100644 --- a/crates/sail-session/src/session_manager/actor/handler.rs +++ b/crates/sail-session/src/session_manager/actor/handler.rs @@ -392,4 +392,20 @@ async fn cleanup_checkpoint_storage(runtime_env: &RuntimeEnv, storage_uri: &str) { 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}"); + } + } } diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 3c10076e7f..36185c2af6 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -837,6 +837,22 @@ async fn cleanup_checkpoint_path(ctx: &SessionContext, checkpoint_uri: &str) { { warn!("failed to remove checkpoint location {checkpoint_uri}: {e}"); } + cleanup_file_checkpoint_directory(&checkpoint_url, checkpoint_uri).await; +} + +async fn cleanup_file_checkpoint_directory(checkpoint_url: &url::Url, checkpoint_uri: &str) { + if checkpoint_url.scheme() != "file" { + return; + } + let Ok(path) = checkpoint_url.to_file_path() else { + warn!("failed to convert checkpoint file URL to path for cleanup: {checkpoint_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 {checkpoint_uri}: {e}"); + } + } } pub(crate) async fn handle_execute_remove_cached_remote_relation_command( From 70332d821b65cc2d2b05662487e9a97c019a0626 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Wed, 13 May 2026 02:06:12 +0000 Subject: [PATCH 22/22] refactor checkpoint cleanup follow-ups Agent-Logs-Url: https://github.com/lakehq/sail/sessions/ed3c4316-4757-49ea-bbf1-6831281f90b9 Co-authored-by: shehabgamin <11789402+shehabgamin@users.noreply.github.com> --- Cargo.lock | 3 ++ crates/sail-catalog/src/manager/tracker.rs | 27 ---------- crates/sail-common-datafusion/Cargo.toml | 3 ++ .../sail-common-datafusion/src/checkpoint.rs | 46 ++++++++++++++++ crates/sail-common-datafusion/src/lib.rs | 1 + .../src/session_manager/actor/handler.rs | 45 +--------------- .../src/service/plan_executor.rs | 54 +++---------------- 7 files changed, 60 insertions(+), 119 deletions(-) create mode 100644 crates/sail-common-datafusion/src/checkpoint.rs diff --git a/Cargo.lock b/Cargo.lock index a7b40026a9..ba1fcef33d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7152,6 +7152,8 @@ dependencies = [ "futures", "lazy_static", "lexical-core", + "log", + "object_store", "parquet-variant-compute", "parquet-variant-json", "pin-project-lite", @@ -7171,6 +7173,7 @@ dependencies = [ "tokio", "tonic", "tonic-prost-build", + "url", ] [[package]] diff --git a/crates/sail-catalog/src/manager/tracker.rs b/crates/sail-catalog/src/manager/tracker.rs index 1f107e732b..fead9d3452 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -130,30 +130,3 @@ impl CatalogObjectTracker { .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); - } - } - } - } - } -} diff --git a/crates/sail-common-datafusion/Cargo.toml b/crates/sail-common-datafusion/Cargo.toml index b6ccb1f2a7..991e301bd6 100644 --- a/crates/sail-common-datafusion/Cargo.toml +++ b/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] diff --git a/crates/sail-common-datafusion/src/checkpoint.rs b/crates/sail-common-datafusion/src/checkpoint.rs new file mode 100644 index 0000000000..550a0045b7 --- /dev/null +++ b/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::>() + .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}"); + } + } +} diff --git a/crates/sail-common-datafusion/src/lib.rs b/crates/sail-common-datafusion/src/lib.rs index 10e66400b7..00fb4898c7 100644 --- a/crates/sail-common-datafusion/src/lib.rs +++ b/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; diff --git a/crates/sail-session/src/session_manager/actor/handler.rs b/crates/sail-session/src/session_manager/actor/handler.rs index 675f437607..81427afbaa 100644 --- a/crates/sail-session/src/session_manager/actor/handler.rs +++ b/crates/sail-session/src/session_manager/actor/handler.rs @@ -5,10 +5,9 @@ use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::prelude::SessionContext; use fastrace::collector::SpanContext; use fastrace::Span; -use futures::{StreamExt, TryStreamExt}; use log::{info, warn}; -use object_store::ObjectStoreScheme; use sail_catalog::manager::CatalogManager; +use sail_common_datafusion::checkpoint::cleanup_checkpoint_storage; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::activity::ActivityTracker; use sail_common_datafusion::session::job::JobService; @@ -367,45 +366,3 @@ async fn cleanup_cached_relation_storage(runtime_env: Arc, storage_u cleanup_checkpoint_storage(&runtime_env, &uri).await; } } - -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::>() - .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}"); - } - } -} diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 36185c2af6..5cae327f45 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -15,13 +15,14 @@ use datafusion::prelude::{ParquetReadOptions, SessionContext}; use fastrace::collector::SpanContext; use fastrace::future::FutureExt; use fastrace::Span; -use futures::{stream, StreamExt, TryStreamExt}; +use futures::{stream, TryStreamExt}; use log::{debug, warn}; use object_store::path::Path as ObjectStorePath; use object_store::ObjectStoreScheme; use sail_catalog::manager::tracker::CatalogCachedRelation; use sail_catalog::manager::CatalogManager; use sail_common::spec; +use sail_common_datafusion::checkpoint::cleanup_checkpoint_storage; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; use sail_plan::resolve_and_execute_plan; @@ -556,7 +557,9 @@ pub(crate) async fn handle_execute_checkpoint_command( plan: logical_plan, fields, } = resolver.resolve_named_plan(plan).await?; - let fields = named_plan_fields(&logical_plan, fields); + let fields = fields.ok_or_else(|| { + SparkError::internal("resolved checkpoint query plan did not include field names") + })?; let use_disk = checkpoint_uses_disk(local, storage_level.as_ref()); let (cached_plan, storage_uri) = if use_disk { @@ -602,16 +605,6 @@ pub(crate) async fn handle_execute_checkpoint_command( )) } -fn named_plan_fields(plan: &LogicalPlan, fields: Option>) -> Vec { - fields.unwrap_or_else(|| { - plan.schema() - .fields() - .iter() - .map(|field| field.name().to_string()) - .collect() - }) -} - /// Determines whether a checkpoint should be materialized to disk. /// /// Non-local (reliable) checkpoints always use disk. For local checkpoints @@ -817,42 +810,7 @@ async fn materialize_logical_plan_to_disk( } async fn cleanup_checkpoint_path(ctx: &SessionContext, checkpoint_uri: &str) { - let Ok(checkpoint_url) = url::Url::parse(checkpoint_uri) else { - warn!("failed to parse checkpoint location for cleanup: {checkpoint_uri}"); - return; - }; - let Ok((object_store, checkpoint_path)) = checkpoint_store_and_path(ctx, &checkpoint_url) - else { - warn!("failed to resolve checkpoint object store for cleanup: {checkpoint_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::>() - .await - { - warn!("failed to remove checkpoint location {checkpoint_uri}: {e}"); - } - cleanup_file_checkpoint_directory(&checkpoint_url, checkpoint_uri).await; -} - -async fn cleanup_file_checkpoint_directory(checkpoint_url: &url::Url, checkpoint_uri: &str) { - if checkpoint_url.scheme() != "file" { - return; - } - let Ok(path) = checkpoint_url.to_file_path() else { - warn!("failed to convert checkpoint file URL to path for cleanup: {checkpoint_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 {checkpoint_uri}: {e}"); - } - } + cleanup_checkpoint_storage(ctx.runtime_env().as_ref(), checkpoint_uri).await; } pub(crate) async fn handle_execute_remove_cached_remote_relation_command(