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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]