diff --git a/Cargo.lock b/Cargo.lock index 8b2b6bb3c3..ba1fcef33d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6975,6 +6975,7 @@ dependencies = [ "serde_arrow", "thiserror 2.0.18", "tokio", + "url", ] [[package]] @@ -7151,6 +7152,8 @@ dependencies = [ "futures", "lazy_static", "lexical-core", + "log", + "object_store", "parquet-variant-compute", "parquet-variant-json", "pin-project-lite", @@ -7170,6 +7173,7 @@ dependencies = [ "tokio", "tonic", "tonic-prost-build", + "url", ] [[package]] @@ -7611,6 +7615,7 @@ dependencies = [ "futures", "indexmap 2.14.0", "log", + "object_store", "readonly", "sail-cache", "sail-catalog", @@ -7641,6 +7646,7 @@ dependencies = [ "thiserror 2.0.18", "tokio", "tonic", + "url", ] [[package]] @@ -7654,6 +7660,7 @@ dependencies = [ "lazy_static", "log", "monostate", + "object_store", "pbjson", "pbjson-build", "pbjson-types", @@ -7666,6 +7673,7 @@ dependencies = [ "quote", "regex", "sail-cache", + "sail-catalog", "sail-common", "sail-common-datafusion", "sail-data-source", @@ -7684,6 +7692,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/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 db94464a22..86f2f0c87f 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; @@ -147,6 +149,29 @@ impl CatalogManager { ) -> CatalogResult> { self.tracker.get_tracked_logical_plan(id) } + + pub fn track_cached_relation( + &self, + relation_id: String, + relation: CatalogCachedRelation, + ) -> CatalogResult<()> { + self.tracker.track_cached_relation(relation_id, relation) + } + + pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult { + self.tracker.get_cached_relation(relation_id) + } + + pub fn remove_cached_relation( + &self, + relation_id: &str, + ) -> 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 b7549055c3..fead9d3452 100644 --- a/crates/sail-catalog/src/manager/tracker.rs +++ b/crates/sail-catalog/src/manager/tracker.rs @@ -12,12 +12,25 @@ pub struct CatalogFunctionId(u64); #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash, PartialOrd, Serialize, Deserialize)] pub struct CatalogLogicalPlanId(u64); +#[derive(Debug, Clone)] +pub struct CatalogCachedRelation { + pub plan: Arc, + pub fields: Vec, + pub storage_uri: Option, +} + #[derive(Default)] struct CatalogObjectTrackerState { next_function_id: u64, next_logical_plan_id: u64, functions: HashMap, logical_plans: HashMap>, + /// Maps relation IDs from `CachedRemoteRelation` to their checkpointed + /// logical plans and user-facing field names. Populated when the server + /// handles a `CheckpointCommand`, and queried when resolving a + /// `CachedRemoteRelation` query node. Entries are removed when the server + /// handles a `RemoveCachedRemoteRelationCommand`. + cached_relations: HashMap, } /// Tracks in-memory objects (UDFs and logical plans) that cannot be serialized directly, @@ -75,4 +88,45 @@ impl CatalogObjectTracker { .cloned() .ok_or_else(|| CatalogError::NotFound(CatalogObject::LogicalPlan, id.0.to_string())) } + + /// Stores a logical plan keyed by an opaque relation ID. + /// This is used to back checkpointed DataFrames referenced via + /// `CachedRemoteRelation` in subsequent Spark Connect requests. + pub fn track_cached_relation( + &self, + relation_id: String, + relation: CatalogCachedRelation, + ) -> CatalogResult<()> { + let mut state = self.state()?; + state.cached_relations.insert(relation_id, relation); + Ok(()) + } + + pub fn get_cached_relation(&self, relation_id: &str) -> CatalogResult { + let state = self.state()?; + state + .cached_relations + .get(relation_id) + .cloned() + .ok_or_else(|| { + CatalogError::NotFound(CatalogObject::CachedRelation, relation_id.to_string()) + }) + } + + pub fn remove_cached_relation( + &self, + relation_id: &str, + ) -> CatalogResult> { + 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()) + } } 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-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..bccfd58a29 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 + location that is shared and readable by all workers, such as an + object-store URI or a shared filesystem path, since the default location + is node-local. Checkpoint files are best-effort cleaned up when the client + releases the cached relation or when the session ends. + experimental: true + - key: flight.session_timeout_secs type: number default: "3600" diff --git a/crates/sail-plan/src/resolver/query/misc.rs b/crates/sail-plan/src/resolver/query/misc.rs index a49c825d8c..e65cdbc940 100644 --- a/crates/sail-plan/src/resolver/query/misc.rs +++ b/crates/sail-plan/src/resolver/query/misc.rs @@ -5,9 +5,13 @@ use datafusion::catalog::MemTable; use datafusion_common::{DFSchema, DFSchemaRef, ParamValues}; use datafusion_expr::{EmptyRelation, Extension, LogicalPlan, UNNAMED_TABLE}; use log::warn; +use sail_catalog::error::{CatalogError, CatalogObject}; +use sail_catalog::manager::CatalogManager; use sail_common::spec; use sail_common_datafusion::array::record_batch::{cast_record_batch, read_record_batches}; +use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::literal::LiteralEvaluator; +use sail_common_datafusion::rename::logical_plan::rename_logical_plan; use sail_logical_plan::range::RangeNode; use crate::error::{PlanError, PlanResult}; @@ -138,6 +142,33 @@ impl PlanResolver<'_> { ) } + /// Resolves a query that refers to a cached remote relation by looking up + /// the previously checkpointed logical plan from the catalog manager. The + /// cached plan was stored with internal field IDs and its user-facing field + /// names, so we register fresh field IDs in the current resolver state and + /// rename the plan accordingly. This preserves duplicate user-facing names + /// without relying on them as `Expr::Column` references inside the cached + /// plan. + pub(super) async fn resolve_query_cached_remote_relation( + &self, + relation_id: String, + state: &mut PlanResolverState, + ) -> PlanResult { + let manager = self.ctx.extension::()?; + let relation = match manager.get_cached_relation(&relation_id) { + Ok(relation) => relation, + Err(CatalogError::NotFound(CatalogObject::CachedRelation, _)) => { + return Err(PlanError::AnalysisError(format!( + "cached remote relation not found: {relation_id}" + ))); + } + Err(error) => return Err(error.into()), + }; + let names = state.register_field_names(relation.fields.clone()); + let plan = rename_logical_plan((*relation.plan).clone(), &names)?; + Ok(plan) + } + pub(super) async fn resolve_query_hint( &self, input: spec::QueryPlan, 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-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..81427afbaa 100644 --- a/crates/sail-session/src/session_manager/actor/handler.rs +++ b/crates/sail-session/src/session_manager/actor/handler.rs @@ -1,10 +1,13 @@ 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 log::{info, warn}; +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; @@ -326,11 +329,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 +360,9 @@ 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; + } +} diff --git a/crates/sail-spark-connect/Cargo.toml b/crates/sail-spark-connect/Cargo.toml index 68045d157a..60dd09c998 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" } @@ -40,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/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..5cae327f45 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -1,17 +1,33 @@ +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::prelude::SessionContext; +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; +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 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; +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 +42,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 +532,66 @@ 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, + // 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")?; + 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 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 { + materialize_logical_plan_to_disk(ctx, logical_plan).await? + } else { + ( + materialize_logical_plan_to_memory(ctx, logical_plan).await?, + None, + ) + }; + + let relation_id = uuid::Uuid::new_v4().to_string(); + let manager = ctx.extension::()?; + if let Err(error) = manager.track_cached_relation( + relation_id.clone(), + CatalogCachedRelation { + plan: Arc::new(cached_plan), + fields, + storage_uri: storage_uri.clone(), + }, + ) { + if let Some(uri) = storage_uri { + cleanup_checkpoint_path(ctx, &uri).await; + } + return Err(SparkError::internal(format!( + "failed to track cached relation: {error}" + ))); + } + + let result = CheckpointCommandResult { + relation: Some(CachedRemoteRelation { relation_id }), + }; let mut output = vec![ExecutorOutput::new(ExecutorBatch::CheckpointCommandResult( Box::new(result), ))]; @@ -535,6 +605,243 @@ pub(crate) async fn handle_execute_checkpoint_command( )) } +/// Determines whether a checkpoint should be materialized to disk. +/// +/// Non-local (reliable) checkpoints always use disk. For local checkpoints +/// the decision follows the explicit `StorageLevel` when provided. When the +/// client omits the storage level (which is the common case – PySpark's +/// `localCheckpoint()` defaults to `MEMORY_AND_DISK` but does **not** send the +/// level over the wire), we default to **disk** to avoid collecting the entire +/// DataFrame into server memory, which could easily OOM for large datasets. +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).unwrap_or(true) +} + +/// 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_stream( + ctx: &SessionContext, + plan: LogicalPlan, +) -> 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(); + 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?; + 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)) +} + +/// 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) = 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), + provider_as_source(table), + None, + )? + .build()?; + Ok(scan) +} + +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) +} + +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. +/// +/// In cluster deployments the operator should configure `spark.checkpoint_dir` +/// 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() { + directory_file_url(std::env::temp_dir().join("sail-checkpoints")) + } 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)) + } +} + +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. +/// +/// 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, +) -> 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)?; + 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)?; + + 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. + let write_result: SparkResult<()> = async { + 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 + .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; + + if let Err(error) = write_result { + cleanup_checkpoint_path(ctx, &checkpoint_uri).await; + return Err(error); + } + + let df = match ctx + .read_parquet(&checkpoint_uri, ParquetReadOptions::default()) + .await + { + Ok(df) => df, + Err(error) => { + cleanup_checkpoint_path(ctx, &checkpoint_uri).await; + return Err(error.into()); + } + }; + Ok((df.into_unoptimized_plan(), Some(checkpoint_uri))) +} + +async fn cleanup_checkpoint_path(ctx: &SessionContext, checkpoint_uri: &str) { + cleanup_checkpoint_storage(ctx.runtime_env().as_ref(), checkpoint_uri).await; +} + +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::()?; + // 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 removed = manager + .remove_cached_relation(&relation.relation_id) + .map_err(|e| SparkError::internal(format!("failed to remove cached relation: {e}")))?; + if let Some(uri) = removed.and_then(|relation| relation.storage_uri) { + cleanup_checkpoint_path(ctx, &uri).await; + } + + 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/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}"))?; diff --git a/python/pysail/tests/spark/test_checkpoint.py b/python/pysail/tests/spark/test_checkpoint.py new file mode 100644 index 0000000000..2fc0ca29c1 --- /dev/null +++ b/python/pysail/tests/spark/test_checkpoint.py @@ -0,0 +1,141 @@ +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 + +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() + + +@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"]}), + ) + + +@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() + rows = sorted(pdf.values.tolist()) + + assert checkpointed.columns == ["id", "id"] + assert pdf.columns.tolist() == ["id", "id"] + assert rows == [[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")) + 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): + 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) == expected_rows + assert 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]}), + ) diff --git a/scripts/spark-tests/plugins/spark.py b/scripts/spark-tests/plugins/spark.py index 70a7ed995f..f6c3296dc1 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 + # 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 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.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.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))