Repository: tobymao/sqlglot

Issue:
Support TRANSFORM in hive dialect
**Is your feature request related to a problem? Please describe.**
Currently, parsing hive query with transform, e.g.: 
```
parse("SELECT TRANSFORM (x) USING 'x' AS (x int) FROM t", read='hive')
```
yields an exception 
```
ParseError: Invalid expression / Unexpected token. Line 1, Col: 22.
  SELECT TRANSFORM (x) USING 'x' AS (x int) FROM t
```

**Describe the solution you'd like**
TRANSFORM clause is supported in hive dialect

**Describe alternatives you've considered**
A clear and concise description of any alternative solutions or features you've considered.

**Additional context**
Hive transform grammar: https://cwiki.apache.org/confluence/display/Hive/LanguageManual+Transform
Spark transform grammar: https://spark.apache.org/docs/3.1.2/sql-ref-syntax-qry-select-transform.html


Discussion:

Patch:
diff --git a/sqlglot/dataframe/sql/functions.py b/sqlglot/dataframe/sql/functions.py
index 4d4f6f11..1549a075 100644
--- a/sqlglot/dataframe/sql/functions.py
+++ b/sqlglot/dataframe/sql/functions.py
@@ -1200,7 +1200,9 @@ def transform(
     f: t.Union[t.Callable[[Column], Column], t.Callable[[Column, Column], Column]],
 ) -> Column:
     f_expression = _get_lambda_from_func(f)
-    return Column.invoke_anonymous_function(col, "TRANSFORM", Column(f_expression))
+    return Column.invoke_expression_over_column(
+        col, expression.Transform, expression=Column(f_expression)
+    )
 
 
 def exists(col: ColumnOrName, f: t.Callable[[Column], Column]) -> Column:
diff --git a/sqlglot/dialects/hive.py b/sqlglot/dialects/hive.py
index 6bb4b0b1..e1314346 100644
--- a/sqlglot/dialects/hive.py
+++ b/sqlglot/dialects/hive.py
@@ -272,6 +272,11 @@ class Hive(Dialect):
             "YEAR": lambda args: exp.Year(this=exp.TsOrDsToDate.from_arg_list(args)),
         }
 
+        FUNCTION_PARSERS = {
+            **parser.Parser.FUNCTION_PARSERS,
+            "TRANSFORM": lambda self: self._parse_transform(),
+        }
+
         PROPERTY_PARSERS = {
             **parser.Parser.PROPERTY_PARSERS,
             "WITH SERDEPROPERTIES": lambda self: exp.SerdeProperties(
@@ -279,6 +284,40 @@ class Hive(Dialect):
             ),
         }
 
+        def _parse_transform(self) -> exp.Transform | exp.QueryTransform:
+            args = self._parse_csv(self._parse_lambda)
+            self._match_r_paren()
+
+            row_format_before = self._parse_row_format(match_row=True)
+
+            record_writer = None
+            if self._match_text_seq("RECORDWRITER"):
+                record_writer = self._parse_string()
+
+            if not self._match(TokenType.USING):
+                return exp.Transform.from_arg_list(args)
+
+            command_script = self._parse_string()
+
+            self._match(TokenType.ALIAS)
+            schema = self._parse_schema()
+
+            row_format_after = self._parse_row_format(match_row=True)
+            record_reader = None
+            if self._match_text_seq("RECORDREADER"):
+                record_reader = self._parse_string()
+
+            return self.expression(
+                exp.QueryTransform,
+                expressions=args,
+                command_script=command_script,
+                schema=schema,
+                row_format_before=row_format_before,
+                record_writer=record_writer,
+                row_format_after=row_format_after,
+                record_reader=record_reader,
+            )
+
         def _parse_types(
             self, check_func: bool = False, schema: bool = False
         ) -> t.Optional[exp.Expression]:
@@ -400,7 +439,6 @@ class Hive(Dialect):
             exp.UnixToTime: rename_func("FROM_UNIXTIME"),
             exp.UnixToTimeStr: rename_func("FROM_UNIXTIME"),
             exp.PartitionedByProperty: lambda self, e: f"PARTITIONED BY {self.sql(e, 'this')}",
-            exp.RowFormatSerdeProperty: lambda self, e: f"ROW FORMAT SERDE {self.sql(e, 'this')}",
             exp.SerdeProperties: lambda self, e: self.properties(e, prefix="WITH SERDEPROPERTIES"),
             exp.NumberToStr: rename_func("FORMAT_NUMBER"),
             exp.LastDateOfMonth: rename_func("LAST_DAY"),
