fix(tesseract): compose grain.include with rolling_window

packages/cubejs-schema-compiler/test/integration/postgres/multi-stage-grain-rolling-window.test.ts

@@ -0,0 +1,170 @@
+import { getEnv } from '@cubejs-backend/shared';
+import { prepareYamlCompiler } from '../../unit/PrepareCompiler';
+import { dbRunner } from './PostgresDBRunner';
+
+// A multi-stage measure that declares both `grain` and `rolling_window`.
+// The grain says at what grain the value inside a bucket is computed; the
+// window says which rows land in the bucket. Both have to hold at once.
+//
+// Inline data: two equally weighted rows per day, so the weighted daily
+// factor equals the row factor — 1.10, 1.20, 0.50. Linked over the three
+// days that is exactly -0.34, and every window below is wide enough to
+// cover all three, so at any grain coarser than a day the windowed measure
+// has to agree with the plain one.
+//
+// The measures round to six digits so the assertions compare exact strings
+// rather than IEEE noise from EXP/LN.
+describe('Multi-Stage grain with rolling window', () => {
+  jest.setTimeout(200000);
+
+  const { compiler, joinGraph, cubeEvaluator } = prepareYamlCompiler(`
+cubes:
+  - name: returns
+    sql: >
+      SELECT '2024-01-01'::date as DAY, 'A' as SECURITY, 10.0 as IRR, 100.0 as WEIGHT
+      union all
+      SELECT '2024-01-01'::date, 'B', 10.0, 100.0
+      union all
+      SELECT '2024-01-02'::date, 'A', 20.0, 100.0
+      union all
+      SELECT '2024-01-02'::date, 'B', 20.0, 100.0
+      union all
+      SELECT '2024-01-03'::date, 'A', -50.0, 100.0
+      union all
+      SELECT '2024-01-03'::date, 'B', -50.0, 100.0
+
+    dimensions:
+      - name: day
+        sql: DAY
+        type: time
+
+      - name: security
+        sql: SECURITY
+        type: string
+
+      - name: irr
+        sql: IRR
+        type: number
+
+      - name: weight
+        sql: WEIGHT
+        type: number
+
+    measures:
+      - name: weight_sum
+        sql: "{CUBE.weight}"
+        type: sum
+
+      - name: irr_weight_sum
+        sql: "{CUBE.irr} * {CUBE.weight}"
+        type: sum
+
+      - name: daily_irr_weight
+        multi_stage: true
+        sql: "{CUBE.irr_weight_sum}"
+        type: number
+
+      - name: daily_weight
+        multi_stage: true
+        sql: "{CUBE.weight_sum}"
+        type: number
+
+      - name: log_return_sum
+        multi_stage: true
+        sql: "LN(1.0 + ({CUBE.daily_irr_weight} / NULLIF({CUBE.daily_weight}, 0)) / 100.0)"
+        type: sum
+        grain:
+          include:
+            - returns.day
+
+      - name: twr
+        multi_stage: true
+        sql: "ROUND((EXP({CUBE.log_return_sum}) - 1)::numeric, 6)"
+        type: number
+
+      - name: log_return_sum_ytd
+        multi_stage: true
+        sql: "LN(1.0 + ({CUBE.daily_irr_weight} / NULLIF({CUBE.daily_weight}, 0)) / 100.0)"
+        type: sum
+        grain:
+          include:
+            - returns.day
+        rolling_window:
+          type: to_date
+          granularity: year
+
+      - name: twr_ytd
+        multi_stage: true
+        sql: "ROUND((EXP({CUBE.log_return_sum_ytd}) - 1)::numeric, 6)"
+        type: number
+
+      - name: log_return_sum_1y
+        multi_stage: true
+        sql: "LN(1.0 + ({CUBE.daily_irr_weight} / NULLIF({CUBE.daily_weight}, 0)) / 100.0)"
+        type: sum
+        grain:
+          include:
+            - returns.day
+        rolling_window:
+          trailing: 1 year
+          offset: end
+
+      - name: twr_1y
+        multi_stage: true
+        sql: "ROUND((EXP({CUBE.log_return_sum_1y}) - 1)::numeric, 6)"
+        type: number
+    `);
+
+  if (getEnv('nativeSqlPlanner')) {
+    it('day granularity: the window accumulates over the declared grain', async () => dbRunner.runQueryTest({
+      measures: ['returns.twr', 'returns.twr_ytd', 'returns.twr_1y'],
+      timeDimensions: [{
+        dimension: 'returns.day',
+        granularity: 'day',
+        dateRange: ['2024-01-01', '2024-01-03'],
+      }],
+      order: [{ id: 'returns.day' }],
+      timezone: 'UTC',
+    }, [
+      // Per-day factors 1.10 / 1.20 / 0.50; the windows link them up.
+      { returns__day_day: '2024-01-01T00:00:00.000Z', returns__twr: '0.100000', returns__twr_ytd: '0.100000', returns__twr_1y: '0.100000' },
+      { returns__day_day: '2024-01-02T00:00:00.000Z', returns__twr: '0.200000', returns__twr_ytd: '0.320000', returns__twr_1y: '0.320000' },
+      { returns__day_day: '2024-01-03T00:00:00.000Z', returns__twr: '-0.500000', returns__twr_ytd: '-0.340000', returns__twr_1y: '-0.340000' },
+    ], { joinGraph, cubeEvaluator, compiler }));
+
+    it('month granularity: query grain does not replace the declared grain', async () => dbRunner.runQueryTest({
+      measures: ['returns.twr', 'returns.twr_ytd', 'returns.twr_1y'],
+      timeDimensions: [{
+        dimension: 'returns.day',
+        granularity: 'month',
+        dateRange: ['2024-01-01', '2024-01-31'],
+      }],
+      order: [{ id: 'returns.day' }],
+      timezone: 'UTC',
+    }, [
+      { returns__day_month: '2024-01-01T00:00:00.000Z', returns__twr: '-0.340000', returns__twr_ytd: '-0.340000', returns__twr_1y: '-0.340000' },
+    ], { joinGraph, cubeEvaluator, compiler }));
+
+    it('no time dimension: the window is one bucket over the whole range', async () => dbRunner.runQueryTest({
+      measures: ['returns.twr', 'returns.twr_ytd', 'returns.twr_1y'],
+      timezone: 'UTC',
+    }, [
+      { returns__twr: '-0.340000', returns__twr_ytd: '-0.340000', returns__twr_1y: '-0.340000' },
+    ], { joinGraph, cubeEvaluator, compiler }));
+
+    it('non-time dimension: the grain survives a plain group by', async () => dbRunner.runQueryTest({
+      measures: ['returns.twr', 'returns.twr_ytd', 'returns.twr_1y'],
+      dimensions: ['returns.security'],
+      order: [{ id: 'returns.security' }],
+      timezone: 'UTC',
+    }, [
+      { returns__security: 'A', returns__twr: '-0.340000', returns__twr_ytd: '-0.340000', returns__twr_1y: '-0.340000' },
+      { returns__security: 'B', returns__twr: '-0.340000', returns__twr_ytd: '-0.340000', returns__twr_1y: '-0.340000' },
+    ], { joinGraph, cubeEvaluator, compiler }));
+  } else {
+    test.skip('day granularity: the window accumulates over the declared grain', () => { expect(1).toBe(1); });
+    test.skip('month granularity: query grain does not replace the declared grain', () => { expect(1).toBe(1); });
+    test.skip('no time dimension: the window is one bucket over the whole range', () => { expect(1).toBe(1); });
+    test.skip('non-time dimension: the grain survives a plain group by', () => { expect(1).toBe(1); });
+  }
+});

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

