feat: add extension to support DataFrame::checkpoint()

ballista/client/Cargo.toml

@@ -39,6 +39,7 @@ log = { workspace = true }
 
 tokio = { workspace = true }
 url = { workspace = true }
+uuid = { workspace = true }
 
 [dev-dependencies]
 ballista-executor = { path = "../executor", version = "54.0.0" }

ballista/client/src/extension.rs

@@ -15,12 +15,22 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use std::sync::Arc;
+
 pub use ballista_core::extension::{SessionConfigExt, SessionStateExt};
-use ballista_core::serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient;
+use ballista_core::{
+    extension::BallistaCheckpointNode,
+    serde::protobuf::scheduler_grpc_client::SchedulerGrpcClient,
+};
 use datafusion::{
-    error::DataFusionError, execution::SessionState, prelude::SessionContext,
+    dataframe::{DataFrame, DataFrameWriteOptions},
+    error::DataFusionError,
+    execution::{SessionState, options::ParquetReadOptions},
+    logical_expr::LogicalPlan,
+    prelude::SessionContext,
 };
 use url::Url;
+use uuid::Uuid;
 
 const DEFAULT_SCHEDULER_PORT: u16 = 50050;
 
@@ -234,3 +244,90 @@ impl Extension {
         Ok(scheduler_url)
     }
 }
+
+/// Providing [DataFrameExt] for an extended functionality on DataFusion DataFrame.
+///
+/// Checkpointing truncates the logical plan by materialising it to the location
+/// configured with `ballista.checkpoint.dir` and continuing from a plain scan of
+/// that location. This mirrors Spark's `Dataset.checkpoint`, including its split
+/// between eager and lazy variants.
+///
+/// # Storage requirements
+///
+/// The checkpoint is written by the executors and read back by the client, so
+/// `ballista.checkpoint.dir` **must** point at storage that is visible to the
+/// whole cluster (an object store URL, or a shared filesystem). A plain local
+/// path only works for a standalone/single node cluster. This matches Spark,
+/// where `setCheckpointDir` is expected to be an HDFS compatible path.
+#[async_trait::async_trait]
+pub trait DataFrameExt {
+    /// Eagerly checkpoints dataframe. Equivalent of Spark's DataFrame.checkpoint(eager=True)
+    /// The plan is executed immediately as a
+    /// distributed job, the result is written to `ballista.checkpoint.dir`, and
+    /// the returned DataFrame is a fresh scan of that data with no lineage back
+    /// to the original plan.
+    async fn checkpoint(self) -> datafusion::error::Result<DataFrame>;
+
+    /// Lazily checkpoints dataframe. Equivalent of Spark's DataFrame.checkpoint(eager=False)
+    /// A checkpoint
+    /// marker is inserted into the logical plan and the materialisation happens
+    /// on the scheduler when the first action (`collect`, `write`, etc) is
+    /// executed against the returned DataFrame.
+    fn checkpoint_lazy(self) -> datafusion::error::Result<DataFrame>;
+}
+
+#[async_trait::async_trait]
+impl DataFrameExt for DataFrame {
+    async fn checkpoint(self) -> datafusion::error::Result<DataFrame> {
+        let (state, plan) = self.into_parts();
+        let path = checkpoint_location(&state)?;
+        let ctx = SessionContext::new_with_state(state);
+
+        // Executes the original plan as a normal distributed job.
+        ctx.execute_logical_plan(plan)
+            .await?
+            .write_parquet(&path, DataFrameWriteOptions::new(), None)
+            .await?;
+
+        // Fresh TableScan-rooted plan: lineage to the original plan is gone.
+        ctx.read_parquet(&path, ParquetReadOptions::default()).await
+    }
+
+    fn checkpoint_lazy(self) -> datafusion::error::Result<DataFrame> {
+        let (state, plan) = self.into_parts();
+        let location = checkpoint_location(&state)?;
+        let session_id = state.session_id().to_string();
+
+        let plan = LogicalPlan::Extension(datafusion::logical_expr::Extension {
+            node: Arc::new(BallistaCheckpointNode::new(
+                Uuid::new_v4().to_string(),
+                session_id,
+                location,
+                plan,
+            )),
+        });
+
+        Ok(DataFrame::new(state, plan))
+    }
+}
+
+/// Resolves a unique location for a single checkpoint under the configured
+/// `ballista.checkpoint.dir`.
+///
+/// The location is derived once, on the client, so that a lazy checkpoint node
+/// is fully self describing by the time it reaches the scheduler.
+fn checkpoint_location(state: &SessionState) -> datafusion::error::Result<String> {
+    let base_dir = state.config().ballista_checkpoint_dir().ok_or_else(|| {
+        DataFusionError::Configuration(
+            "ballista.checkpoint.dir must be set to use DataFrame::checkpoint()"
+                .to_string(),
+        )
+    })?;
+
+    Ok(format!(
+        "{}/{}/{}",
+        base_dir.trim_end_matches('/'),
+        state.session_id(),
+        Uuid::new_v4()
+    ))
+}

