You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "sandugood (via GitHub)" <gi...@apache.org> on 2026/07/11 11:46:35 UTC
[PR] feat: add extension for `DataFrame::checkpoint()` [datafusion-ballista]
sandugood opened a new pull request, #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993
# Which issue does this PR close?
Closes #1937 .
# Rationale for this change
Before this change, there was no `checkpoint()` method for the DataFrame in Ballista (where we break lineage of the plan and store intermediate results to specified location)
# What changes are included in this PR?
Introducing new `DataFrameExt` trait in the client, also introducing `checkpoint()` there.
A new entry to the `BallistaConfig` was added as well. When trying to call `checkpoint() `without the config entry set user will get an error.
# Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
New `checkpoint()` method introduced.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-5553656081
I will try and fix all of the comments asap. Then we can proceed with `.cache()`
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#discussion_r3735341567
##########
ballista/client/src/extension.rs:
##########
@@ -234,3 +237,46 @@ impl Extension {
Ok(scheduler_url)
}
}
+
+/// Providing [DataFrameExt] for an extended functionality on DataFusion DataFrame.
+///
+#[async_trait::async_trait]
+pub trait DataFrameExt {
+ /// Checkpointing DataFrame - storing intermediate result to disk and breaking lineage in the plan
+ async fn checkpoint(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 ctx = SessionContext::new_with_state(state);
+
+ let base_dir =
+ ctx.state()
+ .config()
+ .ballista_checkpoint_dir()
+ .ok_or_else(|| {
+ DataFusionError::Configuration(
+ "ballista.checkpoint.dir must be set to use DataFrame::checkpoint()"
+ .to_string(),
+ )
+ })?;
+
+ let path = format!(
+ "{}/{}/{}",
+ base_dir.trim_end_matches('/'),
+ ctx.state().session_id(),
+ Uuid::new_v4()
+ );
+
+ // Executes the original plan as a normal distributed job.
+ ctx.execute_logical_plan(plan)
Review Comment:
I chose the path of introducing `BallistaCheckpointNode` (alongside the `BallistaCacheNode`), and also adding lazy materialization (to mirror Spark's functionality)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#discussion_r3734935689
##########
ballista/client/src/extension.rs:
##########
@@ -234,3 +237,46 @@ impl Extension {
Ok(scheduler_url)
}
}
+
+/// Providing [DataFrameExt] for an extended functionality on DataFusion DataFrame.
+///
+#[async_trait::async_trait]
+pub trait DataFrameExt {
+ /// Checkpointing DataFrame - storing intermediate result to disk and breaking lineage in the plan
+ async fn checkpoint(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 ctx = SessionContext::new_with_state(state);
+
+ let base_dir =
+ ctx.state()
+ .config()
+ .ballista_checkpoint_dir()
+ .ok_or_else(|| {
+ DataFusionError::Configuration(
+ "ballista.checkpoint.dir must be set to use DataFrame::checkpoint()"
+ .to_string(),
+ )
+ })?;
+
+ let path = format!(
+ "{}/{}/{}",
+ base_dir.trim_end_matches('/'),
+ ctx.state().session_id(),
+ Uuid::new_v4()
+ );
+
+ // Executes the original plan as a normal distributed job.
+ ctx.execute_logical_plan(plan)
Review Comment:
this is one way to implement this, but difference from spark that ballista checkpoint will materialize the dag (up to that point) instead of waiting standard actions, such as write.
this will work for first version, but maybe we should consider adding new logical plan extension which would represent checkpoint, wait for write op and then split plan on scheduler side using planner.
wdyt @sandugood
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood closed pull request #1993: feat: add extension to support `DataFrame::checkpoint()`
URL: https://github.com/apache/datafusion-ballista/pull/1993
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "akshaychitneni (via GitHub)" <gi...@apache.org>.
akshaychitneni commented on code in PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#discussion_r3941386803
##########
ballista/scheduler/src/checkpoint.rs:
##########
@@ -0,0 +1,595 @@
+// 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.
+
+//! Resolution of lazy checkpoints into separate jobs.
+//!
+//! A `BallistaCheckpointNode` in a submitted logical plan means "materialise
+//! everything below me before evaluating anything above me". This module turns
+//! that into two jobs at planning time: a `CopyTo` job that writes the subtree
+//! to the checkpoint location, and the caller's job, rewritten to scan it.
+//!
+//! Nothing here reaches physical planning: by the time the plan is handed to
+//! `create_physical_plan` every checkpoint node has been replaced by a scan.
+
+use async_trait::async_trait;
+use ballista_core::error::{BallistaError, Result};
+use ballista_core::extension::BallistaCheckpointNode;
+use datafusion::common::DFSchema;
+use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion};
+use datafusion::datasource::file_format::DefaultFileType;
+use datafusion::datasource::file_format::parquet::ParquetFormatFactory;
+use datafusion::datasource::listing::ListingTableUrl;
+use datafusion::logical_expr::{Expr, LogicalPlanBuilder, UserDefinedLogicalNodeCore};
+use datafusion::logical_expr::{LogicalPlan, dml::CopyTo};
+use datafusion::prelude::{ParquetReadOptions, SessionContext};
+use futures::StreamExt;
+use futures::future::BoxFuture;
+use std::sync::Arc;
+
+/// Runs a plan to completion as its own job.
+///
+/// Implemented by `SchedulerServer`; the indirection keeps this module free of
+/// the event loop and lets tests substitute a local runner.
+#[async_trait]
+pub(crate) trait CheckpointMaterializer: Send + Sync {
Review Comment:
I think .cache() (#1829) needs this same materialize-then-substitute approach. I think a generic name here like `Materializer`, and materialize() returning the output's partition locations (not just Result<()>) so cache can register/pin where the output landed, would let it reuse this instead of re-deriving. I can propose the generalization as a follow-up so it doesn't grow this PR's scope. cc @sandugood @milenkovicm
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-6185148684
> > I will try and fix all of the comments asap. Then we can proceed with `.cache()`
>
> @sandugood checking to see if you are still looking into it.
Not working on it at the moment. If you can - you can proceee with it. Thanks
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-5217219022
I believe you could get away without physical node if you split into two jobs at planning
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-5231267056
> I believe you could get away without physical node if you split into two jobs at planning
Tried to introduce this concept and it worked out nicely
Added unit tests and an example as well. thanks @milenkovicm
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "akshaychitneni (via GitHub)" <gi...@apache.org>.
akshaychitneni commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-6180219502
> I will try and fix all of the comments asap. Then we can proceed with `.cache()`
@sandugood checking to see if you are still looking into it.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#discussion_r3735074462
##########
ballista/client/src/extension.rs:
##########
@@ -234,3 +237,46 @@ impl Extension {
Ok(scheduler_url)
}
}
+
+/// Providing [DataFrameExt] for an extended functionality on DataFusion DataFrame.
+///
+#[async_trait::async_trait]
+pub trait DataFrameExt {
+ /// Checkpointing DataFrame - storing intermediate result to disk and breaking lineage in the plan
+ async fn checkpoint(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 ctx = SessionContext::new_with_state(state);
+
+ let base_dir =
+ ctx.state()
+ .config()
+ .ballista_checkpoint_dir()
+ .ok_or_else(|| {
+ DataFusionError::Configuration(
+ "ballista.checkpoint.dir must be set to use DataFrame::checkpoint()"
+ .to_string(),
+ )
+ })?;
+
+ let path = format!(
+ "{}/{}/{}",
+ base_dir.trim_end_matches('/'),
+ ctx.state().session_id(),
+ Uuid::new_v4()
+ );
+
+ // Executes the original plan as a normal distributed job.
+ ctx.execute_logical_plan(plan)
Review Comment:
Let's keep this one, if possible
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-4945464768
Is that OK that we'll hold the `DataFrameExt` in the `client/src/extension.rs` or should we move it to a separate file?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#discussion_r3812924643
##########
ballista/scheduler/src/scheduler_server/mod.rs:
##########
@@ -331,6 +334,19 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
plan: &SubmitPlan,
subscriber: Option<JobStatusSubscriber>,
) -> Result<JobId> {
+ // Lazy checkpoints become their own jobs; what reaches the planner is
+ // always checkpoint free. Physical plans cannot carry checkpoint nodes,
+ // since those are resolved before physical planning.
+ let resolved;
+ let plan = match plan {
+ SubmitPlan::Logical(logical) if contains_checkpoint(logical) => {
+ resolved =
+ SubmitPlan::Logical(resolve_checkpoints(&ctx, logical, self).await?);
Review Comment:
if im not mistaken, this will block job submission (of reading part) until checkpoint creation finishes ? can we create two jobs and link them using job notification `subscriber`
##########
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<()>
Review Comment:
`should_insert_checkpoint_lazy` as name?
##########
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()?;
Review Comment:
should we have test to check execution of `checkpoint_lazy` or ive missed it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-5233587647
thanks @sandugood will have a look
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-5231365878
Also, I think it is good to have an example with remote-checkpoint, so the user could see the output of it in the ballista-tui
<img width="2110" height="972" alt="image" src="https://github.com/user-attachments/assets/21025e3b-539d-4dac-9c29-4e775978aa4a" />
During this planning I have:
1. Read .`parquet` and applied a simple filter. Then called `.checkpoint_lazy()`, which split the plan into two jobs
2. Then reused the same D`F`, which didn't read from the original `.parquet`, but rather from the checkpointed part
Hence having 3 separated jobs in the TUI
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#issuecomment-5217099488
I think, same as `.cache()` this needs a physical operator (something like `CheckpointExec`), an extension planner, a physical codec entry
Right now it seems OK with logical planning. wdyt @milenkovicm
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: add extension to support `DataFrame::checkpoint()` [datafusion-ballista]
Posted by "sandugood (via GitHub)" <gi...@apache.org>.
sandugood commented on code in PR #1993:
URL: https://github.com/apache/datafusion-ballista/pull/1993#discussion_r3735021143
##########
ballista/client/src/extension.rs:
##########
@@ -234,3 +237,46 @@ impl Extension {
Ok(scheduler_url)
}
}
+
+/// Providing [DataFrameExt] for an extended functionality on DataFusion DataFrame.
+///
+#[async_trait::async_trait]
+pub trait DataFrameExt {
+ /// Checkpointing DataFrame - storing intermediate result to disk and breaking lineage in the plan
+ async fn checkpoint(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 ctx = SessionContext::new_with_state(state);
+
+ let base_dir =
+ ctx.state()
+ .config()
+ .ballista_checkpoint_dir()
+ .ok_or_else(|| {
+ DataFusionError::Configuration(
+ "ballista.checkpoint.dir must be set to use DataFrame::checkpoint()"
+ .to_string(),
+ )
+ })?;
+
+ let path = format!(
+ "{}/{}/{}",
+ base_dir.trim_end_matches('/'),
+ ctx.state().session_id(),
+ Uuid::new_v4()
+ );
+
+ // Executes the original plan as a normal distributed job.
+ ctx.execute_logical_plan(plan)
Review Comment:
Seems like a reasonable thing to do
Should we open a separate PR for this though?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org