Repository: tobymao/sqlglot

Issue:
Support lineage through a LATERAL FLATTEN from clause in Snowflake
**Is your feature request related to a problem? Please describe.**
I have a Snowflake SQL query like `SELECT VALUE FROM TABLE_A, LATERAL FLATTEN(TABLE_A.VARIANT_COLUMN)`. I would like to be able to use the `lineage` function to get the lineage tree of the column "VALUE" (in case TABLE_A is actually a subquery, cte, etc).

**Describe the solution you'd like**
lineage should understand that VALUE's lineage is through TABLE_A.VARIANT_COLUMN (or through a subquery or CTE in place of TABLE_A). The `FLATTEN` function, when the Snowflake dialect is selected, should be treated as a select expression that adds the columns documented as [its output](https://docs.snowflake.com/en/sql-reference/functions/flatten#output).

The lineage tree of value should look something like VALUE -> FLATTEN(TABLE_A>VARIANT_COLUMN) -> TABLE_A.VARIANT_COLUMN -> TABLE_A. 

**Describe alternatives you've considered**
My plan was to have the Snowflake dialect parser produce a TableAlias expression when it parses a FLATTEN function populated with with the columns produces by the expression. Right now a TableAlias is produced, but it has no columns, so lineage fails raising a ValueError.

I'm happy to work on this feature if it's something you want to support in sqlglot but wanted to open this issue to make sure my suggestion of an implementation strategy make sense and is something you'd accept. 



Discussion:
I see, what you're suggesting looks reasonable. Just one comment:

> My plan was to have the Snowflake dialect parser produce a TableAlias expression when it parses a FLATTEN function populated with with the columns produces by the expression. Right now a TableAlias is produced, but it has no columns, so lineage fails raising a ValueError.

I'm not sure if we want to touch Snowflake's parser for this, even though we could. My initial reaction was that this can be solved by adding extra logic to the `qualify_tables` rule in order to emit a `TableAlias` with additional info, i.e. the output columns specified in Snowflake's docs. The downside to this is that we pollute the rule with Snowflake-specific logic, because there's no other way easy way to parameterize how we do stuff per dialect in the optimizer currently.

@tobymao what do you think?
shouldn't this be done in qualify_columns? it already deals with table column aliases
Not sure, what do we gain from doing it there? I figured that `qualify_tables` would allow it to just work for the rest of the optimizer rules. It also feels more natural because we're trying to add more info to a node associated with a table as opposed to, say, expanding stars etc which fits in `qualify_columns` better.
Doing it in `qualify_tables` makes sense to me after a quick glance. I had looked at the parser since that seemed like the place to switch on dialect.
@edpaget are you planning to work on this?

Resolving patch:
diff --git a/sqlglot/dataframe/sql/functions.py b/sqlglot/dataframe/sql/functions.py
--- a/sqlglot/dataframe/sql/functions.py
+++ b/sqlglot/dataframe/sql/functions.py
@@ -1113,7 +1113,7 @@ def reverse(col: ColumnOrName) -> Column:
 
 
 def flatten(col: ColumnOrName) -> Column:
-    return Column.invoke_anonymous_function(col, "FLATTEN")
+    return Column.invoke_expression_over_column(None, expression.Flatten, expressions=[col])
 
 
 def map_keys(col: ColumnOrName) -> Column:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -262,6 +262,7 @@ class Parser(parser.Parser):
             ),
             "DATEDIFF": _parse_datediff,
             "DIV0": _div0_to_if,
+            "FLATTEN": exp.Explode.from_arg_list,
             "IFF": exp.If.from_arg_list,
             "LISTAGG": exp.GroupConcat.from_arg_list,
             "NULLIFZERO": _nullifzero_to_if,
@@ -325,6 +326,22 @@ class Parser(parser.Parser):
             TokenType.MOD,
             TokenType.SLASH,
         }
