sql-generation parity

Author: waralexrom

packages/cubejs-schema-compiler/test/integration/postgres/sql-generation.test.ts

@@ -2539,7 +2539,11 @@ SELECT 1 AS revenue,  cast('2024-01-01' AS timestamp) as time UNION ALL
       filters: [],
       order: [{
         id: 'visitor_checkins.id'
-      }],
+      },
+      {
+        id: 'visitor_checkins.created_at'
+      }
+      ],
       ungrouped: true
     });
 

rust/cubesqlplanner/\

@@ -0,0 +1,290 @@
+use super::super::{LogicalNodeProcessor, ProcessableNode, PushDownBuilderContext};
+use crate::logical_plan::{
+    pretty_print, pretty_print_rc, FullKeyAggregate, ResolvedMultipliedMeasures, SimpleQuery,
+    SimpleQuerySource,
+};
+use crate::physical_plan_builder::PhysicalPlanBuilder;
+use crate::plan::{
+    Expr, Filter, From, JoinBuilder, JoinCondition, MemberExpression, QualifiedColumnName,
+    QueryPlan, Select, SelectBuilder, Union,
+};
+use crate::planner::query_tools::QueryTools;
+use crate::planner::sql_evaluator::sql_nodes::SqlNodesFactory;
+use crate::planner::sql_evaluator::{MemberSymbol, ReferencesBuilder};
+use crate::planner::sql_templates::PlanSqlTemplates;
+use crate::planner::{BaseMember, MemberSymbolRef};
+use cubenativeutils::CubeError;
+use itertools::Itertools;
+use std::collections::HashMap;
+use std::rc::Rc;
+
+trait FullKeyAggregateStrategy {
+    fn process(
+        &self,
+        full_key_aggregate: &FullKeyAggregate,
+        context: &PushDownBuilderContext,
+    ) -> Result<Rc<From>, CubeError>;
+}
+
+struct KeysFullKeyAggregateStrategy<'a> {
+    builder: &'a PhysicalPlanBuilder,
+}
+
+impl<'a> KeysFullKeyAggregateStrategy<'a> {
+    pub fn new(builder: &'a PhysicalPlanBuilder) -> Rc<Self> {
+        Rc::new(Self { builder })
+    }
+}
+
+impl<'a> FullKeyAggregateStrategy for KeysFullKeyAggregateStrategy<'a> {
+    fn process(
+        &self,
+        full_key_aggregate: &FullKeyAggregate,
+        context: &PushDownBuilderContext,
+    ) -> Result<Rc<From>, CubeError> {
+        let query_tools = self.builder.query_tools();
+        let mut keys_queries = vec![];
+        let mut data_queries = vec![];
+        let mut keys_context = context.clone();
+        keys_context.dimensions_query = true;
+        if let Some(resolved_multiplied_measures) = &full_key_aggregate.multiplied_measures_resolver
+        {
+            match resolved_multiplied_measures {
+                ResolvedMultipliedMeasures::ResolveMultipliedMeasures(
+                    resolve_multiplied_measures,
+                ) => {
+                    for regular_measure_query in resolve_multiplied_measures
+                        .regular_measure_subqueries
+                        .iter()
+                    {
+                        let keys_query = self
+                            .builder
+                            .process_node(regular_measure_query.as_ref(), &keys_context)?;
+                        keys_queries.push(keys_query);
+                        let query = self
+                            .builder
+                            .process_node(regular_measure_query.as_ref(), &context)?;
+                        data_queries.push(query);
+                    }
+                    for multiplied_measure_query in resolve_multiplied_measures
+                        .aggregate_multiplied_subqueries
+                        .iter()
+                    {
+                        let keys_query = self
+                            .builder
+                            .process_node(multiplied_measure_query.as_ref(), &keys_context)?;
+                        keys_queries.push(keys_query);
+                        let query = self
+                            .builder
+                            .process_node(multiplied_measure_query.as_ref(), &context)?;
+                        data_queries.push(query);
+                    }
+                }
+                ResolvedMultipliedMeasures::PreAggregation(simple_query) => todo!(),
+            }
+        }
+        for cte_souce in full_key_aggregate.cte_sources.iter() {
+        }
+        if data_queries.is_empty() {
+            return Err(CubeError::internal(format!(
+                "FullKeyAggregate should have at least one source: {}",
+                pretty_print(full_key_aggregate)
+            )));
+        }
+
+        if data_queries.len() == 1 {
+            let select = data_queries[0].clone();
+            let result = From::new_from_subselect(select, "fk_aggregate".to_string());
+            return Ok(result);
+        }
+
+        let keys_from = From::new_from_union(
+            Rc::new(Union::new_from_subselects(&keys_queries)),
+            "pk_aggregate_keys_source".to_string(),
+        );
+        let references_builder = ReferencesBuilder::new(keys_from.clone());
+        let mut keys_select_builder = SelectBuilder::new(keys_from);
+
+        for member in full_key_aggregate.schema.all_dimensions() {
+            let alias = references_builder.resolve_alias_for_member(&member.full_name(), &None);
+            if alias.is_none() {
+                return Err(CubeError::internal(format!(
+                    "Source for {} not found in full key aggregate subqueries",
+                    member.full_name()
+                )));
+            }
+            let reference = QualifiedColumnName::new(None, alias.unwrap());
+            let member_ref = member.clone().as_base_member(query_tools.clone())?;
+            keys_select_builder.add_projection_member_reference(&member_ref, reference);
+        }
+        keys_select_builder.set_distinct();
+
+        let sql_context = SqlNodesFactory::new();
+        let keys_select = Rc::new(keys_select_builder.build(sql_context));
+
+        let keys_alias = "fk_aggregate_keys".to_string();
+
+        let mut join_builder =
+            JoinBuilder::new_from_subselect(keys_select.clone(), keys_alias.clone());
+
+        for (i, query) in data_queries.into_iter().enumerate() {
+            let query_alias = format!("q_{}", i);
+            let conditions = full_key_aggregate
+                .schema
+                .all_dimensions()
+                .map(|dim| -> Result<_, CubeError> {
+                    let member_ref = dim.clone().as_base_member(query_tools.clone())?;
+                    let alias_in_keys_query =
+                        keys_select.schema().resolve_member_alias(&member_ref);
+                    let keys_query_ref = Expr::Reference(QualifiedColumnName::new(
+                        Some(keys_alias.clone()),
+                        alias_in_keys_query,
+                    ));
+                    let alias_in_data_query = query.schema().resolve_member_alias(&member_ref);
+                    let data_query_ref = Expr::Reference(QualifiedColumnName::new(
+                        Some(query_alias.clone()),
+                        alias_in_data_query,
+                    ));
+
+                    Ok(vec![(keys_query_ref, data_query_ref)])
+                })
+                .collect::<Result<Vec<_>, _>>()?;
+
+            join_builder.left_join_subselect(
+                query,
+                query_alias.clone(),
+                JoinCondition::new_dimension_join(conditions, true),
+            );
+        }
+
+        let result = join_builder.build();
+        Ok(From::new_from_join(result))
+    }
+}
+
+struct InnerJoinFullKeyAggregateStrategy<'a> {
+    builder: &'a PhysicalPlanBuilder,
+}
+
+impl<'a> InnerJoinFullKeyAggregateStrategy<'a> {
+    pub fn new(builder: &'a PhysicalPlanBuilder) -> Rc<Self> {
+        Rc::new(Self { builder })
+    }
+}
+
+impl<'a> FullKeyAggregateStrategy for InnerJoinFullKeyAggregateStrategy<'a> {
+    fn process(
+        &self,
+        full_key_aggregate: &FullKeyAggregate,
+        context: &PushDownBuilderContext,
+    ) -> Result<Rc<From>, CubeError> {
+        let query_tools = self.builder.query_tools();
+        let mut data_queries = vec![];
+        if let Some(resolved_multiplied_measures) = &full_key_aggregate.multiplied_measures_resolver
+        {
+            match resolved_multiplied_measures {
+                ResolvedMultipliedMeasures::ResolveMultipliedMeasures(
+                    resolve_multiplied_measures,
+                ) => {
+                    for regular_measure_query in resolve_multiplied_measures
+                        .regular_measure_subqueries
+                        .iter()
+                    {
+                        let query = self
+                            .builder
+                            .process_node(regular_measure_query.as_ref(), &context)?;
+                        data_queries.push(query);
+                    }
+                    for multiplied_measure_query in resolve_multiplied_measures
+                        .aggregate_multiplied_subqueries
+                        .iter()
+                    {
+                        let query = self
+                            .builder
+                            .process_node(multiplied_measure_query.as_ref(), &context)?;
+                        data_queries.push(query);
+                    }
+                }
+                ResolvedMultipliedMeasures::PreAggregation(simple_query) => todo!(),
+            }
+        }
+        if data_queries.is_empty() {
+            return Err(CubeError::internal(format!(
+                "FullKeyAggregate should have at least one source: {}",
+                pretty_print(full_key_aggregate)
+            )));
+        }
+
+        if data_queries.len() == 1 {
+            let select = data_queries[0].clone();
+            let result = From::new_from_subselect(select, "fk_aggregate".to_string());
+            return Ok(result);
+        }
+
+        let mut join_builder =
+            JoinBuilder::new_from_subselect(data_queries[0].clone(), "q_0".to_string());
+
+        for (i, query) in data_queries.iter().skip(1).enumerate() {
+            let prev_alias = format!("q_{}", i);
+            let query_alias = format!("q_{}", i + 1);
+            let conditions = full_key_aggregate
+                .schema
+                .all_dimensions()
+                .map(|dim| -> Result<_, CubeError> {
+                    let member_ref = dim.clone().as_base_member(query_tools.clone())?;
+                    let alias_in_prev_query =
+                        data_queries[i].schema().resolve_member_alias(&member_ref);
+                    let prev_query_ref = Expr::Reference(QualifiedColumnName::new(
+                        Some(prev_alias.clone()),
+                        alias_in_prev_query,
+                    ));
+                    let alias_in_data_query = query.schema().resolve_member_alias(&member_ref);
+                    let data_query_ref = Expr::Reference(QualifiedColumnName::new(
+                        Some(query_alias.clone()),
+                        alias_in_data_query,
+                    ));
+
+                    Ok(vec![(prev_query_ref, data_query_ref)])
+                })
+                .collect::<Result<Vec<_>, _>>()?;
+
+            join_builder.inner_join_subselect(
+                query.clone(),
+                query_alias.clone(),
+                JoinCondition::new_dimension_join(conditions, true),
+            );
+        }
+
+        let result = join_builder.build();
+        Ok(From::new_from_join(result))
+    }
+}
+
+pub struct FullKeyAggregateProcessor<'a> {
+    builder: &'a PhysicalPlanBuilder,
+}
+
+impl<'a> LogicalNodeProcessor<'a, FullKeyAggregate> for FullKeyAggregateProcessor<'a> {
+    type PhysycalNode = Rc<From>;
+    fn new(builder: &'a PhysicalPlanBuilder) -> Self {
+        Self { builder }
+    }
+
+    fn process(
+        &self,
+        full_key_aggregate: &FullKeyAggregate,
+        context: &PushDownBuilderContext,
+    ) -> Result<Self::PhysycalNode, CubeError> {
+        let strategy: Rc<dyn FullKeyAggregateStrategy> =
+            if full_key_aggregate.schema.has_dimensions() {
+                KeysFullKeyAggregateStrategy::new(self.builder)
+            } else {
+                InnerJoinFullKeyAggregateStrategy::new(self.builder)
+            };
+        strategy.process(full_key_aggregate, context)
+    }
+}
+
+impl ProcessableNode for FullKeyAggregate {
+    type ProcessorType<'a> = FullKeyAggregateProcessor<'a>;
+}

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

