fix(starrocks): exp.Unnest transpilation

use `arrays_zip` to merge multiple `Lateral view explode`

Author: lin.zhang

sqlglot/dialects/hive.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import typing as t
+from functools import partial
 
 from sqlglot import exp, generator, parser, tokens, transforms
 from sqlglot.dialects.dialect import (
@@ -485,7 +486,7 @@ class Generator(generator.Generator):
                 [
                     transforms.eliminate_qualify,
                     transforms.eliminate_distinct_on,
-                    transforms.unnest_to_explode,
+                    partial(transforms.unnest_to_explode, unnest_using_arrays_zip=False),
                 ]
             ),
             exp.Property: _property_sql,

sqlglot/dialects/spark2.py

@@ -190,6 +190,13 @@ class Generator(Hive.Generator):
 
         TRANSFORMS = {
             **Hive.Generator.TRANSFORMS,
+            exp.Select: transforms.preprocess(
+                [
+                    transforms.eliminate_qualify,
+                    transforms.eliminate_distinct_on,
+                    transforms.unnest_to_explode,
+                ]
+            ),
             exp.ApproxDistinct: rename_func("APPROX_COUNT_DISTINCT"),
             exp.ArraySum: lambda self,
             e: f"AGGREGATE({self.sql(e, 'this')}, 0, (acc, x) -> acc + x, acc -> acc)",

sqlglot/expressions.py

@@ -5331,6 +5331,12 @@ class Explode(Func):
     is_var_len_args = True
 
 
+# https://spark.apache.org/docs/latest/api/sql/#inline
+class Inline(Func):
+    arg_types = {"this": True, "expressions": False}
+    is_var_len_args = True
+
+
 class ExplodeOuter(Explode):
     pass
 

sqlglot/transforms.py

@@ -3,6 +3,7 @@
 import typing as t
 
 from sqlglot import expressions as exp
+from sqlglot.expressions import Func
 from sqlglot.helper import find_new_name, name_sequence
 
 if t.TYPE_CHECKING:
@@ -296,12 +297,16 @@ def unqualify_unnest(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
-def unnest_to_explode(expression: exp.Expression) -> exp.Expression:
+def unnest_to_explode(
+    expression: exp.Expression,
+    unnest_using_arrays_zip: bool = True,
+    support_unnest_func: bool = False,
+) -> exp.Expression:
     """Convert cross join unnest into lateral view explode."""
     if isinstance(expression, exp.Select):
         from_ = expression.args.get("from")
 
-        if from_ and isinstance(from_.this, exp.Unnest):
+        if not support_unnest_func and from_ and isinstance(from_.this, exp.Unnest):
             unnest = from_.this
             alias = unnest.args.get("alias")
             udtf = exp.Posexplode if unnest.args.get("offset") else exp.Explode
@@ -325,17 +330,30 @@ def unnest_to_explode(expression: exp.Expression) -> exp.Expression:
 
             if isinstance(unnest, exp.Unnest):
                 alias = join_expr.args.get("alias") if is_lateral else unnest.args.get("alias")
-                udtf = exp.Posexplode if unnest.args.get("offset") else exp.Explode
+                has_multi_expr = len(unnest.expressions) > 1
+                _udtf: type[Func] = exp.Posexplode if unnest.args.get("offset") else exp.Explode
 
                 expression.args["joins"].remove(join)
 
+                if unnest_using_arrays_zip and has_multi_expr:
+                    # Modify the logic to use arrays_zip if there are multiple expressions
+                    # Build arrays_zip with nested expressions correctly
+                    unnest.set(
+                        "expressions",
+                        [exp.Anonymous(this="arrays_zip", expressions=unnest.expressions)],
+                    )
+                    _udtf = exp.Inline
+
                 for e, column in zip(unnest.expressions, alias.columns if alias else []):
                     expression.append(
                         "laterals",
                         exp.Lateral(
-                            this=udtf(this=e),
+                            this=_udtf(this=e),
                             view=True,
-                            alias=exp.TableAlias(this=alias.this, columns=[column]),  # type: ignore
+                            alias=exp.TableAlias(
+                                this=alias.this,  # type: ignore
+                                columns=alias.columns if unnest_using_arrays_zip else [column],  # type: ignore
+                            ),
                         ),
                     )
 

tests/dialects/test_dialect.py

@@ -1399,7 +1399,8 @@ def test_cross_join(self):
             write={
                 "drill": "SELECT a, b FROM x CROSS JOIN UNNEST(y, z) AS t(a, b)",
                 "presto": "SELECT a, b FROM x CROSS JOIN UNNEST(y, z) AS t(a, b)",
-                "spark": "SELECT a, b FROM x LATERAL VIEW EXPLODE(y) t AS a LATERAL VIEW EXPLODE(z) t AS b",
+                "spark": "SELECT a, b FROM x LATERAL VIEW INLINE(ARRAYS_ZIP(y, z)) t AS a, b",
+                "hive": "SELECT a, b FROM x LATERAL VIEW EXPLODE(y) t AS a LATERAL VIEW EXPLODE(z) t AS b",
             },
         )
         self.validate_all(

tests/dialects/test_starrocks.py

@@ -36,26 +36,32 @@ def test_unnest(self):
             "SELECT student, score, t.unnest FROM tests CROSS JOIN LATERAL UNNEST(scores) AS t(unnest)",
         )
         self.validate_identity(
-            "SELECT student, score, unnest.unnest FROM tests CROSS JOIN LATERAL UNNEST(scores)",
-            "SELECT student, score, unnest.unnest FROM tests CROSS JOIN LATERAL UNNEST(scores) AS unnest(unnest)",
+            "SELECT student, score, unnest FROM tests CROSS JOIN LATERAL UNNEST(scores)",
+            "SELECT student, score, unnest FROM tests CROSS JOIN LATERAL UNNEST(scores) AS unnest(unnest)",
         )
         self.validate_all(
             r"""SELECT * FROM UNNEST(array['John','Jane','Jim','Jamie'], array[24,25,26,27]) AS t(name, age)""",
             write={
                 "postgres": "SELECT * FROM UNNEST(ARRAY['John', 'Jane', 'Jim', 'Jamie'], ARRAY[24, 25, 26, 27]) AS t(name, age)",
-                "spark": "SELECT * FROM EXPLODE(ARRAY('John', 'Jane', 'Jim', 'Jamie'), ARRAY(24, 25, 26, 27)) AS "
-                "t(name, age)",
+                "spark": "SELECT * FROM EXPLODE(ARRAY('John', 'Jane', 'Jim', 'Jamie'), ARRAY(24, 25, 26, 27)) AS t(name, age)",
             },
         )
         # Use UNNEST to convert into multiple columns
         # see: https://docs.starrocks.io/docs/sql-reference/sql-functions/array-functions/unnest/
         self.validate_all(
             r"""SELECT id, t.type, t.scores FROM example_table, unnest(split(type, ";"), scores) AS t(type,scores)""",
             write={
-                "postgres": "SELECT id, t.type, t.scores FROM example_table, UNNEST(SPLIT(type, ';'), scores) AS "
-                "t(type, scores)",
-                "spark": "SELECT id, t.type, t.scores FROM example_table LATERAL VIEW EXPLODE(SPLIT(type, CONCAT"
+                "postgres": "SELECT id, t.type, t.scores FROM example_table, UNNEST(SPLIT(type, ';'), scores) AS t(type, scores)",
+                "spark": r"""SELECT id, t.type, t.scores FROM example_table LATERAL VIEW INLINE(ARRAYS_ZIP(SPLIT(type, CONCAT('\\Q', ';')), scores)) t AS type, scores""",
+                "hive": "SELECT id, t.type, t.scores FROM example_table LATERAL VIEW EXPLODE(SPLIT(type, CONCAT"
                 r"""('\\Q', ';'))) t AS type LATERAL VIEW EXPLODE(scores) t AS scores""",
+                "databricks": r"""SELECT id, t.type, t.scores FROM example_table LATERAL VIEW INLINE(ARRAYS_ZIP(SPLIT(type, CONCAT('\\Q', ';')), scores)) t AS type, scores""",
+            },
+        )
+        self.validate_all(
+            r"""SELECT id, t.type, t.scores FROM example_table CROSS JOIN LATERAL unnest(split(type, ";"), scores) AS t(type,scores)""",
+            write={
+                "spark": r"""SELECT id, t.type, t.scores FROM example_table LATERAL VIEW INLINE(ARRAYS_ZIP(SPLIT(type, CONCAT('\\Q', ';')), scores)) t AS type, scores""",
             },
         )