You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "killzoner (via GitHub)" <gi...@apache.org> on 2026/01/28 18:57:13 UTC

[PR] feat: disable DataFrame.cache() for ballista [datafusion-ballista]

killzoner opened a new pull request, #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420

   # Which issue does this PR close?
   
   <!--
   We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
   -->
   
   Closes #1395 (partially if we want to expand on future configurable cache later).
   
    # Rationale for this change
   <!--
    Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
    Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.  
   -->
   
   Disable explicitly `cache()` method with good user level error rather than relying on current `panic` 
   
   # What changes are included in this PR?
   <!--
   There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
   -->
   
   Add cache factory to properly error the unsupported feature (links to https://github.com/apache/datafusion/issues/18893)
   
   # Are there any user-facing changes?
   
   User level error instead of panic
   
   <!--
   If there are user-facing changes then we may require documentation to be updated before approving the PR.
   -->
   
   <!--
   If there are any breaking changes to public APIs, please add the `api change` label.
   -->
   


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795863239


##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {
+            node: Arc::new(BallistaCacheNode::new(
+                Uuid::new_v4().to_string(),
+                session_state.session_id().to_string(),
+                plan,
+            )),
+        }))
+    }
+}
+
+/// Ballista logical Extension for caching.
+#[derive(PartialEq, Eq, PartialOrd, Hash, Debug)]
+pub struct BallistaCacheNode {
+    cache_id: String,
+    session_id: String,
+    input: LogicalPlan,
+    exprs: Vec<Expr>,
+}
+
+impl BallistaCacheNode {
+    /// Create a new cache node from provided logical input plan and cache infos.
+    pub fn new(cache_id: String, session_id: String, input: LogicalPlan) -> Self {
+        Self {
+            cache_id,
+            session_id,
+            input,
+            exprs: vec![],
+        }
+    }
+
+    /// Returns cache id.
+    pub fn cache_id(&self) -> &str {
+        self.cache_id.as_str()
+    }
+
+    /// Returns session id.
+    pub fn session_id(&self) -> &str {
+        self.session_id.as_str()
+    }
+}
+
+impl UserDefinedLogicalNodeCore for BallistaCacheNode {
+    fn name(&self) -> &str {
+        "BallistaCacheNode"
+    }
+
+    fn inputs(&self) -> Vec<&LogicalPlan> {
+        vec![&self.input]
+    }
+
+    fn schema(&self) -> &DFSchemaRef {
+        self.input.schema()
+    }
+
+    fn expressions(&self) -> Vec<Expr> {
+        self.exprs.clone()

Review Comment:
   from what I get from the trait, it could be non-empty if you use `with_exprs_and_inputs` ? maybe it's not intended this way, I'm not very familiar with this



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2789287532


##########
ballista/client/tests/context_unsupported.rs:
##########
@@ -96,4 +96,28 @@ mod unsupported {
 
         assert_batches_eq!(expected, &result);
     }
+
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_support_on_cache_collect(
+        #[future(awt)]
+        #[case]
+        ctx: SessionContext,
+    ) -> datafusion::error::Result<()> {
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+
+        // Collect fails because extension node can't be serialized

Review Comment:
   Need to update this comment



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795897802


##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {

Review Comment:
   Would it not defeat the purpose of the initial issue https://github.com/apache/datafusion-ballista/issues/1395 of informing explicitely the user that the feature is not supported?
   
   Will add this in a dedicated commit



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3880094227

   > I will work on adapting the datafusion example [apache/datafusion#19139](https://github.com/apache/datafusion/pull/19139) with ballista custom query planner once this one is done
   
   no need to provide ballista planner at the moment, its ok to fail 


-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3846142409

   > You need to serialise logical plan extension, showcase project https://github.com/milenkovicm/ballista_extensions
   
   Thanks, I was looking at this part and I was not sure I needed to update ballista proto, this gives more context :pray: 


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2821158685


##########
ballista/core/src/serde/mod.rs:
##########
@@ -186,17 +191,63 @@ impl LogicalExtensionCodec for BallistaLogicalExtensionCodec {
         &self,
         buf: &[u8],
         inputs: &[datafusion::logical_expr::LogicalPlan],
-        ctx: &TaskContext,
+        _ctx: &TaskContext,
     ) -> Result<datafusion::logical_expr::Extension> {
-        self.default_codec.try_decode(buf, inputs, ctx)
+        let plan: protobuf::BallistaLogicalPlanNode =
+            protobuf::BallistaLogicalPlanNode::decode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "Could not deserialize BallistaLogicalPlanNode: {e}"
+                ))
+            })?;
+
+        let plan: protobuf::ballista_logical_plan_node::LogicalPlanType =
+            plan.logical_plan_type.ok_or_else(|| {
+                DataFusionError::Internal(
+                    "Could not deserialize BallistaLogicalPlanNode because it's logical_plan_type is none".to_string()
+                )
+            })?;
+
+        match plan {
+            LogicalPlanType::CacheNode(plan_cache) => Ok(Extension {
+                node: Arc::new(BallistaCacheNode::new(
+                    plan_cache.cache_id,
+                    plan_cache.session_id,
+                    inputs
+                        .first()
+                        .ok_or(DataFusionError::Plan(
+                            "expected input size of 1".to_string(),
+                        ))?
+                        .clone(),
+                )),
+            }),
+        }
     }
 
     fn try_encode(
         &self,
         node: &datafusion::logical_expr::Extension,
         buf: &mut Vec<u8>,
     ) -> Result<()> {
-        self.default_codec.try_encode(node, buf)
+        if let Some(node) = node.node.as_any().downcast_ref::<BallistaCacheNode>() {
+            let proto = protobuf::BallistaLogicalPlanNode {
+                logical_plan_type: Some(LogicalPlanType::CacheNode(
+                    protobuf::LogicalPlanCacheNode {
+                        cache_id: node.cache_id().to_owned(),
+                        session_id: node.session_id().to_owned(),
+                    },
+                )),
+            };
+
+            proto.encode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "failed to encode cache node logical plan: {e:?}"
+                ))
+            })?;
+
+            Ok(())
+        } else {
+            Ok(())

Review Comment:
   this should be `self.default_codec.try_encode(node, buf)`



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795897802


##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {

Review Comment:
   Would it not defeat the purpose of the initial issue https://github.com/apache/datafusion-ballista/issues/1395 of informing explicitely the user that the feature is not supported?
   
   On the other hand I guess the issue is there for a while so we could just provide `cache` functionality  to users already trying to use it transparently by upgrading ballista version that will support this in the future



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2821161537


##########
ballista/core/src/serde/mod.rs:
##########
@@ -186,17 +191,63 @@ impl LogicalExtensionCodec for BallistaLogicalExtensionCodec {
         &self,
         buf: &[u8],
         inputs: &[datafusion::logical_expr::LogicalPlan],
-        ctx: &TaskContext,
+        _ctx: &TaskContext,
     ) -> Result<datafusion::logical_expr::Extension> {
-        self.default_codec.try_decode(buf, inputs, ctx)
+        let plan: protobuf::BallistaLogicalPlanNode =
+            protobuf::BallistaLogicalPlanNode::decode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "Could not deserialize BallistaLogicalPlanNode: {e}"
+                ))
+            })?;

