Preserve alias-based group by dims

Author: nicosuave

test/sql/measures.test

@@ -761,6 +761,34 @@ FROM daily_orders_v
 2	320.0	840.0
 3	270.0	840.0
 
+# =============================================================================
+# Test: GROUP BY alias for expression dimension
+# =============================================================================
+
+statement ok
+CREATE TABLE monthly_sales (order_date DATE, region TEXT, amount DOUBLE);
+
+statement ok
+INSERT INTO monthly_sales VALUES
+    ('2023-01-05', 'US', 100), ('2023-01-12', 'EU', 50),
+    ('2023-02-03', 'US', 200), ('2023-02-20', 'EU', 20);
+
+statement ok
+CREATE VIEW monthly_sales_v AS
+SELECT DATE_TRUNC('month', order_date) AS month, region, SUM(amount) AS MEASURE revenue
+FROM monthly_sales
+GROUP BY DATE_TRUNC('month', order_date), region;
+
+query IIRR rowsort
+SEMANTIC SELECT month, region, AGGREGATE(revenue), AGGREGATE(revenue) AT (ALL region) AS month_total
+FROM monthly_sales_v
+;
+----
+2023-01-01	EU	50.0	150.0
+2023-01-01	US	100.0	150.0
+2023-02-01	EU	20.0	220.0
+2023-02-01	US	200.0	220.0
+
 # =============================================================================
 # Test: Multi-fact JOINs (wide tables)
 # =============================================================================

yardstick-rs/src/sql/measures.rs

@@ -1054,6 +1054,7 @@ fn group_by_matches_view(outer_cols: &[String], view_cols: &[String]) -> bool {
 fn filter_group_by_cols_for_measure(
     outer_cols: &[String],
     view_cols: &[String],
+    dimension_exprs: &HashMap<String, String>,
 ) -> Vec<String> {
     if view_cols.is_empty() {
         return outer_cols.to_vec();
@@ -1064,9 +1065,34 @@ fn filter_group_by_cols_for_measure(
         .map(|col| normalize_group_by_col(col))
         .collect();
 
+    let mut alias_expr_norms: Vec<(String, String)> = Vec::new();
+    alias_expr_norms.reserve(dimension_exprs.len());
+    for (alias, expr) in dimension_exprs.iter() {
+        alias_expr_norms.push((
+            normalize_group_by_col(alias),
+            normalize_group_by_col(expr),
+        ));
+    }
+
     outer_cols
         .iter()
-        .filter(|col| view_set.contains(&normalize_group_by_col(col)))
+        .filter(|col| {
+            let normalized_outer = normalize_group_by_col(col);
+            if view_set.contains(&normalized_outer) {
+                return true;
+            }
+
+            if let Some(expr) = dimension_exprs.get(&normalized_outer) {
+                let normalized_expr = normalize_group_by_col(expr);
+                if view_set.contains(&normalized_expr) {
+                    return true;
+                }
+            }
+
+            alias_expr_norms.iter().any(|(alias_norm, expr_norm)| {
+                expr_norm == &normalized_outer && view_set.contains(alias_norm)
+            })
+        })
         .cloned()
         .collect()
 }
@@ -3436,8 +3462,11 @@ pub fn expand_aggregate_with_at(sql: &str) -> AggregateExpandResult {
     for (measure_name, modifiers, start, end) in patterns {
         // Look up which view contains this measure (for JOIN support)
         let resolved = resolve_measure_source(&measure_name, &primary_table_name);
-        let measure_group_by_cols =
-            filter_group_by_cols_for_measure(&group_by_cols, &resolved.view_group_by_cols);
+        let measure_group_by_cols = filter_group_by_cols_for_measure(
+            &group_by_cols,
+            &resolved.view_group_by_cols,
+            &resolved.dimension_exprs,
+        );
 
         // Non-decomposable measures are recomputed from base rows (including AT modifiers)
 
@@ -3551,8 +3580,11 @@ pub fn expand_aggregate_with_at(sql: &str) -> AggregateExpandResult {
 
     for (measure_name, start, end) in plain_calls {
         let resolved = resolve_measure_source(&measure_name, &primary_table_name);
-        let measure_group_by_cols =
-            filter_group_by_cols_for_measure(&group_by_cols, &resolved.view_group_by_cols);
+        let measure_group_by_cols = filter_group_by_cols_for_measure(
+            &group_by_cols,
+            &resolved.view_group_by_cols,
+            &resolved.dimension_exprs,
+        );
 
 
         // For derived measures, use the expanded expression; otherwise use AGG(measure_name)