@@ -799,6 +799,12 @@ impl MultiStageQueryPlanner {
                     }
                 }
 
+                let grain = measure
+                    .multi_stage()
+                    .map(|ms| ms.grain.clone())
+                    .unwrap_or_default();
+                let grain_include = grain.include.clone().unwrap_or_default();
+
                 let ungrouped = measure.is_rolling_window() && !measure.is_additive();
 
                 let mut time_dimensions = self
@@ -828,8 +834,14 @@ impl MultiStageQueryPlanner {
                             scope,
                         )?
                     } else {
+                        // Without a time dimension there is no series to walk,
+                        // so the window collapses to a single bucket: no frame
+                        // to build, and no outer stage to carry the aggregation.
+                        // What is left is the measure's own multi-stage
+                        // definition — aggregation and grain included — over the
+                        // base state prepared above.
                         self.make_queries_descriptions(
-                            base_member,
+                            MemberSymbol::new_measure(transforms::strip_rolling_window(&measure)),
                             base_state,
                             descriptions,
                             resolved_multi_stage_dimensions,
@@ -853,12 +865,59 @@ impl MultiStageQueryPlanner {
                     GranularityHelper::find_dimension_with_min_granularity(&time_dimensions)?;
                 let time_dimension = MemberSymbol::new_time_dimension(time_dimension);
 
+                // Of the grain keys only `include` reaches the window assembly.
+                // It extends the grain the values inside a bucket are computed
+                // at, which the base CTE carries anyway. `exclude` and
+                // `keep_only` narrow the grain the value is *reported* at, and a
+                // narrowed value has to be broadcast back onto the query grid;
+                // the rolling window node has no side enumerating that grid, so
+                // there is nothing to broadcast from.
+                //
+                // What is rejected is the narrowing actually happening, not the
+                // keys being declared: `exclude` of a member this grain does not
+                // carry subtracts nothing, and `keep_only` listing everything the
+                // query groups by intersects to the same list. Such a key costs
+                // the query nothing, and the value is the one the measure would
+                // have without it. `partition_filter` only ever removes, so a
+                // shorter list is exactly the case that has no answer here.
+                //
+                // The check sits below the branch above on purpose. Without a
+                // time dimension the window has no frame and the measure is
+                // planned through the ordinary multi-stage path, which narrows
+                // the grain and broadcasts it back the usual way.
+                let narrows = Self::partition_filter(state.dimensions(), &grain).len()
+                    != state.dimensions().len()
+                    || Self::partition_filter(state.time_dimensions(), &grain).len()
+                        != state.time_dimensions().len();
+                if narrows {
+                    return Err(CubeError::user(format!(
+                        "Measure {} narrows the grain of this query through `grain.exclude` / \
+                         `reduce_by` or `grain.keep_only` / `group_by` while also declaring a \
+                         `rolling_window` over {}, which is not supported. Drop the narrowing \
+                         keys, drop the window, or query the measure without a time dimension.",
+                        member.full_name(),
+                        time_dimension.full_name(),
+                    )));
+                }
+
                 let (base_rolling_state, base_time_dimension) = self.make_rolling_base_state(
                     time_dimension.clone(),
                     &rolling_window,
                     state.clone(),
                 )?;
 
+                // `grain.include` extends the grain the values inside the window
+                // are computed at. The frame still keys off the base time
+                // dimension, so the extension only splits rows the outer
+                // aggregation merges back together.
+                let base_rolling_state = if grain_include.is_empty() {
+                    base_rolling_state
+                } else {
+                    let mut extended = base_rolling_state.as_ref().clone();
+                    extended.add_dimensions(grain_include);
+                    Rc::new(extended)
+                };
+
                 let time_series =
                     self.add_time_series(time_dimension.clone(), state.clone(), descriptions)?;
 

rust/cube/cubesqlplanner/cubesqlplanner/src/planner/query_properties.rs

@@ -711,13 +711,24 @@ impl QueryProperties {
     }
 
     /// Append `dimensions` to the existing list, deduplicating by
-    /// reference-chain-resolved full name.
+    /// reference-chain-resolved full name. A dimension the grain already
+    /// carries as a time dimension is dropped rather than appended: a time
+    /// dimension's full name pins its granularity, so an entry that matches
+    /// one would render the very same column under the very same alias.
     pub fn add_dimensions(&mut self, dimensions: Vec<Rc<MemberSymbol>>) {
+        let time_dimension_names = self
+            .time_dimensions
+            .iter()
+            .map(|d| d.clone().resolve_reference_chain().full_name())
+            .collect::<HashSet<_>>();
+        let added = dimensions.into_iter().filter(|d| {
+            !time_dimension_names.contains(&d.clone().resolve_reference_chain().full_name())
+        });
         self.dimensions = self
             .dimensions
             .iter()
             .cloned()
-            .chain(dimensions.into_iter())
+            .chain(added)
             .unique_by(|d| d.clone().resolve_reference_chain().full_name())
             .collect_vec();
         self.invalidate_join_groups_cache();

rust/cube/cubesqlplanner/cubesqlplanner/src/planner/symbols/transforms/unroll_rolling.rs

@@ -44,3 +44,17 @@ pub fn unroll_rolling(measure: &MeasureSymbol) -> Rc<MeasureSymbol> {
         render_modifier: None,
     })
 }
+
+/// Returns a copy of the measure with the rolling window dropped and
+/// the rest of its definition — aggregation kind, multi-stage grain —
+/// left in place. What remains is the measure as it would have been
+/// declared without a window, which is what a window that resolves to
+/// a single bucket computes.
+pub fn strip_rolling_window(measure: &MeasureSymbol) -> Rc<MeasureSymbol> {
+    let mut result = measure.clone();
+    result.rolling_window = None;
+    // The render modifier is stamped against the rolling form and has
+    // nothing to modify once the window is gone.
+    result.render_modifier = None;
+    Rc::new(result)
+}