Skip to content

Commit

Permalink
feat(api): support all dtypes in MapGet and MapContains (#8648)
Browse files Browse the repository at this point in the history
Co-authored-by: Phillip Cloud <417981+cpcloud@users.noreply.github.com>
  • Loading branch information
NickCrews and cpcloud authored Mar 21, 2024
1 parent ac86f91 commit 401e0a4
Show file tree
Hide file tree
Showing 4 changed files with 216 additions and 85 deletions.
6 changes: 5 additions & 1 deletion ibis/backends/flink/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion ibis/backends/snowflake/compiler.py
Original file line number Diff line number Diff line change
Expand Up @@ -198,7 +198,7 @@ def visit_RegexSplit(self, op, *, arg, pattern):
def visit_Map(self, op, *, keys, values):
return self.if_(
sg.and_(self.f.is_array(keys), self.f.is_array(values)),
self.f.udf.object_from_arrays(keys, values),
self.f.arrays_to_object(keys, values),
NULL,
)

Expand Down
Loading

0 comments on commit 401e0a4

Please sign in to comment.