fix(tesseract): keep time_shift when a pre-aggregation serves the query

rust/cube/cubesqlplanner/cubesqlplanner/src/logical_plan/optimizers/pre_aggregation/optimizer.rs

@@ -3,6 +3,7 @@ use super::*;
 use crate::logical_plan::visitor::{LogicalPlanRewriter, NodeRewriteResult};
 use crate::logical_plan::*;
 use crate::planner::collectors::{collect_cube_names_from_symbols, has_multi_stage_members};
+use crate::planner::filter::typed_filter::resolve_base_symbol;
 use crate::planner::filter::FilterItem;
 use crate::planner::filter::FilterOp;
 use crate::planner::join_hints::JoinHints;
@@ -140,9 +141,13 @@ impl PreAggregationOptimizer {
             let external = pre_aggregation.external.unwrap_or(false);
             let date_range =
                 Self::extract_date_range(&query.filter(), &self.query_tools, time_shifts, external);
-            if let Some(rewritten) =
-                self.try_rewrite_simple_query(query, pre_aggregation, date_range, is_user_query)?
-            {
+            if let Some(rewritten) = self.try_rewrite_simple_query(
+                query,
+                pre_aggregation,
+                date_range,
+                is_user_query,
+                time_shifts,
+            )? {
                 return Ok(Some(rewritten));
             }
         }
@@ -156,6 +161,7 @@ impl PreAggregationOptimizer {
         pre_aggregation: &Rc<CompiledPreAggregation>,
         date_range: Option<(String, String)>,
         is_user_query: bool,
+        time_shifts: &TimeShiftState,
     ) -> Result<Option<Rc<Query>>, CubeError> {
         // Row identity for an ungrouped read is judged against the join this
         // very node will render, taken from the node itself rather than
@@ -174,6 +180,14 @@ impl PreAggregationOptimizer {
             pre_aggregation,
             row_grain,
         )? {
+            if !Self::can_carry_time_shifts(
+                pre_aggregation,
+                &matched_measures,
+                &Self::read_member_names(&query.schema(), &query.filter()),
+                time_shifts,
+            ) {
+                return Ok(None);
+            }
             let source =
                 self.make_pre_aggregation_source(pre_aggregation, &matched_measures, date_range)?;
             let new_query = Query::builder()
@@ -477,6 +491,88 @@ impl PreAggregationOptimizer {
         }
     }
 
+    // A stored member is shifted by offsetting its column as a whole, which
+    // only reproduces the shifted values when the shift can be attributed to
+    // that column. A column built from several members of which just some are
+    // shifted has no such offset — moving it would carry along rows the shift
+    // must leave in place — and the lookup cannot attribute a shift to it
+    // either, so the two agree: whenever a shift is involved but cannot be
+    // attributed, the pre-aggregation cannot serve the shifted leaf.
+    //
+    // Grouping members — time dimensions, dimensions and segments — are
+    // substituted by column, so a pre-aggregation can only serve a shifted
+    // leaf when every stored member a shift reaches is one whose column can
+    // carry that shift. `shift_for_substituted_column` decides that, and the
+    // rendering node asks it too, so a member admitted here is one that will
+    // actually be offset.
+    //
+    // A measure column holds an aggregate, and a shift changes which rows
+    // feed it rather than the value itself, so no offset applies at all. A
+    // stored measure reading a shifted member is therefore always unusable,
+    // however cleanly the shift could be attributed to it. Only the measures
+    // matching consumed are examined, since the rest are never read.
+    // Resolved names of every member the query reads, so a stored member no
+    // one reads cannot decide anything.
+    fn read_member_names(schema: &LogicalSchema, filter: &LogicalFilter) -> HashSet<String> {
+        let mut symbols: Vec<Rc<MemberSymbol>> = schema
+            .dimensions
+            .iter()
+            .chain(schema.time_dimensions.iter())
+            .chain(schema.measures.iter())
+            .cloned()
+            .collect();
+        for item in filter
+            .dimensions_filters
+            .iter()
+            .chain(filter.time_dimensions_filters.iter())
+            .chain(filter.segments.iter())
+        {
+            item.find_all_member_evaluators(&mut symbols);
+        }
+        symbols
+            .into_iter()
+            .map(|symbol| {
+                resolve_base_symbol(&symbol)
+                    .resolve_reference_chain()
+                    .full_name()
+            })
+            .collect()
+    }
+
+    fn can_carry_time_shifts(
+        pre_aggregation: &CompiledPreAggregation,
+        matched_measures: &HashSet<String>,
+        read_members: &HashSet<String>,
+        time_shifts: &TimeShiftState,
+    ) -> bool {
+        if time_shifts.is_empty() {
+            return true;
+        }
+        let is_read = |member: &Rc<MemberSymbol>| {
+            read_members.contains(
+                &resolve_base_symbol(member)
+                    .resolve_reference_chain()
+                    .full_name(),
+            )
+        };
+        let grouping_members_carry_shift = pre_aggregation
+            .time_dimensions
+            .iter()
+            .chain(pre_aggregation.dimensions.iter())
+            .chain(pre_aggregation.segments.iter())
+            .filter(|member| is_read(member))
+            .all(|member| {
+                !time_shifts.has_shift_under(member)
+                    || time_shifts.shift_for_substituted_column(member).is_some()
+            });
+        grouping_members_carry_shift
+            && pre_aggregation
+                .measures
+                .iter()
+                .filter(|measure| matched_measures.contains(&measure.full_name()))
+                .all(|measure| !time_shifts.has_shift_under(measure))
+    }
+
     fn extract_date_range(
         filter: &LogicalFilter,
         query_tools: &Rc<State>,
@@ -496,8 +592,7 @@ impl PreAggregationOptimizer {
                         // Apply time shift for this dimension if present.
                         // SQL renders `column + interval`, so actual data range is `date - interval`.
                         if let Some(interval) = time_shifts
-                            .dimensions_shifts
-                            .get(&base_filter.member_name())
+                            .get_for_symbol(base_filter.raw_member_evaluator_ref())
                             .and_then(|s| s.interval.as_ref())
                         {
                             let tz = query_tools.timezone();

rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/filter/base_filter.rs

@@ -27,8 +27,7 @@ impl ToSql for BaseFilter {
             {
                 let time_shift = visitor
                     .time_shifts()
-                    .dimensions_shifts
-                    .get(&symbol_to_match.full_name())
+                    .get_for_symbol(&symbol_to_match)
                     .and_then(|shift| shift.interval.as_ref());
                 return self.typed_filter().to_sql_for_filter_params(
                     filter_params_item,

rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/factory.rs

@@ -287,7 +287,11 @@ impl SqlNodesFactory {
         };
 
         let input = if !self.time_shifts.is_empty() {
-            TimeShiftSqlNode::new(self.time_shifts.clone(), input)
+            TimeShiftSqlNode::new(
+                self.time_shifts.clone(),
+                self.pre_aggregation_dimensions_references.clone(),
+                input,
+            )
         } else {
             input
         };

rust/cube/cubesqlplanner/cubesqlplanner/src/physical_plan/sql_nodes/time_shift.rs

@@ -1,4 +1,5 @@
 use super::SqlNode;
+use crate::physical_plan::sql_nodes::render_references::RenderReferences;
 use crate::physical_plan::SqlEvaluatorVisitor;
 use crate::planner::planners::multi_stage::TimeShiftState;
 use crate::planner::query_tools::QueryTools;
@@ -11,14 +12,27 @@ use std::rc::Rc;
 /// Applies a per-dimension time shift to time dimensions whose
 /// full name is in `shifts`, by rendering the dimension expression
 /// shifted by the configured interval.
+///
+/// `substituted` names the dimensions rendered as a stored column instead
+/// of being evaluated. Their SQL is never expanded, so the shift cannot be
+/// picked up further down and has to be applied to the column itself.
 pub struct TimeShiftSqlNode {
     shifts: TimeShiftState,
+    substituted: RenderReferences,
     input: Rc<dyn SqlNode>,
 }
 
 impl TimeShiftSqlNode {
-    pub fn new(shifts: TimeShiftState, input: Rc<dyn SqlNode>) -> Rc<Self> {
-        Rc::new(Self { shifts, input })
+    pub fn new(
+        shifts: TimeShiftState,
+        substituted: RenderReferences,
+        input: Rc<dyn SqlNode>,
+    ) -> Rc<Self> {
+        Rc::new(Self {
+            shifts,
+            substituted,
+            input,
+        })
     }
 
     pub fn input(&self) -> &Rc<dyn SqlNode> {
@@ -38,8 +52,34 @@ impl SqlNode for TimeShiftSqlNode {
         let res = match node.as_ref() {
             MemberSymbol::Dimension(ev) => {
                 if !ev.is_reference() && ev.is_time() {
-                    if let Some(shift) = self.shifts.dimensions_shifts.get(&ev.full_name()) {
-                        let shift = shift.interval.clone().unwrap().to_sql();
+                    // The first probe is by exact name on purpose: a dimension
+                    // that gets evaluated has its shift applied when the
+                    // recursion reaches the owned member it wraps, and matching
+                    // it here as well would add the interval twice. Only a
+                    // substituted dimension, which is never expanded, resolves
+                    // through the chain.
+                    let shift = self
+                        .shifts
+                        .dimensions_shifts
+                        .get(&ev.full_name())
+                        .or_else(|| {
+                            if self.substituted.contains_key(&ev.full_name()) {
+                                self.shifts.shift_for_substituted_column(node)
+                            } else {
+                                None
+                            }
+                        });
+                    if let Some(shift) = shift {
+                        let shift = shift
+                            .interval
+                            .as_ref()
+                            .ok_or_else(|| {
+                                CubeError::internal(format!(
+                                    "Time shift for dimension {} has no interval",
+                                    ev.full_name()
+                                ))
+                            })?
+                            .to_sql();
                         let inner_visitor = visitor.with_arg_needs_paren_safe(false);
                         let input = self.input.to_sql(
                             &inner_visitor,

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

@@ -1,7 +1,11 @@
+use crate::planner::collectors::find_owned_by_cube_child;
+use crate::planner::filter::typed_filter::resolve_base_symbol;
 use crate::planner::symbols::CalendarDimensionTimeShift;
+use crate::planner::symbols::MemberSymbol;
 use crate::planner::DimensionTimeShift;
 use cubenativeutils::CubeError;
 use std::collections::HashMap;
+use std::rc::Rc;
 
 /// Per-dimension time-shift accumulator used during multi-stage
 /// planning. Keyed by dimension full name; aggregates the shifts
@@ -16,6 +20,53 @@ impl TimeShiftState {
         self.dimensions_shifts.is_empty()
     }
 
+    /// Looks up the shift for a symbol that may still be wrapped in a
+    /// `TimeDimension`, be a reference to the shifted member, or wrap it in
+    /// its own SQL. Keys are built either from the chain-resolved dimension
+    /// or, for dimension-specific shifts, from the owned member the declared
+    /// dimension wraps, so both forms are probed.
+    pub fn get_for_symbol(&self, symbol: &Rc<MemberSymbol>) -> Option<&DimensionTimeShift> {
+        let resolved = resolve_base_symbol(symbol).resolve_reference_chain();
+        if let Some(shift) = self.dimensions_shifts.get(&resolved.full_name()) {
+            return Some(shift);
+        }
+        let owned = find_owned_by_cube_child(&resolved).ok()?;
+        self.dimensions_shifts.get(&owned.full_name())
+    }
+
+    /// The shift a stored column standing for this member can carry.
+    ///
+    /// A column is shifted by offsetting it, which only stands in for the
+    /// shifted member when the member is a time dimension evaluated in place:
+    /// a reference is rendered through to what it points at, and a non-time
+    /// member has no meaning under an interval. Both the gate that admits a
+    /// pre-aggregation and the node that renders from one ask this, so the two
+    /// cannot come to different conclusions.
+    pub fn shift_for_substituted_column(
+        &self,
+        symbol: &Rc<MemberSymbol>,
+    ) -> Option<&DimensionTimeShift> {
+        let dimension = resolve_base_symbol(symbol).as_dimension().ok()?;
+        if dimension.is_reference() || !dimension.is_time() {
+            return None;
+        }
+        self.get_for_symbol(symbol)
+    }
+
+    /// True when the symbol itself, or any member it is built from, is
+    /// shifted. Unlike `get_for_symbol` this answers whether a shift is
+    /// involved at all, not whether one can be attributed to the symbol.
+    pub fn has_shift_under(&self, symbol: &Rc<MemberSymbol>) -> bool {
+        let symbol = resolve_base_symbol(symbol);
+        if self.dimensions_shifts.contains_key(&symbol.full_name()) {
+            return true;
+        }
+        symbol
+            .get_dependencies()
+            .iter()
+            .any(|dep| self.has_shift_under(dep))
+    }
+
     /// Splits the accumulated shifts into two maps: regular
     /// `DimensionTimeShift`s applied at render time, and
     /// `CalendarDimensionTimeShift`s that come from a calendar

rust/cube/cubesqlplanner/cubesqlplanner/src/test_fixtures/schemas/yaml_files/common/integration_derived_time_dim_shift_pre_agg.yaml

@@ -0,0 +1,57 @@
+cubes:
+    - name: pa_customers
+      sql: "SELECT * FROM pa_customers"
+      joins:
+          - name: pa_returns
+            relationship: one_to_many
+            sql: "{pa_customers}.id = {pa_returns.customer_id}"
+      dimensions:
+          - name: id
+            type: number
+            sql: id
+            primary_key: true
+
+          # Time dimension derived from another cube's time dimension:
+          # not owned by its cube and not a plain reference, so its shift
+          # is keyed by the owned member it wraps.
+          - name: return_day
+            type: time
+            sql: "DATE_TRUNC('day', {pa_returns.created_at})"
+      measures:
+          - name: total_value
+            type: sum
+            sql: lifetime_value
+
+          - name: total_value_prev_month
+            type: number
+            sql: "{CUBE.total_value}"
+            multi_stage: true
+            time_shift:
+                - interval: "1 month"
+                  type: prior
+                  timeDimension: pa_customers.return_day
+
+      pre_aggregations:
+          - name: value_by_return_day_month
+            type: rollup
+            measures:
+                - total_value
+            time_dimension: pa_customers.return_day
+            granularity: month
+
+    - name: pa_returns
+      sql: "SELECT * FROM pa_returns"
+      dimensions:
+          - name: id
+            type: number
+            sql: id
+            primary_key: true
+          - name: customer_id
+            type: number
+            sql: customer_id
+          - name: created_at
+            type: time
+            sql: created_at
+      measures:
+          - name: count
+            type: count