Fix: make map gen more robust

tobymao · May 26, 2023 · 8f0fbad · 8f0fbad
1 parent 425af88
commit 8f0fbad
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/sqlglot/executor/env.py b/sqlglot/executor/env.py
@@ -169,7 +169,7 @@ def interval(this, unit):
     "LOWER": null_if_any(lambda arg: arg.lower()),
     "LT": null_if_any(lambda this, e: this < e),
     "LTE": null_if_any(lambda this, e: this <= e),
-    "MAP": null_if_any(lambda k, v: dict(zip(k, v))),
+    "MAP": null_if_any(lambda *args: dict(zip(*args))),  # type: ignore
     "MOD": null_if_any(lambda e, this: e % this),
     "MUL": null_if_any(lambda e, this: e * this),
     "NEQ": null_if_any(lambda this, e: this != e),

diff --git a/sqlglot/executor/python.py b/sqlglot/executor/python.py
@@ -360,11 +360,19 @@ def _ordered_py(self, expression):
 
 def _rename(self, e):
     try:
-        if "expressions" in e.args:
-            this = self.sql(e, "this")
-            this = f"{this}, " if this else ""
-            return f"{e.key.upper()}({this}{self.expressions(e)})"
-        return self.func(e.key, *e.args.values())
+        values = list(e.args.values())
+
+        if len(values) == 1:
+            values = values[0]
+            if not isinstance(values, list):
+                return self.func(e.key, values)
+            return self.func(e.key, *values)
+
+        if isinstance(e, exp.Func) and e.is_var_len_args:
+            *head, tail = values
+            return self.func(e.key, *head, *tail)
+
+        return self.func(e.key, *values)
     except Exception as ex:
         raise Exception(f"Could not rename {repr(e)}") from ex
 

diff --git a/tests/test_executor.py b/tests/test_executor.py
@@ -58,6 +58,7 @@ def rename_anonymous(self, source, target):
 
     def test_py_dialect(self):
         self.assertEqual(Python().generate(parse_one("'x '''")), r"'x \''")
+        self.assertEqual(Python().generate(parse_one("MAP([1], [2])")), "MAP([1], [2])")
 
     def test_optimized_tpch(self):
         for i, (sql, optimized) in enumerate(self.sqls[:20], start=1):
@@ -596,6 +597,7 @@ def test_scalar_functions(self):
             ("1::bool", True),
             ("0::bool", False),
             ("MAP(['a'], [1]).a", 1),
+            ("MAP()", {}),
             ("STRFTIME('%j', '2023-03-23 15:00:00')", "082"),
             ("STRFTIME('%j', NULL)", None),
             ("DATESTRTODATE('2022-01-01')", date(2022, 1, 1)),