in work

Author: waralexrom

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/full_key_aggregate.rs

@@ -21,7 +21,7 @@ impl PrettyPrint for MultiStageSubqueryRef {
 #[derive(Clone)]
 pub enum ResolvedMultipliedMeasures {
     ResolveMultipliedMeasures(Rc<ResolveMultipliedMeasures>),
-    PreAggregation(Rc<SimpleQuery>),
+    PreAggregation(Rc<Query>),
 }
 
 impl ResolvedMultipliedMeasures {
@@ -56,6 +56,59 @@ pub struct FullKeyAggregate {
     pub multi_stage_subquery_refs: Vec<Rc<MultiStageSubqueryRef>>,
 }
 
+impl LogicalNode for FullKeyAggregate {
+    type InputsType = OptionNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::LogicalJoin(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        let plan_node = match &self.source {
+            QuerySource::LogicalJoin(join) => SingleNodeInput::new(join.as_plan_node()),
+            QuerySource::FullKeyAggregate(full_key) => {
+                SingleNodeInput::new(full_key.as_plan_node())
+            }
+            QuerySource::PreAggregation(pre_aggregation) => {
+                SingleNodeInput::new(pre_aggregation.as_plan_node())
+            }
+        };
+        plan_node
+    }
+
+    fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        let source = match inputs.item() {
+            PlanNode::LogicalJoin(item) => QuerySource::LogicalJoin(item.clone()),
+            PlanNode::FullKeyAggregate(item) => QuerySource::FullKeyAggregate(item.clone()),
+            PlanNode::PreAggregation(item) => QuerySource::PreAggregation(item.clone()),
+            _ => {
+                return Err(CubeError::internal(format!(
+                    "{} is incorrect input for Query node",
+                    inputs.item().node_name()
+                )))
+            }
+        };
+        Ok(Rc::new(Query {
+            multistage_members: self.multistage_members.clone(),
+            schema: self.schema.clone(),
+            filter: self.filter.clone(),
+            modifers: self.modifers.clone(),
+            source,
+        }))
+    }
+
+    fn node_name() -> &'static str {
+        "Query"
+    }
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::Query(query) = plan_node {
+            Ok(query)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}
+
 impl PrettyPrint for FullKeyAggregate {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println("FullKeyAggregate: ", state);

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/full_key_aggregate_query.rs

@@ -4,39 +4,27 @@ use crate::planner::sql_evaluator::SqlCall;
 use std::rc::Rc;
 
 #[derive(Clone)]
-pub struct CubeJoinItem {
+pub struct LogicalJoinItem {
     pub cube: Rc<Cube>,
     pub on_sql: Rc<SqlCall>,
 }
 
-impl PrettyPrint for CubeJoinItem {
+impl PrettyPrint for LogicalJoinItem {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println(&format!("CubeJoinItem: "), state);
         let details_state = state.new_level();
         self.cube.pretty_print(result, &details_state);
     }
 }
 
-#[derive(Clone)]
-pub enum LogicalJoinItem {
-    CubeJoinItem(CubeJoinItem),
-}
-
-impl PrettyPrint for LogicalJoinItem {
-    fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
-        match self {
-            LogicalJoinItem::CubeJoinItem(item) => item.pretty_print(result, state),
-        }
-    }
-}
-
 #[derive(Clone)]
 pub struct LogicalJoin {
     pub root: Rc<Cube>,
     pub joins: Vec<LogicalJoinItem>,
     pub dimension_subqueries: Vec<Rc<DimensionSubQuery>>,
 }
 
+
 impl PrettyPrint for LogicalJoin {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println(&format!("Join: "), state);

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/join.rs

@@ -0,0 +1,111 @@
+use super::*;
+use cubenativeutils::CubeError;
+use std::rc::Rc;
+
+pub trait NodeInputs {
+    fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_>;
+}
+
+pub struct SingleNodeInput {
+    item: PlanNode,
+}
+
+impl SingleNodeInput {
+    pub fn new(item: PlanNode) -> Self {
+        Self { item }
+    }
+
+    pub fn item(&self) -> &PlanNode {
+        &self.item
+    }
+}
+
+impl NodeInputs for SingleNodeInput {
+    fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
+        Box::new(std::iter::once(&self.item))
+    }
+}
+
+pub struct OptionNodeInput {
+    item: Option<PlanNode>,
+}
+
+impl OptionNodeInput {
+    pub fn new<T: LogicalNode>(item: Option<PlanNode>) -> Self {
+        Self { item }
+    }
+
+    pub fn item(&self) -> &Option<PlanNode> {
+        &self.item
+    }
+}
+
+impl NodeInputs for OptionNodeInput {
+    fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
+        if let Some(item) = &self.item {
+            Box::new(std::iter::once(item))
+        } else {
+            Box::new(std::iter::empty())
+        }
+    }
+}
+
+pub struct VecNodeInput {
+    items: Vec<PlanNode>,
+}
+
+impl VecNodeInput {
+    pub fn new(items: Vec<PlanNode>) -> Self {
+        Self { items }
+    }
+
+    pub fn items(&self) -> &Vec<PlanNode> {
+        &self.items
+    }
+}
+
+impl NodeInputs for VecNodeInput {
+    fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
+        Box::new(self.items.iter())
+    }
+}
+
+pub trait LogicalNode {
+    type InputsType: NodeInputs;
+
+    fn inputs(&self) -> Self::InputsType;
+
+    fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError>;
+
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError>;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode;
+
+    fn node_name() -> &'static str;
+}
+
+pub enum PlanNode {
+    Query(Rc<Query>),
+    LogicalJoin(Rc<LogicalJoin>),
+    FullKeyAggregate(Rc<FullKeyAggregate>),
+    PreAggregation(Rc<PreAggregation>),
+}
+
+impl PlanNode {
+    pub fn node_name(&self) -> &'static str {
+        match self {
+            PlanNode::Query(_) => Query::node_name(),
+            PlanNode::LogicalJoin(_) => LogicalJoin::node_name(),
+            PlanNode::FullKeyAggregate(_) => FullKeyAggregate::node_name(),
+            PlanNode::PreAggregation(_) => PreAggregation::node_name(),
+        }
+    }
+}
+
+pub(super) fn cast_error<T: LogicalNode>(plan_node: &PlanNode) -> CubeError {
+    CubeError::internal(format!(
+        "Can't cast {} PlanNode into {}",
+        plan_node.node_name(),
+        T::node_name(),
+    ))
+}

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/logical_node.rs

