feat: snowflake array_construct_compact to spark

tobymao · May 17, 2024 · c811adb · c811adb
1 parent 5b64475
commit c811adb
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 0 deletions.
diff --git a/sqlglot/dialects/spark.py b/sqlglot/dialects/spark.py
@@ -96,6 +96,9 @@ class Generator(Spark2.Generator):
 
         TRANSFORMS = {
             **Spark2.Generator.TRANSFORMS,
+            exp.ArrayConstructCompact: lambda self, e: self.func(
+                "ARRAY_COMPACT", self.func("ARRAY", *e.expressions)
+            ),
             exp.Create: preprocess(
                 [
                     remove_unique_constraints,

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -4775,6 +4775,11 @@ class ArrayConcat(Func):
     is_var_len_args = True
 
 
+class ArrayConstructCompact(Func):
+    arg_types = {"expressions": True}
+    is_var_len_args = True
+
+
 class ArrayContains(Binary, Func):
     pass
 

diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
@@ -10,6 +10,14 @@ class TestSnowflake(Validator):
     dialect = "snowflake"
 
     def test_snowflake(self):
+        self.validate_all(
+            "ARRAY_CONSTRUCT_COMPACT(1, null, 2)",
+            write={
+                "spark": "ARRAY_COMPACT(ARRAY(1, NULL, 2))",
+                "snowflake": "ARRAY_CONSTRUCT_COMPACT(1, NULL, 2)",
+            },
+        )
+
         expr = parse_one("SELECT APPROX_TOP_K(C4, 3, 5) FROM t")
         expr.selects[0].assert_is(exp.AggFunc)
         self.assertEqual(expr.sql(dialect="snowflake"), "SELECT APPROX_TOP_K(C4, 3, 5) FROM t")