in work

Author: waralexrom

packages/cubejs-schema-compiler/test/integration/postgres/calc-groups.test.ts

@@ -126,6 +126,6 @@ describe('SqlParser', () => {
 
     // Verify table alias extraction still works after grammar changes
     const extractedConditions = sqlParser.extractWhereConditions('t');
-    expect(extractedConditions).toEqual(`t.status = 'active' AND t.created_at > '2024-01-01'`);
+    expect(extractedConditions).toEqual('t.status = \'active\' AND t.created_at > \'2024-01-01\'');
   });
 });

packages/cubejs-schema-compiler/test/unit/sql-parser.test.ts

@@ -31,15 +31,16 @@ impl PrettyPrint for LogicalJoinItem {
 
 #[derive(Clone, TypedBuilder)]
 pub struct LogicalJoin {
-    root: Rc<Cube>,
+    #[builder(default)]
+    root: Option<Rc<Cube>>,
     #[builder(default)]
     joins: Vec<LogicalJoinItem>,
     #[builder(default)]
     dimension_subqueries: Vec<Rc<DimensionSubQuery>>,
 }
 
 impl LogicalJoin {
-    pub fn root(&self) -> &Rc<Cube> {
+    pub fn root(&self) -> &Option<Rc<Cube>> {
         &self.root
     }
 
@@ -68,6 +69,12 @@ impl LogicalNode for LogicalJoin {
             dimension_subqueries,
         } = LogicalJoinInputUnPacker::new(&self, &inputs)?;
 
+        let root = if let Some(r) = root {
+            Some(r.clone().into_logical_node()?)
+        } else {
+            None
+        };
+
         let joins = self
             .joins()
             .iter()
@@ -81,12 +88,14 @@ impl LogicalNode for LogicalJoin {
             .collect::<Result<Vec<_>, _>>()?;
 
         let result = Self::builder()
-            .root(root.clone().into_logical_node()?)
+            .root(root)
             .joins(joins)
-            .dimension_subqueries(dimension_subqueries
-                .iter()
-                .map(|itm| itm.clone().into_logical_node())
-                .collect::<Result<Vec<_>, _>>()?)
+            .dimension_subqueries(
+                dimension_subqueries
+                    .iter()
+                    .map(|itm| itm.clone().into_logical_node())
+                    .collect::<Result<Vec<_>, _>>()?,
+            )
             .build();
 
         Ok(Rc::new(result))
@@ -109,7 +118,9 @@ pub struct LogicalJoinInputPacker;
 impl LogicalJoinInputPacker {
     pub fn pack(join: &LogicalJoin) -> Vec<PlanNode> {
         let mut result = vec![];
-        result.push(join.root().as_plan_node());
+        if let Some(root) = join.root() {
+            result.push(root.as_plan_node());
+        }
         result.extend(join.joins().iter().map(|item| item.cube().as_plan_node()));
         result.extend(
             join.dimension_subqueries()
@@ -121,7 +132,7 @@ impl LogicalJoinInputPacker {
 }
 
 pub struct LogicalJoinInputUnPacker<'a> {
-    root: &'a PlanNode,
+    root: Option<&'a PlanNode>,
     joins: &'a [PlanNode],
     dimension_subqueries: &'a [PlanNode],
 }
@@ -130,8 +141,14 @@ impl<'a> LogicalJoinInputUnPacker<'a> {
     pub fn new(join: &LogicalJoin, inputs: &'a Vec<PlanNode>) -> Result<Self, CubeError> {
         check_inputs_len(&inputs, Self::inputs_len(join), join.node_name())?;
 
-        let root = &inputs[0];
-        let joins_start = 1;
+        let mut joins_start = 0;
+        let root = if join.root.is_some() {
+            joins_start = 1;
+            Some(&inputs[0])
+        } else {
+            None
+        };
+
         let joins_end = joins_start + join.joins().len();
         let joins = &inputs[joins_start..joins_end];
         let dimension_subqueries = &inputs[joins_end..];
@@ -150,22 +167,27 @@ impl<'a> LogicalJoinInputUnPacker<'a> {
 
 impl PrettyPrint for LogicalJoin {
     fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
-        result.println(&format!("Join: "), state);
-        let state = state.new_level();
-        let details_state = state.new_level();
-        result.println(&format!("root: "), &state);
-        self.root().pretty_print(result, &details_state);
-        result.println(&format!("joins: "), &state);
-        let state = state.new_level();
-        for join in self.joins().iter() {
-            join.pretty_print(result, &state);
-        }
-        if !self.dimension_subqueries().is_empty() {
-            result.println("dimension_subqueries:", &state);
+        if let Some(root) = self.root() {
+            result.println(&format!("Join: "), state);
+
+            let state = state.new_level();
             let details_state = state.new_level();
-            for subquery in self.dimension_subqueries().iter() {
-                subquery.pretty_print(result, &details_state);
+            result.println(&format!("root: "), &state);
+            root.pretty_print(result, &details_state);
+            result.println(&format!("joins: "), &state);
+            let state = state.new_level();
+            for join in self.joins().iter() {
+                join.pretty_print(result, &state);
+            }
+            if !self.dimension_subqueries().is_empty() {
+                result.println("dimension_subqueries:", &state);
+                let details_state = state.new_level();
+                for subquery in self.dimension_subqueries().iter() {
+                    subquery.pretty_print(result, &details_state);
+                }
             }
+        } else {
+            result.println(&format!("Empty source"), state);
         }
     }
 }

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/join.rs

@@ -29,6 +29,7 @@ pub enum PlanNode {
     MultiStageGetDateRange(Rc<MultiStageGetDateRange>),
     MultiStageLeafMeasure(Rc<MultiStageLeafMeasure>),
     MultiStageMeasureCalculation(Rc<MultiStageMeasureCalculation>),
+    MultiStageDimensionCalculation(Rc<MultiStageDimensionCalculation>),
     MultiStageTimeSeries(Rc<MultiStageTimeSeries>),
     MultiStageRollingWindow(Rc<MultiStageRollingWindow>),
     LogicalMultiStageMember(Rc<LogicalMultiStageMember>),
@@ -51,6 +52,7 @@ macro_rules! match_plan_node {
             PlanNode::MultiStageGetDateRange($node) => $block,
             PlanNode::MultiStageLeafMeasure($node) => $block,
             PlanNode::MultiStageMeasureCalculation($node) => $block,
+            PlanNode::MultiStageDimensionCalculation($node) => $block,
             PlanNode::MultiStageTimeSeries($node) => $block,
             PlanNode::MultiStageRollingWindow($node) => $block,
             PlanNode::LogicalMultiStageMember($node) => $block,

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

@@ -0,0 +1,118 @@
+use crate::logical_plan::*;
+use crate::planner::query_properties::OrderByItem;
+use crate::planner::sql_evaluator::collectors::has_multi_stage_members;
+use crate::planner::sql_evaluator::MemberSymbol;
+use cubenativeutils::CubeError;
+use itertools::Itertools;
+use std::collections::HashSet;
+use std::rc::Rc;
+use typed_builder::TypedBuilder;
+
+#[derive(TypedBuilder)]
+pub struct MultiStageDimensionCalculation {
+    schema: Rc<LogicalSchema>,
+    multi_stage_dimension: Rc<MemberSymbol>,
+    #[builder(default)]
+    order_by: Vec<OrderByItem>,
+    source: Rc<FullKeyAggregate>,
+}
+
+impl MultiStageDimensionCalculation {
+    pub fn schema(&self) -> &Rc<LogicalSchema> {
+        &self.schema
+    }
+
+    pub fn multi_stage_dimension(&self) -> &Rc<MemberSymbol> {
+        &self.multi_stage_dimension
+    }
+
+    pub fn order_by(&self) -> &Vec<OrderByItem> {
+        &self.order_by
+    }
+
+    pub fn source(&self) -> &Rc<FullKeyAggregate> {
+        &self.source
+    }
+
+    pub fn resolved_dimensions(&self) -> Result<Vec<String>, CubeError> {
+        let mut result = vec![];
+        for dim in self.schema.all_dimensions() {
+            if has_multi_stage_members(dim, false)? {
+                result.push(dim.clone().resolve_reference_chain().full_name());
+            }
+        }
+        result.sort();
+        Ok(result)
+    }
+
+    pub fn join_dimensions(&self) -> Result<Vec<Rc<MemberSymbol>>, CubeError> {
+        let mut result = vec![];
+        for dim in self.schema.all_dimensions() {
+            if !has_multi_stage_members(dim, false)? {
+                result.push(dim.clone());
+            }
+        }
+        Ok(result)
+    }
+}
+
+impl PrettyPrint for MultiStageDimensionCalculation {
+    fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
+        result.println(&format!("Dimension Calculation",), state);
+        let state = state.new_level();
+        let details_state = state.new_level();
+        result.println("schema:", &state);
+        self.schema().pretty_print(result, &details_state);
+        if !self.order_by().is_empty() {
+            result.println("order_by:", &state);
+            for order_by in self.order_by().iter() {
+                result.println(
+                    &format!(
+                        "{} {}",
+                        order_by.name(),
+                        if order_by.desc() { "desc" } else { "asc" }
+                    ),
+                    &details_state,
+                );
+            }
+        }
+        result.println("source:", &state);
+        self.source().pretty_print(result, &details_state);
+    }
+}
+
+impl LogicalNode for MultiStageDimensionCalculation {
+    fn as_plan_node(self: &Rc<Self>) -> PlanNode {
+        PlanNode::MultiStageDimensionCalculation(self.clone())
+    }
+
+    fn inputs(&self) -> Vec<PlanNode> {
+        vec![self.source().as_plan_node()]
+    }
+
+    fn with_inputs(self: Rc<Self>, inputs: Vec<PlanNode>) -> Result<Rc<Self>, CubeError> {
+        check_inputs_len(&inputs, 1, self.node_name())?;
+        let source = &inputs[0];
+
+        Ok(Rc::new(
+            Self::builder()
+                .schema(self.schema().clone())
+                .order_by(self.order_by().clone())
+                .multi_stage_dimension(self.multi_stage_dimension.clone())
+                .source(source.clone().into_logical_node()?)
+                .build(),
+        ))
+    }
+
+    fn node_name(&self) -> &'static str {
+        "MultiStageDimensionCalculation"
+    }
+
+    fn try_from_plan_node(plan_node: PlanNode) -> Result<Rc<Self>, CubeError> {
+        if let PlanNode::MultiStageDimensionCalculation(item) = plan_node {
+            Ok(item)
+        } else {
+            Err(cast_error(&plan_node, "MultiStageMeasureCalculation"))
+        }
+    }
+}

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

@@ -5,6 +5,7 @@ use std::rc::Rc;
 pub enum MultiStageMemberLogicalType {
     LeafMeasure(Rc<MultiStageLeafMeasure>),
     MeasureCalculation(Rc<MultiStageMeasureCalculation>),
+    DimensionCalculation(Rc<MultiStageDimensionCalculation>),
     GetDateRange(Rc<MultiStageGetDateRange>),
     TimeSeries(Rc<MultiStageTimeSeries>),
     RollingWindow(Rc<MultiStageRollingWindow>),
@@ -15,6 +16,7 @@ impl MultiStageMemberLogicalType {
         match self {
             Self::LeafMeasure(item) => item.as_plan_node(),
             Self::MeasureCalculation(item) => item.as_plan_node(),
+            Self::DimensionCalculation(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(),
@@ -25,6 +27,9 @@ impl MultiStageMemberLogicalType {
         Ok(match self {
             Self::LeafMeasure(_) => Self::LeafMeasure(plan_node.into_logical_node()?),
             Self::MeasureCalculation(_) => Self::MeasureCalculation(plan_node.into_logical_node()?),
+            Self::DimensionCalculation(_) => {
+                Self::DimensionCalculation(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()?),
@@ -37,6 +42,7 @@ impl PrettyPrint for MultiStageMemberLogicalType {
         match self {
             Self::LeafMeasure(measure) => measure.pretty_print(result, state),
             Self::MeasureCalculation(calculation) => calculation.pretty_print(result, state),
+            Self::DimensionCalculation(calculation) => calculation.pretty_print(result, state),
             Self::GetDateRange(get_date_range) => get_date_range.pretty_print(result, state),
             Self::TimeSeries(time_series) => time_series.pretty_print(result, state),
             Self::RollingWindow(rolling_window) => rolling_window.pretty_print(result, state),

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

@@ -1,12 +1,14 @@
 mod calculation;
 mod common;
+mod dimension;
 mod get_date_range;
 mod leaf_measure;
 mod member;
 mod rolling_window;
 mod time_series;
 
 pub use calculation::*;
+pub use dimension::*;
 pub use get_date_range::*;
 pub use leaf_measure::*;
 pub use member::*;

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

@@ -270,6 +270,7 @@ impl<'a> DimensionMatcher<'a> {
             filter.member_evaluator().clone()
         };
         let add_to_matched_dimension = add_to_matched_dimension && filter.is_single_value_equal();
-        self.try_match_symbol(&symbol, add_to_matched_dimension)
+        let res = self.try_match_symbol(&symbol, add_to_matched_dimension)?;
+        Ok(res)
     }
 }