@@ -3,9 +3,9 @@ mod cube;
 mod dimension_subquery;
 mod filter;
 mod full_key_aggregate;
-mod full_key_aggregate_query;
 mod join;
 mod keys_subquery;
+mod logical_node;
 mod logical_query_modifers;
 mod measure_subquery;
 mod multistage;
@@ -16,16 +16,15 @@ mod query;
 mod regular_measures_query;
 mod resolve_multiplied_measures;
 mod schema;
-mod simple_query;
 
 pub use aggregate_multiplied_subquery::*;
 pub use cube::*;
 pub use dimension_subquery::*;
 pub use filter::*;
 pub use full_key_aggregate::*;
-pub use full_key_aggregate_query::*;
 pub use join::*;
 pub use keys_subquery::*;
+pub use logical_node::*;
 pub use logical_query_modifers::*;
 pub use measure_subquery::*;
 pub use multistage::*;
@@ -36,4 +35,3 @@ pub use query::*;
 pub use regular_measures_query::*;
 pub use resolve_multiplied_measures::*;
 pub use schema::*;
-pub use simple_query::*;

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/mod.rs

@@ -1,4 +1,5 @@
 use super::*;
+use cubenativeutils::CubeError;
 use std::rc::Rc;
 
 #[derive(Clone)]
@@ -27,6 +28,59 @@ pub struct Query {
     pub source: QuerySource,
 }
 
+impl LogicalNode for Query {
+    type InputsType = SingleNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::Query(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        let plan_node = match &self.source {
+            QuerySource::LogicalJoin(join) => SingleNodeInput::new(join.as_plan_node()),
+            QuerySource::FullKeyAggregate(full_key) => {
+                SingleNodeInput::new(full_key.as_plan_node())
+            }
+            QuerySource::PreAggregation(pre_aggregation) => {
+                SingleNodeInput::new(pre_aggregation.as_plan_node())
+            }
+        };
+        plan_node
+    }
+
+    fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        let source = match inputs.item() {
+            PlanNode::LogicalJoin(item) => QuerySource::LogicalJoin(item.clone()),
+            PlanNode::FullKeyAggregate(item) => QuerySource::FullKeyAggregate(item.clone()),
+            PlanNode::PreAggregation(item) => QuerySource::PreAggregation(item.clone()),
+            _ => {
+                return Err(CubeError::internal(format!(
+                    "{} is incorrect input for Query node",
+                    inputs.item().node_name()
+                )))
+            }
+        };
+        Ok(Rc::new(Query {
+            multistage_members: self.multistage_members.clone(),
+            schema: self.schema.clone(),
+            filter: self.filter.clone(),
+            modifers: self.modifers.clone(),
+            source,
+        }))
+    }
+
+    fn node_name() -> &'static str {
+        "Query"
+    }
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::Query(query) = plan_node {
+            Ok(query)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}
+
 impl PrettyPrint for Query {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println("Query: ", state);