Review Comment:
   there should probably be a fallback to `self.default_codec.try_decode(buf, inputs, ctx)` here



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2789290077


##########
ballista/core/proto/ballista.proto:
##########
@@ -27,6 +27,15 @@ option java_outer_classname = "BallistaProto";
 import "datafusion.proto";
 import "datafusion_common.proto";
 
+///////////////////////////////////////////////////////////////////////////////////////////////////
+// Ballista Logical Plan
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+message BallistaLogicalPlanCacheNode {

Review Comment:
   Is there any specific place we want to put this?



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2821205020


##########
ballista/core/src/serde/mod.rs:
##########
@@ -186,17 +191,63 @@ impl LogicalExtensionCodec for BallistaLogicalExtensionCodec {
         &self,
         buf: &[u8],
         inputs: &[datafusion::logical_expr::LogicalPlan],
-        ctx: &TaskContext,
+        _ctx: &TaskContext,
     ) -> Result<datafusion::logical_expr::Extension> {
-        self.default_codec.try_decode(buf, inputs, ctx)
+        let plan: protobuf::BallistaLogicalPlanNode =
+            protobuf::BallistaLogicalPlanNode::decode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "Could not deserialize BallistaLogicalPlanNode: {e}"
+                ))
+            })?;
+
+        let plan: protobuf::ballista_logical_plan_node::LogicalPlanType =
+            plan.logical_plan_type.ok_or_else(|| {
+                DataFusionError::Internal(
+                    "Could not deserialize BallistaLogicalPlanNode because it's logical_plan_type is none".to_string()
+                )
+            })?;
+
+        match plan {
+            LogicalPlanType::CacheNode(plan_cache) => Ok(Extension {
+                node: Arc::new(BallistaCacheNode::new(
+                    plan_cache.cache_id,
+                    plan_cache.session_id,
+                    inputs
+                        .first()
+                        .ok_or(DataFusionError::Plan(
+                            "expected input size of 1".to_string(),
+                        ))?
+                        .clone(),
+                )),
+            }),
+        }
     }
 
     fn try_encode(
         &self,
         node: &datafusion::logical_expr::Extension,
         buf: &mut Vec<u8>,
     ) -> Result<()> {
-        self.default_codec.try_encode(node, buf)
+        if let Some(node) = node.node.as_any().downcast_ref::<BallistaCacheNode>() {
+            let proto = protobuf::BallistaLogicalPlanNode {
+                logical_plan_type: Some(LogicalPlanType::CacheNode(
+                    protobuf::LogicalPlanCacheNode {
+                        cache_id: node.cache_id().to_owned(),
+                        session_id: node.session_id().to_owned(),
+                    },
+                )),
+            };
+
+            proto.encode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "failed to encode cache node logical plan: {e:?}"
+                ))
+            })?;
+
+            Ok(())
+        } else {
+            Ok(())

Review Comment:
   you might be right 



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795863239


##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {
+            node: Arc::new(BallistaCacheNode::new(
+                Uuid::new_v4().to_string(),
+                session_state.session_id().to_string(),
+                plan,
+            )),
+        }))
+    }
+}
+
+/// Ballista logical Extension for caching.
+#[derive(PartialEq, Eq, PartialOrd, Hash, Debug)]
+pub struct BallistaCacheNode {
+    cache_id: String,
+    session_id: String,
+    input: LogicalPlan,
+    exprs: Vec<Expr>,
+}
+
+impl BallistaCacheNode {
+    /// Create a new cache node from provided logical input plan and cache infos.
+    pub fn new(cache_id: String, session_id: String, input: LogicalPlan) -> Self {
+        Self {
+            cache_id,
+            session_id,
+            input,
+            exprs: vec![],
+        }
+    }
+
+    /// Returns cache id.
+    pub fn cache_id(&self) -> &str {
+        self.cache_id.as_str()
+    }
+
+    /// Returns session id.
+    pub fn session_id(&self) -> &str {
+        self.session_id.as_str()
+    }
+}
+
+impl UserDefinedLogicalNodeCore for BallistaCacheNode {
+    fn name(&self) -> &str {
+        "BallistaCacheNode"
+    }
+
+    fn inputs(&self) -> Vec<&LogicalPlan> {
+        vec![&self.input]
+    }
+
+    fn schema(&self) -> &DFSchemaRef {
+        self.input.schema()
+    }
+
+    fn expressions(&self) -> Vec<Expr> {
+        self.exprs.clone()

Review Comment:
   from what I get from the trait, it could be non-empty if you use `with_exprs_and_inputs` ? maybe it's not intended this way, I'm not very familiar with this
   
   https://github.com/apache/datafusion-ballista/pull/1420/changes#diff-77305d1314ffee1d338be417750116a3501c19274c426a3380e9efd967f25e6dR751-R765



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795897802


##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {

Review Comment:
   will do



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2817522323


##########
ballista/client/tests/context_unsupported.rs:
##########


Review Comment:
   i guess we should move this to other test, as it is not expected to fail



##########
ballista/client/tests/context_unsupported.rs:
##########
@@ -143,4 +149,52 @@ mod unsupported {
 
         Ok(())
     }
+
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_support_on_cache_collect(
+        #[future(awt)]
+        #[case]
+        ctx: SessionContext,
+    ) -> datafusion::error::Result<()> {
+        // opt out case, should fail
+        ctx.sql("SET ballista.cache_noop = false")
+            .await?
+            .show()
+            .await?;
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+
+        // Collect fails because extension node is not handled for now by default query planner
+        let result = cached_df.collect().await;
+
+        assert!(result.is_err());
+        let err_msg = result.unwrap_err().to_string();
+        assert!(
+            err_msg.contains("No installed planner was able to convert the custom node to an execution plan: BallistaCacheNode"),
+            "Expected planner error, got: {err_msg}"
+        );
+
+        // opt in case, should succeed transparently without effective cache
+        ctx.sql("SET ballista.cache_noop = true")
+            .await?
+            .show()
+            .await?;
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+        // Collect does not fail because cache is ignored silently with cache_noop
+        let result = cached_df.collect().await;
+
+        assert!(result.is_ok());
+        let expected = [
+            "+----------+",
+            "| Int64(1) |",
+            "+----------+",
+            "| 1        |",
+            "+----------+",
+        ];
+        assert_batches_eq!(expected, &result?);
+

Review Comment:
   we have other test to test this 



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2817552370


##########
ballista/client/tests/context_unsupported.rs:
##########
@@ -143,4 +149,52 @@ mod unsupported {
 
         Ok(())
     }
+
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_support_on_cache_collect(
+        #[future(awt)]
+        #[case]
+        ctx: SessionContext,
+    ) -> datafusion::error::Result<()> {
+        // opt out case, should fail
+        ctx.sql("SET ballista.cache_noop = false")
+            .await?
+            .show()
+            .await?;
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+
+        // Collect fails because extension node is not handled for now by default query planner
+        let result = cached_df.collect().await;
+
+        assert!(result.is_err());
+        let err_msg = result.unwrap_err().to_string();
+        assert!(
+            err_msg.contains("No installed planner was able to convert the custom node to an execution plan: BallistaCacheNode"),
+            "Expected planner error, got: {err_msg}"
+        );
+
+        // opt in case, should succeed transparently without effective cache
+        ctx.sql("SET ballista.cache_noop = true")
+            .await?
+            .show()
+            .await?;
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+        // Collect does not fail because cache is ignored silently with cache_noop
+        let result = cached_df.collect().await;
+
+        assert!(result.is_ok());
+        let expected = [
+            "+----------+",
+            "| Int64(1) |",
+            "+----------+",
+            "| 1        |",
+            "+----------+",
+        ];
+        assert_batches_eq!(expected, &result?);
+

Review Comment:
   i wanted to group test with both `ballista.cache_noop` values, will move it to specific cases then



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2817579925


##########
ballista/client/tests/context_unsupported.rs:
##########
@@ -143,4 +149,52 @@ mod unsupported {
 
         Ok(())
     }
+
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_support_on_cache_collect(
+        #[future(awt)]
+        #[case]
+        ctx: SessionContext,
+    ) -> datafusion::error::Result<()> {
+        // opt out case, should fail
+        ctx.sql("SET ballista.cache_noop = false")
+            .await?
+            .show()
+            .await?;
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+
+        // Collect fails because extension node is not handled for now by default query planner
+        let result = cached_df.collect().await;
+
+        assert!(result.is_err());
+        let err_msg = result.unwrap_err().to_string();
+        assert!(
+            err_msg.contains("No installed planner was able to convert the custom node to an execution plan: BallistaCacheNode"),
+            "Expected planner error, got: {err_msg}"
+        );
+
+        // opt in case, should succeed transparently without effective cache
+        ctx.sql("SET ballista.cache_noop = true")

Review Comment:
   can we rename config to `ballista.cache.noop` or similar  ?



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2792171758


##########
ballista/client/tests/context_unsupported.rs:
##########
@@ -96,4 +96,28 @@ mod unsupported {
 
         assert_batches_eq!(expected, &result);
     }
+
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_support_on_cache_collect(
+        #[future(awt)]
+        #[case]
+        ctx: SessionContext,
+    ) -> datafusion::error::Result<()> {
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+
+        // Collect fails because extension node can't be serialized

Review Comment:
   done



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3813479568

   no need to implement anything on scheduler side at the moment 


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm merged PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3919387995

   > thanks @killzoner just minor comments about tests, otherwise it looks good
   
   Thank you ! I guess there is only https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795572324 remaining, but I'm not sure on what to do given the API


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3909619722

   Updated, this item is still pending though https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795572324


-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2738088820


##########
ballista/client/tests/context_unsupported.rs:
##########
@@ -96,4 +96,26 @@ mod unsupported {
 
         assert_batches_eq!(expected, &result);
     }
+
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_support_caching_data_frame_without_error(

Review Comment:
   This one is a bit duplicated with https://github.com/killzoner/datafusion-ballista/blob/8d0f8421df3391479153b018da6d9e879cd60124/ballista/client/tests/context_unsupported.rs#L67, however the previous one will still be useful when we want to test the user level caching strategy later



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795572324


##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {
+            node: Arc::new(BallistaCacheNode::new(
+                Uuid::new_v4().to_string(),
+                session_state.session_id().to_string(),
+                plan,
+            )),
+        }))
+    }
+}
+
+/// Ballista logical Extension for caching.
+#[derive(PartialEq, Eq, PartialOrd, Hash, Debug)]
+pub struct BallistaCacheNode {
+    cache_id: String,
+    session_id: String,
+    input: LogicalPlan,
+    exprs: Vec<Expr>,
+}
+
+impl BallistaCacheNode {
+    /// Create a new cache node from provided logical input plan and cache infos.
+    pub fn new(cache_id: String, session_id: String, input: LogicalPlan) -> Self {
+        Self {
+            cache_id,
+            session_id,
+            input,
+            exprs: vec![],
+        }
+    }
+
+    /// Returns cache id.
+    pub fn cache_id(&self) -> &str {
+        self.cache_id.as_str()
+    }
+
+    /// Returns session id.
+    pub fn session_id(&self) -> &str {
+        self.session_id.as_str()
+    }
+}
+
+impl UserDefinedLogicalNodeCore for BallistaCacheNode {
+    fn name(&self) -> &str {
+        "BallistaCacheNode"
+    }
+
+    fn inputs(&self) -> Vec<&LogicalPlan> {
+        vec![&self.input]
+    }
+
+    fn schema(&self) -> &DFSchemaRef {
+        self.input.schema()
+    }
+
+    fn expressions(&self) -> Vec<Expr> {
+        self.exprs.clone()

Review Comment:
   can we just return empty vector instead of having property exprs, will they ever be non empty?



##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {

Review Comment:
   i wonder if we should have a ballista config option to "shortcut" cache. At the moment cache() will fail as there is no cache functionality implemented. But if we flip the switch, instead of returning LogicalPlan extension we can return `plan`. There will be no cache but pipelines relying on cache will not fail (it will be noop), wdyt?



##########
ballista/core/proto/ballista.proto:
##########
@@ -27,6 +27,15 @@ option java_outer_classname = "BallistaProto";
 import "datafusion.proto";
 import "datafusion_common.proto";
 
+///////////////////////////////////////////////////////////////////////////////////////////////////
+// Ballista Logical Plan
+///////////////////////////////////////////////////////////////////////////////////////////////////
+

Review Comment:
   I would suggest to add additional node, `BallistaLogicalPlaneNode` which would have `oneof` where `BallistaLogicalPlanCacheNode` would be part of it, similar to `BallistaPhysicalPlanNode`
   
   maybe `CacheLogicalPlanExtension` would be better name for `BallistaLogicalPlanCacheNode` 
   



##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {

Review Comment:
   ```
   if config.ballista_cache_noop {
    plan
   } else {
     Ok(LogicalPlan::Extension ...
     ..
     }
   }
   
   ```



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2820961020


##########
ballista/client/tests/context_unsupported.rs:
##########
@@ -143,4 +149,52 @@ mod unsupported {
 
         Ok(())
     }
+
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_support_on_cache_collect(
+        #[future(awt)]
+        #[case]
+        ctx: SessionContext,
+    ) -> datafusion::error::Result<()> {
+        // opt out case, should fail
+        ctx.sql("SET ballista.cache_noop = false")
+            .await?
+            .show()
+            .await?;
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+
+        // Collect fails because extension node is not handled for now by default query planner
+        let result = cached_df.collect().await;
+
+        assert!(result.is_err());
+        let err_msg = result.unwrap_err().to_string();
+        assert!(
+            err_msg.contains("No installed planner was able to convert the custom node to an execution plan: BallistaCacheNode"),
+            "Expected planner error, got: {err_msg}"
+        );
+
+        // opt in case, should succeed transparently without effective cache
+        ctx.sql("SET ballista.cache_noop = true")

Review Comment:
   done



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2821295163


##########
ballista/core/src/serde/mod.rs:
##########
@@ -186,17 +191,63 @@ impl LogicalExtensionCodec for BallistaLogicalExtensionCodec {
         &self,
         buf: &[u8],
         inputs: &[datafusion::logical_expr::LogicalPlan],
-        ctx: &TaskContext,
+        _ctx: &TaskContext,
     ) -> Result<datafusion::logical_expr::Extension> {
-        self.default_codec.try_decode(buf, inputs, ctx)
+        let plan: protobuf::BallistaLogicalPlanNode =
+            protobuf::BallistaLogicalPlanNode::decode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "Could not deserialize BallistaLogicalPlanNode: {e}"
+                ))
+            })?;
+
+        let plan: protobuf::ballista_logical_plan_node::LogicalPlanType =
+            plan.logical_plan_type.ok_or_else(|| {
+                DataFusionError::Internal(
+                    "Could not deserialize BallistaLogicalPlanNode because it's logical_plan_type is none".to_string()
+                )
+            })?;
+
+        match plan {
+            LogicalPlanType::CacheNode(plan_cache) => Ok(Extension {
+                node: Arc::new(BallistaCacheNode::new(
+                    plan_cache.cache_id,
+                    plan_cache.session_id,
+                    inputs
+                        .first()
+                        .ok_or(DataFusionError::Plan(
+                            "expected input size of 1".to_string(),
+                        ))?
+                        .clone(),
+                )),
+            }),
+        }
     }
 
     fn try_encode(
         &self,
         node: &datafusion::logical_expr::Extension,
         buf: &mut Vec<u8>,
     ) -> Result<()> {
-        self.default_codec.try_encode(node, buf)
+        if let Some(node) = node.node.as_any().downcast_ref::<BallistaCacheNode>() {
+            let proto = protobuf::BallistaLogicalPlanNode {
+                logical_plan_type: Some(LogicalPlanType::CacheNode(
+                    protobuf::LogicalPlanCacheNode {
+                        cache_id: node.cache_id().to_owned(),
+                        session_id: node.session_id().to_owned(),
+                    },
+                )),
+            };
+
+            proto.encode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "failed to encode cache node logical plan: {e:?}"
+                ))
+            })?;
+
+            Ok(())
+        } else {
+            Ok(())

Review Comment:
   done in 38cd250c9016fb0cf4f73cd9cb6ea5e92260a9de



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2821161537


##########
ballista/core/src/serde/mod.rs:
##########
@@ -186,17 +191,63 @@ impl LogicalExtensionCodec for BallistaLogicalExtensionCodec {
         &self,
         buf: &[u8],
         inputs: &[datafusion::logical_expr::LogicalPlan],
-        ctx: &TaskContext,
+        _ctx: &TaskContext,
     ) -> Result<datafusion::logical_expr::Extension> {
-        self.default_codec.try_decode(buf, inputs, ctx)
+        let plan: protobuf::BallistaLogicalPlanNode =
+            protobuf::BallistaLogicalPlanNode::decode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "Could not deserialize BallistaLogicalPlanNode: {e}"
+                ))
+            })?;

