in work

Author: waralexrom

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

@@ -108,6 +108,11 @@ pub enum PlanNode {
     MeasureSubquery(Rc<MeasureSubquery>),
     DimensionSubQuery(Rc<DimensionSubQuery>),
     KeysSubQuery(Rc<KeysSubQuery>),
+    MultiStageGetDateRange(Rc<MultiStageGetDateRange>),
+    MultiStageLeafMeasure(Rc<MultiStageLeafMeasure>),
+    MultiStageMeasureCalculation(Rc<MultiStageMeasureCalculation>),
+    MultiStageTimeSeries(Rc<MultiStageTimeSeries>),
+    MultiStageRollingWindow(Rc<MultiStageRollingWindow>),
 }
 
 impl PlanNode {
@@ -123,6 +128,11 @@ impl PlanNode {
             PlanNode::MeasureSubquery(_) => MeasureSubquery::node_name(),
             PlanNode::DimensionSubQuery(_) => DimensionSubQuery::node_name(),
             PlanNode::KeysSubQuery(_) => KeysSubQuery::node_name(),
+            PlanNode::MultiStageGetDateRange(_) => MultiStageGetDateRange::node_name(),
+            PlanNode::MultiStageLeafMeasure(_) => MultiStageLeafMeasure::node_name(),
+            PlanNode::MultiStageMeasureCalculation(_) => MultiStageMeasureCalculation::node_name(),
+            PlanNode::MultiStageTimeSeries(_) => MultiStageTimeSeries::node_name(),
+            PlanNode::MultiStageRollingWindow(_) => MultiStageRollingWindow::node_name(),
         }
     }
 

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/multistage/calculation.rs

@@ -1,10 +1,11 @@
 use crate::logical_plan::*;
 use crate::planner::query_properties::OrderByItem;
 use crate::planner::sql_evaluator::MemberSymbol;
+use cubenativeutils::CubeError;
 use itertools::Itertools;
 use std::rc::Rc;
 
-#[derive(PartialEq)]
+#[derive(PartialEq, Clone)]
 pub enum MultiStageCalculationType {
     Rank,
     Aggregate,
@@ -21,7 +22,7 @@ impl ToString for MultiStageCalculationType {
     }
 }
 
-#[derive(PartialEq)]
+#[derive(PartialEq, Clone)]
 pub enum MultiStageCalculationWindowFunction {
     Rank,
     Window,
@@ -96,3 +97,41 @@ impl PrettyPrint for MultiStageMeasureCalculation {
         self.source.pretty_print(result, &details_state);
     }
 }
+
+impl LogicalNode for MultiStageMeasureCalculation {
+    type InputsType = SingleNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::MultiStageMeasureCalculation(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        SingleNodeInput::new(self.source.as_plan_node())
+    }
+
+    fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        let source = inputs.unpack();
+
+        Ok(Rc::new(Self {
+            schema: self.schema.clone(),
+            is_ungrouped: self.is_ungrouped,
+            calculation_type: self.calculation_type.clone(),
+            partition_by: self.partition_by.clone(),
+            window_function_to_use: self.window_function_to_use.clone(),
+            order_by: self.order_by.clone(),
+            source: source.into_logical_node()?,
+        }))
+    }
+
+    fn node_name() -> &'static str {
+        "MultiStageMeasureCalculation"
+    }
+
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::MultiStageMeasureCalculation(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/multistage/get_date_range.rs

@@ -1,12 +1,46 @@
 use crate::logical_plan::*;
 use crate::planner::sql_evaluator::MemberSymbol;
+use cubenativeutils::CubeError;
 use std::rc::Rc;
+
 pub struct MultiStageGetDateRange {
     pub time_dimension: Rc<MemberSymbol>,
-    pub dimension_subqueries: Vec<Rc<DimensionSubQuery>>,
     pub source: Rc<LogicalJoin>,
 }
 
+impl LogicalNode for MultiStageGetDateRange {
+    type InputsType = SingleNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::MultiStageGetDateRange(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        SingleNodeInput::new(self.source.as_plan_node())
+    }
+
+    fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        let source = inputs.unpack();
+
+        Ok(Rc::new(Self {
+            time_dimension: self.time_dimension.clone(),
+            source: source.into_logical_node()?,
+        }))
+    }
+
+    fn node_name() -> &'static str {
+        "MultiStageGetDateRange"
+    }
+
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::MultiStageGetDateRange(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}
+
 impl PrettyPrint for MultiStageGetDateRange {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println("Get Date Range", state);
@@ -16,12 +50,6 @@ impl PrettyPrint for MultiStageGetDateRange {
             &format!("time_dimension: {}", self.time_dimension.full_name()),
             &details_state,
         );
-        if !self.dimension_subqueries.is_empty() {
-            result.println("dimension_subqueries:", &state);
-            for subquery in self.dimension_subqueries.iter() {
-                subquery.pretty_print(result, &details_state);
-            }
-        }
         result.println("source:", &state);
         self.source.pretty_print(result, &details_state);
     }

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/multistage/leaf_measure.rs

