Repository: tobymao/sqlglot
Pull request: lambda fix for single unbracketed arguments
PR to fix #73 

I know this works for the specific use cases I'm aware of where the desired (or at least allowed) syntax is always `x ->`. 

If there is a dialect where `x ->` is incorrect and only `(x) ->` is allowed then this generic application of the fix may not work.

Problem statement:
Arguments to lambda functions wrapped in brackets even with only one argument
```python
sql = """
transform(array('a','b','c'), X -> upper(X))
"""

import sqlglot
sqlglot.transpile(sql)
```
```
> ["TRANSFORM(ARRAY('a', 'b', 'c'), (X) -> UPPER(X))"]
```

I discovered this as it was breaking some pyspark queries. Minimal pyspark example below:
```python
from pyspark.context import SparkContext
from pyspark.sql import SparkSession
sc = SparkContext.getOrCreate()
spark = SparkSession(sc)
from pyspark.sql import Row
data_list = [
    {"array_col": ['a', 'b']},

        ]

df = spark.createDataFrame(Row(**x) for x in data_list)
df.createOrReplaceTempView("df")

import sqlglot 
sql = "select transform(array_col, x -> x) as transformed from df"
spark.sql(sql).toPandas()

sql = sqlglot.transpile(sql, read='spark', write='spark')
spark.sql(sql).toPandas()
```

Will submit a PR to:
- distinguish between `x -> f(x)` and `(x,y,z) -> f(x,y,z)` syntax
- add a new test 
- update `identity.sql` accordingly

Hints from issue discussion:

Commit message: reformatting to pass checks

Patch to review:
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index f0ae91f6f4..74d84cba20 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -423,8 +423,11 @@ def join_sql(self, expression):
         return f"{expression_sql}{op_sql} {this_sql}{on_sql}"
 
     def lambda_sql(self, expression):
+        a, b = (
+            ["(", ")"] if len(self.expressions(expression, flat=True)) > 1 else ["", ""]
+        )
         return self.no_identify(
-            lambda: f"({self.expressions(expression, flat=True)}) -> {self.sql(expression, 'this')}"
+            lambda: f"{a}{self.expressions(expression, flat=True)}{b} -> {self.sql(expression, 'this')}"
         )
 
     def lateral_sql(self, expression):
diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
index 54bcf51ef6..82ce06090b 100644
--- a/tests/fixtures/identity.sql
+++ b/tests/fixtures/identity.sql
@@ -59,7 +59,7 @@ INTERVAL '1' day
 INTERVAL '1' month
 INTERVAL 2 months
 QUANTILE(x, 0.5)
-REGEXP_REPLACE('new york', '(\w)(\w*)', (x) -> UPPER(x[1]) || LOWER(x[2]))
+REGEXP_REPLACE('new york', '(\w)(\w*)', x -> UPPER(x[1]) || LOWER(x[2]))
 REPLACE(1)
 DATE(x) = DATE(y)
 x[y - 1]
@@ -91,11 +91,11 @@ SELECT t.count
 SELECT DISTINCT x FROM test
 SELECT DISTINCT x, y FROM test
 SELECT GREATEST((3 + 1), LEAST(3, 4))
-SELECT TRANSFORM(a, (b) -> b) AS x
+SELECT TRANSFORM(a, b -> b) AS x
 SELECT AGGREGATE(a, (a, b) -> a + b) AS x
-SELECT X((a, b) -> a + b, (z) -> z) AS x
-SELECT X((a) -> "a" + ("z" - 1))
-SELECT EXISTS(ARRAY(2, 3), (x) -> x % 2 = 0)
+SELECT X((a, b) -> a + b, z -> z) AS x
+SELECT X(a -> "a" + ("z" - 1))
+SELECT EXISTS(ARRAY(2, 3), x -> x % 2 = 0)
 SELECT test.* FROM test
 SELECT a AS b FROM test
 SELECT "a"."b" FROM "a"
diff --git a/tests/test_transpile.py b/tests/test_transpile.py
index 7ddd55f423..d289e520db 100644
--- a/tests/test_transpile.py
+++ b/tests/test_transpile.py
@@ -280,3 +280,13 @@ def test_error_level(self, logger):
 
         with self.assertRaises(ParseError):
             transpile("x + 1 (")
+
+    def test_lambda(self):
+        self.validate(
+            "SELECT TRANSFORM(a, x -> UPPER(x)) FROM foo",
+            "SELECT TRANSFORM(a, x -> UPPER(x)) FROM foo",
+        )
+        self.validate(
+            "SELECT ARRAY_SORT(a, (x, y) -> CASE WHEN x < y THEN -1 WHEN x > y THEN 1 ELSE 0 END)",
+            "SELECT ARRAY_SORT(a, (x, y) -> CASE WHEN x < y THEN -1 WHEN x > y THEN 1 ELSE 0 END)",
+        )


Code review comments:
diff_hunk: @@ -280,3 +280,13 @@ def test_error_level(self, logger):
 
         with self.assertRaises(ParseError):
             transpile("x + 1 (")
+
+    def test_lambda(self):
path: tests/test_transpile.py
text: @user1:
can you delete these unit tests, they aren't necessary, or simply add these variants to identity.sql

@author:
Thanks, I thought that might be the case. I wrote the test before I discovered the extensive list in identity.sql

diff_hunk: @@ -423,8 +423,11 @@ def join_sql(self, expression):
         return f"{expression_sql}{op_sql} {this_sql}{on_sql}"
 
     def lambda_sql(self, expression):
+        a, b = (
path: sqlglot/generator.py
text: @user1:
can we clean this up a bit

args = self.expressions(expression, flat=True)
args = f"({args})" if len(args) > 1 else args
return self.no_identify(f"{args} -> ...")

@author:
Much neater, thanks. I wasn't happy with my code but no alternative came to mind.