in work

Author: waralexrom

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/aggregate_multiplied_subquery.rs

@@ -1,20 +1,106 @@
 use super::pretty_print::*;
 use super::*;
+use cubenativeutils::CubeError;
+use itertools::Itertools;
 use std::rc::Rc;
 
 pub enum AggregateMultipliedSubquerySouce {
-    Cube,
+    Cube(Rc<Cube>),
     MeasureSubquery(Rc<MeasureSubquery>),
 }
 
+impl AggregateMultipliedSubquerySouce {
+    fn as_plan_node(&self) -> PlanNode {
+        match self {
+            Self::Cube(item) => item.as_plan_node(),
+            Self::MeasureSubquery(item) => item.as_plan_node(),
+        }
+    }
+    fn with_plan_node(&self, plan_node: PlanNode) -> Result<Self, CubeError> {
+        Ok(match self {
+            Self::Cube(_) => Self::Cube(plan_node.into_logical_node()?),
+            Self::MeasureSubquery(_) => Self::MeasureSubquery(plan_node.into_logical_node()?),
+        })
+    }
+}
+
 pub struct AggregateMultipliedSubquery {
     pub schema: Rc<LogicalSchema>,
     pub keys_subquery: Rc<KeysSubQuery>,
-    pub pk_cube: Rc<Cube>, //FIXME may be duplication with information in keys_subquery
-    pub source: Rc<AggregateMultipliedSubquerySouce>,
+    pub source: AggregateMultipliedSubquerySouce,
     pub dimension_subqueries: Vec<Rc<DimensionSubQuery>>,
 }
 
