Repository: tobymao/sqlglot
Pull request: Add support for snowflake's flatten function
fixes #675 

Problem statement:
Difficulty parsing snowflake query with `split()` nested in `flatten()`
Hello!

I've been unable to get sqlglot to parse the following query and I'm not sure if it's user error or a bug? Would love to get this resolved:

```python

import sqlglot
import sqlglot.expressions as exp

query = f"""
select
  dag_report.acct_id,
  dag_report.report_date,
  dag_report.report_uuid,
  dag_report.airflow_name,
  dag_report.dag_id,
  f.value::varchar as operator
from cs.telescope.dag_report,
table(flatten(input=>split(operators, ','))) f
"""

def get_query_columns(query):
    column_names = []
    for expression in sqlglot.parse_one(query, read="snowflake").find(exp.Select).args["expressions"]:
        if isinstance(expression, exp.Alias):
            column_names.append(expression.text("alias"))
        elif isinstance(expression, exp.Column):
            column_names.append(expression.text("this"))

    return column_names

if __name__ == '__main__':
    column_names = get_query_columns(query)
    print(column_names)

```

Error:

```text
sqlglot.errors.ParseError: Required keyword: 'this' missing for <class 'sqlglot.expressions.GT'>. Line 10, Col: 43.
```

Any help would be greatly appreciated. 🍻

Hints from issue discussion:
Hello, thanks for the report.

This error happens because `=>` is not currently parsed properly to support Snowflake's `flatten` [syntax](https://docs.snowflake.com/en/sql-reference/functions/flatten.html#syntax).

I'll take a look at it soon.

Commit message: Add support for snowflake's flatten function

Patch to review:
diff --git a/sqlglot/dataframe/sql/functions.py b/sqlglot/dataframe/sql/functions.py
index dbfb06f30b..9863d56ddc 100644
--- a/sqlglot/dataframe/sql/functions.py
+++ b/sqlglot/dataframe/sql/functions.py
@@ -1078,7 +1078,7 @@ def reverse(col: ColumnOrName) -> Column:
 
 
 def flatten(col: ColumnOrName) -> Column:
-    return Column.invoke_anonymous_function(col, "FLATTEN")
+    return Column.invoke_expression_over_column(col, glotexp.Flatten)
 
 
 def map_keys(col: ColumnOrName) -> Column:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 9149ef4b3d..2eff6c3172 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -58,7 +58,7 @@ def _snowflake_to_timestamp(args):
     return exp.UnixToTime.from_arg_list(args)
 
 
-def _unix_to_time(self, expression):
+def _unix_to_time_sql(self, expression):
     scale = expression.args.get("scale")
     timestamp = self.sql(expression, "this")
     if scale in [None, exp.UnixToTime.SECONDS]:
@@ -71,6 +71,18 @@ def _unix_to_time(self, expression):
     raise ValueError("Improper scale for timestamp")
 
 
+def _flatten_sql(self, expression):
+    args = [
+        f"{'INPUT' if key == 'this' else key.upper()} => {self.sql(arg)}"
+        for key, arg in expression.args.items()
+    ]
+
+    if self.pretty and self.text_width(args) > self._max_text_width:
+        nl = "\n"
+        return f"FLATTEN({self.indent({nl} + f',{nl}'.join(args) + {nl}, skip_first=True, skip_last=True)})"
+    return f'FLATTEN({", ".join(args)})'
+
+
 # https://docs.snowflake.com/en/sql-reference/functions/date_part.html
 # https://docs.snowflake.com/en/sql-reference/functions-date-time.html#label-supported-date-time-parts
 def _parse_date_part(self):
@@ -100,6 +112,18 @@ def _parse_date_part(self):
     return self.expression(exp.Extract, this=this, expression=expression)
 
 
+# https://docs.snowflake.com/en/sql-reference/functions/flatten.html#syntax
+def _parse_flatten(self):
+    def _parse_flatten_arg():
+        self._match_set(self.FLATTEN_ARGS)
+        arg_name = "this" if self._prev.token_type == TokenType.INPUT else self._prev.text.lower()
+        self._match(TokenType.ARROW)
+        return arg_name, self._parse_expression()
+
+    args = dict(self._parse_csv(_parse_flatten_arg))
+    return self.expression(exp.Flatten, **args)
+
+
 class Snowflake(Dialect):
     null_ordering = "nulls_are_large"
     time_format = "'yyyy-mm-dd hh24:mi:ss'"
@@ -135,6 +159,14 @@ class Snowflake(Dialect):
     }
 
     class Parser(parser.Parser):
