in work

Author: waralexrom

packages/cubejs-backend-native/Cargo.lock

@@ -20,6 +20,7 @@ chrono = "0.4.15"
 chrono-tz = "0.8.2"
 lazy_static = "1.4.0"
 regex = "1.3.9"
+typed-builder = "0.21.2"
 
 [dependencies.neon]
 version = "=1"

rust/cubesqlplanner/Cargo.lock

@@ -3,30 +3,12 @@ use super::*;
 use cubenativeutils::CubeError;
 use std::rc::Rc;
 
-pub enum AggregateMultipliedSubquerySouce {
-    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()?),
-        })
-    }
-}
+logical_source_enum!(AggregateMultipliedSubquerySource, [Cube, MeasureSubquery]);
 
 pub struct AggregateMultipliedSubquery {
     pub schema: Rc<LogicalSchema>,
     pub keys_subquery: Rc<KeysSubQuery>,
-    pub source: AggregateMultipliedSubquerySouce,
+    pub source: AggregateMultipliedSubquerySource,
     pub dimension_subqueries: Vec<Rc<DimensionSubQuery>>,
 }
 
@@ -127,16 +109,7 @@ impl PrettyPrint for AggregateMultipliedSubquery {
         result.println("keys_subquery:", &state);
         self.keys_subquery.pretty_print(result, &details_state);
         result.println("source:", &state);
-        match &self.source {
-            AggregateMultipliedSubquerySouce::Cube(cube) => {
-                result.println("Cube:", &details_state);
-                cube.pretty_print(result, &details_state.new_level());
-            }
-            AggregateMultipliedSubquerySouce::MeasureSubquery(measure_subquery) => {
-                result.println(&format!("MeasureSubquery: "), &details_state);
-                measure_subquery.pretty_print(result, &details_state);
-            }
-        }
+        self.source.pretty_print(result, &details_state);
         if !self.dimension_subqueries.is_empty() {
             result.println("dimension_subqueries:", &state);
             let details_state = state.new_level();

rust/cubesqlplanner/cubesqlplanner/Cargo.toml

@@ -6,8 +6,8 @@ use std::rc::Rc;
 
 #[derive(Clone)]
 pub struct CalcGroupDescription {
-    symbol: Rc<MemberSymbol>,
-    values: Vec<String>,
+    pub symbol: Rc<MemberSymbol>,
+    pub values: Vec<String>,
 }
 
 impl PrettyPrint for CalcGroupDescription {
@@ -20,10 +20,15 @@ impl PrettyPrint for CalcGroupDescription {
     }
 }
 
+logical_source_enum!(
+    CalcGroupsCrossJoinSource,
+    [LogicalJoin, FullKeyAggregate, PreAggregation]
+);
+
 #[derive(Clone)]
 pub struct CalcGroupsCrossJoin {
-    source: QuerySource,
-    calc_groups: Vec<Rc<CalcGroupDescription>>,
+    pub source: CalcGroupsCrossJoinSource,
+    pub calc_groups: Vec<Rc<CalcGroupDescription>>,
 }
 
 impl LogicalNode for CalcGroupsCrossJoin {
@@ -46,7 +51,7 @@ impl LogicalNode for CalcGroupsCrossJoin {
     }
 
     fn node_name(&self) -> &'static str {
-        "MultiStageGetDateRange"
+        "CalcGroupsCrossJoin"
     }
 
     fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {

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

@@ -31,7 +31,9 @@ impl ResolvedMultipliedMeasures {
             ResolvedMultipliedMeasures::ResolveMultipliedMeasures(resolve_multiplied_measures) => {
                 resolve_multiplied_measures.schema.clone()
             }
-            ResolvedMultipliedMeasures::PreAggregation(simple_query) => simple_query.schema.clone(),
+            ResolvedMultipliedMeasures::PreAggregation(simple_query) => {
+                simple_query.schema().clone()
+            }
         }
     }
 }

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/calc_groups_cross_join.rs

@@ -2,14 +2,33 @@ use super::*;
 use crate::planner::sql_evaluator::MemberSymbol;
 use cubenativeutils::CubeError;
 use std::rc::Rc;
+use typed_builder::TypedBuilder;
 
-#[derive(Clone)]
+#[derive(Clone, TypedBuilder)]
 pub struct KeysSubQuery {
-    pub pk_cube: Rc<Cube>,
-    pub schema: Rc<LogicalSchema>,
-    pub primary_keys_dimensions: Vec<Rc<MemberSymbol>>,
-    pub filter: Rc<LogicalFilter>,
-    pub source: Rc<LogicalJoin>,
+    pk_cube: Rc<Cube>,
+    schema: Rc<LogicalSchema>,
+    primary_keys_dimensions: Vec<Rc<MemberSymbol>>,
+    filter: Rc<LogicalFilter>,
+    source: Rc<LogicalJoin>,
+}
+
+impl KeysSubQuery {
+    pub fn pk_cube(&self) -> &Rc<Cube> {
+        &self.pk_cube
+    }
+    pub fn schema(&self) -> &Rc<LogicalSchema> {
+        &self.schema
+    }
+    pub fn primary_keys_dimensions(&self) -> &Vec<Rc<MemberSymbol>> {
+        &self.primary_keys_dimensions
+    }
+    pub fn filter(&self) -> &Rc<LogicalFilter> {
+        &self.filter
+    }
+    pub fn source(&self) -> &Rc<LogicalJoin> {
+        &self.source
+    }
 }
 
 impl LogicalNode for KeysSubQuery {

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

@@ -2,7 +2,7 @@ use super::*;
 use cubenativeutils::CubeError;
 use std::rc::Rc;
 
-pub trait LogicalSource: Sized {
+pub trait LogicalSource: Sized + PrettyPrint {
     fn as_plan_node(&self) -> PlanNode;
     fn with_plan_node(&self, plan_node: PlanNode) -> Result<Self, CubeError>;
 }

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/keys_subquery.rs

@@ -1,3 +1,5 @@
+#[macro_use]
+mod logical_source;
 mod aggregate_multiplied_subquery;
 mod calc_groups_cross_join;
 mod cube;
@@ -8,7 +10,6 @@ mod join;
 mod keys_subquery;
 mod logical_node;
 mod logical_query_modifers;
-mod logical_source;
 mod measure_subquery;
 mod multistage;
 pub mod optimizers;

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/logical_source.rs

@@ -49,7 +49,7 @@ impl PreAggregationOptimizer {
         query: Rc<Query>,
         pre_aggregation: &Rc<CompiledPreAggregation>,
     ) -> Result<Option<Rc<Query>>, CubeError> {
-        if query.multistage_members.is_empty() {
+        if query.multistage_members().is_empty() {
             self.try_rewrite_simple_query(&query, pre_aggregation)
         } else if !self.allow_multi_stage {
             Ok(None)
@@ -63,10 +63,9 @@ impl PreAggregationOptimizer {
         query: &Rc<Query>,
         pre_aggregation: &Rc<CompiledPreAggregation>,
     ) -> Result<Option<Rc<Query>>, CubeError> {
-        if self.is_schema_and_filters_match(&query.schema, &query.filter, pre_aggregation)? {
+        if self.is_schema_and_filters_match(&query.schema(), &query.filter(), pre_aggregation)? {
             let mut new_query = query.as_ref().clone();
-            new_query.source =
-                QuerySource::PreAggregation(self.make_pre_aggregation_source(pre_aggregation)?);
+            new_query.set_source(self.make_pre_aggregation_source(pre_aggregation)?.into());
             Ok(Some(Rc::new(new_query)))
         } else {
             Ok(None)
@@ -82,7 +81,7 @@ impl PreAggregationOptimizer {
         let mut has_unrewritten_leaf = false;
 
         let mut rewritten_multistages = Vec::new();
-        for multi_stage in &query.multistage_members {
+        for multi_stage in query.multistage_members() {
             let rewritten = rewriter.rewrite_top_down_with(multi_stage.clone(), |plan_node| {
                 let res = match plan_node {
                     PlanNode::MultiStageLeafMeasure(multi_stage_leaf_measure) => {
@@ -119,7 +118,7 @@ impl PreAggregationOptimizer {
             return Ok(None);
         }
 
-        let source = if let QuerySource::FullKeyAggregate(full_key_aggregate) = &query.source {
+        let source = if let QuerySource::FullKeyAggregate(full_key_aggregate) = query.source() {
             let fk_source = if let Some(resolver_multiplied_measures) =
                 &full_key_aggregate.multiplied_measures_resolver
             {
@@ -135,18 +134,17 @@ impl PreAggregationOptimizer {
                         let pre_aggregation_source =
                             self.make_pre_aggregation_source(pre_aggregation)?;
 
-                        let pre_aggregation_query = Query {
-                            schema: resolver_multiplied_measures.schema.clone(),
-                            filter: resolver_multiplied_measures.filter.clone(),
-                            modifers: Rc::new(LogicalQueryModifiers {
+                        let pre_aggregation_query = Query::builder()
+                            .schema(resolver_multiplied_measures.schema.clone())
+                            .filter(resolver_multiplied_measures.filter.clone())
+                            .modifers(Rc::new(LogicalQueryModifiers {
                                 offset: None,
                                 limit: None,
                                 ungrouped: false,
                                 order_by: vec![],
-                            }),
-                            source: QuerySource::PreAggregation(pre_aggregation_source),
-                            multistage_members: vec![],
-                        };
+                            }))
+                            .source(pre_aggregation_source.into())
+                            .build();
                         Some(ResolvedMultipliedMeasures::PreAggregation(Rc::new(
                             pre_aggregation_query,
                         )))
@@ -161,18 +159,18 @@ impl PreAggregationOptimizer {
             };
             let mut result = full_key_aggregate.as_ref().clone();
             result.multiplied_measures_resolver = fk_source;
-            QuerySource::FullKeyAggregate(Rc::new(result))
+            Rc::new(result).into()
         } else {
-            query.source.clone()
+            query.source().clone()
         };
 
-        let result = Query {
-            multistage_members: rewritten_multistages,
-            schema: query.schema.clone(),
-            filter: query.filter.clone(),
-            modifers: query.modifers.clone(),
-            source,
-        };
+        let result = Query::builder()
+            .multistage_members(rewritten_multistages)
+            .schema(query.schema().clone())
+            .filter(query.filter().clone())
+            .modifers(query.modifers().clone())
+            .source(source)
+            .build();
 
         Ok(Some(Rc::new(result)))
     }