feat(tesseract): Support multiple time dimensions in pre-aggregations (#10599)

Author: waralexrom

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

@@ -170,6 +170,85 @@ describe('PreAggregationsMultiStage', () => {
 
   })
 
+  cube('coach', {
+    sql: \`
+      SELECT 101 AS id, '2025-01-01'::TIMESTAMP AS time UNION ALL
+      SELECT 102 AS id, '2025-02-01'::TIMESTAMP AS time UNION ALL
+      SELECT 103 AS id, '2025-02-02'::TIMESTAMP AS time UNION ALL
+      SELECT 104 AS id, '2025-03-01'::TIMESTAMP AS time UNION ALL
+      SELECT 105 AS id, '2025-03-02'::TIMESTAMP AS time UNION ALL
+      SELECT 106 AS id, '2025-03-03'::TIMESTAMP AS time UNION ALL
+      SELECT 107 AS id, '2025-04-01'::TIMESTAMP AS time UNION ALL
+      SELECT 108 AS id, '2025-04-02'::TIMESTAMP AS time UNION ALL
+      SELECT 109 AS id, '2025-04-03'::TIMESTAMP AS time UNION ALL
+      SELECT 110 AS id, '2025-04-04'::TIMESTAMP AS time UNION ALL
+      SELECT 111 AS id, '2025-05-01'::TIMESTAMP AS time UNION ALL
+      SELECT 112 AS id, '2025-05-02'::TIMESTAMP AS time UNION ALL
+      SELECT 113 AS id, '2025-05-03'::TIMESTAMP AS time UNION ALL
+      SELECT 114 AS id, '2025-05-04'::TIMESTAMP AS time UNION ALL
+      SELECT 115 AS id, '2025-05-05'::TIMESTAMP AS time UNION ALL
+      SELECT 116 AS id, '2025-06-01'::TIMESTAMP AS time UNION ALL
+      SELECT 117 AS id, '2025-06-02'::TIMESTAMP AS time UNION ALL
+      SELECT 118 AS id, '2025-06-03'::TIMESTAMP AS time UNION ALL
+      SELECT 119 AS id, '2025-06-04'::TIMESTAMP AS time UNION ALL
+      SELECT 120 AS id, '2025-06-05'::TIMESTAMP AS time UNION ALL
+      SELECT 121 AS id, '2025-06-06'::TIMESTAMP AS time UNION ALL
+      SELECT 122 AS id, '2025-07-01'::TIMESTAMP AS time UNION ALL
+      SELECT 123 AS id, '2025-07-02'::TIMESTAMP AS time UNION ALL
+      SELECT 124 AS id, '2025-07-03'::TIMESTAMP AS time UNION ALL
+      SELECT 125 AS id, '2025-07-04'::TIMESTAMP AS time UNION ALL
+      SELECT 126 AS id, '2025-07-05'::TIMESTAMP AS time UNION ALL
+      SELECT 127 AS id, '2025-07-06'::TIMESTAMP AS time UNION ALL
+      SELECT 128 AS id, '2025-07-07'::TIMESTAMP AS time
+    \`,
+
+    dimensions: {
+      time: {
+        sql: 'time',
+        type: 'time',
+        public: false,
+      },
+    },
+
+    measures: {
+      count_distinct: {
+        sql: 'id',
+        type: 'countDistinct',
+        public: false,
+      },
+      count_distinct__sum_by_quarter_aux: {
+        multi_stage: true,
+        sql: \`\${count_distinct}\`,
+        type: 'sum',
+        add_group_by: [time.month],
+        group_by: [time.quarter],
+        public: false,
+      },
+      count_distinct__sum_by_quarter: {
+        multi_stage: true,
+        sql: \`\${count_distinct__sum_by_quarter_aux}\`,
+        type: 'sum',
+        add_group_by: [time.quarter],
+      },
+    },
+
+    preAggregations: {
+      main: {
+        type: 'rollup',
+        measures: [count_distinct],
+        timeDimensions: [
+          {
+            dimension: time,
+            granularity: 'month'
+          },
+          {
+            dimension: time,
+            granularity: 'quarter'
+          }
+        ]
+      },
+    },
+  })
 
    `);
 
@@ -285,6 +364,30 @@ describe('PreAggregationsMultiStage', () => {
         );
       });
     }));
