fix: Keep Computes with Aggregate when Filter follows (#5130) (#5639)

Co-authored-by: Claude <[email protected]>

Author: max-sixty

prqlc/prqlc/src/sql/pq/anchor.rs

@@ -349,7 +349,17 @@ fn is_split_required(transform: &SqlTransform, following: &mut HashSet<String>)
             contains_any(following, ["From", "Join", "Aggregate", "Compute"])
         }
         Super(Filter(_)) => contains_any(following, ["From", "Join"]),
-        Super(Compute(_)) => contains_any(following, ["From", "Join", /* "Aggregate" */ "Filter"]),
+        Super(Compute(_)) => {
+            // Don't split between Compute and Filter when there's an Aggregate in following.
+            // Filter after Aggregate becomes HAVING in the same SELECT, so all preceding
+            // Computes (including those with aggregation functions like SUM) must stay
+            // with the Aggregate to get proper GROUP BY.
+            if following.contains("Aggregate") {
+                contains_any(following, ["From", "Join"])
+            } else {
+                contains_any(following, ["From", "Join", /* "Aggregate" */ "Filter"])
+            }
+        }
 
         // Sort will be pushed down the CTEs, so there is no point in splitting for it.
         // Super(Sort(_)) => contains_any(following, ["From", "Join", "Compute", "Aggregate"]),

prqlc/prqlc/tests/integration/sql.rs

@@ -6151,7 +6151,7 @@ fn test_append_select_multiple() {
     )
     sort { +invoice_id, +total }
     select { total, invoice_id }
-    "###).unwrap(), @r"
+    "###).unwrap(), @"
     WITH table_3 AS (
       SELECT
         *
@@ -6273,7 +6273,10 @@ fn test_append_with_cte() {
 
 #[test]
 fn test_distinct_on_sort_on_compute() {
-    // Test for handling distinct on with sorting on computed columns
+    // Test for handling distinct on with sorting on computed columns.
+    // Note: table_0 includes billing_city and _expr_1 even though they're unused
+    // downstream. This is a side effect of the fix for #5130 which keeps Computes
+    // together with Aggregates when Filter follows.
     assert_snapshot!(compile(r###"
     from invoices
     derive code = case [customer_id < 10 => billing_postal_code, true => null]
@@ -6283,7 +6286,7 @@ fn test_distinct_on_sort_on_compute() {
     )
     filter (customer_id | in [4])
     group {billing_country} (aggregate {total = math.round 2 (sum total)})
-    "###).unwrap(), @r"
+    "###).unwrap(), @"
     WITH table_1 AS (
       SELECT
         billing_country,
@@ -6309,7 +6312,9 @@ fn test_distinct_on_sort_on_compute() {
           billing_country
           ORDER BY
             _expr_1 DESC
-        ) AS _expr_0
+        ) AS _expr_0,
+        billing_city,
+        _expr_1
       FROM
         table_1
     )
@@ -6917,3 +6922,50 @@ fn test_sort_take_before_aggregate() {
       total_sum DESC
     "#);
 }
+
+#[test]
+fn test_aggregate_with_operations_and_filter() {
+    // Issue #5130: Filtering on combined aggregates was generating invalid SQL.
+    // The CTE was missing GROUP BY when aggregate expressions had operations.
+    // Previously, this produced a CTE with SUM() but no GROUP BY clause.
+    assert_snapshot!(compile(r###"
+    from invoices
+    group billing_city (
+      aggregate{
+        sum_c = (sum customer_id) + (sum customer_id)
+      }
+    )
+    filter sum_c > 0
+    "###).unwrap(), @"
+    SELECT
+      billing_city,
+      COALESCE(SUM(customer_id), 0) + COALESCE(SUM(customer_id), 0) AS sum_c
+    FROM
+      invoices
+    GROUP BY
+      billing_city
+    HAVING
+      COALESCE(SUM(customer_id), 0) + COALESCE(SUM(customer_id), 0) > 0
+    ");
+
+    // Also test with scalar operations
+    assert_snapshot!(compile(r###"
+    from invoices
+    group billing_city (
+      aggregate{
+        sum_c = (sum customer_id) * 2
+      }
+    )
+    filter sum_c > 0
+    "###).unwrap(), @"
+    SELECT
+      billing_city,
+      COALESCE(SUM(customer_id), 0) * 2 AS sum_c
+    FROM
+      invoices
+    GROUP BY
+      billing_city
+    HAVING
+      COALESCE(SUM(customer_id), 0) * 2 > 0
+    ");
+}

web/book/tests/documentation/snapshots/documentation__book__README__prql-language-book__1.snap

@@ -2,27 +2,17 @@
 source: web/book/tests/documentation/book.rs
 expression: "from employees\nfilter start_date > @2021-01-01            # Clear date syntax\nderive {                                   # `derive` adds columns / variables\n  gross_salary = salary + (tax ?? 0),      # Terse coalesce\n  gross_cost = gross_salary + benefits,    # Variables can use other variables\n}\nfilter gross_cost > 0\ngroup {title, country} (                   # `group` runs a pipeline over each group\n  aggregate {                              # `aggregate` reduces each group to a value\n    average gross_salary,\n    sum_gross_cost = sum gross_cost,       # `=` sets a column name\n  }\n)\nfilter sum_gross_cost > 100_000            # `filter` replaces both of SQL's `WHERE` & `HAVING`\nderive id = f\"{title}_{country}\"           # F-strings like Python\nderive country_code = s\"LEFT(country, 2)\"  # S-strings permit SQL as an escape hatch\nsort {sum_gross_cost, -country}            # `-country` means descending order\ntake 1..20                                 # Range expressions (also valid as `take 20`)\n"
 ---
-WITH table_1 AS (
+WITH table_0 AS (
   SELECT
     title,
     country,
-    salary + COALESCE(tax, 0) + benefits AS _expr_1,
-    salary + COALESCE(tax, 0) AS _expr_2
+    AVG(salary + COALESCE(tax, 0)) AS _expr_0,
+    COALESCE(SUM(salary + COALESCE(tax, 0) + benefits), 0) AS sum_gross_cost
   FROM
     employees
   WHERE
     start_date > DATE '2021-01-01'
-),
-table_0 AS (
-  SELECT
-    title,
-    country,
-    AVG(_expr_2) AS _expr_0,
-    COALESCE(SUM(_expr_1), 0) AS sum_gross_cost
-  FROM
-    table_1
-  WHERE
-    _expr_1 > 0
+    AND salary + COALESCE(tax, 0) + benefits > 0
   GROUP BY
     title,
     country