in work

Author: waralexrom

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

@@ -113,6 +113,7 @@ pub enum PlanNode {
     MultiStageMeasureCalculation(Rc<MultiStageMeasureCalculation>),
     MultiStageTimeSeries(Rc<MultiStageTimeSeries>),
     MultiStageRollingWindow(Rc<MultiStageRollingWindow>),
+    LogicalMultiStageMember(Rc<LogicalMultiStageMember>),
 }
 
 impl PlanNode {
@@ -133,6 +134,7 @@ impl PlanNode {
             PlanNode::MultiStageMeasureCalculation(_) => MultiStageMeasureCalculation::node_name(),
             PlanNode::MultiStageTimeSeries(_) => MultiStageTimeSeries::node_name(),
             PlanNode::MultiStageRollingWindow(_) => MultiStageRollingWindow::node_name(),
+            PlanNode::LogicalMultiStageMember(_) => LogicalMultiStageMember::node_name(),
         }
     }
 

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

@@ -15,6 +15,7 @@ pub mod pretty_print;
 mod query;
 mod resolve_multiplied_measures;
 mod schema;
+pub mod visitor;
 
 pub use aggregate_multiplied_subquery::*;
 pub use cube::*;

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

@@ -1,11 +1,35 @@
 use crate::logical_plan::*;
+use cubenativeutils::CubeError;
+use std::rc::Rc;
 
 pub enum MultiStageMemberLogicalType {
-    LeafMeasure(MultiStageLeafMeasure),
-    MeasureCalculation(MultiStageMeasureCalculation),
-    GetDateRange(MultiStageGetDateRange),
-    TimeSeries(MultiStageTimeSeries),
-    RollingWindow(MultiStageRollingWindow),
+    LeafMeasure(Rc<MultiStageLeafMeasure>),
+    MeasureCalculation(Rc<MultiStageMeasureCalculation>),
+    GetDateRange(Rc<MultiStageGetDateRange>),
+    TimeSeries(Rc<MultiStageTimeSeries>),
+    RollingWindow(Rc<MultiStageRollingWindow>),
+}
+
+impl MultiStageMemberLogicalType {
+    fn as_plan_node(&self) -> PlanNode {
+        match self {
+            Self::LeafMeasure(item) => item.as_plan_node(),
+            Self::MeasureCalculation(item) => item.as_plan_node(),
+            Self::GetDateRange(item) => item.as_plan_node(),
+            Self::TimeSeries(item) => item.as_plan_node(),
+            Self::RollingWindow(item) => item.as_plan_node(),
+        }
+    }
+    
+    fn with_plan_node(&self, plan_node: PlanNode) -> Result<Self, CubeError> {
+        Ok(match self {
+            Self::LeafMeasure(_) => Self::LeafMeasure(plan_node.into_logical_node()?),
+            Self::MeasureCalculation(_) => Self::MeasureCalculation(plan_node.into_logical_node()?),
+            Self::GetDateRange(_) => Self::GetDateRange(plan_node.into_logical_node()?),
+            Self::TimeSeries(_) => Self::TimeSeries(plan_node.into_logical_node()?),
+            Self::RollingWindow(_) => Self::RollingWindow(plan_node.into_logical_node()?),
+        })
+    }
 }
 
 impl PrettyPrint for MultiStageMemberLogicalType {
@@ -25,6 +49,39 @@ pub struct LogicalMultiStageMember {
     pub member_type: MultiStageMemberLogicalType,
 }
 
+impl LogicalNode for LogicalMultiStageMember {
+    type InputsType = SingleNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::LogicalMultiStageMember(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        SingleNodeInput::new(self.member_type.as_plan_node())
+    }
+
+    fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        let input = inputs.unpack();
+        
+        Ok(Rc::new(Self {
+            name: self.name.clone(),
+            member_type: self.member_type.with_plan_node(input)?,
+        }))
+    }
+
+    fn node_name() -> &'static str {
+        "LogicalMultiStageMember"
+    }
+
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::LogicalMultiStageMember(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}
+
 impl PrettyPrint for LogicalMultiStageMember {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println(&format!("MultiStageMember `{}`: ", self.name), state);

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/query.rs

@@ -47,24 +47,47 @@ pub struct Query {
 }
 
 impl LogicalNode for Query {
-    type InputsType = SingleNodeInput;
+    type InputsType = QueryInput;
 
     fn as_plan_node(self: &Rc<Self>) -> PlanNode {
         PlanNode::Query(self.clone())
     }
 
     fn inputs(&self) -> Self::InputsType {
-        SingleNodeInput::new(self.source.as_plan_node())
+        let source = self.source.as_plan_node();
+        let multistage_members = self
+            .multistage_members
+            .iter()
+            .map(|member| member.as_plan_node())
+            .collect();
+        
+        QueryInput {
+            source,
+            multistage_members,
+        }
     }
 
     fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
-        let input = inputs.unpack();
+        let QueryInput {
+            source,
+            multistage_members,
+        } = inputs;
+
+        check_inputs_len::<Self>(
+            "multistage_members",
+            &multistage_members,
+            self.multistage_members.len(),
+        )?;
+
         Ok(Rc::new(Self {
-            multistage_members: self.multistage_members.clone(),
+            multistage_members: multistage_members
+                .into_iter()
+                .map(|member| member.into_logical_node())
+                .collect::<Result<Vec<_>, _>>()?,
             schema: self.schema.clone(),
             filter: self.filter.clone(),
             modifers: self.modifers.clone(),
-            source: self.source.with_plan_node(self.source.as_plan_node())?,
+            source: self.source.with_plan_node(source)?,
         }))
     }
 
@@ -80,6 +103,20 @@ impl LogicalNode for Query {
     }
 }
 
