feat(tesseract): Allow named calendar timeshifts for common intervals (#9777)

* throw error if named timeshift not found

* feat(tesseract): Allow named calendar timeshifts for common intervals

* add tests

* cargo fmt

* update snapshots

* fix

* simplify builder by moving extract_time_shifts to time_shift_state

Author: KSDaemon

packages/cubejs-schema-compiler/src/compiler/CubeValidator.ts

@@ -626,15 +626,26 @@ const MeasuresSchema = Joi.object().pattern(identifierRegex, Joi.alternatives().
   ]
 ));
 
-const CalendarTimeShiftItem = Joi.object({
-  name: identifier,
-  interval: regexTimeInterval,
-  type: Joi.string().valid('next', 'prior'),
-  sql: Joi.func().required(),
-})
-  .or('name', 'interval')
-  .with('interval', 'type')
-  .with('type', 'interval');
+const CalendarTimeShiftItem = Joi.alternatives().try(
+  Joi.object({
+    name: identifier.required(),
+    interval: regexTimeInterval.required(),
+    type: Joi.string().valid('next', 'prior').required(),
+    sql: Joi.forbidden()
+  }),
+  Joi.object({
+    name: identifier.required(),
+    sql: Joi.func().required(),
+    interval: Joi.forbidden(),
+    type: Joi.forbidden()
+  }),
+  Joi.object({
+    interval: regexTimeInterval.required(),
+    type: Joi.string().valid('next', 'prior').required(),
+    sql: Joi.func().required(),
+    name: Joi.forbidden()
+  })
+);
 
 const DimensionsSchema = Joi.object().pattern(identifierRegex, Joi.alternatives().try(
   inherit(BaseDimensionWithoutSubQuery, {

packages/cubejs-schema-compiler/test/integration/postgres/calendars.test.ts

@@ -80,6 +80,13 @@ cubes:
         time_shift:
           - name: one_year
 
+      - name: count_shifted_y_named_common_interval
+        type: number
+        multi_stage: true
+        sql: "{count}"
+        time_shift:
+          - name: one_year_common_interval
+
       - name: count_shifted_y1d_named
         type: number
         multi_stage: true
@@ -261,6 +268,10 @@ cubes:
           - name: one_year
             sql: "{CUBE}.retail_date_prev_year"
 
+          - name: one_year_common_interval
+            interval: 1 year
+            type: prior
+
           - name: one_year_and_one_day
             sql: "({CUBE}.retail_date_prev_year + interval '1 day')"
 
@@ -305,6 +316,10 @@ cubes:
           - name: one_year
             sql: "{CUBE}.retail_date_prev_year"
 
+          - name: one_year_common_interval
+            interval: 1 year
+            type: prior
+
           - name: one_year_and_one_day
             sql: "({CUBE}.retail_date_prev_year + interval '1 day')"
 
@@ -736,6 +751,22 @@ cubes:
           custom_calendar__retail_date_year: '2025-02-02T00:00:00.000Z',
         },
       ]));
+
+      it('Count shifted by year (custom named shift with common interval + custom granularity)', async () => runQueryTest({
+        measures: ['calendar_orders.count', 'calendar_orders.count_shifted_y_named_common_interval'],
+        timeDimensions: [{
+          dimension: 'custom_calendar.retail_date',
+          granularity: 'year',
+          dateRange: ['2025-02-02', '2026-02-01']
+        }],
+        order: [{ id: 'custom_calendar.retail_date' }]
+      }, [
+        {
+          calendar_orders__count: '37',
+          calendar_orders__count_shifted_y_named_common_interval: '39',
+          custom_calendar__retail_date_year: '2025-02-02T00:00:00.000Z',
+        },
+      ]));
     });
 
     describe('PK dimension time-shifts', () => {
@@ -889,6 +920,22 @@ cubes:
           custom_calendar__date_val_year: '2025-02-02T00:00:00.000Z',
         },
       ]));
+
+      it.skip('Count shifted by year (custom named shift with common interval + custom granularity)', async () => runQueryTest({
+        measures: ['calendar_orders.count', 'calendar_orders.count_shifted_y_named_common_interval'],
+        timeDimensions: [{
+          dimension: 'custom_calendar.date_val',
+          granularity: 'year',
+          dateRange: ['2025-02-02', '2026-02-01']
+        }],
+        order: [{ id: 'custom_calendar.date_val' }]
+      }, [
+        {
+          calendar_orders__count: '37',
+          calendar_orders__count_shifted_y_named_common_interval: '39',
+          custom_calendar__retail_date_year: '2025-02-02T00:00:00.000Z',
+        },
+      ]));
     });
   });
 });