+        FLATTEN_COLUMNS = ["SEQ", "KEY", "PATH", "INDEX", "VALUE", "THIS"]
+
+        def _parse_lateral(self) -> t.Optional[exp.Lateral]:
+            lateral = super()._parse_lateral()
+            if not lateral:
+                return lateral
+
+            if isinstance(lateral.this, exp.Explode):
+                table_alias = lateral.args.get("alias")
+                columns = [exp.to_identifier(col) for col in self.FLATTEN_COLUMNS]
+                if table_alias and not table_alias.args.get("columns"):
+                    table_alias.set("columns", columns)
+                elif not table_alias:
+                    exp.alias_(lateral, "_flattened", table=columns, copy=False)
+
+            return lateral
 
         def _parse_table_parts(self, schema: bool = False) -> exp.Table:
             # https://docs.snowflake.com/en/user-guide/querying-stage
@@ -452,6 +469,7 @@ class Generator(generator.Generator):
             exp.DateStrToDate: datestrtodate_sql,
             exp.DataType: _datatype_sql,
             exp.DayOfWeek: rename_func("DAYOFWEEK"),
+            exp.Explode: rename_func("FLATTEN"),
             exp.Extract: rename_func("DATE_PART"),
             exp.GenerateSeries: lambda self, e: self.func(
                 "ARRAY_GENERATE_RANGE", e.args["start"], e.args["end"] + 1, e.args.get("step")
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -4168,6 +4168,11 @@ class Abs(Func):
     pass
 
 
+class Flatten(Func):
+    arg_types = {"expressions": True}
+    is_var_len_args = True
+
+
 # https://spark.apache.org/docs/latest/api/sql/index.html#transform
 class Transform(Func):
     arg_types = {"this": True, "expression": True}
@@ -4540,8 +4545,10 @@ class Exp(Func):
     pass
 
 
+# https://docs.snowflake.com/en/sql-reference/functions/flatten
 class Explode(Func):
-    pass
+    arg_types = {"this": True, "expressions": False}
+    is_var_len_args = True
 
 
 class ExplodeOuter(Explode):


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
@@ -911,7 +911,23 @@ def test_flatten(self):
   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') AS f, LATERAL FLATTEN(input => f.value['business']) AS f1""",
+FROM persons AS p, LATERAL FLATTEN(input => p.c, path => 'contact') AS f(SEQ, KEY, PATH, INDEX, VALUE, THIS), LATERAL FLATTEN(input => f.value['business']) AS f1(SEQ, KEY, PATH, INDEX, VALUE, THIS)""",
+            },
+            pretty=True,
+        )
+
+        self.validate_all(
+            """
+            SELECT id as "ID",
+              value AS "Contact"
+            FROM persons p,
+              lateral flatten(input => p.c, path => 'contact')
+            """,
+            write={
+                "snowflake": """SELECT
+  id AS "ID",
+  value AS "Contact"
+FROM persons AS p, LATERAL FLATTEN(input => p.c, path => 'contact') AS _flattened(SEQ, KEY, PATH, INDEX, VALUE, THIS)""",
             },
             pretty=True,
         )
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -199,3 +199,22 @@ def test_lineage_union(self) -> None:
             "SELECT x FROM (SELECT ax AS x FROM a UNION SELECT bx FROM b UNION SELECT cx FROM c)",
         )
         assert len(node.downstream) == 3
+
+    def test_lineage_lateral_flatten(self) -> None:
+        node = lineage(
+            "VALUE",
+            "SELECT FLATTENED.VALUE FROM TEST_TABLE, LATERAL FLATTEN(INPUT => RESULT, OUTER => TRUE) FLATTENED",
+            dialect="snowflake",
+        )
+        self.assertEqual(node.name, "VALUE")
+
+        downstream = node.downstream[0]
+        self.assertEqual(downstream.name, "FLATTENED.VALUE")
+        self.assertEqual(
+            downstream.source.sql(dialect="snowflake"),
+            "LATERAL FLATTEN(INPUT => TEST_TABLE.RESULT, OUTER => TRUE) AS FLATTENED(SEQ, KEY, PATH, INDEX, VALUE, THIS)",
+        )
+        self.assertEqual(
+            downstream.expression.sql(dialect="snowflake"),
+            "VALUE",
+        )