Fix(optimizer): expand refs in QUALIFY into corresponding projections (…

…tobymao#1665)

* Fix(optimizer): expand refs in QUALIFY into corresponding projections

* Add test to cover execution path where we have an AggFunc

sqlglot/optimizer/qualify_columns.py

@@ -194,6 +194,7 @@ def replace_columns(
     replace_columns(expression.args.get("where"))
     replace_columns(expression.args.get("group"))
     replace_columns(expression.args.get("having"), resolve_agg=True)
+    replace_columns(expression.args.get("qualify"), resolve_agg=True)
     replace_columns(expression.args.get("order"), expand=False, resolve_agg=True)
     scope.clear_cache()
 

tests/fixtures/optimizer/pushdown_projections.sql

@@ -61,9 +61,13 @@ SELECT q.x AS x FROM (VALUES (1, 2)) AS q(x, y);
 SELECT i.a FROM x AS i LEFT JOIN (SELECT a, b FROM (SELECT a, b FROM x)) AS j ON i.a = j.a;
 SELECT i.a AS a FROM x AS i LEFT JOIN (SELECT _q_0.a AS a FROM (SELECT x.a AS a FROM x AS x) AS _q_0) AS j ON i.a = j.a;
 
+WITH cte AS (SELECT source.a AS a, ROW_NUMBER() OVER (PARTITION BY source.id, source.timestamp ORDER BY source.a DESC) AS index FROM source AS source QUALIFY index) SELECT cte.a AS a FROM cte;
+WITH cte AS (SELECT source.a AS a FROM source AS source QUALIFY ROW_NUMBER() OVER (PARTITION BY source.id, source.timestamp ORDER BY source.a DESC)) SELECT cte.a AS a FROM cte;
+
 --------------------------------------
 -- Unknown Star Expansion
 --------------------------------------
+
 SELECT a FROM (SELECT * FROM zz) WHERE b = 1;
 SELECT _q_0.a AS a FROM (SELECT zz.a AS a, zz.b AS b FROM zz AS zz) AS _q_0 WHERE _q_0.b = 1;
 

tests/fixtures/optimizer/qualify_columns.sql

@@ -396,16 +396,20 @@ SELECT x.a + 1 AS a, ROW_NUMBER() OVER (PARTITION BY x.b ORDER BY x.a) AS row_nu
 
 # dialect: bigquery
 SELECT ROW_NUMBER() OVER (PARTITION BY a ORDER BY b) AS row_num FROM x QUALIFY row_num = 1;
-SELECT ROW_NUMBER() OVER (PARTITION BY x.a ORDER BY x.b) AS row_num FROM x AS x QUALIFY row_num = 1;
+SELECT ROW_NUMBER() OVER (PARTITION BY x.a ORDER BY x.b) AS row_num FROM x AS x QUALIFY ROW_NUMBER() OVER (PARTITION BY x.a ORDER BY x.b) = 1;
 
 # dialect: bigquery
 SELECT x.b, x.a FROM x LEFT JOIN y ON x.b = y.b QUALIFY ROW_NUMBER() OVER(PARTITION BY x.b ORDER BY x.a DESC) = 1;
 SELECT x.b AS b, x.a AS a FROM x AS x LEFT JOIN y AS y ON x.b = y.b QUALIFY ROW_NUMBER() OVER (PARTITION BY x.b ORDER BY x.a DESC) = 1;
 
+# dialect: snowflake
+SELECT * FROM x QUALIFY COUNT(a) OVER (PARTITION BY b) > 1;
+SELECT x.a AS a, x.b AS b FROM x AS x QUALIFY COUNT(x.a) OVER (PARTITION BY x.b) > 1;
 
 --------------------------------------
 -- Expand laterals
 --------------------------------------
+
 # title: expand alias reference
 SELECT
   x.a + 1 AS i,