From c0e376135ce522ec318a94c9cf4d2fea77b2139c Mon Sep 17 00:00:00 2001
From: Nick Crews <nicholas.b.crews@gmail.com>
Date: Fri, 15 Mar 2024 08:45:47 -0800
Subject: [PATCH] fix: add workaround for NULL map bugs

partially addresses https://github.com/ibis-project/ibis/issues/8632

postgres started failing these new test too,
so had to adjust that visit_Map()
---
 ibis/backends/duckdb/compiler.py          |  6 ++++++
 ibis/backends/duckdb/tests/test_client.py |  8 ++++++++
 ibis/backends/postgres/compiler.py        |  5 ++++-
 ibis/backends/tests/test_map.py           | 24 ++++++++++++++++++++++-
 4 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/ibis/backends/duckdb/compiler.py b/ibis/backends/duckdb/compiler.py
index d2f8280adb597..d3f6ca6dabc19 100644
--- a/ibis/backends/duckdb/compiler.py
+++ b/ibis/backends/duckdb/compiler.py
@@ -199,6 +199,12 @@ def visit_ArrayZip(self, op, *, arg):
         any_arg_null = sg.or_(*(arr.is_(NULL) for arr in arg))
         return self.if_(any_arg_null, NULL, zipped_arrays)
 
+    def visit_Map(self, op, *, keys, values):
+        # workaround for https://github.com/ibis-project/ibis/issues/8632
+        regular = self.f.map(keys, values)
+        either_null = sg.or_(keys.is_(NULL), values.is_(NULL))
+        return self.if_(either_null, NULL, regular)
+
     def visit_MapGet(self, op, *, arg, key, default):
         return self.f.ifnull(
             self.f.list_extract(self.f.element_at(arg, key), 1), default
diff --git a/ibis/backends/duckdb/tests/test_client.py b/ibis/backends/duckdb/tests/test_client.py
index 6f82c48c2f841..6d8b1fd9ddafe 100644
--- a/ibis/backends/duckdb/tests/test_client.py
+++ b/ibis/backends/duckdb/tests/test_client.py
@@ -276,3 +276,11 @@ def test_invalid_connect(tmp_path):
     url = f"duckdb://{tmp_path}?read_only=invalid_value"
     with pytest.raises(ValueError):
         ibis.connect(url)
+
+
+@pytest.mark.xfail(
+    raises=(duckdb.InvalidInputException, duckdb.BinderException),
+    reason="https://github.com/ibis-project/ibis/issues/8632",
+)
+def test_map_null_workaround_xfail(con):
+    con.con.sql("SELECT MAP([0, 1, 2], NULL::INT[]);").fetchone()
diff --git a/ibis/backends/postgres/compiler.py b/ibis/backends/postgres/compiler.py
index 7d794984375b3..12b89829a3f92 100644
--- a/ibis/backends/postgres/compiler.py
+++ b/ibis/backends/postgres/compiler.py
@@ -330,7 +330,10 @@ def visit_ToJSONArray(self, op, *, arg):
         )
 
     def visit_Map(self, op, *, keys, values):
-        return self.f.map(self.f.array(*keys), self.f.array(*values))
+        # map(["a", "b"], NULL) results in {"a": NULL, "b": NULL} in regular postgres,
+        # so we need to modify it to return NULL instead
+        regular = self.f.map(keys, values)
+        return self.if_(values.is_(NULL), NULL, regular)
 
     def visit_MapLength(self, op, *, arg):
         return self.f.cardinality(self.f.akeys(arg))
diff --git a/ibis/backends/tests/test_map.py b/ibis/backends/tests/test_map.py
index 28ac4010ab82b..6f375bbaec5b0 100644
--- a/ibis/backends/tests/test_map.py
+++ b/ibis/backends/tests/test_map.py
@@ -10,7 +10,7 @@
 import ibis
 import ibis.common.exceptions as exc
 import ibis.expr.datatypes as dt
-from ibis.backends.tests.errors import Py4JJavaError
+from ibis.backends.tests.errors import PsycoPg2InternalError, Py4JJavaError
 
 pytestmark = [
     pytest.mark.never(
@@ -39,6 +39,28 @@
 )
 
 
+@pytest.mark.notyet("clickhouse", reason="nested types can't be NULL")
+@pytest.mark.broken(["pandas", "dask"], reason="TypeError: iteration over a 0-d array")
+@pytest.mark.notimpl(
+    ["risingwave"],
+    raises=PsycoPg2InternalError,
+    reason="function hstore(character varying[], character varying[]) does not exist",
+)
+@pytest.mark.parametrize(
+    ("k", "v"),
+    [
+        param(["a", "b"], None, id="null_values"),
+        param(None, ["c", "d"], id="null_keys"),
+        param(None, None, id="null_both"),
+    ],
+)
+def test_map_nulls(con, k, v):
+    k = ibis.literal(k, type="array<string>")
+    v = ibis.literal(v, type="array<string>")
+    m = ibis.map(k, v)
+    assert con.execute(m) is None
+
+
 @pytest.mark.notimpl(["pandas", "dask"])
 def test_map_table(backend):
     table = backend.map