@@ -414,6 +452,11 @@ class Hive(Dialect):
             exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
         }
 
+        def rowformatserdeproperty_sql(self, expression: exp.RowFormatSerdeProperty) -> str:
+            serde_props = self.sql(expression, "serde_properties")
+            serde_props = f" {serde_props}" if serde_props else ""
+            return f"ROW FORMAT SERDE {self.sql(expression, 'this')}{serde_props}"
+
         def arrayagg_sql(self, expression: exp.ArrayAgg) -> str:
             return self.func(
                 "COLLECT_LIST",
diff --git a/sqlglot/dialects/spark2.py b/sqlglot/dialects/spark2.py
index e1a113ca..016118e0 100644
--- a/sqlglot/dialects/spark2.py
+++ b/sqlglot/dialects/spark2.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import typing as t
 
-from sqlglot import exp, parser, transforms
+from sqlglot import exp, transforms
 from sqlglot.dialects.dialect import (
     create_with_partitions_sql,
     format_time_lambda,
@@ -142,7 +142,7 @@ class Spark2(Hive):
         }
 
         FUNCTION_PARSERS = {
-            **parser.Parser.FUNCTION_PARSERS,
+            **Hive.Parser.FUNCTION_PARSERS,
             "BROADCAST": lambda self: self._parse_join_hint("BROADCAST"),
             "BROADCASTJOIN": lambda self: self._parse_join_hint("BROADCASTJOIN"),
             "MAPJOIN": lambda self: self._parse_join_hint("MAPJOIN"),
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 64c0c0d0..401b1a0c 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2031,7 +2031,20 @@ class RowFormatDelimitedProperty(Property):
 
 
 class RowFormatSerdeProperty(Property):
-    arg_types = {"this": True}
+    arg_types = {"this": True, "serde_properties": False}
+
+
+# https://spark.apache.org/docs/3.1.2/sql-ref-syntax-qry-select-transform.html
+class QueryTransform(Expression):
+    arg_types = {
+        "expressions": True,
+        "command_script": True,
+        "schema": False,
+        "row_format_before": False,
+        "record_writer": False,
+        "row_format_after": False,
+        "record_reader": False,
+    }
 
 
 class SchemaCommentProperty(Property):
@@ -3876,6 +3889,11 @@ class Abs(Func):
     pass
 
 
+# https://spark.apache.org/docs/latest/api/sql/index.html#transform
+class Transform(Func):
+    arg_types = {"this": True, "expression": True}
+
+
 class Anonymous(Func):
     arg_types = {"this": True, "expressions": False}
     is_var_len_args = True
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index a2fe2a13..6142fc9a 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -2508,6 +2508,21 @@ class Generator:
 
         return self.func("ANY_VALUE", this)
 
+    def querytransform_sql(self, expression: exp.QueryTransform) -> str:
+        transform = self.func("TRANSFORM", *expression.expressions)
+        row_format_before = self.sql(expression, "row_format_before")
+        row_format_before = f" {row_format_before}" if row_format_before else ""
+        record_writer = self.sql(expression, "record_writer")
+        record_writer = f" RECORDWRITER {record_writer}" if record_writer else ""
+        using = f" USING {self.sql(expression, 'command_script')}"
+        schema = self.sql(expression, "schema")
+        schema = f" AS {schema}" if schema else ""
+        row_format_after = self.sql(expression, "row_format_after")
+        row_format_after = f" {row_format_after}" if row_format_after else ""
+        record_reader = self.sql(expression, "record_reader")
+        record_reader = f" RECORDREADER {record_reader}" if record_reader else ""
+        return f"{transform}{row_format_before}{record_writer}{using}{schema}{row_format_after}{record_reader}"
+
 
 def cached_generator(
     cache: t.Optional[t.Dict[int, str]] = None
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index ef696349..e784acb7 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -716,7 +716,7 @@ class Parser(metaclass=_Parser):
 
     FUNCTIONS_WITH_ALIASED_ARGS = {"STRUCT"}
 
-    FUNCTION_PARSERS: t.Dict[str, t.Callable] = {
+    FUNCTION_PARSERS = {
         "ANY_VALUE": lambda self: self._parse_any_value(),
         "CAST": lambda self: self._parse_cast(self.STRICT_CAST),
         "CONCAT": lambda self: self._parse_concat(),
@@ -1784,7 +1784,17 @@ class Parser(metaclass=_Parser):
             return None
 
         if self._match_text_seq("SERDE"):
-            return self.expression(exp.RowFormatSerdeProperty, this=self._parse_string())
+            this = self._parse_string()
+
+            serde_properties = None
+            if self._match(TokenType.SERDE_PROPERTIES):
+                serde_properties = self.expression(
+                    exp.SerdeProperties, expressions=self._parse_wrapped_csv(self._parse_property)
+                )
+
+            return self.expression(
+                exp.RowFormatSerdeProperty, this=this, serde_properties=serde_properties
+            )
 
         self._match_text_seq("DELIMITED")
 
@@ -3331,7 +3341,7 @@ class Parser(metaclass=_Parser):
             else:
                 this = self.expression(exp.Anonymous, this=this, expressions=args)
 
-        self._match_r_paren(this)
+        self._match(TokenType.R_PAREN, expression=this)
         return self._parse_window(this)
 
     def _parse_function_parameter(self) -> t.Optional[exp.Expression]:


Test patch:
diff --git a/tests/dialects/test_spark.py b/tests/dialects/test_spark.py
index 494cc154..2afa868d 100644
--- a/tests/dialects/test_spark.py
+++ b/tests/dialects/test_spark.py
@@ -224,6 +224,8 @@ TBLPROPERTIES (
         )
 
     def test_spark(self):
+        self.validate_identity("SELECT TRANSFORM(ARRAY(1, 2, 3), x -> x + 1)")
+        self.validate_identity("SELECT TRANSFORM(ARRAY(1, 2, 3), (x, i) -> x + i)")
         self.validate_identity("REFRESH table a.b.c")
         self.validate_identity("INTERVAL -86 days")
         self.validate_identity("SELECT UNIX_TIMESTAMP()")
@@ -502,3 +504,21 @@ TBLPROPERTIES (
             "CURRENT_USER()",
             write={"spark": "CURRENT_USER()"},
         )
+
+    def test_transform_query(self):
+        self.validate_identity("SELECT TRANSFORM(x) USING 'x' AS (x INT) FROM t")
+        self.validate_identity(
+            "SELECT TRANSFORM(zip_code, name, age) USING 'cat' AS (a, b, c) FROM person WHERE zip_code > 94511"
+        )
+        self.validate_identity(
+            "SELECT TRANSFORM(zip_code, name, age) USING 'cat' AS (a STRING, b STRING, c STRING) FROM person WHERE zip_code > 94511"
+        )
+        self.validate_identity(
+            "SELECT TRANSFORM(name, age) ROW FORMAT DELIMITED FIELDS TERMINATED BY ',' LINES TERMINATED BY '\n' NULL DEFINED AS 'NULL' USING 'cat' AS (name_age STRING) ROW FORMAT DELIMITED FIELDS TERMINATED BY '@' LINES TERMINATED BY '\n' NULL DEFINED AS 'NULL' FROM person"
+        )
+        self.validate_identity(
+            "SELECT TRANSFORM(zip_code, name, age) ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe' WITH SERDEPROPERTIES ('field.delim'='\t') USING 'cat' AS (a STRING, b STRING, c STRING) ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe' WITH SERDEPROPERTIES ('field.delim'='\t') FROM person WHERE zip_code > 94511"
+        )
+        self.validate_identity(
+            "SELECT TRANSFORM(zip_code, name, age) USING 'cat' FROM person WHERE zip_code > 94500"
+        )
diff --git a/tests/test_expressions.py b/tests/test_expressions.py
index 989220e0..1aab1c05 100644
--- a/tests/test_expressions.py
+++ b/tests/test_expressions.py
@@ -556,6 +556,7 @@ class TestExpressions(unittest.TestCase):
         self.assertIsInstance(parse_one("HEX(foo)"), exp.Hex)
         self.assertIsInstance(parse_one("TO_HEX(foo)", read="bigquery"), exp.Hex)
         self.assertIsInstance(parse_one("TO_HEX(MD5(foo))", read="bigquery"), exp.MD5)
+        self.assertIsInstance(parse_one("TRANSFORM(a, b)", read="spark"), exp.Transform)
 
     def test_column(self):
         column = parse_one("a.b.c.d")