Repository: tobymao/sqlglot

Issue:
Lambda expression with data types in Snowflake dialect
**Before you file an issue**
- Make sure you specify the "read" dialect eg. `parse_one(sql, read="spark")`
- Make sure you specify the "write" dialect eg. `ast.sql(dialect="duckdb")`
- Check if the issue still exists on main

**Fully reproducible code snippet**
Please include a fully reproducible code snippet or the input sql, dialect, and expected output.

Lambda expression with type definition are incorrectly handled.
```
import sqlglot
sql = 'SELECT TRANSFORM([1, 2, 3], a INT -> a * 2) AS "Multiply by Two"'
sqlglot.parse_one(sql)
```

Without type definitions query is parsed correctly.
```
SELECT TRANSFORM([1, 2, 3], a -> a * 2) AS "Multiply by Two";
```
```
Select(
  expressions=[
    Alias(
      this=Transform(
        this=Array(
          expressions=[
            Literal(this=1, is_string=False),
            Literal(this=2, is_string=False),
            Literal(this=3, is_string=False)]),
        expression=Lambda(
          this=Mul(
            this=Identifier(this=a, quoted=False),
            expression=Literal(this=2, is_string=False)),
          expressions=[
            Identifier(this=a, quoted=False)])),
      alias=Identifier(this=Multiply by Two, quoted=True))])
```

For details, see:
https://docs.snowflake.com/en/sql-reference/functions/transform#multiply-each-element-in-an-array-by-a-value



**Official Documentation**
Please include links to official SQL documentation related to your issue.



Discussion:

Resolving patch:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -482,6 +482,18 @@ class Parser(parser.Parser):
 
         SCHEMA_KINDS = {"OBJECTS", "TABLES", "VIEWS", "SEQUENCES", "UNIQUE KEYS", "IMPORTED KEYS"}
 
+        LAMBDAS = {
+            **parser.Parser.LAMBDAS,
+            TokenType.ARROW: lambda self, expressions: self.expression(
+                exp.Lambda,
+                this=self._replace_lambda(
+                    self._parse_conjunction(),
+                    expressions,
+                ),
+                expressions=[e.this if isinstance(e, exp.Cast) else e for e in expressions],
+            ),
+        }
+
         def _parse_create(self) -> exp.Create | exp.Command:
             expression = super()._parse_create()
             if isinstance(expression, exp.Create) and expression.kind == "TAG":
@@ -716,6 +728,19 @@ def _parse_location_path(self) -> exp.Var:
 
             return exp.var("".join(part.text for part in parts if part))
 
+        def _parse_lambda_arg(self) -> t.Optional[exp.Expression]:
+            this = super()._parse_lambda_arg()
+
+            if not this:
+                return this
+
+            typ = self._parse_types()
+
+            if typ:
+                return self.expression(exp.Cast, this=this, to=typ)
+
+            return this
+
     class Tokenizer(tokens.Tokenizer):
         STRING_ESCAPES = ["\\", "'"]
         HEX_STRINGS = [("x'", "'"), ("X'", "'")]
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -580,7 +580,7 @@ class Parser(metaclass=_Parser):
             exp.Lambda,
             this=self._replace_lambda(
                 self._parse_conjunction(),
-                {node.name for node in expressions},
+                expressions,
             ),
             expressions=expressions,
         ),
@@ -4680,18 +4680,21 @@ def _parse_session_parameter(self) -> exp.SessionParameter:
 
         return self.expression(exp.SessionParameter, this=this, kind=kind)
 
+    def _parse_lambda_arg(self) -> t.Optional[exp.Expression]:
+        return self._parse_id_var()
+
     def _parse_lambda(self, alias: bool = False) -> t.Optional[exp.Expression]:
         index = self._index
 
         if self._match(TokenType.L_PAREN):
             expressions = t.cast(
-                t.List[t.Optional[exp.Expression]], self._parse_csv(self._parse_id_var)
+                t.List[t.Optional[exp.Expression]], self._parse_csv(self._parse_lambda_arg)
             )
 
             if not self._match(TokenType.R_PAREN):
                 self._retreat(index)
         else:
-            expressions = [self._parse_id_var()]
+            expressions = [self._parse_lambda_arg()]
 
         if self._match_set(self.LAMBDAS):
             return self.LAMBDAS[self._prev.token_type](self, expressions)
@@ -6433,14 +6436,25 @@ def _match_text_seq(self, *texts, advance=True):
         return True
 
     def _replace_lambda(
-        self, node: t.Optional[exp.Expression], lambda_variables: t.Set[str]
+        self, node: t.Optional[exp.Expression], expressions: t.List[exp.Expression]
     ) -> t.Optional[exp.Expression]:
         if not node:
             return node
 
+        lambda_types = {e.name: e.args.get("to") or False for e in expressions}
+
         for column in node.find_all(exp.Column):
-            if column.parts[0].name in lambda_variables:
+            typ = lambda_types.get(column.parts[0].name)
+            if typ is not None:
                 dot_or_id = column.to_dot() if column.table else column.this
+
+                if typ:
+                    dot_or_id = self.expression(
+                        exp.Cast,
+                        this=dot_or_id,
+                        to=typ,
+                    )
+
                 parent = column.parent
 
                 while isinstance(parent, exp.Dot):


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -10,6 +10,11 @@ class TestSnowflake(Validator):
     dialect = "snowflake"
 
     def test_snowflake(self):
+        self.validate_identity(
+            "transform(x, a int -> a + a + 1)",
+            "TRANSFORM(x, a -> CAST(a AS INT) + CAST(a AS INT) + 1)",
+        )
+
         self.validate_all(
             "ARRAY_CONSTRUCT_COMPACT(1, null, 2)",
             write={