fix(flink): cast inputs in map contains impl

ibis-project · Mar 21, 2024 · 3e9b9df · 3e9b9df
1 parent b8f7296
commit 3e9b9df
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/ibis/backends/flink/compiler.py b/ibis/backends/flink/compiler.py
@@ -546,7 +546,11 @@ def visit_CountDistinct(self, op, *, arg, where):
         return self.f.count(sge.Distinct(expressions=[arg]))
 
     def visit_MapContains(self, op: ops.MapContains, *, arg, key):
-        return self.f.array_contains(self.f.map_keys(arg), key)
+        key_type = op.arg.dtype.key_type
+        return self.f.array_contains(
+            self.cast(self.f.map_keys(arg), dt.Array(value_type=key_type)),
+            self.cast(key, key_type),
+        )
 
     def visit_Map(self, op: ops.Map, *, keys, values):
         return self.cast(self.f.map_from_arrays(keys, values), op.dtype)

diff --git a/ibis/backends/tests/test_map.py b/ibis/backends/tests/test_map.py
@@ -196,7 +196,6 @@ def test_literal_map_getitem_broadcast(backend, alltypes, df):
                 pytest.mark.notyet(
                     "clickhouse", reason="only supports str,int,bool,timestamp keys"
                 ),
-                pytest.mark.notyet("flink", reason="doesn't support float keys"),
                 mark_notyet_postgres,
             ],
             id="float",
@@ -227,7 +226,6 @@ def test_literal_map_getitem_broadcast(backend, alltypes, df):
                 ),
                 pytest.mark.notyet(["pandas", "dask"]),
                 mark_notyet_postgres,
-                pytest.mark.notimpl("flink"),
             ],
             id="array",
         ),