packages/cubejs-schema-compiler/test/unit/__snapshots__/schema.test.ts.snap

@@ -275,6 +275,11 @@ Object {
           "name": "retail_date_prev_year",
           "sql": [Function],
         },
+        Object {
+          "interval": "1 year",
+          "name": "prev_year_named_common_interval",
+          "type": "prior",
+        },
         Object {
           "interval": "2 year",
           "sql": [Function],

packages/cubejs-schema-compiler/test/unit/fixtures/custom_calendar.yml

@@ -48,6 +48,10 @@ cubes:
           - name: retail_date_prev_year
             sql: "{CUBE.retail_date_prev_year}"
 
+          - name: prev_year_named_common_interval
+            interval: 1 year
+            type: prior
+
           # All needed intervals should be defined explicitly
           - interval: 2 year
             type: prior

rust/cubesqlplanner/cubesqlplanner/src/physical_plan_builder/builder.rs

@@ -5,15 +5,14 @@ use crate::planner::planners::multi_stage::TimeShiftState;
 use crate::planner::query_properties::OrderByItem;
 use crate::planner::query_tools::QueryTools;
 use crate::planner::sql_evaluator::sql_nodes::SqlNodesFactory;
-use crate::planner::sql_evaluator::symbols::CalendarDimensionTimeShift;
+use crate::planner::sql_evaluator::MemberSymbol;
 use crate::planner::sql_evaluator::ReferencesBuilder;
-use crate::planner::sql_evaluator::{DimensionTimeShift, MemberSymbol};
 use crate::planner::sql_templates::PlanSqlTemplates;
 use crate::planner::BaseMemberHelper;
 use crate::planner::SqlJoinCondition;
 use crate::planner::{BaseMember, MemberSymbolRef};
 use cubenativeutils::CubeError;
-use itertools::{Either, Itertools};
+use itertools::Itertools;
 use std::collections::HashMap;
 use std::collections::HashSet;
 use std::rc::Rc;
@@ -31,41 +30,10 @@ struct PhysicalPlanBuilderContext {
 }
 
 impl PhysicalPlanBuilderContext {
-    pub fn make_sql_nodes_factory(&self) -> SqlNodesFactory {
+    pub fn make_sql_nodes_factory(&self) -> Result<SqlNodesFactory, CubeError> {
         let mut factory = SqlNodesFactory::new();
 
-        let (time_shifts, calendar_time_shifts): (
-            HashMap<String, DimensionTimeShift>,
-            HashMap<String, CalendarDimensionTimeShift>,
-        ) = self
-            .time_shifts
-            .dimensions_shifts
-            .iter()
-            .partition_map(|(key, shift)| {
-                if let Ok(dimension) = shift.dimension.as_dimension() {
-                    if let Some(dim_shift_name) = &shift.name {
-                        if let Some((dim_key, cts)) =
-                            dimension.calendar_time_shift_for_named_interval(dim_shift_name)
-                        {
-                            return Either::Right((dim_key.clone(), cts.clone()));
-                        } else if let Some(_calendar_pk) = dimension.time_shift_pk_full_name() {
-                            // TODO: Handle case when named shift is not found
-                        }
-                    } else if let Some(dim_shift_interval) = &shift.interval {
-                        if let Some((dim_key, cts)) =
-                            dimension.calendar_time_shift_for_interval(dim_shift_interval)
-                        {
-                            return Either::Right((dim_key.clone(), cts.clone()));
-                        } else if let Some(calendar_pk) = dimension.time_shift_pk_full_name() {
-                            let mut shift = shift.clone();
-                            shift.interval = Some(dim_shift_interval.inverse());
-                            return Either::Left((calendar_pk, shift.clone()));
-                        }
-                    }
-                }
-                Either::Left((key.clone(), shift.clone()))
-            });
-
+        let (time_shifts, calendar_time_shifts) = self.time_shifts.extract_time_shifts()?;
         let common_time_shifts = TimeShiftState {
             dimensions_shifts: time_shifts,
         };
@@ -75,7 +43,7 @@ impl PhysicalPlanBuilderContext {
         factory.set_count_approx_as_state(self.render_measure_as_state);
         factory.set_ungrouped_measure(self.render_measure_for_ungrouped);
         factory.set_original_sql_pre_aggregations(self.original_sql_pre_aggregations.clone());
-        factory
+        Ok(factory)
     }
 }
 
@@ -117,7 +85,7 @@ impl PhysicalPlanBuilder {
         let from = From::new_from_subselect(source.clone(), ORIGINAL_QUERY.to_string());
         let mut select_builder = SelectBuilder::new(from);
         select_builder.add_count_all(TOTAL_COUNT.to_string());
-        let context_factory = context.make_sql_nodes_factory();
+        let context_factory = context.make_sql_nodes_factory()?;
         Ok(Rc::new(select_builder.build(context_factory)))
     }
 
@@ -142,7 +110,7 @@ impl PhysicalPlanBuilder {
         let mut render_references = HashMap::new();
         let mut measure_references = HashMap::new();
         let mut dimensions_references = HashMap::new();
-        let mut context_factory = context.make_sql_nodes_factory();
+        let mut context_factory = context.make_sql_nodes_factory()?;
         let from = match &logical_plan.source {
             SimpleQuerySource::LogicalJoin(join) => self.process_logical_join(
                 &join,
@@ -400,7 +368,7 @@ impl PhysicalPlanBuilder {
         select_builder.set_offset(logical_plan.offset);
         select_builder.set_ctes(ctes);
 
-        let mut context_factory = context.make_sql_nodes_factory();
+        let mut context_factory = context.make_sql_nodes_factory()?;
         context_factory.set_render_references(render_references);
 
         Ok(Rc::new(select_builder.build(context_factory)))
@@ -735,7 +703,7 @@ impl PhysicalPlanBuilder {
         let mut join_builder =
             JoinBuilder::new_from_subselect(keys_query.clone(), keys_query_alias.clone());
 
-        let mut context_factory = context.make_sql_nodes_factory();
+        let mut context_factory = context.make_sql_nodes_factory()?;
         let primary_keys_dimensions = &aggregate_multiplied_subquery
             .keys_subquery
             .primary_keys_dimensions;
@@ -873,7 +841,7 @@ impl PhysicalPlanBuilder {
             &measure_subquery.dimension_subqueries,
             &mut render_references,
         )?;
-        let mut context_factory = context.make_sql_nodes_factory();
+        let mut context_factory = context.make_sql_nodes_factory()?;
         let mut select_builder = SelectBuilder::new(from);
 
         context_factory.set_ungrouped_measure(true);
@@ -936,7 +904,7 @@ impl PhysicalPlanBuilder {
 
         select_builder.set_distinct();
         select_builder.set_filter(keys_subquery.filter.all_filters());
-        let mut context_factory = context.make_sql_nodes_factory();
+        let mut context_factory = context.make_sql_nodes_factory()?;
         context_factory.set_render_references(render_references);
         let res = Rc::new(select_builder.build(context_factory));
         Ok(res)
@@ -1022,7 +990,7 @@ impl PhysicalPlanBuilder {
             &mut render_references,
         )?;
         let mut select_builder = SelectBuilder::new(from);
-        let mut context_factory = context.make_sql_nodes_factory();
+        let mut context_factory = context.make_sql_nodes_factory()?;
         let args = vec![get_date_range
             .time_dimension
             .clone()
@@ -1181,7 +1149,7 @@ impl PhysicalPlanBuilder {
             on,
         );
 
-        let mut context_factory = context.make_sql_nodes_factory();
+        let mut context_factory = context.make_sql_nodes_factory()?;
         context_factory.set_rolling_window(true);
         let from = From::new_from_join(join_builder.build());
         let references_builder = ReferencesBuilder::new(from.clone());
@@ -1320,7 +1288,7 @@ impl PhysicalPlanBuilder {
             );
         }
 
-        let mut context_factory = context.make_sql_nodes_factory();
+        let mut context_factory = context.make_sql_nodes_factory()?;
         let partition_by = measure_calculation
             .partition_by
             .iter()

rust/cubesqlplanner/cubesqlplanner/src/planner/planners/multi_stage/applied_state.rs

@@ -1,25 +1,14 @@
 use crate::plan::{FilterGroup, FilterItem};
 use crate::planner::filter::FilterOperator;
+use crate::planner::planners::multi_stage::time_shift_state::TimeShiftState;
 use crate::planner::sql_evaluator::{DimensionTimeShift, MeasureTimeShifts, MemberSymbol};
 use crate::planner::{BaseDimension, BaseMember, BaseTimeDimension};
 use cubenativeutils::CubeError;
 use itertools::Itertools;
 use std::cmp::PartialEq;
-use std::collections::HashMap;
 use std::fmt::Debug;
 use std::rc::Rc;
 
-#[derive(Clone, Default, Debug)]
-pub struct TimeShiftState {
-    pub dimensions_shifts: HashMap<String, DimensionTimeShift>,
-}
-
-impl TimeShiftState {
-    pub fn is_empty(&self) -> bool {
-        self.dimensions_shifts.is_empty()
-    }
-}
-
 #[derive(Clone)]
 pub struct MultiStageAppliedState {
     time_dimensions: Vec<Rc<BaseTimeDimension>>,

rust/cubesqlplanner/cubesqlplanner/src/planner/planners/multi_stage/mod.rs

@@ -3,9 +3,11 @@ mod member;
 mod member_query_planner;
 mod multi_stage_query_planner;
 mod query_description;
+mod time_shift_state;
 
 pub use applied_state::*;
 pub use member::*;
 pub use member_query_planner::MultiStageMemberQueryPlanner;
 pub use multi_stage_query_planner::MultiStageQueryPlanner;
 pub use query_description::MultiStageQueryDescription;
+pub use time_shift_state::TimeShiftState;

rust/cubesqlplanner/cubesqlplanner/src/planner/planners/multi_stage/time_shift_state.rs

@@ -0,0 +1,67 @@
+use crate::planner::sql_evaluator::symbols::CalendarDimensionTimeShift;
+use crate::planner::sql_evaluator::DimensionTimeShift;
+use cubenativeutils::CubeError;
+use std::collections::HashMap;
+
+#[derive(Clone, Default, Debug)]
+pub struct TimeShiftState {
+    pub dimensions_shifts: HashMap<String, DimensionTimeShift>,
+}
+
+impl TimeShiftState {
+    pub fn is_empty(&self) -> bool {
+        self.dimensions_shifts.is_empty()
+    }
+
+    pub fn extract_time_shifts(
+        &self,
+    ) -> Result<
+        (
+            HashMap<String, DimensionTimeShift>,
+            HashMap<String, CalendarDimensionTimeShift>,
+        ),
+        CubeError,
+    > {
+        let mut time_shifts = HashMap::new();
+        let mut calendar_time_shifts = HashMap::new();
+
+        for (key, shift) in self.dimensions_shifts.iter() {
+            if let Ok(dimension) = shift.dimension.as_dimension() {
+                // 1. Shift might be referenced by name or by interval
+                // 2. Shift body might be defined in calendar dimension as:
+                //      * sql reference
+                //      * interval + type
+
+                if let Some(dim_shift_name) = &shift.name {
+                    if let Some((dim_key, cts)) =
+                        dimension.calendar_time_shift_for_named_interval(dim_shift_name)
+                    {
+                        calendar_time_shifts.insert(dim_key.clone(), cts.clone());
+                    } else if let Some(_calendar_pk) = dimension.time_shift_pk_full_name() {
+                        return Err(CubeError::internal(format!(
+                            "Time shift with name {} not found for dimension {}",
+                            dim_shift_name,
+                            dimension.full_name()
+                        )));
+                    }
+                } else if let Some(dim_shift_interval) = &shift.interval {
+                    if let Some((dim_key, cts)) =
+                        dimension.calendar_time_shift_for_interval(dim_shift_interval)
+                    {
+                        calendar_time_shifts.insert(dim_key.clone(), cts.clone());
+                    } else if let Some(calendar_pk) = dimension.time_shift_pk_full_name() {
+                        let mut shift = shift.clone();
+                        shift.interval = Some(dim_shift_interval.inverse());
+                        time_shifts.insert(calendar_pk, shift.clone());
+                    } else {
+                        time_shifts.insert(key.clone(), shift.clone());
+                    }
+                }
+            } else {
+                time_shifts.insert(key.clone(), shift.clone());
+            }
+        }
+
+        Ok((time_shifts, calendar_time_shifts))
+    }
+}