@@ -9,10 +9,10 @@ pub enum AggregateMultipliedSubquerySouce {
 
 pub struct AggregateMultipliedSubquery {
     pub schema: Rc<LogicalSchema>,
-    pub dimension_subqueries: Vec<Rc<DimensionSubQuery>>,
     pub keys_subquery: Rc<KeysSubQuery>,
     pub pk_cube: Rc<Cube>, //FIXME may be duplication with information in keys_subquery
     pub source: Rc<AggregateMultipliedSubquerySouce>,
+    pub dimension_subqueries: Vec<Rc<DimensionSubQuery>>,
 }
 
 impl PrettyPrint for AggregateMultipliedSubquery {
@@ -22,12 +22,6 @@ impl PrettyPrint for AggregateMultipliedSubquery {
         let details_state = state.new_level();
         result.println("schema:", &state);
         self.schema.pretty_print(result, &details_state);
-        if !self.dimension_subqueries.is_empty() {
-            result.println("dimension_subqueries:", &state);
-            for subquery in self.dimension_subqueries.iter() {
-                subquery.pretty_print(result, &details_state);
-            }
-        }
         result.println("keys_subquery:", &state);
         self.keys_subquery.pretty_print(result, &details_state);
         result.println("source:", &state);
@@ -38,12 +32,16 @@ impl PrettyPrint for AggregateMultipliedSubquery {
                     .pretty_print(result, &details_state.new_level());
             }
             AggregateMultipliedSubquerySouce::MeasureSubquery(measure_subquery) => {
-                result.println(
-                    &format!("MeasureSubquery: {}", measure_subquery.measures.len()),
-                    &details_state,
-                );
+                result.println(&format!("MeasureSubquery: "), &details_state);
                 measure_subquery.pretty_print(result, &details_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);
+            }
+        }
     }
 }

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

@@ -1,5 +1,5 @@
 use super::pretty_print::*;
-use super::Cube;
+use super::*;
 use crate::planner::sql_evaluator::SqlCall;
 use std::rc::Rc;
 
@@ -34,6 +34,7 @@ impl PrettyPrint for LogicalJoinItem {
 pub struct LogicalJoin {
     pub root: Rc<Cube>,
     pub joins: Vec<LogicalJoinItem>,
+    pub dimension_subqueries: Vec<Rc<DimensionSubQuery>>,
 }
 
 impl PrettyPrint for LogicalJoin {
@@ -48,5 +49,12 @@ impl PrettyPrint for LogicalJoin {
         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);
+            }
+        }
     }
 }