+impl LogicalNode for AggregateMultipliedSubquery {
+    type InputsType = AggregateMultipliedSubqueryInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::AggregateMultipliedSubquery(self.clone())
+    }
+
+    fn inputs(&self) -> Self::InputsType {
+        let keys_subquery = self.keys_subquery.as_plan_node();
+        let source = self.source.as_plan_node();
+        let dimension_subqueries = self
+            .dimension_subqueries
+            .iter()
+            .map(|itm| itm.as_plan_node())
+            .collect_vec();
+        AggregateMultipliedSubqueryInput {
+            keys_subquery,
+            source,
+            dimension_subqueries,
+        }
+    }
+
+    fn with_inputs(self: Rc<Self>, inputs: Self::InputsType) -> Result<Rc<Self>, CubeError> {
+        let AggregateMultipliedSubqueryInput {
+            keys_subquery,
+            source,
+            dimension_subqueries,
+        } = inputs;
+
+        let result = Self {
+            schema: self.schema.clone(),
+            keys_subquery: keys_subquery.into_logical_node()?,
+            source: self.source.with_plan_node(source)?,
+            dimension_subqueries: dimension_subqueries
+                .into_iter()
+                .map(|itm| itm.into_logical_node())
+                .collect::<Result<Vec<_>, _>>()?,
+        };
+
+        Ok(Rc::new(result))
+    }
+
+    fn node_name() -> &'static str {
+        "AggregateMultipliedSubquery"
+    }
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::AggregateMultipliedSubquery(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}
+
+pub struct AggregateMultipliedSubqueryInput {
+    pub keys_subquery: PlanNode,
+    pub source: PlanNode,
+    pub dimension_subqueries: Vec<PlanNode>,
+}
+
+impl NodeInputs for AggregateMultipliedSubqueryInput {
+    fn iter(&self) -> Box<dyn Iterator<Item = &PlanNode> + '_> {
+        Box::new(
+            std::iter::once(&self.keys_subquery)
+                .chain(std::iter::once(&self.source))
+                .chain(self.dimension_subqueries.iter()),
+        )
+    }
+}
+
 impl PrettyPrint for AggregateMultipliedSubquery {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println("AggregateMultipliedSubquery: ", state);
@@ -25,11 +111,10 @@ impl PrettyPrint for AggregateMultipliedSubquery {
         result.println("keys_subquery:", &state);
         self.keys_subquery.pretty_print(result, &details_state);
         result.println("source:", &state);
-        match self.source.as_ref() {
-            AggregateMultipliedSubquerySouce::Cube => {
+        match &self.source {
+            AggregateMultipliedSubquerySouce::Cube(cube) => {
                 result.println("Cube:", &details_state);
-                self.pk_cube
-                    .pretty_print(result, &details_state.new_level());
+                cube.pretty_print(result, &details_state.new_level());
             }
             AggregateMultipliedSubquerySouce::MeasureSubquery(measure_subquery) => {
                 result.println(&format!("MeasureSubquery: "), &details_state);

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/cube.rs

@@ -1,5 +1,6 @@
 use super::*;
 use crate::planner::BaseCube;
+use cubenativeutils::CubeError;
 use std::rc::Rc;
 
 #[derive(Clone)]
@@ -49,3 +50,30 @@ impl Cube {
         })
     }
 }
+
+impl LogicalNode for Cube {
+    type InputsType = EmptyNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::Cube(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 {
+        "Cube"
+    }
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::Cube(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/dimension_subquery.rs

@@ -1,6 +1,7 @@
 use super::pretty_print::*;
 use super::*;
 use crate::planner::sql_evaluator::MemberSymbol;
+use cubenativeutils::CubeError;
 use std::rc::Rc;
 
 pub struct DimensionSubQuery {
@@ -10,6 +11,39 @@ pub struct DimensionSubQuery {
     pub measure_for_subquery_dimension: Rc<MemberSymbol>,
 }
 
+impl LogicalNode for DimensionSubQuery {
+    type InputsType = SingleNodeInput;
+
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::DimensionSubQuery(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 {
+            query: query.into_logical_node()?,
+            primary_keys_dimensions: self.primary_keys_dimensions.clone(),
+            subquery_dimension: self.subquery_dimension.clone(),
+            measure_for_subquery_dimension: self.measure_for_subquery_dimension.clone(),
+        }))
+    }
+
+    fn node_name() -> &'static str {
+        "DimensionSubQuery"
+    }
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::DimensionSubQuery(query) = plan_node {
+            Ok(query)
+        } else {
+            Err(cast_error::<Self>(&plan_node))
+        }
+    }
+}
+
 impl PrettyPrint for DimensionSubQuery {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
         result.println("DimensionSubQuery: ", state);

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

@@ -1,5 +1,6 @@
 use super::*;
 use crate::planner::sql_evaluator::MemberSymbol;
+use cubenativeutils::CubeError;
 use std::rc::Rc;
 
 pub struct MultiStageSubqueryRef {
@@ -60,49 +61,59 @@ impl LogicalNode for FullKeyAggregate {
     type InputsType = OptionNodeInput;
 
     fn as_plan_node(self: &Rc<Self>) -> PlanNode {
-        PlanNode::LogicalJoin(self.clone())
+        PlanNode::FullKeyAggregate(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
+        let plan_node = self
+            .multiplied_measures_resolver
+            .as_ref()
+            .map(|resolver| match resolver {
+                ResolvedMultipliedMeasures::ResolveMultipliedMeasures(item) => item.as_plan_node(),
+                ResolvedMultipliedMeasures::PreAggregation(item) => item.as_plan_node(),
+            });
+        OptionNodeInput::new(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()
-                )))
-            }
+        let input = inputs.unpack();
+        let multiplied_measures_resolver = if self.multiplied_measures_resolver.is_none()
+            && input.is_none()
+        {
+            None
+        } else if let (Some(self_source), Some(input_source)) =
+            (&self.multiplied_measures_resolver, input)
+        {
+            Some(match self_source {
+                ResolvedMultipliedMeasures::ResolveMultipliedMeasures(_) => {
+                    ResolvedMultipliedMeasures::ResolveMultipliedMeasures(
+                        input_source.into_logical_node()?,
+                    )
+                }
+                ResolvedMultipliedMeasures::PreAggregation(_) => {
+                    ResolvedMultipliedMeasures::PreAggregation(input_source.into_logical_node()?)
+                }
+            })
+        } else {
+            return Err(CubeError::internal(format!(
+                "Wrong inputs for FullKeyAggregate node"
+            )));
         };
-        Ok(Rc::new(Query {
-            multistage_members: self.multistage_members.clone(),
+
+        Ok(Rc::new(Self {
             schema: self.schema.clone(),
-            filter: self.filter.clone(),
-            modifers: self.modifers.clone(),
-            source,
+            use_full_join_and_coalesce: self.use_full_join_and_coalesce,
+            multiplied_measures_resolver,
+            multi_stage_subquery_refs: self.multi_stage_subquery_refs.clone(),
         }))
     }
 
     fn node_name() -> &'static str {
-        "Query"
+        "FullKeyAggregate"
     }
     fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
-        if let PlanNode::Query(query) = plan_node {
-            Ok(query)
+        if let PlanNode::FullKeyAggregate(item) = plan_node {
+            Ok(item)
         } else {
             Err(cast_error::<Self>(&plan_node))
         }