Repository: tobymao/sqlglot
Pull request: Fix(optimizer): wrap scalar subquery replacement in a MAX call
Fixes #1987

Tested this in SQL fiddle, seems to be good.

Problem statement:
Nested Query Bug - Column '_u_0._col_0' is invalid in the select list because it is not contained in either an aggregate function or the GROUP BY clause.
<h3>Code snippet:</h3>

```python
import sqlglot
from sqlglot.optimizer import optimize, RULES as optimize_rules
DCT_RULES = {x.__name__ : x for x in optimize_rules}
DCT_RULES.pop('quote_identifiers')
optimize_rules = tuple(list(DCT_RULES.values()))

if __name__ == '__main__':

  sqlglot_metadata = {
    "catalogo_proveedores": {
      "ID_PROVEEDOR": "INT",
      "PROVEEDOR": "STRING"
    },
    "compra_real": {
      "ANNUAL": "STRING",
      "TRIMESTRE": "STRING",
      "PIEZAS": "FLOAT",
      "PRECIO": "FLOAT",
      "IMPORTE": "FLOAT",
      "CLAVE_CBYCM": "INT",
      "ID_PROVEEDOR": "INT"
    },
  }
  # This query is working in this form
  sql_query = """SELECT PROVEEDOR, SUM(IMPORTE) / (SELECT SUM(IMPORTE) FROM compra_real WHERE CLAVE_CBYCM = 2133 AND ANNUAL = '2022') * 100 AS Percentage_Imports FROM compra_real INNER JOIN catalogo_proveedores ON compra_real.ID_PROVEEDOR = catalogo_proveedores.ID_PROVEEDOR WHERE CLAVE_CBYCM = 2133 AND ANNUAL = '2022' GROUP BY PROVEEDOR"""

  expr = optimize(
    expression=sqlglot.parse_one(sql_query, read="tsql"),
    schema=sqlglot_metadata,
    rules=optimize_rules,
  )

  print(expr.sql(pretty=True, dialect="tsql"))
```


<h3>Optimised query</h3>

```
WITH _u_0 AS (
  SELECT
    SUM(compra_real.importe) AS _col_0
  FROM compra_real AS compra_real
  WHERE
    compra_real.annual = '2022' AND compra_real.clave_cbycm = 2133
)
SELECT
  catalogo_proveedores.proveedor AS proveedor,
  SUM(compra_real.importe) / _u_0._col_0 * 100 AS percentage_imports
FROM compra_real AS compra_real
CROSS JOIN _u_0 AS _u_0
JOIN catalogo_proveedores AS catalogo_proveedores
  ON compra_real.id_proveedor = catalogo_proveedores.id_proveedor
WHERE
  compra_real.annual = '2022' AND compra_real.clave_cbycm = 2133
GROUP BY
  catalogo_proveedores.proveedor
```

<h3>The error:</h3>

```
Column '_u_0._col_0' is invalid in the select list because it is not contained in either an aggregate function or the GROUP BY clause.
```

<h3>Expected output:</h3>
If the GROUP BY clause contains also _u_0._col_0, then the query will be executed successfully.

Hints from issue discussion:
Thanks for the report, we'll take a look

Commit message: Fix(optimizer): wrap scalar subquery replacement in a MAX call

Patch to review:
diff --git a/sqlglot/optimizer/unnest_subqueries.py b/sqlglot/optimizer/unnest_subqueries.py
index 09e3f2aa4e..9d17b25ab8 100644
--- a/sqlglot/optimizer/unnest_subqueries.py
+++ b/sqlglot/optimizer/unnest_subqueries.py
@@ -46,20 +46,20 @@ def unnest(select, parent_select, next_alias_name):
     if not predicate or parent_select is not predicate.parent_select:
         return
 
-    # this subquery returns a scalar and can just be converted to a cross join
+    # This subquery returns a scalar and can just be converted to a cross join
     if not isinstance(predicate, (exp.In, exp.Any)):
         having = predicate.find_ancestor(exp.Having)
         column = exp.column(select.selects[0].alias_or_name, alias)
-        if having and having.parent_select is parent_select:
+
+        if (
+            (having and having.parent_select is parent_select)
+            or any(projection.find(exp.AggFunc) for projection in parent_select.selects)
+            or parent_select.args.get("group")
+        ):
             column = exp.Max(this=column)
-        _replace(select.parent, column)
 
-        parent_select.join(
-            select,
-            join_type="CROSS",
-            join_alias=alias,
-            copy=False,
-        )
+        _replace(select.parent, column)
+        parent_select.join(select, join_type="CROSS", join_alias=alias, copy=False)
         return
 
     if select.find(exp.Limit, exp.Offset):
diff --git a/tests/fixtures/optimizer/optimizer.sql b/tests/fixtures/optimizer/optimizer.sql
index 53708a4e0e..981e0520a9 100644
--- a/tests/fixtures/optimizer/optimizer.sql
+++ b/tests/fixtures/optimizer/optimizer.sql
@@ -890,3 +890,20 @@ FROM (
     JOIN "y" AS "y"
       ON "x"."a" = "y"."c"
 );