Review Comment:
   there should probably be a fallback to `self.default_codec.try_decode(buf, inputs, ctx)` here



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2789680788


##########
ballista/core/proto/ballista.proto:
##########
@@ -27,6 +27,15 @@ option java_outer_classname = "BallistaProto";
 import "datafusion.proto";
 import "datafusion_common.proto";
 
+///////////////////////////////////////////////////////////////////////////////////////////////////
+// Ballista Logical Plan
+///////////////////////////////////////////////////////////////////////////////////////////////////
+
+message BallistaLogicalPlanCacheNode {

Review Comment:
   ballista proto is ok for now 



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3887650292

   > thansk @killzoner few comments and one idea.
   > 
   > what do you think if we add configuration option which would make cache operation a no op instead of throwing exception?
   
   Thanks for the review! All the points make sense to me, putting this to draft and addressing the points in the future days


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795852719


##########
ballista/core/proto/ballista.proto:
##########
@@ -27,6 +27,15 @@ option java_outer_classname = "BallistaProto";
 import "datafusion.proto";
 import "datafusion_common.proto";
 
+///////////////////////////////////////////////////////////////////////////////////////////////////
+// Ballista Logical Plan
+///////////////////////////////////////////////////////////////////////////////////////////////////
+

Review Comment:
   I was also wondering for the oneof, thanks will do



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3912861569

   will review it asap, thanks @killzoner 


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3919405378

   > > thanks @killzoner just minor comments about tests, otherwise it looks good
   > 
   > Thank you ! I guess there is only [#1420 (comment)](https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795572324) remaining, but I'm not sure on what to do given the trait API
   
   thats ok, we can leave it as it is


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2821211316


##########
ballista/core/src/serde/mod.rs:
##########
@@ -186,17 +191,63 @@ impl LogicalExtensionCodec for BallistaLogicalExtensionCodec {
         &self,
         buf: &[u8],
         inputs: &[datafusion::logical_expr::LogicalPlan],
-        ctx: &TaskContext,
+        _ctx: &TaskContext,
     ) -> Result<datafusion::logical_expr::Extension> {
-        self.default_codec.try_decode(buf, inputs, ctx)
+        let plan: protobuf::BallistaLogicalPlanNode =
+            protobuf::BallistaLogicalPlanNode::decode(buf).map_err(|e| {
+                DataFusionError::Internal(format!(
+                    "Could not deserialize BallistaLogicalPlanNode: {e}"
+                ))
+            })?;

Review Comment:
   yes 



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3845805527

   You need to serialise logical plan extension, showcase project https://github.com/milenkovicm/ballista_extensions


-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2795897802


##########
ballista/core/src/extension.rs:
##########
@@ -665,6 +671,100 @@ impl BallistaConfigGrpcEndpoint {
 #[derive(Clone, Copy)]
 pub struct BallistaUseTls(pub bool);
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {
+    fn create(
+        &self,
+        plan: LogicalPlan,
+        session_state: &SessionState,
+    ) -> datafusion::error::Result<LogicalPlan> {
+        Ok(LogicalPlan::Extension(Extension {

Review Comment:
   Would it not defeat the purpose of the initial issue https://github.com/apache/datafusion-ballista/issues/1395 of informing explicitely the user that the feature is not supported?
   
   On the other hand I guess the issue is there for a while so we could just provide `cache` functionality  to users already trying to use it transparently when support is added



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2738218552


##########
ballista/core/src/extension.rs:
##########
@@ -528,6 +531,27 @@ impl BallistaQueryPlannerExtension {
     }
 }
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {

Review Comment:
   can we make cache factory return a logical plan extension, something like https://github.com/apache/datafusion/blob/20d5f06a554f39407e52dd1d4848267dbc1467ce/datafusion-examples/examples/dataframe/cache_factory.rs#L103
   
   but with two additional properties, session_id and cache_id (populated from random id), we can then serialize it and send it to scheduler. Scheduler will then fail as there is no implementation. So when users want to implement cache they just need to add new query planner 



-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2738288003


##########
ballista/core/src/extension.rs:
##########
@@ -528,6 +531,27 @@ impl BallistaQueryPlannerExtension {
     }
 }
 
+#[derive(Debug)]
+struct BallistaCacheFactory;
+
+impl BallistaCacheFactory {
+    fn new() -> Self {
+        Self {}
+    }
+}
+
+impl CacheFactory for BallistaCacheFactory {

Review Comment:
   Ah i see the link clearer with https://github.com/apache/datafusion-ballista/issues/1395#issuecomment-3798590795 now.
   Didn't know about this behaviour with the scheduler, will look into this, 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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on code in PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#discussion_r2820965372


##########
ballista/client/tests/context_unsupported.rs:
##########
@@ -143,4 +149,52 @@ mod unsupported {
 
         Ok(())
     }
+
+    #[rstest]
+    #[case::standalone(standalone_context())]
+    #[case::remote(remote_context())]
+    #[tokio::test]
+    async fn should_support_on_cache_collect(
+        #[future(awt)]
+        #[case]
+        ctx: SessionContext,
+    ) -> datafusion::error::Result<()> {
+        // opt out case, should fail
+        ctx.sql("SET ballista.cache_noop = false")
+            .await?
+            .show()
+            .await?;
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+
+        // Collect fails because extension node is not handled for now by default query planner
+        let result = cached_df.collect().await;
+
+        assert!(result.is_err());
+        let err_msg = result.unwrap_err().to_string();
+        assert!(
+            err_msg.contains("No installed planner was able to convert the custom node to an execution plan: BallistaCacheNode"),
+            "Expected planner error, got: {err_msg}"
+        );
+
+        // opt in case, should succeed transparently without effective cache
+        ctx.sql("SET ballista.cache_noop = true")
+            .await?
+            .show()
+            .await?;
+        let cached_df = ctx.sql("SELECT 1").await?.cache().await?;
+        // Collect does not fail because cache is ignored silently with cache_noop
+        let result = cached_df.collect().await;
+
+        assert!(result.is_ok());
+        let expected = [
+            "+----------+",
+            "| Int64(1) |",
+            "+----------+",
+            "| 1        |",
+            "+----------+",
+        ];
+        assert_batches_eq!(expected, &result?);
+

Review Comment:
   removed some parts and moved the remainder next to `should_support_caching_data_frame` test



-- 
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 `Dataframe.cache()` factory (no planner handling) [datafusion-ballista]

Posted by "milenkovicm (via GitHub)" <gi...@apache.org>.
milenkovicm commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3920846943

   merging this, thanks @killzoner 


-- 
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: disable DataFrame.cache() for ballista [datafusion-ballista]

Posted by "killzoner (via GitHub)" <gi...@apache.org>.
killzoner commented on PR #1420:
URL: https://github.com/apache/datafusion-ballista/pull/1420#issuecomment-3879574839

   I will work on adapting the datafusion example with the custom query planner abstraction one this one is done


-- 
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