fix(tesseract): reduce_by/grain matches time dimensions regardless of granularity (cube-js#11188)

Author: waralexrom

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

@@ -542,15 +542,15 @@ impl MultiStageMemberQueryPlanner {
         let dimensions = if let Some(exclude) = &grain.exclude {
             dimensions
                 .into_iter()
-                .filter(|d| !exclude.iter().any(|m| d.has_member_in_reference_chain(m)))
+                .filter(|d| !exclude.iter().any(|m| d.matches_grain_reference(m)))
                 .collect_vec()
         } else {
             dimensions
         };
         let dimensions = if let Some(keep_only) = &grain.keep_only {
             dimensions
                 .into_iter()
-                .filter(|d| keep_only.iter().any(|m| d.has_member_in_reference_chain(m)))
+                .filter(|d| keep_only.iter().any(|m| d.matches_grain_reference(m)))
                 .collect_vec()
         } else {
             dimensions

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

@@ -294,15 +294,15 @@ impl MultiStageQueryPlanner {
     ) -> Vec<Rc<MemberSymbol>> {
         let dims: Vec<Rc<MemberSymbol>> = if let Some(exclude) = &grain.exclude {
             dims.iter()
-                .filter(|d| !exclude.iter().any(|m| d.has_member_in_reference_chain(m)))
+                .filter(|d| !exclude.iter().any(|m| d.matches_grain_reference(m)))
                 .cloned()
                 .collect()
         } else {
             dims.clone()
         };
         if let Some(keep_only) = &grain.keep_only {
             dims.into_iter()
-                .filter(|d| keep_only.iter().any(|m| d.has_member_in_reference_chain(m)))
+                .filter(|d| keep_only.iter().any(|m| d.matches_grain_reference(m)))
                 .collect()
         } else {
             dims

rust/cube/cubesqlplanner/cubesqlplanner/src/planner/symbols/member_symbol.rs

@@ -236,6 +236,23 @@ impl MemberSymbol {
         false
     }
 
+    /// Whether this projected member is targeted by a grain reference
+    /// (`reduce_by` / `group_by`). Behaves like `has_member_in_reference_chain`,
+    /// but a time dimension also matches through its underlying base dimension:
+    /// a projected time dimension carries a granularity suffix in its full name
+    /// (`created_at_month`) while grain references point at the bare dimension
+    /// (`created_at`), so a grain reference removes the time dimension from the
+    /// grain regardless of the granularity it is queried at.
+    pub fn matches_grain_reference(&self, member: &Rc<MemberSymbol>) -> bool {
+        if self.has_member_in_reference_chain(member) {
+            return true;
+        }
+        if let Self::TimeDimension(td) = self {
+            return td.base_symbol().has_member_in_reference_chain(member);
+        }
+        false
+    }
+
     /// Returns a copy of this symbol with the path reduced to just the owning cube,
     /// stripping any join chain prefix (e.g. from views or cross-cube references).
     pub fn with_stripped_join_prefix(&self) -> Rc<Self> {

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

@@ -186,6 +186,17 @@ cubes:
             reduce_by:
                 - orders.status
 
+          # reduce_by targets the time dimension itself (bare, no
+          # granularity). When the query selects created_at at some
+          # granularity, that granular dimension must be dropped from the
+          # partition regardless of its granularity suffix (CORE-550).
+          - name: amount_reduce_time
+            type: sum
+            sql: "{CUBE.total_amount}"
+            multi_stage: true
+            reduce_by:
+                - orders.created_at
+
           - name: amount_reduce_all
             type: sum
             sql: "{CUBE.total_amount}"
@@ -290,6 +301,20 @@ cubes:
                 - sql: "{CUBE.total_amount}"
                   dir: desc
 
+          # Rank reduced by the time dimension (bare) while ordered by a
+          # measure. With created_at selected at a granularity the partition
+          # must be "all outer dims except created_at" (CORE-550); the buggy
+          # planner left created_at_<gran> in PARTITION BY so every partition
+          # held one row and all ranks collapsed to 1.
+          - name: amount_rank_reduce_time
+            type: rank
+            multi_stage: true
+            reduce_by:
+                - orders.created_at
+            order_by:
+                - sql: "{CUBE.total_amount}"
+                  dir: desc
+
           # Rank by the raw time dimension itself (order_by is a
           # dimension, not a measure) reduced by created_at — partition
           # is "all outer dims except created_at". First link in the
@@ -533,3 +558,5 @@ views:
                 - total_amount
                 - amount_prev_month
                 - mom_growth
+                - amount_reduce_time
+                - amount_rank_reduce_time

rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/rank.rs

@@ -9,6 +9,66 @@ fn create_context() -> TestContext {
 
 const SEED: &str = "integration_multi_stage_tables.sql";
 
+/// PARTITION BY clause bodies (text between `PARTITION BY` and the following
+/// `ORDER BY`) of every window function in `sql`.
+fn partition_by_clauses(sql: &str) -> Vec<String> {
+    sql.match_indices("PARTITION BY")
+        .map(|(i, _)| {
+            let rest = &sql[i + "PARTITION BY".len()..];
+            let end = rest.find("ORDER BY").unwrap_or(rest.len());
+            rest[..end].to_string()
+        })
+        .collect()
+}
+
+// CORE-550: a rank reduced by a bare time dimension, ordered by a measure,
+// queried with that time dimension at a granularity. Partition must be
+// [status] alone; the buggy planner kept created_at_month in PARTITION BY so
+// each (status, month) partition held one row and every rank was 1. Correct
+// ranks per status by descending monthly total (completed Mar>Feb>Jan = 1,2,3;
+// pending 1,2,3; cancelled Mar=1, Jan/Feb tie at 2).
+#[tokio::test(flavor = "multi_thread")]
+async fn test_rank_reduce_by_time_dimension() {
+    let ctx = create_context();
+
+    let query = indoc! {r#"
+        measures:
+          - orders.total_amount
+          - orders.amount_rank_reduce_time
+        dimensions:
+          - orders.status
+        time_dimensions:
+          - dimension: orders.created_at
+            granularity: month
+            dateRange:
+              - "2024-01-01"
+              - "2024-03-31"
+        order:
+          - id: orders.status
+          - id: orders.created_at
+    "#};
+
+    let sql = ctx.build_sql(query).unwrap();
+    for clause in partition_by_clauses(&sql) {
+        assert!(
+            !clause.contains("created_at"),
+            "reduce_by(created_at) must drop the time dimension from PARTITION BY,\n\
+             partition clause was: {}\nfull SQL:\n{}",
+            clause,
+            sql,
+        );
+        assert!(
+            clause.contains("status"),
+            "status must remain in PARTITION BY, partition clause was: {}",
+            clause,
+        );
+    }
+
+    if let Some(result) = ctx.try_execute_pg(query, SEED).await {
+        insta::assert_snapshot!(result);
+    }
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_basic_rank() {
     let ctx = create_context();

rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/reduce_by.rs

@@ -27,6 +27,18 @@ fn assert_no_window(sql: &str) {
     );
 }
 
+/// PARTITION BY clause bodies (text between `PARTITION BY` and the following
+/// `ORDER BY`) of every window function in `sql`.
+fn partition_by_clauses(sql: &str) -> Vec<String> {
+    sql.match_indices("PARTITION BY")
+        .map(|(i, _)| {
+            let rest = &sql[i + "PARTITION BY".len()..];
+            let end = rest.find("ORDER BY").unwrap_or(rest.len());
+            rest[..end].to_string()
+        })
+        .collect()
+}
+
 // add_group_by + reduce_by together: leaf grain extends with customer_id
 // while partition grain shrinks by removing category. Three distinct grains:
 //   leaf       = (status, category, customer_id)  ← per-customer sum(amount)
@@ -317,6 +329,57 @@ async fn test_reduce_by_count_distinct() {
     }
 }
 
+// CORE-550: reduce_by a bare time dimension must drop it from the window's
+// PARTITION BY even when the query selects that dimension at a granularity
+// (auto-suffixed `_month`). The buggy planner matched reduce_by references by
+// full_name, so `orders.created_at` never matched the granular
+// `orders.created_at_month`; created_at stayed in the partition, the window
+// collapsed to plain group-by and the "reduced" value equalled total_amount.
+// Correct: partition = [status] → per-status total broadcast across months
+// (completed=1400, pending=650, cancelled=200).
+#[tokio::test(flavor = "multi_thread")]
+async fn test_reduce_by_time_dimension() {
+    let ctx = create_context();
+
+    let query = indoc! {r#"
+        measures:
+          - orders.total_amount
+          - orders.amount_reduce_time
+        dimensions:
+          - orders.status
+        time_dimensions:
+          - dimension: orders.created_at
+            granularity: month
+            dateRange:
+              - "2024-01-01"
+              - "2024-03-31"
+        order:
+          - id: orders.status
+          - id: orders.created_at
+    "#};
+
+    let sql = ctx.build_sql(query).unwrap();
+    assert_uses_window(&sql);
+    for clause in partition_by_clauses(&sql) {
+        assert!(
+            !clause.contains("created_at"),
+            "reduce_by(created_at) must drop the time dimension from PARTITION BY,\n\
+             partition clause was: {}\nfull SQL:\n{}",
+            clause,
+            sql,
+        );
+        assert!(
+            clause.contains("status"),
+            "status must remain in PARTITION BY, partition clause was: {}",
+            clause,
+        );
+    }
+
+    if let Some(result) = ctx.try_execute_pg(query, SEED).await {
+        insta::assert_snapshot!(result);
+    }
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_reduce_by_with_time() {
     let ctx = create_context();

rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/snapshots/cubesqlplanner__tests__integration__multi_stage__rank__rank_reduce_by_time_dimension.snap

@@ -0,0 +1,15 @@
+---
+source: cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/rank.rs
+expression: result
+---
+orders__status | orders__created_at_month | orders__total_amount | orders__amount_rank_reduce_time
+---------------+--------------------------+----------------------+--------------------------------
+cancelled      | 2024-01-01 00:00:00      | 50.00                | 2                              
+cancelled      | 2024-02-01 00:00:00      | 50.00                | 2                              
+cancelled      | 2024-03-01 00:00:00      | 100.00               | 1                              
+completed      | 2024-01-01 00:00:00      | 300.00               | 3                              
+completed      | 2024-02-01 00:00:00      | 500.00               | 2                              
+completed      | 2024-03-01 00:00:00      | 600.00               | 1                              
+pending        | 2024-01-01 00:00:00      | 150.00               | 3                              
+pending        | 2024-02-01 00:00:00      | 200.00               | 2                              
+pending        | 2024-03-01 00:00:00      | 300.00               | 1

rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/snapshots/cubesqlplanner__tests__integration__multi_stage__reduce_by__reduce_by_time_dimension.snap

@@ -0,0 +1,15 @@
+---
+source: cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/reduce_by.rs
+expression: result
+---
+orders__status | orders__created_at_month | orders__total_amount | orders__amount_reduce_time
+---------------+--------------------------+----------------------+---------------------------
+cancelled      | 2024-01-01 00:00:00      | 50.00                | 200.00                    
+cancelled      | 2024-02-01 00:00:00      | 50.00                | 200.00                    
+cancelled      | 2024-03-01 00:00:00      | 100.00               | 200.00                    
+completed      | 2024-01-01 00:00:00      | 300.00               | 1400.00                   
+completed      | 2024-02-01 00:00:00      | 500.00               | 1400.00                   
+completed      | 2024-03-01 00:00:00      | 600.00               | 1400.00                   
+pending        | 2024-01-01 00:00:00      | 150.00               | 650.00                    
+pending        | 2024-02-01 00:00:00      | 200.00               | 650.00                    
+pending        | 2024-03-01 00:00:00      | 300.00               | 650.00

rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/snapshots/cubesqlplanner__tests__integration__multi_stage__views__view_reduce_by_time_dimension.snap

@@ -0,0 +1,15 @@
+---
+source: cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/views.rs
+expression: result
+---
+orders_ms_view__status | orders_ms_view__created_at_month | orders_ms_view__total_amount | orders_ms_view__amount_reduce_time
+-----------------------+----------------------------------+------------------------------+-----------------------------------
+cancelled              | 2024-01-01 00:00:00              | 50.00                        | 200.00                            
+cancelled              | 2024-02-01 00:00:00              | 50.00                        | 200.00                            
+cancelled              | 2024-03-01 00:00:00              | 100.00                       | 200.00                            
+completed              | 2024-01-01 00:00:00              | 300.00                       | 1400.00                           
+completed              | 2024-02-01 00:00:00              | 500.00                       | 1400.00                           
+completed              | 2024-03-01 00:00:00              | 600.00                       | 1400.00                           
+pending                | 2024-01-01 00:00:00              | 150.00                       | 650.00                            
+pending                | 2024-02-01 00:00:00              | 200.00                       | 650.00                            
+pending                | 2024-03-01 00:00:00              | 300.00                       | 650.00

rust/cube/cubesqlplanner/cubesqlplanner/src/tests/integration/multi_stage/views.rs

@@ -9,6 +9,76 @@ fn create_context() -> TestContext {
 
 const SEED: &str = "integration_multi_stage_tables.sql";
 
+fn assert_uses_window(sql: &str) {
+    assert!(
+        sql.contains("OVER (PARTITION BY"),
+        "expected SQL to use a window function (`... OVER (PARTITION BY ...)`),\n\
+         got:\n{}",
+        sql,
+    );
+}
+
+/// PARTITION BY clause bodies (text between `PARTITION BY` and the following
+/// `ORDER BY`) of every window function in `sql`.
+fn partition_by_clauses(sql: &str) -> Vec<String> {
+    sql.match_indices("PARTITION BY")
+        .map(|(i, _)| {
+            let rest = &sql[i + "PARTITION BY".len()..];
+            let end = rest.find("ORDER BY").unwrap_or(rest.len());
+            rest[..end].to_string()
+        })
+        .collect()
+}
+
+// CORE-550 through a view: the query groups by view dimensions
+// (orders_ms_view.status / .created_at) while the measure's reduce_by targets
+// the underlying base cube member (orders.created_at). The granular view time
+// dimension (created_at_month) must still be dropped from the window's
+// PARTITION BY, leaving [status]. Guards that the fix resolves the reduce_by
+// reference through the view→base reference chain, not just for plain cubes.
+#[tokio::test(flavor = "multi_thread")]
+async fn test_view_reduce_by_time_dimension() {
+    let ctx = create_context();
+
+    let query = indoc! {r#"
+        measures:
+          - orders_ms_view.total_amount
+          - orders_ms_view.amount_reduce_time
+        dimensions:
+          - orders_ms_view.status
+        time_dimensions:
+          - dimension: orders_ms_view.created_at
+            granularity: month
+            dateRange:
+              - "2024-01-01"
+              - "2024-03-31"
+        order:
+          - id: orders_ms_view.status
+          - id: orders_ms_view.created_at
+    "#};
+
+    let sql = ctx.build_sql(query).unwrap();
+    assert_uses_window(&sql);
+    for clause in partition_by_clauses(&sql) {
+        assert!(
+            !clause.contains("created_at"),
+            "reduce_by(created_at) must drop the view time dimension from PARTITION BY,\n\
+             partition clause was: {}\nfull SQL:\n{}",
+            clause,
+            sql,
+        );
+        assert!(
+            clause.contains("status"),
+            "status must remain in PARTITION BY, partition clause was: {}",
+            clause,
+        );
+    }
+
+    if let Some(result) = ctx.try_execute_pg(query, SEED).await {
+        insta::assert_snapshot!(result);
+    }
+}
+
 #[tokio::test(flavor = "multi_thread")]
 async fn test_view_with_time_shift() {
     let ctx = create_context();