ballista/client/src/prelude.rs

@@ -17,4 +17,4 @@
 
 //! Ballista Prelude (common imports)
 
-pub use crate::extension::{SessionConfigExt, SessionContextExt};
+pub use crate::extension::{DataFrameExt, SessionConfigExt, SessionContextExt};

ballista/client/tests/checkpoint.rs

@@ -0,0 +1,164 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+mod common;
+
+use ballista::prelude::DataFrameExt;
+use datafusion::assert_batches_eq;
+use datafusion::logical_expr::LogicalPlan;
+use tempfile::TempDir;
+
+#[tokio::test]
+async fn should_checkpoint_dataframe() -> datafusion::error::Result<()> {
+    let checkpoint_dir = TempDir::new().unwrap();
+    let checkpoint_path = checkpoint_dir.path().to_str().unwrap().to_string();
+
+    let ctx = common::standalone_context_with_checkpoint_dir(&checkpoint_path).await;
+
+    let test_data = common::example_test_data();
+    ctx.register_parquet(
+        "test",
+        &format!("{test_data}/alltypes_plain.parquet"),
+        Default::default(),
+    )
+    .await?;
+
+    let df = ctx
+        .sql("select string_col, timestamp_col from test where id > 4")
+        .await?;
+
+    let expected = [
+        "+------------+---------------------+",
+        "| string_col | timestamp_col       |",
+        "+------------+---------------------+",
+        "| 31         | 2009-03-01T00:01:00 |",
+        "| 30         | 2009-04-01T00:00:00 |",
+        "| 31         | 2009-04-01T00:01:00 |",
+        "+------------+---------------------+",
+    ];
+
+    // sanity check the plan pre-checkpoint isn't already a bare TableScan
+    assert!(!matches!(df.logical_plan(), LogicalPlan::TableScan(_)));
+
+    let checkpointed = df.checkpoint().await?;
+
+    // lineage is broken: the new DataFrame is a scan over the checkpoint files,
+    // not the filter/projection chain that produced them
+    assert!(matches!(
+        checkpointed.logical_plan(),
+        LogicalPlan::TableScan(_)
+    ));
+
+    let result = checkpointed.collect().await?;
+    assert_batches_eq!(expected, &result);
+
+    // the checkpoint actually landed on the configured storage
+    let written_any_files = std::fs::read_dir(checkpoint_dir.path())
+        .unwrap()
+        .next()
+        .is_some();
+    assert!(written_any_files, "expected checkpoint files to be written");
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn checkpoint_without_configured_dir_errors() -> datafusion::error::Result<()> {
+    // no with_ballista_checkpoint_dir() call
+    let ctx = common::standalone_context_with_state().await;
+
+    let test_data = common::example_test_data();
+    ctx.register_parquet(
+        "test",
+        &format!("{test_data}/alltypes_plain.parquet"),
+        Default::default(),
+    )
+    .await?;
+
+    let df = ctx.sql("select * from test").await?;
+
+    let err = df.checkpoint().await.unwrap_err();
+    assert!(err.to_string().contains("ballista.checkpoint.dir"));
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn should_insert_checkpoint_node_without_executing() -> datafusion::error::Result<()>
+{
+    let checkpoint_dir = TempDir::new().unwrap();
+    let checkpoint_path = checkpoint_dir.path().to_str().unwrap().to_string();
+
+    let ctx = common::standalone_context_with_checkpoint_dir(&checkpoint_path).await;
+
+    let test_data = common::example_test_data();
+    ctx.register_parquet(
+        "test",
+        &format!("{test_data}/alltypes_plain.parquet"),
+        Default::default(),
+    )
+    .await?;
+
+    let df = ctx
+        .sql("select string_col, timestamp_col from test where id > 4")
+        .await?;
+    let schema_before = df.schema().clone();
+
+    let checkpointed = df.checkpoint_lazy()?;
+
+    // The marker sits at the root and the schema is unchanged: the node is a
+    // pass through until the scheduler materialises it.
+    assert!(matches!(
+        checkpointed.logical_plan(),
+        LogicalPlan::Extension(_)
+    ));
+    assert_eq!(checkpointed.schema(), &schema_before);
+    assert!(
+        format!("{}", checkpointed.logical_plan().display_indent())
+            .contains("BallistaCheckpointNode")
+    );
+
+    // Nothing has been materialised: a lazy checkpoint performs no I/O.
+    let written_any_files = std::fs::read_dir(checkpoint_dir.path())
+        .unwrap()
+        .next()
+        .is_some();
+    assert!(!written_any_files, "lazy checkpoint must not write eagerly");
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn lazy_checkpoint_without_configured_dir_errors() -> datafusion::error::Result<()>
+{
+    let ctx = common::standalone_context_with_state().await;
+
+    let test_data = common::example_test_data();
+    ctx.register_parquet(
+        "test",
+        &format!("{test_data}/alltypes_plain.parquet"),
+        Default::default(),
+    )
+    .await?;
+
+    let df = ctx.sql("select * from test").await?;
+
+    let err = df.checkpoint_lazy().unwrap_err();
+    assert!(err.to_string().contains("ballista.checkpoint.dir"));
+
+    Ok(())
+}

ballista/client/tests/common/mod.rs

@@ -156,6 +156,17 @@ pub async fn setup_test_cluster_with_state(session_state: SessionState) -> (Stri
     (host, addr.port())
 }
 
+#[allow(dead_code)]
+pub async fn standalone_context_with_checkpoint_dir(dir: &str) -> SessionContext {
+    let config =
+        SessionConfig::new_with_ballista().with_ballista_checkpoint_dir(dir.to_string());
+    let state = SessionStateBuilder::new()
+        .with_config(config)
+        .with_default_features()
+        .build();
+    SessionContext::standalone_with_state(state).await.unwrap()
+}
+
 #[allow(dead_code)]
 pub async fn setup_test_cluster_with_builders(
     config_producer: ConfigProducer,

ballista/core/proto/ballista.proto

@@ -33,6 +33,7 @@ import "datafusion_common.proto";
 message BallistaLogicalPlanNode {
   oneof LogicalPlanType {
     LogicalPlanCacheNode cache_node = 1;
+    LogicalPlanCheckpointNode checkpoint_node = 2;
   }
 }
 
@@ -41,6 +42,12 @@ message LogicalPlanCacheNode {
   string session_id = 2;
 }
 
+message LogicalPlanCheckpointNode {
+  string checkpoint_id = 1;
+  string session_id = 2;
+  string location = 3;
+}
+
 ///////////////////////////////////////////////////////////////////////////////////////////////////
 // Ballista Physical Plan
 ///////////////////////////////////////////////////////////////////////////////////////////////////

ballista/core/src/config.rs

@@ -176,6 +176,9 @@ pub const BALLISTA_SHUFFLE_COMPRESSION_CODEC: &str = "ballista.shuffle.compressi
 pub const BALLISTA_SCHEDULER_MAX_PARTITIONS_PER_TASK: &str =
     "ballista.scheduler.max_partitions_per_task";
 
+/// Configuration key for the base object-store location used by `DataFrame::checkpoint()`.
+pub const BALLISTA_CHECKPOINT_DIR: &str = "ballista.checkpoint.dir";
+
 /// Result type for configuration parsing operations.
 pub type ParseResult<T> = result::Result<T, String>;
 use std::sync::LazyLock;
@@ -425,6 +428,16 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
             DataType::UInt64,
             Some(1.to_string()),
         ),
+        ConfigEntry::new(
+            BALLISTA_CHECKPOINT_DIR.to_string(),
+            "Optional location where DataFrame::checkpoint() results are stored. \
+             Must be reachable from every node in the cluster, so an object store URL \
+             or a shared filesystem should be used for anything but a standalone cluster. \
+             If it is not set, calling DataFrame::checkpoint() returns an error.".to_string(),
+            DataType::Utf8,
+            Some("".to_string())
+        )
+        .with_doc_default("(none)"),
     ];
     entries
         .into_iter()
@@ -811,6 +824,12 @@ impl BallistaConfig {
         self.get_usize_setting(BALLISTA_CLIENT_IO_RETRY_WAIT_TIME_MS)
     }
 
+    /// Returns configured checkpoint directory (local FS, HDFS, S3 etc.)
+    pub fn checkpoint_dir(&self) -> Option<String> {
+        let dir = self.get_string_setting(BALLISTA_CHECKPOINT_DIR);
+        if dir.is_empty() { None } else { Some(dir) }
+    }
+
     fn get_usize_setting(&self, key: &str) -> usize {
         if let Some(v) = self.settings.get(key) {
             // infallible because we validate all configs in the constructor