+
+# title: replace scalar subquery, wrap resulting column in a MAX
+SELECT a, SUM(c) / (SELECT SUM(c) FROM y) * 100 AS foo FROM y INNER JOIN x ON y.b = x.b GROUP BY a;
+WITH "_u_0" AS (
+  SELECT
+    SUM("y"."c") AS "_col_0"
+  FROM "y" AS "y"
+)
+SELECT
+  "x"."a" AS "a",
+  SUM("y"."c") / MAX("_u_0"."_col_0") * 100 AS "foo"
+FROM "y" AS "y"
+CROSS JOIN "_u_0" AS "_u_0"
+JOIN "x" AS "x"
+  ON "y"."b" = "x"."b"
+GROUP BY
+  "x"."a";
diff --git a/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql b/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql
index 8aaf50cb3a..46715da81d 100644
--- a/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql
+++ b/tests/fixtures/optimizer/tpc-ds/tpc-ds.sql
@@ -956,7 +956,7 @@ JOIN "date_dim" AS "d"
 JOIN "item" AS "i"
   ON "s"."ss_item_sk" = "i"."i_item_sk"
 JOIN "_u_0" AS "_u_0"
-  ON "d"."d_month_seq" = "_u_0"."d_month_seq"
+  ON "d"."d_month_seq" = MAX("_u_0"."d_month_seq")
 LEFT JOIN "_u_1" AS "_u_1"
   ON "_u_1"."_u_2" = "i"."i_category"
 WHERE
@@ -1449,11 +1449,31 @@ WITH "_u_0" AS (
     "store_sales"."ss_quantity" <= 80 AND "store_sales"."ss_quantity" >= 61
 )
 SELECT
-  CASE WHEN "_u_0"."_col_0" > 3672 THEN "_u_1"."_col_0" ELSE "_u_2"."_col_0" END AS "bucket1",
-  CASE WHEN "_u_3"."_col_0" > 3392 THEN "_u_4"."_col_0" ELSE "_u_5"."_col_0" END AS "bucket2",
-  CASE WHEN "_u_6"."_col_0" > 32784 THEN "_u_7"."_col_0" ELSE "_u_8"."_col_0" END AS "bucket3",
-  CASE WHEN "_u_9"."_col_0" > 26032 THEN "_u_10"."_col_0" ELSE "_u_11"."_col_0" END AS "bucket4",
-  CASE WHEN "_u_12"."_col_0" > 23982 THEN "_u_13"."_col_0" ELSE "_u_14"."_col_0" END AS "bucket5"
+  CASE
+    WHEN MAX("_u_0"."_col_0") > 3672
+    THEN MAX("_u_1"."_col_0")
+    ELSE MAX("_u_2"."_col_0")
+  END AS "bucket1",
+  CASE
+    WHEN MAX("_u_3"."_col_0") > 3392
+    THEN MAX("_u_4"."_col_0")
+    ELSE MAX("_u_5"."_col_0")
+  END AS "bucket2",
+  CASE
+    WHEN MAX("_u_6"."_col_0") > 32784
+    THEN MAX("_u_7"."_col_0")
+    ELSE MAX("_u_8"."_col_0")
+  END AS "bucket3",
+  CASE
+    WHEN MAX("_u_9"."_col_0") > 26032
+    THEN MAX("_u_10"."_col_0")
+    ELSE MAX("_u_11"."_col_0")
+  END AS "bucket4",
+  CASE
+    WHEN MAX("_u_12"."_col_0") > 23982
+    THEN MAX("_u_13"."_col_0")
+    ELSE MAX("_u_14"."_col_0")
+  END AS "bucket5"
 FROM "reason" AS "reason"
 CROSS JOIN "_u_0" AS "_u_0"
 CROSS JOIN "_u_1" AS "_u_1"
@@ -6589,9 +6609,9 @@ WITH "cs_or_ws_sales" AS (
     ON "customer_address"."ca_county" = "store"."s_county"
     AND "customer_address"."ca_state" = "store"."s_state"
   JOIN "_u_0" AS "_u_0"
-    ON "date_dim"."d_month_seq" >= "_u_0"."_col_0"
+    ON "date_dim"."d_month_seq" >= MAX("_u_0"."_col_0")
   JOIN "_u_1" AS "_u_1"
-    ON "date_dim"."d_month_seq" <= "_u_1"."_col_0"
+    ON "date_dim"."d_month_seq" <= MAX("_u_1"."_col_0")
   GROUP BY
     "my_customers"."c_customer_sk"
 )


Code review comments:
diff_hunk: @@ -46,20 +46,20 @@ def unnest(select, parent_select, next_alias_name):
     if not predicate or parent_select is not predicate.parent_select:
         return
 
-    # this subquery returns a scalar and can just be converted to a cross join
+    # This subquery returns a scalar and can just be converted to a cross join
     if not isinstance(predicate, (exp.In, exp.Any)):
         having = predicate.find_ancestor(exp.Having)
         column = exp.column(select.selects[0].alias_or_name, alias)
-        if having and having.parent_select is parent_select:
+
+        if (
+            (having and having.parent_select is parent_select)
+            or any(projection.find(exp.AggFunc) for projection in parent_select.selects)
+            or parent_select.args.get("group")
path: sqlglot/optimizer/unnest_subqueries.py
text: @user1:
put the group by first

@author:
Good point, thanks

diff_hunk: @@ -6589,9 +6609,9 @@ WITH "cs_or_ws_sales" AS (
     ON "customer_address"."ca_county" = "store"."s_county"
     AND "customer_address"."ca_state" = "store"."s_state"
   JOIN "_u_0" AS "_u_0"
-    ON "date_dim"."d_month_seq" >= "_u_0"."_col_0"
+    ON "date_dim"."d_month_seq" >= MAX("_u_0"."_col_0")
path: tests/fixtures/optimizer/tpc-ds/tpc-ds.sql
text: @user1:
i don't think this one is correct

@author:
Made a few changes to address this, let me know when you take another look