in work

Author: waralexrom

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

@@ -84,6 +84,17 @@ cubes:
         type: string
         add_group_by: [orders.customerId]
 
+      - name: changeTypeComplexWithJoin
+        sql: >
+          CASE
+            WHEN {revenueYearAgo} IS NULL THEN 'New'
+            WHEN {revenue} > {revenueYearAgo} THEN 'Grow'
+            ELSE 'Down'
+          END
+        multi_stage: true
+        type: string
+        add_group_by: [first_date.customerId]
+
       - name: changeTypeConcat
         sql: "CONCAT({changeTypeComplex}, '-test')"
         type: string
@@ -115,6 +126,45 @@ cubes:
           END
         type: string
 
+  - name: first_date
+    sql: >
+      SELECT  1 AS id, '2023-03-01T00:00:00Z'::timestamptz AS createdAt, 1 AS customerId UNION ALL
+      SELECT  8 AS id, '2023-09-01T00:00:00Z'::timestamptz AS createdAt, 2 AS customerId UNION ALL
+      SELECT 16 AS id, '2024-09-01T00:00:00Z'::timestamptz AS createdAt, 3 AS customerId UNION ALL
+      SELECT 23 AS id, '2025-03-01T00:00:00Z'::timestamptz AS createdAt, 4 AS customerId UNION ALL
+      SELECT 29 AS id, '2025-03-01T00:00:00Z'::timestamptz AS createdAt, 5 AS customerId UNION ALL
+      SELECT 36 AS id, '2025-09-01T00:00:00Z'::timestamptz AS createdAt, 6 AS customerId
+
+    joins:
+    - name: orders
+      sql: "{first_date.customerId} = {orders.customerId}"
+      relationship: one_to_many
+
+    dimensions:
+      - name: customerId
+        sql: customerId
+        type: number
+
+      - name: createdAt
+        sql: createdAt
+        type: time
+
+      - name: customerType
+        sql: >
+          CASE
+            WHEN {orders.revenue} < 10000 THEN 'Low'
+            WHEN {orders.revenue} < 20000 THEN 'Medium'
+            ELSE 'Top'
+          END
+        multi_stage: true
+        type: string
+        add_group_by: [first_date.customerId]
+
+      - name: customerTypeConcat
+        sql: "CONCAT('Customer type: ', {customerType})"
+        multi_stage: true
+        type: string
+        add_group_by: [first_date.customerId]
 
 
     `);
@@ -288,6 +338,62 @@ cubes:
       },
     ],
     { joinGraph, cubeEvaluator, compiler }));
+    it('bucketing with join and bucket dimension', async () => dbRunner.runQueryTest({
+      dimensions: ['orders.changeTypeComplexWithJoin'],
+      measures: ['orders.revenue', 'orders.revenueYearAgo'],
+      timeDimensions: [
+        {
+          dimension: 'orders.createdAt',
+          granularity: 'year',
+          dateRange: ['2024-01-02T00:00:00', '2026-01-01T00:00:00']
+        }
+      ],
+      timezone: 'UTC',
+      order: [{
+        id: 'orders.changeTypeComplexWithJoin'
+      }, { id: 'orders.createdAt' }],
+    },
+    [
+      {
+        orders__change_type_complex_with_join: 'Down',
+        orders__created_at_year: '2024-01-01T00:00:00.000Z',
+        orders__revenue: '20400',
+        orders__revenue_year_ago: '22800'
+      },
+      {
+        orders__change_type_complex_with_join: 'Down',
+        orders__created_at_year: '2025-01-01T00:00:00.000Z',
+        orders__revenue: '17800',
+        orders__revenue_year_ago: '20400'
+      },
+      {
+        orders__change_type_complex_with_join: 'Grow',
+        orders__created_at_year: '2024-01-01T00:00:00.000Z',
+        orders__revenue: '11700',
+        orders__revenue_year_ago: '9400'
+      },
+      {
+        orders__change_type_complex_with_join: 'Grow',
+        orders__created_at_year: '2025-01-01T00:00:00.000Z',
+        orders__revenue: '14100',
+        orders__revenue_year_ago: '11700'
+      },
+    ],
+    { joinGraph, cubeEvaluator, compiler }));
+    it('bucketing dim reference other cube measure', async () => dbRunner.runQueryTest({
+      dimensions: ['first_date.customerType'],
+      measures: ['orders.revenue'],
+      timezone: 'UTC',
+      order: [{
+        id: 'first_date.customerType'
+      }],
+    },
+    [
+      { first_date__customer_type: 'Low', orders__revenue: '8100' },
+      { first_date__customer_type: 'Medium', orders__revenue: '41700' },
+      { first_date__customer_type: 'Top', orders__revenue: '46400' }
+    ],
+    { joinGraph, cubeEvaluator, compiler }));
   } else {
     // This test is working only in tesseract
     test.skip('multi stage over sub query', () => { expect(1).toBe(1); });

rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/collectors/join_hints_collector.rs

@@ -27,7 +27,20 @@ impl TraversalVisitor for JoinHintsCollector {
         _: &Self::State,
     ) -> Result<Option<Self::State>, CubeError> {
         if node.is_multi_stage() {
-            //We don't add multi-stage members childs to join hints
+            if let Ok(dim) = node.as_dimension() {
+                if let Some(add_group_by) = dim.add_group_by() {
+                    for item in add_group_by.iter() {
+                        self.apply(item, &())?;
+                    }
+                    for (dep, path) in dim.get_dependencies_with_path().into_iter() {
+                        if let Ok(dim) = dep.as_dimension() {
+                            if dim.is_multi_stage() {
+                                self.on_node_traverse(&dep, &path, &())?;
+                            }
+                        }
+                    }
+                }
+            }
             return Ok(None);
         }
 

rust/cubesqlplanner/cubesqlplanner/src/planner/sql_evaluator/symbols/dimension_symbol.rs

@@ -238,11 +238,6 @@ impl DimensionSymbol {
         if let Some(member_sql) = &self.longitude {
             member_sql.extract_symbol_deps(&mut deps);
         }
-        if let Some(add_group_by) = &self.add_group_by {
-            for member_sql in add_group_by {
-                deps.extend(member_sql.get_dependencies().into_iter());
-            }
-        }
         if let Some(case) = &self.case {
             case.extract_symbol_deps(&mut deps);
         }
@@ -263,11 +258,6 @@ impl DimensionSymbol {
         if let Some(case) = &self.case {
             case.extract_symbol_deps_with_path(&mut deps);
         }
-        if let Some(add_group_by) = &self.add_group_by {
-            for member_sql in add_group_by {
-                deps.extend(member_sql.get_dependencies_with_path().into_iter());
-            }
-        }
         deps
     }