in work

Author: waralexrom

packages/cubejs-schema-compiler/src/adapter/BaseQuery.js

@@ -715,6 +715,9 @@ export class BaseQuery {
     const res = buildResult.result;
     // FIXME
     res[1] = [...res[1]];
+    if (res[2]) {
+      this.preAggregations.preAggregationForQuery = res[2];
+    }
     return res;
   }
 
@@ -738,6 +741,10 @@ export class BaseQuery {
     return timeSeriesFromCustomInterval(granularityInterval, dateRange, moment(origin), { timestampPrecision: 3 });
   }
 
+  getPreAggregationByName(cube, preAggregationName) {
+    return this.preAggregations.getRollupPreAggregationByName(cube, preAggregationName);
+  }
+
   get shouldReuseParams() {
     return false;
   }

packages/cubejs-schema-compiler/src/adapter/PreAggregations.js

@@ -819,6 +819,29 @@ export class PreAggregations {
     )(preAggregations);
   }
 
+  getRollupPreAggregationByName(cube, preAggregationName) {
+    const canUsePreAggregation = () => true;
+    const preAggregation = R.pipe(
+      R.toPairs,
+      // eslint-disable-next-line no-unused-vars
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      R.filter(([k, a]) => a.type === 'rollup' || a.type === 'rollupJoin' || a.type === 'rollupLambda'),
+      // eslint-disable-next-line no-unused-vars
+      // eslint-disable-next-line @typescript-eslint/no-unused-vars
+      R.find(([k, a]) => k === preAggregationName)
+    )(this.query.cubeEvaluator.preAggregationsForCube(cube));
+    if (preAggregation) {
+      const tableName = this.preAggregationTableName(cube, preAggregation[0], preAggregation[1]);
+      const preAggObj = preAggregation ? this.evaluatedPreAggregationObj(cube, preAggregation[0], preAggregation[1], canUsePreAggregation) : {};
+      return {
+        tableName,
+        ...preAggObj
+      };
+    } else {
+      return {};
+    }
+  }
+
   // TODO check multiplication factor didn't change
   buildRollupJoin(preAggObj, preAggObjsToJoin) {
     return this.query.cacheValue(

packages/cubejs-schema-compiler/test/integration/postgres/pre-aggregations.test.ts

@@ -231,7 +231,7 @@ describe('PreAggregations', () => {
           granularity: 'hour',
           partitionGranularity: 'month'
         },
-/*         countCustomGranularity: {
+        countCustomGranularity: {
           measures: [count],
           timeDimension: createdAt,
           granularity: 'hourTenMinOffset',
@@ -517,7 +517,6 @@ describe('PreAggregations', () => {
     console.log(queryAndParams);
     console.log("!!!! pre aggrs", query.preAggregations?.preAggregationForQuery);
     console.log("!!!! pre aggrs fun", query.preAggregations?.preAggregationForQuery.preAggregation);
-    console.log("!!!! pre aggrs fun 2", query.preAggregations?.preAggregationForQuery.preAggregation.dimensionReferences.toString());
     expect(query.preAggregations?.preAggregationForQuery?.canUsePreAggregation).toEqual(true);
 
     return dbRunner.evaluateQueryWithPreAggregations(query).then(res => {
@@ -887,7 +886,7 @@ describe('PreAggregations', () => {
     });
   }));
 
-  it('multiplied measure match', () => compiler.compile().then(() => {
+  it('multiplied measure match 1', () => compiler.compile().then(() => {
     const query = new PostgresQuery({ joinGraph, cubeEvaluator, compiler }, {
       measures: [
         'visitors.count'

packages/cubejs-schema-compiler/test/integration/utils/BaseDbRunner.ts

@@ -171,8 +171,6 @@ export class BaseDbRunner {
 
   public async evaluateQueryWithPreAggregations(query, seed = this.nextSeed++) {
     const preAggregationsDescription = query.preAggregations?.preAggregationsDescription();
-    // console.log(preAggregationsDescription);
-
     await Promise.all(preAggregationsDescription.map(
       async desc => {
         if (desc.partitionGranularity) {

rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/base_tools.rs

@@ -1,8 +1,10 @@
 use super::base_query_options::FilterItem;
 use super::filter_group::{FilterGroup, NativeFilterGroup};
 use super::filter_params::{FilterParams, NativeFilterParams};
+use cubenativeutils::wrappers::NativeArray;
 use super::security_context::{NativeSecurityContext, SecurityContext};
 use super::sql_templates_render::{NativeSqlTemplatesRender, SqlTemplatesRender};
+use super::pre_aggregation_obj::{NativePreAggregationObj, PreAggregationObj};
 use super::sql_utils::{NativeSqlUtils, SqlUtils};
 use cubenativeutils::wrappers::serializer::{
     NativeDeserialize, NativeDeserializer, NativeSerialize,
@@ -58,4 +60,6 @@ pub trait BaseTools {
         source: String,
         origin: String,
     ) -> Result<String, CubeError>;
+
+    fn get_pre_aggregation_by_name(&self, cube_name: String, name: String) -> Result<Rc<dyn PreAggregationObj>, CubeError>;
 }

rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/mod.rs

@@ -27,3 +27,4 @@ pub mod segment_definition;
 pub mod sql_templates_render;
 pub mod sql_utils;
 pub mod struct_with_sql_member;
+pub mod pre_aggregation_obj;

rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/pre_aggregation_obj.rs

@@ -0,0 +1,26 @@
+use super::member_sql::{MemberSql, NativeMemberSql};
+use cubenativeutils::wrappers::serializer::{
+    NativeDeserialize, NativeDeserializer, NativeSerialize,
+};
+use cubenativeutils::wrappers::NativeArray;
+use cubenativeutils::wrappers::NativeContextHolder;
+use cubenativeutils::wrappers::NativeObjectHandle;
+use cubenativeutils::CubeError;
+use serde::{Deserialize, Serialize};
+use std::any::Any;
+use std::rc::Rc;
+
+#[derive(Serialize, Deserialize, Debug)]
+pub struct PreAggregationObjStatic {
+    #[serde(rename = "tableName")]
+    pub table_name: Option<String>,
+    #[serde(rename = "preAggregationName")]
+    pub pre_aggregation_name: Option<String>,
+    pub cube: Option<String>,
+    #[serde(rename = "preAggregationId")]
+    pub pre_aggregation_id: Option<String>,
+}
+
+#[nativebridge::native_bridge(PreAggregationObjStatic)]
+pub trait PreAggregationObj {
+}

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/full_key_aggregate_query.rs

@@ -1,6 +1,6 @@
 use super::*;
-use crate::planner::query_properties::OrderByItem;
 use std::rc::Rc;
+use crate::planner::query_properties::OrderByItem;
 
 pub struct FullKeyAggregateQuery {
     pub multistage_members: Vec<Rc<LogicalMultiStageMember>>,
@@ -14,7 +14,7 @@ pub struct FullKeyAggregateQuery {
 }
 
 impl PrettyPrint for FullKeyAggregateQuery {
-    fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {
+    fn pretty_print(&self, result: &mut PrettyPrintResult, state: &PrettyPrintState) {  
         result.println("FullKeyAggregateQuery: ", state);
         let state = state.new_level();
         let details_state = state.new_level();
@@ -39,17 +39,10 @@ impl PrettyPrint for FullKeyAggregateQuery {
         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(&format!("{} {}", order_by.name(), if order_by.desc() { "desc" } else { "asc" }), &details_state);
             }
         }
         result.println("source:", &state);
         self.source.pretty_print(result, &details_state);
     }
-}
+}

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

@@ -8,12 +8,14 @@ mod join;
 mod keys_subquery;
 mod measure_subquery;
 mod multistage;
-mod pretty_print;
+pub mod pretty_print;
 mod query;
 mod regular_measures_query;
 mod resolve_multiplied_measures;
 mod schema;
 mod simple_query;
+mod pre_aggregation;
+pub mod optimizers;
 
 pub use aggregate_multiplied_subquery::*;
 pub use cube::*;
@@ -31,3 +33,5 @@ pub use regular_measures_query::*;
 pub use resolve_multiplied_measures::*;
 pub use schema::*;
 pub use simple_query::*;
+pub use optimizers::*;
+pub use pre_aggregation::*;

rust/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/common/cube_names_collector.rs

@@ -0,0 +1,135 @@
+use crate::logical_plan::*;
+use cubenativeutils::CubeError;
+use itertools::Itertools;
+use std::collections::HashSet;
+use std::rc::Rc;
+
+pub struct CubeNamesCollector {
+    cube_names: HashSet<String>,
+}
+
+impl CubeNamesCollector {
+    pub fn new() -> Self {
+        Self {
+            cube_names: HashSet::new(),
+        }
+    }
+
+    pub fn collect(&mut self, query: &Query) -> Result<(), CubeError> {
+        match query {
+            Query::SimpleQuery(query) => self.collect_from_simple_query(query),
+            Query::FullKeyAggregateQuery(query) => {
+                self.collect_from_full_key_aggregate_query(query)
+            }
+        }
+    }
+
+    pub fn result(self) -> Vec<String> {
+        self.cube_names.into_iter().collect_vec()
+    }
+
+    fn collect_from_simple_query(&mut self, query: &SimpleQuery) -> Result<(), CubeError> {
+        self.collect_from_simple_query_source(&query.source)?;
+        self.collect_from_dimension_subqueries(&query.dimension_subqueries)?;
+        Ok(())
+    }
+
+    fn collect_from_full_key_aggregate_query(
+        &mut self,
+        query: &FullKeyAggregateQuery,
+    ) -> Result<(), CubeError> {
+        self.collect_from_full_key_aggregate(&query.source)?;
+        Ok(())
+    }
+
+    fn collect_from_measure_subquery(
+        &mut self,
+        subquery: &Rc<MeasureSubquery>,
+    ) -> Result<(), CubeError> {
+        self.collect_from_logical_join(&subquery.source)?;
+        self.collect_from_dimension_subqueries(&subquery.dimension_subqueries)?;
+        Ok(())
+    }
+
+    fn collect_from_full_key_aggregate(
+        &mut self,
+        full_key_aggregate: &Rc<FullKeyAggregate>,
+    ) -> Result<(), CubeError> {
+        for source in full_key_aggregate.sources.iter() {
+            self.collect_from_full_key_aggregate_source(source)?;
+        }
+        Ok(())
+    }
+    fn collect_from_full_key_aggregate_source(
+        &mut self,
+        source: &FullKeyAggregateSource,
+    ) -> Result<(), CubeError> {
+        match source {
+            FullKeyAggregateSource::ResolveMultipliedMeasures(resolve_multiplied_measures) => {
+                self.collect_from_multiplied_measures_resolver(resolve_multiplied_measures)
+            }
+            FullKeyAggregateSource::MultiStageSubqueryRef(multi_stage_subquery_ref) => Ok(()),
+        }
+    }
+
+    fn collect_from_multiplied_measures_resolver(
+        &mut self,
+        resolver: &ResolveMultipliedMeasures,
+    ) -> Result<(), CubeError> {
+        for regular_subquery in resolver.regular_measure_subqueries.iter() {
+            self.collect_from_simple_query(&regular_subquery)?;
+        }
+        for aggregate_multiplied_subquery in resolver.aggregate_multiplied_subqueries.iter() {
+            self.collect_from_aggregate_multiplied_subquery(&aggregate_multiplied_subquery)?;
+        }
+        Ok(())
+    }
+
+    fn collect_from_aggregate_multiplied_subquery(
+        &mut self,
+        subquery: &Rc<AggregateMultipliedSubquery>,
+    ) -> Result<(), CubeError> {
+        self.collect_from_logical_join(&subquery.keys_subquery.source)?;
+        match subquery.source.as_ref() {
+            AggregateMultipliedSubquerySouce::Cube => {
+                self.cube_names.insert(subquery.pk_cube.name().clone());
+            }
+            AggregateMultipliedSubquerySouce::MeasureSubquery(measure_subquery) => {
+                self.collect_from_measure_subquery(&measure_subquery)?;
+            }
+        }
+        Ok(())
+    }
+
+    fn collect_from_simple_query_source(
+        &mut self,
+        source: &SimpleQuerySource,
+    ) -> Result<(), CubeError> {
+        match source {
+            SimpleQuerySource::LogicalJoin(join) => self.collect_from_logical_join(join),
+            SimpleQuerySource::PreAggregation(_) => Ok(()),
+        }
+    }
+
+    fn collect_from_logical_join(&mut self, join: &Rc<LogicalJoin>) -> Result<(), CubeError> {
+        self.cube_names.insert(join.root.name.clone());
+        for join_item in join.joins.iter() {
+            match join_item {
+                LogicalJoinItem::CubeJoinItem(cube_join_item) => {
+                    self.cube_names.insert(cube_join_item.cube.name.clone());
+                }
+            }
+        }
+        Ok(())
+    }
+
+    fn collect_from_dimension_subqueries(
+        &mut self,
+        dimension_subqueries: &Vec<Rc<DimensionSubQuery>>,
+    ) -> Result<(), CubeError> {
+        for subquery in dimension_subqueries.iter() {
+            self.collect(&subquery.query)?;
+        }
+        Ok(())
+    }
+}