+        FLATTEN_ARGS = {
+            TokenType.INPUT,
+            TokenType.PATH,
+            TokenType.OUTER,
+            TokenType.RECURSIVE,
+            TokenType.MODE,
+        }
+
         FUNCTIONS = {
             **parser.Parser.FUNCTIONS,
             "ARRAYAGG": exp.ArrayAgg.from_arg_list,
@@ -147,6 +179,7 @@ class Parser(parser.Parser):
         FUNCTION_PARSERS = {
             **parser.Parser.FUNCTION_PARSERS,
             "DATE_PART": _parse_date_part,
+            "FLATTEN": _parse_flatten,
         }
 
         FUNC_TOKENS = {
@@ -180,6 +213,7 @@ class Tokenizer(tokens.Tokenizer):
 
         KEYWORDS = {
             **tokens.Tokenizer.KEYWORDS,
+            "=>": TokenType.ARROW,
             "QUALIFY": TokenType.QUALIFY,
             "DOUBLE PRECISION": TokenType.DOUBLE,
             "TIMESTAMP_LTZ": TokenType.TIMESTAMPLTZ,
@@ -197,12 +231,13 @@ class Generator(generator.Generator):
             exp.ArrayConcat: rename_func("ARRAY_CAT"),
             exp.If: rename_func("IFF"),
             exp.StrToTime: lambda self, e: f"TO_TIMESTAMP({self.sql(e, 'this')}, {self.format_time(e)})",
-            exp.UnixToTime: _unix_to_time,
+            exp.UnixToTime: _unix_to_time_sql,
             exp.TimeToUnix: lambda self, e: f"EXTRACT(epoch_second FROM {self.sql(e, 'this')})",
             exp.Array: inline_array_sql,
             exp.StrPosition: rename_func("POSITION"),
             exp.Parameter: lambda self, e: f"${self.sql(e, 'this')}",
             exp.PartitionedByProperty: lambda self, e: f"PARTITION BY {self.sql(e, 'value')}",
+            exp.Flatten: _flatten_sql,
         }
 
         TYPE_MAPPING = {
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 3b64f06a94..bf180a9bbb 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2553,6 +2553,11 @@ class Explode(Func):
     pass
 
 
+# https://docs.snowflake.com/en/sql-reference/functions/flatten.html#syntax
+class Flatten(Func):
+    arg_types = {"this": True, "path": False, "outer": False, "recursive": False, "mode": False}
+
+
 class Floor(Func):
     arg_types = {"this": True, "decimals": False}
 
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
index b349828842..545e6aed03 100644
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -184,6 +184,7 @@ class TokenType(AutoName):
     IN = auto()
     INDEX = auto()
     INNER = auto()
+    INPUT = auto()
     INSERT = auto()
     INTERSECT = auto()
     INTERVAL = auto()
@@ -205,6 +206,7 @@ class TokenType(AutoName):
     MAP = auto()
     MATERIALIZED = auto()
     MOD = auto()
+    MODE = auto()
     NATURAL = auto()
     NEXT = auto()
     NO_ACTION = auto()
@@ -227,6 +229,7 @@ class TokenType(AutoName):
     PARTITION = auto()
     PARTITION_BY = auto()
     PARTITIONED_BY = auto()
+    PATH = auto()
     PERCENT = auto()
     PIVOT = auto()
     PLACEHOLDER = auto()
@@ -509,6 +512,7 @@ class Tokenizer(metaclass=_Tokenizer):
         "IN": TokenType.IN,
         "INDEX": TokenType.INDEX,
         "INNER": TokenType.INNER,
+        "INPUT": TokenType.INPUT,
         "INSERT": TokenType.INSERT,
         "INTERVAL": TokenType.INTERVAL,
         "INTERSECT": TokenType.INTERSECT,
@@ -527,6 +531,7 @@ class Tokenizer(metaclass=_Tokenizer):
         "LOCAL": TokenType.LOCAL,
         "LOCATION": TokenType.LOCATION,
         "MATERIALIZED": TokenType.MATERIALIZED,
+        "MODE": TokenType.MODE,
         "NATURAL": TokenType.NATURAL,
         "NEXT": TokenType.NEXT,
         "NO ACTION": TokenType.NO_ACTION,
@@ -552,6 +557,7 @@ class Tokenizer(metaclass=_Tokenizer):
         "PARTITION BY": TokenType.PARTITION_BY,
         "PARTITIONED BY": TokenType.PARTITIONED_BY,
         "PARTITIONED_BY": TokenType.PARTITIONED_BY,
+        "PATH": TokenType.PATH,
         "PERCENT": TokenType.PERCENT,
         "PIVOT": TokenType.PIVOT,
         "PRECEDING": TokenType.PRECEDING,
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 5702986cdc..757b229fec 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -371,3 +371,105 @@ def test_table_literal(self):
             r"""SELECT * FROM TABLE($MYVAR) WHERE COL1 = 10""",
             write={"snowflake": r"""SELECT * FROM TABLE($MYVAR) WHERE COL1 = 10"""},
         )
+
+    def test_flatten(self):
+        self.validate_all(
+            """
+            select
+              dag_report.acct_id,
+              dag_report.report_date,
+              dag_report.report_uuid,
+              dag_report.airflow_name,
+              dag_report.dag_id,
+              f.value::varchar as operator
+            from cs.telescope.dag_report,
+            table(flatten(input=>split(operators, ','))) f
+            """,
+            write={
+                "snowflake": """SELECT
+  dag_report.acct_id,
+  dag_report.report_date,
+  dag_report.report_uuid,
+  dag_report.airflow_name,
+  dag_report.dag_id,
+  CAST(f.value AS VARCHAR) AS operator
+FROM cs.telescope.dag_report, TABLE(FLATTEN(INPUT => SPLIT(operators, ','))) AS f"""
+            },
+            pretty=True,
+        )
+
+        # All examples from https://docs.snowflake.com/en/sql-reference/functions/flatten.html#syntax
+        self.validate_all(
+            "SELECT * FROM TABLE(FLATTEN(input => parse_json('[1, ,77]'))) f",
+            write={
+                "snowflake": "SELECT * FROM TABLE(FLATTEN(INPUT => PARSE_JSON('[1, ,77]'))) AS f"
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88]}'), outer => true)) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(INPUT => PARSE_JSON('{"a":1, "b":[77,88]}'), OUTER => TRUE)) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88]}'), path => 'b')) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(INPUT => PARSE_JSON('{"a":1, "b":[77,88]}'), PATH => 'b')) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('[]'))) f""",
+            write={"snowflake": """SELECT * FROM TABLE(FLATTEN(INPUT => PARSE_JSON('[]'))) AS f"""},
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('[]'), outer => true)) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(INPUT => PARSE_JSON('[]'), OUTER => TRUE)) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88], "c": {"d":"X"}}'))) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(INPUT => PARSE_JSON('{"a":1, "b":[77,88], "c": {"d":"X"}}'))) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88], "c": {"d":"X"}}'), recursive => true)) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(INPUT => PARSE_JSON('{"a":1, "b":[77,88], "c": {"d":"X"}}'), RECURSIVE => TRUE)) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """SELECT * FROM TABLE(FLATTEN(input => parse_json('{"a":1, "b":[77,88], "c": {"d":"X"}}'), recursive => true, mode => 'object')) f""",
+            write={
+                "snowflake": """SELECT * FROM TABLE(FLATTEN(INPUT => PARSE_JSON('{"a":1, "b":[77,88], "c": {"d":"X"}}'), RECURSIVE => TRUE, MODE => 'object')) AS f"""
+            },
+        )
+
+        self.validate_all(
+            """
+            SELECT id as "ID",
+              f.value AS "Contact",
+              f1.value:type AS "Type",
+              f1.value:content AS "Details"
+            FROM persons p,
+              lateral flatten(input => p.c, path => 'contact') f,
+              lateral flatten(input => f.value:business) f1
+            """,
+            write={
+                "snowflake": """SELECT
+  id AS "ID",
+  f.value AS "Contact",
+  f1.value['type'] AS "Type",
+  f1.value['content'] AS "Details"
+FROM persons AS p, LATERAL FLATTEN(INPUT => p.c, PATH => 'contact') f, LATERAL FLATTEN(INPUT => f.value['business']) f1""",
+            },
+            pretty=True,
+        )


Code review comments:
diff_hunk: @@ -71,6 +71,18 @@ def _unix_to_time(self, expression):
     raise ValueError("Improper scale for timestamp")
 
 
+def _flatten_sql(self, expression):
+    args = [
+        f"{'INPUT' if key == 'this' else key.upper()} => {self.sql(arg)}"
+        for key, arg in expression.args.items()
+    ]
+
+    if self.pretty and self.text_width(args) > self._max_text_width:
+        nl = "\n"
+        return f"FLATTEN({self.indent({nl} + f',{nl}'.join(args) + {nl}, skip_first=True, skip_last=True)})"
+    return f'FLATTEN({", ".join(args)})'
path: sqlglot/dialects/snowflake.py
text: @author:
I was trying to use `format_args` for this, because it's essentially the same logic, but with different sql formatting. Maybe we can dry this up by adding a `sql_lambda` argument to the `format_args` which takes care of the formatting, so we can reuse it for situations like this. For example, perhaps we could do:

```python
    def format_args(self, *args, sql_lambda=None):
        sql_lambda = sql_lambda or self.sql
        args = tuple(sql_lambda(arg) for arg in args if arg is not None)
        if self.pretty and self.text_width(args) > self._max_text_width:
            return self.indent("\n" + f",\n".join(args) + "\n", skip_first=True, skip_last=True)
        return ", ".join(args)
```
So then `_flatten_sql` would become:

```python
def _flatten_sql(self, expression):
    sql_lambda = lambda key_arg: f"{key_arg[0]} => {key_arg[1]}"
    return self.format_args(*[(key, arg) for key, arg in expression.args.items()], sql_lambda=sql_lambda)
```