fix(snowflake): Fix annotate_types crashes for FLATTEN over STRUCT and WithinGroup (#7434)

Author: georgesittas

sqlglot/optimizer/annotate_types.py

@@ -655,19 +655,21 @@ def _annotate_by_args(
             for expr in ensure_list(expressions):
                 expr_type = expr.type
 
-                # Stop at the first nested data type found - we don't want to _maybe_coerce nested types
+                if expr_type.is_type(exp.DType.UNKNOWN):
+                    self._set_type(expression, exp.DType.UNKNOWN)
+                    return expression
+
+                if nested_type:
+                    continue
+
+                # Stop coercing at the first nested data type found
                 if expr_type.args.get("nested"):
                     nested_type = expr_type
-                    break
-
-                if isinstance(expr, exp.Literal):
+                elif isinstance(expr, exp.Literal):
                     literal_type = self._maybe_coerce(literal_type or expr_type, expr_type)
                 else:
                     non_literal_type = self._maybe_coerce(non_literal_type or expr_type, expr_type)
 
-            if nested_type:
-                break
-
         result_type = None
 
         if nested_type:

sqlglot/typing/snowflake.py

@@ -102,6 +102,8 @@ def _annotate_within_group(self: TypeAnnotator, expression: exp.WithinGroup) ->
         and isinstance(ordered_expr := order_expr.expressions[0], exp.Ordered)
     ):
         self._set_type(expression, ordered_expr.this.type)
+    else:
+        self._set_type(expression, expression.this.type)
 
     return expression
 

tests/fixtures/optimizer/annotate_functions.sql

@@ -5287,6 +5287,14 @@ DOUBLE;
 PERCENTILE_CONT(0.75) WITHIN GROUP (ORDER BY tbl.bigint_col) OVER (PARTITION BY 1);
 BIGINT;
 
+# dialect: snowflake
+ARRAY_AGG(tbl.int_col) WITHIN GROUP (ORDER BY tbl.int_col);
+ARRAY;
+
+# dialect: snowflake
+ARRAY_AGG(tbl.int_col) WITHIN GROUP (ORDER BY tbl.int_col) OVER (PARTITION BY tbl.text_col);
+ARRAY;
+
 # dialect: snowflake
 PARSE_IP('192.168.1.1', 'INET');
 OBJECT;

tests/fixtures/optimizer/optimizer.sql

@@ -1566,4 +1566,27 @@ RIGHT JOIN GENERATE_SERIES((
   FROM "x" AS "x"
 ), 10, 1) AS "t1"("c1")
   ON "t1"."c1" > "z"."c"
-;
+;
+
+# title: flatten over object_construct
+# dialect: snowflake
+# execute: false
+WITH obj AS (SELECT object_construct('a', '1', 'b', '2') AS data), flattened AS (SELECT f.key, f.value FROM obj, lateral flatten(input => obj.data) AS f) SELECT key::varchar, value::varchar FROM flattened;
+WITH "OBJ" AS (
+  SELECT
+    OBJECT_CONSTRUCT('a', '1', 'b', '2') AS "DATA"
+)
+SELECT
+  CAST("F"."KEY" AS VARCHAR) AS "KEY",
+  CAST("F"."VALUE" AS VARCHAR) AS "VALUE"
+FROM "OBJ" AS "OBJ"
+CROSS JOIN LATERAL FLATTEN(input => "OBJ"."DATA") AS "F"("SEQ", "KEY", "PATH", "INDEX", "VALUE", "THIS");
+
+# title: array_agg within group over
+# dialect: snowflake
+# execute: false
+SELECT array_agg(id) WITHIN GROUP (ORDER BY id) OVER (PARTITION BY grp) FROM t;
+SELECT
+  ARRAY_AGG("T"."ID") WITHIN GROUP (ORDER BY
+    "T"."ID") OVER (PARTITION BY "T"."GRP") AS "_col_0"
+FROM "T" AS "T";