@@ -1,6 +1,7 @@
 use crate::logical_plan::*;
 use crate::planner::planners::multi_stage::TimeShiftState;
 use crate::planner::sql_evaluator::MemberSymbol;
+use cubenativeutils::CubeError;
 use std::rc::Rc;
 
 pub struct MultiStageLeafMeasure {
@@ -47,3 +48,39 @@ impl PrettyPrint for MultiStageLeafMeasure {
         self.query.pretty_print(result, &details_state);
     }
 }
+
+impl LogicalNode for MultiStageLeafMeasure {
+    type InputsType = SingleNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::MultiStageLeafMeasure(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        SingleNodeInput::new(self.query.as_plan_node())
+    }
+
+    fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        let query = inputs.unpack();
+        
+        Ok(Rc::new(Self {
+            measure: self.measure.clone(),
+            render_measure_as_state: self.render_measure_as_state,
+            render_measure_for_ungrouped: self.render_measure_for_ungrouped,
+            time_shifts: self.time_shifts.clone(),
+            query: query.into_logical_node()?,
+        }))
+    }
+
+    fn node_name() -> &'static str {
+        "MultiStageLeafMeasure"
+    }
+
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::MultiStageLeafMeasure(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/multistage/rolling_window.rs

@@ -2,6 +2,7 @@ use crate::logical_plan::*;
 use crate::planner::query_properties::OrderByItem;
 use crate::planner::sql_evaluator::MemberSymbol;
 use crate::planner::Granularity;
+use cubenativeutils::CubeError;
 use std::rc::Rc;
 
 pub struct MultiStageRegularRollingWindow {
@@ -110,3 +111,31 @@ impl PrettyPrint for MultiStageRollingWindow {
         );
     }
 }
+
+impl LogicalNode for MultiStageRollingWindow {
+    type InputsType = EmptyNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::MultiStageRollingWindow(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        EmptyNodeInput::new()
+    }
+
+    fn with_inputs(self: Rc<Self>, _inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        Ok(self)
+    }
+
+    fn node_name() -> &'static str {
+        "MultiStageRollingWindow"
+    }
+
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::MultiStageRollingWindow(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/multistage/time_series.rs

@@ -1,5 +1,6 @@
 use crate::logical_plan::*;
 use crate::planner::sql_evaluator::MemberSymbol;
+use cubenativeutils::CubeError;
 use std::rc::Rc;
 pub struct MultiStageTimeSeries {
     pub time_dimension: Rc<MemberSymbol>,
@@ -32,3 +33,31 @@ impl PrettyPrint for MultiStageTimeSeries {
         }
     }
 }
+
+impl LogicalNode for MultiStageTimeSeries {
+    type InputsType = EmptyNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::MultiStageTimeSeries(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        EmptyNodeInput::new()
+    }
+
+    fn with_inputs(self: Rc<Self>, _inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        Ok(self)
+    }
+
+    fn node_name() -> &'static str {
+        "MultiStageTimeSeries"
+    }
+
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::MultiStageTimeSeries(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}

rust/cubesqlplanner/cubesqlplanner/src/planner/planners/multi_stage/member_query_planner.rs

@@ -82,7 +82,6 @@ impl MultiStageMemberQueryPlanner {
 
         let result = MultiStageGetDateRange {
             time_dimension: time_dimension.clone(),
-            dimension_subqueries: source.dimension_subqueries.clone(),
             source,
         };
         let member = LogicalMultiStageMember {