+
+    it('multi stage count_distinct sum by quarter with pre-aggregation', () => compiler.compile().then(() => {
+      const query = new PostgresQuery({ joinGraph, cubeEvaluator, compiler }, {
+        measures: [
+          'coach.count_distinct__sum_by_quarter'
+        ],
+        timezone: 'UTC',
+        preAggregationsSchema: '',
+        cubestoreSupportMultistage: true
+      });
+
+      const preAggregationsDescription: any = query.preAggregations?.preAggregationsDescription();
+      const sqlAndParams = query.buildSqlAndParams();
+      expect(preAggregationsDescription[0].tableName).toEqual('coach_main');
+      expect(sqlAndParams[0]).toContain('coach_main');
+
+      return dbRunner.evaluateQueryWithPreAggregations(query).then(res => {
+        expect(res).toEqual(
+          [
+            { coach__count_distinct__sum_by_quarter: '28' },
+          ]
+        );
+      });
+    }));
   } else {
     it.skip('multi stage pre-aggregations', () => {
       // Skipping because it works only in Tesseract

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

@@ -29,6 +29,7 @@ pub mod member_sql;
 pub mod options_member;
 pub mod pre_aggregation_description;
 pub mod pre_aggregation_obj;
+pub mod pre_aggregation_time_dimension;
 pub mod security_context;
 pub mod segment_definition;
 pub mod sql_templates_render;

rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/pre_aggregation_description.rs

@@ -1,7 +1,11 @@
 use super::member_sql::{MemberSql, NativeMemberSql};
+use super::pre_aggregation_time_dimension::{
+    NativePreAggregationTimeDimension, PreAggregationTimeDimension,
+};
 use cubenativeutils::wrappers::serializer::{
     NativeDeserialize, NativeDeserializer, NativeSerialize,
 };
+use cubenativeutils::wrappers::NativeArray;
 use cubenativeutils::wrappers::NativeContextHolder;
 use cubenativeutils::wrappers::NativeObjectHandle;
 use cubenativeutils::CubeError;
@@ -38,4 +42,9 @@ pub trait PreAggregationDescription {
 
     #[nbridge(field, optional)]
     fn rollup_references(&self) -> Result<Option<Rc<dyn MemberSql>>, CubeError>;
+
+    #[nbridge(field, vec, optional)]
+    fn time_dimension_references(
+        &self,
+    ) -> Result<Option<Vec<Rc<dyn PreAggregationTimeDimension>>>, CubeError>;
 }

rust/cubesqlplanner/cubesqlplanner/src/cube_bridge/pre_aggregation_time_dimension.rs

@@ -0,0 +1,21 @@
+use super::member_sql::{MemberSql, NativeMemberSql};
+use cubenativeutils::wrappers::serializer::{
+    NativeDeserialize, NativeDeserializer, NativeSerialize,
+};
+use cubenativeutils::wrappers::NativeContextHolder;
+use cubenativeutils::wrappers::NativeObjectHandle;
+use cubenativeutils::CubeError;
+use serde::{Deserialize, Serialize};
+use std::any::Any;
+use std::rc::Rc;
+
+#[derive(Deserialize, Serialize, Clone, Debug)]
+pub struct PreAggregationTimeDimensionStatic {
+    pub granularity: String,
+}
+
+#[nativebridge::native_bridge(PreAggregationTimeDimensionStatic)]
+pub trait PreAggregationTimeDimension {
+    #[nbridge(field)]
+    fn dimension(&self) -> Result<Rc<dyn MemberSql>, CubeError>;
+}

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

@@ -34,7 +34,7 @@ pub struct DimensionMatcher<'a> {
     query_tools: Rc<QueryTools>,
     pre_aggregation: &'a CompiledPreAggregation,
     pre_aggregation_dimensions: HashMap<String, bool>,
-    pre_aggregation_time_dimensions: HashMap<String, (Option<Rc<TimeDimensionSymbol>>, bool)>,
+    pre_aggregation_time_dimensions: HashMap<String, Vec<(Rc<TimeDimensionSymbol>, bool)>>,
     pre_aggregation_segments: HashMap<String, bool>,
     result: MatchState,
 }
@@ -46,17 +46,17 @@ impl<'a> DimensionMatcher<'a> {
             .iter()
             .map(|d| (d.full_name(), false))
             .collect();
-        let pre_aggregation_time_dimensions = pre_aggregation
-            .time_dimensions
-            .iter()
-            .map(|dim| {
-                if let Ok(td) = dim.as_time_dimension() {
-                    (td.base_symbol().full_name(), (Some(td), false))
-                } else {
-                    (dim.full_name(), (None, false))
-                }
-            })
-            .collect::<HashMap<_, _>>();
+        let mut pre_aggregation_time_dimensions =
+            HashMap::<String, Vec<(Rc<TimeDimensionSymbol>, bool)>>::new();
+        for dim in pre_aggregation.time_dimensions.iter() {
+            if let Ok(td) = dim.as_time_dimension() {
+                let key = td.base_symbol().full_name();
+                pre_aggregation_time_dimensions
+                    .entry(key)
+                    .or_default()
+                    .push((td, false));
+            }
+        }
         let pre_aggregation_segments = pre_aggregation
             .segments
             .iter()
@@ -128,12 +128,15 @@ impl<'a> DimensionMatcher<'a> {
             MatchState::Partial
         };
         self.result = self.result.combine(&dimension_coverage_result);
-        let time_dimension_coverage_result =
-            if self.pre_aggregation_time_dimensions.values().all(|v| v.1) {
-                MatchState::Full
-            } else {
-                MatchState::Partial
-            };
+        let time_dimension_coverage_result = if self
+            .pre_aggregation_time_dimensions
+            .values()
+            .all(|entries| entries.iter().all(|(_, matched)| *matched))
+        {
+            MatchState::Full
+        } else {
+            MatchState::Partial
+        };
         self.result = self.result.combine(&time_dimension_coverage_result);
         let segment_coverage_result = if self.pre_aggregation_segments.values().all(|v| *v) {
             MatchState::Full
@@ -233,39 +236,41 @@ impl<'a> DimensionMatcher<'a> {
         };
         let base_symbol_name = time_dimension.base_symbol().full_name();
 
-        if let Some(found) = self
+        if let Some(entries) = self
             .pre_aggregation_time_dimensions
             .get_mut(&base_symbol_name)
         {
-            if add_to_matched_dimension {
-                found.1 = true;
+            // First, look for exact granularity match
+            let exact_match = entries
+                .iter_mut()
+                .find(|(td, _)| granularity.is_none() || td.granularity() == &granularity);
+            if let Some((_, matched)) = exact_match {
+                if add_to_matched_dimension {
+                    *matched = true;
+                }
+                return Ok(MatchState::Full);
             }
 
-            let pre_agg_td = &found.0;
-            let pre_aggr_granularity = if let Some(pre_agg_td) = pre_agg_td {
-                pre_agg_td.granularity().clone()
-            } else {
-                None
-            };
-
-            if granularity.is_none() || pre_aggr_granularity == granularity {
-                Ok(MatchState::Full)
-            } else if pre_aggr_granularity.is_none() {
-                Ok(MatchState::NotMatched)
-            } else if let Some(pre_agg_td) = pre_agg_td {
+            // No exact match — find the finest pre-agg granularity that covers the query
+            let mut best_match = MatchState::NotMatched;
+            for (pre_agg_td, matched) in entries.iter_mut() {
+                let pre_aggr_granularity = pre_agg_td.granularity();
+                if pre_aggr_granularity.is_none() {
+                    continue;
+                }
                 let min_granularity = GranularityHelper::min_granularity_for_time_dimensions(
                     (&granularity, time_dimension),
-                    (&pre_aggr_granularity, &pre_agg_td),
+                    (pre_aggr_granularity, pre_agg_td),
                 )?;
-
-                if min_granularity == pre_aggr_granularity {
-                    Ok(MatchState::Partial)
-                } else {
-                    Ok(MatchState::NotMatched)
+                if min_granularity == *pre_aggr_granularity {
+                    if add_to_matched_dimension {
+                        *matched = true;
+                    }
+                    best_match = MatchState::Partial;
+                    break;
                 }
-            } else {
-                Ok(MatchState::NotMatched)
             }
+            Ok(best_match)
         } else {
             if time_dimension.owned_by_cube() {
                 Ok(MatchState::NotMatched)