Fix(optimizer): handle TableFromRows properly in annotate_types (#4889)

Author: georgesittas

sqlglot/lineage.py

@@ -237,7 +237,7 @@ def to_node(
     # Find all columns that went into creating this one to list their lineage nodes.
     source_columns = set(find_all_in_scope(select, exp.Column))
 
-    # If the source is a UDTF find columns used in the UTDF to generate the table
+    # If the source is a UDTF find columns used in the UDTF to generate the table
     if isinstance(source, exp.UDTF):
         source_columns |= set(source.find_all(exp.Column))
         derived_tables = [

sqlglot/optimizer/annotate_types.py

@@ -223,7 +223,7 @@ def annotate_scope(self, scope: Scope) -> None:
                         values = [expression.this.this]
                 elif isinstance(expression, exp.Unnest):
                     values = [expression]
-                else:
+                elif not isinstance(expression, exp.TableFromRows):
                     values = expression.expressions[0].expressions
 
                 if not values:

tests/test_optimizer.py

@@ -1434,6 +1434,16 @@ def test_union_annotation(self):
         """
         self.assertEqual(optimizer.optimize(sql).selects[0].type.this, exp.DataType.Type.VARCHAR)
 
+    def test_udtf_annotation(self):
+        table_udtf = parse_one(
+            "SELECT * FROM TABLE(GENERATOR(ROWCOUNT => 100000))",
+            read="snowflake",
+        )
+        self.assertEqual(
+            annotate_types(table_udtf, dialect="snowflake").sql("snowflake"),
+            "SELECT * FROM TABLE(GENERATOR(ROWCOUNT => 100000))",
+        )
+
     def test_recursive_cte(self):
         query = parse_one(
             """