simplify builder by moving extract_time_shifts to time_shift_state

Author: KSDaemon

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

@@ -33,46 +33,7 @@ impl PhysicalPlanBuilderContext {
     pub fn make_sql_nodes_factory(&self) -> Result<SqlNodesFactory, CubeError> {
         let mut factory = SqlNodesFactory::new();
 
-        let mut time_shifts = HashMap::new();
-        let mut calendar_time_shifts = HashMap::new();
-
-        for (key, shift) in self.time_shifts.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());
-            }
-        }
-
+        let (time_shifts, calendar_time_shifts) = self.time_shifts.extract_time_shifts()?;
         let common_time_shifts = TimeShiftState {
             dimensions_shifts: time_shifts,
         };

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))
+    }
+}