From 05caceb5531bdadc7d698465d4ebd7f0f06f8803 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Thu, 21 May 2026 23:54:28 +0800 Subject: [PATCH 01/40] base --- .../src/cached_relation.rs | 75 +++++++ crates/sail-common-datafusion/src/lib.rs | 1 + crates/sail-plan/src/resolver/query/mod.rs | 14 +- .../src/session_factory/server.rs | 2 + crates/sail-spark-connect/src/server.rs | 5 +- .../src/service/plan_executor.rs | 201 +++++++++++++++++- .../pysail/data/compatibility/dataframe.json | 6 +- python/pysail/tests/spark/test_dataframe.py | 92 ++++++++ 8 files changed, 381 insertions(+), 15 deletions(-) create mode 100644 crates/sail-common-datafusion/src/cached_relation.rs diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs new file mode 100644 index 0000000000..df1c19ae48 --- /dev/null +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -0,0 +1,75 @@ +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; + +use datafusion_common::{internal_datafusion_err, Result}; +use datafusion_expr::LogicalPlan; + +use crate::extension::SessionExtension; + +#[derive(Debug, Clone)] +pub enum CachedRelationCleanup { + LocalPath(String), +} + +#[derive(Debug, Clone)] +pub struct CachedRelation { + plan: Arc, + cleanup: Option, +} + +impl CachedRelation { + pub fn new(plan: Arc, cleanup: Option) -> Self { + Self { plan, cleanup } + } + + pub fn plan(&self) -> &Arc { + &self.plan + } + + pub fn into_cleanup(self) -> Option { + self.cleanup + } +} + +#[derive(Debug, Default)] +pub struct CachedRelationRegistry { + // TODO: Clean cached local checkpoint relations when a session is closed. + relations: RwLock>, +} + +impl CachedRelationRegistry { + pub fn insert(&self, relation_id: String, relation: CachedRelation) -> Result<()> { + let mut relations = self + .relations + .write() + .map_err(|e| internal_datafusion_err!("{e}"))?; + if relations.insert(relation_id.clone(), relation).is_some() { + return Err(internal_datafusion_err!( + "cached relation already exists: {relation_id}" + )); + } + Ok(()) + } + + pub fn get(&self, relation_id: &str) -> Result> { + let relations = self + .relations + .read() + .map_err(|e| internal_datafusion_err!("{e}"))?; + Ok(relations.get(relation_id).cloned()) + } + + pub fn remove(&self, relation_id: &str) -> Result> { + let mut relations = self + .relations + .write() + .map_err(|e| internal_datafusion_err!("{e}"))?; + Ok(relations.remove(relation_id)) + } +} + +impl SessionExtension for CachedRelationRegistry { + fn name() -> &'static str { + "cached relation registry" + } +} diff --git a/crates/sail-common-datafusion/src/lib.rs b/crates/sail-common-datafusion/src/lib.rs index 10e66400b7..01fb551837 100644 --- a/crates/sail-common-datafusion/src/lib.rs +++ b/crates/sail-common-datafusion/src/lib.rs @@ -1,4 +1,5 @@ pub mod array; +pub mod cached_relation; pub mod catalog; pub mod column_features; pub mod datasource; diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index 13dbb0a01a..b067d81612 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/crates/sail-plan/src/resolver/query/mod.rs @@ -1,6 +1,9 @@ use async_recursion::async_recursion; use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; use sail_common::spec; +use sail_common_datafusion::cached_relation::CachedRelationRegistry; +use sail_common_datafusion::extension::SessionExtensionAccessor; +use sail_common_datafusion::rename::logical_plan::rename_logical_plan; use crate::error::{PlanError, PlanResult}; use crate::resolver::state::PlanResolverState; @@ -227,8 +230,15 @@ impl PlanResolver<'_> { QueryNode::CachedLocalRelation { .. } => { return Err(PlanError::todo("cached local relation")); } - QueryNode::CachedRemoteRelation { .. } => { - return Err(PlanError::todo("cached remote relation")); + QueryNode::CachedRemoteRelation { relation_id } => { + let registry = self.ctx.extension::()?; + let relation = registry.get(&relation_id)?.ok_or_else(|| { + PlanError::invalid(format!("cached remote relation not found: {relation_id}")) + })?; + let plan = relation.plan().as_ref().clone(); + let names = state.register_fields(plan.schema().inner().fields()); + // TODO: Preserve Spark attribute metadata/expr IDs for full checkpoint parity. + rename_logical_plan(plan, &names)? } QueryNode::CommonInlineUserDefinedTableFunction(udtf) => { self.resolve_query_common_inline_udtf(udtf, state).await? diff --git a/crates/sail-session/src/session_factory/server.rs b/crates/sail-session/src/session_factory/server.rs index d4a7f4c515..617c74e95d 100644 --- a/crates/sail-session/src/session_factory/server.rs +++ b/crates/sail-session/src/session_factory/server.rs @@ -13,6 +13,7 @@ use sail_catalog::provider::CatalogCacheManager; use sail_catalog_system::service::SystemTableService; use sail_common::config::{AppConfig, ExecutionMode}; use sail_common::runtime::RuntimeHandle; +use sail_common_datafusion::cached_relation::CachedRelationRegistry; use sail_common_datafusion::session::activity::ActivityTracker; use sail_common_datafusion::session::job::{JobRunner, JobService}; use sail_delta_lake::session_extension::DeltaTableCache; @@ -120,6 +121,7 @@ impl ServerSessionFactory { self.catalog_cache_manager.clone(), )?)) .with_extension(Arc::new(ActivityTracker::new())) + .with_extension(Arc::new(CachedRelationRegistry::default())) .with_extension(Arc::new(JobService::new(job_runner))) .with_extension(Arc::new(self.create_system_table_service(info)?)) .with_extension(Arc::new(DeltaTableCache::default())); 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..2effe6f799 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -1,20 +1,26 @@ +use std::path::{Component, Path}; 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::prelude::{ParquetReadOptions, SessionContext}; use fastrace::collector::SpanContext; use fastrace::future::FutureExt; use fastrace::Span; use futures::stream; use log::{debug, warn}; use sail_common::spec; +use sail_common_datafusion::cached_relation::{ + CachedRelation, CachedRelationCleanup, CachedRelationRegistry, +}; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; use sail_plan::resolve_and_execute_plan; use tonic::codegen::tokio_stream::wrappers::ReceiverStream; use tonic::codegen::tokio_stream::Stream; use tonic::Status; +use uuid::Uuid; use crate::error::{ProtoFieldExt, SparkError, SparkResult}; use crate::executor::{ @@ -26,16 +32,19 @@ 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, }; use crate::streaming::timeout_millis; +const SPARK_CHECKPOINT_DIR: &str = "spark.checkpoint.dir"; + pub struct ExecutePlanResponseStream { session_id: String, operation_id: String, @@ -515,13 +524,63 @@ 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 spark = ctx.extension::()?; - let result = CheckpointCommandResult { relation: None }; + let CheckpointCommand { + relation, + local, + eager, + storage_level, + } = checkpoint; + if !eager { + // TODO: Implement Spark-compatible lazy checkpoint materialization on first action. + return Err(SparkError::unsupported( + "DataFrame checkpoint with eager=false is not supported", + )); + } + if local { + validate_local_checkpoint_storage_level(storage_level)?; + } else if storage_level.is_some() { + return Err(SparkError::invalid( + "StorageLevel is only valid for localCheckpoint", + )); + } + + let relation = relation.required("checkpoint relation")?; + let plan: spec::Plan = relation.try_into()?; + let relation_id = Uuid::new_v4().to_string(); + let path = checkpoint_path(&spark, local, &relation_id)?; + let (physical_plan, _) = resolve_and_execute_plan(ctx, spark.plan_config()?, plan).await?; + let schema = physical_plan.schema(); + // TODO: Use Sail's file-write physical path through JobRunner for cluster-native checkpointing. + if let Err(error) = ctx.write_parquet(physical_plan, &path, None).await { + return Err(cleanup_checkpoint_after_error(&path, error.into()).await); + } + let register_result: SparkResult<()> = async { + let read_options = ParquetReadOptions::new().schema(schema.as_ref()); + let df = ctx.read_parquet(&path, read_options).await?; + let (_, read_plan) = df.into_parts(); + let cleanup = if local { + Some(CachedRelationCleanup::LocalPath(path.clone())) + } else { + // TODO: Add Spark-like reliable checkpoint cleanup/reference tracking. + None + }; + ctx.extension::()?.insert( + relation_id.clone(), + CachedRelation::new(Arc::new(read_plan), cleanup), + )?; + Ok(()) + } + .await; + if let Err(error) = register_result { + return Err(cleanup_checkpoint_after_error(&path, error).await); + } + let result = CheckpointCommandResult { + relation: Some(CachedRemoteRelation { relation_id }), + }; let mut output = vec![ExecutorOutput::new(ExecutorBatch::CheckpointCommandResult( Box::new(result), ))]; @@ -535,6 +594,130 @@ pub(crate) async fn handle_execute_checkpoint_command( )) } +pub(crate) async fn handle_execute_remove_cached_remote_relation_command( + ctx: &SessionContext, + command: RemoveCachedRemoteRelationCommand, + metadata: ExecutorMetadata, +) -> SparkResult { + let spark = ctx.extension::()?; + let relation = command + .relation + .required("remove cached remote relation relation")?; + if let Some(relation) = ctx + .extension::()? + .remove(&relation.relation_id)? + { + cleanup_cached_relation(relation).await?; + } + let output = if metadata.reattachable { + vec![ExecutorOutput::complete()] + } else { + vec![] + }; + Ok(ExecutePlanResponseStream::new( + spark.session_id().to_string(), + metadata.operation_id, + Box::pin(stream::iter(output)), + )) +} + +fn validate_local_checkpoint_storage_level( + storage_level: Option, +) -> SparkResult<()> { + let Some(storage_level) = storage_level else { + return Ok(()); + }; + let storage_level: spec::StorageLevel = storage_level.try_into()?; + let is_default = storage_level.use_disk + && storage_level.use_memory + && !storage_level.use_off_heap + && !storage_level.deserialized + && storage_level.replication == 1; + if !is_default { + // TODO: Map Spark StorageLevel once Sail cache/persist has real storage semantics. + return Err(SparkError::unsupported( + "non-default StorageLevel for localCheckpoint is not supported", + )); + } + Ok(()) +} + +fn checkpoint_path(spark: &SparkSession, local: bool, relation_id: &str) -> SparkResult { + let root = if local { + // TODO: Use Sail runtime temporary-file directories and quota for local checkpoint storage. + std::env::temp_dir().to_string_lossy().into_owned() + } else { + // TODO: Wire Spark-compatible startup config / setCheckpointDir semantics. + spark + .get_config_option(vec![SPARK_CHECKPOINT_DIR.to_string()])? + .into_iter() + .next() + .and_then(|kv| kv.value) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + SparkError::invalid(format!( + "checkpoint directory is not set; set {SPARK_CHECKPOINT_DIR}" + )) + })? + }; + validate_checkpoint_directory(&root)?; + validate_checkpoint_path_segment("session ID", spark.session_id())?; + validate_checkpoint_path_segment("relation ID", relation_id)?; + Ok(format!( + "{}/sail-checkpoints/{}/{}", + root.trim_end_matches('/'), + spark.session_id(), + relation_id + )) +} + +fn validate_checkpoint_directory(root: &str) -> SparkResult<()> { + if Path::new(root) + .components() + .any(|component| matches!(component, Component::ParentDir)) + { + return Err(SparkError::invalid( + "checkpoint directory cannot contain parent path components", + )); + } + Ok(()) +} + +fn validate_checkpoint_path_segment(name: &str, value: &str) -> SparkResult<()> { + let mut components = Path::new(value).components(); + match (components.next(), components.next()) { + (Some(Component::Normal(_)), None) => Ok(()), + _ => Err(SparkError::invalid(format!( + "checkpoint {name} must be a single path segment" + ))), + } +} + +async fn cleanup_checkpoint_after_error(path: &str, error: SparkError) -> SparkError { + match cleanup_checkpoint_path(path).await { + Ok(()) => error, + Err(cleanup_error) => SparkError::internal(format!( + "checkpoint failed: {error}; additionally failed to clean checkpoint path {path}: {cleanup_error}" + )), + } +} + +async fn cleanup_checkpoint_path(path: &str) -> SparkResult<()> { + match tokio::fs::remove_dir_all(path).await { + Ok(()) => Ok(()), + Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), + Err(e) => Err(e.into()), + } +} + +async fn cleanup_cached_relation(relation: CachedRelation) -> SparkResult<()> { + match relation.into_cleanup() { + Some(CachedRelationCleanup::LocalPath(path)) => cleanup_checkpoint_path(&path).await, + None => Ok(()), + } +} + pub(crate) async fn handle_interrupt_all(ctx: &SessionContext) -> SparkResult> { let spark = ctx.extension::()?; let mut results = vec![]; diff --git a/python/pysail/data/compatibility/dataframe.json b/python/pysail/data/compatibility/dataframe.json index 9624100524..754f10a462 100644 --- a/python/pysail/data/compatibility/dataframe.json +++ b/python/pysail/data/compatibility/dataframe.json @@ -23,7 +23,8 @@ { "module": "pyspark.sql.DataFrame", "function": "checkpoint", - "status": "planned", + "status": "partially supported", + "note": "Supports eager=True with spark.checkpoint.dir configured. eager=False is not supported yet.", "ref": "https://github.com/lakehq/sail/issues/476" }, { @@ -253,7 +254,8 @@ { "module": "pyspark.sql.DataFrame", "function": "localCheckpoint", - "status": "planned", + "status": "partially supported", + "note": "Supports eager=True with default storage level. eager=False and non-default storage levels are not supported yet.", "ref": "https://github.com/lakehq/sail/issues/482" }, { diff --git a/python/pysail/tests/spark/test_dataframe.py b/python/pysail/tests/spark/test_dataframe.py index 034753a06a..9ad61288b1 100644 --- a/python/pysail/tests/spark/test_dataframe.py +++ b/python/pysail/tests/spark/test_dataframe.py @@ -1,9 +1,12 @@ 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 is_jvm_spark + def test_dataframe_drop(spark): df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) @@ -72,6 +75,95 @@ def test_dataframe_drop(spark): ) +def test_dataframe_local_checkpoint(spark): + df = spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b")], + ) + + checkpointed = df.where(col("id") >= 1).localCheckpoint() + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + target_id = 2 + assert_frame_equal( + checkpointed.where(col("id") == target_id).select("value").toPandas(), + pd.DataFrame({"value": ["b"]}), + ) + + +@pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") +def test_dataframe_checkpoint(spark, tmp_path): + df = spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b")], + ) + spark.conf.set("spark.checkpoint.dir", str(tmp_path / "checkpoints")) + try: + checkpointed = df.where(col("id") >= 1).checkpoint() + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + assert_frame_equal( + checkpointed.where(col("id") == 1).select("value").toPandas(), + pd.DataFrame({"value": ["a"]}), + ) + finally: + spark.conf.unset("spark.checkpoint.dir") + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific missing checkpoint dir coverage") +def test_dataframe_checkpoint_requires_directory(spark_session_factory): + spark = spark_session_factory() + df = spark.createDataFrame( + schema="id INT", + data=[(1,)], + ) + + with pytest.raises(Exception, match=r"spark\.checkpoint\.dir"): + df.checkpoint() + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint path validation") +def test_dataframe_checkpoint_rejects_parent_path_components(spark, tmp_path): + df = spark.createDataFrame( + schema="id INT", + data=[(1,)], + ) + spark.conf.set("spark.checkpoint.dir", str(tmp_path / ".." / "checkpoints")) + try: + with pytest.raises(Exception, match="parent path"): + df.checkpoint() + finally: + spark.conf.unset("spark.checkpoint.dir") + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail MVP limitation") +def test_dataframe_checkpoint_lazy_is_not_supported(spark): + df = spark.createDataFrame( + schema="id INT", + data=[(1,)], + ) + + with pytest.raises(Exception, match="eager=false"): + df.localCheckpoint(eager=False) + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail MVP limitation") +def test_dataframe_local_checkpoint_non_default_storage_level_is_not_supported(spark): + df = spark.createDataFrame( + schema="id INT", + data=[(1,)], + ) + + with pytest.raises(Exception, match="StorageLevel"): + df.localCheckpoint(storageLevel=StorageLevel.MEMORY_ONLY) + + def test_dataframe_with_column_alias(spark): df = spark.createDataFrame( schema="id INTEGER, value STRING", From cf4ebfad992cba571fa7a66ec4f2ad4b4ad14526 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 22 May 2026 00:54:25 +0800 Subject: [PATCH 02/40] update --- .../src/cached_relation.rs | 63 ++++++- crates/sail-plan/src/resolver/query/mod.rs | 4 +- .../src/service/plan_executor.rs | 168 +++++++++++------- .../pysail/data/compatibility/dataframe.json | 2 +- python/pysail/tests/spark/test_dataframe.py | 22 ++- 5 files changed, 183 insertions(+), 76 deletions(-) diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index df1c19ae48..f0f1e4a32f 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -1,8 +1,12 @@ use std::collections::HashMap; use std::sync::{Arc, RwLock}; -use datafusion_common::{internal_datafusion_err, Result}; -use datafusion_expr::LogicalPlan; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::datasource::{provider_as_source, MemTable}; +use datafusion_common::{internal_datafusion_err, Result, TableReference}; +use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; +use sail_common::spec; use crate::extension::SessionExtension; @@ -13,17 +17,64 @@ pub enum CachedRelationCleanup { #[derive(Debug, Clone)] pub struct CachedRelation { - plan: Arc, + data: CachedRelationData, cleanup: Option, } +#[derive(Debug, Clone)] +enum CachedRelationData { + LogicalPlan(Arc), + Memory { + schema: SchemaRef, + partitions: Arc>>, + storage_level: Option, + }, +} + impl CachedRelation { pub fn new(plan: Arc, cleanup: Option) -> Self { - Self { plan, cleanup } + Self { + data: CachedRelationData::LogicalPlan(plan), + cleanup, + } } - pub fn plan(&self) -> &Arc { - &self.plan + pub fn new_memory( + schema: SchemaRef, + partitions: Vec>, + storage_level: Option, + ) -> Self { + Self { + data: CachedRelationData::Memory { + schema, + partitions: Arc::new(partitions), + storage_level, + }, + cleanup: None, + } + } + + pub fn to_logical_plan(&self, relation_id: &str) -> Result { + match &self.data { + CachedRelationData::LogicalPlan(plan) => Ok(plan.as_ref().clone()), + CachedRelationData::Memory { + schema, + partitions, + storage_level, + } => { + let _ = storage_level; + // TODO: Use Sail's full cache storage tiers once memory/disk/off-heap cache exists. + // Sail's physical-plan codec serializes MemorySourceConfig batches, so this cache + // representation can still be executed by local-cluster workers. + let table = MemTable::try_new(Arc::clone(schema), partitions.as_ref().clone())?; + LogicalPlanBuilder::scan( + TableReference::bare(format!("__sail_cached_relation_{relation_id}")), + provider_as_source(Arc::new(table)), + None, + )? + .build() + } + } } pub fn into_cleanup(self) -> Option { diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index b067d81612..21fd920578 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/crates/sail-plan/src/resolver/query/mod.rs @@ -233,9 +233,9 @@ impl PlanResolver<'_> { QueryNode::CachedRemoteRelation { relation_id } => { let registry = self.ctx.extension::()?; let relation = registry.get(&relation_id)?.ok_or_else(|| { - PlanError::invalid(format!("cached remote relation not found: {relation_id}")) + PlanError::invalid(format!("No DataFrame with id {relation_id} is found")) })?; - let plan = relation.plan().as_ref().clone(); + let plan = relation.to_logical_plan(&relation_id)?; let names = state.register_fields(plan.schema().inner().fields()); // TODO: Preserve Spark attribute metadata/expr IDs for full checkpoint parity. rename_logical_plan(plan, &names)? diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 2effe6f799..95a5f488c8 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -540,43 +540,37 @@ pub(crate) async fn handle_execute_checkpoint_command( "DataFrame checkpoint with eager=false is not supported", )); } - if local { - validate_local_checkpoint_storage_level(storage_level)?; + let storage_level = if local { + Some(resolve_local_checkpoint_storage_level(storage_level)?) } else if storage_level.is_some() { return Err(SparkError::invalid( "StorageLevel is only valid for localCheckpoint", )); - } + } else { + None + }; let relation = relation.required("checkpoint relation")?; let plan: spec::Plan = relation.try_into()?; let relation_id = Uuid::new_v4().to_string(); - let path = checkpoint_path(&spark, local, &relation_id)?; let (physical_plan, _) = resolve_and_execute_plan(ctx, spark.plan_config()?, plan).await?; - let schema = physical_plan.schema(); - // TODO: Use Sail's file-write physical path through JobRunner for cluster-native checkpointing. - if let Err(error) = ctx.write_parquet(physical_plan, &path, None).await { - return Err(cleanup_checkpoint_after_error(&path, error.into()).await); - } - let register_result: SparkResult<()> = async { - let read_options = ParquetReadOptions::new().schema(schema.as_ref()); - let df = ctx.read_parquet(&path, read_options).await?; - let (_, read_plan) = df.into_parts(); - let cleanup = if local { - Some(CachedRelationCleanup::LocalPath(path.clone())) - } else { - // TODO: Add Spark-like reliable checkpoint cleanup/reference tracking. - None - }; - ctx.extension::()?.insert( - relation_id.clone(), - CachedRelation::new(Arc::new(read_plan), cleanup), - )?; - Ok(()) - } - .await; - if let Err(error) = register_result { - return Err(cleanup_checkpoint_after_error(&path, error).await); + let (relation, cleanup_on_insert_error) = if local { + ( + materialize_local_checkpoint(ctx, physical_plan, storage_level).await?, + None, + ) + } else { + materialize_reliable_checkpoint(ctx, &spark, &relation_id, physical_plan).await? + }; + if let Err(error) = ctx + .extension::()? + .insert(relation_id.clone(), relation) + { + let error = SparkError::from(error); + if let Some(path) = cleanup_on_insert_error { + return Err(cleanup_checkpoint_after_error(&path, error).await); + } + return Err(error); } let result = CheckpointCommandResult { relation: Some(CachedRemoteRelation { relation_id }), @@ -594,6 +588,50 @@ pub(crate) async fn handle_execute_checkpoint_command( )) } +async fn materialize_local_checkpoint( + ctx: &SessionContext, + physical_plan: Arc, + storage_level: Option, +) -> SparkResult { + let schema = physical_plan.schema(); + let service = ctx.extension::()?; + let stream = service.runner().execute(ctx, physical_plan).await?; + let batches = read_stream(stream).await?; + Ok(CachedRelation::new_memory( + schema, + vec![batches], + storage_level, + )) +} + +async fn materialize_reliable_checkpoint( + ctx: &SessionContext, + spark: &SparkSession, + relation_id: &str, + physical_plan: Arc, +) -> SparkResult<(CachedRelation, Option)> { + let path = checkpoint_path(spark, relation_id)?; + let schema = physical_plan.schema(); + // TODO: Use Sail's file-write physical path through JobRunner for cluster-native checkpointing. + if let Err(error) = ctx.write_parquet(physical_plan, &path, None).await { + return Err(cleanup_checkpoint_after_error(&path, error.into()).await); + } + let register_result: SparkResult = async { + let read_options = ParquetReadOptions::new().schema(schema.as_ref()); + let df = ctx.read_parquet(&path, read_options).await?; + let (_, read_plan) = df.into_parts(); + Ok(CachedRelation::new(Arc::new(read_plan), None)) + } + .await; + match register_result { + Ok(relation) => { + // TODO: Add Spark-like reliable checkpoint cleanup/reference tracking. + Ok((relation, Some(path))) + } + Err(error) => Err(cleanup_checkpoint_after_error(&path, error).await), + } +} + pub(crate) async fn handle_execute_remove_cached_remote_relation_command( ctx: &SessionContext, command: RemoveCachedRemoteRelationCommand, @@ -621,46 +659,50 @@ pub(crate) async fn handle_execute_remove_cached_remote_relation_command( )) } -fn validate_local_checkpoint_storage_level( +fn resolve_local_checkpoint_storage_level( storage_level: Option, -) -> SparkResult<()> { - let Some(storage_level) = storage_level else { - return Ok(()); - }; - let storage_level: spec::StorageLevel = storage_level.try_into()?; - let is_default = storage_level.use_disk - && storage_level.use_memory - && !storage_level.use_off_heap - && !storage_level.deserialized - && storage_level.replication == 1; - if !is_default { - // TODO: Map Spark StorageLevel once Sail cache/persist has real storage semantics. +) -> SparkResult { + let storage_level: spec::StorageLevel = storage_level + .map(TryInto::try_into) + .transpose()? + .map(Ok) + .unwrap_or_else(|| "MEMORY_AND_DISK".parse().map_err(SparkError::from))?; + + if !storage_level.use_memory { return Err(SparkError::unsupported( - "non-default StorageLevel for localCheckpoint is not supported", + "localCheckpoint StorageLevel without memory is not supported", + )); + } + if storage_level.use_off_heap { + return Err(SparkError::unsupported( + "localCheckpoint off-heap StorageLevel is not supported", + )); + } + if storage_level.replication != 1 { + return Err(SparkError::unsupported( + "localCheckpoint replicated StorageLevel is not supported", )); } - Ok(()) -} -fn checkpoint_path(spark: &SparkSession, local: bool, relation_id: &str) -> SparkResult { - let root = if local { - // TODO: Use Sail runtime temporary-file directories and quota for local checkpoint storage. - std::env::temp_dir().to_string_lossy().into_owned() - } else { - // TODO: Wire Spark-compatible startup config / setCheckpointDir semantics. - spark - .get_config_option(vec![SPARK_CHECKPOINT_DIR.to_string()])? - .into_iter() - .next() - .and_then(|kv| kv.value) - .map(|value| value.trim().to_string()) - .filter(|value| !value.is_empty()) - .ok_or_else(|| { - SparkError::invalid(format!( - "checkpoint directory is not set; set {SPARK_CHECKPOINT_DIR}" - )) - })? - }; + // TODO: Enforce disk spill and serialized/deserialized storage tiers once Sail has a + // DataFrame/block cache subsystem. For now we only support memory-backed local checkpoint. + Ok(storage_level) +} + +fn checkpoint_path(spark: &SparkSession, relation_id: &str) -> SparkResult { + // TODO: Wire Spark-compatible startup config / setCheckpointDir semantics. + let root = spark + .get_config_option(vec![SPARK_CHECKPOINT_DIR.to_string()])? + .into_iter() + .next() + .and_then(|kv| kv.value) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + SparkError::invalid(format!( + "checkpoint directory is not set; set {SPARK_CHECKPOINT_DIR}" + )) + })?; validate_checkpoint_directory(&root)?; validate_checkpoint_path_segment("session ID", spark.session_id())?; validate_checkpoint_path_segment("relation ID", relation_id)?; diff --git a/python/pysail/data/compatibility/dataframe.json b/python/pysail/data/compatibility/dataframe.json index 754f10a462..24e3016533 100644 --- a/python/pysail/data/compatibility/dataframe.json +++ b/python/pysail/data/compatibility/dataframe.json @@ -255,7 +255,7 @@ "module": "pyspark.sql.DataFrame", "function": "localCheckpoint", "status": "partially supported", - "note": "Supports eager=True with default storage level. eager=False and non-default storage levels are not supported yet.", + "note": "Supports eager=True with memory-backed cache materialization. StorageLevel values without memory, off-heap storage, replication, eager=False, and full storage tier semantics are not supported yet.", "ref": "https://github.com/lakehq/sail/issues/482" }, { diff --git a/python/pysail/tests/spark/test_dataframe.py b/python/pysail/tests/spark/test_dataframe.py index 9ad61288b1..8a7a4b66a8 100644 --- a/python/pysail/tests/spark/test_dataframe.py +++ b/python/pysail/tests/spark/test_dataframe.py @@ -153,15 +153,29 @@ def test_dataframe_checkpoint_lazy_is_not_supported(spark): df.localCheckpoint(eager=False) -@pytest.mark.skipif(is_jvm_spark(), reason="Sail MVP limitation") -def test_dataframe_local_checkpoint_non_default_storage_level_is_not_supported(spark): +def test_dataframe_local_checkpoint_with_memory_storage_level(spark): + df = spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b")], + ) + + checkpointed = df.localCheckpoint(storageLevel=StorageLevel.MEMORY_ONLY) + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail cache storage tier limitation") +def test_dataframe_local_checkpoint_disk_only_storage_level_is_not_supported(spark): df = spark.createDataFrame( schema="id INT", data=[(1,)], ) - with pytest.raises(Exception, match="StorageLevel"): - df.localCheckpoint(storageLevel=StorageLevel.MEMORY_ONLY) + with pytest.raises(Exception, match="without memory"): + df.localCheckpoint(storageLevel=StorageLevel.DISK_ONLY) def test_dataframe_with_column_alias(spark): From 129d1a26500d3678960f4bc964f90fdf5dbc995a Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 22 May 2026 12:33:18 +0800 Subject: [PATCH 03/40] update --- Cargo.lock | 2 + crates/sail-common-datafusion/Cargo.toml | 2 + .../src/cached_relation.rs | 415 ++++++++++++++++-- .../src/rename/logical_plan.rs | 11 +- crates/sail-plan/src/lib.rs | 10 + crates/sail-plan/src/resolver/query/mod.rs | 3 +- .../src/service/plan_executor.rs | 153 +++---- .../pysail/data/compatibility/dataframe.json | 8 +- python/pysail/tests/spark/test_dataframe.py | 48 +- 9 files changed, 509 insertions(+), 143 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 5cb0b768a9..e321f8f26a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7165,6 +7165,8 @@ dependencies = [ "futures", "lazy_static", "lexical-core", + "object_store", + "parquet", "parquet-variant-compute", "parquet-variant-json", "pin-project-lite", diff --git a/crates/sail-common-datafusion/Cargo.toml b/crates/sail-common-datafusion/Cargo.toml index b6ccb1f2a7..f6fda9a7aa 100644 --- a/crates/sail-common-datafusion/Cargo.toml +++ b/crates/sail-common-datafusion/Cargo.toml @@ -26,6 +26,8 @@ tonic = { workspace = true } pin-project-lite = { workspace = true } parquet-variant-compute = { workspace = true } parquet-variant-json = { workspace = true } +object_store = { workspace = true } +parquet = { workspace = true } tokio = { workspace = true } serde_arrow = { workspace = true } diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index f0f1e4a32f..e37d915c56 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -3,12 +3,22 @@ use std::sync::{Arc, RwLock}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::datasource::listing::ListingTableUrl; use datafusion::datasource::{provider_as_source, MemTable}; -use datafusion_common::{internal_datafusion_err, Result, TableReference}; +use datafusion::execution::disk_manager::RefCountedTempFile; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::{ParquetReadOptions, SessionContext}; +use datafusion_common::{internal_datafusion_err, DataFusionError, Result, TableReference}; use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; +use futures::{StreamExt, TryStreamExt}; +use object_store::ObjectStoreExt; +use parquet::arrow::async_writer::{AsyncArrowWriter, ParquetObjectWriter}; use sail_common::spec; -use crate::extension::SessionExtension; +use crate::array::record_batch::{read_record_batches, write_record_batches}; +use crate::extension::{SessionExtension, SessionExtensionAccessor}; +use crate::rename::physical_plan::rename_physical_plan; +use crate::session::job::JobService; #[derive(Debug, Clone)] pub enum CachedRelationCleanup { @@ -17,63 +27,129 @@ pub enum CachedRelationCleanup { #[derive(Debug, Clone)] pub struct CachedRelation { - data: CachedRelationData, + data: Arc>, cleanup: Option, } #[derive(Debug, Clone)] enum CachedRelationData { LogicalPlan(Arc), - Memory { - schema: SchemaRef, - partitions: Arc>>, - storage_level: Option, - }, + Materialized(CachedRelationMaterialized), + Pending(CachedRelationPending), +} + +#[derive(Debug, Clone)] +struct CachedRelationMaterialized { + schema: SchemaRef, + memory_partitions: Option>>>, + serialized_memory_partitions: Option>>>, + disk_partitions: Option>>, + storage_level: Option, +} + +#[derive(Debug, Clone)] +struct CachedRelationDiskPartition { + files: Vec, +} + +#[derive(Debug, Clone)] +struct CachedRelationPending { + plan: Arc, + fields: Option>, + target: CachedRelationPendingTarget, +} + +#[derive(Debug, Clone)] +enum CachedRelationPendingTarget { + Local { storage_level: spec::StorageLevel }, + Reliable { path: String }, } impl CachedRelation { pub fn new(plan: Arc, cleanup: Option) -> Self { Self { - data: CachedRelationData::LogicalPlan(plan), + data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::LogicalPlan( + plan, + ))), cleanup, } } - pub fn new_memory( - schema: SchemaRef, - partitions: Vec>, - storage_level: Option, + pub async fn new_local_checkpoint( + ctx: &SessionContext, + plan: Arc, + storage_level: spec::StorageLevel, + ) -> Result { + let data = materialize_local_checkpoint(ctx, plan, storage_level).await?; + Ok(Self { + data: Arc::new(tokio::sync::Mutex::new(data)), + cleanup: None, + }) + } + + pub async fn new_reliable_checkpoint( + ctx: &SessionContext, + plan: Arc, + path: &str, + ) -> Result { + let data = materialize_reliable_checkpoint(ctx, plan, path).await?; + Ok(Self { + data: Arc::new(tokio::sync::Mutex::new(data)), + cleanup: None, + }) + } + + pub fn new_pending_local_checkpoint( + plan: Arc, + fields: Option>, + storage_level: spec::StorageLevel, ) -> Self { Self { - data: CachedRelationData::Memory { - schema, - partitions: Arc::new(partitions), - storage_level, - }, + data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( + CachedRelationPending { + plan, + fields, + target: CachedRelationPendingTarget::Local { storage_level }, + }, + ))), cleanup: None, } } - pub fn to_logical_plan(&self, relation_id: &str) -> Result { - match &self.data { + pub fn new_pending_reliable_checkpoint( + plan: Arc, + fields: Option>, + path: String, + ) -> Self { + Self { + data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( + CachedRelationPending { + plan, + fields, + target: CachedRelationPendingTarget::Reliable { path }, + }, + ))), + cleanup: None, + } + } + + pub async fn to_logical_plan( + &self, + ctx: &SessionContext, + relation_id: &str, + ) -> Result { + let mut data = self.data.lock().await; + if let CachedRelationData::Pending(pending) = &*data { + *data = pending.materialize(ctx).await?; + } + match &*data { CachedRelationData::LogicalPlan(plan) => Ok(plan.as_ref().clone()), - CachedRelationData::Memory { - schema, - partitions, - storage_level, - } => { - let _ = storage_level; - // TODO: Use Sail's full cache storage tiers once memory/disk/off-heap cache exists. - // Sail's physical-plan codec serializes MemorySourceConfig batches, so this cache - // representation can still be executed by local-cluster workers. - let table = MemTable::try_new(Arc::clone(schema), partitions.as_ref().clone())?; - LogicalPlanBuilder::scan( - TableReference::bare(format!("__sail_cached_relation_{relation_id}")), - provider_as_source(Arc::new(table)), - None, - )? - .build() + CachedRelationData::Materialized(materialized) => { + materialized.to_logical_plan(relation_id).await } + CachedRelationData::Pending(_) => Err(internal_datafusion_err!( + "cached relation materialization did not complete" + )), } } @@ -84,7 +160,6 @@ impl CachedRelation { #[derive(Debug, Default)] pub struct CachedRelationRegistry { - // TODO: Clean cached local checkpoint relations when a session is closed. relations: RwLock>, } @@ -124,3 +199,269 @@ impl SessionExtension for CachedRelationRegistry { "cached relation registry" } } + +impl CachedRelationPending { + async fn materialize(&self, ctx: &SessionContext) -> Result { + let physical_plan = + create_physical_plan(ctx, self.plan.as_ref().clone(), self.fields.as_ref()).await?; + match &self.target { + CachedRelationPendingTarget::Local { storage_level } => { + materialize_local_checkpoint(ctx, physical_plan, storage_level.clone()).await + } + CachedRelationPendingTarget::Reliable { path } => { + materialize_reliable_checkpoint(ctx, physical_plan, path).await + } + } + } +} + +impl CachedRelationMaterialized { + async fn try_new( + ctx: &SessionContext, + schema: SchemaRef, + partitions: Vec>, + storage_level: spec::StorageLevel, + ) -> Result { + let use_memory = storage_level.use_memory; + let use_serialized_memory = use_memory && !storage_level.deserialized; + let use_deserialized_memory = use_memory && storage_level.deserialized; + let use_disk = storage_level.use_disk; + if !use_deserialized_memory && !use_serialized_memory && !use_disk { + return Err(internal_datafusion_err!( + "local checkpoint storage level must use memory or disk" + )); + } + + let mut serialized_partitions = if use_serialized_memory { + Some(Vec::with_capacity(partitions.len())) + } else { + None + }; + let mut disk_partitions = if use_disk { + Some(Vec::with_capacity(partitions.len())) + } else { + None + }; + + for partition in &partitions { + let bytes = if use_serialized_memory || use_disk { + Some(write_record_batches(partition, schema.as_ref())?) + } else { + None + }; + + if let Some(serialized_partitions) = serialized_partitions.as_mut() { + let bytes = bytes + .as_ref() + .ok_or_else(|| internal_datafusion_err!("missing serialized partition"))?; + serialized_partitions.push(bytes.clone()); + } + + if let Some(disk_partitions) = disk_partitions.as_mut() { + let bytes = bytes + .as_ref() + .ok_or_else(|| internal_datafusion_err!("missing serialized partition"))?; + let replication = storage_level.replication.max(1); + let mut files = Vec::with_capacity(replication); + for _ in 0..replication { + files.push(write_disk_partition(ctx, bytes).await?); + } + disk_partitions.push(CachedRelationDiskPartition { files }); + } + } + + let memory_partitions = if use_deserialized_memory { + Some(Arc::new(partitions)) + } else { + None + }; + let serialized_memory_partitions = serialized_partitions.map(Arc::new); + let disk_partitions = disk_partitions.map(Arc::new); + + Ok(Self { + schema, + memory_partitions, + serialized_memory_partitions, + disk_partitions, + storage_level: Some(storage_level), + }) + } + + async fn to_logical_plan(&self, relation_id: &str) -> Result { + let _ = &self.storage_level; + let partitions = self.load_partitions().await?; + let table = MemTable::try_new(Arc::clone(&self.schema), partitions)?; + LogicalPlanBuilder::scan( + TableReference::bare(format!("__sail_cached_relation_{relation_id}")), + provider_as_source(Arc::new(table)), + None, + )? + .build() + } + + async fn load_partitions(&self) -> Result>> { + if let Some(partitions) = &self.memory_partitions { + return Ok(partitions.as_ref().clone()); + } + if let Some(partitions) = &self.serialized_memory_partitions { + return partitions + .iter() + .map(|bytes| read_record_batches(bytes)) + .collect(); + } + if let Some(partitions) = &self.disk_partitions { + let mut batches = Vec::with_capacity(partitions.len()); + for partition in partitions.iter() { + batches.push(partition.read().await?); + } + return Ok(batches); + } + Err(internal_datafusion_err!( + "cached relation has no materialized data" + )) + } +} + +impl CachedRelationDiskPartition { + async fn read(&self) -> Result> { + let mut last_error = None; + for file in &self.files { + match tokio::fs::read(file.path()) + .await + .map_err(DataFusionError::IoError) + .and_then(|bytes| read_record_batches(&bytes)) + { + Ok(batches) => return Ok(batches), + Err(error) => last_error = Some(error), + } + } + Err(last_error.unwrap_or_else(|| { + DataFusionError::Internal("cached relation disk partition has no files".to_string()) + })) + } +} + +async fn create_physical_plan( + ctx: &SessionContext, + plan: LogicalPlan, + fields: Option<&Vec>, +) -> Result> { + let df = ctx.execute_logical_plan(plan).await?; + let (session_state, plan) = df.into_parts(); + let plan = session_state.optimize(&plan)?; + let plan = session_state + .query_planner() + .create_physical_plan(&plan, &session_state) + .await?; + if let Some(fields) = fields { + rename_physical_plan(plan, fields) + } else { + Ok(plan) + } +} + +async fn materialize_local_checkpoint( + ctx: &SessionContext, + plan: Arc, + storage_level: spec::StorageLevel, +) -> Result { + let schema = plan.schema(); + let partitions = collect_checkpoint_partitions(ctx, plan).await?; + let materialized = + CachedRelationMaterialized::try_new(ctx, schema, partitions, storage_level).await?; + Ok(CachedRelationData::Materialized(materialized)) +} + +async fn materialize_reliable_checkpoint( + ctx: &SessionContext, + plan: Arc, + path: &str, +) -> Result { + let schema = plan.schema(); + if let Err(error) = write_reliable_checkpoint(ctx, plan, path).await { + let _ = cleanup_checkpoint_path(ctx, path).await; + return Err(error); + } + let read_options = ParquetReadOptions::new().schema(schema.as_ref()); + let df = ctx.read_parquet(path, read_options).await?; + let (_, read_plan) = df.into_parts(); + Ok(CachedRelationData::LogicalPlan(Arc::new(read_plan))) +} + +pub async fn cleanup_checkpoint_path(ctx: &SessionContext, path: &str) -> Result<()> { + let checkpoint_dir = format!("{}/", path.trim_end_matches('/')); + let checkpoint_url = ListingTableUrl::parse(&checkpoint_dir)?; + let store = ctx + .runtime_env() + .object_store(checkpoint_url.object_store())?; + let prefix = checkpoint_url.prefix().clone(); + let files = store + .list(Some(&prefix)) + .try_collect::>() + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + for file in files { + match store.delete(&file.location).await { + Ok(()) | Err(object_store::Error::NotFound { .. }) => {} + Err(error) => return Err(DataFusionError::External(Box::new(error))), + } + } + Ok(()) +} + +async fn collect_checkpoint_partitions( + ctx: &SessionContext, + plan: Arc, +) -> Result>> { + let service = ctx.extension::()?; + let mut stream = service.runner().execute(ctx, plan).await?; + let mut batches = vec![]; + while let Some(batch) = stream.next().await { + batches.push(batch?); + } + Ok(vec![batches]) +} + +async fn write_disk_partition(ctx: &SessionContext, bytes: &[u8]) -> Result { + let mut file = ctx + .runtime_env() + .disk_manager + .create_tmp_file("writing local checkpoint cache partition")?; + tokio::fs::write(file.path(), bytes) + .await + .map_err(DataFusionError::IoError)?; + file.update_disk_usage()?; + Ok(file) +} + +async fn write_reliable_checkpoint( + ctx: &SessionContext, + plan: Arc, + path: &str, +) -> Result<()> { + let schema = plan.schema(); + let checkpoint_dir = format!("{}/", path.trim_end_matches('/')); + let checkpoint_url = ListingTableUrl::parse(&checkpoint_dir)?; + let store = ctx + .runtime_env() + .object_store(checkpoint_url.object_store())?; + let file_path = checkpoint_url.prefix().clone().join("part-00000.parquet"); + let object_writer = ParquetObjectWriter::new(store, file_path); + let mut writer = AsyncArrowWriter::try_new(object_writer, schema.clone(), None) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + + let service = ctx.extension::()?; + let mut stream = service.runner().execute(ctx, plan).await?; + while let Some(batch) = stream.next().await { + writer + .write(&batch?) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + } + writer + .close() + .await + .map_err(|e| DataFusionError::External(Box::new(e)))?; + + Ok(()) +} diff --git a/crates/sail-common-datafusion/src/rename/logical_plan.rs b/crates/sail-common-datafusion/src/rename/logical_plan.rs index b1a7dd3b41..416133ce48 100644 --- a/crates/sail-common-datafusion/src/rename/logical_plan.rs +++ b/crates/sail-common-datafusion/src/rename/logical_plan.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use datafusion_common::exec_err; +use datafusion_common::metadata::FieldMetadata; use datafusion_expr::{Expr, LogicalPlan, Projection}; /// Wraps a logical plan in a projection that aliases each column to a new name. @@ -19,10 +20,16 @@ pub fn rename_logical_plan( .schema() .columns() .into_iter() + .zip(plan.schema().fields()) .zip(names.iter()) - .map(|(column, name)| { + .map(|((column, field), name)| { let relation = column.relation.clone(); - Expr::Column(column).alias_qualified(relation, name) + let metadata = if field.metadata().is_empty() { + None + } else { + Some(FieldMetadata::from(field.metadata().clone())) + }; + Expr::Column(column).alias_qualified_with_metadata(relation, name, metadata) }) .collect(); // The logical plan schema requires field names to be unique. diff --git a/crates/sail-plan/src/lib.rs b/crates/sail-plan/src/lib.rs index c7e4f51a9a..1ab197784c 100644 --- a/crates/sail-plan/src/lib.rs +++ b/crates/sail-plan/src/lib.rs @@ -64,3 +64,13 @@ pub async fn resolve_and_execute_plan( )); Ok((plan, info)) } + +pub async fn resolve_logical_plan( + ctx: &SessionContext, + config: Arc, + plan: spec::Plan, +) -> PlanResult { + PlanResolver::new(ctx, config) + .resolve_named_plan(plan) + .await +} diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index 21fd920578..77d13cc01b 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/crates/sail-plan/src/resolver/query/mod.rs @@ -235,9 +235,8 @@ impl PlanResolver<'_> { let relation = registry.get(&relation_id)?.ok_or_else(|| { PlanError::invalid(format!("No DataFrame with id {relation_id} is found")) })?; - let plan = relation.to_logical_plan(&relation_id)?; + let plan = relation.to_logical_plan(self.ctx, &relation_id).await?; let names = state.register_fields(plan.schema().inner().fields()); - // TODO: Preserve Spark attribute metadata/expr IDs for full checkpoint parity. rename_logical_plan(plan, &names)? } QueryNode::CommonInlineUserDefinedTableFunction(udtf) => { diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 95a5f488c8..14798f69cd 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use datafusion::arrow::compute::concat_batches; -use datafusion::prelude::{ParquetReadOptions, SessionContext}; +use datafusion::prelude::SessionContext; use fastrace::collector::SpanContext; use fastrace::future::FutureExt; use fastrace::Span; @@ -12,11 +12,12 @@ use futures::stream; use log::{debug, warn}; use sail_common::spec; use sail_common_datafusion::cached_relation::{ - CachedRelation, CachedRelationCleanup, CachedRelationRegistry, + cleanup_checkpoint_path as cleanup_cached_checkpoint_path, CachedRelation, + CachedRelationCleanup, CachedRelationRegistry, }; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; -use sail_plan::resolve_and_execute_plan; +use sail_plan::{resolve_and_execute_plan, resolve_logical_plan}; use tonic::codegen::tokio_stream::wrappers::ReceiverStream; use tonic::codegen::tokio_stream::Stream; use tonic::Status; @@ -534,12 +535,6 @@ pub(crate) async fn handle_execute_checkpoint_command( eager, storage_level, } = checkpoint; - if !eager { - // TODO: Implement Spark-compatible lazy checkpoint materialization on first action. - return Err(SparkError::unsupported( - "DataFrame checkpoint with eager=false is not supported", - )); - } let storage_level = if local { Some(resolve_local_checkpoint_storage_level(storage_level)?) } else if storage_level.is_some() { @@ -553,14 +548,46 @@ pub(crate) async fn handle_execute_checkpoint_command( let relation = relation.required("checkpoint relation")?; let plan: spec::Plan = relation.try_into()?; let relation_id = Uuid::new_v4().to_string(); - let (physical_plan, _) = resolve_and_execute_plan(ctx, spark.plan_config()?, plan).await?; - let (relation, cleanup_on_insert_error) = if local { + let (relation, cleanup_on_insert_error) = if eager { + let (physical_plan, _) = resolve_and_execute_plan(ctx, spark.plan_config()?, plan).await?; + if local { + ( + CachedRelation::new_local_checkpoint( + ctx, + physical_plan, + storage_level.required("local checkpoint storage level")?, + ) + .await?, + None, + ) + } else { + let path = checkpoint_path(&spark, &relation_id)?; + ( + CachedRelation::new_reliable_checkpoint(ctx, physical_plan, &path).await?, + Some(path), + ) + } + } else if local { + let named_plan = resolve_logical_plan(ctx, spark.plan_config()?, plan).await?; ( - materialize_local_checkpoint(ctx, physical_plan, storage_level).await?, + CachedRelation::new_pending_local_checkpoint( + Arc::new(named_plan.plan), + named_plan.fields, + storage_level.required("local checkpoint storage level")?, + ), None, ) } else { - materialize_reliable_checkpoint(ctx, &spark, &relation_id, physical_plan).await? + let path = checkpoint_path(&spark, &relation_id)?; + let named_plan = resolve_logical_plan(ctx, spark.plan_config()?, plan).await?; + ( + CachedRelation::new_pending_reliable_checkpoint( + Arc::new(named_plan.plan), + named_plan.fields, + path, + ), + None, + ) }; if let Err(error) = ctx .extension::()? @@ -568,7 +595,7 @@ pub(crate) async fn handle_execute_checkpoint_command( { let error = SparkError::from(error); if let Some(path) = cleanup_on_insert_error { - return Err(cleanup_checkpoint_after_error(&path, error).await); + return Err(cleanup_checkpoint_after_error(ctx, &path, error).await); } return Err(error); } @@ -588,50 +615,6 @@ pub(crate) async fn handle_execute_checkpoint_command( )) } -async fn materialize_local_checkpoint( - ctx: &SessionContext, - physical_plan: Arc, - storage_level: Option, -) -> SparkResult { - let schema = physical_plan.schema(); - let service = ctx.extension::()?; - let stream = service.runner().execute(ctx, physical_plan).await?; - let batches = read_stream(stream).await?; - Ok(CachedRelation::new_memory( - schema, - vec![batches], - storage_level, - )) -} - -async fn materialize_reliable_checkpoint( - ctx: &SessionContext, - spark: &SparkSession, - relation_id: &str, - physical_plan: Arc, -) -> SparkResult<(CachedRelation, Option)> { - let path = checkpoint_path(spark, relation_id)?; - let schema = physical_plan.schema(); - // TODO: Use Sail's file-write physical path through JobRunner for cluster-native checkpointing. - if let Err(error) = ctx.write_parquet(physical_plan, &path, None).await { - return Err(cleanup_checkpoint_after_error(&path, error.into()).await); - } - let register_result: SparkResult = async { - let read_options = ParquetReadOptions::new().schema(schema.as_ref()); - let df = ctx.read_parquet(&path, read_options).await?; - let (_, read_plan) = df.into_parts(); - Ok(CachedRelation::new(Arc::new(read_plan), None)) - } - .await; - match register_result { - Ok(relation) => { - // TODO: Add Spark-like reliable checkpoint cleanup/reference tracking. - Ok((relation, Some(path))) - } - Err(error) => Err(cleanup_checkpoint_after_error(&path, error).await), - } -} - pub(crate) async fn handle_execute_remove_cached_remote_relation_command( ctx: &SessionContext, command: RemoveCachedRemoteRelationCommand, @@ -645,7 +628,7 @@ pub(crate) async fn handle_execute_remove_cached_remote_relation_command( .extension::()? .remove(&relation.relation_id)? { - cleanup_cached_relation(relation).await?; + cleanup_cached_relation(ctx, relation).await?; } let output = if metadata.reattachable { vec![ExecutorOutput::complete()] @@ -662,35 +645,29 @@ pub(crate) async fn handle_execute_remove_cached_remote_relation_command( fn resolve_local_checkpoint_storage_level( storage_level: Option, ) -> SparkResult { - let storage_level: spec::StorageLevel = storage_level + let mut storage_level: spec::StorageLevel = storage_level .map(TryInto::try_into) .transpose()? .map(Ok) .unwrap_or_else(|| "MEMORY_AND_DISK".parse().map_err(SparkError::from))?; - if !storage_level.use_memory { - return Err(SparkError::unsupported( - "localCheckpoint StorageLevel without memory is not supported", - )); - } - if storage_level.use_off_heap { - return Err(SparkError::unsupported( - "localCheckpoint off-heap StorageLevel is not supported", - )); + if !storage_level.use_disk && !storage_level.use_memory && !storage_level.use_off_heap { + storage_level = "MEMORY_AND_DISK".parse().map_err(SparkError::from)?; + } else { + // Spark localCheckpoint transforms every user-provided StorageLevel to one that + // uses disk, so cached partitions remain recoverable after memory eviction. + storage_level.use_disk = true; + storage_level.use_off_heap = false; } - if storage_level.replication != 1 { - return Err(SparkError::unsupported( - "localCheckpoint replicated StorageLevel is not supported", + if storage_level.replication == 0 { + return Err(SparkError::invalid( + "localCheckpoint StorageLevel replication must be at least 1", )); } - - // TODO: Enforce disk spill and serialized/deserialized storage tiers once Sail has a - // DataFrame/block cache subsystem. For now we only support memory-backed local checkpoint. Ok(storage_level) } fn checkpoint_path(spark: &SparkSession, relation_id: &str) -> SparkResult { - // TODO: Wire Spark-compatible startup config / setCheckpointDir semantics. let root = spark .get_config_option(vec![SPARK_CHECKPOINT_DIR.to_string()])? .into_iter() @@ -736,8 +713,12 @@ fn validate_checkpoint_path_segment(name: &str, value: &str) -> SparkResult<()> } } -async fn cleanup_checkpoint_after_error(path: &str, error: SparkError) -> SparkError { - match cleanup_checkpoint_path(path).await { +async fn cleanup_checkpoint_after_error( + ctx: &SessionContext, + path: &str, + error: SparkError, +) -> SparkError { + match cleanup_checkpoint_path(ctx, path).await { Ok(()) => error, Err(cleanup_error) => SparkError::internal(format!( "checkpoint failed: {error}; additionally failed to clean checkpoint path {path}: {cleanup_error}" @@ -745,17 +726,17 @@ async fn cleanup_checkpoint_after_error(path: &str, error: SparkError) -> SparkE } } -async fn cleanup_checkpoint_path(path: &str) -> SparkResult<()> { - match tokio::fs::remove_dir_all(path).await { - Ok(()) => Ok(()), - Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), - Err(e) => Err(e.into()), - } +async fn cleanup_checkpoint_path(ctx: &SessionContext, path: &str) -> SparkResult<()> { + cleanup_cached_checkpoint_path(ctx, path).await?; + Ok(()) } -async fn cleanup_cached_relation(relation: CachedRelation) -> SparkResult<()> { +async fn cleanup_cached_relation( + ctx: &SessionContext, + relation: CachedRelation, +) -> SparkResult<()> { match relation.into_cleanup() { - Some(CachedRelationCleanup::LocalPath(path)) => cleanup_checkpoint_path(&path).await, + Some(CachedRelationCleanup::LocalPath(path)) => cleanup_checkpoint_path(ctx, &path).await, None => Ok(()), } } diff --git a/python/pysail/data/compatibility/dataframe.json b/python/pysail/data/compatibility/dataframe.json index 24e3016533..c1e6e708aa 100644 --- a/python/pysail/data/compatibility/dataframe.json +++ b/python/pysail/data/compatibility/dataframe.json @@ -23,8 +23,8 @@ { "module": "pyspark.sql.DataFrame", "function": "checkpoint", - "status": "partially supported", - "note": "Supports eager=True with spark.checkpoint.dir configured. eager=False is not supported yet.", + "status": "supported", + "note": "Supports eager=True and eager=False with spark.checkpoint.dir configured.", "ref": "https://github.com/lakehq/sail/issues/476" }, { @@ -254,8 +254,8 @@ { "module": "pyspark.sql.DataFrame", "function": "localCheckpoint", - "status": "partially supported", - "note": "Supports eager=True with memory-backed cache materialization. StorageLevel values without memory, off-heap storage, replication, eager=False, and full storage tier semantics are not supported yet.", + "status": "supported", + "note": "Supports eager=True and eager=False with cache-backed StorageLevel handling.", "ref": "https://github.com/lakehq/sail/issues/482" }, { diff --git a/python/pysail/tests/spark/test_dataframe.py b/python/pysail/tests/spark/test_dataframe.py index 8a7a4b66a8..584f9e85fd 100644 --- a/python/pysail/tests/spark/test_dataframe.py +++ b/python/pysail/tests/spark/test_dataframe.py @@ -142,15 +142,36 @@ def test_dataframe_checkpoint_rejects_parent_path_components(spark, tmp_path): spark.conf.unset("spark.checkpoint.dir") -@pytest.mark.skipif(is_jvm_spark(), reason="Sail MVP limitation") -def test_dataframe_checkpoint_lazy_is_not_supported(spark): +def test_dataframe_local_checkpoint_lazy(spark): df = spark.createDataFrame( - schema="id INT", - data=[(1,)], + schema="id INT, value STRING", + data=[(1, "a"), (2, "b")], + ) + + checkpointed = df.where(col("id") >= 1).localCheckpoint(eager=False) + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + + +@pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") +def test_dataframe_checkpoint_lazy(spark, tmp_path): + df = spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b")], ) + spark.conf.set("spark.checkpoint.dir", str(tmp_path / "checkpoints")) + try: + checkpointed = df.where(col("id") >= 1).checkpoint(eager=False) - with pytest.raises(Exception, match="eager=false"): - df.localCheckpoint(eager=False) + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + finally: + spark.conf.unset("spark.checkpoint.dir") def test_dataframe_local_checkpoint_with_memory_storage_level(spark): @@ -167,15 +188,18 @@ def test_dataframe_local_checkpoint_with_memory_storage_level(spark): ) -@pytest.mark.skipif(is_jvm_spark(), reason="Sail cache storage tier limitation") -def test_dataframe_local_checkpoint_disk_only_storage_level_is_not_supported(spark): +def test_dataframe_local_checkpoint_disk_only_storage_level(spark): df = spark.createDataFrame( - schema="id INT", - data=[(1,)], + schema="id INT, value STRING", + data=[(1, "a"), (2, "b")], ) - with pytest.raises(Exception, match="without memory"): - df.localCheckpoint(storageLevel=StorageLevel.DISK_ONLY) + checkpointed = df.localCheckpoint(storageLevel=StorageLevel.DISK_ONLY) + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) def test_dataframe_with_column_alias(spark): From 3812b4a0ae1595a3d18bc9c6331920731c0885ab Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 22 May 2026 12:45:03 +0800 Subject: [PATCH 04/40] update --- python/pysail/tests/spark/test_dataframe.py | 159 +++++++++++--------- 1 file changed, 86 insertions(+), 73 deletions(-) diff --git a/python/pysail/tests/spark/test_dataframe.py b/python/pysail/tests/spark/test_dataframe.py index 584f9e85fd..3bf817bed4 100644 --- a/python/pysail/tests/spark/test_dataframe.py +++ b/python/pysail/tests/spark/test_dataframe.py @@ -1,3 +1,5 @@ +import shutil + import pandas as pd import pytest from pandas.testing import assert_frame_equal @@ -7,6 +9,38 @@ from pysail.testing.spark.utils.common import is_jvm_spark +_CHECKPOINT_SOURCE_MAX_ID = 2 +_CHECKPOINT_SOURCE_ROW_COUNT = 2 + + +def _read_checkpoint_source(spark, path): + spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b"), (3, "c")], + ).write.mode("overwrite").parquet(str(path)) + return spark.read.parquet(str(path)).where(col("id") <= _CHECKPOINT_SOURCE_MAX_ID) + + +def _assert_checkpoint_source_result(df): + assert_frame_equal( + df.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + assert_frame_equal( + df.where(col("id") == _CHECKPOINT_SOURCE_MAX_ID).select("value").toPandas(), + pd.DataFrame({"value": ["b"]}), + ) + + +def _checkpoint_parquet_files(path): + if not path.exists(): + return [] + return list(path.rglob("*.parquet")) + + +def _uuid_values(df): + return [row.value for row in df.sort("id").collect()] + def test_dataframe_drop(spark): df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) @@ -75,43 +109,38 @@ def test_dataframe_drop(spark): ) -def test_dataframe_local_checkpoint(spark): - df = spark.createDataFrame( - schema="id INT, value STRING", - data=[(1, "a"), (2, "b")], - ) +def test_dataframe_local_checkpoint_survives_source_removal(spark, tmp_path): + source_path = tmp_path / "source" + df = _read_checkpoint_source(spark, source_path) - checkpointed = df.where(col("id") >= 1).localCheckpoint() + checkpointed = df.localCheckpoint() + shutil.rmtree(source_path) - assert_frame_equal( - checkpointed.sort("id").toPandas(), - pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), - ) - target_id = 2 - assert_frame_equal( - checkpointed.where(col("id") == target_id).select("value").toPandas(), - pd.DataFrame({"value": ["b"]}), - ) + _assert_checkpoint_source_result(checkpointed) + + +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +def test_dataframe_local_checkpoint_freezes_nondeterministic_values(spark, eager): + df = spark.sql("SELECT id, uuid() AS value FROM range(3)") + assert _uuid_values(df) != _uuid_values(df) + + checkpointed = df.localCheckpoint(eager=eager) + + assert _uuid_values(checkpointed) == _uuid_values(checkpointed) @pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") def test_dataframe_checkpoint(spark, tmp_path): - df = spark.createDataFrame( - schema="id INT, value STRING", - data=[(1, "a"), (2, "b")], - ) - spark.conf.set("spark.checkpoint.dir", str(tmp_path / "checkpoints")) + source_path = tmp_path / "source" + checkpoint_path = tmp_path / "checkpoints" + df = _read_checkpoint_source(spark, source_path) + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) try: - checkpointed = df.where(col("id") >= 1).checkpoint() + checkpointed = df.checkpoint() + assert _checkpoint_parquet_files(checkpoint_path) + shutil.rmtree(source_path) - assert_frame_equal( - checkpointed.sort("id").toPandas(), - pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), - ) - assert_frame_equal( - checkpointed.where(col("id") == 1).select("value").toPandas(), - pd.DataFrame({"value": ["a"]}), - ) + _assert_checkpoint_source_result(checkpointed) finally: spark.conf.unset("spark.checkpoint.dir") @@ -142,64 +171,48 @@ def test_dataframe_checkpoint_rejects_parent_path_components(spark, tmp_path): spark.conf.unset("spark.checkpoint.dir") -def test_dataframe_local_checkpoint_lazy(spark): - df = spark.createDataFrame( - schema="id INT, value STRING", - data=[(1, "a"), (2, "b")], - ) +def test_dataframe_local_checkpoint_lazy_survives_source_removal_after_first_action(spark, tmp_path): + source_path = tmp_path / "source" + df = _read_checkpoint_source(spark, source_path) - checkpointed = df.where(col("id") >= 1).localCheckpoint(eager=False) + checkpointed = df.localCheckpoint(eager=False) + assert checkpointed.count() == _CHECKPOINT_SOURCE_ROW_COUNT + shutil.rmtree(source_path) - assert_frame_equal( - checkpointed.sort("id").toPandas(), - pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), - ) + _assert_checkpoint_source_result(checkpointed) @pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") def test_dataframe_checkpoint_lazy(spark, tmp_path): - df = spark.createDataFrame( - schema="id INT, value STRING", - data=[(1, "a"), (2, "b")], - ) - spark.conf.set("spark.checkpoint.dir", str(tmp_path / "checkpoints")) + source_path = tmp_path / "source" + checkpoint_path = tmp_path / "checkpoints" + df = _read_checkpoint_source(spark, source_path) + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) try: - checkpointed = df.where(col("id") >= 1).checkpoint(eager=False) + checkpointed = df.checkpoint(eager=False) + assert not _checkpoint_parquet_files(checkpoint_path) + assert checkpointed.count() == _CHECKPOINT_SOURCE_ROW_COUNT + assert _checkpoint_parquet_files(checkpoint_path) + shutil.rmtree(source_path) - assert_frame_equal( - checkpointed.sort("id").toPandas(), - pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), - ) + _assert_checkpoint_source_result(checkpointed) finally: spark.conf.unset("spark.checkpoint.dir") -def test_dataframe_local_checkpoint_with_memory_storage_level(spark): - df = spark.createDataFrame( - schema="id INT, value STRING", - data=[(1, "a"), (2, "b")], - ) - - checkpointed = df.localCheckpoint(storageLevel=StorageLevel.MEMORY_ONLY) - - assert_frame_equal( - checkpointed.sort("id").toPandas(), - pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), - ) +@pytest.mark.parametrize( + ("storage_level", "source_name"), + [(StorageLevel.MEMORY_ONLY, "memory-only"), (StorageLevel.DISK_ONLY, "disk-only")], + ids=["memory-only", "disk-only"], +) +def test_dataframe_local_checkpoint_storage_level_survives_source_removal(spark, tmp_path, storage_level, source_name): + source_path = tmp_path / f"source-{source_name}" + df = _read_checkpoint_source(spark, source_path) + checkpointed = df.localCheckpoint(storageLevel=storage_level) + shutil.rmtree(source_path) -def test_dataframe_local_checkpoint_disk_only_storage_level(spark): - df = spark.createDataFrame( - schema="id INT, value STRING", - data=[(1, "a"), (2, "b")], - ) - - checkpointed = df.localCheckpoint(storageLevel=StorageLevel.DISK_ONLY) - - assert_frame_equal( - checkpointed.sort("id").toPandas(), - pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), - ) + _assert_checkpoint_source_result(checkpointed) def test_dataframe_with_column_alias(spark): From ccdf643fec51acba070033f98c5e946f42222a02 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Sun, 21 Jun 2026 20:59:49 +0800 Subject: [PATCH 05/40] refactor: update checkpoint handling to use Arrow-backed storage --- Cargo.lock | 1 - crates/sail-common-datafusion/Cargo.toml | 1 - .../src/cached_relation.rs | 89 +++++++++++++++---- .../src/service/plan_executor.rs | 4 +- .../pysail/data/compatibility/dataframe.json | 4 +- .../tests/spark/dataframe/test_dataframe.py | 10 +-- 6 files changed, 82 insertions(+), 27 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index d773ed122b..c362ba48c9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7578,7 +7578,6 @@ dependencies = [ "lazy_static", "lexical-core", "object_store", - "parquet", "parquet-variant", "parquet-variant-compute", "parquet-variant-json", diff --git a/crates/sail-common-datafusion/Cargo.toml b/crates/sail-common-datafusion/Cargo.toml index da0a0505e5..057f80a6ac 100644 --- a/crates/sail-common-datafusion/Cargo.toml +++ b/crates/sail-common-datafusion/Cargo.toml @@ -29,7 +29,6 @@ parquet-variant-compute = { workspace = true } parquet-variant = { workspace = true } parquet-variant-json = { workspace = true } object_store = { workspace = true } -parquet = { workspace = true } tokio = { workspace = true } serde_arrow = { workspace = true } diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index e37d915c56..bc69ffae15 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -6,13 +6,13 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::listing::ListingTableUrl; use datafusion::datasource::{provider_as_source, MemTable}; use datafusion::execution::disk_manager::RefCountedTempFile; +use datafusion::execution::options::ArrowReadOptions; use datafusion::physical_plan::ExecutionPlan; -use datafusion::prelude::{ParquetReadOptions, SessionContext}; +use datafusion::prelude::SessionContext; use datafusion_common::{internal_datafusion_err, DataFusionError, Result, TableReference}; use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; use futures::{StreamExt, TryStreamExt}; use object_store::ObjectStoreExt; -use parquet::arrow::async_writer::{AsyncArrowWriter, ParquetObjectWriter}; use sail_common::spec; use crate::array::record_batch::{read_record_batches, write_record_batches}; @@ -22,7 +22,7 @@ use crate::session::job::JobService; #[derive(Debug, Clone)] pub enum CachedRelationCleanup { - LocalPath(String), + ObjectStorePath(String), } #[derive(Debug, Clone)] @@ -95,7 +95,7 @@ impl CachedRelation { let data = materialize_reliable_checkpoint(ctx, plan, path).await?; Ok(Self { data: Arc::new(tokio::sync::Mutex::new(data)), - cleanup: None, + cleanup: Some(CachedRelationCleanup::ObjectStorePath(path.to_string())), }) } @@ -121,6 +121,7 @@ impl CachedRelation { fields: Option>, path: String, ) -> Self { + let cleanup_path = path.clone(); Self { data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( CachedRelationPending { @@ -129,7 +130,7 @@ impl CachedRelation { target: CachedRelationPendingTarget::Reliable { path }, }, ))), - cleanup: None, + cleanup: Some(CachedRelationCleanup::ObjectStorePath(cleanup_path)), } } @@ -382,8 +383,8 @@ async fn materialize_reliable_checkpoint( let _ = cleanup_checkpoint_path(ctx, path).await; return Err(error); } - let read_options = ParquetReadOptions::new().schema(schema.as_ref()); - let df = ctx.read_parquet(path, read_options).await?; + let read_options = ArrowReadOptions::default().schema(schema.as_ref()); + let df = ctx.read_arrow(path, read_options).await?; let (_, read_plan) = df.into_parts(); Ok(CachedRelationData::LogicalPlan(Arc::new(read_plan))) } @@ -445,23 +446,77 @@ async fn write_reliable_checkpoint( let store = ctx .runtime_env() .object_store(checkpoint_url.object_store())?; - let file_path = checkpoint_url.prefix().clone().join("part-00000.parquet"); - let object_writer = ParquetObjectWriter::new(store, file_path); - let mut writer = AsyncArrowWriter::try_new(object_writer, schema.clone(), None) - .map_err(|e| DataFusionError::External(Box::new(e)))?; let service = ctx.extension::()?; let mut stream = service.runner().execute(ctx, plan).await?; + let mut batches = vec![]; while let Some(batch) = stream.next().await { - writer - .write(&batch?) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; + batches.push(batch?); } - writer - .close() + let bytes = write_record_batches(&batches, schema.as_ref())?; + let file_path = checkpoint_url.prefix().clone().join("part-00000.arrow"); + store + .put(&file_path, bytes.into()) .await .map_err(|e| DataFusionError::External(Box::new(e)))?; Ok(()) } + +#[cfg(test)] +mod tests { + use std::time::{SystemTime, UNIX_EPOCH}; + + use super::*; + + #[test] + fn pending_reliable_checkpoint_tracks_cleanup_path() -> Result<()> { + let path = "memory://checkpoint-root/session/relation".to_string(); + let plan = LogicalPlanBuilder::empty(false).build()?; + let relation = + CachedRelation::new_pending_reliable_checkpoint(Arc::new(plan), None, path.clone()); + + match relation.into_cleanup() { + Some(CachedRelationCleanup::ObjectStorePath(actual)) => { + assert_eq!(actual, path); + } + other => panic!("unexpected cleanup value: {other:?}"), + } + Ok(()) + } + + #[tokio::test] + async fn cleanup_checkpoint_path_deletes_prefix_files() -> Result<()> { + let suffix = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|e| internal_datafusion_err!("{e}"))? + .as_nanos(); + let root = std::env::temp_dir().join(format!( + "sail-common-datafusion-checkpoint-cleanup-{suffix}" + )); + let checkpoint = root.join("checkpoint"); + let nested = checkpoint.join("nested"); + tokio::fs::create_dir_all(&nested) + .await + .map_err(DataFusionError::IoError)?; + tokio::fs::write(checkpoint.join("part-00000.arrow"), b"checkpoint") + .await + .map_err(DataFusionError::IoError)?; + tokio::fs::write(nested.join("part-00001.arrow"), b"checkpoint") + .await + .map_err(DataFusionError::IoError)?; + tokio::fs::write(root.join("outside.arrow"), b"outside") + .await + .map_err(DataFusionError::IoError)?; + + let ctx = SessionContext::new(); + cleanup_checkpoint_path(&ctx, checkpoint.to_string_lossy().as_ref()).await?; + + assert!(!checkpoint.join("part-00000.arrow").exists()); + assert!(!nested.join("part-00001.arrow").exists()); + assert!(root.join("outside.arrow").exists()); + + let _ = tokio::fs::remove_dir_all(root).await; + Ok(()) + } +} diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 14798f69cd..ab5447bb0a 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -736,7 +736,9 @@ async fn cleanup_cached_relation( relation: CachedRelation, ) -> SparkResult<()> { match relation.into_cleanup() { - Some(CachedRelationCleanup::LocalPath(path)) => cleanup_checkpoint_path(ctx, &path).await, + Some(CachedRelationCleanup::ObjectStorePath(path)) => { + cleanup_checkpoint_path(ctx, &path).await + } None => Ok(()), } } diff --git a/python/pysail/data/compatibility/dataframe.json b/python/pysail/data/compatibility/dataframe.json index 637027ebaf..bdb4c9b267 100644 --- a/python/pysail/data/compatibility/dataframe.json +++ b/python/pysail/data/compatibility/dataframe.json @@ -24,7 +24,7 @@ "module": "pyspark.sql.DataFrame", "function": "checkpoint", "status": "supported", - "note": "Supports eager=True and eager=False with spark.checkpoint.dir configured.", + "note": "Supports eager=True and eager=False with spark.checkpoint.dir configured, using Arrow-backed temporary storage.", "ref": "https://github.com/lakehq/sail/issues/476" }, { @@ -561,4 +561,4 @@ "function": "writeTo", "status": "supported" } -] \ No newline at end of file +] diff --git a/python/pysail/tests/spark/dataframe/test_dataframe.py b/python/pysail/tests/spark/dataframe/test_dataframe.py index 92267b8123..8ad194a08a 100644 --- a/python/pysail/tests/spark/dataframe/test_dataframe.py +++ b/python/pysail/tests/spark/dataframe/test_dataframe.py @@ -33,10 +33,10 @@ def _assert_checkpoint_source_result(df): ) -def _checkpoint_parquet_files(path): +def _checkpoint_arrow_files(path): if not path.exists(): return [] - return list(path.rglob("*.parquet")) + return list(path.rglob("*.arrow")) def _uuid_values(df): @@ -138,7 +138,7 @@ def test_dataframe_checkpoint(spark, tmp_path): spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) try: checkpointed = df.checkpoint() - assert _checkpoint_parquet_files(checkpoint_path) + assert _checkpoint_arrow_files(checkpoint_path) shutil.rmtree(source_path) _assert_checkpoint_source_result(checkpointed) @@ -191,9 +191,9 @@ def test_dataframe_checkpoint_lazy(spark, tmp_path): spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) try: checkpointed = df.checkpoint(eager=False) - assert not _checkpoint_parquet_files(checkpoint_path) + assert not _checkpoint_arrow_files(checkpoint_path) assert checkpointed.count() == _CHECKPOINT_SOURCE_ROW_COUNT - assert _checkpoint_parquet_files(checkpoint_path) + assert _checkpoint_arrow_files(checkpoint_path) shutil.rmtree(source_path) _assert_checkpoint_source_result(checkpointed) From 95cefd722efe0743bc4ba6c64c680da8065ed327 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Sun, 21 Jun 2026 21:24:15 +0800 Subject: [PATCH 06/40] feat: materialize checkpoints as arrow object-store plans --- Cargo.lock | 3 +- crates/sail-common-datafusion/Cargo.toml | 2 + .../src/cached_relation.rs | 237 +++++++++++++----- crates/sail-object-store/src/lib.rs | 2 + crates/sail-object-store/src/path.rs | 128 ++++++++++ crates/sail-session/src/planner.rs | 13 +- crates/sail-telemetry/Cargo.toml | 1 - .../src/execution/physical_plan.rs | 5 +- .../tests/spark/dataframe/test_dataframe.py | 13 + 9 files changed, 341 insertions(+), 63 deletions(-) create mode 100644 crates/sail-object-store/src/path.rs diff --git a/Cargo.lock b/Cargo.lock index c362ba48c9..f81dc92e1c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7572,6 +7572,7 @@ dependencies = [ "chrono", "datafusion", "datafusion-common", + "datafusion-datasource", "datafusion-expr", "either", "futures", @@ -7589,6 +7590,7 @@ dependencies = [ "regex", "ryu", "sail-common", + "sail-object-store", "serde", "serde_arrow", "serde_json", @@ -8214,7 +8216,6 @@ dependencies = [ "quote", "regex", "sail-common", - "sail-common-datafusion", "serde", "serde_yaml", "syn 2.0.118", diff --git a/crates/sail-common-datafusion/Cargo.toml b/crates/sail-common-datafusion/Cargo.toml index 057f80a6ac..4bf954def5 100644 --- a/crates/sail-common-datafusion/Cargo.toml +++ b/crates/sail-common-datafusion/Cargo.toml @@ -8,11 +8,13 @@ workspace = true [dependencies] sail-common = { path = "../sail-common" } +sail-object-store = { path = "../sail-object-store" } chrono = { workspace = true } ryu = { workspace = true } datafusion = { workspace = true } datafusion-common = { workspace = true } +datafusion-datasource = { workspace = true } datafusion-expr = { workspace = true } arrow-schema = { workspace = true } thiserror = { workspace = true } diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index bc69ffae15..4d1b3dbe61 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -1,19 +1,28 @@ +use std::cmp::Ordering; use std::collections::HashMap; +use std::fmt::Formatter; use std::sync::{Arc, RwLock}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; -use datafusion::datasource::listing::ListingTableUrl; -use datafusion::datasource::{provider_as_source, MemTable}; +use datafusion::datasource::physical_plan::ArrowSource; use datafusion::execution::disk_manager::RefCountedTempFile; -use datafusion::execution::options::ArrowReadOptions; +use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; -use datafusion_common::{internal_datafusion_err, DataFusionError, Result, TableReference}; -use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; -use futures::{StreamExt, TryStreamExt}; -use object_store::ObjectStoreExt; +use datafusion_common::{ + internal_datafusion_err, DFSchema, DFSchemaRef, DataFusionError, Result, Statistics, +}; +use datafusion_datasource::file_groups::FileGroup; +use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_datasource::memory::MemorySourceConfig; +use datafusion_datasource::source::DataSourceExec; +use datafusion_datasource::{PartitionedFile, TableSchema}; +use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use futures::StreamExt; +use object_store::ObjectMeta; use sail_common::spec; +use sail_object_store::{delete_object_store_prefix, resolve_object_store_path}; use crate::array::record_batch::{read_record_batches, write_record_batches}; use crate::extension::{SessionExtension, SessionExtensionAccessor}; @@ -25,6 +34,68 @@ pub enum CachedRelationCleanup { ObjectStorePath(String), } +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +pub struct CachedRelationNode { + relation_id: String, + schema: DFSchemaRef, +} + +impl CachedRelationNode { + fn try_new(relation_id: String, schema: SchemaRef) -> Result { + let schema = Arc::new(DFSchema::try_from(schema.as_ref().clone())?); + Ok(Self { + relation_id, + schema, + }) + } + + pub fn relation_id(&self) -> &str { + &self.relation_id + } +} + +impl PartialOrd for CachedRelationNode { + fn partial_cmp(&self, other: &Self) -> Option { + self.relation_id.partial_cmp(&other.relation_id) + } +} + +impl UserDefinedLogicalNodeCore for CachedRelationNode { + fn name(&self) -> &str { + "CachedRelation" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CachedRelation: relation_id={}", self.relation_id) + } + + fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { + if !exprs.is_empty() { + return Err(internal_datafusion_err!( + "CachedRelation does not support expressions" + )); + } + if !inputs.is_empty() { + return Err(internal_datafusion_err!( + "CachedRelation does not support inputs" + )); + } + Ok(self.clone()) + } +} + #[derive(Debug, Clone)] pub struct CachedRelation { data: Arc>, @@ -39,7 +110,16 @@ enum CachedRelationData { } #[derive(Debug, Clone)] -struct CachedRelationMaterialized { +enum CachedRelationMaterialized { + Local(CachedRelationLocalMaterialized), + Physical { + schema: SchemaRef, + plan: Arc, + }, +} + +#[derive(Debug, Clone)] +struct CachedRelationLocalMaterialized { schema: SchemaRef, memory_partitions: Option>>>, serialized_memory_partitions: Option>>>, @@ -146,7 +226,7 @@ impl CachedRelation { match &*data { CachedRelationData::LogicalPlan(plan) => Ok(plan.as_ref().clone()), CachedRelationData::Materialized(materialized) => { - materialized.to_logical_plan(relation_id).await + materialized.to_logical_plan(relation_id) } CachedRelationData::Pending(_) => Err(internal_datafusion_err!( "cached relation materialization did not complete" @@ -154,6 +234,19 @@ impl CachedRelation { } } + pub async fn to_physical_plan(&self) -> Result> { + let data = self.data.lock().await; + match &*data { + CachedRelationData::Materialized(materialized) => materialized.to_physical_plan().await, + CachedRelationData::LogicalPlan(_) => Err(internal_datafusion_err!( + "cached relation is not materialized as a physical plan" + )), + CachedRelationData::Pending(_) => Err(internal_datafusion_err!( + "cached relation materialization did not complete" + )), + } + } + pub fn into_cleanup(self) -> Option { self.cleanup } @@ -217,6 +310,31 @@ impl CachedRelationPending { } impl CachedRelationMaterialized { + fn schema(&self) -> SchemaRef { + match self { + Self::Local(local) => Arc::clone(&local.schema), + Self::Physical { schema, .. } => Arc::clone(schema), + } + } + + fn to_logical_plan(&self, relation_id: &str) -> Result { + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(CachedRelationNode::try_new( + relation_id.to_string(), + self.schema(), + )?), + })) + } + + async fn to_physical_plan(&self) -> Result> { + match self { + Self::Local(local) => local.to_physical_plan().await, + Self::Physical { plan, .. } => Ok(Arc::clone(plan)), + } + } +} + +impl CachedRelationLocalMaterialized { async fn try_new( ctx: &SessionContext, schema: SchemaRef, @@ -288,16 +406,12 @@ impl CachedRelationMaterialized { }) } - async fn to_logical_plan(&self, relation_id: &str) -> Result { + async fn to_physical_plan(&self) -> Result> { let _ = &self.storage_level; let partitions = self.load_partitions().await?; - let table = MemTable::try_new(Arc::clone(&self.schema), partitions)?; - LogicalPlanBuilder::scan( - TableReference::bare(format!("__sail_cached_relation_{relation_id}")), - provider_as_source(Arc::new(table)), - None, - )? - .build() + let plan: Arc = + MemorySourceConfig::try_new_exec(&partitions, Arc::clone(&self.schema), None)?; + Ok(plan) } async fn load_partitions(&self) -> Result>> { @@ -369,8 +483,10 @@ async fn materialize_local_checkpoint( let schema = plan.schema(); let partitions = collect_checkpoint_partitions(ctx, plan).await?; let materialized = - CachedRelationMaterialized::try_new(ctx, schema, partitions, storage_level).await?; - Ok(CachedRelationData::Materialized(materialized)) + CachedRelationLocalMaterialized::try_new(ctx, schema, partitions, storage_level).await?; + Ok(CachedRelationData::Materialized( + CachedRelationMaterialized::Local(materialized), + )) } async fn materialize_reliable_checkpoint( @@ -379,35 +495,40 @@ async fn materialize_reliable_checkpoint( path: &str, ) -> Result { let schema = plan.schema(); - if let Err(error) = write_reliable_checkpoint(ctx, plan, path).await { - let _ = cleanup_checkpoint_path(ctx, path).await; - return Err(error); - } - let read_options = ArrowReadOptions::default().schema(schema.as_ref()); - let df = ctx.read_arrow(path, read_options).await?; - let (_, read_plan) = df.into_parts(); - Ok(CachedRelationData::LogicalPlan(Arc::new(read_plan))) + let (object_store_url, object_meta) = match write_reliable_checkpoint(ctx, plan, path).await { + Ok(file) => file, + Err(error) => { + let _ = cleanup_checkpoint_path(ctx, path).await; + return Err(error); + } + }; + let physical_plan = create_arrow_checkpoint_scan(object_store_url, object_meta, &schema)?; + Ok(CachedRelationData::Materialized( + CachedRelationMaterialized::Physical { + schema, + plan: physical_plan, + }, + )) +} + +fn create_arrow_checkpoint_scan( + object_store_url: ObjectStoreUrl, + object_meta: ObjectMeta, + schema: &SchemaRef, +) -> Result> { + let source = ArrowSource::new_stream_file_source(TableSchema::new(Arc::clone(schema), vec![])); + let config = FileScanConfigBuilder::new(object_store_url, Arc::new(source)) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new_from_meta( + object_meta, + )])]) + .with_statistics(Statistics::new_unknown(schema)) + .build(); + Ok(DataSourceExec::from_data_source(config)) } pub async fn cleanup_checkpoint_path(ctx: &SessionContext, path: &str) -> Result<()> { - let checkpoint_dir = format!("{}/", path.trim_end_matches('/')); - let checkpoint_url = ListingTableUrl::parse(&checkpoint_dir)?; - let store = ctx - .runtime_env() - .object_store(checkpoint_url.object_store())?; - let prefix = checkpoint_url.prefix().clone(); - let files = store - .list(Some(&prefix)) - .try_collect::>() - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; - for file in files { - match store.delete(&file.location).await { - Ok(()) | Err(object_store::Error::NotFound { .. }) => {} - Err(error) => return Err(DataFusionError::External(Box::new(error))), - } - } - Ok(()) + let runtime_env = ctx.runtime_env(); + delete_object_store_prefix(runtime_env.as_ref(), path).await } async fn collect_checkpoint_partitions( @@ -439,13 +560,10 @@ async fn write_reliable_checkpoint( ctx: &SessionContext, plan: Arc, path: &str, -) -> Result<()> { +) -> Result<(ObjectStoreUrl, ObjectMeta)> { let schema = plan.schema(); - let checkpoint_dir = format!("{}/", path.trim_end_matches('/')); - let checkpoint_url = ListingTableUrl::parse(&checkpoint_dir)?; - let store = ctx - .runtime_env() - .object_store(checkpoint_url.object_store())?; + let runtime_env = ctx.runtime_env(); + let checkpoint_path = resolve_object_store_path(runtime_env.as_ref(), path)?; let service = ctx.extension::()?; let mut stream = service.runner().execute(ctx, plan).await?; @@ -454,19 +572,18 @@ async fn write_reliable_checkpoint( batches.push(batch?); } let bytes = write_record_batches(&batches, schema.as_ref())?; - let file_path = checkpoint_url.prefix().clone().join("part-00000.arrow"); - store - .put(&file_path, bytes.into()) - .await - .map_err(|e| DataFusionError::External(Box::new(e)))?; + let file_path = checkpoint_path.child("part-00000.arrow"); + let object_meta = checkpoint_path.put_bytes(&file_path, bytes).await?; - Ok(()) + Ok((checkpoint_path.object_store_url().clone(), object_meta)) } #[cfg(test)] mod tests { use std::time::{SystemTime, UNIX_EPOCH}; + use datafusion_expr::LogicalPlanBuilder; + use super::*; #[test] @@ -480,7 +597,11 @@ mod tests { Some(CachedRelationCleanup::ObjectStorePath(actual)) => { assert_eq!(actual, path); } - other => panic!("unexpected cleanup value: {other:?}"), + other => { + return Err(internal_datafusion_err!( + "unexpected cleanup value: {other:?}" + )); + } } Ok(()) } diff --git a/crates/sail-object-store/src/lib.rs b/crates/sail-object-store/src/lib.rs index c9bc0541ae..31c2ccf367 100644 --- a/crates/sail-object-store/src/lib.rs +++ b/crates/sail-object-store/src/lib.rs @@ -1,6 +1,8 @@ mod hugging_face; mod layers; +mod path; mod registry; mod s3; +pub use path::{delete_object_store_prefix, resolve_object_store_path, ResolvedObjectStorePath}; pub use registry::DynamicObjectStoreRegistry; diff --git a/crates/sail-object-store/src/path.rs b/crates/sail-object-store/src/path.rs new file mode 100644 index 0000000000..5552e72d32 --- /dev/null +++ b/crates/sail-object-store/src/path.rs @@ -0,0 +1,128 @@ +use std::sync::Arc; + +use datafusion::datasource::listing::ListingTableUrl; +use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::execution::runtime_env::RuntimeEnv; +use datafusion_common::{DataFusionError, Result}; +use futures::TryStreamExt; +use object_store::path::Path; +use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt, PutPayload}; + +#[derive(Clone)] +pub struct ResolvedObjectStorePath { + object_store_url: ObjectStoreUrl, + prefix: Path, + store: Arc, +} + +impl ResolvedObjectStorePath { + pub fn object_store_url(&self) -> &ObjectStoreUrl { + &self.object_store_url + } + + pub fn prefix(&self) -> &Path { + &self.prefix + } + + pub fn store(&self) -> &Arc { + &self.store + } + + pub fn child(&self, name: &str) -> Path { + self.prefix.clone().join(name) + } + + pub async fn put_bytes(&self, location: &Path, bytes: Vec) -> Result { + self.store + .put(location, PutPayload::from(bytes)) + .await + .map_err(|e| DataFusionError::ObjectStore(Box::new(e)))?; + self.store + .head(location) + .await + .map_err(|e| DataFusionError::ObjectStore(Box::new(e))) + } + + pub async fn delete_prefix(&self) -> Result<()> { + let files = self + .store + .list(Some(&self.prefix)) + .try_collect::>() + .await + .map_err(|e| DataFusionError::ObjectStore(Box::new(e)))?; + for file in files { + match self.store.delete(&file.location).await { + Ok(()) | Err(object_store::Error::NotFound { .. }) => {} + Err(error) => return Err(DataFusionError::ObjectStore(Box::new(error))), + } + } + Ok(()) + } +} + +pub fn resolve_object_store_path( + runtime_env: &RuntimeEnv, + path: &str, +) -> Result { + let directory = format!("{}/", path.trim_end_matches('/')); + let url = ListingTableUrl::parse(&directory)?; + let object_store_url = url.object_store(); + let store = runtime_env.object_store(&object_store_url)?; + Ok(ResolvedObjectStorePath { + object_store_url, + prefix: url.prefix().clone(), + store, + }) +} + +pub async fn delete_object_store_prefix(runtime_env: &RuntimeEnv, path: &str) -> Result<()> { + resolve_object_store_path(runtime_env, path)? + .delete_prefix() + .await +} + +#[cfg(test)] +mod tests { + use std::time::{SystemTime, UNIX_EPOCH}; + + use datafusion::execution::runtime_env::RuntimeEnv; + use object_store::memory::InMemory; + use url::Url; + + use super::*; + + #[tokio::test] + async fn resolved_path_puts_and_deletes_prefix() -> Result<()> { + let runtime_env = RuntimeEnv::default(); + runtime_env.register_object_store( + &Url::parse("memory:///").map_err(|e| DataFusionError::External(Box::new(e)))?, + Arc::new(InMemory::new()), + ); + let suffix = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|e| DataFusionError::External(Box::new(e)))? + .as_nanos(); + let path = format!("memory:///{suffix}/checkpoint"); + let resolved = resolve_object_store_path(&runtime_env, &path)?; + let first = resolved.child("part-00000.arrow"); + let nested = resolved.child("nested/part-00001.arrow"); + let outside = Path::from(format!("{suffix}/outside.arrow")); + + resolved.put_bytes(&first, b"first".to_vec()).await?; + resolved.put_bytes(&nested, b"nested".to_vec()).await?; + resolved.put_bytes(&outside, b"outside".to_vec()).await?; + + resolved.delete_prefix().await?; + + assert!(matches!( + resolved.store().head(&first).await, + Err(object_store::Error::NotFound { .. }) + )); + assert!(matches!( + resolved.store().head(&nested).await, + Err(object_store::Error::NotFound { .. }) + )); + assert!(resolved.store().head(&outside).await.is_ok()); + Ok(()) + } +} diff --git a/crates/sail-session/src/planner.rs b/crates/sail-session/src/planner.rs index 02cd978a43..ea71aa0cc5 100644 --- a/crates/sail-session/src/planner.rs +++ b/crates/sail-session/src/planner.rs @@ -11,6 +11,8 @@ use datafusion_common::{internal_err, DFSchema}; use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use datafusion_physical_expr::{create_physical_sort_exprs, Partitioning}; use sail_catalog_system::planner::SystemTablePhysicalPlanner; +use sail_common_datafusion::cached_relation::{CachedRelationNode, CachedRelationRegistry}; +use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::logical_rewriter::LogicalRewriter; use sail_common_datafusion::rename::physical_plan::rename_projected_physical_plan; use sail_common_datafusion::streaming::event::schema::{ @@ -98,8 +100,17 @@ impl ExtensionPlanner for ExtensionPhysicalPlanner { session_state: &SessionState, ) -> datafusion_common::Result>> { let plan: Arc = if let Some(node) = - node.as_any().downcast_ref::() + node.as_any().downcast_ref::() { + let registry = session_state.extension::()?; + let relation = registry.get(node.relation_id())?.ok_or_else(|| { + datafusion_common::DataFusionError::Internal(format!( + "No DataFrame with id {} is found", + node.relation_id() + )) + })?; + relation.to_physical_plan().await? + } else if let Some(node) = node.as_any().downcast_ref::() { let schema = UserDefinedLogicalNode::schema(node).inner().clone(); let projection = (0..schema.fields().len()).collect(); Arc::new(RangeExec::try_new( diff --git a/crates/sail-telemetry/Cargo.toml b/crates/sail-telemetry/Cargo.toml index b4859dd382..7805de9c78 100644 --- a/crates/sail-telemetry/Cargo.toml +++ b/crates/sail-telemetry/Cargo.toml @@ -8,7 +8,6 @@ workspace = true [dependencies] sail-common = { path = "../sail-common" } -sail-common-datafusion = { path = "../sail-common-datafusion" } thiserror = { workspace = true } env_logger = { workspace = true } diff --git a/crates/sail-telemetry/src/execution/physical_plan.rs b/crates/sail-telemetry/src/execution/physical_plan.rs index 6586790d6b..4d0ff127b3 100644 --- a/crates/sail-telemetry/src/execution/physical_plan.rs +++ b/crates/sail-telemetry/src/execution/physical_plan.rs @@ -29,7 +29,6 @@ use fastrace::Span; use fastrace_futures::StreamExt; use futures::Stream; use pin_project_lite::pin_project; -use sail_common_datafusion::utils::items::ItemTaker; use crate::common::{KeyValue, SpanAttribute}; use crate::execution::metrics::MetricEmitter; @@ -149,7 +148,9 @@ impl ExecutionPlan for TracingExec { self: Arc, children: Vec>, ) -> Result> { - let child = children.one()?; + let [child] = children.try_into().map_err(|_| { + datafusion::common::internal_datafusion_err!("TracingExec requires exactly one child") + })?; Ok(Arc::new(TracingExec::new(child, self.options.clone()))) } diff --git a/python/pysail/tests/spark/dataframe/test_dataframe.py b/python/pysail/tests/spark/dataframe/test_dataframe.py index 8ad194a08a..f25120506a 100644 --- a/python/pysail/tests/spark/dataframe/test_dataframe.py +++ b/python/pysail/tests/spark/dataframe/test_dataframe.py @@ -8,6 +8,7 @@ from pyspark.sql.functions import col, lit, row_number from pyspark.sql.window import Window +from pysail.testing.spark.steps.plan import normalize_plan_text from pysail.testing.spark.utils.common import is_jvm_spark _CHECKPOINT_SOURCE_MAX_ID = 2 @@ -130,6 +131,18 @@ def test_dataframe_local_checkpoint_freezes_nondeterministic_values(spark, eager assert _uuid_values(checkpointed) == _uuid_values(checkpointed) +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names") +def test_dataframe_local_checkpoint_explain_truncates_plan(spark): + df = spark.range(0, 3).withColumn("value", lit(1)).filter("id >= 0") + original = normalize_plan_text(df._explain_string()) # noqa: SLF001 + checkpointed = df.localCheckpoint() + checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + + assert "RangeExec" in original + assert "RangeExec" not in checkpointed_plan + assert "DataSourceExec" in checkpointed_plan + + @pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") def test_dataframe_checkpoint(spark, tmp_path): source_path = tmp_path / "source" From 45544d70c4aa92e194159b90947b3d2c203feb8d Mon Sep 17 00:00:00 2001 From: XL Liang Date: Sun, 21 Jun 2026 22:05:08 +0800 Subject: [PATCH 07/40] refactor: route checkpoint storage through session object store --- Cargo.lock | 1 - crates/sail-common-datafusion/Cargo.toml | 1 - .../src/cached_relation.rs | 134 +++++++++--------- .../src/session/checkpoint.rs | 80 +++++++++++ .../sail-common-datafusion/src/session/mod.rs | 1 + crates/sail-session/src/checkpoint.rs | 102 +++++++++++++ crates/sail-session/src/lib.rs | 1 + .../src/session_factory/server.rs | 5 + .../tests/spark/dataframe/test_dataframe.py | 31 ++++ 9 files changed, 290 insertions(+), 66 deletions(-) create mode 100644 crates/sail-common-datafusion/src/session/checkpoint.rs create mode 100644 crates/sail-session/src/checkpoint.rs diff --git a/Cargo.lock b/Cargo.lock index f81dc92e1c..bb2f22e296 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7590,7 +7590,6 @@ dependencies = [ "regex", "ryu", "sail-common", - "sail-object-store", "serde", "serde_arrow", "serde_json", diff --git a/crates/sail-common-datafusion/Cargo.toml b/crates/sail-common-datafusion/Cargo.toml index 4bf954def5..072691e108 100644 --- a/crates/sail-common-datafusion/Cargo.toml +++ b/crates/sail-common-datafusion/Cargo.toml @@ -8,7 +8,6 @@ workspace = true [dependencies] sail-common = { path = "../sail-common" } -sail-object-store = { path = "../sail-object-store" } chrono = { workspace = true } ryu = { workspace = true } diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index 4d1b3dbe61..78d32a883c 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -22,11 +22,11 @@ use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; use futures::StreamExt; use object_store::ObjectMeta; use sail_common::spec; -use sail_object_store::{delete_object_store_prefix, resolve_object_store_path}; use crate::array::record_batch::{read_record_batches, write_record_batches}; use crate::extension::{SessionExtension, SessionExtensionAccessor}; use crate::rename::physical_plan::rename_physical_plan; +use crate::session::checkpoint::CheckpointStoreService; use crate::session::job::JobService; #[derive(Debug, Clone)] @@ -495,14 +495,22 @@ async fn materialize_reliable_checkpoint( path: &str, ) -> Result { let schema = plan.schema(); - let (object_store_url, object_meta) = match write_reliable_checkpoint(ctx, plan, path).await { - Ok(file) => file, + let service = ctx.extension::()?; + let checkpoint = match service + .write_reliable_checkpoint(ctx, plan, path, Arc::clone(&schema)) + .await + { + Ok(checkpoint) => checkpoint, Err(error) => { let _ = cleanup_checkpoint_path(ctx, path).await; return Err(error); } }; - let physical_plan = create_arrow_checkpoint_scan(object_store_url, object_meta, &schema)?; + let physical_plan = create_arrow_checkpoint_scan( + checkpoint.object_store_url().clone(), + checkpoint.object_meta().to_vec(), + &schema, + )?; Ok(CachedRelationData::Materialized( CachedRelationMaterialized::Physical { schema, @@ -513,22 +521,29 @@ async fn materialize_reliable_checkpoint( fn create_arrow_checkpoint_scan( object_store_url: ObjectStoreUrl, - object_meta: ObjectMeta, + object_meta: Vec, schema: &SchemaRef, ) -> Result> { + if object_meta.is_empty() { + return Err(internal_datafusion_err!( + "reliable checkpoint did not produce any files" + )); + } let source = ArrowSource::new_stream_file_source(TableSchema::new(Arc::clone(schema), vec![])); + let file_groups = object_meta + .into_iter() + .map(|object_meta| FileGroup::new(vec![PartitionedFile::new_from_meta(object_meta)])) + .collect(); let config = FileScanConfigBuilder::new(object_store_url, Arc::new(source)) - .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new_from_meta( - object_meta, - )])]) + .with_file_groups(file_groups) .with_statistics(Statistics::new_unknown(schema)) .build(); Ok(DataSourceExec::from_data_source(config)) } pub async fn cleanup_checkpoint_path(ctx: &SessionContext, path: &str) -> Result<()> { - let runtime_env = ctx.runtime_env(); - delete_object_store_prefix(runtime_env.as_ref(), path).await + let service = ctx.extension::()?; + service.cleanup_checkpoint_path(ctx, path).await } async fn collect_checkpoint_partitions( @@ -556,35 +571,15 @@ async fn write_disk_partition(ctx: &SessionContext, bytes: &[u8]) -> Result, - path: &str, -) -> Result<(ObjectStoreUrl, ObjectMeta)> { - let schema = plan.schema(); - let runtime_env = ctx.runtime_env(); - let checkpoint_path = resolve_object_store_path(runtime_env.as_ref(), path)?; - - let service = ctx.extension::()?; - let mut stream = service.runner().execute(ctx, plan).await?; - let mut batches = vec![]; - while let Some(batch) = stream.next().await { - batches.push(batch?); - } - let bytes = write_record_batches(&batches, schema.as_ref())?; - let file_path = checkpoint_path.child("part-00000.arrow"); - let object_meta = checkpoint_path.put_bytes(&file_path, bytes).await?; - - Ok((checkpoint_path.object_store_url().clone(), object_meta)) -} - #[cfg(test)] mod tests { - use std::time::{SystemTime, UNIX_EPOCH}; + use std::sync::Mutex; + use async_trait::async_trait; use datafusion_expr::LogicalPlanBuilder; use super::*; + use crate::session::checkpoint::{CheckpointStore, ReliableCheckpoint}; #[test] fn pending_reliable_checkpoint_tracks_cleanup_path() -> Result<()> { @@ -606,38 +601,49 @@ mod tests { Ok(()) } + #[derive(Default)] + struct TestCheckpointStore { + cleanup_paths: Mutex>, + } + + #[async_trait] + impl CheckpointStore for TestCheckpointStore { + async fn write_reliable_checkpoint( + &self, + _ctx: &SessionContext, + _plan: Arc, + _path: &str, + _schema: SchemaRef, + ) -> Result { + Err(internal_datafusion_err!( + "test store does not write checkpoints" + )) + } + + async fn cleanup_checkpoint_path(&self, _ctx: &SessionContext, path: &str) -> Result<()> { + self.cleanup_paths + .lock() + .map_err(|e| internal_datafusion_err!("{e}"))? + .push(path.to_string()); + Ok(()) + } + } + #[tokio::test] - async fn cleanup_checkpoint_path_deletes_prefix_files() -> Result<()> { - let suffix = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_err(|e| internal_datafusion_err!("{e}"))? - .as_nanos(); - let root = std::env::temp_dir().join(format!( - "sail-common-datafusion-checkpoint-cleanup-{suffix}" - )); - let checkpoint = root.join("checkpoint"); - let nested = checkpoint.join("nested"); - tokio::fs::create_dir_all(&nested) - .await - .map_err(DataFusionError::IoError)?; - tokio::fs::write(checkpoint.join("part-00000.arrow"), b"checkpoint") - .await - .map_err(DataFusionError::IoError)?; - tokio::fs::write(nested.join("part-00001.arrow"), b"checkpoint") - .await - .map_err(DataFusionError::IoError)?; - tokio::fs::write(root.join("outside.arrow"), b"outside") - .await - .map_err(DataFusionError::IoError)?; - - let ctx = SessionContext::new(); - cleanup_checkpoint_path(&ctx, checkpoint.to_string_lossy().as_ref()).await?; - - assert!(!checkpoint.join("part-00000.arrow").exists()); - assert!(!nested.join("part-00001.arrow").exists()); - assert!(root.join("outside.arrow").exists()); - - let _ = tokio::fs::remove_dir_all(root).await; + async fn cleanup_checkpoint_path_delegates_to_store() -> Result<()> { + let store = Arc::new(TestCheckpointStore::default()); + let config = datafusion::prelude::SessionConfig::new() + .with_extension(Arc::new(CheckpointStoreService::new(store.clone()))); + let ctx = SessionContext::new_with_config(config); + let path = "memory:///checkpoint-root/session/relation"; + + cleanup_checkpoint_path(&ctx, path).await?; + + let paths = store + .cleanup_paths + .lock() + .map_err(|e| internal_datafusion_err!("{e}"))?; + assert_eq!(paths.as_slice(), &[path.to_string()]); Ok(()) } } diff --git a/crates/sail-common-datafusion/src/session/checkpoint.rs b/crates/sail-common-datafusion/src/session/checkpoint.rs new file mode 100644 index 0000000000..0fbec8b010 --- /dev/null +++ b/crates/sail-common-datafusion/src/session/checkpoint.rs @@ -0,0 +1,80 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use datafusion_common::Result; +use object_store::ObjectMeta; + +use crate::extension::SessionExtension; + +#[derive(Debug, Clone)] +pub struct ReliableCheckpoint { + object_store_url: ObjectStoreUrl, + object_meta: Vec, +} + +impl ReliableCheckpoint { + pub fn new(object_store_url: ObjectStoreUrl, object_meta: Vec) -> Self { + Self { + object_store_url, + object_meta, + } + } + + pub fn object_store_url(&self) -> &ObjectStoreUrl { + &self.object_store_url + } + + pub fn object_meta(&self) -> &[ObjectMeta] { + &self.object_meta + } +} + +#[async_trait] +pub trait CheckpointStore: Send + Sync + 'static { + async fn write_reliable_checkpoint( + &self, + ctx: &SessionContext, + plan: Arc, + path: &str, + schema: SchemaRef, + ) -> Result; + + async fn cleanup_checkpoint_path(&self, ctx: &SessionContext, path: &str) -> Result<()>; +} + +#[derive(Clone)] +pub struct CheckpointStoreService { + store: Arc, +} + +impl CheckpointStoreService { + pub fn new(store: Arc) -> Self { + Self { store } + } + + pub async fn write_reliable_checkpoint( + &self, + ctx: &SessionContext, + plan: Arc, + path: &str, + schema: SchemaRef, + ) -> Result { + self.store + .write_reliable_checkpoint(ctx, plan, path, schema) + .await + } + + pub async fn cleanup_checkpoint_path(&self, ctx: &SessionContext, path: &str) -> Result<()> { + self.store.cleanup_checkpoint_path(ctx, path).await + } +} + +impl SessionExtension for CheckpointStoreService { + fn name() -> &'static str { + "CheckpointStoreService" + } +} diff --git a/crates/sail-common-datafusion/src/session/mod.rs b/crates/sail-common-datafusion/src/session/mod.rs index 2bde515e56..ee8a2b3ed6 100644 --- a/crates/sail-common-datafusion/src/session/mod.rs +++ b/crates/sail-common-datafusion/src/session/mod.rs @@ -1,4 +1,5 @@ pub mod activity; +pub mod checkpoint; pub mod job; pub mod plan; pub mod repartition; diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs new file mode 100644 index 0000000000..f153a9f4ac --- /dev/null +++ b/crates/sail-session/src/checkpoint.rs @@ -0,0 +1,102 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::SessionContext; +use datafusion_common::Result; +use futures::StreamExt; +use sail_common_datafusion::array::record_batch::write_record_batches; +use sail_common_datafusion::extension::SessionExtensionAccessor; +use sail_common_datafusion::session::checkpoint::{CheckpointStore, ReliableCheckpoint}; +use sail_common_datafusion::session::job::JobService; +use sail_object_store::{delete_object_store_prefix, resolve_object_store_path}; + +#[derive(Debug, Default)] +pub struct ObjectStoreCheckpointStore; + +#[async_trait] +impl CheckpointStore for ObjectStoreCheckpointStore { + async fn write_reliable_checkpoint( + &self, + ctx: &SessionContext, + plan: Arc, + path: &str, + schema: SchemaRef, + ) -> Result { + let runtime_env = ctx.runtime_env(); + let checkpoint_path = resolve_object_store_path(runtime_env.as_ref(), path)?; + let service = ctx.extension::()?; + let mut stream = service.runner().execute(ctx, plan).await?; + let mut files = vec![]; + let mut file_index = 0usize; + + while let Some(batch) = stream.next().await { + let batch = batch?; + let bytes = write_record_batches(&[batch], schema.as_ref())?; + let file_path = checkpoint_path.child(&format!("part-{file_index:05}.arrow")); + files.push(checkpoint_path.put_bytes(&file_path, bytes).await?); + file_index += 1; + } + + if files.is_empty() { + let bytes = write_record_batches(&[], schema.as_ref())?; + let file_path = checkpoint_path.child("part-00000.arrow"); + files.push(checkpoint_path.put_bytes(&file_path, bytes).await?); + } + + Ok(ReliableCheckpoint::new( + checkpoint_path.object_store_url().clone(), + files, + )) + } + + async fn cleanup_checkpoint_path(&self, ctx: &SessionContext, path: &str) -> Result<()> { + let runtime_env = ctx.runtime_env(); + delete_object_store_prefix(runtime_env.as_ref(), path).await + } +} + +#[cfg(test)] +mod tests { + use std::time::{SystemTime, UNIX_EPOCH}; + + use datafusion_common::{internal_datafusion_err, DataFusionError}; + + use super::*; + + #[tokio::test] + async fn cleanup_checkpoint_path_deletes_prefix_files() -> Result<()> { + let suffix = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|e| internal_datafusion_err!("{e}"))? + .as_nanos(); + let root = std::env::temp_dir().join(format!("sail-session-checkpoint-cleanup-{suffix}")); + let checkpoint = root.join("checkpoint"); + let nested = checkpoint.join("nested"); + tokio::fs::create_dir_all(&nested) + .await + .map_err(DataFusionError::IoError)?; + tokio::fs::write(checkpoint.join("part-00000.arrow"), b"checkpoint") + .await + .map_err(DataFusionError::IoError)?; + tokio::fs::write(nested.join("part-00001.arrow"), b"checkpoint") + .await + .map_err(DataFusionError::IoError)?; + tokio::fs::write(root.join("outside.arrow"), b"outside") + .await + .map_err(DataFusionError::IoError)?; + + let ctx = SessionContext::new(); + ObjectStoreCheckpointStore + .cleanup_checkpoint_path(&ctx, checkpoint.to_string_lossy().as_ref()) + .await?; + + assert!(!checkpoint.join("part-00000.arrow").exists()); + assert!(!nested.join("part-00001.arrow").exists()); + assert!(root.join("outside.arrow").exists()); + + let _ = tokio::fs::remove_dir_all(root).await; + Ok(()) + } +} diff --git a/crates/sail-session/src/lib.rs b/crates/sail-session/src/lib.rs index 7af87b9130..28325626dc 100644 --- a/crates/sail-session/src/lib.rs +++ b/crates/sail-session/src/lib.rs @@ -1,4 +1,5 @@ pub mod catalog; +pub mod checkpoint; pub mod error; pub mod formats; pub mod observable; diff --git a/crates/sail-session/src/session_factory/server.rs b/crates/sail-session/src/session_factory/server.rs index a40b9708e5..e496b29c6d 100644 --- a/crates/sail-session/src/session_factory/server.rs +++ b/crates/sail-session/src/session_factory/server.rs @@ -15,6 +15,7 @@ use sail_common::config::{AppConfig, ExecutionMode}; use sail_common::runtime::RuntimeHandle; use sail_common_datafusion::cached_relation::CachedRelationRegistry; use sail_common_datafusion::session::activity::ActivityTracker; +use sail_common_datafusion::session::checkpoint::CheckpointStoreService; use sail_common_datafusion::session::job::{JobRunner, JobService}; use sail_common_datafusion::session::repartition::RepartitionBufferConfig; use sail_delta_lake::session_extension::DeltaTableCache; @@ -27,6 +28,7 @@ use sail_physical_optimizer::{get_physical_optimizers, PhysicalOptimizerOptions} use sail_server::actor::{ActorHandle, ActorSystem}; use crate::catalog::create_catalog_manager; +use crate::checkpoint::ObjectStoreCheckpointStore; use crate::formats::create_table_format_registry; use crate::observable::SessionManagerHandle; use crate::optimizer::{default_analyzer_rules, default_optimizer_rules}; @@ -123,6 +125,9 @@ impl ServerSessionFactory { )?)) .with_extension(Arc::new(ActivityTracker::new())) .with_extension(Arc::new(CachedRelationRegistry::default())) + .with_extension(Arc::new(CheckpointStoreService::new(Arc::new( + ObjectStoreCheckpointStore, + )))) .with_extension(Arc::new(JobService::new(job_runner))) .with_extension(Arc::new(RepartitionBufferConfig::new( self.config.cluster.task_stream_buffer, diff --git a/python/pysail/tests/spark/dataframe/test_dataframe.py b/python/pysail/tests/spark/dataframe/test_dataframe.py index f25120506a..2bb6eb3d08 100644 --- a/python/pysail/tests/spark/dataframe/test_dataframe.py +++ b/python/pysail/tests/spark/dataframe/test_dataframe.py @@ -13,6 +13,7 @@ _CHECKPOINT_SOURCE_MAX_ID = 2 _CHECKPOINT_SOURCE_ROW_COUNT = 2 +_CHECKPOINT_DEEP_PLAN_ROW_COUNT = 10 def _read_checkpoint_source(spark, path): @@ -159,6 +160,22 @@ def test_dataframe_checkpoint(spark, tmp_path): spark.conf.unset("spark.checkpoint.dir") +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific object-store checkpoint URL") +def test_dataframe_checkpoint_with_memory_object_store(spark, tmp_path): + source_path = tmp_path / "source" + df = _read_checkpoint_source(spark, source_path) + spark.conf.set("spark.checkpoint.dir", "memory:///dataframe-checkpoint-test") + try: + checkpointed = df.checkpoint() + checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + assert "DataSourceExec" in checkpointed_plan + + shutil.rmtree(source_path) + _assert_checkpoint_source_result(checkpointed) + finally: + spark.conf.unset("spark.checkpoint.dir") + + @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific missing checkpoint dir coverage") def test_dataframe_checkpoint_requires_directory(spark_session_factory): spark = spark_session_factory() @@ -214,6 +231,20 @@ def test_dataframe_checkpoint_lazy(spark, tmp_path): spark.conf.unset("spark.checkpoint.dir") +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names and stack configuration") +def test_dataframe_local_checkpoint_deep_plan_explain_truncates(spark): + df = spark.range(0, _CHECKPOINT_DEEP_PLAN_ROW_COUNT) + for i in range(20): + df = df.withColumn(f"value_{i}", col("id") + lit(i)).filter(col("id") >= 0) + + checkpointed = df.localCheckpoint() + checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + + assert checkpointed.count() == _CHECKPOINT_DEEP_PLAN_ROW_COUNT + assert "RangeExec" not in checkpointed_plan + assert "DataSourceExec" in checkpointed_plan + + @pytest.mark.parametrize( ("storage_level", "source_name"), [(StorageLevel.MEMORY_ONLY, "memory-only"), (StorageLevel.DISK_ONLY, "disk-only")], From d87f3b5e36c932620717cbe21043a8afba3df2b8 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Sun, 21 Jun 2026 22:12:35 +0800 Subject: [PATCH 08/40] test: split dataframe checkpoint tests --- Cargo.lock | 1 + crates/sail-telemetry/Cargo.toml | 1 + .../src/execution/physical_plan.rs | 5 +- .../tests/spark/dataframe/test_checkpoint.py | 235 ++++++++++++++++++ .../tests/spark/dataframe/test_dataframe.py | 187 -------------- 5 files changed, 239 insertions(+), 190 deletions(-) create mode 100644 python/pysail/tests/spark/dataframe/test_checkpoint.py diff --git a/Cargo.lock b/Cargo.lock index bb2f22e296..d6e0ff2d76 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8215,6 +8215,7 @@ dependencies = [ "quote", "regex", "sail-common", + "sail-common-datafusion", "serde", "serde_yaml", "syn 2.0.118", diff --git a/crates/sail-telemetry/Cargo.toml b/crates/sail-telemetry/Cargo.toml index 7805de9c78..b4859dd382 100644 --- a/crates/sail-telemetry/Cargo.toml +++ b/crates/sail-telemetry/Cargo.toml @@ -8,6 +8,7 @@ workspace = true [dependencies] sail-common = { path = "../sail-common" } +sail-common-datafusion = { path = "../sail-common-datafusion" } thiserror = { workspace = true } env_logger = { workspace = true } diff --git a/crates/sail-telemetry/src/execution/physical_plan.rs b/crates/sail-telemetry/src/execution/physical_plan.rs index 4d0ff127b3..6586790d6b 100644 --- a/crates/sail-telemetry/src/execution/physical_plan.rs +++ b/crates/sail-telemetry/src/execution/physical_plan.rs @@ -29,6 +29,7 @@ use fastrace::Span; use fastrace_futures::StreamExt; use futures::Stream; use pin_project_lite::pin_project; +use sail_common_datafusion::utils::items::ItemTaker; use crate::common::{KeyValue, SpanAttribute}; use crate::execution::metrics::MetricEmitter; @@ -148,9 +149,7 @@ impl ExecutionPlan for TracingExec { self: Arc, children: Vec>, ) -> Result> { - let [child] = children.try_into().map_err(|_| { - datafusion::common::internal_datafusion_err!("TracingExec requires exactly one child") - })?; + let child = children.one()?; Ok(Arc::new(TracingExec::new(child, self.options.clone()))) } diff --git a/python/pysail/tests/spark/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py new file mode 100644 index 0000000000..a06ae504ba --- /dev/null +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -0,0 +1,235 @@ +import shutil + +import pandas as pd +import pytest +from pandas.testing import assert_frame_equal +from pyspark import StorageLevel +from pyspark.sql.functions import col, lit + +from pysail.testing.spark.steps.plan import normalize_plan_text +from pysail.testing.spark.utils.common import is_jvm_spark + + +def test_dataframe_local_checkpoint_survives_source_removal(spark, tmp_path): + source_path = tmp_path / "source" + spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b"), (3, "c")], + ).write.mode("overwrite").parquet(str(source_path)) + df = spark.read.parquet(str(source_path)).where(col("id") <= 2) # noqa: PLR2004 + + checkpointed = df.localCheckpoint() + shutil.rmtree(source_path) + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + assert_frame_equal( + checkpointed.where(col("id") == 2).select("value").toPandas(), # noqa: PLR2004 + pd.DataFrame({"value": ["b"]}), + ) + + +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +def test_dataframe_local_checkpoint_freezes_nondeterministic_values(spark, eager): + df = spark.sql("SELECT id, uuid() AS value FROM range(3)") + assert [row.value for row in df.sort("id").collect()] != [row.value for row in df.sort("id").collect()] + + checkpointed = df.localCheckpoint(eager=eager) + + assert [row.value for row in checkpointed.sort("id").collect()] == [ + row.value for row in checkpointed.sort("id").collect() + ] + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names") +def test_dataframe_local_checkpoint_explain_truncates_plan(spark): + df = spark.range(0, 3).withColumn("value", lit(1)).filter("id >= 0") + original = normalize_plan_text(df._explain_string()) # noqa: SLF001 + checkpointed = df.localCheckpoint() + checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + + assert "RangeExec" in original + assert original.count("RangeExec") >= 1 + assert "RangeExec" not in checkpointed_plan + assert checkpointed_plan.count("RangeExec") == 0 + assert checkpointed_plan.count("FilterExec") < 1 + assert checkpointed_plan.count("ProjectionExec") < 3 # noqa: PLR2004 + assert "DataSourceExec" in checkpointed_plan + assert checkpointed_plan.count("DataSourceExec") <= 2 # noqa: PLR2004 + + +@pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") +def test_dataframe_checkpoint(spark, tmp_path): + source_path = tmp_path / "source" + checkpoint_path = tmp_path / "checkpoints" + spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b"), (3, "c")], + ).write.mode("overwrite").parquet(str(source_path)) + df = spark.read.parquet(str(source_path)).where(col("id") <= 2) # noqa: PLR2004 + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + checkpointed = df.checkpoint() + assert list(checkpoint_path.rglob("*.arrow")) + shutil.rmtree(source_path) + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + assert_frame_equal( + checkpointed.where(col("id") == 2).select("value").toPandas(), # noqa: PLR2004 + pd.DataFrame({"value": ["b"]}), + ) + finally: + spark.conf.unset("spark.checkpoint.dir") + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific object-store checkpoint URL") +def test_dataframe_checkpoint_with_memory_object_store(spark, tmp_path): + source_path = tmp_path / "source" + spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b"), (3, "c")], + ).write.mode("overwrite").parquet(str(source_path)) + df = spark.read.parquet(str(source_path)).where(col("id") <= 2) # noqa: PLR2004 + spark.conf.set("spark.checkpoint.dir", "memory:///dataframe-checkpoint-test") + try: + checkpointed = df.checkpoint() + checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + assert "DataSourceExec" in checkpointed_plan + + shutil.rmtree(source_path) + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + assert_frame_equal( + checkpointed.where(col("id") == 2).select("value").toPandas(), # noqa: PLR2004 + pd.DataFrame({"value": ["b"]}), + ) + finally: + spark.conf.unset("spark.checkpoint.dir") + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific missing checkpoint dir coverage") +def test_dataframe_checkpoint_requires_directory(spark_session_factory): + spark = spark_session_factory() + df = spark.createDataFrame( + schema="id INT", + data=[(1,)], + ) + + with pytest.raises(Exception, match=r"spark\.checkpoint\.dir"): + df.checkpoint() + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint path validation") +def test_dataframe_checkpoint_rejects_parent_path_components(spark, tmp_path): + df = spark.createDataFrame( + schema="id INT", + data=[(1,)], + ) + spark.conf.set("spark.checkpoint.dir", str(tmp_path / ".." / "checkpoints")) + try: + with pytest.raises(Exception, match="parent path"): + df.checkpoint() + finally: + spark.conf.unset("spark.checkpoint.dir") + + +def test_dataframe_local_checkpoint_lazy_survives_source_removal_after_first_action(spark, tmp_path): + source_path = tmp_path / "source" + spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b"), (3, "c")], + ).write.mode("overwrite").parquet(str(source_path)) + df = spark.read.parquet(str(source_path)).where(col("id") <= 2) # noqa: PLR2004 + + checkpointed = df.localCheckpoint(eager=False) + assert checkpointed.count() == 2 # noqa: PLR2004 + shutil.rmtree(source_path) + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + assert_frame_equal( + checkpointed.where(col("id") == 2).select("value").toPandas(), # noqa: PLR2004 + pd.DataFrame({"value": ["b"]}), + ) + + +@pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") +def test_dataframe_checkpoint_lazy(spark, tmp_path): + source_path = tmp_path / "source" + checkpoint_path = tmp_path / "checkpoints" + spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b"), (3, "c")], + ).write.mode("overwrite").parquet(str(source_path)) + df = spark.read.parquet(str(source_path)).where(col("id") <= 2) # noqa: PLR2004 + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + checkpointed = df.checkpoint(eager=False) + assert not list(checkpoint_path.rglob("*.arrow")) + assert checkpointed.count() == 2 # noqa: PLR2004 + assert list(checkpoint_path.rglob("*.arrow")) + shutil.rmtree(source_path) + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + assert_frame_equal( + checkpointed.where(col("id") == 2).select("value").toPandas(), # noqa: PLR2004 + pd.DataFrame({"value": ["b"]}), + ) + finally: + spark.conf.unset("spark.checkpoint.dir") + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names and stack configuration") +def test_dataframe_local_checkpoint_deep_plan_explain_truncates(spark): + df = spark.range(0, 10) + for i in range(20): + df = df.withColumn(f"value_{i}", col("id") + lit(i)).filter(col("id") >= 0) + + checkpointed = df.localCheckpoint() + checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + + assert checkpointed.count() == 10 # noqa: PLR2004 + assert "RangeExec" not in checkpointed_plan + assert checkpointed_plan.count("RangeExec") == 0 + assert checkpointed_plan.count("FilterExec") < 1 + assert checkpointed_plan.count("ProjectionExec") < 3 # noqa: PLR2004 + assert "DataSourceExec" in checkpointed_plan + assert checkpointed_plan.count("DataSourceExec") <= 2 # noqa: PLR2004 + + +@pytest.mark.parametrize( + ("storage_level", "source_name"), + [(StorageLevel.MEMORY_ONLY, "memory-only"), (StorageLevel.DISK_ONLY, "disk-only")], + ids=["memory-only", "disk-only"], +) +def test_dataframe_local_checkpoint_storage_level_survives_source_removal(spark, tmp_path, storage_level, source_name): + source_path = tmp_path / f"source-{source_name}" + spark.createDataFrame( + schema="id INT, value STRING", + data=[(1, "a"), (2, "b"), (3, "c")], + ).write.mode("overwrite").parquet(str(source_path)) + df = spark.read.parquet(str(source_path)).where(col("id") <= 2) # noqa: PLR2004 + + checkpointed = df.localCheckpoint(storageLevel=storage_level) + shutil.rmtree(source_path) + + assert_frame_equal( + checkpointed.sort("id").toPandas(), + pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), + ) + assert_frame_equal( + checkpointed.where(col("id") == 2).select("value").toPandas(), # noqa: PLR2004 + pd.DataFrame({"value": ["b"]}), + ) diff --git a/python/pysail/tests/spark/dataframe/test_dataframe.py b/python/pysail/tests/spark/dataframe/test_dataframe.py index 2bb6eb3d08..efa3af38bd 100644 --- a/python/pysail/tests/spark/dataframe/test_dataframe.py +++ b/python/pysail/tests/spark/dataframe/test_dataframe.py @@ -1,49 +1,10 @@ -import shutil - 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, row_number from pyspark.sql.window import Window -from pysail.testing.spark.steps.plan import normalize_plan_text -from pysail.testing.spark.utils.common import is_jvm_spark - -_CHECKPOINT_SOURCE_MAX_ID = 2 -_CHECKPOINT_SOURCE_ROW_COUNT = 2 -_CHECKPOINT_DEEP_PLAN_ROW_COUNT = 10 - - -def _read_checkpoint_source(spark, path): - spark.createDataFrame( - schema="id INT, value STRING", - data=[(1, "a"), (2, "b"), (3, "c")], - ).write.mode("overwrite").parquet(str(path)) - return spark.read.parquet(str(path)).where(col("id") <= _CHECKPOINT_SOURCE_MAX_ID) - - -def _assert_checkpoint_source_result(df): - assert_frame_equal( - df.sort("id").toPandas(), - pd.DataFrame({"id": [1, 2], "value": ["a", "b"]}).astype({"id": "int32"}), - ) - assert_frame_equal( - df.where(col("id") == _CHECKPOINT_SOURCE_MAX_ID).select("value").toPandas(), - pd.DataFrame({"value": ["b"]}), - ) - - -def _checkpoint_arrow_files(path): - if not path.exists(): - return [] - return list(path.rglob("*.arrow")) - - -def _uuid_values(df): - return [row.value for row in df.sort("id").collect()] - def test_dataframe_drop(spark): df = spark.createDataFrame([(14, "Tom"), (23, "Alice"), (16, "Bob")], ["age", "name"]) @@ -112,154 +73,6 @@ def test_dataframe_drop(spark): ) -def test_dataframe_local_checkpoint_survives_source_removal(spark, tmp_path): - source_path = tmp_path / "source" - df = _read_checkpoint_source(spark, source_path) - - checkpointed = df.localCheckpoint() - shutil.rmtree(source_path) - - _assert_checkpoint_source_result(checkpointed) - - -@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) -def test_dataframe_local_checkpoint_freezes_nondeterministic_values(spark, eager): - df = spark.sql("SELECT id, uuid() AS value FROM range(3)") - assert _uuid_values(df) != _uuid_values(df) - - checkpointed = df.localCheckpoint(eager=eager) - - assert _uuid_values(checkpointed) == _uuid_values(checkpointed) - - -@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names") -def test_dataframe_local_checkpoint_explain_truncates_plan(spark): - df = spark.range(0, 3).withColumn("value", lit(1)).filter("id >= 0") - original = normalize_plan_text(df._explain_string()) # noqa: SLF001 - checkpointed = df.localCheckpoint() - checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 - - assert "RangeExec" in original - assert "RangeExec" not in checkpointed_plan - assert "DataSourceExec" in checkpointed_plan - - -@pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") -def test_dataframe_checkpoint(spark, tmp_path): - source_path = tmp_path / "source" - checkpoint_path = tmp_path / "checkpoints" - df = _read_checkpoint_source(spark, source_path) - spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) - try: - checkpointed = df.checkpoint() - assert _checkpoint_arrow_files(checkpoint_path) - shutil.rmtree(source_path) - - _assert_checkpoint_source_result(checkpointed) - finally: - spark.conf.unset("spark.checkpoint.dir") - - -@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific object-store checkpoint URL") -def test_dataframe_checkpoint_with_memory_object_store(spark, tmp_path): - source_path = tmp_path / "source" - df = _read_checkpoint_source(spark, source_path) - spark.conf.set("spark.checkpoint.dir", "memory:///dataframe-checkpoint-test") - try: - checkpointed = df.checkpoint() - checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 - assert "DataSourceExec" in checkpointed_plan - - shutil.rmtree(source_path) - _assert_checkpoint_source_result(checkpointed) - finally: - spark.conf.unset("spark.checkpoint.dir") - - -@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific missing checkpoint dir coverage") -def test_dataframe_checkpoint_requires_directory(spark_session_factory): - spark = spark_session_factory() - df = spark.createDataFrame( - schema="id INT", - data=[(1,)], - ) - - with pytest.raises(Exception, match=r"spark\.checkpoint\.dir"): - df.checkpoint() - - -@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint path validation") -def test_dataframe_checkpoint_rejects_parent_path_components(spark, tmp_path): - df = spark.createDataFrame( - schema="id INT", - data=[(1,)], - ) - spark.conf.set("spark.checkpoint.dir", str(tmp_path / ".." / "checkpoints")) - try: - with pytest.raises(Exception, match="parent path"): - df.checkpoint() - finally: - spark.conf.unset("spark.checkpoint.dir") - - -def test_dataframe_local_checkpoint_lazy_survives_source_removal_after_first_action(spark, tmp_path): - source_path = tmp_path / "source" - df = _read_checkpoint_source(spark, source_path) - - checkpointed = df.localCheckpoint(eager=False) - assert checkpointed.count() == _CHECKPOINT_SOURCE_ROW_COUNT - shutil.rmtree(source_path) - - _assert_checkpoint_source_result(checkpointed) - - -@pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") -def test_dataframe_checkpoint_lazy(spark, tmp_path): - source_path = tmp_path / "source" - checkpoint_path = tmp_path / "checkpoints" - df = _read_checkpoint_source(spark, source_path) - spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) - try: - checkpointed = df.checkpoint(eager=False) - assert not _checkpoint_arrow_files(checkpoint_path) - assert checkpointed.count() == _CHECKPOINT_SOURCE_ROW_COUNT - assert _checkpoint_arrow_files(checkpoint_path) - shutil.rmtree(source_path) - - _assert_checkpoint_source_result(checkpointed) - finally: - spark.conf.unset("spark.checkpoint.dir") - - -@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names and stack configuration") -def test_dataframe_local_checkpoint_deep_plan_explain_truncates(spark): - df = spark.range(0, _CHECKPOINT_DEEP_PLAN_ROW_COUNT) - for i in range(20): - df = df.withColumn(f"value_{i}", col("id") + lit(i)).filter(col("id") >= 0) - - checkpointed = df.localCheckpoint() - checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 - - assert checkpointed.count() == _CHECKPOINT_DEEP_PLAN_ROW_COUNT - assert "RangeExec" not in checkpointed_plan - assert "DataSourceExec" in checkpointed_plan - - -@pytest.mark.parametrize( - ("storage_level", "source_name"), - [(StorageLevel.MEMORY_ONLY, "memory-only"), (StorageLevel.DISK_ONLY, "disk-only")], - ids=["memory-only", "disk-only"], -) -def test_dataframe_local_checkpoint_storage_level_survives_source_removal(spark, tmp_path, storage_level, source_name): - source_path = tmp_path / f"source-{source_name}" - df = _read_checkpoint_source(spark, source_path) - - checkpointed = df.localCheckpoint(storageLevel=storage_level) - shutil.rmtree(source_path) - - _assert_checkpoint_source_result(checkpointed) - - def test_dataframe_with_column_alias(spark): df = spark.createDataFrame( schema="id INTEGER, value STRING", From 1ad1e4fb8771f7961ce3bb0b5250670bd9a03b8a Mon Sep 17 00:00:00 2001 From: XL Liang Date: Sun, 21 Jun 2026 22:43:22 +0800 Subject: [PATCH 09/40] fix: harden checkpoint materialization lifecycle --- .../src/cached_relation.rs | 296 +++++++++++++----- crates/sail-plan/src/lib.rs | 11 + crates/sail-plan/src/resolver/query/mod.rs | 2 +- crates/sail-session/src/checkpoint.rs | 18 +- crates/sail-session/src/planner.rs | 2 +- .../src/session_manager/actor/handler.rs | 5 + .../src/service/plan_executor.rs | 34 +- .../tests/spark/dataframe/test_checkpoint.py | 57 ++++ 8 files changed, 310 insertions(+), 115 deletions(-) diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index 78d32a883c..21f0ae87a7 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -1,4 +1,5 @@ use std::cmp::Ordering; +use std::collections::hash_map::Entry; use std::collections::HashMap; use std::fmt::Formatter; use std::sync::{Arc, RwLock}; @@ -8,7 +9,13 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::physical_plan::ArrowSource; use datafusion::execution::disk_manager::RefCountedTempFile; use datafusion::execution::object_store::ObjectStoreUrl; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::{ + collect_partitioned, with_new_children_if_necessary, DisplayAs, DisplayFormatType, + ExecutionPlan, ExecutionPlanProperties, PlanProperties, +}; use datafusion::prelude::SessionContext; use datafusion_common::{ internal_datafusion_err, DFSchema, DFSchemaRef, DataFusionError, Result, Statistics, @@ -19,15 +26,13 @@ use datafusion_datasource::memory::MemorySourceConfig; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::{PartitionedFile, TableSchema}; use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; -use futures::StreamExt; +use futures::future::BoxFuture; use object_store::ObjectMeta; use sail_common::spec; use crate::array::record_batch::{read_record_batches, write_record_batches}; use crate::extension::{SessionExtension, SessionExtensionAccessor}; -use crate::rename::physical_plan::rename_physical_plan; -use crate::session::checkpoint::CheckpointStoreService; -use crate::session::job::JobService; +use crate::session::checkpoint::{CheckpointStoreService, ReliableCheckpoint}; #[derive(Debug, Clone)] pub enum CachedRelationCleanup { @@ -112,9 +117,9 @@ enum CachedRelationData { #[derive(Debug, Clone)] enum CachedRelationMaterialized { Local(CachedRelationLocalMaterialized), - Physical { + Reliable { schema: SchemaRef, - plan: Arc, + checkpoint: ReliableCheckpoint, }, } @@ -134,8 +139,7 @@ struct CachedRelationDiskPartition { #[derive(Debug, Clone)] struct CachedRelationPending { - plan: Arc, - fields: Option>, + plan: Arc, target: CachedRelationPendingTarget, } @@ -145,6 +149,85 @@ enum CachedRelationPendingTarget { Reliable { path: String }, } +#[derive(Debug, Clone)] +pub struct PendingCachedRelationExec { + relation_id: String, + relation: CachedRelation, + properties: Arc, +} + +impl PendingCachedRelationExec { + fn new(relation_id: String, relation: CachedRelation, plan: Arc) -> Self { + let schema = plan.schema(); + let partition_count = plan.output_partitioning().partition_count(); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(partition_count), + EmissionType::Final, + Boundedness::Bounded, + )); + Self { + relation_id, + relation, + properties, + } + } + + pub fn relation_id(&self) -> &str { + &self.relation_id + } + + async fn materialize(&self, ctx: &SessionContext) -> Result> { + self.relation.materialize_pending(ctx).await + } +} + +impl DisplayAs for PendingCachedRelationExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "PendingCachedRelationExec: relation_id={}", + self.relation_id + ) + } +} + +impl ExecutionPlan for PendingCachedRelationExec { + fn name(&self) -> &str { + "PendingCachedRelationExec" + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if !children.is_empty() { + return Err(internal_datafusion_err!( + "PendingCachedRelationExec should have no children" + )); + } + Ok(self) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Err(internal_datafusion_err!( + "PendingCachedRelationExec should be materialized before execution" + )) + } +} + impl CachedRelation { pub fn new(plan: Arc, cleanup: Option) -> Self { Self { @@ -180,15 +263,13 @@ impl CachedRelation { } pub fn new_pending_local_checkpoint( - plan: Arc, - fields: Option>, + plan: Arc, storage_level: spec::StorageLevel, ) -> Self { Self { data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( CachedRelationPending { plan, - fields, target: CachedRelationPendingTarget::Local { storage_level }, }, ))), @@ -196,17 +277,12 @@ impl CachedRelation { } } - pub fn new_pending_reliable_checkpoint( - plan: Arc, - fields: Option>, - path: String, - ) -> Self { + pub fn new_pending_reliable_checkpoint(plan: Arc, path: String) -> Self { let cleanup_path = path.clone(); Self { data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( CachedRelationPending { plan, - fields, target: CachedRelationPendingTarget::Reliable { path }, }, ))), @@ -214,28 +290,37 @@ impl CachedRelation { } } - pub async fn to_logical_plan( - &self, - ctx: &SessionContext, - relation_id: &str, - ) -> Result { - let mut data = self.data.lock().await; - if let CachedRelationData::Pending(pending) = &*data { - *data = pending.materialize(ctx).await?; - } + pub async fn to_logical_plan(&self, relation_id: &str) -> Result { + let data = self.data.lock().await; match &*data { CachedRelationData::LogicalPlan(plan) => Ok(plan.as_ref().clone()), CachedRelationData::Materialized(materialized) => { materialized.to_logical_plan(relation_id) } - CachedRelationData::Pending(_) => Err(internal_datafusion_err!( - "cached relation materialization did not complete" - )), + CachedRelationData::Pending(pending) => pending.to_logical_plan(relation_id), } } - pub async fn to_physical_plan(&self) -> Result> { + pub async fn to_physical_plan(&self, relation_id: &str) -> Result> { let data = self.data.lock().await; + match &*data { + CachedRelationData::Materialized(materialized) => materialized.to_physical_plan().await, + CachedRelationData::LogicalPlan(_) => Err(internal_datafusion_err!( + "cached relation is not materialized as a physical plan" + )), + CachedRelationData::Pending(pending) => Ok(Arc::new(PendingCachedRelationExec::new( + relation_id.to_string(), + self.clone(), + Arc::clone(&pending.plan), + ))), + } + } + + async fn materialize_pending(&self, ctx: &SessionContext) -> Result> { + let mut data = self.data.lock().await; + if let CachedRelationData::Pending(pending) = &*data { + *data = pending.materialize(ctx).await?; + } match &*data { CachedRelationData::Materialized(materialized) => materialized.to_physical_plan().await, CachedRelationData::LogicalPlan(_) => Err(internal_datafusion_err!( @@ -263,12 +348,16 @@ impl CachedRelationRegistry { .relations .write() .map_err(|e| internal_datafusion_err!("{e}"))?; - if relations.insert(relation_id.clone(), relation).is_some() { - return Err(internal_datafusion_err!( - "cached relation already exists: {relation_id}" - )); + match relations.entry(relation_id) { + Entry::Occupied(entry) => Err(internal_datafusion_err!( + "cached relation already exists: {}", + entry.key() + )), + Entry::Vacant(entry) => { + entry.insert(relation); + Ok(()) + } } - Ok(()) } pub fn get(&self, relation_id: &str) -> Result> { @@ -286,6 +375,14 @@ impl CachedRelationRegistry { .map_err(|e| internal_datafusion_err!("{e}"))?; Ok(relations.remove(relation_id)) } + + pub fn drain(&self) -> Result> { + let mut relations = self + .relations + .write() + .map_err(|e| internal_datafusion_err!("{e}"))?; + Ok(relations.drain().map(|(_, relation)| relation).collect()) + } } impl SessionExtension for CachedRelationRegistry { @@ -296,8 +393,7 @@ impl SessionExtension for CachedRelationRegistry { impl CachedRelationPending { async fn materialize(&self, ctx: &SessionContext) -> Result { - let physical_plan = - create_physical_plan(ctx, self.plan.as_ref().clone(), self.fields.as_ref()).await?; + let physical_plan = materialize_cached_relations(ctx, Arc::clone(&self.plan)).await?; match &self.target { CachedRelationPendingTarget::Local { storage_level } => { materialize_local_checkpoint(ctx, physical_plan, storage_level.clone()).await @@ -307,13 +403,22 @@ impl CachedRelationPending { } } } + + fn to_logical_plan(&self, relation_id: &str) -> Result { + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(CachedRelationNode::try_new( + relation_id.to_string(), + self.plan.schema(), + )?), + })) + } } impl CachedRelationMaterialized { fn schema(&self) -> SchemaRef { match self { Self::Local(local) => Arc::clone(&local.schema), - Self::Physical { schema, .. } => Arc::clone(schema), + Self::Reliable { schema, .. } => Arc::clone(schema), } } @@ -329,7 +434,11 @@ impl CachedRelationMaterialized { async fn to_physical_plan(&self) -> Result> { match self { Self::Local(local) => local.to_physical_plan().await, - Self::Physical { plan, .. } => Ok(Arc::clone(plan)), + Self::Reliable { schema, checkpoint } => create_arrow_checkpoint_scan( + checkpoint.object_store_url().clone(), + checkpoint.object_meta().to_vec(), + schema, + ), } } } @@ -456,23 +565,22 @@ impl CachedRelationDiskPartition { } } -async fn create_physical_plan( - ctx: &SessionContext, - plan: LogicalPlan, - fields: Option<&Vec>, -) -> Result> { - let df = ctx.execute_logical_plan(plan).await?; - let (session_state, plan) = df.into_parts(); - let plan = session_state.optimize(&plan)?; - let plan = session_state - .query_planner() - .create_physical_plan(&plan, &session_state) - .await?; - if let Some(fields) = fields { - rename_physical_plan(plan, fields) - } else { - Ok(plan) - } +pub fn materialize_cached_relations<'a>( + ctx: &'a SessionContext, + plan: Arc, +) -> BoxFuture<'a, Result>> { + Box::pin(async move { + let mut children = Vec::with_capacity(plan.children().len()); + for child in plan.children() { + children.push(materialize_cached_relations(ctx, Arc::clone(child)).await?); + } + let plan = with_new_children_if_necessary(plan, children)?; + if let Some(pending) = plan.downcast_ref::() { + pending.materialize(ctx).await + } else { + Ok(plan) + } + }) } async fn materialize_local_checkpoint( @@ -480,6 +588,7 @@ async fn materialize_local_checkpoint( plan: Arc, storage_level: spec::StorageLevel, ) -> Result { + let plan = materialize_cached_relations(ctx, plan).await?; let schema = plan.schema(); let partitions = collect_checkpoint_partitions(ctx, plan).await?; let materialized = @@ -494,6 +603,7 @@ async fn materialize_reliable_checkpoint( plan: Arc, path: &str, ) -> Result { + let plan = materialize_cached_relations(ctx, plan).await?; let schema = plan.schema(); let service = ctx.extension::()?; let checkpoint = match service @@ -506,16 +616,8 @@ async fn materialize_reliable_checkpoint( return Err(error); } }; - let physical_plan = create_arrow_checkpoint_scan( - checkpoint.object_store_url().clone(), - checkpoint.object_meta().to_vec(), - &schema, - )?; Ok(CachedRelationData::Materialized( - CachedRelationMaterialized::Physical { - schema, - plan: physical_plan, - }, + CachedRelationMaterialized::Reliable { schema, checkpoint }, )) } @@ -546,17 +648,28 @@ pub async fn cleanup_checkpoint_path(ctx: &SessionContext, path: &str) -> Result service.cleanup_checkpoint_path(ctx, path).await } +pub async fn cleanup_cached_relation(ctx: &SessionContext, relation: CachedRelation) -> Result<()> { + match relation.into_cleanup() { + Some(CachedRelationCleanup::ObjectStorePath(path)) => { + cleanup_checkpoint_path(ctx, &path).await + } + None => Ok(()), + } +} + +pub async fn cleanup_cached_relations(ctx: &SessionContext) -> Result<()> { + let relations = ctx.extension::()?.drain()?; + for relation in relations { + cleanup_cached_relation(ctx, relation).await?; + } + Ok(()) +} + async fn collect_checkpoint_partitions( ctx: &SessionContext, plan: Arc, ) -> Result>> { - let service = ctx.extension::()?; - let mut stream = service.runner().execute(ctx, plan).await?; - let mut batches = vec![]; - while let Some(batch) = stream.next().await { - batches.push(batch?); - } - Ok(vec![batches]) + collect_partitioned(plan, ctx.task_ctx()).await } async fn write_disk_partition(ctx: &SessionContext, bytes: &[u8]) -> Result { @@ -576,7 +689,8 @@ mod tests { use std::sync::Mutex; use async_trait::async_trait; - use datafusion_expr::LogicalPlanBuilder; + use datafusion::arrow::datatypes::Schema; + use datafusion::physical_plan::empty::EmptyExec; use super::*; use crate::session::checkpoint::{CheckpointStore, ReliableCheckpoint}; @@ -584,9 +698,8 @@ mod tests { #[test] fn pending_reliable_checkpoint_tracks_cleanup_path() -> Result<()> { let path = "memory://checkpoint-root/session/relation".to_string(); - let plan = LogicalPlanBuilder::empty(false).build()?; - let relation = - CachedRelation::new_pending_reliable_checkpoint(Arc::new(plan), None, path.clone()); + let plan: Arc = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + let relation = CachedRelation::new_pending_reliable_checkpoint(plan, path.clone()); match relation.into_cleanup() { Some(CachedRelationCleanup::ObjectStorePath(actual)) => { @@ -601,6 +714,35 @@ mod tests { Ok(()) } + #[test] + fn registry_insert_does_not_replace_existing_relation() -> Result<()> { + let registry = CachedRelationRegistry::default(); + let first_path = "memory:///checkpoint-root/first".to_string(); + let second_path = "memory:///checkpoint-root/second".to_string(); + let plan: Arc = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + let first = + CachedRelation::new_pending_reliable_checkpoint(Arc::clone(&plan), first_path.clone()); + let second = CachedRelation::new_pending_reliable_checkpoint(plan, second_path); + + registry.insert("relation".to_string(), first)?; + assert!(registry.insert("relation".to_string(), second).is_err()); + + let relation = registry.get("relation")?.ok_or_else(|| { + internal_datafusion_err!("cached relation missing after duplicate insert") + })?; + match relation.into_cleanup() { + Some(CachedRelationCleanup::ObjectStorePath(actual)) => { + assert_eq!(actual, first_path); + } + other => { + return Err(internal_datafusion_err!( + "unexpected cleanup value: {other:?}" + )); + } + } + Ok(()) + } + #[derive(Default)] struct TestCheckpointStore { cleanup_paths: Mutex>, diff --git a/crates/sail-plan/src/lib.rs b/crates/sail-plan/src/lib.rs index 1ab197784c..5559a8fdd7 100644 --- a/crates/sail-plan/src/lib.rs +++ b/crates/sail-plan/src/lib.rs @@ -7,6 +7,7 @@ use datafusion_common::display::{PlanType, StringifiedPlan, ToStringifiedPlan}; use datafusion_common::Result; use datafusion_expr::LogicalPlan; use sail_common::spec; +use sail_common_datafusion::cached_relation::materialize_cached_relations; use sail_common_datafusion::rename::physical_plan::rename_physical_plan; use crate::config::PlanConfig; @@ -35,6 +36,16 @@ pub async fn resolve_and_execute_plan( ctx: &SessionContext, config: Arc, plan: spec::Plan, +) -> PlanResult<(Arc, Vec)> { + let (plan, info) = resolve_physical_plan(ctx, config, plan).await?; + let plan = materialize_cached_relations(ctx, plan).await?; + Ok((plan, info)) +} + +pub async fn resolve_physical_plan( + ctx: &SessionContext, + config: Arc, + plan: spec::Plan, ) -> PlanResult<(Arc, Vec)> { let mut info = vec![]; let resolver = PlanResolver::new(ctx, config); diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index 480037f88a..b313b17913 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/crates/sail-plan/src/resolver/query/mod.rs @@ -235,7 +235,7 @@ impl PlanResolver<'_> { let relation = registry.get(&relation_id)?.ok_or_else(|| { PlanError::invalid(format!("No DataFrame with id {relation_id} is found")) })?; - let plan = relation.to_logical_plan(self.ctx, &relation_id).await?; + let plan = relation.to_logical_plan(&relation_id).await?; let names = state.register_fields(plan.schema().inner().fields()); rename_logical_plan(plan, &names)? } diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs index f153a9f4ac..dcb932dbfb 100644 --- a/crates/sail-session/src/checkpoint.rs +++ b/crates/sail-session/src/checkpoint.rs @@ -2,14 +2,12 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::{execute_stream_partitioned, ExecutionPlan}; use datafusion::prelude::SessionContext; use datafusion_common::Result; use futures::StreamExt; use sail_common_datafusion::array::record_batch::write_record_batches; -use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::checkpoint::{CheckpointStore, ReliableCheckpoint}; -use sail_common_datafusion::session::job::JobService; use sail_object_store::{delete_object_store_prefix, resolve_object_store_path}; #[derive(Debug, Default)] @@ -26,17 +24,17 @@ impl CheckpointStore for ObjectStoreCheckpointStore { ) -> Result { let runtime_env = ctx.runtime_env(); let checkpoint_path = resolve_object_store_path(runtime_env.as_ref(), path)?; - let service = ctx.extension::()?; - let mut stream = service.runner().execute(ctx, plan).await?; + let streams = execute_stream_partitioned(plan, ctx.task_ctx())?; let mut files = vec![]; - let mut file_index = 0usize; - while let Some(batch) = stream.next().await { - let batch = batch?; - let bytes = write_record_batches(&[batch], schema.as_ref())?; + for (file_index, mut stream) in streams.into_iter().enumerate() { + let mut batches = vec![]; + while let Some(batch) = stream.next().await { + batches.push(batch?); + } + let bytes = write_record_batches(&batches, schema.as_ref())?; let file_path = checkpoint_path.child(&format!("part-{file_index:05}.arrow")); files.push(checkpoint_path.put_bytes(&file_path, bytes).await?); - file_index += 1; } if files.is_empty() { diff --git a/crates/sail-session/src/planner.rs b/crates/sail-session/src/planner.rs index ea71aa0cc5..d51d698342 100644 --- a/crates/sail-session/src/planner.rs +++ b/crates/sail-session/src/planner.rs @@ -109,7 +109,7 @@ impl ExtensionPlanner for ExtensionPhysicalPlanner { node.relation_id() )) })?; - relation.to_physical_plan().await? + relation.to_physical_plan(node.relation_id()).await? } else if let Some(node) = node.as_any().downcast_ref::() { let schema = UserDefinedLogicalNode::schema(node).inner().clone(); let projection = (0..schema.fields().len()).collect(); diff --git a/crates/sail-session/src/session_manager/actor/handler.rs b/crates/sail-session/src/session_manager/actor/handler.rs index 900bff95d1..422c78895d 100644 --- a/crates/sail-session/src/session_manager/actor/handler.rs +++ b/crates/sail-session/src/session_manager/actor/handler.rs @@ -5,6 +5,7 @@ use datafusion::prelude::SessionContext; use fastrace::collector::SpanContext; use fastrace::Span; use log::{info, warn}; +use sail_common_datafusion::cached_relation::cleanup_cached_relations; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::activity::ActivityTracker; use sail_common_datafusion::session::job::JobService; @@ -328,8 +329,12 @@ impl SessionManagerActor { }; let handle = ctx.handle().clone(); let (tx, rx) = oneshot::channel(); + let context = context.clone(); ctx.spawn(async move { service.runner().stop(tx).await; + if let Err(error) = cleanup_cached_relations(&context).await { + warn!("failed to clean cached relations for session {session_id}: {error}"); + } let message = match rx.await { Ok(x) => SessionManagerEvent::SetSessionHistory { session_id, diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index ab5447bb0a..1da0352134 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -1,6 +1,5 @@ use std::path::{Component, Path}; use std::pin::Pin; -use std::sync::Arc; use std::task::{Context, Poll}; use datafusion::arrow::compute::concat_batches; @@ -12,12 +11,12 @@ use futures::stream; use log::{debug, warn}; use sail_common::spec; use sail_common_datafusion::cached_relation::{ - cleanup_checkpoint_path as cleanup_cached_checkpoint_path, CachedRelation, - CachedRelationCleanup, CachedRelationRegistry, + cleanup_cached_relation, cleanup_checkpoint_path as cleanup_cached_checkpoint_path, + CachedRelation, CachedRelationRegistry, }; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; -use sail_plan::{resolve_and_execute_plan, resolve_logical_plan}; +use sail_plan::{resolve_and_execute_plan, resolve_physical_plan}; use tonic::codegen::tokio_stream::wrappers::ReceiverStream; use tonic::codegen::tokio_stream::Stream; use tonic::Status; @@ -549,7 +548,7 @@ pub(crate) async fn handle_execute_checkpoint_command( let plan: spec::Plan = relation.try_into()?; let relation_id = Uuid::new_v4().to_string(); let (relation, cleanup_on_insert_error) = if eager { - let (physical_plan, _) = resolve_and_execute_plan(ctx, spark.plan_config()?, plan).await?; + let (physical_plan, _) = resolve_physical_plan(ctx, spark.plan_config()?, plan).await?; if local { ( CachedRelation::new_local_checkpoint( @@ -568,24 +567,19 @@ pub(crate) async fn handle_execute_checkpoint_command( ) } } else if local { - let named_plan = resolve_logical_plan(ctx, spark.plan_config()?, plan).await?; + let (physical_plan, _) = resolve_physical_plan(ctx, spark.plan_config()?, plan).await?; ( CachedRelation::new_pending_local_checkpoint( - Arc::new(named_plan.plan), - named_plan.fields, + physical_plan, storage_level.required("local checkpoint storage level")?, ), None, ) } else { let path = checkpoint_path(&spark, &relation_id)?; - let named_plan = resolve_logical_plan(ctx, spark.plan_config()?, plan).await?; + let (physical_plan, _) = resolve_physical_plan(ctx, spark.plan_config()?, plan).await?; ( - CachedRelation::new_pending_reliable_checkpoint( - Arc::new(named_plan.plan), - named_plan.fields, - path, - ), + CachedRelation::new_pending_reliable_checkpoint(physical_plan, path), None, ) }; @@ -731,18 +725,6 @@ async fn cleanup_checkpoint_path(ctx: &SessionContext, path: &str) -> SparkResul Ok(()) } -async fn cleanup_cached_relation( - ctx: &SessionContext, - relation: CachedRelation, -) -> SparkResult<()> { - match relation.into_cleanup() { - Some(CachedRelationCleanup::ObjectStorePath(path)) => { - cleanup_checkpoint_path(ctx, &path).await - } - None => Ok(()), - } -} - 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/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index a06ae504ba..af4fdc6bb5 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -1,4 +1,5 @@ import shutil +import time import pandas as pd import pytest @@ -191,6 +192,62 @@ def test_dataframe_checkpoint_lazy(spark, tmp_path): spark.conf.unset("spark.checkpoint.dir") +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific lazy checkpoint explain behavior") +def test_dataframe_checkpoint_lazy_explain_does_not_materialize(spark, tmp_path): + checkpoint_path = tmp_path / "checkpoints" + df = spark.range(0, 5).where(col("id") <= 2) # noqa: PLR2004 + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + checkpointed = df.checkpoint(eager=False) + checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + + assert "PendingCachedRelationExec" in checkpointed_plan + assert not list(checkpoint_path.rglob("*.arrow")) + + assert checkpointed.count() == 3 # noqa: PLR2004 + assert list(checkpoint_path.rglob("*.arrow")) + + materialized_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + assert "PendingCachedRelationExec" not in materialized_plan + assert "DataSourceExec" in materialized_plan + finally: + spark.conf.unset("spark.checkpoint.dir") + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint session cleanup") +def test_dataframe_checkpoint_cleanup_on_session_stop(spark_session_factory, tmp_path): + spark = spark_session_factory() + checkpoint_path = tmp_path / "checkpoints" + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + + _checkpointed = spark.range(0, 5).checkpoint() + assert list(checkpoint_path.rglob("*.arrow")) + + spark.stop() + deadline = time.monotonic() + 5 + while time.monotonic() < deadline and list(checkpoint_path.rglob("*.arrow")): + time.sleep(0.1) + assert not list(checkpoint_path.rglob("*.arrow")) + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint partition preservation") +@pytest.mark.parametrize("local", [True, False], ids=["local", "reliable"]) +def test_dataframe_checkpoint_preserves_multiple_output_partitions(spark, tmp_path, local): + df = spark.range(0, 100, numPartitions=4).repartition(4) + if local: + checkpointed = df.localCheckpoint() + else: + checkpoint_path = tmp_path / "checkpoints" + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + checkpointed = df.checkpoint() + finally: + spark.conf.unset("spark.checkpoint.dir") + + partition_ids = {row.pid for row in checkpointed.selectExpr("spark_partition_id() AS pid").distinct().collect()} + assert len(partition_ids) > 1 + + @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names and stack configuration") def test_dataframe_local_checkpoint_deep_plan_explain_truncates(spark): df = spark.range(0, 10) From 6c2fad142d0ee9da28205dbd38063f600b96e592 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Sun, 21 Jun 2026 22:51:10 +0800 Subject: [PATCH 10/40] refactor: clean checkpoint implementation --- .../src/cached_relation.rs | 36 ++++--------------- .../src/service/plan_executor.rs | 8 +---- .../tests/spark/dataframe/test_checkpoint.py | 16 ++++----- 3 files changed, 15 insertions(+), 45 deletions(-) diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index 21f0ae87a7..952c22313d 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -35,7 +35,7 @@ use crate::extension::{SessionExtension, SessionExtensionAccessor}; use crate::session::checkpoint::{CheckpointStoreService, ReliableCheckpoint}; #[derive(Debug, Clone)] -pub enum CachedRelationCleanup { +enum CachedRelationCleanup { ObjectStorePath(String), } @@ -109,7 +109,6 @@ pub struct CachedRelation { #[derive(Debug, Clone)] enum CachedRelationData { - LogicalPlan(Arc), Materialized(CachedRelationMaterialized), Pending(CachedRelationPending), } @@ -129,7 +128,6 @@ struct CachedRelationLocalMaterialized { memory_partitions: Option>>>, serialized_memory_partitions: Option>>>, disk_partitions: Option>>, - storage_level: Option, } #[derive(Debug, Clone)] @@ -150,7 +148,7 @@ enum CachedRelationPendingTarget { } #[derive(Debug, Clone)] -pub struct PendingCachedRelationExec { +struct PendingCachedRelationExec { relation_id: String, relation: CachedRelation, properties: Arc, @@ -173,10 +171,6 @@ impl PendingCachedRelationExec { } } - pub fn relation_id(&self) -> &str { - &self.relation_id - } - async fn materialize(&self, ctx: &SessionContext) -> Result> { self.relation.materialize_pending(ctx).await } @@ -229,15 +223,6 @@ impl ExecutionPlan for PendingCachedRelationExec { } impl CachedRelation { - pub fn new(plan: Arc, cleanup: Option) -> Self { - Self { - data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::LogicalPlan( - plan, - ))), - cleanup, - } - } - pub async fn new_local_checkpoint( ctx: &SessionContext, plan: Arc, @@ -293,7 +278,6 @@ impl CachedRelation { pub async fn to_logical_plan(&self, relation_id: &str) -> Result { let data = self.data.lock().await; match &*data { - CachedRelationData::LogicalPlan(plan) => Ok(plan.as_ref().clone()), CachedRelationData::Materialized(materialized) => { materialized.to_logical_plan(relation_id) } @@ -305,9 +289,6 @@ impl CachedRelation { let data = self.data.lock().await; match &*data { CachedRelationData::Materialized(materialized) => materialized.to_physical_plan().await, - CachedRelationData::LogicalPlan(_) => Err(internal_datafusion_err!( - "cached relation is not materialized as a physical plan" - )), CachedRelationData::Pending(pending) => Ok(Arc::new(PendingCachedRelationExec::new( relation_id.to_string(), self.clone(), @@ -323,16 +304,13 @@ impl CachedRelation { } match &*data { CachedRelationData::Materialized(materialized) => materialized.to_physical_plan().await, - CachedRelationData::LogicalPlan(_) => Err(internal_datafusion_err!( - "cached relation is not materialized as a physical plan" - )), CachedRelationData::Pending(_) => Err(internal_datafusion_err!( "cached relation materialization did not complete" )), } } - pub fn into_cleanup(self) -> Option { + fn into_cleanup(self) -> Option { self.cleanup } } @@ -393,13 +371,13 @@ impl SessionExtension for CachedRelationRegistry { impl CachedRelationPending { async fn materialize(&self, ctx: &SessionContext) -> Result { - let physical_plan = materialize_cached_relations(ctx, Arc::clone(&self.plan)).await?; match &self.target { CachedRelationPendingTarget::Local { storage_level } => { - materialize_local_checkpoint(ctx, physical_plan, storage_level.clone()).await + materialize_local_checkpoint(ctx, Arc::clone(&self.plan), storage_level.clone()) + .await } CachedRelationPendingTarget::Reliable { path } => { - materialize_reliable_checkpoint(ctx, physical_plan, path).await + materialize_reliable_checkpoint(ctx, Arc::clone(&self.plan), path).await } } } @@ -511,12 +489,10 @@ impl CachedRelationLocalMaterialized { memory_partitions, serialized_memory_partitions, disk_partitions, - storage_level: Some(storage_level), }) } async fn to_physical_plan(&self) -> Result> { - let _ = &self.storage_level; let partitions = self.load_partitions().await?; let plan: Arc = MemorySourceConfig::try_new_exec(&partitions, Arc::clone(&self.schema), None)?; diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 1da0352134..973a9227ad 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -11,8 +11,7 @@ use futures::stream; use log::{debug, warn}; use sail_common::spec; use sail_common_datafusion::cached_relation::{ - cleanup_cached_relation, cleanup_checkpoint_path as cleanup_cached_checkpoint_path, - CachedRelation, CachedRelationRegistry, + cleanup_cached_relation, cleanup_checkpoint_path, CachedRelation, CachedRelationRegistry, }; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; @@ -720,11 +719,6 @@ async fn cleanup_checkpoint_after_error( } } -async fn cleanup_checkpoint_path(ctx: &SessionContext, path: &str) -> SparkResult<()> { - cleanup_cached_checkpoint_path(ctx, path).await?; - Ok(()) -} - 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/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index af4fdc6bb5..4f3951f028 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -73,7 +73,7 @@ def test_dataframe_checkpoint(spark, tmp_path): spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) try: checkpointed = df.checkpoint() - assert list(checkpoint_path.rglob("*.arrow")) + assert any(checkpoint_path.rglob("*.arrow")) shutil.rmtree(source_path) assert_frame_equal( @@ -175,9 +175,9 @@ def test_dataframe_checkpoint_lazy(spark, tmp_path): spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) try: checkpointed = df.checkpoint(eager=False) - assert not list(checkpoint_path.rglob("*.arrow")) + assert not any(checkpoint_path.rglob("*.arrow")) assert checkpointed.count() == 2 # noqa: PLR2004 - assert list(checkpoint_path.rglob("*.arrow")) + assert any(checkpoint_path.rglob("*.arrow")) shutil.rmtree(source_path) assert_frame_equal( @@ -202,10 +202,10 @@ def test_dataframe_checkpoint_lazy_explain_does_not_materialize(spark, tmp_path) checkpointed_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 assert "PendingCachedRelationExec" in checkpointed_plan - assert not list(checkpoint_path.rglob("*.arrow")) + assert not any(checkpoint_path.rglob("*.arrow")) assert checkpointed.count() == 3 # noqa: PLR2004 - assert list(checkpoint_path.rglob("*.arrow")) + assert any(checkpoint_path.rglob("*.arrow")) materialized_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 assert "PendingCachedRelationExec" not in materialized_plan @@ -221,13 +221,13 @@ def test_dataframe_checkpoint_cleanup_on_session_stop(spark_session_factory, tmp spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) _checkpointed = spark.range(0, 5).checkpoint() - assert list(checkpoint_path.rglob("*.arrow")) + assert any(checkpoint_path.rglob("*.arrow")) spark.stop() deadline = time.monotonic() + 5 - while time.monotonic() < deadline and list(checkpoint_path.rglob("*.arrow")): + while time.monotonic() < deadline and any(checkpoint_path.rglob("*.arrow")): time.sleep(0.1) - assert not list(checkpoint_path.rglob("*.arrow")) + assert not any(checkpoint_path.rglob("*.arrow")) @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint partition preservation") From 54166c22b42a6abb02fca7b59c15273ad7046edf Mon Sep 17 00:00:00 2001 From: XL Liang Date: Sun, 21 Jun 2026 23:04:17 +0800 Subject: [PATCH 11/40] test: skip checkpoint tests on unsupported PySpark --- python/pysail/tests/spark/dataframe/test_checkpoint.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/pysail/tests/spark/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index 4f3951f028..2c7faa3c7f 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -8,7 +8,12 @@ from pyspark.sql.functions import col, lit from pysail.testing.spark.steps.plan import normalize_plan_text -from pysail.testing.spark.utils.common import is_jvm_spark +from pysail.testing.spark.utils.common import is_jvm_spark, pyspark_version + +pytestmark = pytest.mark.skipif( + pyspark_version() < (4,), + reason="checkpoint and localCheckpoint require PySpark Connect 4+", +) def test_dataframe_local_checkpoint_survives_source_removal(spark, tmp_path): From 3c9908ce7dec11ed374e010957880fbea75dfbdc Mon Sep 17 00:00:00 2001 From: XL Liang Date: Sun, 21 Jun 2026 23:24:33 +0800 Subject: [PATCH 12/40] test: skip JVM-specific checkpoint parity cases --- scripts/spark-tests/plugins/spark.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/scripts/spark-tests/plugins/spark.py b/scripts/spark-tests/plugins/spark.py index 1eb419d5d8..4d0df614f5 100644 --- a/scripts/spark-tests/plugins/spark.py +++ b/scripts/spark-tests/plugins/spark.py @@ -340,6 +340,18 @@ def _spark_major_version() -> int: keywords=["pyspark.sql.functions.java_method"], reason="JVM-dependent test", ), + TestMarker( + keywords=["test_parity_dataframe.py", "test_local_checkpoint_dataframe"], + reason="JVM physical plan assertion; checkpoint behavior is covered by PySail tests", + ), + TestMarker( + keywords=["test_connect_basic.py", "test_garbage_collection_checkpoint"], + reason="JVM Spark Connect dataFrameCache internals", + ), + TestMarker( + keywords=["test_connect_basic.py", "test_garbage_collection_derived_checkpoint"], + reason="JVM Spark Connect dataFrameCache internals", + ), # We skip all the streaming tests since some of them are slow, # and some of them test behaviors that are tied to the specific JVM implementation # of Spark Structured Streaming. From 7ad572a419c5a0e33b2053a0bdd94424f06a21f1 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Thu, 25 Jun 2026 11:06:13 +0800 Subject: [PATCH 13/40] fix: enable checkpoint execution on workers --- .../sail-common-datafusion/src/array/record_batch.rs | 12 +++++++++++- crates/sail-common-datafusion/src/cached_relation.rs | 2 +- crates/sail-session/src/checkpoint.rs | 9 ++++++--- crates/sail-session/src/session_factory/worker.rs | 7 +++++++ .../pysail/tests/spark/dataframe/test_checkpoint.py | 12 ++++++++++++ 5 files changed, 37 insertions(+), 5 deletions(-) diff --git a/crates/sail-common-datafusion/src/array/record_batch.rs b/crates/sail-common-datafusion/src/array/record_batch.rs index 8a528766cb..03fbf505f2 100644 --- a/crates/sail-common-datafusion/src/array/record_batch.rs +++ b/crates/sail-common-datafusion/src/array/record_batch.rs @@ -12,7 +12,7 @@ use datafusion::arrow::datatypes::{ TimestampSecondType, }; use datafusion::arrow::ipc::reader::StreamReader; -use datafusion::arrow::ipc::writer::StreamWriter; +use datafusion::arrow::ipc::writer::{FileWriter, StreamWriter}; use datafusion_common::{DataFusionError, Result}; pub fn cast_record_batch_positionally( @@ -373,6 +373,16 @@ pub fn write_record_batches(batches: &[RecordBatch], schema: &Schema) -> Result< Ok(output) } +pub fn write_record_batches_file(batches: &[RecordBatch], schema: &Schema) -> Result> { + let mut output = Vec::new(); + let mut writer = FileWriter::try_new(&mut output, schema)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + Ok(output) +} + pub fn record_batch_with_schema(batch: RecordBatch, schema: &SchemaRef) -> Result { Ok(RecordBatch::try_new_with_options( schema.clone(), diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index 952c22313d..42d9738b37 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -607,7 +607,7 @@ fn create_arrow_checkpoint_scan( "reliable checkpoint did not produce any files" )); } - let source = ArrowSource::new_stream_file_source(TableSchema::new(Arc::clone(schema), vec![])); + let source = ArrowSource::new_file_source(TableSchema::new(Arc::clone(schema), vec![])); let file_groups = object_meta .into_iter() .map(|object_meta| FileGroup::new(vec![PartitionedFile::new_from_meta(object_meta)])) diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs index dcb932dbfb..a662a083b0 100644 --- a/crates/sail-session/src/checkpoint.rs +++ b/crates/sail-session/src/checkpoint.rs @@ -6,7 +6,7 @@ use datafusion::physical_plan::{execute_stream_partitioned, ExecutionPlan}; use datafusion::prelude::SessionContext; use datafusion_common::Result; use futures::StreamExt; -use sail_common_datafusion::array::record_batch::write_record_batches; +use sail_common_datafusion::array::record_batch::write_record_batches_file; use sail_common_datafusion::session::checkpoint::{CheckpointStore, ReliableCheckpoint}; use sail_object_store::{delete_object_store_prefix, resolve_object_store_path}; @@ -32,13 +32,16 @@ impl CheckpointStore for ObjectStoreCheckpointStore { while let Some(batch) = stream.next().await { batches.push(batch?); } - let bytes = write_record_batches(&batches, schema.as_ref())?; + if batches.iter().all(|batch| batch.num_rows() == 0) { + continue; + } + let bytes = write_record_batches_file(&batches, schema.as_ref())?; let file_path = checkpoint_path.child(&format!("part-{file_index:05}.arrow")); files.push(checkpoint_path.put_bytes(&file_path, bytes).await?); } if files.is_empty() { - let bytes = write_record_batches(&[], schema.as_ref())?; + let bytes = write_record_batches_file(&[], schema.as_ref())?; let file_path = checkpoint_path.child("part-00000.arrow"); files.push(checkpoint_path.put_bytes(&file_path, bytes).await?); } diff --git a/crates/sail-session/src/session_factory/worker.rs b/crates/sail-session/src/session_factory/worker.rs index c608f28388..eacfb7bbce 100644 --- a/crates/sail-session/src/session_factory/worker.rs +++ b/crates/sail-session/src/session_factory/worker.rs @@ -5,9 +5,12 @@ use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{SessionConfig, SessionContext}; use sail_common::config::AppConfig; use sail_common::runtime::RuntimeHandle; +use sail_common_datafusion::cached_relation::CachedRelationRegistry; +use sail_common_datafusion::session::checkpoint::CheckpointStoreService; use sail_common_datafusion::session::repartition::RepartitionBufferConfig; use sail_delta_lake::session_extension::DeltaTableCache; +use crate::checkpoint::ObjectStoreCheckpointStore; use crate::runtime::RuntimeEnvFactory; use crate::session_factory::SessionFactory; @@ -35,6 +38,10 @@ impl SessionFactory<()> for WorkerSessionFactory { // when decoding the execution plan. let config = SessionConfig::default() .with_extension(Arc::new(DeltaTableCache::default())) + .with_extension(Arc::new(CachedRelationRegistry::default())) + .with_extension(Arc::new(CheckpointStoreService::new(Arc::new( + ObjectStoreCheckpointStore, + )))) .with_extension(Arc::new(RepartitionBufferConfig::new( self.repartition_buffer_size, ))); diff --git a/python/pysail/tests/spark/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index 2c7faa3c7f..fdb4a9a3c6 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -1,3 +1,4 @@ +import os import shutil import time @@ -7,6 +8,7 @@ from pyspark import StorageLevel from pyspark.sql.functions import col, lit +from pysail.testing.spark.session import spark_session_factory from pysail.testing.spark.steps.plan import normalize_plan_text from pysail.testing.spark.utils.common import is_jvm_spark, pyspark_version @@ -16,6 +18,12 @@ ) +@pytest.fixture(name="spark_session_factory") +def spark_session_factory_fixture(remote): + with spark_session_factory(remote) as sessions: + yield sessions.create + + def test_dataframe_local_checkpoint_survives_source_removal(spark, tmp_path): source_path = tmp_path / "source" spark.createDataFrame( @@ -94,6 +102,10 @@ def test_dataframe_checkpoint(spark, tmp_path): @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific object-store checkpoint URL") +@pytest.mark.skipif( + os.environ.get("SAIL_MODE") == "local-cluster", + reason="memory object stores are scoped to one runtime and are not shared with workers", +) def test_dataframe_checkpoint_with_memory_object_store(spark, tmp_path): source_path = tmp_path / "source" spark.createDataFrame( From 4b885a6fb6451f660b99e622db4a4acbff02dd04 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Wed, 8 Jul 2026 22:25:40 +0800 Subject: [PATCH 14/40] tmp --- .../src/cached_relation.rs | 80 +++++++++++++------ crates/sail-plan/src/resolver/query/mod.rs | 30 ++++--- 2 files changed, 74 insertions(+), 36 deletions(-) diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index 42d9738b37..cd493dbcb8 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -103,6 +103,7 @@ impl UserDefinedLogicalNodeCore for CachedRelationNode { #[derive(Debug, Clone)] pub struct CachedRelation { + schema: SchemaRef, data: Arc>, cleanup: Option, } @@ -229,7 +230,9 @@ impl CachedRelation { storage_level: spec::StorageLevel, ) -> Result { let data = materialize_local_checkpoint(ctx, plan, storage_level).await?; + let schema = data.schema(); Ok(Self { + schema, data: Arc::new(tokio::sync::Mutex::new(data)), cleanup: None, }) @@ -241,7 +244,9 @@ impl CachedRelation { path: &str, ) -> Result { let data = materialize_reliable_checkpoint(ctx, plan, path).await?; + let schema = data.schema(); Ok(Self { + schema, data: Arc::new(tokio::sync::Mutex::new(data)), cleanup: Some(CachedRelationCleanup::ObjectStorePath(path.to_string())), }) @@ -251,7 +256,9 @@ impl CachedRelation { plan: Arc, storage_level: spec::StorageLevel, ) -> Self { + let schema = plan.schema(); Self { + schema, data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( CachedRelationPending { plan, @@ -264,7 +271,9 @@ impl CachedRelation { pub fn new_pending_reliable_checkpoint(plan: Arc, path: String) -> Self { let cleanup_path = path.clone(); + let schema = plan.schema(); Self { + schema, data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( CachedRelationPending { plan, @@ -275,14 +284,13 @@ impl CachedRelation { } } - pub async fn to_logical_plan(&self, relation_id: &str) -> Result { - let data = self.data.lock().await; - match &*data { - CachedRelationData::Materialized(materialized) => { - materialized.to_logical_plan(relation_id) - } - CachedRelationData::Pending(pending) => pending.to_logical_plan(relation_id), - } + pub fn to_logical_plan(&self, relation_id: &str) -> Result { + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(CachedRelationNode::try_new( + relation_id.to_string(), + Arc::clone(&self.schema), + )?), + })) } pub async fn to_physical_plan(&self, relation_id: &str) -> Result> { @@ -369,6 +377,15 @@ impl SessionExtension for CachedRelationRegistry { } } +impl CachedRelationData { + fn schema(&self) -> SchemaRef { + match self { + Self::Materialized(materialized) => materialized.schema(), + Self::Pending(pending) => pending.schema(), + } + } +} + impl CachedRelationPending { async fn materialize(&self, ctx: &SessionContext) -> Result { match &self.target { @@ -382,13 +399,8 @@ impl CachedRelationPending { } } - fn to_logical_plan(&self, relation_id: &str) -> Result { - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(CachedRelationNode::try_new( - relation_id.to_string(), - self.plan.schema(), - )?), - })) + fn schema(&self) -> SchemaRef { + self.plan.schema() } } @@ -400,15 +412,6 @@ impl CachedRelationMaterialized { } } - fn to_logical_plan(&self, relation_id: &str) -> Result { - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(CachedRelationNode::try_new( - relation_id.to_string(), - self.schema(), - )?), - })) - } - async fn to_physical_plan(&self) -> Result> { match self { Self::Local(local) => local.to_physical_plan().await, @@ -665,7 +668,7 @@ mod tests { use std::sync::Mutex; use async_trait::async_trait; - use datafusion::arrow::datatypes::Schema; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_plan::empty::EmptyExec; use super::*; @@ -719,6 +722,33 @@ mod tests { Ok(()) } + #[test] + fn pending_cached_relation_logical_plan_uses_plan_schema() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + true, + )])); + let plan: Arc = Arc::new(EmptyExec::new(schema)); + let relation = CachedRelation::new_pending_local_checkpoint( + plan, + spec::StorageLevel { + use_disk: false, + use_memory: true, + use_off_heap: false, + deserialized: true, + replication: 1, + }, + ); + + let plan = relation.to_logical_plan("relation")?; + let fields = plan.schema().fields(); + + assert_eq!(fields.len(), 1); + assert_eq!(fields[0].name(), "value"); + Ok(()) + } + #[derive(Default)] struct TestCheckpointStore { cleanup_paths: Mutex>, diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index b313b17913..4524a46caf 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/crates/sail-plan/src/resolver/query/mod.rs @@ -82,8 +82,8 @@ impl PlanResolver<'_> { } }, QueryNode::Project { input, expressions } => { - self.resolve_query_project(input.map(|x| *x), expressions, state) - .await? + // Keep large helper futures out of this recursive dispatcher frame. + Box::pin(self.resolve_query_project(input.map(|x| *x), expressions, state)).await? } QueryNode::Filter { input, condition } => { self.resolve_query_filter(*input, condition, state).await? @@ -183,7 +183,7 @@ impl PlanResolver<'_> { self.resolve_query_hint(*input, name, parameters, state) .await? } - QueryNode::Pivot(pivot) => self.resolve_query_pivot(pivot, state).await?, + QueryNode::Pivot(pivot) => Box::pin(self.resolve_query_pivot(pivot, state)).await?, QueryNode::Unpivot(unpivot) => self.resolve_query_unpivot(unpivot, state).await?, QueryNode::ToSchema { input, schema } => { self.resolve_query_to_schema(*input, schema, state).await? @@ -231,13 +231,7 @@ impl PlanResolver<'_> { return Err(PlanError::todo("cached local relation")); } QueryNode::CachedRemoteRelation { relation_id } => { - let registry = self.ctx.extension::()?; - let relation = registry.get(&relation_id)?.ok_or_else(|| { - PlanError::invalid(format!("No DataFrame with id {relation_id} is found")) - })?; - let plan = relation.to_logical_plan(&relation_id).await?; - let names = state.register_fields(plan.schema().inner().fields()); - rename_logical_plan(plan, &names)? + self.resolve_query_cached_remote_relation(relation_id, state)? } QueryNode::CommonInlineUserDefinedTableFunction(udtf) => { self.resolve_query_common_inline_udtf(udtf, state).await? @@ -315,7 +309,7 @@ impl PlanResolver<'_> { .await? } QueryNode::Empty { produce_one_row } => self.resolve_query_empty(produce_one_row)?, - QueryNode::Values(values) => self.resolve_query_values(values, state).await?, + QueryNode::Values(values) => Box::pin(self.resolve_query_values(values, state)).await?, QueryNode::TableAlias { input, name, @@ -375,6 +369,20 @@ impl PlanResolver<'_> { Ok(plan) } + fn resolve_query_cached_remote_relation( + &self, + relation_id: String, + state: &mut PlanResolverState, + ) -> PlanResult { + let registry = self.ctx.extension::()?; + let relation = registry.get(&relation_id)?.ok_or_else(|| { + PlanError::invalid(format!("No DataFrame with id {relation_id} is found")) + })?; + let plan = relation.to_logical_plan(&relation_id)?; + let names = state.register_fields(plan.schema().inner().fields()); + rename_logical_plan(plan, &names).map_err(PlanError::from) + } + fn remove_hidden_fields( &self, plan: LogicalPlan, From fd7129358d4874d85e4189060710bad2601dedcc Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 10 Jul 2026 20:03:51 +0800 Subject: [PATCH 15/40] ci: increase Rust test stack --- .github/workflows/rust-tests.yml | 1 + .../src/cached_relation.rs | 80 ++++++------------- crates/sail-plan/src/resolver/query/mod.rs | 30 +++---- 3 files changed, 37 insertions(+), 74 deletions(-) diff --git a/.github/workflows/rust-tests.yml b/.github/workflows/rust-tests.yml index 114dc91ce3..43effbb70c 100644 --- a/.github/workflows/rust-tests.yml +++ b/.github/workflows/rust-tests.yml @@ -19,6 +19,7 @@ jobs: # Coverage environment variables RUSTC_WORKSPACE_WRAPPER: ${{ github.workspace }}/.github/scripts/rustc-workspace-wrapper.sh LLVM_PROFILE_FILE: ${{ github.workspace }}/target/llvm-profiles/rust-unit/sail-%p-%m.profraw + RUST_MIN_STACK: 8388608 steps: - uses: actions/checkout@v7 diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index cd493dbcb8..42d9738b37 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -103,7 +103,6 @@ impl UserDefinedLogicalNodeCore for CachedRelationNode { #[derive(Debug, Clone)] pub struct CachedRelation { - schema: SchemaRef, data: Arc>, cleanup: Option, } @@ -230,9 +229,7 @@ impl CachedRelation { storage_level: spec::StorageLevel, ) -> Result { let data = materialize_local_checkpoint(ctx, plan, storage_level).await?; - let schema = data.schema(); Ok(Self { - schema, data: Arc::new(tokio::sync::Mutex::new(data)), cleanup: None, }) @@ -244,9 +241,7 @@ impl CachedRelation { path: &str, ) -> Result { let data = materialize_reliable_checkpoint(ctx, plan, path).await?; - let schema = data.schema(); Ok(Self { - schema, data: Arc::new(tokio::sync::Mutex::new(data)), cleanup: Some(CachedRelationCleanup::ObjectStorePath(path.to_string())), }) @@ -256,9 +251,7 @@ impl CachedRelation { plan: Arc, storage_level: spec::StorageLevel, ) -> Self { - let schema = plan.schema(); Self { - schema, data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( CachedRelationPending { plan, @@ -271,9 +264,7 @@ impl CachedRelation { pub fn new_pending_reliable_checkpoint(plan: Arc, path: String) -> Self { let cleanup_path = path.clone(); - let schema = plan.schema(); Self { - schema, data: Arc::new(tokio::sync::Mutex::new(CachedRelationData::Pending( CachedRelationPending { plan, @@ -284,13 +275,14 @@ impl CachedRelation { } } - pub fn to_logical_plan(&self, relation_id: &str) -> Result { - Ok(LogicalPlan::Extension(Extension { - node: Arc::new(CachedRelationNode::try_new( - relation_id.to_string(), - Arc::clone(&self.schema), - )?), - })) + pub async fn to_logical_plan(&self, relation_id: &str) -> Result { + let data = self.data.lock().await; + match &*data { + CachedRelationData::Materialized(materialized) => { + materialized.to_logical_plan(relation_id) + } + CachedRelationData::Pending(pending) => pending.to_logical_plan(relation_id), + } } pub async fn to_physical_plan(&self, relation_id: &str) -> Result> { @@ -377,15 +369,6 @@ impl SessionExtension for CachedRelationRegistry { } } -impl CachedRelationData { - fn schema(&self) -> SchemaRef { - match self { - Self::Materialized(materialized) => materialized.schema(), - Self::Pending(pending) => pending.schema(), - } - } -} - impl CachedRelationPending { async fn materialize(&self, ctx: &SessionContext) -> Result { match &self.target { @@ -399,8 +382,13 @@ impl CachedRelationPending { } } - fn schema(&self) -> SchemaRef { - self.plan.schema() + fn to_logical_plan(&self, relation_id: &str) -> Result { + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(CachedRelationNode::try_new( + relation_id.to_string(), + self.plan.schema(), + )?), + })) } } @@ -412,6 +400,15 @@ impl CachedRelationMaterialized { } } + fn to_logical_plan(&self, relation_id: &str) -> Result { + Ok(LogicalPlan::Extension(Extension { + node: Arc::new(CachedRelationNode::try_new( + relation_id.to_string(), + self.schema(), + )?), + })) + } + async fn to_physical_plan(&self) -> Result> { match self { Self::Local(local) => local.to_physical_plan().await, @@ -668,7 +665,7 @@ mod tests { use std::sync::Mutex; use async_trait::async_trait; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::datatypes::Schema; use datafusion::physical_plan::empty::EmptyExec; use super::*; @@ -722,33 +719,6 @@ mod tests { Ok(()) } - #[test] - fn pending_cached_relation_logical_plan_uses_plan_schema() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new( - "value", - DataType::Int32, - true, - )])); - let plan: Arc = Arc::new(EmptyExec::new(schema)); - let relation = CachedRelation::new_pending_local_checkpoint( - plan, - spec::StorageLevel { - use_disk: false, - use_memory: true, - use_off_heap: false, - deserialized: true, - replication: 1, - }, - ); - - let plan = relation.to_logical_plan("relation")?; - let fields = plan.schema().fields(); - - assert_eq!(fields.len(), 1); - assert_eq!(fields[0].name(), "value"); - Ok(()) - } - #[derive(Default)] struct TestCheckpointStore { cleanup_paths: Mutex>, diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index d43b9533cf..226296f025 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/crates/sail-plan/src/resolver/query/mod.rs @@ -82,8 +82,8 @@ impl PlanResolver<'_> { } }, QueryNode::Project { input, expressions } => { - // Keep large helper futures out of this recursive dispatcher frame. - Box::pin(self.resolve_query_project(input.map(|x| *x), expressions, state)).await? + self.resolve_query_project(input.map(|x| *x), expressions, state) + .await? } QueryNode::Filter { input, condition } => { self.resolve_query_filter(*input, condition, state).await? @@ -183,7 +183,7 @@ impl PlanResolver<'_> { self.resolve_query_hint(*input, name, parameters, state) .await? } - QueryNode::Pivot(pivot) => Box::pin(self.resolve_query_pivot(pivot, state)).await?, + QueryNode::Pivot(pivot) => self.resolve_query_pivot(pivot, state).await?, QueryNode::Unpivot(unpivot) => self.resolve_query_unpivot(unpivot, state).await?, QueryNode::ToSchema { input, schema } => { self.resolve_query_to_schema(*input, schema, state).await? @@ -231,7 +231,13 @@ impl PlanResolver<'_> { return Err(PlanError::todo("cached local relation")); } QueryNode::CachedRemoteRelation { relation_id } => { - self.resolve_query_cached_remote_relation(relation_id, state)? + let registry = self.ctx.extension::()?; + let relation = registry.get(&relation_id)?.ok_or_else(|| { + PlanError::invalid(format!("No DataFrame with id {relation_id} is found")) + })?; + let plan = relation.to_logical_plan(&relation_id).await?; + let names = state.register_fields(plan.schema().inner().fields()); + rename_logical_plan(plan, &names)? } QueryNode::CommonInlineUserDefinedTableFunction(udtf) => { self.resolve_query_common_inline_udtf(udtf, state).await? @@ -309,7 +315,7 @@ impl PlanResolver<'_> { .await? } QueryNode::Empty { produce_one_row } => self.resolve_query_empty(produce_one_row)?, - QueryNode::Values(values) => Box::pin(self.resolve_query_values(values, state)).await?, + QueryNode::Values(values) => self.resolve_query_values(values, state).await?, QueryNode::TableAlias { input, name, @@ -369,20 +375,6 @@ impl PlanResolver<'_> { Ok(plan) } - fn resolve_query_cached_remote_relation( - &self, - relation_id: String, - state: &mut PlanResolverState, - ) -> PlanResult { - let registry = self.ctx.extension::()?; - let relation = registry.get(&relation_id)?.ok_or_else(|| { - PlanError::invalid(format!("No DataFrame with id {relation_id} is found")) - })?; - let plan = relation.to_logical_plan(&relation_id)?; - let names = state.register_fields(plan.schema().inner().fields()); - rename_logical_plan(plan, &names).map_err(PlanError::from) - } - fn remove_hidden_fields( &self, plan: LogicalPlan, From c1489fd598b86b51e705bea960aa853f9fc7e722 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 10 Jul 2026 20:04:22 +0800 Subject: [PATCH 16/40] style: format Rust imports --- crates/sail-common-datafusion/src/cached_relation.rs | 8 ++++---- crates/sail-object-store/src/lib.rs | 2 +- crates/sail-session/src/checkpoint.rs | 4 ++-- crates/sail-spark-connect/src/service/plan_executor.rs | 2 +- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index 42d9738b37..470d16ca1c 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -1,6 +1,6 @@ use std::cmp::Ordering; -use std::collections::hash_map::Entry; use std::collections::HashMap; +use std::collections::hash_map::Entry; use std::fmt::Formatter; use std::sync::{Arc, RwLock}; @@ -13,12 +13,12 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::{ - collect_partitioned, with_new_children_if_necessary, DisplayAs, DisplayFormatType, - ExecutionPlan, ExecutionPlanProperties, PlanProperties, + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, + collect_partitioned, with_new_children_if_necessary, }; use datafusion::prelude::SessionContext; use datafusion_common::{ - internal_datafusion_err, DFSchema, DFSchemaRef, DataFusionError, Result, Statistics, + DFSchema, DFSchemaRef, DataFusionError, Result, Statistics, internal_datafusion_err, }; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; diff --git a/crates/sail-object-store/src/lib.rs b/crates/sail-object-store/src/lib.rs index 31c2ccf367..afd3d3fafe 100644 --- a/crates/sail-object-store/src/lib.rs +++ b/crates/sail-object-store/src/lib.rs @@ -4,5 +4,5 @@ mod path; mod registry; mod s3; -pub use path::{delete_object_store_prefix, resolve_object_store_path, ResolvedObjectStorePath}; +pub use path::{ResolvedObjectStorePath, delete_object_store_prefix, resolve_object_store_path}; pub use registry::DynamicObjectStoreRegistry; diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs index a662a083b0..b4f72ec8c8 100644 --- a/crates/sail-session/src/checkpoint.rs +++ b/crates/sail-session/src/checkpoint.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::physical_plan::{execute_stream_partitioned, ExecutionPlan}; +use datafusion::physical_plan::{ExecutionPlan, execute_stream_partitioned}; use datafusion::prelude::SessionContext; use datafusion_common::Result; use futures::StreamExt; @@ -62,7 +62,7 @@ impl CheckpointStore for ObjectStoreCheckpointStore { mod tests { use std::time::{SystemTime, UNIX_EPOCH}; - use datafusion_common::{internal_datafusion_err, DataFusionError}; + use datafusion_common::{DataFusionError, internal_datafusion_err}; use super::*; diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index fab7aefe4f..88ced5ea06 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -11,7 +11,7 @@ use futures::stream; use log::{debug, warn}; use sail_common::spec; use sail_common_datafusion::cached_relation::{ - cleanup_cached_relation, cleanup_checkpoint_path, CachedRelation, CachedRelationRegistry, + CachedRelation, CachedRelationRegistry, cleanup_cached_relation, cleanup_checkpoint_path, }; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; From 15d0cb5107e9f2b58159f502b240df74228713c9 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 10 Jul 2026 20:06:33 +0800 Subject: [PATCH 17/40] test: make checkpoint plans cluster-compatible --- python/pysail/tests/spark/dataframe/test_checkpoint.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/pysail/tests/spark/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index fdb4a9a3c6..536edb33d0 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -69,7 +69,6 @@ def test_dataframe_local_checkpoint_explain_truncates_plan(spark): assert "RangeExec" not in checkpointed_plan assert checkpointed_plan.count("RangeExec") == 0 assert checkpointed_plan.count("FilterExec") < 1 - assert checkpointed_plan.count("ProjectionExec") < 3 # noqa: PLR2004 assert "DataSourceExec" in checkpointed_plan assert checkpointed_plan.count("DataSourceExec") <= 2 # noqa: PLR2004 @@ -278,7 +277,6 @@ def test_dataframe_local_checkpoint_deep_plan_explain_truncates(spark): assert "RangeExec" not in checkpointed_plan assert checkpointed_plan.count("RangeExec") == 0 assert checkpointed_plan.count("FilterExec") < 1 - assert checkpointed_plan.count("ProjectionExec") < 3 # noqa: PLR2004 assert "DataSourceExec" in checkpointed_plan assert checkpointed_plan.count("DataSourceExec") <= 2 # noqa: PLR2004 From d83b4c3b34608f8571452740427c19fca30ca0d8 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 10 Jul 2026 20:32:31 +0800 Subject: [PATCH 18/40] fix: preserve checkpoint lifecycle semantics --- .../src/cached_relation.rs | 167 +++++++++++++++--- crates/sail-session/src/checkpoint.rs | 3 - .../src/service/plan_executor.rs | 9 +- .../tests/spark/dataframe/test_checkpoint.py | 72 +++++++- 4 files changed, 217 insertions(+), 34 deletions(-) diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index 470d16ca1c..db62ba23dc 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -309,10 +309,6 @@ impl CachedRelation { )), } } - - fn into_cleanup(self) -> Option { - self.cleanup - } } #[derive(Debug, Default)] @@ -354,12 +350,12 @@ impl CachedRelationRegistry { Ok(relations.remove(relation_id)) } - pub fn drain(&self) -> Result> { + pub fn drain(&self) -> Result> { let mut relations = self .relations .write() .map_err(|e| internal_datafusion_err!("{e}"))?; - Ok(relations.drain().map(|(_, relation)| relation).collect()) + Ok(relations.drain().collect()) } } @@ -614,6 +610,7 @@ fn create_arrow_checkpoint_scan( .collect(); let config = FileScanConfigBuilder::new(object_store_url, Arc::new(source)) .with_file_groups(file_groups) + .with_partitioned_by_file_group(true) .with_statistics(Statistics::new_unknown(schema)) .build(); Ok(DataSourceExec::from_data_source(config)) @@ -624,19 +621,55 @@ pub async fn cleanup_checkpoint_path(ctx: &SessionContext, path: &str) -> Result service.cleanup_checkpoint_path(ctx, path).await } -pub async fn cleanup_cached_relation(ctx: &SessionContext, relation: CachedRelation) -> Result<()> { - match relation.into_cleanup() { +pub async fn cleanup_cached_relation( + ctx: &SessionContext, + relation: &CachedRelation, +) -> Result<()> { + match relation.cleanup.as_ref() { Some(CachedRelationCleanup::ObjectStorePath(path)) => { - cleanup_checkpoint_path(ctx, &path).await + cleanup_checkpoint_path(ctx, path).await } None => Ok(()), } } +pub async fn remove_cached_relation(ctx: &SessionContext, relation_id: &str) -> Result<()> { + let registry = ctx.extension::()?; + let Some(relation) = registry.remove(relation_id)? else { + return Ok(()); + }; + if let Err(cleanup_error) = cleanup_cached_relation(ctx, &relation).await { + if let Err(restore_error) = registry.insert(relation_id.to_string(), relation) { + return Err(internal_datafusion_err!( + "failed to clean cached relation {relation_id}: {cleanup_error}; additionally failed to restore it: {restore_error}" + )); + } + return Err(cleanup_error); + } + Ok(()) +} + pub async fn cleanup_cached_relations(ctx: &SessionContext) -> Result<()> { - let relations = ctx.extension::()?.drain()?; - for relation in relations { - cleanup_cached_relation(ctx, relation).await?; + let registry = ctx.extension::()?; + let relations = registry.drain()?; + let mut errors = vec![]; + for (relation_id, relation) in relations { + if let Err(cleanup_error) = cleanup_cached_relation(ctx, &relation).await { + if let Err(restore_error) = registry.insert(relation_id.clone(), relation) { + errors.push(format!( + "{relation_id}: {cleanup_error}; additionally failed to restore it: {restore_error}" + )); + } else { + errors.push(format!("{relation_id}: {cleanup_error}")); + } + } + } + if !errors.is_empty() { + return Err(internal_datafusion_err!( + "failed to clean {} cached relation(s): {}", + errors.len(), + errors.join("; ") + )); } Ok(()) } @@ -662,6 +695,7 @@ async fn write_disk_partition(ctx: &SessionContext, bytes: &[u8]) -> Result = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); let relation = CachedRelation::new_pending_reliable_checkpoint(plan, path.clone()); - match relation.into_cleanup() { + match relation.cleanup.as_ref() { Some(CachedRelationCleanup::ObjectStorePath(actual)) => { - assert_eq!(actual, path); + assert_eq!(actual, &path); } other => { return Err(internal_datafusion_err!( @@ -706,9 +740,9 @@ mod tests { let relation = registry.get("relation")?.ok_or_else(|| { internal_datafusion_err!("cached relation missing after duplicate insert") })?; - match relation.into_cleanup() { + match relation.cleanup.as_ref() { Some(CachedRelationCleanup::ObjectStorePath(actual)) => { - assert_eq!(actual, first_path); + assert_eq!(actual, &first_path); } other => { return Err(internal_datafusion_err!( @@ -722,6 +756,22 @@ mod tests { #[derive(Default)] struct TestCheckpointStore { cleanup_paths: Mutex>, + cleanup_failures: Mutex>, + } + + impl TestCheckpointStore { + fn set_cleanup_failure(&self, path: &str, fail: bool) -> Result<()> { + let mut failures = self + .cleanup_failures + .lock() + .map_err(|e| internal_datafusion_err!("{e}"))?; + if fail { + failures.insert(path.to_string()); + } else { + failures.remove(path); + } + Ok(()) + } } #[async_trait] @@ -739,20 +789,36 @@ mod tests { } async fn cleanup_checkpoint_path(&self, _ctx: &SessionContext, path: &str) -> Result<()> { + let should_fail = self + .cleanup_failures + .lock() + .map_err(|e| internal_datafusion_err!("{e}"))? + .contains(path); self.cleanup_paths .lock() .map_err(|e| internal_datafusion_err!("{e}"))? .push(path.to_string()); - Ok(()) + if should_fail { + Err(internal_datafusion_err!( + "failed to clean checkpoint path {path}" + )) + } else { + Ok(()) + } } } + fn checkpoint_context(store: Arc) -> SessionContext { + let config = datafusion::prelude::SessionConfig::new() + .with_extension(Arc::new(CheckpointStoreService::new(store))) + .with_extension(Arc::new(CachedRelationRegistry::default())); + SessionContext::new_with_config(config) + } + #[tokio::test] async fn cleanup_checkpoint_path_delegates_to_store() -> Result<()> { let store = Arc::new(TestCheckpointStore::default()); - let config = datafusion::prelude::SessionConfig::new() - .with_extension(Arc::new(CheckpointStoreService::new(store.clone()))); - let ctx = SessionContext::new_with_config(config); + let ctx = checkpoint_context(store.clone()); let path = "memory:///checkpoint-root/session/relation"; cleanup_checkpoint_path(&ctx, path).await?; @@ -764,4 +830,65 @@ mod tests { assert_eq!(paths.as_slice(), &[path.to_string()]); Ok(()) } + + #[tokio::test] + async fn failed_cached_relation_cleanup_can_be_retried() -> Result<()> { + let store = Arc::new(TestCheckpointStore::default()); + let ctx = checkpoint_context(store.clone()); + let registry = ctx.extension::()?; + let relation_id = "relation"; + let path = "memory:///checkpoint-root/session/relation"; + let plan: Arc = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + registry.insert( + relation_id.to_string(), + CachedRelation::new_pending_reliable_checkpoint(plan, path.to_string()), + )?; + store.set_cleanup_failure(path, true)?; + + assert!(remove_cached_relation(&ctx, relation_id).await.is_err()); + assert!(registry.get(relation_id)?.is_some()); + + store.set_cleanup_failure(path, false)?; + remove_cached_relation(&ctx, relation_id).await?; + assert!(registry.get(relation_id)?.is_none()); + Ok(()) + } + + #[tokio::test] + async fn session_cleanup_attempts_every_cached_relation() -> Result<()> { + let store = Arc::new(TestCheckpointStore::default()); + let ctx = checkpoint_context(store.clone()); + let registry = ctx.extension::()?; + let first_path = "memory:///checkpoint-root/session/first"; + let second_path = "memory:///checkpoint-root/session/second"; + let plan: Arc = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + registry.insert( + "first".to_string(), + CachedRelation::new_pending_reliable_checkpoint( + Arc::clone(&plan), + first_path.to_string(), + ), + )?; + registry.insert( + "second".to_string(), + CachedRelation::new_pending_reliable_checkpoint(plan, second_path.to_string()), + )?; + store.set_cleanup_failure(first_path, true)?; + + assert!(cleanup_cached_relations(&ctx).await.is_err()); + let mut paths = store + .cleanup_paths + .lock() + .map_err(|e| internal_datafusion_err!("{e}"))? + .clone(); + paths.sort(); + assert_eq!(paths, vec![first_path.to_string(), second_path.to_string()]); + assert!(registry.get("first")?.is_some()); + assert!(registry.get("second")?.is_none()); + + store.set_cleanup_failure(first_path, false)?; + cleanup_cached_relations(&ctx).await?; + assert!(registry.get("first")?.is_none()); + Ok(()) + } } diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs index b4f72ec8c8..f0124ce0a6 100644 --- a/crates/sail-session/src/checkpoint.rs +++ b/crates/sail-session/src/checkpoint.rs @@ -32,9 +32,6 @@ impl CheckpointStore for ObjectStoreCheckpointStore { while let Some(batch) = stream.next().await { batches.push(batch?); } - if batches.iter().all(|batch| batch.num_rows() == 0) { - continue; - } let bytes = write_record_batches_file(&batches, schema.as_ref())?; let file_path = checkpoint_path.child(&format!("part-{file_index:05}.arrow")); files.push(checkpoint_path.put_bytes(&file_path, bytes).await?); diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 88ced5ea06..ed3d88704a 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -11,7 +11,7 @@ use futures::stream; use log::{debug, warn}; use sail_common::spec; use sail_common_datafusion::cached_relation::{ - CachedRelation, CachedRelationRegistry, cleanup_cached_relation, cleanup_checkpoint_path, + CachedRelation, CachedRelationRegistry, cleanup_checkpoint_path, remove_cached_relation, }; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; @@ -617,12 +617,7 @@ pub(crate) async fn handle_execute_remove_cached_remote_relation_command( let relation = command .relation .required("remove cached remote relation relation")?; - if let Some(relation) = ctx - .extension::()? - .remove(&relation.relation_id)? - { - cleanup_cached_relation(ctx, relation).await?; - } + remove_cached_relation(ctx, &relation.relation_id).await?; let output = if metadata.reattachable { vec![ExecutorOutput::complete()] } else { diff --git a/python/pysail/tests/spark/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index 536edb33d0..9d723efd9c 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -1,3 +1,4 @@ +import gc import os import shutil import time @@ -24,6 +25,16 @@ def spark_session_factory_fixture(remote): yield sessions.create +def _wait_for_checkpoint_files(checkpoint_path, *, present): + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + files_present = any(checkpoint_path.rglob("*.arrow")) + if files_present == present: + return + time.sleep(0.1) + assert any(checkpoint_path.rglob("*.arrow")) == present + + def test_dataframe_local_checkpoint_survives_source_removal(spark, tmp_path): source_path = tmp_path / "source" spark.createDataFrame( @@ -240,10 +251,42 @@ def test_dataframe_checkpoint_cleanup_on_session_stop(spark_session_factory, tmp assert any(checkpoint_path.rglob("*.arrow")) spark.stop() - deadline = time.monotonic() + 5 - while time.monotonic() < deadline and any(checkpoint_path.rglob("*.arrow")): - time.sleep(0.1) - assert not any(checkpoint_path.rglob("*.arrow")) + _wait_for_checkpoint_files(checkpoint_path, present=False) + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint DataFrame cleanup") +def test_dataframe_checkpoint_cleanup_on_dataframe_gc(spark, tmp_path): + checkpoint_path = tmp_path / "checkpoints" + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + checkpointed = spark.range(0, 5).checkpoint() + assert any(checkpoint_path.rglob("*.arrow")) + + del checkpointed + gc.collect() + _wait_for_checkpoint_files(checkpoint_path, present=False) + finally: + spark.conf.unset("spark.checkpoint.dir") + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific derived checkpoint cleanup") +def test_dataframe_checkpoint_cleanup_waits_for_derived_dataframe(spark, tmp_path): + checkpoint_path = tmp_path / "checkpoints" + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + checkpointed = spark.range(0, 5).checkpoint() + derived = checkpointed.repartition(2) + assert any(checkpoint_path.rglob("*.arrow")) + + del checkpointed + gc.collect() + _wait_for_checkpoint_files(checkpoint_path, present=True) + + del derived + gc.collect() + _wait_for_checkpoint_files(checkpoint_path, present=False) + finally: + spark.conf.unset("spark.checkpoint.dir") @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint partition preservation") @@ -264,6 +307,27 @@ def test_dataframe_checkpoint_preserves_multiple_output_partitions(spark, tmp_pa assert len(partition_ids) > 1 +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint partition preservation") +@pytest.mark.parametrize("local", [True, False], ids=["local", "reliable"]) +def test_dataframe_checkpoint_preserves_empty_partition_ids(spark, tmp_path, local): + df = spark.range(0, 1, numPartitions=4) + expected = df.selectExpr("id", "spark_partition_id() AS pid").collect() + assert expected[0].pid != 0 + + if local: + checkpointed = df.localCheckpoint() + else: + checkpoint_path = tmp_path / "checkpoints" + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + checkpointed = df.checkpoint() + finally: + spark.conf.unset("spark.checkpoint.dir") + + actual = checkpointed.selectExpr("id", "spark_partition_id() AS pid").collect() + assert actual == expected + + @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names and stack configuration") def test_dataframe_local_checkpoint_deep_plan_explain_truncates(spark): df = spark.range(0, 10) From 90b0c2e3527448219155ea7cda9d43eb4c9b0e78 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Tue, 14 Jul 2026 11:45:45 +0800 Subject: [PATCH 19/40] test: cover checkpoint lifecycle and properties --- .../src/cached_relation.rs | 41 +++++++++++++++ .../tests/spark/dataframe/test_checkpoint.py | 52 +++++++++++++++++++ 2 files changed, 93 insertions(+) diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index db62ba23dc..784f0fb1c1 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -854,6 +854,47 @@ mod tests { Ok(()) } + #[tokio::test] + async fn cached_relation_cleanup_waits_for_physical_plan_lease() -> Result<()> { + let store = Arc::new(TestCheckpointStore::default()); + let ctx = checkpoint_context(store.clone()); + let registry = ctx.extension::()?; + let relation_id = "relation"; + let path = "memory:///checkpoint-root/session/relation"; + let plan: Arc = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + registry.insert( + relation_id.to_string(), + CachedRelation::new_pending_reliable_checkpoint(plan, path.to_string()), + )?; + let relation = registry + .get(relation_id)? + .ok_or_else(|| internal_datafusion_err!("cached relation missing"))?; + let physical_plan = relation.to_physical_plan(relation_id).await?; + drop(relation); + + remove_cached_relation(&ctx, relation_id).await?; + assert!(registry.get(relation_id)?.is_none()); + assert!( + store + .cleanup_paths + .lock() + .map_err(|e| internal_datafusion_err!("{e}"))? + .is_empty() + ); + + drop(physical_plan); + cleanup_cached_relations(&ctx).await?; + assert_eq!( + store + .cleanup_paths + .lock() + .map_err(|e| internal_datafusion_err!("{e}"))? + .as_slice(), + &[path.to_string()] + ); + Ok(()) + } + #[tokio::test] async fn session_cleanup_attempts_every_cached_relation() -> Result<()> { let store = Arc::new(TestCheckpointStore::default()); diff --git a/python/pysail/tests/spark/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index 9d723efd9c..52e40a5c5b 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -219,6 +219,23 @@ def test_dataframe_checkpoint_lazy(spark, tmp_path): spark.conf.unset("spark.checkpoint.dir") +@pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") +def test_dataframe_checkpoint_lazy_retains_checkpoint_dependency(spark, tmp_path): + checkpoint_path = tmp_path / "checkpoints" + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + first = spark.range(0, 10).checkpoint() + second = first.checkpoint(eager=False) + + del first + gc.collect() + + assert second.count() == 10 # noqa: PLR2004 + assert [row.id for row in second.sort("id").collect()] == list(range(10)) + finally: + spark.conf.unset("spark.checkpoint.dir") + + @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific lazy checkpoint explain behavior") def test_dataframe_checkpoint_lazy_explain_does_not_materialize(spark, tmp_path): checkpoint_path = tmp_path / "checkpoints" @@ -328,6 +345,33 @@ def test_dataframe_checkpoint_preserves_empty_partition_ids(spark, tmp_path, loc assert actual == expected +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint physical properties") +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("local", [True, False], ids=["local", "reliable"]) +def test_dataframe_checkpoint_preserves_partitioning_and_ordering(spark, tmp_path, local, eager): + df = ( + spark.range(0, 100, numPartitions=4) + .withColumn("key", col("id") % 4) + .repartition(4, "key") + .sortWithinPartitions("key", "id") + ) + if local: + checkpointed = df.localCheckpoint(eager=eager) + else: + checkpoint_path = tmp_path / f"checkpoints-{eager}" + spark.conf.set("spark.checkpoint.dir", str(checkpoint_path)) + try: + checkpointed = df.checkpoint(eager=eager) + finally: + spark.conf.unset("spark.checkpoint.dir") + + aggregate_plan = normalize_plan_text(checkpointed.groupBy("key").count()._explain_string()) # noqa: SLF001 + ordered_plan = normalize_plan_text(checkpointed.sortWithinPartitions("key", "id")._explain_string()) # noqa: SLF001 + + assert "RepartitionExec" not in aggregate_plan + assert "SortExec" not in ordered_plan + + @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names and stack configuration") def test_dataframe_local_checkpoint_deep_plan_explain_truncates(spark): df = spark.range(0, 10) @@ -369,3 +413,11 @@ def test_dataframe_local_checkpoint_storage_level_survives_source_removal(spark, checkpointed.where(col("id") == 2).select("value").toPandas(), # noqa: PLR2004 pd.DataFrame({"value": ["b"]}), ) + + +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific local checkpoint validation") +def test_dataframe_local_checkpoint_rejects_excessive_replication(spark): + storage_level = StorageLevel(True, False, False, False, 40) + + with pytest.raises(Exception, match="replication must be less than 40"): + spark.range(0, 1).localCheckpoint(storageLevel=storage_level) From 29c1d3be6b5be1d6c298a731aa7f1a9cd950a53a Mon Sep 17 00:00:00 2001 From: XL Liang Date: Tue, 14 Jul 2026 11:45:59 +0800 Subject: [PATCH 20/40] fix: preserve checkpoint execution semantics --- .../src/cached_relation.rs | 489 +++++++++++++++--- .../sail-common-datafusion/src/checkpoint.rs | 294 +++++++++++ crates/sail-common-datafusion/src/lib.rs | 1 + .../proto/sail/plan/physical.proto | 19 + crates/sail-execution/src/proto/codec.rs | 64 +++ crates/sail-session/src/checkpoint.rs | 92 +++- .../src/service/plan_executor.rs | 5 + 7 files changed, 875 insertions(+), 89 deletions(-) create mode 100644 crates/sail-common-datafusion/src/checkpoint.rs diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-common-datafusion/src/cached_relation.rs index 784f0fb1c1..8e62b4d214 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-common-datafusion/src/cached_relation.rs @@ -1,20 +1,21 @@ use std::cmp::Ordering; -use std::collections::HashMap; use std::collections::hash_map::Entry; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::fmt::Formatter; use std::sync::{Arc, RwLock}; +use datafusion::arrow::array::{Array, LargeBinaryArray, UInt64Array}; use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::physical_plan::ArrowSource; use datafusion::execution::disk_manager::RefCountedTempFile; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - collect_partitioned, with_new_children_if_necessary, + with_new_children_if_necessary, }; use datafusion::prelude::SessionContext; use datafusion_common::{ @@ -26,13 +27,16 @@ use datafusion_datasource::memory::MemorySourceConfig; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::{PartitionedFile, TableSchema}; use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use futures::StreamExt; use futures::future::BoxFuture; use object_store::ObjectMeta; use sail_common::spec; -use crate::array::record_batch::{read_record_batches, write_record_batches}; +use crate::array::record_batch::read_record_batches; +use crate::checkpoint::LocalCheckpointExec; use crate::extension::{SessionExtension, SessionExtensionAccessor}; use crate::session::checkpoint::{CheckpointStoreService, ReliableCheckpoint}; +use crate::session::job::JobService; #[derive(Debug, Clone)] enum CachedRelationCleanup { @@ -119,19 +123,26 @@ enum CachedRelationMaterialized { Reliable { schema: SchemaRef, checkpoint: ReliableCheckpoint, + properties: Arc, }, } #[derive(Debug, Clone)] struct CachedRelationLocalMaterialized { schema: SchemaRef, + properties: Arc, memory_partitions: Option>>>, - serialized_memory_partitions: Option>>>, + serialized_memory_partitions: Option>>>>, disk_partitions: Option>>, } #[derive(Debug, Clone)] struct CachedRelationDiskPartition { + chunks: Vec, +} + +#[derive(Debug, Clone)] +struct CachedRelationDiskChunk { files: Vec, } @@ -147,6 +158,104 @@ enum CachedRelationPendingTarget { Reliable { path: String }, } +#[derive(Debug)] +pub struct CachedRelationExec { + input: Arc, + properties: Arc, + relation_lease: Option, +} + +impl CachedRelationExec { + pub fn new(input: Arc, properties: Arc) -> Self { + Self { + input, + properties, + relation_lease: None, + } + } + + fn with_relation_lease( + input: Arc, + properties: Arc, + relation: CachedRelation, + ) -> Self { + Self { + input, + properties, + relation_lease: Some(relation), + } + } + + pub fn input(&self) -> &Arc { + &self.input + } +} + +impl DisplayAs for CachedRelationExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CachedRelationExec") + } +} + +impl ExecutionPlan for CachedRelationExec { + fn name(&self) -> &str { + Self::static_name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn maintains_input_order(&self) -> Vec { + vec![true] + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(internal_datafusion_err!( + "CachedRelationExec must have exactly one child" + )); + } + Ok(Arc::new(Self { + input: Arc::clone(&children[0]), + properties: Arc::clone(&self.properties), + relation_lease: self.relation_lease.clone(), + })) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let stream = self.input.execute(partition, context)?; + let Some(relation_lease) = self.relation_lease.clone() else { + return Ok(stream); + }; + let schema = stream.schema(); + let stream = stream.map(move |batch| { + let _ = &relation_lease; + batch + }); + Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.input.partition_statistics(partition) + } +} + #[derive(Debug, Clone)] struct PendingCachedRelationExec { relation_id: String, @@ -156,14 +265,7 @@ struct PendingCachedRelationExec { impl PendingCachedRelationExec { fn new(relation_id: String, relation: CachedRelation, plan: Arc) -> Self { - let schema = plan.schema(); - let partition_count = plan.output_partitioning().partition_count(); - let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(schema), - Partitioning::UnknownPartitioning(partition_count), - EmissionType::Final, - Boundedness::Bounded, - )); + let properties = Arc::clone(plan.properties()); Self { relation_id, relation, @@ -288,7 +390,14 @@ impl CachedRelation { pub async fn to_physical_plan(&self, relation_id: &str) -> Result> { let data = self.data.lock().await; match &*data { - CachedRelationData::Materialized(materialized) => materialized.to_physical_plan().await, + CachedRelationData::Materialized(materialized) => { + let input = materialized.to_physical_plan().await?; + Ok(Arc::new(CachedRelationExec::with_relation_lease( + input, + materialized.properties(), + self.clone(), + ))) + } CachedRelationData::Pending(pending) => Ok(Arc::new(PendingCachedRelationExec::new( relation_id.to_string(), self.clone(), @@ -303,26 +412,52 @@ impl CachedRelation { *data = pending.materialize(ctx).await?; } match &*data { - CachedRelationData::Materialized(materialized) => materialized.to_physical_plan().await, + CachedRelationData::Materialized(materialized) => { + let input = materialized.to_physical_plan().await?; + Ok(Arc::new(CachedRelationExec::with_relation_lease( + input, + materialized.properties(), + self.clone(), + ))) + } CachedRelationData::Pending(_) => Err(internal_datafusion_err!( "cached relation materialization did not complete" )), } } + + fn is_exclusively_owned(&self) -> bool { + Arc::strong_count(&self.data) == 1 + } } #[derive(Debug, Default)] pub struct CachedRelationRegistry { - relations: RwLock>, + state: RwLock, +} + +#[derive(Debug, Default)] +struct CachedRelationRegistryState { + relations: HashMap, + retired: Vec<(String, CachedRelation)>, } impl CachedRelationRegistry { pub fn insert(&self, relation_id: String, relation: CachedRelation) -> Result<()> { - let mut relations = self - .relations + let mut state = self + .state .write() .map_err(|e| internal_datafusion_err!("{e}"))?; - match relations.entry(relation_id) { + if state + .retired + .iter() + .any(|(retired_id, _)| retired_id == &relation_id) + { + return Err(internal_datafusion_err!( + "cached relation already exists: {relation_id}" + )); + } + match state.relations.entry(relation_id) { Entry::Occupied(entry) => Err(internal_datafusion_err!( "cached relation already exists: {}", entry.key() @@ -335,27 +470,51 @@ impl CachedRelationRegistry { } pub fn get(&self, relation_id: &str) -> Result> { - let relations = self - .relations + let state = self + .state .read() .map_err(|e| internal_datafusion_err!("{e}"))?; - Ok(relations.get(relation_id).cloned()) + Ok(state.relations.get(relation_id).cloned()) } pub fn remove(&self, relation_id: &str) -> Result> { - let mut relations = self - .relations + let mut state = self + .state + .write() + .map_err(|e| internal_datafusion_err!("{e}"))?; + Ok(state.relations.remove(relation_id)) + } + + fn retire(&self, relation_id: String, relation: CachedRelation) -> Result<()> { + let mut state = self + .state + .write() + .map_err(|e| internal_datafusion_err!("{e}"))?; + state.retired.push((relation_id, relation)); + Ok(()) + } + + fn take_cleanup_ready(&self) -> Result> { + let mut state = self + .state .write() .map_err(|e| internal_datafusion_err!("{e}"))?; - Ok(relations.remove(relation_id)) + let retired = std::mem::take(&mut state.retired); + let (ready, retained) = retired + .into_iter() + .partition(|(_, relation)| relation.is_exclusively_owned()); + state.retired = retained; + Ok(ready) } pub fn drain(&self) -> Result> { - let mut relations = self - .relations + let mut state = self + .state .write() .map_err(|e| internal_datafusion_err!("{e}"))?; - Ok(relations.drain().collect()) + let mut relations: Vec<_> = state.relations.drain().collect(); + relations.append(&mut state.retired); + Ok(relations) } } @@ -396,6 +555,13 @@ impl CachedRelationMaterialized { } } + fn properties(&self) -> Arc { + match self { + Self::Local(local) => Arc::clone(&local.properties), + Self::Reliable { properties, .. } => Arc::clone(properties), + } + } + fn to_logical_plan(&self, relation_id: &str) -> Result { Ok(LogicalPlan::Extension(Extension { node: Arc::new(CachedRelationNode::try_new( @@ -408,7 +574,9 @@ impl CachedRelationMaterialized { async fn to_physical_plan(&self) -> Result> { match self { Self::Local(local) => local.to_physical_plan().await, - Self::Reliable { schema, checkpoint } => create_arrow_checkpoint_scan( + Self::Reliable { + schema, checkpoint, .. + } => create_arrow_checkpoint_scan( checkpoint.object_store_url().clone(), checkpoint.object_meta().to_vec(), schema, @@ -421,7 +589,9 @@ impl CachedRelationLocalMaterialized { async fn try_new( ctx: &SessionContext, schema: SchemaRef, - partitions: Vec>, + properties: Arc, + partition_count: usize, + mut stream: SendableRecordBatchStream, storage_level: spec::StorageLevel, ) -> Result { let use_memory = storage_level.use_memory; @@ -433,55 +603,125 @@ impl CachedRelationLocalMaterialized { "local checkpoint storage level must use memory or disk" )); } + if !(1..40).contains(&storage_level.replication) { + return Err(internal_datafusion_err!( + "local checkpoint storage level replication must be between 1 and 39" + )); + } + let mut memory_partitions = if use_deserialized_memory { + Some(partition_maps(partition_count)) + } else { + None + }; let mut serialized_partitions = if use_serialized_memory { - Some(Vec::with_capacity(partitions.len())) + Some(partition_maps(partition_count)) } else { None }; let mut disk_partitions = if use_disk { - Some(Vec::with_capacity(partitions.len())) + Some(partition_maps(partition_count)) } else { None }; + let mut sequences: Vec> = + (0..partition_count).map(|_| BTreeSet::new()).collect(); - for partition in &partitions { - let bytes = if use_serialized_memory || use_disk { - Some(write_record_batches(partition, schema.as_ref())?) - } else { - None - }; - - if let Some(serialized_partitions) = serialized_partitions.as_mut() { - let bytes = bytes - .as_ref() - .ok_or_else(|| internal_datafusion_err!("missing serialized partition"))?; - serialized_partitions.push(bytes.clone()); + while let Some(batch) = stream.next().await { + let batch = batch?; + if batch.num_columns() != 3 { + return Err(internal_datafusion_err!( + "local checkpoint returned {} columns instead of 3", + batch.num_columns() + )); } + let partition_array = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("invalid checkpoint partition column"))?; + let sequence_array = batch + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("invalid checkpoint sequence column"))?; + let data_array = batch + .column(2) + .as_any() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("invalid checkpoint data column"))?; + + for row in 0..batch.num_rows() { + if partition_array.is_null(row) + || sequence_array.is_null(row) + || data_array.is_null(row) + { + return Err(internal_datafusion_err!( + "local checkpoint returned a null metadata value" + )); + } + let partition = usize::try_from(partition_array.value(row)).map_err(|_| { + internal_datafusion_err!("local checkpoint partition index is too large") + })?; + if partition >= partition_count { + return Err(internal_datafusion_err!( + "local checkpoint returned invalid partition {partition}" + )); + } + let sequence = sequence_array.value(row); + if !sequences[partition].insert(sequence) { + return Err(internal_datafusion_err!( + "local checkpoint returned duplicate partition {partition} sequence {sequence}" + )); + } + let bytes = data_array.value(row); - if let Some(disk_partitions) = disk_partitions.as_mut() { - let bytes = bytes - .as_ref() - .ok_or_else(|| internal_datafusion_err!("missing serialized partition"))?; - let replication = storage_level.replication.max(1); - let mut files = Vec::with_capacity(replication); - for _ in 0..replication { - files.push(write_disk_partition(ctx, bytes).await?); + if let Some(partitions) = memory_partitions.as_mut() { + partitions[partition].insert(sequence, read_record_batches(bytes)?); + } + if let Some(partitions) = serialized_partitions.as_mut() { + partitions[partition].insert(sequence, bytes.to_vec()); + } + if let Some(partitions) = disk_partitions.as_mut() { + partitions[partition].insert( + sequence, + write_disk_chunk(ctx, bytes, storage_level.replication).await?, + ); } - disk_partitions.push(CachedRelationDiskPartition { files }); } } - let memory_partitions = if use_deserialized_memory { - Some(Arc::new(partitions)) - } else { - None - }; - let serialized_memory_partitions = serialized_partitions.map(Arc::new); - let disk_partitions = disk_partitions.map(Arc::new); + validate_partition_sequences(&sequences)?; + let memory_partitions = memory_partitions.map(|partitions| { + Arc::new( + partitions + .into_iter() + .map(|partition| partition.into_values().flatten().collect()) + .collect(), + ) + }); + let serialized_memory_partitions = serialized_partitions.map(|partitions| { + Arc::new( + partitions + .into_iter() + .map(|partition| partition.into_values().collect()) + .collect(), + ) + }); + let disk_partitions = disk_partitions.map(|partitions| { + Arc::new( + partitions + .into_iter() + .map(|partition| CachedRelationDiskPartition { + chunks: partition.into_values().collect(), + }) + .collect(), + ) + }); Ok(Self { schema, + properties, memory_partitions, serialized_memory_partitions, disk_partitions, @@ -500,10 +740,15 @@ impl CachedRelationLocalMaterialized { return Ok(partitions.as_ref().clone()); } if let Some(partitions) = &self.serialized_memory_partitions { - return partitions - .iter() - .map(|bytes| read_record_batches(bytes)) - .collect(); + let mut batches = Vec::with_capacity(partitions.len()); + for partition in partitions.iter() { + let mut partition_batches = vec![]; + for bytes in partition { + partition_batches.extend(read_record_batches(bytes)?); + } + batches.push(partition_batches); + } + return Ok(batches); } if let Some(partitions) = &self.disk_partitions { let mut batches = Vec::with_capacity(partitions.len()); @@ -519,6 +764,16 @@ impl CachedRelationLocalMaterialized { } impl CachedRelationDiskPartition { + async fn read(&self) -> Result> { + let mut batches = vec![]; + for chunk in &self.chunks { + batches.extend(chunk.read().await?); + } + Ok(batches) + } +} + +impl CachedRelationDiskChunk { async fn read(&self) -> Result> { let mut last_error = None; for file in &self.files { @@ -532,11 +787,33 @@ impl CachedRelationDiskPartition { } } Err(last_error.unwrap_or_else(|| { - DataFusionError::Internal("cached relation disk partition has no files".to_string()) + DataFusionError::Internal("cached relation disk chunk has no files".to_string()) })) } } +fn partition_maps(partition_count: usize) -> Vec> { + (0..partition_count).map(|_| BTreeMap::new()).collect() +} + +fn validate_partition_sequences(sequences: &[BTreeSet]) -> Result<()> { + for (partition, sequences) in sequences.iter().enumerate() { + if sequences.is_empty() { + return Err(internal_datafusion_err!( + "local checkpoint did not return partition {partition}" + )); + } + let expected = 0..u64::try_from(sequences.len()) + .map_err(|_| internal_datafusion_err!("checkpoint sequence count is too large"))?; + if !sequences.iter().copied().eq(expected) { + return Err(internal_datafusion_err!( + "local checkpoint returned incomplete sequence for partition {partition}" + )); + } + } + Ok(()) +} + pub fn materialize_cached_relations<'a>( ctx: &'a SessionContext, plan: Arc, @@ -562,9 +839,22 @@ async fn materialize_local_checkpoint( ) -> Result { let plan = materialize_cached_relations(ctx, plan).await?; let schema = plan.schema(); - let partitions = collect_checkpoint_partitions(ctx, plan).await?; - let materialized = - CachedRelationLocalMaterialized::try_new(ctx, schema, partitions, storage_level).await?; + let properties = checkpoint_plan_properties(&plan); + let partition_count = plan.output_partitioning().partition_count(); + let service = ctx.extension::()?; + let stream = service + .runner() + .execute(ctx, Arc::new(LocalCheckpointExec::new(plan))) + .await?; + let materialized = CachedRelationLocalMaterialized::try_new( + ctx, + schema, + properties, + partition_count, + stream, + storage_level, + ) + .await?; Ok(CachedRelationData::Materialized( CachedRelationMaterialized::Local(materialized), )) @@ -577,6 +867,7 @@ async fn materialize_reliable_checkpoint( ) -> Result { let plan = materialize_cached_relations(ctx, plan).await?; let schema = plan.schema(); + let properties = checkpoint_plan_properties(&plan); let service = ctx.extension::()?; let checkpoint = match service .write_reliable_checkpoint(ctx, plan, path, Arc::clone(&schema)) @@ -589,7 +880,20 @@ async fn materialize_reliable_checkpoint( } }; Ok(CachedRelationData::Materialized( - CachedRelationMaterialized::Reliable { schema, checkpoint }, + CachedRelationMaterialized::Reliable { + schema, + checkpoint, + properties, + }, + )) +} + +fn checkpoint_plan_properties(plan: &Arc) -> Arc { + Arc::new(PlanProperties::new( + plan.properties().eq_properties.clone(), + plan.output_partitioning().clone(), + EmissionType::Incremental, + Boundedness::Bounded, )) } @@ -638,6 +942,10 @@ pub async fn remove_cached_relation(ctx: &SessionContext, relation_id: &str) -> let Some(relation) = registry.remove(relation_id)? else { return Ok(()); }; + if !relation.is_exclusively_owned() { + registry.retire(relation_id.to_string(), relation)?; + return cleanup_retired_cached_relations(ctx, ®istry).await; + } if let Err(cleanup_error) = cleanup_cached_relation(ctx, &relation).await { if let Err(restore_error) = registry.insert(relation_id.to_string(), relation) { return Err(internal_datafusion_err!( @@ -646,6 +954,34 @@ pub async fn remove_cached_relation(ctx: &SessionContext, relation_id: &str) -> } return Err(cleanup_error); } + drop(relation); + cleanup_retired_cached_relations(ctx, ®istry).await +} + +async fn cleanup_retired_cached_relations( + ctx: &SessionContext, + registry: &CachedRelationRegistry, +) -> Result<()> { + let relations = registry.take_cleanup_ready()?; + let mut errors = vec![]; + for (relation_id, relation) in relations { + if let Err(cleanup_error) = cleanup_cached_relation(ctx, &relation).await { + if let Err(restore_error) = registry.retire(relation_id.clone(), relation) { + errors.push(format!( + "{relation_id}: {cleanup_error}; additionally failed to retain it for retry: {restore_error}" + )); + } else { + errors.push(format!("{relation_id}: {cleanup_error}")); + } + } + } + if !errors.is_empty() { + return Err(internal_datafusion_err!( + "failed to clean {} retired cached relation(s): {}", + errors.len(), + errors.join("; ") + )); + } Ok(()) } @@ -674,14 +1010,19 @@ pub async fn cleanup_cached_relations(ctx: &SessionContext) -> Result<()> { Ok(()) } -async fn collect_checkpoint_partitions( +async fn write_disk_chunk( ctx: &SessionContext, - plan: Arc, -) -> Result>> { - collect_partitioned(plan, ctx.task_ctx()).await + bytes: &[u8], + replication: usize, +) -> Result { + let mut files = Vec::with_capacity(replication); + for _ in 0..replication { + files.push(write_disk_file(ctx, bytes).await?); + } + Ok(CachedRelationDiskChunk { files }) } -async fn write_disk_partition(ctx: &SessionContext, bytes: &[u8]) -> Result { +async fn write_disk_file(ctx: &SessionContext, bytes: &[u8]) -> Result { let mut file = ctx .runtime_env() .disk_manager diff --git a/crates/sail-common-datafusion/src/checkpoint.rs b/crates/sail-common-datafusion/src/checkpoint.rs new file mode 100644 index 0000000000..a486f9774b --- /dev/null +++ b/crates/sail-common-datafusion/src/checkpoint.rs @@ -0,0 +1,294 @@ +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, LargeBinaryArray, RecordBatch, UInt64Array}; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::execution::object_store::ObjectStoreUrl; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning}; +use datafusion::physical_plan::execution_plan::EmissionType; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, +}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use futures::{StreamExt, stream}; +use object_store::path::Path; +use object_store::{ObjectStoreExt, PutPayload}; + +use crate::array::record_batch::{write_record_batches, write_record_batches_file}; + +const PARTITION_COLUMN: &str = "partition"; +const SEQUENCE_COLUMN: &str = "sequence"; +const DATA_COLUMN: &str = "data"; + +#[derive(Debug)] +pub struct LocalCheckpointExec { + input: Arc, + properties: Arc, +} + +impl LocalCheckpointExec { + pub fn new(input: Arc) -> Self { + let schema = Arc::new(Schema::new(vec![ + Field::new(PARTITION_COLUMN, DataType::UInt64, false), + Field::new(SEQUENCE_COLUMN, DataType::UInt64, false), + Field::new(DATA_COLUMN, DataType::LargeBinary, false), + ])); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(input.output_partitioning().partition_count()), + EmissionType::Incremental, + input.boundedness(), + )); + Self { input, properties } + } + + pub fn input(&self) -> &Arc { + &self.input + } +} + +impl DisplayAs for LocalCheckpointExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "LocalCheckpointExec") + } +} + +impl ExecutionPlan for LocalCheckpointExec { + fn name(&self) -> &str { + Self::static_name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(internal_datafusion_err!( + "LocalCheckpointExec must have exactly one child" + )); + } + Ok(Arc::new(Self::new(Arc::clone(&children[0])))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let input = self.input.execute(partition, context)?; + let input_schema = self.input.schema(); + let output_schema = self.schema(); + let stream_schema = Arc::clone(&output_schema); + let stream = stream::unfold( + LocalCheckpointStreamState { + input, + sequence: 0, + emitted: false, + finished: false, + }, + move |mut state| { + let input_schema = Arc::clone(&input_schema); + let output_schema = Arc::clone(&output_schema); + async move { + if state.finished { + return None; + } + let result = match state.input.next().await { + Some(Ok(batch)) => { + let bytes = write_record_batches(&[batch], input_schema.as_ref()); + bytes.and_then(|bytes| { + local_checkpoint_batch( + partition, + state.sequence, + &bytes, + output_schema, + ) + }) + } + Some(Err(error)) => { + state.finished = true; + Err(error) + } + None if !state.emitted => { + state.finished = true; + let bytes = write_record_batches(&[], input_schema.as_ref()); + bytes.and_then(|bytes| { + local_checkpoint_batch(partition, 0, &bytes, output_schema) + }) + } + None => return None, + }; + if result.is_err() { + state.finished = true; + } + state.sequence += 1; + state.emitted = true; + Some((result, state)) + } + }, + ); + Ok(Box::pin(RecordBatchStreamAdapter::new( + stream_schema, + stream, + ))) + } +} + +struct LocalCheckpointStreamState { + input: SendableRecordBatchStream, + sequence: u64, + emitted: bool, + finished: bool, +} + +fn local_checkpoint_batch( + partition: usize, + sequence: u64, + bytes: &[u8], + schema: SchemaRef, +) -> Result { + let partition = u64::try_from(partition) + .map_err(|_| internal_datafusion_err!("checkpoint partition index is too large"))?; + let columns: Vec = vec![ + Arc::new(UInt64Array::from(vec![partition])), + Arc::new(UInt64Array::from(vec![sequence])), + Arc::new(LargeBinaryArray::from_vec(vec![bytes])), + ]; + Ok(RecordBatch::try_new(schema, columns)?) +} + +#[derive(Debug)] +pub struct ReliableCheckpointExec { + input: Arc, + object_store_url: ObjectStoreUrl, + path: Path, + properties: Arc, +} + +impl ReliableCheckpointExec { + pub fn new( + input: Arc, + object_store_url: ObjectStoreUrl, + path: Path, + ) -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + PARTITION_COLUMN, + DataType::UInt64, + false, + )])); + let properties = Arc::new(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::UnknownPartitioning(input.output_partitioning().partition_count()), + EmissionType::Final, + input.boundedness(), + )); + Self { + input, + object_store_url, + path, + properties, + } + } + + pub fn input(&self) -> &Arc { + &self.input + } + + pub fn object_store_url(&self) -> &ObjectStoreUrl { + &self.object_store_url + } + + pub fn path(&self) -> &Path { + &self.path + } +} + +impl DisplayAs for ReliableCheckpointExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!( + f, + "ReliableCheckpointExec: object_store_url={}, path={}", + self.object_store_url, self.path + ) + } +} + +impl ExecutionPlan for ReliableCheckpointExec { + fn name(&self) -> &str { + Self::static_name() + } + + fn properties(&self) -> &Arc { + &self.properties + } + + fn benefits_from_input_partitioning(&self) -> Vec { + vec![false] + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + if children.len() != 1 { + return Err(internal_datafusion_err!( + "ReliableCheckpointExec must have exactly one child" + )); + } + Ok(Arc::new(Self::new( + Arc::clone(&children[0]), + self.object_store_url.clone(), + self.path.clone(), + ))) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + let mut input = self.input.execute(partition, Arc::clone(&context))?; + let input_schema = self.input.schema(); + let output_schema = self.schema(); + let stream_schema = Arc::clone(&output_schema); + let store = context.runtime_env().object_store(&self.object_store_url)?; + let location = self.path.clone().join(format!("part-{partition:05}.arrow")); + let output = stream::once(async move { + let mut batches = vec![]; + while let Some(batch) = input.next().await { + batches.push(batch?); + } + let bytes = write_record_batches_file(&batches, input_schema.as_ref())?; + store + .put(&location, PutPayload::from(bytes)) + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?; + let partition = u64::try_from(partition) + .map_err(|_| internal_datafusion_err!("checkpoint partition index is too large"))?; + let columns: Vec = vec![Arc::new(UInt64Array::from(vec![partition]))]; + Ok(RecordBatch::try_new(output_schema, columns)?) + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + stream_schema, + output, + ))) + } +} diff --git a/crates/sail-common-datafusion/src/lib.rs b/crates/sail-common-datafusion/src/lib.rs index ad6fa365ab..3b34e83d78 100644 --- a/crates/sail-common-datafusion/src/lib.rs +++ b/crates/sail-common-datafusion/src/lib.rs @@ -1,6 +1,7 @@ pub mod array; pub mod cached_relation; pub mod catalog; +pub mod checkpoint; pub mod column_features; pub mod datasource; pub mod display; diff --git a/crates/sail-execution/proto/sail/plan/physical.proto b/crates/sail-execution/proto/sail/plan/physical.proto index 9ea7495411..a006d87fbf 100644 --- a/crates/sail-execution/proto/sail/plan/physical.proto +++ b/crates/sail-execution/proto/sail/plan/physical.proto @@ -67,6 +67,9 @@ message ExtendedPhysicalPlanNode { NoopSinkExecNode noop_sink = 52; CsvExecNode csv = 53; ParquetExecNode parquet = 54; + CachedRelationExecNode cached_relation = 55; + LocalCheckpointExecNode local_checkpoint = 56; + ReliableCheckpointExecNode reliable_checkpoint = 57; } } @@ -589,6 +592,22 @@ message MemoryExecNode { optional uint64 limit = 6; } +message CachedRelationExecNode { + bytes input = 1; + EquivalenceProperties eq_properties = 2; + bytes partitioning = 3; +} + +message LocalCheckpointExecNode { + bytes input = 1; +} + +message ReliableCheckpointExecNode { + bytes input = 1; + string object_store_url = 2; + string path = 3; +} + message ValuesExecNode { bytes data = 1; bytes schema = 2; diff --git a/crates/sail-execution/src/proto/codec.rs b/crates/sail-execution/src/proto/codec.rs index dab1e4c0c0..c2a9f1888a 100644 --- a/crates/sail-execution/src/proto/codec.rs +++ b/crates/sail-execution/src/proto/codec.rs @@ -80,9 +80,11 @@ use datafusion_spark::function::url::url_encode::UrlEncode; use prost::Message; use sail_catalog_system::physical_plan::SystemTableExec; use sail_common_datafusion::array::record_batch::{read_record_batches, write_record_batches}; +use sail_common_datafusion::cached_relation::CachedRelationExec; use sail_common_datafusion::catalog::{ CatalogPartitionField, LakehouseExecutionContext, PartitionTransform, }; +use sail_common_datafusion::checkpoint::{LocalCheckpointExec, ReliableCheckpointExec}; use sail_common_datafusion::datasource::PhysicalSinkMode; use sail_common_datafusion::schema_evolution::SchemaEvolutionCastColumnExpr; use sail_common_datafusion::system::catalog::SystemTable; @@ -450,6 +452,46 @@ impl PhysicalExtensionCodec for RemoteExecutionCodec { .with_limit(limit.map(|x| x as usize)); Ok(Arc::new(DataSourceExec::new(Arc::new(source)))) } + NodeKind::CachedRelation(r#gen::CachedRelationExecNode { + input, + eq_properties, + partitioning, + }) => { + let input = try_decode_physical_plan(ctx, self, &input)?; + let eq_properties = match eq_properties { + Some(properties) => self.try_decode_equivalence_properties(&properties, ctx)?, + None => return plan_err!("no equivalence properties for CachedRelationExec"), + }; + let partitioning = + self.try_decode_partitioning(&partitioning, eq_properties.schema(), ctx)?; + let properties = Arc::new(PlanProperties::new( + eq_properties, + partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Arc::new(CachedRelationExec::new(input, properties))) + } + NodeKind::LocalCheckpoint(r#gen::LocalCheckpointExecNode { input }) => { + let input = try_decode_physical_plan(ctx, self, &input)?; + Ok(Arc::new(LocalCheckpointExec::new(input))) + } + NodeKind::ReliableCheckpoint(r#gen::ReliableCheckpointExecNode { + input, + object_store_url, + path, + }) => { + let input = try_decode_physical_plan(ctx, self, &input)?; + let object_store_url = + datafusion::execution::object_store::ObjectStoreUrl::parse(object_store_url)?; + let path = object_store::path::Path::parse(path) + .map_err(|e| plan_datafusion_err!("invalid checkpoint path: {e}"))?; + Ok(Arc::new(ReliableCheckpointExec::new( + input, + object_store_url, + path, + ))) + } NodeKind::Values(r#gen::ValuesExecNode { data, schema }) => { let schema = try_decode_schema(&schema)?; let data = read_record_batches(&data)?; @@ -1506,6 +1548,28 @@ impl PhysicalExtensionCodec for RemoteExecutionCodec { udf: Some(udf), schema, }) + } else if let Some(cached_relation) = node.downcast_ref::() { + let input = try_encode_physical_plan(self, cached_relation.input().clone())?; + let eq_properties = self.try_encode_equivalence_properties( + cached_relation.properties().equivalence_properties(), + )?; + let partitioning = + self.try_encode_partitioning(cached_relation.properties().output_partitioning())?; + NodeKind::CachedRelation(r#gen::CachedRelationExecNode { + input, + eq_properties: Some(eq_properties), + partitioning, + }) + } else if let Some(local_checkpoint) = node.downcast_ref::() { + let input = try_encode_physical_plan(self, local_checkpoint.input().clone())?; + NodeKind::LocalCheckpoint(r#gen::LocalCheckpointExecNode { input }) + } else if let Some(reliable_checkpoint) = node.downcast_ref::() { + let input = try_encode_physical_plan(self, reliable_checkpoint.input().clone())?; + NodeKind::ReliableCheckpoint(r#gen::ReliableCheckpointExecNode { + input, + object_store_url: reliable_checkpoint.object_store_url().as_str().to_string(), + path: reliable_checkpoint.path().to_string(), + }) } else if let Some(work_table) = node.downcast_ref::() { let name = work_table.name().to_string(); let schema = try_encode_schema(work_table.schema().as_ref())?; diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs index f0124ce0a6..708a9c6f5e 100644 --- a/crates/sail-session/src/checkpoint.rs +++ b/crates/sail-session/src/checkpoint.rs @@ -1,13 +1,18 @@ use std::sync::Arc; use async_trait::async_trait; +use datafusion::arrow::array::{Array, UInt64Array}; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::physical_plan::{ExecutionPlan, execute_stream_partitioned}; +use datafusion::object_store::ObjectStoreExt; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion::prelude::SessionContext; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use futures::StreamExt; use sail_common_datafusion::array::record_batch::write_record_batches_file; +use sail_common_datafusion::checkpoint::ReliableCheckpointExec; +use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::checkpoint::{CheckpointStore, ReliableCheckpoint}; +use sail_common_datafusion::session::job::JobService; use sail_object_store::{delete_object_store_prefix, resolve_object_store_path}; #[derive(Debug, Default)] @@ -24,23 +29,80 @@ impl CheckpointStore for ObjectStoreCheckpointStore { ) -> Result { let runtime_env = ctx.runtime_env(); let checkpoint_path = resolve_object_store_path(runtime_env.as_ref(), path)?; - let streams = execute_stream_partitioned(plan, ctx.task_ctx())?; - let mut files = vec![]; + let partition_count = plan.output_partitioning().partition_count(); + if partition_count == 0 { + let bytes = write_record_batches_file(&[], schema.as_ref())?; + let file_path = checkpoint_path.child("part-00000.arrow"); + let file = checkpoint_path.put_bytes(&file_path, bytes).await?; + return Ok(ReliableCheckpoint::new( + checkpoint_path.object_store_url().clone(), + vec![file], + )); + } - for (file_index, mut stream) in streams.into_iter().enumerate() { - let mut batches = vec![]; - while let Some(batch) = stream.next().await { - batches.push(batch?); + let checkpoint_exec = ReliableCheckpointExec::new( + plan, + checkpoint_path.object_store_url().clone(), + checkpoint_path.prefix().clone(), + ); + let service = ctx.extension::()?; + let mut stream = service + .runner() + .execute(ctx, Arc::new(checkpoint_exec)) + .await?; + let mut completed = vec![false; partition_count]; + while let Some(batch) = stream.next().await { + let batch = batch?; + if batch.num_columns() != 1 { + return Err(internal_datafusion_err!( + "reliable checkpoint returned {} columns instead of 1", + batch.num_columns() + )); + } + let partitions = batch + .column(0) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("invalid reliable checkpoint partition column") + })?; + for row in 0..batch.num_rows() { + if partitions.is_null(row) { + return Err(internal_datafusion_err!( + "reliable checkpoint returned a null partition" + )); + } + let partition = usize::try_from(partitions.value(row)).map_err(|_| { + internal_datafusion_err!("reliable checkpoint partition index is too large") + })?; + let Some(done) = completed.get_mut(partition) else { + return Err(internal_datafusion_err!( + "reliable checkpoint returned invalid partition {partition}" + )); + }; + if std::mem::replace(done, true) { + return Err(internal_datafusion_err!( + "reliable checkpoint returned duplicate partition {partition}" + )); + } } - let bytes = write_record_batches_file(&batches, schema.as_ref())?; - let file_path = checkpoint_path.child(&format!("part-{file_index:05}.arrow")); - files.push(checkpoint_path.put_bytes(&file_path, bytes).await?); } - if files.is_empty() { - let bytes = write_record_batches_file(&[], schema.as_ref())?; - let file_path = checkpoint_path.child("part-00000.arrow"); - files.push(checkpoint_path.put_bytes(&file_path, bytes).await?); + let mut files = Vec::with_capacity(partition_count); + for (partition, completed) in completed.into_iter().enumerate() { + if !completed { + return Err(internal_datafusion_err!( + "reliable checkpoint did not return partition {partition}" + )); + } + let file_path = checkpoint_path.child(&format!("part-{partition:05}.arrow")); + files.push( + checkpoint_path + .store() + .head(&file_path) + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?, + ); } Ok(ReliableCheckpoint::new( diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index ed3d88704a..ed33b8effa 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -652,6 +652,11 @@ fn resolve_local_checkpoint_storage_level( "localCheckpoint StorageLevel replication must be at least 1", )); } + if storage_level.replication >= 40 { + return Err(SparkError::invalid( + "localCheckpoint StorageLevel replication must be less than 40", + )); + } Ok(storage_level) } From 0438b3525bf332e7e88f77bacc68dfba52638f5f Mon Sep 17 00:00:00 2001 From: XL Liang Date: Tue, 14 Jul 2026 13:55:28 +0800 Subject: [PATCH 21/40] remove RUST_MIN_STACK environment variable --- .github/workflows/rust-tests.yml | 1 - 1 file changed, 1 deletion(-) diff --git a/.github/workflows/rust-tests.yml b/.github/workflows/rust-tests.yml index 65b34586d1..af94daa12e 100644 --- a/.github/workflows/rust-tests.yml +++ b/.github/workflows/rust-tests.yml @@ -19,7 +19,6 @@ jobs: # Coverage environment variables RUSTC_WORKSPACE_WRAPPER: ${{ github.workspace }}/.github/scripts/rustc-workspace-wrapper.sh LLVM_PROFILE_FILE: ${{ github.workspace }}/target/llvm-profiles/rust-unit/sail-%p-%m.profraw - RUST_MIN_STACK: 8388608 steps: - uses: actions/checkout@v7 From fc0cc6082162c6203ea5c79e008db291b4272597 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Wed, 15 Jul 2026 20:41:49 +0800 Subject: [PATCH 22/40] refactor --- Cargo.lock | 14 +++- crates/sail-cache/Cargo.toml | 9 ++ .../src/cached_relation.rs | 82 ++----------------- .../session => sail-cache/src}/checkpoint.rs | 3 +- crates/sail-cache/src/lib.rs | 2 + crates/sail-common-datafusion/Cargo.toml | 2 - crates/sail-common-datafusion/src/lib.rs | 2 - .../sail-common-datafusion/src/session/mod.rs | 1 - crates/sail-execution/Cargo.toml | 1 + crates/sail-execution/src/proto/codec.rs | 4 +- .../sail-logical-plan/src/cached_relation.rs | 69 ++++++++++++++++ crates/sail-logical-plan/src/lib.rs | 1 + crates/sail-physical-plan/Cargo.toml | 1 + .../src/checkpoint.rs | 5 +- crates/sail-physical-plan/src/lib.rs | 1 + crates/sail-plan/Cargo.toml | 1 + crates/sail-plan/src/lib.rs | 2 +- crates/sail-plan/src/resolver/query/mod.rs | 2 +- crates/sail-session/src/checkpoint.rs | 4 +- crates/sail-session/src/planner.rs | 3 +- .../src/session_factory/server.rs | 4 +- .../src/session_factory/worker.rs | 7 -- .../src/session_manager/actor/handler.rs | 2 +- .../src/service/plan_executor.rs | 4 +- 24 files changed, 123 insertions(+), 103 deletions(-) rename crates/{sail-common-datafusion => sail-cache}/src/cached_relation.rs (95%) rename crates/{sail-common-datafusion/src/session => sail-cache/src}/checkpoint.rs (97%) create mode 100644 crates/sail-logical-plan/src/cached_relation.rs rename crates/{sail-common-datafusion => sail-physical-plan}/src/checkpoint.rs (98%) diff --git a/Cargo.lock b/Cargo.lock index 89dd1af2d1..f52aa89fd1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7340,13 +7340,22 @@ dependencies = [ name = "sail-cache" version = "0.6.6" dependencies = [ + "async-trait", "chrono", "datafusion", + "datafusion-common", + "datafusion-datasource", + "datafusion-expr", + "futures", "log", "moka", "object_store", "sail-common", + "sail-common-datafusion", + "sail-logical-plan", + "sail-physical-plan", "thiserror 2.0.18", + "tokio", ] [[package]] @@ -7545,13 +7554,11 @@ dependencies = [ "chrono", "datafusion", "datafusion-common", - "datafusion-datasource", "datafusion-expr", "either", "futures", "lazy_static", "lexical-core", - "object_store", "parquet-variant", "parquet-variant-compute", "parquet-variant-json", @@ -7689,6 +7696,7 @@ dependencies = [ "prost-build", "rand 0.10.2", "readonly", + "sail-cache", "sail-catalog", "sail-catalog-system", "sail-common", @@ -7917,6 +7925,7 @@ dependencies = [ "datafusion-expr", "futures", "lazy_static", + "object_store", "sail-catalog", "sail-common-datafusion", "sail-logical-plan", @@ -7952,6 +7961,7 @@ dependencies = [ "rand 0.10.2", "regex", "ryu", + "sail-cache", "sail-catalog", "sail-catalog-memory", "sail-common", diff --git a/crates/sail-cache/Cargo.toml b/crates/sail-cache/Cargo.toml index cd6b240eef..59b0aaf646 100644 --- a/crates/sail-cache/Cargo.toml +++ b/crates/sail-cache/Cargo.toml @@ -8,10 +8,19 @@ workspace = true [dependencies] sail-common = { path = "../sail-common" } +sail-common-datafusion = { path = "../sail-common-datafusion" } +sail-logical-plan = { path = "../sail-logical-plan" } +sail-physical-plan = { path = "../sail-physical-plan" } +async-trait = { workspace = true } thiserror = { workspace = true } object_store = { workspace = true } datafusion = { workspace = true } +datafusion-common = { workspace = true } +datafusion-datasource = { workspace = true } +datafusion-expr = { workspace = true } chrono = { workspace = true } +futures = { workspace = true } log = { workspace = true } moka = { workspace = true } +tokio = { workspace = true } diff --git a/crates/sail-common-datafusion/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs similarity index 95% rename from crates/sail-common-datafusion/src/cached_relation.rs rename to crates/sail-cache/src/cached_relation.rs index 8e62b4d214..58c7493aa5 100644 --- a/crates/sail-common-datafusion/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -1,4 +1,3 @@ -use std::cmp::Ordering; use std::collections::hash_map::Entry; use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::fmt::Formatter; @@ -18,93 +17,30 @@ use datafusion::physical_plan::{ with_new_children_if_necessary, }; use datafusion::prelude::SessionContext; -use datafusion_common::{ - DFSchema, DFSchemaRef, DataFusionError, Result, Statistics, internal_datafusion_err, -}; +use datafusion_common::{DataFusionError, Result, Statistics, internal_datafusion_err}; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_datasource::memory::MemorySourceConfig; use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::{PartitionedFile, TableSchema}; -use datafusion_expr::{Expr, Extension, LogicalPlan, UserDefinedLogicalNodeCore}; +use datafusion_expr::{Extension, LogicalPlan}; use futures::StreamExt; use futures::future::BoxFuture; use object_store::ObjectMeta; use sail_common::spec; +use sail_common_datafusion::array::record_batch::read_record_batches; +use sail_common_datafusion::extension::{SessionExtension, SessionExtensionAccessor}; +use sail_common_datafusion::session::job::JobService; +use sail_logical_plan::cached_relation::CachedRelationNode; +use sail_physical_plan::checkpoint::LocalCheckpointExec; -use crate::array::record_batch::read_record_batches; -use crate::checkpoint::LocalCheckpointExec; -use crate::extension::{SessionExtension, SessionExtensionAccessor}; -use crate::session::checkpoint::{CheckpointStoreService, ReliableCheckpoint}; -use crate::session::job::JobService; +use crate::checkpoint::{CheckpointStoreService, ReliableCheckpoint}; #[derive(Debug, Clone)] enum CachedRelationCleanup { ObjectStorePath(String), } -#[derive(Clone, Debug, Eq, PartialEq, Hash)] -pub struct CachedRelationNode { - relation_id: String, - schema: DFSchemaRef, -} - -impl CachedRelationNode { - fn try_new(relation_id: String, schema: SchemaRef) -> Result { - let schema = Arc::new(DFSchema::try_from(schema.as_ref().clone())?); - Ok(Self { - relation_id, - schema, - }) - } - - pub fn relation_id(&self) -> &str { - &self.relation_id - } -} - -impl PartialOrd for CachedRelationNode { - fn partial_cmp(&self, other: &Self) -> Option { - self.relation_id.partial_cmp(&other.relation_id) - } -} - -impl UserDefinedLogicalNodeCore for CachedRelationNode { - fn name(&self) -> &str { - "CachedRelation" - } - - fn inputs(&self) -> Vec<&LogicalPlan> { - vec![] - } - - fn schema(&self) -> &DFSchemaRef { - &self.schema - } - - fn expressions(&self) -> Vec { - vec![] - } - - fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { - write!(f, "CachedRelation: relation_id={}", self.relation_id) - } - - fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { - if !exprs.is_empty() { - return Err(internal_datafusion_err!( - "CachedRelation does not support expressions" - )); - } - if !inputs.is_empty() { - return Err(internal_datafusion_err!( - "CachedRelation does not support inputs" - )); - } - Ok(self.clone()) - } -} - #[derive(Debug, Clone)] pub struct CachedRelation { data: Arc>, @@ -1044,7 +980,7 @@ mod tests { use datafusion::physical_plan::empty::EmptyExec; use super::*; - use crate::session::checkpoint::{CheckpointStore, ReliableCheckpoint}; + use crate::checkpoint::{CheckpointStore, ReliableCheckpoint}; #[test] fn pending_reliable_checkpoint_tracks_cleanup_path() -> Result<()> { diff --git a/crates/sail-common-datafusion/src/session/checkpoint.rs b/crates/sail-cache/src/checkpoint.rs similarity index 97% rename from crates/sail-common-datafusion/src/session/checkpoint.rs rename to crates/sail-cache/src/checkpoint.rs index 0fbec8b010..886c18d380 100644 --- a/crates/sail-common-datafusion/src/session/checkpoint.rs +++ b/crates/sail-cache/src/checkpoint.rs @@ -7,8 +7,7 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionContext; use datafusion_common::Result; use object_store::ObjectMeta; - -use crate::extension::SessionExtension; +use sail_common_datafusion::extension::SessionExtension; #[derive(Debug, Clone)] pub struct ReliableCheckpoint { diff --git a/crates/sail-cache/src/lib.rs b/crates/sail-cache/src/lib.rs index a8fe1223ac..fc4bf398d3 100644 --- a/crates/sail-cache/src/lib.rs +++ b/crates/sail-cache/src/lib.rs @@ -1,5 +1,7 @@ use log::error; +pub mod cached_relation; +pub mod checkpoint; pub mod error; pub mod file_listing_cache; pub mod file_metadata_cache; diff --git a/crates/sail-common-datafusion/Cargo.toml b/crates/sail-common-datafusion/Cargo.toml index 072691e108..b000371f13 100644 --- a/crates/sail-common-datafusion/Cargo.toml +++ b/crates/sail-common-datafusion/Cargo.toml @@ -13,7 +13,6 @@ chrono = { workspace = true } ryu = { workspace = true } datafusion = { workspace = true } datafusion-common = { workspace = true } -datafusion-datasource = { workspace = true } datafusion-expr = { workspace = true } arrow-schema = { workspace = true } thiserror = { workspace = true } @@ -29,7 +28,6 @@ pin-project-lite = { workspace = true } parquet-variant-compute = { workspace = true } parquet-variant = { workspace = true } parquet-variant-json = { workspace = true } -object_store = { workspace = true } tokio = { workspace = true } serde_arrow = { workspace = true } diff --git a/crates/sail-common-datafusion/src/lib.rs b/crates/sail-common-datafusion/src/lib.rs index 3b34e83d78..892aad6df8 100644 --- a/crates/sail-common-datafusion/src/lib.rs +++ b/crates/sail-common-datafusion/src/lib.rs @@ -1,7 +1,5 @@ pub mod array; -pub mod cached_relation; pub mod catalog; -pub mod checkpoint; pub mod column_features; pub mod datasource; pub mod display; diff --git a/crates/sail-common-datafusion/src/session/mod.rs b/crates/sail-common-datafusion/src/session/mod.rs index ee8a2b3ed6..2bde515e56 100644 --- a/crates/sail-common-datafusion/src/session/mod.rs +++ b/crates/sail-common-datafusion/src/session/mod.rs @@ -1,5 +1,4 @@ pub mod activity; -pub mod checkpoint; pub mod job; pub mod plan; pub mod repartition; diff --git a/crates/sail-execution/Cargo.toml b/crates/sail-execution/Cargo.toml index ea3a17c4e5..f5c3ae06b8 100644 --- a/crates/sail-execution/Cargo.toml +++ b/crates/sail-execution/Cargo.toml @@ -7,6 +7,7 @@ edition = { workspace = true } workspace = true [dependencies] +sail-cache = { path = "../sail-cache" } sail-common = { path = "../sail-common" } sail-common-datafusion = { path = "../sail-common-datafusion" } sail-catalog = { path = "../sail-catalog" } diff --git a/crates/sail-execution/src/proto/codec.rs b/crates/sail-execution/src/proto/codec.rs index c2a9f1888a..9de9706153 100644 --- a/crates/sail-execution/src/proto/codec.rs +++ b/crates/sail-execution/src/proto/codec.rs @@ -78,13 +78,12 @@ use datafusion_spark::function::url::try_url_decode::TryUrlDecode; use datafusion_spark::function::url::url_decode::UrlDecode; use datafusion_spark::function::url::url_encode::UrlEncode; use prost::Message; +use sail_cache::cached_relation::CachedRelationExec; use sail_catalog_system::physical_plan::SystemTableExec; use sail_common_datafusion::array::record_batch::{read_record_batches, write_record_batches}; -use sail_common_datafusion::cached_relation::CachedRelationExec; use sail_common_datafusion::catalog::{ CatalogPartitionField, LakehouseExecutionContext, PartitionTransform, }; -use sail_common_datafusion::checkpoint::{LocalCheckpointExec, ReliableCheckpointExec}; use sail_common_datafusion::datasource::PhysicalSinkMode; use sail_common_datafusion::schema_evolution::SchemaEvolutionCastColumnExpr; use sail_common_datafusion::system::catalog::SystemTable; @@ -251,6 +250,7 @@ use sail_logical_plan::range::Range; use sail_logical_plan::show_string::{ShowStringFormat, ShowStringStyle}; use sail_physical_plan::barrier::BarrierExec; use sail_physical_plan::catalog_command::CatalogCommandExec; +use sail_physical_plan::checkpoint::{LocalCheckpointExec, ReliableCheckpointExec}; use sail_physical_plan::coalesce::CoalesceExec; use sail_physical_plan::data_source::RemoteDataSourceExec; use sail_physical_plan::map_partitions::MapPartitionsExec; diff --git a/crates/sail-logical-plan/src/cached_relation.rs b/crates/sail-logical-plan/src/cached_relation.rs new file mode 100644 index 0000000000..3a319d8291 --- /dev/null +++ b/crates/sail-logical-plan/src/cached_relation.rs @@ -0,0 +1,69 @@ +use std::cmp::Ordering; +use std::fmt::Formatter; +use std::sync::Arc; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion_common::{DFSchema, DFSchemaRef, Result, internal_datafusion_err}; +use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNodeCore}; + +#[derive(Clone, Debug, Eq, PartialEq, Hash)] +pub struct CachedRelationNode { + relation_id: String, + schema: DFSchemaRef, +} + +impl CachedRelationNode { + pub fn try_new(relation_id: String, schema: SchemaRef) -> Result { + let schema = Arc::new(DFSchema::try_from(schema.as_ref().clone())?); + Ok(Self { + relation_id, + schema, + }) + } + + pub fn relation_id(&self) -> &str { + &self.relation_id + } +} + +impl PartialOrd for CachedRelationNode { + fn partial_cmp(&self, other: &Self) -> Option { + self.relation_id.partial_cmp(&other.relation_id) + } +} + +impl UserDefinedLogicalNodeCore for CachedRelationNode { + fn name(&self) -> &str { + "CachedRelation" + } + + fn inputs(&self) -> Vec<&LogicalPlan> { + vec![] + } + + fn schema(&self) -> &DFSchemaRef { + &self.schema + } + + fn expressions(&self) -> Vec { + vec![] + } + + fn fmt_for_explain(&self, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CachedRelation: relation_id={}", self.relation_id) + } + + fn with_exprs_and_inputs(&self, exprs: Vec, inputs: Vec) -> Result { + if !exprs.is_empty() { + return Err(internal_datafusion_err!( + "CachedRelation does not support expressions" + )); + } + if !inputs.is_empty() { + return Err(internal_datafusion_err!( + "CachedRelation does not support inputs" + )); + } + Ok(self.clone()) + } +} diff --git a/crates/sail-logical-plan/src/lib.rs b/crates/sail-logical-plan/src/lib.rs index fa6d6a904c..05accafe59 100644 --- a/crates/sail-logical-plan/src/lib.rs +++ b/crates/sail-logical-plan/src/lib.rs @@ -1,4 +1,5 @@ pub mod barrier; +pub mod cached_relation; pub mod check_constraints; pub mod map_partitions; pub mod merge; diff --git a/crates/sail-physical-plan/Cargo.toml b/crates/sail-physical-plan/Cargo.toml index 5a76c127c1..81da84e01a 100644 --- a/crates/sail-physical-plan/Cargo.toml +++ b/crates/sail-physical-plan/Cargo.toml @@ -15,6 +15,7 @@ datafusion = { workspace = true } datafusion-common = { workspace = true } datafusion-expr = { workspace = true } futures = { workspace = true } +object_store = { workspace = true } tokio = { workspace = true } tokio-stream = { workspace = true } lazy_static = { workspace = true } diff --git a/crates/sail-common-datafusion/src/checkpoint.rs b/crates/sail-physical-plan/src/checkpoint.rs similarity index 98% rename from crates/sail-common-datafusion/src/checkpoint.rs rename to crates/sail-physical-plan/src/checkpoint.rs index a486f9774b..fcf0806f27 100644 --- a/crates/sail-common-datafusion/src/checkpoint.rs +++ b/crates/sail-physical-plan/src/checkpoint.rs @@ -15,8 +15,9 @@ use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use futures::{StreamExt, stream}; use object_store::path::Path; use object_store::{ObjectStoreExt, PutPayload}; - -use crate::array::record_batch::{write_record_batches, write_record_batches_file}; +use sail_common_datafusion::array::record_batch::{ + write_record_batches, write_record_batches_file, +}; const PARTITION_COLUMN: &str = "partition"; const SEQUENCE_COLUMN: &str = "sequence"; diff --git a/crates/sail-physical-plan/src/lib.rs b/crates/sail-physical-plan/src/lib.rs index 524382d39f..7db8f3c26e 100644 --- a/crates/sail-physical-plan/src/lib.rs +++ b/crates/sail-physical-plan/src/lib.rs @@ -1,5 +1,6 @@ pub mod barrier; pub mod catalog_command; +pub mod checkpoint; pub mod coalesce; pub mod data_source; pub mod map_partitions; diff --git a/crates/sail-plan/Cargo.toml b/crates/sail-plan/Cargo.toml index 9d687b4a6e..c6b7843d38 100644 --- a/crates/sail-plan/Cargo.toml +++ b/crates/sail-plan/Cargo.toml @@ -7,6 +7,7 @@ edition = { workspace = true } workspace = true [dependencies] +sail-cache = { path = "../sail-cache" } sail-common = { path = "../sail-common" } sail-common-datafusion = { path = "../sail-common-datafusion" } sail-python-udf = { path = "../sail-python-udf" } diff --git a/crates/sail-plan/src/lib.rs b/crates/sail-plan/src/lib.rs index c76fd9100e..c805b62226 100644 --- a/crates/sail-plan/src/lib.rs +++ b/crates/sail-plan/src/lib.rs @@ -6,8 +6,8 @@ use datafusion::prelude::SessionContext; use datafusion_common::Result; use datafusion_common::display::{PlanType, StringifiedPlan, ToStringifiedPlan}; use datafusion_expr::LogicalPlan; +use sail_cache::cached_relation::materialize_cached_relations; use sail_common::spec; -use sail_common_datafusion::cached_relation::materialize_cached_relations; use sail_common_datafusion::rename::physical_plan::rename_physical_plan; use crate::config::PlanConfig; diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index 226296f025..78cf4c4bee 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/crates/sail-plan/src/resolver/query/mod.rs @@ -1,7 +1,7 @@ use async_recursion::async_recursion; use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; +use sail_cache::cached_relation::CachedRelationRegistry; use sail_common::spec; -use sail_common_datafusion::cached_relation::CachedRelationRegistry; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::rename::logical_plan::rename_logical_plan; diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs index 708a9c6f5e..30ca4add01 100644 --- a/crates/sail-session/src/checkpoint.rs +++ b/crates/sail-session/src/checkpoint.rs @@ -8,12 +8,12 @@ use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion::prelude::SessionContext; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; use futures::StreamExt; +use sail_cache::checkpoint::{CheckpointStore, ReliableCheckpoint}; use sail_common_datafusion::array::record_batch::write_record_batches_file; -use sail_common_datafusion::checkpoint::ReliableCheckpointExec; use sail_common_datafusion::extension::SessionExtensionAccessor; -use sail_common_datafusion::session::checkpoint::{CheckpointStore, ReliableCheckpoint}; use sail_common_datafusion::session::job::JobService; use sail_object_store::{delete_object_store_prefix, resolve_object_store_path}; +use sail_physical_plan::checkpoint::ReliableCheckpointExec; #[derive(Debug, Default)] pub struct ObjectStoreCheckpointStore; diff --git a/crates/sail-session/src/planner.rs b/crates/sail-session/src/planner.rs index 91fb31fb51..5899d7e33b 100644 --- a/crates/sail-session/src/planner.rs +++ b/crates/sail-session/src/planner.rs @@ -11,8 +11,8 @@ use datafusion::physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, Phy use datafusion_common::{DFSchema, internal_err}; use datafusion_expr::{Expr, LogicalPlan, UserDefinedLogicalNode}; use datafusion_physical_expr::{Partitioning, create_physical_sort_exprs}; +use sail_cache::cached_relation::CachedRelationRegistry; use sail_catalog_system::planner::SystemTablePhysicalPlanner; -use sail_common_datafusion::cached_relation::{CachedRelationNode, CachedRelationRegistry}; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::logical_rewriter::LogicalRewriter; use sail_common_datafusion::rename::physical_plan::rename_projected_physical_plan; @@ -26,6 +26,7 @@ use sail_data_source::listing::planner::ListingPhysicalPlanner; use sail_delta_lake::physical::DeltaPhysicalPlanner; use sail_iceberg::IcebergPhysicalPlanner; use sail_logical_plan::barrier::BarrierNode; +use sail_logical_plan::cached_relation::CachedRelationNode; use sail_logical_plan::map_partitions::MapPartitionsNode; use sail_logical_plan::monotonic_id::MonotonicIdNode; use sail_logical_plan::range::RangeNode; diff --git a/crates/sail-session/src/session_factory/server.rs b/crates/sail-session/src/session_factory/server.rs index 298339ae96..e5b6ecbd08 100644 --- a/crates/sail-session/src/session_factory/server.rs +++ b/crates/sail-session/src/session_factory/server.rs @@ -9,13 +9,13 @@ use datafusion::execution::{SessionState, SessionStateBuilder}; use datafusion::functions_aggregate::first_last::first_value_udaf; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_expr::registry::FunctionRegistry; +use sail_cache::cached_relation::CachedRelationRegistry; +use sail_cache::checkpoint::CheckpointStoreService; use sail_catalog::provider::CatalogCacheManager; use sail_catalog_system::service::SystemTableService; use sail_common::config::{AppConfig, ExecutionMode}; use sail_common::runtime::RuntimeHandle; -use sail_common_datafusion::cached_relation::CachedRelationRegistry; use sail_common_datafusion::session::activity::ActivityTracker; -use sail_common_datafusion::session::checkpoint::CheckpointStoreService; use sail_common_datafusion::session::job::{JobRunner, JobService}; use sail_common_datafusion::session::repartition::RepartitionBufferConfig; use sail_delta_lake::session_extension::DeltaTableCache; diff --git a/crates/sail-session/src/session_factory/worker.rs b/crates/sail-session/src/session_factory/worker.rs index eacfb7bbce..c608f28388 100644 --- a/crates/sail-session/src/session_factory/worker.rs +++ b/crates/sail-session/src/session_factory/worker.rs @@ -5,12 +5,9 @@ use datafusion::execution::SessionStateBuilder; use datafusion::prelude::{SessionConfig, SessionContext}; use sail_common::config::AppConfig; use sail_common::runtime::RuntimeHandle; -use sail_common_datafusion::cached_relation::CachedRelationRegistry; -use sail_common_datafusion::session::checkpoint::CheckpointStoreService; use sail_common_datafusion::session::repartition::RepartitionBufferConfig; use sail_delta_lake::session_extension::DeltaTableCache; -use crate::checkpoint::ObjectStoreCheckpointStore; use crate::runtime::RuntimeEnvFactory; use crate::session_factory::SessionFactory; @@ -38,10 +35,6 @@ impl SessionFactory<()> for WorkerSessionFactory { // when decoding the execution plan. let config = SessionConfig::default() .with_extension(Arc::new(DeltaTableCache::default())) - .with_extension(Arc::new(CachedRelationRegistry::default())) - .with_extension(Arc::new(CheckpointStoreService::new(Arc::new( - ObjectStoreCheckpointStore, - )))) .with_extension(Arc::new(RepartitionBufferConfig::new( self.repartition_buffer_size, ))); diff --git a/crates/sail-session/src/session_manager/actor/handler.rs b/crates/sail-session/src/session_manager/actor/handler.rs index 6b94c1fe8a..b2764fa029 100644 --- a/crates/sail-session/src/session_manager/actor/handler.rs +++ b/crates/sail-session/src/session_manager/actor/handler.rs @@ -5,7 +5,7 @@ use datafusion::prelude::SessionContext; use fastrace::Span; use fastrace::collector::SpanContext; use log::{info, warn}; -use sail_common_datafusion::cached_relation::cleanup_cached_relations; +use sail_cache::cached_relation::cleanup_cached_relations; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::activity::ActivityTracker; use sail_common_datafusion::session::job::JobService; diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index ed33b8effa..6e8dbbc21a 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -9,10 +9,10 @@ use fastrace::collector::SpanContext; use fastrace::future::FutureExt; use futures::stream; use log::{debug, warn}; -use sail_common::spec; -use sail_common_datafusion::cached_relation::{ +use sail_cache::cached_relation::{ CachedRelation, CachedRelationRegistry, cleanup_checkpoint_path, remove_cached_relation, }; +use sail_common::spec; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; use sail_plan::{resolve_and_execute_plan, resolve_physical_plan}; From 473752a69117973440027a069d6b3bdf3dfbfde0 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Thu, 16 Jul 2026 23:48:38 +0800 Subject: [PATCH 23/40] test: capture checkpoint parity gaps --- crates/sail-cache/src/cached_relation.rs | 1 + crates/sail-physical-plan/src/checkpoint.rs | 2 + .../src/session_manager/actor/handler.rs | 1 + .../tests/spark/execution/test_checkpoint.py | 249 ++++++++++++++++++ 4 files changed, 253 insertions(+) create mode 100644 python/pysail/tests/spark/execution/test_checkpoint.py diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index 58c7493aa5..3e3bfdd730 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -98,6 +98,7 @@ enum CachedRelationPendingTarget { pub struct CachedRelationExec { input: Arc, properties: Arc, + // FIXME: Ensure completed distributed jobs release this lease so retired checkpoints can be cleaned. relation_lease: Option, } diff --git a/crates/sail-physical-plan/src/checkpoint.rs b/crates/sail-physical-plan/src/checkpoint.rs index fcf0806f27..ff660fe407 100644 --- a/crates/sail-physical-plan/src/checkpoint.rs +++ b/crates/sail-physical-plan/src/checkpoint.rs @@ -271,8 +271,10 @@ impl ExecutionPlan for ReliableCheckpointExec { let output_schema = self.schema(); let stream_schema = Arc::clone(&output_schema); let store = context.runtime_env().object_store(&self.object_store_url)?; + // FIXME: Publish attempt-scoped temporary files atomically so stale retries cannot overwrite successful output. let location = self.path.clone().join(format!("part-{partition:05}.arrow")); let output = stream::once(async move { + // FIXME: Stream Arrow IPC to object storage instead of buffering an entire partition. let mut batches = vec![]; while let Some(batch) = input.next().await { batches.push(batch?); diff --git a/crates/sail-session/src/session_manager/actor/handler.rs b/crates/sail-session/src/session_manager/actor/handler.rs index b2764fa029..8c9e8a5493 100644 --- a/crates/sail-session/src/session_manager/actor/handler.rs +++ b/crates/sail-session/src/session_manager/actor/handler.rs @@ -329,6 +329,7 @@ impl SessionManagerActor { let context = context.clone(); ctx.spawn(async move { service.runner().stop(tx).await; + // FIXME: Await runner shutdown before removing checkpoint data still visible to tasks. if let Err(error) = cleanup_cached_relations(&context).await { warn!("failed to clean cached relations for session {session_id}: {error}"); } diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py new file mode 100644 index 0000000000..b024497c90 --- /dev/null +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -0,0 +1,249 @@ +import contextlib +import gc +import shutil +import time +import uuid + +import pytest +from pyspark import StorageLevel +from pyspark.errors import PySparkException +from pyspark.sql import SparkSession +from pyspark.sql.connect import proto +from pyspark.sql.connect.plan import Checkpoint + +from pysail.testing.spark.session import ( + configure_spark_session, + patch_spark_connect_session, + spark_connect_server, +) +from pysail.testing.spark.utils.common import is_jvm_spark, pyspark_version + +pytestmark = [ + pytest.mark.skipif( + pyspark_version() < (4,), + reason="checkpoint and localCheckpoint require PySpark Connect 4+", + ), + pytest.mark.xfail( + not is_jvm_spark(), + reason="Known Sail checkpoint parity bugs", + strict=True, + ), +] + + +SMALL_PAYLOAD_ROWS = 16 * 1024 +LARGE_PAYLOAD_ROWS = 64 * 1024 +PAYLOAD_BYTES = 1024 +DUPLICATE_COLUMN_ROWS = 3 +CLEANUP_ROWS = 32 + + +@contextlib.contextmanager +def _checkpoint_spark_session(remote, checkpoint_path): + builder = ( + SparkSession.builder.appName(f"test-{uuid.uuid4().hex}") + .remote(remote) + .config("spark.checkpoint.dir", str(checkpoint_path)) + ) + session = builder.getOrCreate() if remote.startswith("local") else builder.create() + configure_spark_session(session) + patch_spark_connect_session(session) + try: + yield session + finally: + session.stop() + + +@pytest.fixture(scope="module") +def remote(): + envs = None if is_jvm_spark() else {"SAIL_MODE": "local-cluster"} + with spark_connect_server(envs=envs) as server: + yield server.remote + + +@pytest.fixture(scope="module") +def checkpoint_path(tmp_path_factory): + return tmp_path_factory.mktemp("execution-checkpoints") + + +@pytest.fixture(scope="module") +def spark(remote, checkpoint_path): + with _checkpoint_spark_session(remote, checkpoint_path) as session: + yield session + + +def _payload_dataframe(spark, rows, partitions=4): + return spark.range(rows, numPartitions=partitions).selectExpr( + "id", + f"repeat('x', {PAYLOAD_BYTES}) AS payload", + ) + + +def _assert_payload(checkpointed, rows): + assert checkpointed.count() == rows + actual = checkpointed.selectExpr("sum(length(payload)) AS payload_bytes").first() + assert actual.payload_bytes == rows * PAYLOAD_BYTES + + +def _assert_large_payload(checkpointed, rows): + actual = checkpointed.selectExpr( + "count(*) AS row_count", + "sum(length(payload)) AS payload_bytes", + ).first() + assert actual.row_count == rows + assert actual.payload_bytes == rows * PAYLOAD_BYTES + + +@pytest.fixture(scope="module") +def retry_spark(spark, checkpoint_path): + if is_jvm_spark(): + yield spark + return + + # The retry state is embedded-only, so close the cluster session before switching servers. + gc.collect() + spark.stop() + with ( + spark_connect_server(envs={"SAIL_MODE": "local"}) as server, + _checkpoint_spark_session(server.remote, checkpoint_path / "retry") as session, + ): + yield session + + +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("kind", ["local", "reliable"]) +def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): + source = spark.range(DUPLICATE_COLUMN_ROWS, numPartitions=2).selectExpr( + "id AS value", + "id + 10 AS value", + ) + + checkpointed = source.localCheckpoint(eager=eager) if kind == "local" else source.checkpoint(eager=eager) + + assert [field.name for field in checkpointed.schema.fields] == ["value", "value"] + assert checkpointed.count() == DUPLICATE_COLUMN_ROWS + + +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize( + "storage_level", + [ + pytest.param(None, id="default"), + pytest.param(StorageLevel.MEMORY_ONLY, id="memory-only"), + pytest.param(StorageLevel.DISK_ONLY, id="disk-only"), + pytest.param(StorageLevel.MEMORY_AND_DISK, id="memory-and-disk"), + ], +) +def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, storage_level): + source = _payload_dataframe(spark, SMALL_PAYLOAD_ROWS, partitions=1) + assert source.count() == SMALL_PAYLOAD_ROWS + + if storage_level is None: + checkpointed = source.localCheckpoint(eager=eager) + else: + checkpointed = source.localCheckpoint(eager=eager, storageLevel=storage_level) + + _assert_payload(checkpointed, SMALL_PAYLOAD_ROWS) + + +def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark): + checkpointed = _payload_dataframe(spark, LARGE_PAYLOAD_ROWS).localCheckpoint() + + _assert_large_payload(checkpointed, LARGE_PAYLOAD_ROWS) + + +def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): + if is_jvm_spark(): + # JVM Spark fixes the checkpoint root at server startup. + safe = spark.range(CLEANUP_ROWS, numPartitions=4).checkpoint() + assert safe.count() == CLEANUP_ROWS + + other = spark.range(1).checkpoint(eager=False) + del other + gc.collect() + + assert safe.count() == CLEANUP_ROWS + return + + original_checkpoint_path = spark.conf.get("spark.checkpoint.dir") + root = tmp_path / "checkpoint-scope" + safe_root = f"{root.as_uri()}/safe" + overlapping_root = f"{root.as_uri()}?scope=other" + + try: + spark.conf.set("spark.checkpoint.dir", safe_root) + safe = spark.range(CLEANUP_ROWS, numPartitions=4).checkpoint() + assert safe.count() == CLEANUP_ROWS + safe_files = list(root.rglob("*.arrow")) + + spark.conf.set("spark.checkpoint.dir", overlapping_root) + other = spark.range(1).checkpoint(eager=False) + del other + gc.collect() + + deadline = time.monotonic() + 5 + while any(path.exists() for path in safe_files) and time.monotonic() < deadline: + time.sleep(0.1) + + assert safe.count() == CLEANUP_ROWS + finally: + spark.conf.set("spark.checkpoint.dir", original_checkpoint_path) + + +def test_checkpoint_command_can_reattach(spark): + client = spark._client # noqa: SLF001 + command = Checkpoint(spark.range(3)._plan, local=True, eager=False).command(client) # noqa: SLF001 + request = client._execute_plan_request_with_metadata() # noqa: SLF001 + request.operation_id = str(uuid.uuid4()) + request.plan.command.CopyFrom(command) + request.request_options.append( + proto.ExecutePlanRequest.RequestOption( + reattach_options=proto.ReattachOptions(reattachable=True), + ) + ) + + call = client._stub.ExecutePlan(request, metadata=client._builder.metadata()) # noqa: SLF001 + first = next(iter(call)) + assert first.HasField("checkpoint_command_result") + call.cancel() + + reattach = proto.ReattachExecuteRequest( + session_id=request.session_id, + client_observed_server_side_session_id=first.server_side_session_id, + user_context=request.user_context, + operation_id=request.operation_id, + client_type=request.client_type, + last_response_id=first.response_id, + ) + responses = list( + client._stub.ReattachExecute( # noqa: SLF001 + reattach, + metadata=client._builder.metadata(), # noqa: SLF001 + ) + ) + + assert any(response.HasField("result_complete") for response in responses) + + +@pytest.mark.parametrize("kind", ["local", "reliable"]) +def test_lazy_checkpoint_can_retry_after_source_recovers(retry_spark, tmp_path, kind): + row_count = 50_000 + source_path = tmp_path / f"source-{kind}" + retry_spark.range(row_count, numPartitions=8).write.parquet(str(source_path)) + source = retry_spark.read.parquet(str(source_path)) + + checkpointed = source.localCheckpoint(eager=False) if kind == "local" else source.checkpoint(eager=False) + + source_files = sorted(source_path.glob("*.parquet")) + assert len(source_files) > 1 + missing = source_files[len(source_files) // 2] + hidden = tmp_path / f"{kind}-{missing.name}.missing" + shutil.move(missing, hidden) + try: + with pytest.raises(PySparkException): + checkpointed.count() + finally: + shutil.move(hidden, missing) + + assert source.count() == row_count + assert checkpointed.count() == row_count From 276fe882aa3fe660821da8092c61050d3a97b8b3 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Thu, 16 Jul 2026 23:56:03 +0800 Subject: [PATCH 24/40] fix: isolate checkpoint cleanup paths --- .../src/service/plan_executor.rs | 65 ++++++++++++++++--- .../tests/spark/execution/test_checkpoint.py | 27 ++++---- 2 files changed, 71 insertions(+), 21 deletions(-) diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 6e8dbbc21a..fd3471e45e 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -3,6 +3,7 @@ use std::pin::Pin; use std::task::{Context, Poll}; use datafusion::arrow::compute::concat_batches; +use datafusion::datasource::listing::ListingTableUrl; use datafusion::prelude::SessionContext; use fastrace::Span; use fastrace::collector::SpanContext; @@ -673,15 +674,35 @@ fn checkpoint_path(spark: &SparkSession, relation_id: &str) -> SparkResult SparkResult { + validate_checkpoint_directory(root)?; + validate_checkpoint_path_segment("session ID", session_id)?; validate_checkpoint_path_segment("relation ID", relation_id)?; - Ok(format!( - "{}/sail-checkpoints/{}/{}", - root.trim_end_matches('/'), - spark.session_id(), - relation_id - )) + + let directory = format!("{}/", root.trim_end_matches('/')); + let checkpoint_root = ListingTableUrl::parse(directory)?; + let root_url = checkpoint_root.get_url(); + if checkpoint_root.get_glob().is_some() + || root_url.query().is_some() + || root_url.fragment().is_some() + { + return Err(SparkError::invalid( + "checkpoint directory cannot contain a query, fragment, or glob", + )); + } + + let mut checkpoint_url = root_url.clone(); + checkpoint_url + .path_segments_mut() + .map_err(|()| SparkError::invalid("checkpoint directory must be a hierarchical URL"))? + .pop_if_empty() + .push("sail-checkpoints") + .push(session_id) + .push(relation_id); + Ok(checkpoint_url.into()) } fn validate_checkpoint_directory(root: &str) -> SparkResult<()> { @@ -706,6 +727,34 @@ fn validate_checkpoint_path_segment(name: &str, value: &str) -> SparkResult<()> } } +#[cfg(test)] +mod checkpoint_path_tests { + use super::*; + + #[test] + fn checkpoint_path_appends_encoded_segments() -> SparkResult<()> { + let path = build_checkpoint_path("memory:///checkpoint-root", "session?one", "relation")?; + + assert_eq!( + path, + "memory:///checkpoint-root/sail-checkpoints/session%3Fone/relation" + ); + Ok(()) + } + + #[test] + fn checkpoint_path_rejects_url_query() { + let error = build_checkpoint_path( + "memory:///checkpoint-root?scope=other", + "session", + "relation", + ) + .expect_err("checkpoint root query must be rejected"); + + assert!(error.to_string().contains("cannot contain a query")); + } +} + async fn cleanup_checkpoint_after_error( ctx: &SessionContext, path: &str, diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index b024497c90..9fc76af461 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -1,7 +1,6 @@ import contextlib import gc import shutil -import time import uuid import pytest @@ -23,13 +22,14 @@ pyspark_version() < (4,), reason="checkpoint and localCheckpoint require PySpark Connect 4+", ), - pytest.mark.xfail( - not is_jvm_spark(), - reason="Known Sail checkpoint parity bugs", - strict=True, - ), ] +SAIL_XFAIL = pytest.mark.xfail( + not is_jvm_spark(), + reason="Known Sail checkpoint parity bug", + strict=True, +) + SMALL_PAYLOAD_ROWS = 16 * 1024 LARGE_PAYLOAD_ROWS = 64 * 1024 @@ -112,6 +112,7 @@ def retry_spark(spark, checkpoint_path): @pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) @pytest.mark.parametrize("kind", ["local", "reliable"]) +@SAIL_XFAIL def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): source = spark.range(DUPLICATE_COLUMN_ROWS, numPartitions=2).selectExpr( "id AS value", @@ -134,6 +135,7 @@ def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): pytest.param(StorageLevel.MEMORY_AND_DISK, id="memory-and-disk"), ], ) +@SAIL_XFAIL def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, storage_level): source = _payload_dataframe(spark, SMALL_PAYLOAD_ROWS, partitions=1) assert source.count() == SMALL_PAYLOAD_ROWS @@ -146,6 +148,7 @@ def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, s _assert_payload(checkpointed, SMALL_PAYLOAD_ROWS) +@SAIL_XFAIL def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark): checkpointed = _payload_dataframe(spark, LARGE_PAYLOAD_ROWS).localCheckpoint() @@ -177,19 +180,16 @@ def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): safe_files = list(root.rglob("*.arrow")) spark.conf.set("spark.checkpoint.dir", overlapping_root) - other = spark.range(1).checkpoint(eager=False) - del other - gc.collect() - - deadline = time.monotonic() + 5 - while any(path.exists() for path in safe_files) and time.monotonic() < deadline: - time.sleep(0.1) + with pytest.raises(PySparkException, match="checkpoint directory cannot contain"): + spark.range(1).checkpoint(eager=False) + assert all(path.exists() for path in safe_files) assert safe.count() == CLEANUP_ROWS finally: spark.conf.set("spark.checkpoint.dir", original_checkpoint_path) +@SAIL_XFAIL def test_checkpoint_command_can_reattach(spark): client = spark._client # noqa: SLF001 command = Checkpoint(spark.range(3)._plan, local=True, eager=False).command(client) # noqa: SLF001 @@ -226,6 +226,7 @@ def test_checkpoint_command_can_reattach(spark): @pytest.mark.parametrize("kind", ["local", "reliable"]) +@SAIL_XFAIL def test_lazy_checkpoint_can_retry_after_source_recovers(retry_spark, tmp_path, kind): row_count = 50_000 source_path = tmp_path / f"source-{kind}" From 4ee4a384acbb24b7a6cde3f2ed0f319652ea87b1 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 00:02:07 +0800 Subject: [PATCH 25/40] fix: refresh cached relation partitioning after rewrites --- crates/sail-cache/src/cached_relation.rs | 26 +++++++++++++++++-- .../tests/spark/execution/test_checkpoint.py | 6 ++--- 2 files changed, 27 insertions(+), 5 deletions(-) diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index 3e3bfdd730..bf74382e19 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -164,9 +164,15 @@ impl ExecutionPlan for CachedRelationExec { "CachedRelationExec must have exactly one child" )); } + let input = Arc::clone(&children[0]); Ok(Arc::new(Self { - input: Arc::clone(&children[0]), - properties: Arc::clone(&self.properties), + properties: Arc::new( + self.properties + .as_ref() + .clone() + .with_partitioning(input.output_partitioning().clone()), + ), + input, relation_lease: self.relation_lease.clone(), })) } @@ -983,6 +989,22 @@ mod tests { use super::*; use crate::checkpoint::{CheckpointStore, ReliableCheckpoint}; + #[test] + fn cached_relation_rewrite_updates_partitioning() -> Result<()> { + let schema = Arc::new(Schema::empty()); + let original: Arc = + Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(2)); + let properties = checkpoint_plan_properties(&original); + let cached = Arc::new(CachedRelationExec::new(original, properties)); + let rewritten_input: Arc = + Arc::new(EmptyExec::new(schema).with_partitions(10)); + + let rewritten = cached.with_new_children(vec![rewritten_input])?; + + assert_eq!(rewritten.output_partitioning().partition_count(), 10); + Ok(()) + } + #[test] fn pending_reliable_checkpoint_tracks_cleanup_path() -> Result<()> { let path = "memory://checkpoint-root/session/relation".to_string(); diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index 9fc76af461..33470385e4 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -126,6 +126,7 @@ def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): @pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("partitions", [1, 2, 4, 8, 10]) @pytest.mark.parametrize( "storage_level", [ @@ -135,9 +136,8 @@ def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): pytest.param(StorageLevel.MEMORY_AND_DISK, id="memory-and-disk"), ], ) -@SAIL_XFAIL -def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, storage_level): - source = _payload_dataframe(spark, SMALL_PAYLOAD_ROWS, partitions=1) +def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, partitions, storage_level): + source = _payload_dataframe(spark, SMALL_PAYLOAD_ROWS, partitions=partitions) assert source.count() == SMALL_PAYLOAD_ROWS if storage_level is None: From 952f57fec8b98124a9b199e40ab85f55d2b36b3b Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 00:06:05 +0800 Subject: [PATCH 26/40] fix: reset lazy checkpoint plans before retry --- crates/sail-cache/src/cached_relation.rs | 8 ++++---- python/pysail/tests/spark/execution/test_checkpoint.py | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index bf74382e19..b37da07330 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -10,7 +10,7 @@ use datafusion::datasource::physical_plan::ArrowSource; use datafusion::execution::disk_manager::RefCountedTempFile; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; -use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType, reset_plan_states}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, @@ -469,13 +469,13 @@ impl SessionExtension for CachedRelationRegistry { impl CachedRelationPending { async fn materialize(&self, ctx: &SessionContext) -> Result { + let plan = reset_plan_states(Arc::clone(&self.plan))?; match &self.target { CachedRelationPendingTarget::Local { storage_level } => { - materialize_local_checkpoint(ctx, Arc::clone(&self.plan), storage_level.clone()) - .await + materialize_local_checkpoint(ctx, plan, storage_level.clone()).await } CachedRelationPendingTarget::Reliable { path } => { - materialize_reliable_checkpoint(ctx, Arc::clone(&self.plan), path).await + materialize_reliable_checkpoint(ctx, plan, path).await } } } diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index 33470385e4..be82f4a1eb 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -226,7 +226,6 @@ def test_checkpoint_command_can_reattach(spark): @pytest.mark.parametrize("kind", ["local", "reliable"]) -@SAIL_XFAIL def test_lazy_checkpoint_can_retry_after_source_recovers(retry_spark, tmp_path, kind): row_count = 50_000 source_path = tmp_path / f"source-{kind}" From 15a3024bdf9f39ea4dd8ce7b598330446493a161 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 00:11:29 +0800 Subject: [PATCH 27/40] fix: resolve duplicate checkpoint columns by position --- crates/sail-cache/src/cached_relation.rs | 58 ++++++++++++++++--- crates/sail-plan/src/resolver/query/mod.rs | 7 +-- crates/sail-session/src/planner.rs | 12 +++- .../tests/spark/execution/test_checkpoint.py | 17 ++++-- 4 files changed, 76 insertions(+), 18 deletions(-) diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index b37da07330..8452e01165 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -30,6 +30,7 @@ use object_store::ObjectMeta; use sail_common::spec; use sail_common_datafusion::array::record_batch::read_record_batches; use sail_common_datafusion::extension::{SessionExtension, SessionExtensionAccessor}; +use sail_common_datafusion::rename::schema::rename_schema; use sail_common_datafusion::session::job::JobService; use sail_logical_plan::cached_relation::CachedRelationNode; use sail_physical_plan::checkpoint::LocalCheckpointExec; @@ -320,13 +321,25 @@ impl CachedRelation { } } - pub async fn to_logical_plan(&self, relation_id: &str) -> Result { + pub async fn schema(&self) -> SchemaRef { + let data = self.data.lock().await; + match &*data { + CachedRelationData::Materialized(materialized) => materialized.schema(), + CachedRelationData::Pending(pending) => pending.plan.schema(), + } + } + + pub async fn to_logical_plan( + &self, + relation_id: &str, + names: &[String], + ) -> Result { let data = self.data.lock().await; match &*data { CachedRelationData::Materialized(materialized) => { - materialized.to_logical_plan(relation_id) + materialized.to_logical_plan(relation_id, names) } - CachedRelationData::Pending(pending) => pending.to_logical_plan(relation_id), + CachedRelationData::Pending(pending) => pending.to_logical_plan(relation_id, names), } } @@ -480,11 +493,12 @@ impl CachedRelationPending { } } - fn to_logical_plan(&self, relation_id: &str) -> Result { + fn to_logical_plan(&self, relation_id: &str, names: &[String]) -> Result { + let schema = rename_schema(self.plan.schema().as_ref(), names)?; Ok(LogicalPlan::Extension(Extension { node: Arc::new(CachedRelationNode::try_new( relation_id.to_string(), - self.plan.schema(), + schema, )?), })) } @@ -505,11 +519,12 @@ impl CachedRelationMaterialized { } } - fn to_logical_plan(&self, relation_id: &str) -> Result { + fn to_logical_plan(&self, relation_id: &str, names: &[String]) -> Result { + let schema = rename_schema(self.schema().as_ref(), names)?; Ok(LogicalPlan::Extension(Extension { node: Arc::new(CachedRelationNode::try_new( relation_id.to_string(), - self.schema(), + schema, )?), })) } @@ -983,7 +998,7 @@ mod tests { use std::sync::Mutex; use async_trait::async_trait; - use datafusion::arrow::datatypes::Schema; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_plan::empty::EmptyExec; use super::*; @@ -1005,6 +1020,33 @@ mod tests { Ok(()) } + #[tokio::test] + async fn cached_relation_logical_plan_renames_duplicate_fields_by_position() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("value", DataType::Int64, true), + Field::new("value", DataType::Int64, true), + ])); + let plan: Arc = Arc::new(EmptyExec::new(schema)); + let storage_level = "MEMORY_AND_DISK" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?; + let relation = CachedRelation::new_pending_local_checkpoint(plan, storage_level); + let names = vec!["field-0".to_string(), "field-1".to_string()]; + + let logical_plan = relation.to_logical_plan("relation", &names).await?; + + assert_eq!( + logical_plan + .schema() + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(), + names + ); + Ok(()) + } + #[test] fn pending_reliable_checkpoint_tracks_cleanup_path() -> Result<()> { let path = "memory://checkpoint-root/session/relation".to_string(); diff --git a/crates/sail-plan/src/resolver/query/mod.rs b/crates/sail-plan/src/resolver/query/mod.rs index 78cf4c4bee..0f62a41b04 100644 --- a/crates/sail-plan/src/resolver/query/mod.rs +++ b/crates/sail-plan/src/resolver/query/mod.rs @@ -3,7 +3,6 @@ use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; use sail_cache::cached_relation::CachedRelationRegistry; use sail_common::spec; use sail_common_datafusion::extension::SessionExtensionAccessor; -use sail_common_datafusion::rename::logical_plan::rename_logical_plan; use crate::error::{PlanError, PlanResult}; use crate::resolver::PlanResolver; @@ -235,9 +234,9 @@ impl PlanResolver<'_> { let relation = registry.get(&relation_id)?.ok_or_else(|| { PlanError::invalid(format!("No DataFrame with id {relation_id} is found")) })?; - let plan = relation.to_logical_plan(&relation_id).await?; - let names = state.register_fields(plan.schema().inner().fields()); - rename_logical_plan(plan, &names)? + let schema = relation.schema().await; + let names = state.register_fields(schema.fields()); + relation.to_logical_plan(&relation_id, &names).await? } QueryNode::CommonInlineUserDefinedTableFunction(udtf) => { self.resolve_query_common_inline_udtf(udtf, state).await? diff --git a/crates/sail-session/src/planner.rs b/crates/sail-session/src/planner.rs index 5899d7e33b..126f97880c 100644 --- a/crates/sail-session/src/planner.rs +++ b/crates/sail-session/src/planner.rs @@ -15,7 +15,9 @@ use sail_cache::cached_relation::CachedRelationRegistry; use sail_catalog_system::planner::SystemTablePhysicalPlanner; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::logical_rewriter::LogicalRewriter; -use sail_common_datafusion::rename::physical_plan::rename_projected_physical_plan; +use sail_common_datafusion::rename::physical_plan::{ + rename_physical_plan, rename_projected_physical_plan, +}; use sail_common_datafusion::streaming::event::schema::{ to_flow_event_field_names, to_flow_event_projection, }; @@ -111,7 +113,13 @@ impl ExtensionPlanner for ExtensionPhysicalPlanner { node.relation_id() )) })?; - relation.to_physical_plan(node.relation_id()).await? + let plan = relation.to_physical_plan(node.relation_id()).await?; + let names = UserDefinedLogicalNode::schema(node) + .fields() + .iter() + .map(|field| field.name().to_string()) + .collect::>(); + rename_physical_plan(plan, &names)? } else if let Some(node) = node.as_any().downcast_ref::() { let schema = UserDefinedLogicalNode::schema(node).inner().clone(); let projection = (0..schema.fields().len()).collect(); diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index be82f4a1eb..b3ec655f3d 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -110,10 +110,7 @@ def retry_spark(spark, checkpoint_path): yield session -@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) -@pytest.mark.parametrize("kind", ["local", "reliable"]) -@SAIL_XFAIL -def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): +def _assert_checkpoint_preserves_duplicate_column_names(spark, eager, kind): source = spark.range(DUPLICATE_COLUMN_ROWS, numPartitions=2).selectExpr( "id AS value", "id + 10 AS value", @@ -125,6 +122,12 @@ def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): assert checkpointed.count() == DUPLICATE_COLUMN_ROWS +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("kind", ["local", "reliable"]) +def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): + _assert_checkpoint_preserves_duplicate_column_names(spark, eager, kind) + + @pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) @pytest.mark.parametrize("partitions", [1, 2, 4, 8, 10]) @pytest.mark.parametrize( @@ -225,6 +228,12 @@ def test_checkpoint_command_can_reattach(spark): assert any(response.HasField("result_complete") for response in responses) +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("kind", ["local", "reliable"]) +def test_embedded_checkpoint_preserves_duplicate_column_names(retry_spark, eager, kind): + _assert_checkpoint_preserves_duplicate_column_names(retry_spark, eager, kind) + + @pytest.mark.parametrize("kind", ["local", "reliable"]) def test_lazy_checkpoint_can_retry_after_source_recovers(retry_spark, tmp_path, kind): row_count = 50_000 From 7b6de7dd7fb4ce2c039b686114244e9ed4eea59d Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 00:20:57 +0800 Subject: [PATCH 28/40] fix: release checkpoint leases from completed jobs --- crates/sail-cache/src/cached_relation.rs | 44 ++++++++++++++++++- .../src/driver/job_scheduler/core.rs | 9 ++++ crates/sail-execution/src/job_graph/mod.rs | 4 ++ .../tests/spark/execution/test_checkpoint.py | 33 ++++++++++++++ 4 files changed, 89 insertions(+), 1 deletion(-) diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index 8452e01165..9320e1edc9 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -17,6 +17,7 @@ use datafusion::physical_plan::{ with_new_children_if_necessary, }; use datafusion::prelude::SessionContext; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{DataFusionError, Result, Statistics, internal_datafusion_err}; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; @@ -99,7 +100,6 @@ enum CachedRelationPendingTarget { pub struct CachedRelationExec { input: Arc, properties: Arc, - // FIXME: Ensure completed distributed jobs release this lease so retired checkpoints can be cleaned. relation_lease: Option, } @@ -200,6 +200,21 @@ impl ExecutionPlan for CachedRelationExec { } } +pub fn release_cached_relation_leases( + plan: Arc, +) -> Result> { + plan.transform_up(|plan| { + let Some(cached) = plan.downcast_ref::() else { + return Ok(Transformed::no(plan)); + }; + Ok(Transformed::yes(Arc::new(CachedRelationExec::new( + Arc::clone(cached.input()), + Arc::clone(cached.properties()), + )))) + }) + .data() +} + #[derive(Debug, Clone)] struct PendingCachedRelationExec { relation_id: String, @@ -1020,6 +1035,33 @@ mod tests { Ok(()) } + #[test] + fn completed_plan_can_release_cached_relation_lease() -> Result<()> { + let schema = Arc::new(Schema::empty()); + let input: Arc = Arc::new(EmptyExec::new(schema)); + let properties = checkpoint_plan_properties(&input); + let relation = CachedRelation::new_pending_local_checkpoint( + Arc::clone(&input), + "MEMORY_AND_DISK" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?, + ); + let mut completed_plan: Arc = Arc::new( + CachedRelationExec::with_relation_lease(input, properties, relation.clone()), + ); + assert!(!relation.is_exclusively_owned()); + + completed_plan = release_cached_relation_leases(completed_plan)?; + + assert!(relation.is_exclusively_owned()); + assert!( + completed_plan + .downcast_ref::() + .is_some_and(|cached| cached.relation_lease.is_none()) + ); + Ok(()) + } + #[tokio::test] async fn cached_relation_logical_plan_renames_duplicate_fields_by_position() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/crates/sail-execution/src/driver/job_scheduler/core.rs b/crates/sail-execution/src/driver/job_scheduler/core.rs index e9363d6676..70bee01711 100644 --- a/crates/sail-execution/src/driver/job_scheduler/core.rs +++ b/crates/sail-execution/src/driver/job_scheduler/core.rs @@ -7,6 +7,7 @@ use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use indexmap::{IndexMap, IndexSet}; use log::{debug, warn}; +use sail_cache::cached_relation::release_cached_relation_leases; use sail_common_datafusion::error::CommonErrorCause; use sail_python_udf::error::PyErrExtractor; use sail_server::actor::ActorContext; @@ -496,6 +497,14 @@ impl JobScheduler { job.state = JobState::Canceled; } job.stopped_at = Some(Utc::now()); + for stage in job.graph.stages_mut() { + match release_cached_relation_leases(Arc::clone(&stage.plan)) { + Ok(plan) => stage.plan = plan, + Err(error) => { + warn!("failed to release cached relation leases for job {job_id}: {error}") + } + } + } actions } diff --git a/crates/sail-execution/src/job_graph/mod.rs b/crates/sail-execution/src/job_graph/mod.rs index cdefa9b11b..c52b945bce 100644 --- a/crates/sail-execution/src/job_graph/mod.rs +++ b/crates/sail-execution/src/job_graph/mod.rs @@ -28,6 +28,10 @@ impl JobGraph { &self.stages } + pub(crate) fn stages_mut(&mut self) -> &mut [Stage] { + &mut self.stages + } + pub fn schema(&self) -> &SchemaRef { &self.schema } diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index b3ec655f3d..64e4183001 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -1,6 +1,7 @@ import contextlib import gc import shutil +import time import uuid import pytest @@ -36,6 +37,7 @@ PAYLOAD_BYTES = 1024 DUPLICATE_COLUMN_ROWS = 3 CLEANUP_ROWS = 32 +COMPLETED_JOB_PARTITIONS = 4 @contextlib.contextmanager @@ -192,6 +194,37 @@ def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): spark.conf.set("spark.checkpoint.dir", original_checkpoint_path) +def test_checkpoint_gc_releases_completed_cluster_job_lease(spark, tmp_path): + if is_jvm_spark(): + checkpointed = spark.range(CLEANUP_ROWS, numPartitions=4).checkpoint() + assert checkpointed.count() == CLEANUP_ROWS + del checkpointed + gc.collect() + return + + original_checkpoint_path = spark.conf.get("spark.checkpoint.dir") + checkpoint_root = tmp_path / "completed-job" + try: + spark.conf.set("spark.checkpoint.dir", checkpoint_root.as_uri()) + checkpointed = spark.range( + CLEANUP_ROWS, + numPartitions=COMPLETED_JOB_PARTITIONS, + ).checkpoint() + assert checkpointed.count() == CLEANUP_ROWS + checkpoint_files = list(checkpoint_root.rglob("*.arrow")) + assert len(checkpoint_files) == COMPLETED_JOB_PARTITIONS + + del checkpointed + gc.collect() + deadline = time.monotonic() + 5 + while any(path.exists() for path in checkpoint_files) and time.monotonic() < deadline: + time.sleep(0.1) + + assert all(not path.exists() for path in checkpoint_files) + finally: + spark.conf.set("spark.checkpoint.dir", original_checkpoint_path) + + @SAIL_XFAIL def test_checkpoint_command_can_reattach(spark): client = spark._client # noqa: SLF001 From 594688c2b6b0f4720ef94cf9e0ddc7729edf5f6b Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 00:35:26 +0800 Subject: [PATCH 29/40] fix: keep checkpoint data out of worker plans --- .../src/driver/job_scheduler/core.rs | 15 ++- .../sail-execution/src/job_graph/planner.rs | 52 ++++++++++ crates/sail-execution/src/proto/encode.rs | 99 ++++++++++++++++++- crates/sail-execution/src/proto/mod.rs | 4 +- .../tests/spark/execution/test_checkpoint.py | 10 +- 5 files changed, 171 insertions(+), 9 deletions(-) diff --git a/crates/sail-execution/src/driver/job_scheduler/core.rs b/crates/sail-execution/src/driver/job_scheduler/core.rs index 70bee01711..c60016b5cb 100644 --- a/crates/sail-execution/src/driver/job_scheduler/core.rs +++ b/crates/sail-execution/src/driver/job_scheduler/core.rs @@ -24,7 +24,9 @@ use crate::id::{JobId, TaskKey, TaskKeyDisplay, TaskStreamKey}; use crate::job_graph::{ InputMode, JobGraph, OutputDistribution, OutputMode, Stage, StageInput, TaskPlacement, }; -use crate::proto::{encode_remote_physical_expr, encode_remote_physical_plan}; +use crate::proto::{ + encode_driver_physical_plan, encode_remote_physical_expr, encode_remote_physical_plan, +}; use crate::task::definition::{ TaskDefinition, TaskInput, TaskInputKey, TaskInputLocator, TaskOutput, TaskOutputDistribution, TaskOutputLocator, @@ -533,7 +535,16 @@ impl JobScheduler { ))); }; - let plan = encode_remote_physical_plan(self.codec.as_ref(), stage.plan.clone())?; + let plan = match stage.placement { + TaskPlacement::Driver => encode_driver_physical_plan( + self.codec.as_ref(), + Arc::clone(&stage.plan), + key.partition, + )?, + TaskPlacement::Worker => { + encode_remote_physical_plan(self.codec.as_ref(), Arc::clone(&stage.plan))? + } + }; let inputs = stage .inputs .iter() diff --git a/crates/sail-execution/src/job_graph/planner.rs b/crates/sail-execution/src/job_graph/planner.rs index 30ab4f58d6..99c4bb0935 100644 --- a/crates/sail-execution/src/job_graph/planner.rs +++ b/crates/sail-execution/src/job_graph/planner.rs @@ -2,6 +2,8 @@ use std::sync::Arc; use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion::common::{JoinType, Result, plan_datafusion_err}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::execution_props::ScalarSubqueryResults; use datafusion::physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion::physical_expr::{Partitioning, PhysicalExpr}; @@ -22,6 +24,7 @@ use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use datafusion::physical_plan::{ ExecutionPlan, ExecutionPlanProperties, PlanProperties, with_new_children_if_necessary, }; +use sail_cache::cached_relation::CachedRelationExec; use sail_catalog_system::physical_plan::SystemTableExec; use sail_common_datafusion::utils::items::ItemTaker; use sail_data_source::listing::delete::FileDeleteExec; @@ -456,6 +459,18 @@ fn plan_job_graph_stages( let plan = create_rescale_input(child, coalesce.output_partitions(), graph, scalar_context)?; PlannedSubtree::without_pending_scalar_subquery_expr(plan) + } else if subtree + .plan + .downcast_ref::() + .is_some_and(|cached| { + cached + .input() + .downcast_ref::() + .is_some_and(|source| source.data_source().is::()) + }) + { + let plan = create_driver_stage(subtree.into_planned_subtree(), graph, scalar_context)?; + PlannedSubtree::without_pending_scalar_subquery_expr(plan) } else if subtree.plan.is::() || subtree.plan.is::() || subtree.plan.is::() @@ -911,7 +926,11 @@ fn create_driver_stage( mod tests { use std::sync::Arc; + use datafusion::arrow::array::Int32Array; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; use datafusion::functions_aggregate::sum::sum_udaf; use datafusion::logical_expr::Operator; use datafusion::logical_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; @@ -927,6 +946,7 @@ mod tests { use datafusion::physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable}; + use sail_cache::cached_relation::CachedRelationExec; use sail_catalog::command::CatalogCommand; use sail_physical_plan::barrier::BarrierExec; use sail_physical_plan::catalog_command::CatalogCommandExec; @@ -945,6 +965,38 @@ mod tests { Arc::new(EmptyExec::new(schema())) } + #[test] + fn test_job_graph_reads_cached_memory_from_driver_stage() { + let batch = + RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); + let memory: Arc = + MemorySourceConfig::try_new_exec(&[vec![batch.clone()], vec![batch]], schema(), None) + .unwrap(); + let cached: Arc = Arc::new(CachedRelationExec::new( + Arc::clone(&memory), + Arc::clone(memory.properties()), + )); + + let graph = JobGraph::try_new(cached).unwrap(); + + assert_eq!(graph.stages().len(), 2); + assert_eq!(graph.stages()[0].placement, TaskPlacement::Driver); + assert_eq!(graph.stages()[1].placement, TaskPlacement::Worker); + assert!(matches!( + graph.stages()[1].inputs.as_slice(), + [StageInput { + stage: 0, + mode: InputMode::Forward, + }] + )); + let cached = graph.stages()[0] + .plan + .downcast_ref::() + .expect("driver stage should retain the cached relation wrapper"); + assert!(cached.input().is::()); + assert!(graph.stages()[1].plan.is::>()); + } + #[test] fn test_job_graph_distinguishes_batch_and_explicit_round_robin_shuffle() { let generic_graph = JobGraph::try_new(Arc::new( diff --git a/crates/sail-execution/src/proto/encode.rs b/crates/sail-execution/src/proto/encode.rs index ec5dea4c13..c363043095 100644 --- a/crates/sail-execution/src/proto/encode.rs +++ b/crates/sail-execution/src/proto/encode.rs @@ -3,7 +3,8 @@ use std::sync::Arc; use datafusion::arrow::datatypes::{FieldRef, Schema}; use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion::common::{Result, plan_err}; -use datafusion::datasource::source::DataSourceExec; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::{DataSource, DataSourceExec}; use datafusion::physical_expr::{HigherOrderFunctionExpr, PhysicalExpr}; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::generated::datafusion_common as gen_datafusion_common; @@ -40,6 +41,53 @@ pub fn encode_remote_physical_plan( try_encode_physical_plan(codec, plan) } +pub fn encode_driver_physical_plan( + codec: &dyn PhysicalExtensionCodec, + plan: Arc, + partition: usize, +) -> Result> { + let plan = select_memory_source_partition(plan, partition)?; + encode_remote_physical_plan(codec, plan) +} + +fn select_memory_source_partition( + plan: Arc, + partition: usize, +) -> Result> { + plan.transform(|node| { + let Some(data_source) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(memory) = data_source + .data_source() + .downcast_ref::() + else { + return Ok(Transformed::no(node)); + }; + if partition >= memory.partitions().len() { + return plan_err!( + "memory source partition {partition} is out of bounds for {} partitions", + memory.partitions().len() + ); + } + + let mut partitions = vec![vec![]; memory.partitions().len()]; + partitions[partition] = memory.partitions()[partition].clone(); + let source = MemorySourceConfig::try_new( + &partitions, + memory.original_schema(), + memory.projection().clone(), + )? + .with_show_sizes(memory.show_sizes()) + .try_with_sort_information(memory.sort_information().to_vec())? + .with_limit(memory.fetch()); + Ok(Transformed::yes( + DataSourceExec::from_data_source(source) as Arc + )) + }) + .data() +} + pub fn encode_remote_physical_expr( codec: &dyn PhysicalExtensionCodec, expr: &Arc, @@ -126,3 +174,52 @@ pub(super) fn try_encode_higher_order_udf( higher_order_udf_kind: Some(udf_kind), }) } + +#[cfg(test)] +#[expect(clippy::unwrap_used)] +mod tests { + use std::sync::Arc; + + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSourceExec; + use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; + + use super::select_memory_source_partition; + + #[test] + fn memory_source_encoding_keeps_only_the_driver_task_partition() { + let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let first = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + let second = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![2]))], + ) + .unwrap(); + let plan: Arc = MemorySourceConfig::try_new_exec( + &[vec![first], vec![second]], + Arc::clone(&schema), + None, + ) + .unwrap(); + + let selected = select_memory_source_partition(plan, 1).unwrap(); + assert_eq!(selected.output_partitioning().partition_count(), 2); + let selected = selected + .downcast_ref::() + .unwrap() + .data_source() + .downcast_ref::() + .unwrap(); + + assert_eq!(selected.partitions().len(), 2); + assert!(selected.partitions()[0].is_empty()); + assert_eq!(selected.partitions()[1][0].num_rows(), 1); + } +} diff --git a/crates/sail-execution/src/proto/mod.rs b/crates/sail-execution/src/proto/mod.rs index c40bdeae72..46eb1a523c 100644 --- a/crates/sail-execution/src/proto/mod.rs +++ b/crates/sail-execution/src/proto/mod.rs @@ -5,4 +5,6 @@ mod encode; pub use codec::RemoteExecutionCodec; pub use decode::{decode_remote_physical_expr, decode_remote_physical_plan}; -pub use encode::{encode_remote_physical_expr, encode_remote_physical_plan}; +pub use encode::{ + encode_driver_physical_plan, encode_remote_physical_expr, encode_remote_physical_plan, +}; diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index 64e4183001..91079d000b 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -33,7 +33,7 @@ SMALL_PAYLOAD_ROWS = 16 * 1024 -LARGE_PAYLOAD_ROWS = 64 * 1024 +LARGE_PAYLOAD_ROW_COUNTS = [64 * 1024, 132 * 1024] PAYLOAD_BYTES = 1024 DUPLICATE_COLUMN_ROWS = 3 CLEANUP_ROWS = 32 @@ -153,11 +153,11 @@ def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, p _assert_payload(checkpointed, SMALL_PAYLOAD_ROWS) -@SAIL_XFAIL -def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark): - checkpointed = _payload_dataframe(spark, LARGE_PAYLOAD_ROWS).localCheckpoint() +@pytest.mark.parametrize("rows", LARGE_PAYLOAD_ROW_COUNTS, ids=["64-mib", "132-mib"]) +def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark, rows): + checkpointed = _payload_dataframe(spark, rows).localCheckpoint() - _assert_large_payload(checkpointed, LARGE_PAYLOAD_ROWS) + _assert_large_payload(checkpointed, rows) def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): From d2d746dc33ccd2bd80fc199648bbd61c6e587daf Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 00:45:06 +0800 Subject: [PATCH 30/40] fix: commit reliable checkpoint attempt winners --- Cargo.lock | 1 + crates/sail-physical-plan/Cargo.toml | 1 + crates/sail-physical-plan/src/checkpoint.rs | 98 ++++++++++- crates/sail-session/src/checkpoint.rs | 174 +++++++++++++++++--- 4 files changed, 238 insertions(+), 36 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index f52aa89fd1..d46e27cf15 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7931,6 +7931,7 @@ dependencies = [ "sail-logical-plan", "tokio", "tokio-stream", + "uuid", ] [[package]] diff --git a/crates/sail-physical-plan/Cargo.toml b/crates/sail-physical-plan/Cargo.toml index 81da84e01a..bc611f6f57 100644 --- a/crates/sail-physical-plan/Cargo.toml +++ b/crates/sail-physical-plan/Cargo.toml @@ -19,3 +19,4 @@ object_store = { workspace = true } tokio = { workspace = true } tokio-stream = { workspace = true } lazy_static = { workspace = true } +uuid = { workspace = true } diff --git a/crates/sail-physical-plan/src/checkpoint.rs b/crates/sail-physical-plan/src/checkpoint.rs index ff660fe407..683ac90913 100644 --- a/crates/sail-physical-plan/src/checkpoint.rs +++ b/crates/sail-physical-plan/src/checkpoint.rs @@ -1,7 +1,7 @@ use std::fmt::Formatter; use std::sync::Arc; -use datafusion::arrow::array::{ArrayRef, LargeBinaryArray, RecordBatch, UInt64Array}; +use datafusion::arrow::array::{ArrayRef, LargeBinaryArray, RecordBatch, StringArray, UInt64Array}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; @@ -18,10 +18,12 @@ use object_store::{ObjectStoreExt, PutPayload}; use sail_common_datafusion::array::record_batch::{ write_record_batches, write_record_batches_file, }; +use uuid::Uuid; const PARTITION_COLUMN: &str = "partition"; const SEQUENCE_COLUMN: &str = "sequence"; const DATA_COLUMN: &str = "data"; +const LOCATION_COLUMN: &str = "location"; #[derive(Debug)] pub struct LocalCheckpointExec { @@ -186,11 +188,10 @@ impl ReliableCheckpointExec { object_store_url: ObjectStoreUrl, path: Path, ) -> Self { - let schema = Arc::new(Schema::new(vec![Field::new( - PARTITION_COLUMN, - DataType::UInt64, - false, - )])); + let schema = Arc::new(Schema::new(vec![ + Field::new(PARTITION_COLUMN, DataType::UInt64, false), + Field::new(LOCATION_COLUMN, DataType::Utf8, false), + ])); let properties = Arc::new(PlanProperties::new( EquivalenceProperties::new(schema), Partitioning::UnknownPartitioning(input.output_partitioning().partition_count()), @@ -271,8 +272,12 @@ impl ExecutionPlan for ReliableCheckpointExec { let output_schema = self.schema(); let stream_schema = Arc::clone(&output_schema); let store = context.runtime_env().object_store(&self.object_store_url)?; - // FIXME: Publish attempt-scoped temporary files atomically so stale retries cannot overwrite successful output. - let location = self.path.clone().join(format!("part-{partition:05}.arrow")); + let location = self + .path + .clone() + .join("_temporary") + .join(Uuid::new_v4().to_string()) + .join(format!("part-{partition:05}.arrow")); let output = stream::once(async move { // FIXME: Stream Arrow IPC to object storage instead of buffering an entire partition. let mut batches = vec![]; @@ -286,7 +291,10 @@ impl ExecutionPlan for ReliableCheckpointExec { .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?; let partition = u64::try_from(partition) .map_err(|_| internal_datafusion_err!("checkpoint partition index is too large"))?; - let columns: Vec = vec![Arc::new(UInt64Array::from(vec![partition]))]; + let columns: Vec = vec![ + Arc::new(UInt64Array::from(vec![partition])), + Arc::new(StringArray::from(vec![location.as_ref()])), + ]; Ok(RecordBatch::try_new(output_schema, columns)?) }); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -295,3 +303,75 @@ impl ExecutionPlan for ReliableCheckpointExec { ))) } } + +#[cfg(test)] +mod tests { + use datafusion::arrow::array::{Array, Int32Array, StringArray}; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::object_store::memory::InMemory; + use datafusion::prelude::SessionContext; + + use super::*; + + async fn execute_checkpoint_attempt( + checkpoint: &ReliableCheckpointExec, + context: Arc, + ) -> Result { + let mut stream = checkpoint.execute(0, context)?; + let batch = stream.next().await.ok_or_else(|| { + internal_datafusion_err!("reliable checkpoint attempt returned no metadata") + })??; + let locations = batch + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("invalid checkpoint location column"))?; + if locations.is_null(0) { + return Err(internal_datafusion_err!( + "checkpoint attempt returned null location" + )); + } + Path::parse(locations.value(0)) + .map_err(|error| internal_datafusion_err!("invalid checkpoint location: {error}")) + } + + #[tokio::test] + async fn reliable_checkpoint_attempts_write_distinct_temporary_objects() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + )?; + let input: Arc = + MemorySourceConfig::try_new_exec(&[vec![batch]], schema, None)?; + let object_store_url = ObjectStoreUrl::parse("memory://")?; + let checkpoint = + ReliableCheckpointExec::new(input, object_store_url.clone(), Path::from("checkpoint")); + let ctx = SessionContext::new(); + ctx.runtime_env() + .register_object_store(object_store_url.as_ref(), Arc::new(InMemory::new())); + + let first = execute_checkpoint_attempt(&checkpoint, ctx.task_ctx()).await?; + let second = execute_checkpoint_attempt(&checkpoint, ctx.task_ctx()).await?; + + assert_ne!(first, second); + let temporary = Path::from("checkpoint/_temporary"); + assert!(first.prefix_matches(&temporary)); + assert!(second.prefix_matches(&temporary)); + let store = ctx.runtime_env().object_store(&object_store_url)?; + store + .head(&first) + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?; + store + .head(&second) + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?; + Ok(()) + } +} diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs index 30ca4add01..3ce27b13ff 100644 --- a/crates/sail-session/src/checkpoint.rs +++ b/crates/sail-session/src/checkpoint.rs @@ -1,18 +1,21 @@ use std::sync::Arc; use async_trait::async_trait; -use datafusion::arrow::array::{Array, UInt64Array}; +use datafusion::arrow::array::{Array, StringArray, UInt64Array}; use datafusion::arrow::datatypes::SchemaRef; -use datafusion::object_store::ObjectStoreExt; +use datafusion::object_store::path::Path; +use datafusion::object_store::{ObjectMeta, ObjectStoreExt}; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion::prelude::SessionContext; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; -use futures::StreamExt; +use futures::{StreamExt, TryStreamExt}; use sail_cache::checkpoint::{CheckpointStore, ReliableCheckpoint}; use sail_common_datafusion::array::record_batch::write_record_batches_file; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; -use sail_object_store::{delete_object_store_prefix, resolve_object_store_path}; +use sail_object_store::{ + ResolvedObjectStorePath, delete_object_store_prefix, resolve_object_store_path, +}; use sail_physical_plan::checkpoint::ReliableCheckpointExec; #[derive(Debug, Default)] @@ -50,12 +53,12 @@ impl CheckpointStore for ObjectStoreCheckpointStore { .runner() .execute(ctx, Arc::new(checkpoint_exec)) .await?; - let mut completed = vec![false; partition_count]; + let mut attempts = vec![None; partition_count]; while let Some(batch) = stream.next().await { let batch = batch?; - if batch.num_columns() != 1 { + if batch.num_columns() != 2 { return Err(internal_datafusion_err!( - "reliable checkpoint returned {} columns instead of 1", + "reliable checkpoint returned {} columns instead of 2", batch.num_columns() )); } @@ -66,21 +69,34 @@ impl CheckpointStore for ObjectStoreCheckpointStore { .ok_or_else(|| { internal_datafusion_err!("invalid reliable checkpoint partition column") })?; + let locations = batch + .column(1) + .as_any() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("invalid reliable checkpoint location column") + })?; for row in 0..batch.num_rows() { - if partitions.is_null(row) { + if partitions.is_null(row) || locations.is_null(row) { return Err(internal_datafusion_err!( - "reliable checkpoint returned a null partition" + "reliable checkpoint returned null metadata" )); } let partition = usize::try_from(partitions.value(row)).map_err(|_| { internal_datafusion_err!("reliable checkpoint partition index is too large") })?; - let Some(done) = completed.get_mut(partition) else { + let location = Path::parse(locations.value(row)).map_err(|error| { + internal_datafusion_err!( + "reliable checkpoint returned invalid location: {error}" + ) + })?; + validate_checkpoint_attempt_location(&checkpoint_path, partition, &location)?; + let Some(attempt) = attempts.get_mut(partition) else { return Err(internal_datafusion_err!( "reliable checkpoint returned invalid partition {partition}" )); }; - if std::mem::replace(done, true) { + if attempt.replace(location).is_some() { return Err(internal_datafusion_err!( "reliable checkpoint returned duplicate partition {partition}" )); @@ -88,22 +104,18 @@ impl CheckpointStore for ObjectStoreCheckpointStore { } } - let mut files = Vec::with_capacity(partition_count); - for (partition, completed) in completed.into_iter().enumerate() { - if !completed { - return Err(internal_datafusion_err!( - "reliable checkpoint did not return partition {partition}" - )); - } - let file_path = checkpoint_path.child(&format!("part-{partition:05}.arrow")); - files.push( - checkpoint_path - .store() - .head(&file_path) - .await - .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?, - ); - } + let attempts = attempts + .into_iter() + .enumerate() + .map(|(partition, attempt)| { + attempt.ok_or_else(|| { + internal_datafusion_err!( + "reliable checkpoint did not return partition {partition}" + ) + }) + }) + .collect::>>()?; + let files = commit_checkpoint_attempts(&checkpoint_path, attempts).await?; Ok(ReliableCheckpoint::new( checkpoint_path.object_store_url().clone(), @@ -117,14 +129,122 @@ impl CheckpointStore for ObjectStoreCheckpointStore { } } +fn validate_checkpoint_attempt_location( + checkpoint_path: &ResolvedObjectStorePath, + partition: usize, + location: &Path, +) -> Result<()> { + let temporary = checkpoint_path.child("_temporary"); + let filename = format!("part-{partition:05}.arrow"); + if !location.prefix_matches(&temporary) + || location.parts_count() != temporary.parts_count() + 2 + || location.filename() != Some(filename.as_str()) + { + return Err(internal_datafusion_err!( + "reliable checkpoint returned unexpected attempt location {location} for partition {partition}" + )); + } + Ok(()) +} + +async fn commit_checkpoint_attempts( + checkpoint_path: &ResolvedObjectStorePath, + attempts: Vec, +) -> Result> { + let mut files = Vec::with_capacity(attempts.len()); + for (partition, attempt) in attempts.into_iter().enumerate() { + let file_path = checkpoint_path.child(&format!("part-{partition:05}.arrow")); + checkpoint_path + .store() + .copy(&attempt, &file_path) + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?; + files.push( + checkpoint_path + .store() + .head(&file_path) + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?, + ); + } + + let temporary = checkpoint_path.child("_temporary"); + let temporary_files = checkpoint_path + .store() + .list(Some(&temporary)) + .try_collect::>() + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?; + for file in temporary_files { + match checkpoint_path.store().delete(&file.location).await { + Ok(()) | Err(datafusion::object_store::Error::NotFound { .. }) => {} + Err(error) => return Err(DataFusionError::ObjectStore(Box::new(error))), + } + } + Ok(files) +} + #[cfg(test)] mod tests { use std::time::{SystemTime, UNIX_EPOCH}; + use datafusion::execution::object_store::ObjectStoreUrl; + use datafusion::object_store::memory::InMemory; use datafusion_common::{DataFusionError, internal_datafusion_err}; use super::*; + #[tokio::test] + async fn late_checkpoint_attempt_cannot_overwrite_committed_winner() -> Result<()> { + let suffix = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_err(|e| internal_datafusion_err!("{e}"))? + .as_nanos(); + let ctx = SessionContext::new(); + let object_store_url = ObjectStoreUrl::parse("memory://")?; + ctx.runtime_env() + .register_object_store(object_store_url.as_ref(), Arc::new(InMemory::new())); + let checkpoint = resolve_object_store_path( + ctx.runtime_env().as_ref(), + &format!("memory:///{suffix}/checkpoint"), + )?; + let winner = checkpoint + .child("_temporary") + .join("winner") + .join("part-00000.arrow"); + let stale = checkpoint + .child("_temporary") + .join("stale") + .join("part-00000.arrow"); + checkpoint.put_bytes(&winner, b"winner".to_vec()).await?; + checkpoint.put_bytes(&stale, b"stale".to_vec()).await?; + + let files = commit_checkpoint_attempts(&checkpoint, vec![winner.clone()]).await?; + + let committed = checkpoint.child("part-00000.arrow"); + assert_eq!(files[0].location, committed); + let winner_head = checkpoint.store().head(&winner).await; + assert!( + matches!( + &winner_head, + Err(datafusion::object_store::Error::NotFound { .. }) + ), + "winner temporary object was not deleted: {winner_head:?}" + ); + + checkpoint.put_bytes(&stale, b"late stale".to_vec()).await?; + let bytes = checkpoint + .store() + .get(&committed) + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))? + .bytes() + .await + .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?; + assert_eq!(bytes.as_ref(), b"winner"); + Ok(()) + } + #[tokio::test] async fn cleanup_checkpoint_path_deletes_prefix_files() -> Result<()> { let suffix = SystemTime::now() From 5b1157a7f3a389297dd40ab423717102ae0dc5ea Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 00:49:47 +0800 Subject: [PATCH 31/40] Revert "fix: keep checkpoint data out of worker plans" --- .../src/driver/job_scheduler/core.rs | 15 +-- .../sail-execution/src/job_graph/planner.rs | 52 ---------- crates/sail-execution/src/proto/encode.rs | 99 +------------------ crates/sail-execution/src/proto/mod.rs | 4 +- .../tests/spark/execution/test_checkpoint.py | 11 ++- 5 files changed, 10 insertions(+), 171 deletions(-) diff --git a/crates/sail-execution/src/driver/job_scheduler/core.rs b/crates/sail-execution/src/driver/job_scheduler/core.rs index c60016b5cb..70bee01711 100644 --- a/crates/sail-execution/src/driver/job_scheduler/core.rs +++ b/crates/sail-execution/src/driver/job_scheduler/core.rs @@ -24,9 +24,7 @@ use crate::id::{JobId, TaskKey, TaskKeyDisplay, TaskStreamKey}; use crate::job_graph::{ InputMode, JobGraph, OutputDistribution, OutputMode, Stage, StageInput, TaskPlacement, }; -use crate::proto::{ - encode_driver_physical_plan, encode_remote_physical_expr, encode_remote_physical_plan, -}; +use crate::proto::{encode_remote_physical_expr, encode_remote_physical_plan}; use crate::task::definition::{ TaskDefinition, TaskInput, TaskInputKey, TaskInputLocator, TaskOutput, TaskOutputDistribution, TaskOutputLocator, @@ -535,16 +533,7 @@ impl JobScheduler { ))); }; - let plan = match stage.placement { - TaskPlacement::Driver => encode_driver_physical_plan( - self.codec.as_ref(), - Arc::clone(&stage.plan), - key.partition, - )?, - TaskPlacement::Worker => { - encode_remote_physical_plan(self.codec.as_ref(), Arc::clone(&stage.plan))? - } - }; + let plan = encode_remote_physical_plan(self.codec.as_ref(), stage.plan.clone())?; let inputs = stage .inputs .iter() diff --git a/crates/sail-execution/src/job_graph/planner.rs b/crates/sail-execution/src/job_graph/planner.rs index 99c4bb0935..30ab4f58d6 100644 --- a/crates/sail-execution/src/job_graph/planner.rs +++ b/crates/sail-execution/src/job_graph/planner.rs @@ -2,8 +2,6 @@ use std::sync::Arc; use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion::common::{JoinType, Result, plan_datafusion_err}; -use datafusion::datasource::memory::MemorySourceConfig; -use datafusion::datasource::source::DataSourceExec; use datafusion::logical_expr::execution_props::ScalarSubqueryResults; use datafusion::physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion::physical_expr::{Partitioning, PhysicalExpr}; @@ -24,7 +22,6 @@ use datafusion::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use datafusion::physical_plan::{ ExecutionPlan, ExecutionPlanProperties, PlanProperties, with_new_children_if_necessary, }; -use sail_cache::cached_relation::CachedRelationExec; use sail_catalog_system::physical_plan::SystemTableExec; use sail_common_datafusion::utils::items::ItemTaker; use sail_data_source::listing::delete::FileDeleteExec; @@ -459,18 +456,6 @@ fn plan_job_graph_stages( let plan = create_rescale_input(child, coalesce.output_partitions(), graph, scalar_context)?; PlannedSubtree::without_pending_scalar_subquery_expr(plan) - } else if subtree - .plan - .downcast_ref::() - .is_some_and(|cached| { - cached - .input() - .downcast_ref::() - .is_some_and(|source| source.data_source().is::()) - }) - { - let plan = create_driver_stage(subtree.into_planned_subtree(), graph, scalar_context)?; - PlannedSubtree::without_pending_scalar_subquery_expr(plan) } else if subtree.plan.is::() || subtree.plan.is::() || subtree.plan.is::() @@ -926,11 +911,7 @@ fn create_driver_stage( mod tests { use std::sync::Arc; - use datafusion::arrow::array::Int32Array; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; - use datafusion::arrow::record_batch::RecordBatch; - use datafusion::datasource::memory::MemorySourceConfig; - use datafusion::datasource::source::DataSourceExec; use datafusion::functions_aggregate::sum::sum_udaf; use datafusion::logical_expr::Operator; use datafusion::logical_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; @@ -946,7 +927,6 @@ mod tests { use datafusion::physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable}; - use sail_cache::cached_relation::CachedRelationExec; use sail_catalog::command::CatalogCommand; use sail_physical_plan::barrier::BarrierExec; use sail_physical_plan::catalog_command::CatalogCommandExec; @@ -965,38 +945,6 @@ mod tests { Arc::new(EmptyExec::new(schema())) } - #[test] - fn test_job_graph_reads_cached_memory_from_driver_stage() { - let batch = - RecordBatch::try_new(schema(), vec![Arc::new(Int32Array::from(vec![1]))]).unwrap(); - let memory: Arc = - MemorySourceConfig::try_new_exec(&[vec![batch.clone()], vec![batch]], schema(), None) - .unwrap(); - let cached: Arc = Arc::new(CachedRelationExec::new( - Arc::clone(&memory), - Arc::clone(memory.properties()), - )); - - let graph = JobGraph::try_new(cached).unwrap(); - - assert_eq!(graph.stages().len(), 2); - assert_eq!(graph.stages()[0].placement, TaskPlacement::Driver); - assert_eq!(graph.stages()[1].placement, TaskPlacement::Worker); - assert!(matches!( - graph.stages()[1].inputs.as_slice(), - [StageInput { - stage: 0, - mode: InputMode::Forward, - }] - )); - let cached = graph.stages()[0] - .plan - .downcast_ref::() - .expect("driver stage should retain the cached relation wrapper"); - assert!(cached.input().is::()); - assert!(graph.stages()[1].plan.is::>()); - } - #[test] fn test_job_graph_distinguishes_batch_and_explicit_round_robin_shuffle() { let generic_graph = JobGraph::try_new(Arc::new( diff --git a/crates/sail-execution/src/proto/encode.rs b/crates/sail-execution/src/proto/encode.rs index c363043095..ec5dea4c13 100644 --- a/crates/sail-execution/src/proto/encode.rs +++ b/crates/sail-execution/src/proto/encode.rs @@ -3,8 +3,7 @@ use std::sync::Arc; use datafusion::arrow::datatypes::{FieldRef, Schema}; use datafusion::common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion::common::{Result, plan_err}; -use datafusion::datasource::memory::MemorySourceConfig; -use datafusion::datasource::source::{DataSource, DataSourceExec}; +use datafusion::datasource::source::DataSourceExec; use datafusion::physical_expr::{HigherOrderFunctionExpr, PhysicalExpr}; use datafusion::physical_plan::ExecutionPlan; use datafusion_proto::generated::datafusion_common as gen_datafusion_common; @@ -41,53 +40,6 @@ pub fn encode_remote_physical_plan( try_encode_physical_plan(codec, plan) } -pub fn encode_driver_physical_plan( - codec: &dyn PhysicalExtensionCodec, - plan: Arc, - partition: usize, -) -> Result> { - let plan = select_memory_source_partition(plan, partition)?; - encode_remote_physical_plan(codec, plan) -} - -fn select_memory_source_partition( - plan: Arc, - partition: usize, -) -> Result> { - plan.transform(|node| { - let Some(data_source) = node.downcast_ref::() else { - return Ok(Transformed::no(node)); - }; - let Some(memory) = data_source - .data_source() - .downcast_ref::() - else { - return Ok(Transformed::no(node)); - }; - if partition >= memory.partitions().len() { - return plan_err!( - "memory source partition {partition} is out of bounds for {} partitions", - memory.partitions().len() - ); - } - - let mut partitions = vec![vec![]; memory.partitions().len()]; - partitions[partition] = memory.partitions()[partition].clone(); - let source = MemorySourceConfig::try_new( - &partitions, - memory.original_schema(), - memory.projection().clone(), - )? - .with_show_sizes(memory.show_sizes()) - .try_with_sort_information(memory.sort_information().to_vec())? - .with_limit(memory.fetch()); - Ok(Transformed::yes( - DataSourceExec::from_data_source(source) as Arc - )) - }) - .data() -} - pub fn encode_remote_physical_expr( codec: &dyn PhysicalExtensionCodec, expr: &Arc, @@ -174,52 +126,3 @@ pub(super) fn try_encode_higher_order_udf( higher_order_udf_kind: Some(udf_kind), }) } - -#[cfg(test)] -#[expect(clippy::unwrap_used)] -mod tests { - use std::sync::Arc; - - use datafusion::arrow::array::Int32Array; - use datafusion::arrow::datatypes::{DataType, Field, Schema}; - use datafusion::arrow::record_batch::RecordBatch; - use datafusion::datasource::memory::MemorySourceConfig; - use datafusion::datasource::source::DataSourceExec; - use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; - - use super::select_memory_source_partition; - - #[test] - fn memory_source_encoding_keeps_only_the_driver_task_partition() { - let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); - let first = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int32Array::from(vec![1]))], - ) - .unwrap(); - let second = RecordBatch::try_new( - Arc::clone(&schema), - vec![Arc::new(Int32Array::from(vec![2]))], - ) - .unwrap(); - let plan: Arc = MemorySourceConfig::try_new_exec( - &[vec![first], vec![second]], - Arc::clone(&schema), - None, - ) - .unwrap(); - - let selected = select_memory_source_partition(plan, 1).unwrap(); - assert_eq!(selected.output_partitioning().partition_count(), 2); - let selected = selected - .downcast_ref::() - .unwrap() - .data_source() - .downcast_ref::() - .unwrap(); - - assert_eq!(selected.partitions().len(), 2); - assert!(selected.partitions()[0].is_empty()); - assert_eq!(selected.partitions()[1][0].num_rows(), 1); - } -} diff --git a/crates/sail-execution/src/proto/mod.rs b/crates/sail-execution/src/proto/mod.rs index 46eb1a523c..c40bdeae72 100644 --- a/crates/sail-execution/src/proto/mod.rs +++ b/crates/sail-execution/src/proto/mod.rs @@ -5,6 +5,4 @@ mod encode; pub use codec::RemoteExecutionCodec; pub use decode::{decode_remote_physical_expr, decode_remote_physical_plan}; -pub use encode::{ - encode_driver_physical_plan, encode_remote_physical_expr, encode_remote_physical_plan, -}; +pub use encode::{encode_remote_physical_expr, encode_remote_physical_plan}; diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index 91079d000b..aa816a2613 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -33,7 +33,7 @@ SMALL_PAYLOAD_ROWS = 16 * 1024 -LARGE_PAYLOAD_ROW_COUNTS = [64 * 1024, 132 * 1024] +LARGE_PAYLOAD_ROWS = 64 * 1024 PAYLOAD_BYTES = 1024 DUPLICATE_COLUMN_ROWS = 3 CLEANUP_ROWS = 32 @@ -153,11 +153,12 @@ def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, p _assert_payload(checkpointed, SMALL_PAYLOAD_ROWS) -@pytest.mark.parametrize("rows", LARGE_PAYLOAD_ROW_COUNTS, ids=["64-mib", "132-mib"]) -def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark, rows): - checkpointed = _payload_dataframe(spark, rows).localCheckpoint() +# FIXME: Keep local checkpoint payloads out of worker task definitions without changing distributed stage placement. +@SAIL_XFAIL +def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark): + checkpointed = _payload_dataframe(spark, LARGE_PAYLOAD_ROWS).localCheckpoint() - _assert_large_payload(checkpointed, rows) + _assert_large_payload(checkpointed, LARGE_PAYLOAD_ROWS) def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): From 492c0bcf78ba051f4bea0dd20198c1ed2322971c Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 00:51:10 +0800 Subject: [PATCH 32/40] Revert "fix: release checkpoint leases from completed jobs" --- crates/sail-cache/src/cached_relation.rs | 44 +------------------ .../src/driver/job_scheduler/core.rs | 9 ---- crates/sail-execution/src/job_graph/mod.rs | 4 -- .../tests/spark/execution/test_checkpoint.py | 4 +- 4 files changed, 4 insertions(+), 57 deletions(-) diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index 9320e1edc9..e06bd302df 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -17,7 +17,6 @@ use datafusion::physical_plan::{ with_new_children_if_necessary, }; use datafusion::prelude::SessionContext; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{DataFusionError, Result, Statistics, internal_datafusion_err}; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; @@ -100,6 +99,7 @@ enum CachedRelationPendingTarget { pub struct CachedRelationExec { input: Arc, properties: Arc, + // FIXME: Completed distributed jobs retain this lease and delay retired checkpoint cleanup. relation_lease: Option, } @@ -200,21 +200,6 @@ impl ExecutionPlan for CachedRelationExec { } } -pub fn release_cached_relation_leases( - plan: Arc, -) -> Result> { - plan.transform_up(|plan| { - let Some(cached) = plan.downcast_ref::() else { - return Ok(Transformed::no(plan)); - }; - Ok(Transformed::yes(Arc::new(CachedRelationExec::new( - Arc::clone(cached.input()), - Arc::clone(cached.properties()), - )))) - }) - .data() -} - #[derive(Debug, Clone)] struct PendingCachedRelationExec { relation_id: String, @@ -1035,33 +1020,6 @@ mod tests { Ok(()) } - #[test] - fn completed_plan_can_release_cached_relation_lease() -> Result<()> { - let schema = Arc::new(Schema::empty()); - let input: Arc = Arc::new(EmptyExec::new(schema)); - let properties = checkpoint_plan_properties(&input); - let relation = CachedRelation::new_pending_local_checkpoint( - Arc::clone(&input), - "MEMORY_AND_DISK" - .parse() - .map_err(|error| internal_datafusion_err!("{error}"))?, - ); - let mut completed_plan: Arc = Arc::new( - CachedRelationExec::with_relation_lease(input, properties, relation.clone()), - ); - assert!(!relation.is_exclusively_owned()); - - completed_plan = release_cached_relation_leases(completed_plan)?; - - assert!(relation.is_exclusively_owned()); - assert!( - completed_plan - .downcast_ref::() - .is_some_and(|cached| cached.relation_lease.is_none()) - ); - Ok(()) - } - #[tokio::test] async fn cached_relation_logical_plan_renames_duplicate_fields_by_position() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/crates/sail-execution/src/driver/job_scheduler/core.rs b/crates/sail-execution/src/driver/job_scheduler/core.rs index 70bee01711..e9363d6676 100644 --- a/crates/sail-execution/src/driver/job_scheduler/core.rs +++ b/crates/sail-execution/src/driver/job_scheduler/core.rs @@ -7,7 +7,6 @@ use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use indexmap::{IndexMap, IndexSet}; use log::{debug, warn}; -use sail_cache::cached_relation::release_cached_relation_leases; use sail_common_datafusion::error::CommonErrorCause; use sail_python_udf::error::PyErrExtractor; use sail_server::actor::ActorContext; @@ -497,14 +496,6 @@ impl JobScheduler { job.state = JobState::Canceled; } job.stopped_at = Some(Utc::now()); - for stage in job.graph.stages_mut() { - match release_cached_relation_leases(Arc::clone(&stage.plan)) { - Ok(plan) => stage.plan = plan, - Err(error) => { - warn!("failed to release cached relation leases for job {job_id}: {error}") - } - } - } actions } diff --git a/crates/sail-execution/src/job_graph/mod.rs b/crates/sail-execution/src/job_graph/mod.rs index c52b945bce..cdefa9b11b 100644 --- a/crates/sail-execution/src/job_graph/mod.rs +++ b/crates/sail-execution/src/job_graph/mod.rs @@ -28,10 +28,6 @@ impl JobGraph { &self.stages } - pub(crate) fn stages_mut(&mut self) -> &mut [Stage] { - &mut self.stages - } - pub fn schema(&self) -> &SchemaRef { &self.schema } diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index aa816a2613..febdd3f841 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -153,7 +153,7 @@ def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, p _assert_payload(checkpointed, SMALL_PAYLOAD_ROWS) -# FIXME: Keep local checkpoint payloads out of worker task definitions without changing distributed stage placement. +# FIXME: Local checkpoint payloads are embedded in every worker task definition. @SAIL_XFAIL def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark): checkpointed = _payload_dataframe(spark, LARGE_PAYLOAD_ROWS).localCheckpoint() @@ -195,6 +195,8 @@ def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): spark.conf.set("spark.checkpoint.dir", original_checkpoint_path) +# FIXME: Completed local-cluster jobs retain checkpoint relation leases after DataFrame GC. +@SAIL_XFAIL def test_checkpoint_gc_releases_completed_cluster_job_lease(spark, tmp_path): if is_jvm_spark(): checkpointed = spark.range(CLEANUP_ROWS, numPartitions=4).checkpoint() From d7c0a50d25f4273bd6f6a0016f68c2bc42607895 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 01:02:34 +0800 Subject: [PATCH 33/40] test: simplify checkpoint test values --- .../tests/spark/execution/test_checkpoint.py | 53 ++++++++----------- 1 file changed, 21 insertions(+), 32 deletions(-) diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index febdd3f841..0fa8620495 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -32,14 +32,6 @@ ) -SMALL_PAYLOAD_ROWS = 16 * 1024 -LARGE_PAYLOAD_ROWS = 64 * 1024 -PAYLOAD_BYTES = 1024 -DUPLICATE_COLUMN_ROWS = 3 -CLEANUP_ROWS = 32 -COMPLETED_JOB_PARTITIONS = 4 - - @contextlib.contextmanager def _checkpoint_spark_session(remote, checkpoint_path): builder = ( @@ -77,14 +69,14 @@ def spark(remote, checkpoint_path): def _payload_dataframe(spark, rows, partitions=4): return spark.range(rows, numPartitions=partitions).selectExpr( "id", - f"repeat('x', {PAYLOAD_BYTES}) AS payload", + "repeat('x', 1024) AS payload", ) def _assert_payload(checkpointed, rows): assert checkpointed.count() == rows actual = checkpointed.selectExpr("sum(length(payload)) AS payload_bytes").first() - assert actual.payload_bytes == rows * PAYLOAD_BYTES + assert actual.payload_bytes == rows * 1024 def _assert_large_payload(checkpointed, rows): @@ -93,7 +85,7 @@ def _assert_large_payload(checkpointed, rows): "sum(length(payload)) AS payload_bytes", ).first() assert actual.row_count == rows - assert actual.payload_bytes == rows * PAYLOAD_BYTES + assert actual.payload_bytes == rows * 1024 @pytest.fixture(scope="module") @@ -113,7 +105,7 @@ def retry_spark(spark, checkpoint_path): def _assert_checkpoint_preserves_duplicate_column_names(spark, eager, kind): - source = spark.range(DUPLICATE_COLUMN_ROWS, numPartitions=2).selectExpr( + source = spark.range(3, numPartitions=2).selectExpr( "id AS value", "id + 10 AS value", ) @@ -121,7 +113,7 @@ def _assert_checkpoint_preserves_duplicate_column_names(spark, eager, kind): checkpointed = source.localCheckpoint(eager=eager) if kind == "local" else source.checkpoint(eager=eager) assert [field.name for field in checkpointed.schema.fields] == ["value", "value"] - assert checkpointed.count() == DUPLICATE_COLUMN_ROWS + assert checkpointed.count() == 3 # noqa: PLR2004 @pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) @@ -142,36 +134,36 @@ def test_checkpoint_preserves_duplicate_column_names(spark, eager, kind): ], ) def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, partitions, storage_level): - source = _payload_dataframe(spark, SMALL_PAYLOAD_ROWS, partitions=partitions) - assert source.count() == SMALL_PAYLOAD_ROWS + source = _payload_dataframe(spark, 16 * 1024, partitions=partitions) + assert source.count() == 16 * 1024 if storage_level is None: checkpointed = source.localCheckpoint(eager=eager) else: checkpointed = source.localCheckpoint(eager=eager, storageLevel=storage_level) - _assert_payload(checkpointed, SMALL_PAYLOAD_ROWS) + _assert_payload(checkpointed, 16 * 1024) # FIXME: Local checkpoint payloads are embedded in every worker task definition. @SAIL_XFAIL def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark): - checkpointed = _payload_dataframe(spark, LARGE_PAYLOAD_ROWS).localCheckpoint() + checkpointed = _payload_dataframe(spark, 64 * 1024).localCheckpoint() - _assert_large_payload(checkpointed, LARGE_PAYLOAD_ROWS) + _assert_large_payload(checkpointed, 64 * 1024) def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): if is_jvm_spark(): # JVM Spark fixes the checkpoint root at server startup. - safe = spark.range(CLEANUP_ROWS, numPartitions=4).checkpoint() - assert safe.count() == CLEANUP_ROWS + safe = spark.range(32, numPartitions=4).checkpoint() + assert safe.count() == 32 # noqa: PLR2004 other = spark.range(1).checkpoint(eager=False) del other gc.collect() - assert safe.count() == CLEANUP_ROWS + assert safe.count() == 32 # noqa: PLR2004 return original_checkpoint_path = spark.conf.get("spark.checkpoint.dir") @@ -181,8 +173,8 @@ def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): try: spark.conf.set("spark.checkpoint.dir", safe_root) - safe = spark.range(CLEANUP_ROWS, numPartitions=4).checkpoint() - assert safe.count() == CLEANUP_ROWS + safe = spark.range(32, numPartitions=4).checkpoint() + assert safe.count() == 32 # noqa: PLR2004 safe_files = list(root.rglob("*.arrow")) spark.conf.set("spark.checkpoint.dir", overlapping_root) @@ -190,7 +182,7 @@ def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): spark.range(1).checkpoint(eager=False) assert all(path.exists() for path in safe_files) - assert safe.count() == CLEANUP_ROWS + assert safe.count() == 32 # noqa: PLR2004 finally: spark.conf.set("spark.checkpoint.dir", original_checkpoint_path) @@ -199,8 +191,8 @@ def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): @SAIL_XFAIL def test_checkpoint_gc_releases_completed_cluster_job_lease(spark, tmp_path): if is_jvm_spark(): - checkpointed = spark.range(CLEANUP_ROWS, numPartitions=4).checkpoint() - assert checkpointed.count() == CLEANUP_ROWS + checkpointed = spark.range(32, numPartitions=4).checkpoint() + assert checkpointed.count() == 32 # noqa: PLR2004 del checkpointed gc.collect() return @@ -209,13 +201,10 @@ def test_checkpoint_gc_releases_completed_cluster_job_lease(spark, tmp_path): checkpoint_root = tmp_path / "completed-job" try: spark.conf.set("spark.checkpoint.dir", checkpoint_root.as_uri()) - checkpointed = spark.range( - CLEANUP_ROWS, - numPartitions=COMPLETED_JOB_PARTITIONS, - ).checkpoint() - assert checkpointed.count() == CLEANUP_ROWS + checkpointed = spark.range(32, numPartitions=4).checkpoint() + assert checkpointed.count() == 32 # noqa: PLR2004 checkpoint_files = list(checkpoint_root.rglob("*.arrow")) - assert len(checkpoint_files) == COMPLETED_JOB_PARTITIONS + assert len(checkpoint_files) == 4 # noqa: PLR2004 del checkpointed gc.collect() From aeda7d2c6524b0047dd0234cde02d6c6e3242730 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 11:41:40 +0800 Subject: [PATCH 34/40] test: document checkpoint xfail limits --- crates/sail-cache/src/cached_relation.rs | 3 ++- python/pysail/tests/spark/execution/test_checkpoint.py | 7 +++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index e06bd302df..621d6bcea1 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -99,7 +99,8 @@ enum CachedRelationPendingTarget { pub struct CachedRelationExec { input: Arc, properties: Arc, - // FIXME: Completed distributed jobs retain this lease and delay retired checkpoint cleanup. + // FIXME: Completed distributed jobs retain this lease after DataFrame GC, so checkpoint files + // remain until the job plans are released during session shutdown. relation_lease: Option, } diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index 0fa8620495..7394ec74f0 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -145,7 +145,9 @@ def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, p _assert_payload(checkpointed, 16 * 1024) -# FIXME: Local checkpoint payloads are embedded in every worker task definition. +# FIXME: In local-cluster, this 64 MiB local checkpoint is embedded in every worker task definition and +# serializes to 135,869,682 bytes, exceeding the 128 MiB gRPC decoding limit. A 132 MiB checkpoint +# serialized to 280,225,527 bytes and raised peak process RSS to about 3.29 GB. @SAIL_XFAIL def test_local_checkpoint_large_payload_remains_executable_in_cluster(spark): checkpointed = _payload_dataframe(spark, 64 * 1024).localCheckpoint() @@ -187,7 +189,8 @@ def test_checkpoint_cleanup_is_scoped_to_one_relation(spark, tmp_path): spark.conf.set("spark.checkpoint.dir", original_checkpoint_path) -# FIXME: Completed local-cluster jobs retain checkpoint relation leases after DataFrame GC. +# FIXME: In local-cluster, a completed job retains the checkpoint relation lease after DataFrame GC. All +# four Arrow files remain after the five-second cleanup window and are released only when the session stops. @SAIL_XFAIL def test_checkpoint_gc_releases_completed_cluster_job_lease(spark, tmp_path): if is_jvm_spark(): From 80f9fb34b2dca26b6ffbc7ab4069927499495aa3 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 12:43:48 +0800 Subject: [PATCH 35/40] fix: preserve checkpoint properties across rewrites --- crates/sail-cache/src/cached_relation.rs | 145 ++++++++++++++++- .../tests/spark/dataframe/test_checkpoint.py | 88 +++++++++- .../tests/spark/execution/test_checkpoint.py | 153 +++++++++++++++++- 3 files changed, 375 insertions(+), 11 deletions(-) diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index 621d6bcea1..aef896078e 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -10,6 +10,7 @@ use datafusion::datasource::physical_plan::ArrowSource; use datafusion::execution::disk_manager::RefCountedTempFile; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::Partitioning; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType, reset_plan_states}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ @@ -167,12 +168,24 @@ impl ExecutionPlan for CachedRelationExec { )); } let input = Arc::clone(&children[0]); + let input_partitioning = input.output_partitioning(); + let stored_partitioning = self.properties.output_partitioning(); + // Checkpoint scans do not report their stored distribution, so retain it while the + // optimizer leaves the physical partition count unchanged. + let partitioning = match input_partitioning { + Partitioning::UnknownPartitioning(input_count) + if *input_count == stored_partitioning.partition_count() => + { + stored_partitioning.clone() + } + _ => input_partitioning.clone(), + }; Ok(Arc::new(Self { properties: Arc::new( self.properties .as_ref() .clone() - .with_partitioning(input.output_partitioning().clone()), + .with_partitioning(partitioning), ), input, relation_lease: self.relation_lease.clone(), @@ -999,25 +1012,143 @@ mod tests { use std::sync::Mutex; use async_trait::async_trait; + use datafusion::arrow::array::ArrayRef; use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::physical_expr::expressions::Column; + use datafusion::physical_expr::{ + EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, + }; use datafusion::physical_plan::empty::EmptyExec; + use datafusion_common::{Constraint, Constraints}; + use futures::stream; + use sail_common_datafusion::array::record_batch::write_record_batches; use super::*; use crate::checkpoint::{CheckpointStore, ReliableCheckpoint}; #[test] - fn cached_relation_rewrite_updates_partitioning() -> Result<()> { - let schema = Arc::new(Schema::empty()); + fn cached_relation_rewrite_preserves_exact_properties() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let column = Arc::new(Column::new("value", 0)) as Arc; + let ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new_default(Arc::clone(&column)).desc(), + ]) + .ok_or_else(|| internal_datafusion_err!("checkpoint ordering is empty"))?; + let equivalence = + EquivalenceProperties::new_with_orderings(Arc::clone(&schema), [ordering.clone()]) + .with_constraints(Constraints::new_unverified(vec![Constraint::Unique(vec![ + 0, + ])])); + let properties = Arc::new(PlanProperties::new( + equivalence, + Partitioning::Hash(vec![column], 2), + EmissionType::Incremental, + Boundedness::Bounded, + )); let original: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(2)); - let properties = checkpoint_plan_properties(&original); let cached = Arc::new(CachedRelationExec::new(original, properties)); - let rewritten_input: Arc = + + let same_partition_count: Arc = + Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(2)); + let rewritten = Arc::clone(&cached).with_new_children(vec![same_partition_count])?; + + assert!(matches!( + rewritten.output_partitioning(), + Partitioning::Hash(_, 2) + )); + assert_eq!(rewritten.output_ordering(), Some(&ordering)); + assert_eq!( + rewritten.properties().eq_properties.constraints(), + &Constraints::new_unverified(vec![Constraint::Unique(vec![0])]) + ); + + let changed_partition_count: Arc = Arc::new(EmptyExec::new(schema).with_partitions(10)); + let rewritten = cached.with_new_children(vec![changed_partition_count])?; + + assert!(matches!( + rewritten.output_partitioning(), + Partitioning::UnknownPartitioning(10) + )); + assert_eq!(rewritten.output_ordering(), Some(&ordering)); + Ok(()) + } + + #[tokio::test] + async fn pending_local_checkpoint_tracks_storage_level() -> Result<()> { + let plan: Arc = Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + let storage_level: spec::StorageLevel = "DISK_ONLY" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?; + let relation = CachedRelation::new_pending_local_checkpoint(plan, storage_level.clone()); + let data = relation.data.lock().await; - let rewritten = cached.with_new_children(vec![rewritten_input])?; + match &*data { + CachedRelationData::Pending(CachedRelationPending { + target: + CachedRelationPendingTarget::Local { + storage_level: actual, + }, + .. + }) => assert_eq!(actual, &storage_level), + other => { + return Err(internal_datafusion_err!( + "unexpected cached relation data: {other:?}" + )); + } + } + Ok(()) + } + + #[tokio::test] + async fn disk_only_local_checkpoint_materializes_only_on_disk() -> Result<()> { + let data_schema = Arc::new(Schema::empty()); + let bytes = write_record_batches(&[], data_schema.as_ref())?; + let checkpoint_schema = Arc::new(Schema::new(vec![ + Field::new("partition", DataType::UInt64, false), + Field::new("sequence", DataType::UInt64, false), + Field::new("data", DataType::LargeBinary, false), + ])); + let columns: Vec = vec![ + Arc::new(UInt64Array::from(vec![0])), + Arc::new(UInt64Array::from(vec![0])), + Arc::new(LargeBinaryArray::from_vec(vec![bytes.as_slice()])), + ]; + let batch = RecordBatch::try_new(Arc::clone(&checkpoint_schema), columns)?; + let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( + checkpoint_schema, + stream::iter(vec![Ok(batch)]), + )); + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&data_schema))); + let storage_level: spec::StorageLevel = "DISK_ONLY" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?; + + let materialized = CachedRelationLocalMaterialized::try_new( + &SessionContext::new(), + data_schema, + checkpoint_plan_properties(&input), + 1, + stream, + storage_level, + ) + .await?; - assert_eq!(rewritten.output_partitioning().partition_count(), 10); + assert!(materialized.memory_partitions.is_none()); + assert!(materialized.serialized_memory_partitions.is_none()); + let disk_partitions = materialized + .disk_partitions + .as_ref() + .ok_or_else(|| internal_datafusion_err!("disk checkpoint partitions are missing"))?; + assert_eq!(disk_partitions.len(), 1); + assert_eq!(disk_partitions[0].chunks.len(), 1); + assert_eq!(disk_partitions[0].chunks[0].files.len(), 1); + assert_eq!(materialized.load_partitions().await?, vec![vec![]]); Ok(()) } diff --git a/python/pysail/tests/spark/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index 52e40a5c5b..4b1cf7d9d3 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -7,6 +7,9 @@ import pytest from pandas.testing import assert_frame_equal from pyspark import StorageLevel +from pyspark.errors import PySparkException +from pyspark.sql.connect.dataframe import DataFrame +from pyspark.sql.connect.plan import CachedRemoteRelation, RemoveRemoteCachedRelation from pyspark.sql.functions import col, lit from pysail.testing.spark.session import spark_session_factory @@ -35,6 +38,30 @@ def _wait_for_checkpoint_files(checkpoint_path, *, present): assert any(checkpoint_path.rglob("*.arrow")) == present +def _cached_relation_id(dataframe): + relation_id = dataframe._plan._relation_id # noqa: SLF001 + assert isinstance(relation_id, str) + assert relation_id + return relation_id + + +def _wait_for_cached_relation_removal(spark, relation_id): + deadline = time.monotonic() + 5 + while time.monotonic() < deadline: + plan = CachedRemoteRelation(relation_id, spark) + probe = DataFrame(plan, spark) + try: + probe.count() + except PySparkException as error: + if f"No DataFrame with id {relation_id} is found" in str(error): + return + raise + finally: + plan._relation_id = None # noqa: SLF001 + time.sleep(0.1) + pytest.fail(f"cached relation {relation_id} was not removed") + + def test_dataframe_local_checkpoint_survives_source_removal(spark, tmp_path): source_path = tmp_path / "source" spark.createDataFrame( @@ -190,6 +217,22 @@ def test_dataframe_local_checkpoint_lazy_survives_source_removal_after_first_act ) +@pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint materialization plan") +def test_dataframe_local_checkpoint_lazy_storage_materialization_state(spark): + checkpointed = spark.range(0, 10, numPartitions=2).localCheckpoint( + eager=False, + storageLevel=StorageLevel.DISK_ONLY, + ) + + pending_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + assert "PendingCachedRelationExec" in pending_plan + + assert checkpointed.count() == 10 # noqa: PLR2004 + materialized_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + assert "PendingCachedRelationExec" not in materialized_plan + assert "DataSourceExec" in materialized_plan + + @pytest.mark.skipif(is_jvm_spark(), reason="JVM Spark Connect requires checkpoint dir at session startup") def test_dataframe_checkpoint_lazy(spark, tmp_path): source_path = tmp_path / "source" @@ -286,6 +329,42 @@ def test_dataframe_checkpoint_cleanup_on_dataframe_gc(spark, tmp_path): spark.conf.unset("spark.checkpoint.dir") +def test_dataframe_local_checkpoint_gc_removes_cached_relation(spark): + checkpointed = spark.range(0, 10).localCheckpoint() + relation_id = _cached_relation_id(checkpointed) + + del checkpointed + gc.collect() + + _wait_for_cached_relation_removal(spark, relation_id) + + +def test_dataframe_local_checkpoint_explicit_cleanup_removes_cached_relation(spark): + checkpointed = spark.range(0, 10).localCheckpoint() + relation_id = _cached_relation_id(checkpointed) + relation = checkpointed._plan # noqa: SLF001 + client = spark._client # noqa: SLF001 + + client.execute_command(RemoveRemoteCachedRelation(relation).command(client)) + relation._relation_id = None # noqa: SLF001 + + _wait_for_cached_relation_removal(spark, relation_id) + + +def test_dataframe_local_checkpoint_cleanup_waits_for_derived_dataframe(spark): + checkpointed = spark.range(0, 10).localCheckpoint() + relation_id = _cached_relation_id(checkpointed) + derived = checkpointed.repartition(2) + + del checkpointed + gc.collect() + assert derived.count() == 10 # noqa: PLR2004 + + del derived + gc.collect() + _wait_for_cached_relation_removal(spark, relation_id) + + @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific derived checkpoint cleanup") def test_dataframe_checkpoint_cleanup_waits_for_derived_dataframe(spark, tmp_path): checkpoint_path = tmp_path / "checkpoints" @@ -309,7 +388,7 @@ def test_dataframe_checkpoint_cleanup_waits_for_derived_dataframe(spark, tmp_pat @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint partition preservation") @pytest.mark.parametrize("local", [True, False], ids=["local", "reliable"]) def test_dataframe_checkpoint_preserves_multiple_output_partitions(spark, tmp_path, local): - df = spark.range(0, 100, numPartitions=4).repartition(4) + df = spark.range(0, 1000, numPartitions=4).repartition(4) if local: checkpointed = df.localCheckpoint() else: @@ -321,7 +400,7 @@ def test_dataframe_checkpoint_preserves_multiple_output_partitions(spark, tmp_pa spark.conf.unset("spark.checkpoint.dir") partition_ids = {row.pid for row in checkpointed.selectExpr("spark_partition_id() AS pid").distinct().collect()} - assert len(partition_ids) > 1 + assert partition_ids == set(range(4)) @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific checkpoint partition preservation") @@ -365,11 +444,16 @@ def test_dataframe_checkpoint_preserves_partitioning_and_ordering(spark, tmp_pat finally: spark.conf.unset("spark.checkpoint.dir") + assert checkpointed.count() == 100 # noqa: PLR2004 aggregate_plan = normalize_plan_text(checkpointed.groupBy("key").count()._explain_string()) # noqa: SLF001 ordered_plan = normalize_plan_text(checkpointed.sortWithinPartitions("key", "id")._explain_string()) # noqa: SLF001 + rows = checkpointed.selectExpr("id", "key", "spark_partition_id() AS pid").collect() assert "RepartitionExec" not in aggregate_plan assert "SortExec" not in ordered_plan + for partition_id in range(4): + partition_rows = [(row.key, row.id) for row in rows if row.pid == partition_id] + assert partition_rows == sorted(partition_rows) @pytest.mark.skipif(is_jvm_spark(), reason="Sail-specific physical plan names and stack configuration") diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index 7394ec74f0..e79d74e368 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -9,13 +9,14 @@ from pyspark.errors import PySparkException from pyspark.sql import SparkSession from pyspark.sql.connect import proto -from pyspark.sql.connect.plan import Checkpoint +from pyspark.sql.connect.plan import CachedRemoteRelation, Checkpoint, RemoveRemoteCachedRelation from pysail.testing.spark.session import ( configure_spark_session, patch_spark_connect_session, spark_connect_server, ) +from pysail.testing.spark.steps.plan import normalize_plan_text from pysail.testing.spark.utils.common import is_jvm_spark, pyspark_version pytestmark = [ @@ -88,6 +89,25 @@ def _assert_large_payload(checkpointed, rows): assert actual.payload_bytes == rows * 1024 +def _checkpoint(dataframe, kind, *, eager=True): + if kind == "local": + return dataframe.localCheckpoint(eager=eager) + return dataframe.checkpoint(eager=eager) + + +def _remove_cached_relation(spark, dataframe): + relation = dataframe._plan # noqa: SLF001 + client = spark._client # noqa: SLF001 + client.execute_command(RemoveRemoteCachedRelation(relation).command(client)) + relation._relation_id = None # noqa: SLF001 + + +def _hash_shuffle_count(plan): + if is_jvm_spark(): + return plan.count("Exchange hashpartitioning") + return plan.count("RepartitionExec: partitioning=Hash") + + @pytest.fixture(scope="module") def retry_spark(spark, checkpoint_path): if is_jvm_spark(): @@ -142,7 +162,10 @@ def test_local_checkpoint_preserves_rows_after_cache_repartition(spark, eager, p else: checkpointed = source.localCheckpoint(eager=eager, storageLevel=storage_level) - _assert_payload(checkpointed, 16 * 1024) + try: + _assert_payload(checkpointed, 16 * 1024) + finally: + _remove_cached_relation(spark, checkpointed) # FIXME: In local-cluster, this 64 MiB local checkpoint is embedded in every worker task definition and @@ -256,6 +279,132 @@ def test_checkpoint_command_can_reattach(spark): assert any(response.HasField("result_complete") for response in responses) +def test_checkpoint_command_preserves_operation_id(spark): + client = spark._client # noqa: SLF001 + command = Checkpoint(spark.range(3)._plan, local=True, eager=True).command(client) # noqa: SLF001 + operation_id = str(uuid.uuid4()) + request = client._execute_plan_request_with_metadata() # noqa: SLF001 + request.operation_id = operation_id + request.plan.command.CopyFrom(command) + + responses = list(client._stub.ExecutePlan(request, metadata=client._builder.metadata())) # noqa: SLF001 + + assert responses + assert all(response.operation_id == operation_id for response in responses) + result = next( + response.checkpoint_command_result for response in responses if response.HasField("checkpoint_command_result") + ) + relation = CachedRemoteRelation(result.relation.relation_id, spark) + try: + client.execute_command(RemoveRemoteCachedRelation(relation).command(client)) + finally: + relation._relation_id = None # noqa: SLF001 + + +def test_local_checkpoint_does_not_write_reliable_checkpoint_directory(spark, checkpoint_path): + files_before = {path.relative_to(checkpoint_path) for path in checkpoint_path.rglob("*") if path.is_file()} + + checkpointed = spark.range(0, 32, numPartitions=4).localCheckpoint( + eager=True, + storageLevel=StorageLevel.DISK_ONLY, + ) + + assert checkpointed.count() == 32 # noqa: PLR2004 + files_after = {path.relative_to(checkpoint_path) for path in checkpoint_path.rglob("*") if path.is_file()} + assert not files_after.difference(files_before) + + +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("kind", ["local", "reliable"]) +def test_checkpoint_preserves_exact_partition_count(spark, eager, kind): + source = spark.range(0, 1000, numPartitions=4).repartition(4) + + checkpointed = _checkpoint(source, kind, eager=eager) + + partition_ids = {row.pid for row in checkpointed.selectExpr("spark_partition_id() AS pid").distinct().collect()} + assert partition_ids == set(range(4)) + + +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("kind", ["local", "reliable"]) +def test_checkpoint_preserves_rows_and_order_within_each_partition(spark, eager, kind): + source = ( + spark.range(0, 100, numPartitions=4) + .selectExpr("id", "id % 4 AS key") + .repartition(4) + .sortWithinPartitions("key", "id") + ) + + checkpointed = _checkpoint(source, kind, eager=eager) + rows = checkpointed.selectExpr("id", "key", "spark_partition_id() AS pid").collect() + + assert len(rows) == 100 # noqa: PLR2004 + assert {row.pid for row in rows} == set(range(4)) + for partition_id in range(4): + partition_rows = [(row.key, row.id) for row in rows if row.pid == partition_id] + assert partition_rows == sorted(partition_rows) + + +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("kind", ["local", "reliable"]) +def test_checkpoint_scalar_subquery_does_not_survive_plan_truncation(spark, eager, kind): + source = spark.sql("SELECT id FROM range(1000) WHERE id = (SELECT max(id) FROM range(10))") + + checkpointed = _checkpoint(source, kind, eager=eager) + assert [row.id for row in checkpointed.collect()] == [9] + checkpoint_plan = normalize_plan_text(checkpointed._explain_string()) # noqa: SLF001 + + assert "Subquery" not in checkpoint_plan + assert [row.id for row in checkpointed.collect()] == [9] + + +@pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) +@pytest.mark.parametrize("kind", ["local", "reliable"]) +def test_checkpointed_aggregate_uses_preserved_hash_partitioning(spark, eager, kind): + source = spark.range(0, 128, numPartitions=4).selectExpr("id", "id % 16 AS key").repartition(4, "key") + + checkpointed = _checkpoint(source, kind, eager=eager) + assert checkpointed.count() == 128 # noqa: PLR2004 + aggregate = checkpointed.groupBy("key").count() + rows = aggregate.sort("key").collect() + + assert [(row.key, row["count"]) for row in rows] == [(key, 8) for key in range(16)] + if not is_jvm_spark(): + plan = normalize_plan_text(aggregate._explain_string()) # noqa: SLF001 + assert "RepartitionExec" not in plan + + +@pytest.mark.parametrize( + ("right_kind", "right_partitions", "sail_hash_shuffles"), + [("checkpoint", 4, 0), ("checkpoint", 3, 1), ("plain", None, 1)], + ids=["compatible", "incompatible", "unpartitioned-sibling"], +) +def test_checkpointed_join_inputs_keep_independent_partitioning( + spark, right_kind, right_partitions, sail_hash_shuffles +): + original_threshold = spark.conf.get("spark.sql.autoBroadcastJoinThreshold") + spark.conf.set("spark.sql.autoBroadcastJoinThreshold", "-1") + try: + left = ( + spark.range(0, 64, numPartitions=4) + .selectExpr("id AS key", "id AS left_value") + .repartition(4, "key") + .checkpoint() + ) + right = spark.range(0, 64, numPartitions=4).selectExpr("id AS key", "id AS right_value") + if right_kind == "checkpoint": + right = right.repartition(right_partitions, "key").checkpoint() + + joined = left.join(right, "key").select("key", "left_value", "right_value") + plan = normalize_plan_text(joined._explain_string()) # noqa: SLF001 + + assert joined.count() == 64 # noqa: PLR2004 + expected_hash_shuffles = 2 if is_jvm_spark() else sail_hash_shuffles + assert _hash_shuffle_count(plan) == expected_hash_shuffles + finally: + spark.conf.set("spark.sql.autoBroadcastJoinThreshold", original_threshold) + + @pytest.mark.parametrize("eager", [True, False], ids=["eager", "lazy"]) @pytest.mark.parametrize("kind", ["local", "reliable"]) def test_embedded_checkpoint_preserves_duplicate_column_names(retry_spark, eager, kind): From 00a3df97a43a6e853f830e06195d1c62ec5f5ae6 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 14:39:36 +0800 Subject: [PATCH 36/40] refactor(cache): account local checkpoints in memory pool --- crates/sail-cache/src/cached_relation.rs | 598 ++++++++++++++++++----- 1 file changed, 473 insertions(+), 125 deletions(-) diff --git a/crates/sail-cache/src/cached_relation.rs b/crates/sail-cache/src/cached_relation.rs index aef896078e..ffdb2babda 100644 --- a/crates/sail-cache/src/cached_relation.rs +++ b/crates/sail-cache/src/cached_relation.rs @@ -8,6 +8,7 @@ use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::physical_plan::ArrowSource; use datafusion::execution::disk_manager::RefCountedTempFile; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::physical_expr::Partitioning; @@ -18,6 +19,7 @@ use datafusion::physical_plan::{ with_new_children_if_necessary, }; use datafusion::prelude::SessionContext; +use datafusion_common::utils::memory::get_record_batch_memory_size; use datafusion_common::{DataFusionError, Result, Statistics, internal_datafusion_err}; use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; @@ -69,14 +71,20 @@ enum CachedRelationMaterialized { struct CachedRelationLocalMaterialized { schema: SchemaRef, properties: Arc, - memory_partitions: Option>>>, - serialized_memory_partitions: Option>>>>, - disk_partitions: Option>>, + partitions: Arc>>, } #[derive(Debug, Clone)] -struct CachedRelationDiskPartition { - chunks: Vec, +enum CachedRelationLocalChunk { + Deserialized { + batches: Vec, + _memory_reservation: Arc, + }, + Serialized { + bytes: Vec, + _memory_reservation: Arc, + }, + Disk(CachedRelationDiskChunk), } #[derive(Debug, Clone)] @@ -581,21 +589,11 @@ impl CachedRelationLocalMaterialized { )); } - let mut memory_partitions = if use_deserialized_memory { - Some(partition_maps(partition_count)) - } else { - None - }; - let mut serialized_partitions = if use_serialized_memory { - Some(partition_maps(partition_count)) - } else { - None - }; - let mut disk_partitions = if use_disk { - Some(partition_maps(partition_count)) - } else { - None - }; + let runtime_env = ctx.runtime_env(); + let memory_registration = Arc::new( + MemoryConsumer::new("local checkpoint cache").register(&runtime_env.memory_pool), + ); + let mut partitions = partition_maps(partition_count); let mut sequences: Vec> = (0..partition_count).map(|_| BTreeSet::new()).collect(); @@ -647,99 +645,63 @@ impl CachedRelationLocalMaterialized { )); } let bytes = data_array.value(row); - - if let Some(partitions) = memory_partitions.as_mut() { - partitions[partition].insert(sequence, read_record_batches(bytes)?); - } - if let Some(partitions) = serialized_partitions.as_mut() { - partitions[partition].insert(sequence, bytes.to_vec()); - } - if let Some(partitions) = disk_partitions.as_mut() { - partitions[partition].insert( - sequence, - write_disk_chunk(ctx, bytes, storage_level.replication).await?, - ); - } + let chunk = materialize_local_checkpoint_chunk( + ctx, + bytes, + use_deserialized_memory, + use_serialized_memory, + use_disk, + storage_level.replication, + memory_registration.as_ref(), + ) + .await?; + partitions[partition].insert(sequence, chunk); } } validate_partition_sequences(&sequences)?; - let memory_partitions = memory_partitions.map(|partitions| { - Arc::new( - partitions - .into_iter() - .map(|partition| partition.into_values().flatten().collect()) - .collect(), - ) - }); - let serialized_memory_partitions = serialized_partitions.map(|partitions| { - Arc::new( - partitions - .into_iter() - .map(|partition| partition.into_values().collect()) - .collect(), - ) - }); - let disk_partitions = disk_partitions.map(|partitions| { - Arc::new( - partitions - .into_iter() - .map(|partition| CachedRelationDiskPartition { - chunks: partition.into_values().collect(), - }) - .collect(), - ) - }); + let partitions = Arc::new( + partitions + .into_iter() + .map(|partition| partition.into_values().collect()) + .collect(), + ); Ok(Self { schema, properties, - memory_partitions, - serialized_memory_partitions, - disk_partitions, + partitions, }) } async fn to_physical_plan(&self) -> Result> { let partitions = self.load_partitions().await?; + // FIXME: Serialized and disk chunks are eagerly decoded into MemorySource without a + // bounded scan reservation. A distributed cache source is needed for streaming scans. let plan: Arc = MemorySourceConfig::try_new_exec(&partitions, Arc::clone(&self.schema), None)?; Ok(plan) } async fn load_partitions(&self) -> Result>> { - if let Some(partitions) = &self.memory_partitions { - return Ok(partitions.as_ref().clone()); - } - if let Some(partitions) = &self.serialized_memory_partitions { - let mut batches = Vec::with_capacity(partitions.len()); - for partition in partitions.iter() { - let mut partition_batches = vec![]; - for bytes in partition { - partition_batches.extend(read_record_batches(bytes)?); + let mut batches = Vec::with_capacity(self.partitions.len()); + for partition in self.partitions.iter() { + let mut partition_batches = vec![]; + for chunk in partition { + match chunk { + CachedRelationLocalChunk::Deserialized { + batches: chunk_batches, + .. + } => partition_batches.extend(chunk_batches.iter().cloned()), + CachedRelationLocalChunk::Serialized { bytes, .. } => { + partition_batches.extend(read_record_batches(bytes)?); + } + CachedRelationLocalChunk::Disk(chunk) => { + partition_batches.extend(chunk.read().await?); + } } - batches.push(partition_batches); - } - return Ok(batches); - } - if let Some(partitions) = &self.disk_partitions { - let mut batches = Vec::with_capacity(partitions.len()); - for partition in partitions.iter() { - batches.push(partition.read().await?); } - return Ok(batches); - } - Err(internal_datafusion_err!( - "cached relation has no materialized data" - )) - } -} - -impl CachedRelationDiskPartition { - async fn read(&self) -> Result> { - let mut batches = vec![]; - for chunk in &self.chunks { - batches.extend(chunk.read().await?); + batches.push(partition_batches); } Ok(batches) } @@ -768,6 +730,91 @@ fn partition_maps(partition_count: usize) -> Vec> { (0..partition_count).map(|_| BTreeMap::new()).collect() } +async fn materialize_local_checkpoint_chunk( + ctx: &SessionContext, + bytes: &[u8], + use_deserialized_memory: bool, + use_serialized_memory: bool, + use_disk: bool, + replication: usize, + memory_registration: &MemoryReservation, +) -> Result { + if use_deserialized_memory { + let memory_reservation = Arc::new(memory_registration.new_empty()); + match read_record_batches_with_reservation(bytes, memory_reservation.as_ref()) { + Ok(batches) => { + return Ok(CachedRelationLocalChunk::Deserialized { + batches, + _memory_reservation: memory_reservation, + }); + } + Err(DataFusionError::ResourcesExhausted(_)) if use_disk => {} + Err(error) => return Err(error), + } + } else if use_serialized_memory { + let memory_reservation = Arc::new(memory_registration.new_empty()); + match memory_reservation.try_grow(bytes.len()) { + Ok(()) => { + return Ok(CachedRelationLocalChunk::Serialized { + bytes: bytes.to_vec(), + _memory_reservation: memory_reservation, + }); + } + Err(error) if !use_disk => return Err(error), + Err(_) => {} + } + } + + if use_disk { + return Ok(CachedRelationLocalChunk::Disk( + write_disk_chunk(ctx, bytes, replication).await?, + )); + } + Err(internal_datafusion_err!( + "local checkpoint chunk has no available storage" + )) +} + +fn read_record_batches_with_reservation( + bytes: &[u8], + memory_reservation: &MemoryReservation, +) -> Result> { + let anticipated_size = bytes.len(); + memory_reservation.try_grow(anticipated_size)?; + let batches = match read_record_batches(bytes) { + Ok(batches) => batches, + Err(error) => { + memory_reservation.shrink(anticipated_size); + return Err(error); + } + }; + let actual_size = match record_batches_memory_size(&batches) { + Ok(size) => size, + Err(error) => { + memory_reservation.shrink(anticipated_size); + return Err(error); + } + }; + match actual_size.cmp(&anticipated_size) { + std::cmp::Ordering::Greater => { + if let Err(error) = memory_reservation.try_grow(actual_size - anticipated_size) { + memory_reservation.shrink(anticipated_size); + return Err(error); + } + } + std::cmp::Ordering::Less => memory_reservation.shrink(anticipated_size - actual_size), + std::cmp::Ordering::Equal => {} + } + Ok(batches) +} + +fn record_batches_memory_size(batches: &[RecordBatch]) -> Result { + batches.iter().try_fold(0usize, |size, batch| { + size.checked_add(get_record_batch_memory_size(batch)) + .ok_or_else(|| internal_datafusion_err!("checkpoint batch memory size is too large")) + }) +} + fn validate_partition_sequences(sequences: &[BTreeSet]) -> Result<()> { for (partition, sequences) in sequences.iter().enumerate() { if sequences.is_empty() { @@ -861,12 +908,13 @@ async fn materialize_reliable_checkpoint( } fn checkpoint_plan_properties(plan: &Arc) -> Arc { - Arc::new(PlanProperties::new( - plan.properties().eq_properties.clone(), - plan.output_partitioning().clone(), - EmissionType::Incremental, - Boundedness::Bounded, - )) + Arc::new( + plan.properties() + .as_ref() + .clone() + .with_emission_type(EmissionType::Incremental) + .with_boundedness(Boundedness::Bounded), + ) } fn create_arrow_checkpoint_scan( @@ -1012,8 +1060,11 @@ mod tests { use std::sync::Mutex; use async_trait::async_trait; - use datafusion::arrow::array::ArrayRef; + use datafusion::arrow::array::{ArrayRef, Int64Array}; use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::execution::context::SessionConfig; + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool}; + use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_expr::expressions::Column; use datafusion::physical_expr::{ EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, @@ -1026,6 +1077,43 @@ mod tests { use super::*; use crate::checkpoint::{CheckpointStore, ReliableCheckpoint}; + fn local_checkpoint_stream( + data_schema: &SchemaRef, + batches: &[RecordBatch], + ) -> Result { + local_checkpoint_chunk_stream(data_schema, &[batches.to_vec()]) + } + + fn local_checkpoint_chunk_stream( + data_schema: &SchemaRef, + chunks: &[Vec], + ) -> Result { + let bytes = chunks + .iter() + .map(|batches| write_record_batches(batches, data_schema.as_ref())) + .collect::>>()?; + let checkpoint_schema = Arc::new(Schema::new(vec![ + Field::new("partition", DataType::UInt64, false), + Field::new("sequence", DataType::UInt64, false), + Field::new("data", DataType::LargeBinary, false), + ])); + let columns: Vec = vec![ + Arc::new(UInt64Array::from(vec![0; bytes.len()])), + Arc::new(UInt64Array::from_iter_values( + 0..u64::try_from(bytes.len()) + .map_err(|_| internal_datafusion_err!("checkpoint chunk count is too large"))?, + )), + Arc::new(LargeBinaryArray::from_iter_values( + bytes.iter().map(Vec::as_slice), + )), + ]; + let batch = RecordBatch::try_new(Arc::clone(&checkpoint_schema), columns)?; + Ok(Box::pin(RecordBatchStreamAdapter::new( + checkpoint_schema, + stream::iter(vec![Ok(batch)]), + ))) + } + #[test] fn cached_relation_rewrite_preserves_exact_properties() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new( @@ -1108,29 +1196,16 @@ mod tests { #[tokio::test] async fn disk_only_local_checkpoint_materializes_only_on_disk() -> Result<()> { let data_schema = Arc::new(Schema::empty()); - let bytes = write_record_batches(&[], data_schema.as_ref())?; - let checkpoint_schema = Arc::new(Schema::new(vec![ - Field::new("partition", DataType::UInt64, false), - Field::new("sequence", DataType::UInt64, false), - Field::new("data", DataType::LargeBinary, false), - ])); - let columns: Vec = vec![ - Arc::new(UInt64Array::from(vec![0])), - Arc::new(UInt64Array::from(vec![0])), - Arc::new(LargeBinaryArray::from_vec(vec![bytes.as_slice()])), - ]; - let batch = RecordBatch::try_new(Arc::clone(&checkpoint_schema), columns)?; - let stream: SendableRecordBatchStream = Box::pin(RecordBatchStreamAdapter::new( - checkpoint_schema, - stream::iter(vec![Ok(batch)]), - )); + let stream = local_checkpoint_stream(&data_schema, &[])?; let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&data_schema))); let storage_level: spec::StorageLevel = "DISK_ONLY" .parse() .map_err(|error| internal_datafusion_err!("{error}"))?; + let ctx = SessionContext::new(); + let memory_pool = Arc::clone(&ctx.runtime_env().memory_pool); let materialized = CachedRelationLocalMaterialized::try_new( - &SessionContext::new(), + &ctx, data_schema, checkpoint_plan_properties(&input), 1, @@ -1139,19 +1214,292 @@ mod tests { ) .await?; - assert!(materialized.memory_partitions.is_none()); - assert!(materialized.serialized_memory_partitions.is_none()); - let disk_partitions = materialized - .disk_partitions - .as_ref() - .ok_or_else(|| internal_datafusion_err!("disk checkpoint partitions are missing"))?; - assert_eq!(disk_partitions.len(), 1); - assert_eq!(disk_partitions[0].chunks.len(), 1); - assert_eq!(disk_partitions[0].chunks[0].files.len(), 1); + assert_eq!(memory_pool.reserved(), 0); + assert_eq!(materialized.partitions.len(), 1); + assert_eq!(materialized.partitions[0].len(), 1); + let CachedRelationLocalChunk::Disk(chunk) = &materialized.partitions[0][0] else { + return Err(internal_datafusion_err!( + "disk-only checkpoint chunk was not stored on disk" + )); + }; + assert_eq!(chunk.files.len(), 1); assert_eq!(materialized.load_partitions().await?, vec![vec![]]); Ok(()) } + #[tokio::test] + async fn memory_and_disk_local_checkpoint_uses_memory_without_disk_duplication() -> Result<()> { + let data_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let data_batch = RecordBatch::try_new( + Arc::clone(&data_schema), + vec![Arc::new(Int64Array::from(vec![1, 2, 3]))], + )?; + let stream = local_checkpoint_stream(&data_schema, std::slice::from_ref(&data_batch))?; + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&data_schema))); + let storage_level: spec::StorageLevel = "MEMORY_AND_DISK" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?; + let ctx = SessionContext::new(); + let memory_pool = Arc::clone(&ctx.runtime_env().memory_pool); + + let materialized = CachedRelationLocalMaterialized::try_new( + &ctx, + data_schema, + checkpoint_plan_properties(&input), + 1, + stream, + storage_level, + ) + .await?; + + assert!(memory_pool.reserved() > 0); + let CachedRelationLocalChunk::Deserialized { + _memory_reservation: first_reservation, + .. + } = &materialized.partitions[0][0] + else { + return Err(internal_datafusion_err!( + "first checkpoint chunk was not stored in memory" + )); + }; + assert_eq!(memory_pool.reserved(), first_reservation.size()); + assert_eq!( + materialized.load_partitions().await?, + vec![vec![data_batch]] + ); + drop(materialized); + assert_eq!(memory_pool.reserved(), 0); + Ok(()) + } + + #[tokio::test] + async fn serialized_memory_local_checkpoint_tracks_reservation() -> Result<()> { + let data_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let data_batch = RecordBatch::try_new( + Arc::clone(&data_schema), + vec![Arc::new(Int64Array::from(vec![1, 2, 3]))], + )?; + let encoded_size = + write_record_batches(std::slice::from_ref(&data_batch), data_schema.as_ref())?.len(); + let stream = local_checkpoint_stream(&data_schema, std::slice::from_ref(&data_batch))?; + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&data_schema))); + let storage_level: spec::StorageLevel = "MEMORY_ONLY_SER" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?; + let ctx = SessionContext::new(); + let memory_pool = Arc::clone(&ctx.runtime_env().memory_pool); + + let materialized = CachedRelationLocalMaterialized::try_new( + &ctx, + data_schema, + checkpoint_plan_properties(&input), + 1, + stream, + storage_level, + ) + .await?; + + let CachedRelationLocalChunk::Serialized { + _memory_reservation: reservation, + .. + } = &materialized.partitions[0][0] + else { + return Err(internal_datafusion_err!( + "serialized checkpoint chunk was not stored in memory" + )); + }; + assert_eq!(reservation.size(), encoded_size); + assert_eq!(memory_pool.reserved(), encoded_size); + assert_eq!( + materialized.load_partitions().await?, + vec![vec![data_batch]] + ); + drop(materialized); + assert_eq!(memory_pool.reserved(), 0); + Ok(()) + } + + #[tokio::test] + async fn memory_and_disk_local_checkpoint_falls_back_to_disk_at_memory_limit() -> Result<()> { + let data_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let data_batch = RecordBatch::try_new( + Arc::clone(&data_schema), + vec![Arc::new(Int64Array::from(vec![1, 2, 3]))], + )?; + let stream = local_checkpoint_stream(&data_schema, std::slice::from_ref(&data_batch))?; + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&data_schema))); + let storage_level: spec::StorageLevel = "MEMORY_AND_DISK" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?; + let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(1)); + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&memory_pool)) + .build()?, + ); + let ctx = SessionContext::new_with_config_rt(SessionConfig::new(), runtime); + + let materialized = CachedRelationLocalMaterialized::try_new( + &ctx, + data_schema, + checkpoint_plan_properties(&input), + 1, + stream, + storage_level, + ) + .await?; + + assert_eq!(memory_pool.reserved(), 0); + assert!(matches!( + materialized.partitions[0][0], + CachedRelationLocalChunk::Disk(_) + )); + drop(data_batch); + drop(materialized); + assert_eq!(memory_pool.reserved(), 0); + Ok(()) + } + + #[tokio::test] + async fn memory_and_disk_local_checkpoint_can_mix_memory_and_disk_chunks() -> Result<()> { + let data_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let first_batch = RecordBatch::try_new( + Arc::clone(&data_schema), + vec![Arc::new(Int64Array::from(vec![1, 2, 3]))], + )?; + let second_batch = RecordBatch::try_new( + Arc::clone(&data_schema), + vec![Arc::new(Int64Array::from(vec![4, 5, 6]))], + )?; + let encoded = + write_record_batches(std::slice::from_ref(&first_batch), data_schema.as_ref())?; + let encoded_size = encoded.len(); + let decoded_size = record_batches_memory_size(&read_record_batches(&encoded)?)?; + let memory_limit = decoded_size + .checked_add(encoded_size.checked_mul(2).ok_or_else(|| { + internal_datafusion_err!("checkpoint test memory limit is invalid") + })?) + .ok_or_else(|| internal_datafusion_err!("checkpoint test memory limit is invalid"))?; + let stream = local_checkpoint_chunk_stream( + &data_schema, + &[vec![first_batch.clone()], vec![second_batch.clone()]], + )?; + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&data_schema))); + let storage_level: spec::StorageLevel = "MEMORY_AND_DISK" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?; + let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(memory_limit)); + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&memory_pool)) + .build()?, + ); + let ctx = SessionContext::new_with_config_rt(SessionConfig::new(), runtime); + let execution_memory = + MemoryConsumer::new("checkpoint materialization input").register(&memory_pool); + execution_memory.try_grow(encoded_size + 1)?; + + let materialized = CachedRelationLocalMaterialized::try_new( + &ctx, + data_schema, + checkpoint_plan_properties(&input), + 1, + stream, + storage_level, + ) + .await?; + + let CachedRelationLocalChunk::Deserialized { + _memory_reservation: first_reservation, + .. + } = &materialized.partitions[0][0] + else { + return Err(internal_datafusion_err!( + "first checkpoint chunk was not stored in memory" + )); + }; + assert!(matches!( + materialized.partitions[0][1], + CachedRelationLocalChunk::Disk(_) + )); + assert_eq!( + memory_pool.reserved(), + first_reservation.size() + execution_memory.size() + ); + drop(execution_memory); + assert_eq!( + materialized.load_partitions().await?, + vec![vec![first_batch, second_batch]] + ); + assert_eq!(memory_pool.reserved(), first_reservation.size()); + drop(materialized); + assert_eq!(memory_pool.reserved(), 0); + Ok(()) + } + + #[tokio::test] + async fn memory_only_checkpoint_releases_failed_reservation() -> Result<()> { + let data_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let data_batch = RecordBatch::try_new( + Arc::clone(&data_schema), + vec![Arc::new(Int64Array::from(vec![1, 2, 3]))], + )?; + let stream = local_checkpoint_stream(&data_schema, &[data_batch])?; + let input: Arc = Arc::new(EmptyExec::new(Arc::clone(&data_schema))); + let storage_level: spec::StorageLevel = "MEMORY_ONLY" + .parse() + .map_err(|error| internal_datafusion_err!("{error}"))?; + let memory_pool: Arc = Arc::new(GreedyMemoryPool::new(1)); + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&memory_pool)) + .build()?, + ); + let ctx = SessionContext::new_with_config_rt(SessionConfig::new(), runtime); + + let error = match CachedRelationLocalMaterialized::try_new( + &ctx, + data_schema, + checkpoint_plan_properties(&input), + 1, + stream, + storage_level, + ) + .await + { + Ok(_) => { + return Err(internal_datafusion_err!( + "memory-only checkpoint did not respect the memory limit" + )); + } + Err(error) => error, + }; + + assert!(matches!(error, DataFusionError::ResourcesExhausted(_))); + assert_eq!(memory_pool.reserved(), 0); + Ok(()) + } + #[tokio::test] async fn cached_relation_logical_plan_renames_duplicate_fields_by_position() -> Result<()> { let schema = Arc::new(Schema::new(vec![ From c5512e21da2df66a82c29c9db407bd26e9f8b2a9 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 14:40:10 +0800 Subject: [PATCH 37/40] fix(object-store): reject URL user information --- crates/sail-object-store/src/path.rs | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/crates/sail-object-store/src/path.rs b/crates/sail-object-store/src/path.rs index 5552e72d32..42383da949 100644 --- a/crates/sail-object-store/src/path.rs +++ b/crates/sail-object-store/src/path.rs @@ -66,6 +66,11 @@ pub fn resolve_object_store_path( ) -> Result { let directory = format!("{}/", path.trim_end_matches('/')); let url = ListingTableUrl::parse(&directory)?; + if !url.get_url().username().is_empty() || url.get_url().password().is_some() { + return Err(DataFusionError::Plan( + "object store URL cannot contain user information".to_string(), + )); + } let object_store_url = url.object_store(); let store = runtime_env.object_store(&object_store_url)?; Ok(ResolvedObjectStorePath { @@ -125,4 +130,27 @@ mod tests { assert!(resolved.store().head(&outside).await.is_ok()); Ok(()) } + + #[test] + fn resolved_path_rejects_url_user_information() -> Result<()> { + let runtime_env = RuntimeEnv::default(); + let error = match resolve_object_store_path( + &runtime_env, + "s3://checkpoint-user:checkpoint-secret@bucket/checkpoint", + ) { + Ok(_) => { + return Err(DataFusionError::Plan( + "object store credentials were accepted in a URL".to_string(), + )); + } + Err(error) => error, + }; + + assert!( + error + .to_string() + .contains("object store URL cannot contain user information") + ); + Ok(()) + } } From 649759f952ce68c76b9ff9a528e9f2d4b13b974c Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 14:40:14 +0800 Subject: [PATCH 38/40] refactor(object-store): stream checkpoint prefix cleanup --- crates/sail-object-store/src/lib.rs | 5 +++- crates/sail-object-store/src/path.rs | 36 +++++++++++++++++---------- crates/sail-session/src/checkpoint.rs | 18 +++----------- 3 files changed, 31 insertions(+), 28 deletions(-) diff --git a/crates/sail-object-store/src/lib.rs b/crates/sail-object-store/src/lib.rs index afd3d3fafe..eab69ee8cd 100644 --- a/crates/sail-object-store/src/lib.rs +++ b/crates/sail-object-store/src/lib.rs @@ -4,5 +4,8 @@ mod path; mod registry; mod s3; -pub use path::{ResolvedObjectStorePath, delete_object_store_prefix, resolve_object_store_path}; +pub use path::{ + ResolvedObjectStorePath, delete_object_store_prefix, delete_object_store_prefix_objects, + resolve_object_store_path, +}; pub use registry::DynamicObjectStoreRegistry; diff --git a/crates/sail-object-store/src/path.rs b/crates/sail-object-store/src/path.rs index 42383da949..fd2c9833de 100644 --- a/crates/sail-object-store/src/path.rs +++ b/crates/sail-object-store/src/path.rs @@ -4,7 +4,7 @@ use datafusion::datasource::listing::ListingTableUrl; use datafusion::execution::object_store::ObjectStoreUrl; use datafusion::execution::runtime_env::RuntimeEnv; use datafusion_common::{DataFusionError, Result}; -use futures::TryStreamExt; +use futures::StreamExt; use object_store::path::Path; use object_store::{ObjectMeta, ObjectStore, ObjectStoreExt, PutPayload}; @@ -44,20 +44,30 @@ impl ResolvedObjectStorePath { } pub async fn delete_prefix(&self) -> Result<()> { - let files = self - .store - .list(Some(&self.prefix)) - .try_collect::>() - .await - .map_err(|e| DataFusionError::ObjectStore(Box::new(e)))?; - for file in files { - match self.store.delete(&file.location).await { - Ok(()) | Err(object_store::Error::NotFound { .. }) => {} - Err(error) => return Err(DataFusionError::ObjectStore(Box::new(error))), - } + delete_object_store_prefix_objects(self.store.as_ref(), &self.prefix).await + } +} + +pub async fn delete_object_store_prefix_objects( + store: &dyn ObjectStore, + prefix: &Path, +) -> Result<()> { + let locations = store.list(Some(prefix)).map(|result| { + result + .map(|object| object.location) + .map_err(|source| object_store::Error::Generic { + store: "object store listing", + source: Box::new(source), + }) + }); + let mut deleted = store.delete_stream(Box::pin(locations)); + while let Some(result) = deleted.next().await { + match result { + Ok(_) | Err(object_store::Error::NotFound { .. }) => {} + Err(error) => return Err(DataFusionError::ObjectStore(Box::new(error))), } - Ok(()) } + Ok(()) } pub fn resolve_object_store_path( diff --git a/crates/sail-session/src/checkpoint.rs b/crates/sail-session/src/checkpoint.rs index 3ce27b13ff..73eec15322 100644 --- a/crates/sail-session/src/checkpoint.rs +++ b/crates/sail-session/src/checkpoint.rs @@ -8,13 +8,14 @@ use datafusion::object_store::{ObjectMeta, ObjectStoreExt}; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion::prelude::SessionContext; use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; -use futures::{StreamExt, TryStreamExt}; +use futures::StreamExt; use sail_cache::checkpoint::{CheckpointStore, ReliableCheckpoint}; use sail_common_datafusion::array::record_batch::write_record_batches_file; use sail_common_datafusion::extension::SessionExtensionAccessor; use sail_common_datafusion::session::job::JobService; use sail_object_store::{ - ResolvedObjectStorePath, delete_object_store_prefix, resolve_object_store_path, + ResolvedObjectStorePath, delete_object_store_prefix, delete_object_store_prefix_objects, + resolve_object_store_path, }; use sail_physical_plan::checkpoint::ReliableCheckpointExec; @@ -169,18 +170,7 @@ async fn commit_checkpoint_attempts( } let temporary = checkpoint_path.child("_temporary"); - let temporary_files = checkpoint_path - .store() - .list(Some(&temporary)) - .try_collect::>() - .await - .map_err(|error| DataFusionError::ObjectStore(Box::new(error)))?; - for file in temporary_files { - match checkpoint_path.store().delete(&file.location).await { - Ok(()) | Err(datafusion::object_store::Error::NotFound { .. }) => {} - Err(error) => return Err(DataFusionError::ObjectStore(Box::new(error))), - } - } + delete_object_store_prefix_objects(checkpoint_path.store().as_ref(), &temporary).await?; Ok(files) } From 5b9019fa6ff855ee38e9fa0f0aba3b8cd2591c07 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 14:40:19 +0800 Subject: [PATCH 39/40] chore(spark-connect): note checkpoint executor follow-up --- crates/sail-spark-connect/src/service/plan_executor.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index fd3471e45e..9ef0e92427 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -596,6 +596,8 @@ pub(crate) async fn handle_execute_checkpoint_command( let result = CheckpointCommandResult { relation: Some(CachedRemoteRelation { relation_id }), }; + // TODO: Run checkpoint commands through Executor so Spark Connect reattach and interrupt + // share the normal operation lifecycle. let mut output = vec![ExecutorOutput::new(ExecutorBatch::CheckpointCommandResult( Box::new(result), ))]; From 63a030a3fa03ea2eb99de225772526641e564b69 Mon Sep 17 00:00:00 2001 From: XL Liang Date: Fri, 17 Jul 2026 14:57:43 +0800 Subject: [PATCH 40/40] linter --- .../src/service/plan_executor.rs | 28 ------------------- .../tests/spark/dataframe/test_checkpoint.py | 6 ++-- .../tests/spark/execution/test_checkpoint.py | 12 ++++---- 3 files changed, 9 insertions(+), 37 deletions(-) diff --git a/crates/sail-spark-connect/src/service/plan_executor.rs b/crates/sail-spark-connect/src/service/plan_executor.rs index 9ef0e92427..58e69c6f34 100644 --- a/crates/sail-spark-connect/src/service/plan_executor.rs +++ b/crates/sail-spark-connect/src/service/plan_executor.rs @@ -729,34 +729,6 @@ fn validate_checkpoint_path_segment(name: &str, value: &str) -> SparkResult<()> } } -#[cfg(test)] -mod checkpoint_path_tests { - use super::*; - - #[test] - fn checkpoint_path_appends_encoded_segments() -> SparkResult<()> { - let path = build_checkpoint_path("memory:///checkpoint-root", "session?one", "relation")?; - - assert_eq!( - path, - "memory:///checkpoint-root/sail-checkpoints/session%3Fone/relation" - ); - Ok(()) - } - - #[test] - fn checkpoint_path_rejects_url_query() { - let error = build_checkpoint_path( - "memory:///checkpoint-root?scope=other", - "session", - "relation", - ) - .expect_err("checkpoint root query must be rejected"); - - assert!(error.to_string().contains("cannot contain a query")); - } -} - async fn cleanup_checkpoint_after_error( ctx: &SessionContext, path: &str, diff --git a/python/pysail/tests/spark/dataframe/test_checkpoint.py b/python/pysail/tests/spark/dataframe/test_checkpoint.py index 4b1cf7d9d3..a02ce34033 100644 --- a/python/pysail/tests/spark/dataframe/test_checkpoint.py +++ b/python/pysail/tests/spark/dataframe/test_checkpoint.py @@ -8,8 +8,8 @@ from pandas.testing import assert_frame_equal from pyspark import StorageLevel from pyspark.errors import PySparkException +from pyspark.sql.connect import plan as connect_plan from pyspark.sql.connect.dataframe import DataFrame -from pyspark.sql.connect.plan import CachedRemoteRelation, RemoveRemoteCachedRelation from pyspark.sql.functions import col, lit from pysail.testing.spark.session import spark_session_factory @@ -48,7 +48,7 @@ def _cached_relation_id(dataframe): def _wait_for_cached_relation_removal(spark, relation_id): deadline = time.monotonic() + 5 while time.monotonic() < deadline: - plan = CachedRemoteRelation(relation_id, spark) + plan = connect_plan.CachedRemoteRelation(relation_id, spark) probe = DataFrame(plan, spark) try: probe.count() @@ -345,7 +345,7 @@ def test_dataframe_local_checkpoint_explicit_cleanup_removes_cached_relation(spa relation = checkpointed._plan # noqa: SLF001 client = spark._client # noqa: SLF001 - client.execute_command(RemoveRemoteCachedRelation(relation).command(client)) + client.execute_command(connect_plan.RemoveRemoteCachedRelation(relation).command(client)) relation._relation_id = None # noqa: SLF001 _wait_for_cached_relation_removal(spark, relation_id) diff --git a/python/pysail/tests/spark/execution/test_checkpoint.py b/python/pysail/tests/spark/execution/test_checkpoint.py index e79d74e368..77cac3d357 100644 --- a/python/pysail/tests/spark/execution/test_checkpoint.py +++ b/python/pysail/tests/spark/execution/test_checkpoint.py @@ -8,8 +8,8 @@ from pyspark import StorageLevel from pyspark.errors import PySparkException from pyspark.sql import SparkSession +from pyspark.sql.connect import plan as connect_plan from pyspark.sql.connect import proto -from pyspark.sql.connect.plan import CachedRemoteRelation, Checkpoint, RemoveRemoteCachedRelation from pysail.testing.spark.session import ( configure_spark_session, @@ -98,7 +98,7 @@ def _checkpoint(dataframe, kind, *, eager=True): def _remove_cached_relation(spark, dataframe): relation = dataframe._plan # noqa: SLF001 client = spark._client # noqa: SLF001 - client.execute_command(RemoveRemoteCachedRelation(relation).command(client)) + client.execute_command(connect_plan.RemoveRemoteCachedRelation(relation).command(client)) relation._relation_id = None # noqa: SLF001 @@ -246,7 +246,7 @@ def test_checkpoint_gc_releases_completed_cluster_job_lease(spark, tmp_path): @SAIL_XFAIL def test_checkpoint_command_can_reattach(spark): client = spark._client # noqa: SLF001 - command = Checkpoint(spark.range(3)._plan, local=True, eager=False).command(client) # noqa: SLF001 + command = connect_plan.Checkpoint(spark.range(3)._plan, local=True, eager=False).command(client) # noqa: SLF001 request = client._execute_plan_request_with_metadata() # noqa: SLF001 request.operation_id = str(uuid.uuid4()) request.plan.command.CopyFrom(command) @@ -281,7 +281,7 @@ def test_checkpoint_command_can_reattach(spark): def test_checkpoint_command_preserves_operation_id(spark): client = spark._client # noqa: SLF001 - command = Checkpoint(spark.range(3)._plan, local=True, eager=True).command(client) # noqa: SLF001 + command = connect_plan.Checkpoint(spark.range(3)._plan, local=True, eager=True).command(client) # noqa: SLF001 operation_id = str(uuid.uuid4()) request = client._execute_plan_request_with_metadata() # noqa: SLF001 request.operation_id = operation_id @@ -294,9 +294,9 @@ def test_checkpoint_command_preserves_operation_id(spark): result = next( response.checkpoint_command_result for response in responses if response.HasField("checkpoint_command_result") ) - relation = CachedRemoteRelation(result.relation.relation_id, spark) + relation = connect_plan.CachedRemoteRelation(result.relation.relation_id, spark) try: - client.execute_command(RemoveRemoteCachedRelation(relation).command(client)) + client.execute_command(connect_plan.RemoveRemoteCachedRelation(relation).command(client)) finally: relation._relation_id = None # noqa: SLF001