+pub struct QueryInput {
+    pub source: PlanNode,
+    pub multistage_members: Vec<PlanNode>,
+}
+
+impl NodeInputs for QueryInput {
+    fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
+        Box::new(
+            std::iter::once(&self.source)
+                .chain(self.multistage_members.iter())
+        )
+    }
+}
+
 impl PrettyPrint for Query {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println("Query: ", state);

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

@@ -0,0 +1,2 @@
+mod visitor;
+pub use visitor::*;

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/visitor/visitor.rs

@@ -1,6 +1,6 @@
 use super::super::context::PushDownBuilderContext;
 use super::super::{LogicalNodeProcessor, ProcessableNode};
-use crate::logical_plan::MultiStageMemberLogicalType;
+use crate::logical_plan::{LogicalNode, MultiStageMemberLogicalType};
 use crate::physical_plan_builder::PhysicalPlanBuilder;
 use crate::plan::QueryPlan;
 use cubenativeutils::CubeError;
@@ -24,19 +24,19 @@ impl<'a> LogicalNodeProcessor<'a, MultiStageMemberLogicalType>
     ) -> Result<Self::PhysycalNode, CubeError> {
         match multi_stage_member {
             MultiStageMemberLogicalType::LeafMeasure(measure) => {
-                self.builder.process_node(measure, context)
+                self.builder.process_node(measure.as_ref(), context)
             }
             MultiStageMemberLogicalType::MeasureCalculation(calculation) => {
-                self.builder.process_node(calculation, context)
+                self.builder.process_node(calculation.as_ref(), context)
             }
             MultiStageMemberLogicalType::GetDateRange(get_date_range) => {
-                self.builder.process_node(get_date_range, context)
+                self.builder.process_node(get_date_range.as_ref(), context)
             }
             MultiStageMemberLogicalType::TimeSeries(time_series) => {
-                self.builder.process_node(time_series, context)
+                self.builder.process_node(time_series.as_ref(), context)
             }
             MultiStageMemberLogicalType::RollingWindow(rolling_window) => {
-                self.builder.process_node(rolling_window, context)
+                self.builder.process_node(rolling_window.as_ref(), context)
             }
         }
     }

rust/cubesqlplanner/cubesqlplanner/src/physical_plan_builder/processors/multi_stage_member_type.rs

@@ -86,7 +86,7 @@ impl MultiStageMemberQueryPlanner {
         };
         let member = LogicalMultiStageMember {
             name: self.description.alias().clone(),
-            member_type: MultiStageMemberLogicalType::GetDateRange(result),
+            member_type: MultiStageMemberLogicalType::GetDateRange(Rc::new(result)),
         };
 
         Ok(Rc::new(member))
@@ -104,7 +104,7 @@ impl MultiStageMemberQueryPlanner {
         };
         Ok(Rc::new(LogicalMultiStageMember {
             name: self.description.alias().clone(),
-            member_type: MultiStageMemberLogicalType::TimeSeries(result),
+            member_type: MultiStageMemberLogicalType::TimeSeries(Rc::new(result)),
         }))
     }
 
@@ -176,7 +176,7 @@ impl MultiStageMemberQueryPlanner {
         };
         Ok(Rc::new(LogicalMultiStageMember {
             name: self.description.alias().clone(),
-            member_type: MultiStageMemberLogicalType::RollingWindow(result),
+            member_type: MultiStageMemberLogicalType::RollingWindow(Rc::new(result)),
         }))
     }
 
@@ -248,7 +248,7 @@ impl MultiStageMemberQueryPlanner {
 
         let result = LogicalMultiStageMember {
             name: self.description.alias().clone(),
-            member_type: MultiStageMemberLogicalType::MeasureCalculation(result),
+            member_type: MultiStageMemberLogicalType::MeasureCalculation(Rc::new(result)),
         };
         Ok(Rc::new(result))
     }
@@ -292,7 +292,7 @@ impl MultiStageMemberQueryPlanner {
         };
         let result = LogicalMultiStageMember {
             name: self.description.alias().clone(),
-            member_type: MultiStageMemberLogicalType::LeafMeasure(leaf_measure_plan),
+            member_type: MultiStageMemberLogicalType::LeafMeasure(Rc::new(leaf_measure_plan)),
         };
         Ok(Rc::new(result))
     }