in work

Author: waralexrom

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

@@ -34,58 +34,6 @@ pub enum PlanNode {
     LogicalMultiStageMember(Rc<LogicalMultiStageMember>),
 }
 
-impl PlanNode {
-    pub fn node_name(&self) -> &'static str {
-        match self {
-            PlanNode::Query(node) => node.node_name(),
-            PlanNode::LogicalJoin(node) => node.node_name(),
-            PlanNode::FullKeyAggregate(node) => node.node_name(),
-            PlanNode::PreAggregation(node) => node.node_name(),
-            PlanNode::ResolveMultipliedMeasures(node) => node.node_name(),
-            PlanNode::AggregateMultipliedSubquery(node) => node.node_name(),
-            PlanNode::Cube(node) => node.node_name(),
-            PlanNode::MeasureSubquery(node) => node.node_name(),
-            PlanNode::DimensionSubQuery(node) => node.node_name(),
-            PlanNode::KeysSubQuery(node) => node.node_name(),
-            PlanNode::MultiStageGetDateRange(node) => node.node_name(),
-            PlanNode::MultiStageLeafMeasure(node) => node.node_name(),
-            PlanNode::MultiStageMeasureCalculation(node) => node.node_name(),
-            PlanNode::MultiStageTimeSeries(node) => node.node_name(),
-            PlanNode::MultiStageRollingWindow(node) => node.node_name(),
-            PlanNode::LogicalMultiStageMember(node) => node.node_name(),
-        }
-    }
-
-    pub fn into_logical_node<T: LogicalNode>(self) -> Result<Rc<T>, CubeError> {
-        T::try_from_plan_node(self)
-    }
-}
-
-pub(super) fn cast_error(plan_node: &PlanNode, target_type: &str) -> CubeError {
-    CubeError::internal(format!(
-        "Can't cast {} PlanNode into {}",
-        plan_node.node_name(),
-        target_type,
-    ))
-}
-
-pub(super) fn check_inputs_len(
-    inputs: &Vec<PlanNode>,
-    expected: usize,
-    node_type: &str,
-) -> Result<(), CubeError> {
-    if inputs.len() == expected {
-        Ok(())
-    } else {
-        Err(CubeError::internal(format!(
-            "For node {} expected {} inputs but received {}",
-            node_type,
-            expected,
-            inputs.len()
-        )))
-    }
-}
-
 // Macro for applying a block to all enum variants
 macro_rules! match_plan_node {
     ($self:expr, $node:ident => $block:block) => {
@@ -110,31 +58,52 @@ macro_rules! match_plan_node {
     };
 }
 
-impl LogicalNode for PlanNode {
-    fn inputs(&self) -> Vec<PlanNode> {
+impl PlanNode {
+    pub fn into_logical_node<T: LogicalNode>(self) -> Result<Rc<T>, CubeError> {
+        T::try_from_plan_node(self)
+    }
+
+    pub fn inputs(&self) -> Vec<PlanNode> {
         match_plan_node!(self, node => {
             node.inputs()
         })
     }
 
-    fn node_name(&self) -> &'static str {
+    pub fn node_name(&self) -> &'static str {
         match_plan_node!(self, node => {
             node.node_name()
         })
     }
 
-    fn with_inputs(self: Rc<Self>, inputs: Vec<PlanNode>) -> Result<Rc<Self>, CubeError> {
-        let result = match_plan_node!(self.as_ref(), node => {
-            node.clone().with_inputs(inputs)?.as_plan_node()
+    pub fn with_inputs(self, inputs: Vec<PlanNode>) -> Result<Self, CubeError> {
+        let result = match_plan_node!(self, node => {
+            node.with_inputs(inputs)?.as_plan_node()
         });
-        Ok(Rc::new(result))
+        Ok(result)
     }
+}
 
-    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
-        Ok(Rc::new(plan_node))
-    }
+pub(super) fn cast_error(plan_node: &PlanNode, target_type: &str) -> CubeError {
+    CubeError::internal(format!(
+        "Can't cast {} PlanNode into {}",
+        plan_node.node_name(),
+        target_type,
+    ))
+}
 
-    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
-        (**self).clone()
+pub(super) fn check_inputs_len(
+    inputs: &Vec<PlanNode>,
+    expected: usize,
+    node_type: &str,
+) -> Result<(), CubeError> {
+    if inputs.len() == expected {
+        Ok(())
+    } else {
+        Err(CubeError::internal(format!(
+            "For node {} expected {} inputs but received {}",
+            node_type,
+            expected,
+            inputs.len()
+        )))
     }
 }

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/dimension_matcher.rs

@@ -1,5 +1,4 @@
 use super::CompiledPreAggregation;
-use super::MatchState;
 use crate::plan::filter::FilterGroupOperator;
 use crate::plan::FilterItem;
 use crate::planner::filter::BaseFilter;
@@ -12,6 +11,25 @@ use cubenativeutils::CubeError;
 use std::collections::HashMap;
 use std::rc::Rc;
 
+#[derive(Clone, Debug, PartialEq)]
+pub enum MatchState {
+    Partial,
+    Full,
+    NotMatched,
+}
+
+impl MatchState {
+    pub fn combine(&self, other: &MatchState) -> MatchState {
+        if matches!(self, MatchState::NotMatched) || matches!(other, MatchState::NotMatched) {
+            return MatchState::NotMatched;
+        }
+        if matches!(self, MatchState::Partial) || matches!(other, MatchState::Partial) {
+            return MatchState::Partial;
+        }
+        MatchState::Full
+    }
+}
+
 pub struct DimensionMatcher<'a> {
     query_tools: Rc<QueryTools>,
     pre_aggregation: &'a CompiledPreAggregation,

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

@@ -1,13 +1,13 @@
 mod compiled_pre_aggregation;
-//mod dimension_matcher;
+mod dimension_matcher;
 //mod measure_matcher;
 //mod optimizer;
 //mod original_sql_collector;
 //mod original_sql_optimizer;
 //mod pre_aggregations_compiler;
 
 pub use compiled_pre_aggregation::*;
-//use dimension_matcher::*;
+use dimension_matcher::*;
 //use measure_matcher::*;
 //pub use optimizer::*;
 //pub use original_sql_collector::*;

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs

@@ -8,25 +8,6 @@ use cubenativeutils::CubeError;
 use std::collections::{HashMap, HashSet};
 use std::rc::Rc;
 
-#[derive(Clone, Debug, PartialEq)]
-pub enum MatchState {
-    Partial,
-    Full,
-    NotMatched,
-}
-
-impl MatchState {
-    pub fn combine(&self, other: &MatchState) -> MatchState {
-        if matches!(self, MatchState::NotMatched) || matches!(other, MatchState::NotMatched) {
-            return MatchState::NotMatched;
-        }
-        if matches!(self, MatchState::Partial) || matches!(other, MatchState::Partial) {
-            return MatchState::Partial;
-        }
-        MatchState::Full
-    }
-}
-
 pub struct PreAggregationOptimizer {
     query_tools: Rc<QueryTools>,
     allow_multi_stage: bool,

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

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

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

@@ -0,0 +1,55 @@
+use crate::logical_plan::{LogicalNode, PlanNode};
+use cubenativeutils::CubeError;
+use std::rc::Rc;
+
+pub trait LogicalNodeRewriter {
+    fn process_node(&mut self, node: &PlanNode) -> Result<Option<PlanNode>, CubeError>;
+}
+
+pub struct LogicalPlanRewriter {}
+
+impl LogicalPlanRewriter {
+    pub fn new() -> Self {
+        Self {}
+    }
+
+    pub fn rewrite_top_down<T: LogicalNodeRewriter, N: LogicalNode>(
+        &self,
+        node_visitor: &mut T,
+        node: Rc<N>,
+    ) -> Result<Rc<N>, CubeError> {
+        let res = if let Some(rewrited) =
+            self.rewrite_top_down_impl(node_visitor, node.as_plan_node())?
+        {
+            rewrited.into_logical_node()?
+        } else {
+            node
+        };
+        Ok(res)
+    }
+
+    fn rewrite_top_down_impl<T: LogicalNodeRewriter>(
+        &self,
+        node_visitor: &mut T,
+        node: PlanNode,
+    ) -> Result<Option<PlanNode>, CubeError> {
+        if let Some(rewrited) = node_visitor.process_node(&node)? {
+            return Ok(Some(rewrited));
+        }
+        let mut has_changes = false;
+        let mut inputs = node.inputs();
+        for input in inputs.iter_mut() {
+            if let Some(rewrited) = self.rewrite_top_down_impl(node_visitor, input.clone())? {
+                *input = rewrited;
+                has_changes = true;
+            }
+        }
+        let res = if has_changes {
+            Some(node.with_inputs(inputs)?)
+        } else {
+            None
+        };
+
+        Ok(res)
+    }
+}

rust/cubesqlplanner/cubesqlplanner/src/plan/builder/select.rs

@@ -167,7 +167,7 @@ impl SelectBuilder {
     }
     pub fn add_projection_reference_member(
         &mut self,
-        member: &Rc<dyn BaseMember>,
+        member: &Rc<MemberSymbol>,
         reference: QualifiedColumnName,
         alias: Option<String>,
     ) {