Skip to content

Commit

Permalink
chore: compose inferring map dtypes using ibis.array()
Browse files Browse the repository at this point in the history
  • Loading branch information
NickCrews committed Mar 15, 2024
1 parent 5ec7ee6 commit 9207925
Showing 1 changed file with 8 additions and 3 deletions.
11 changes: 8 additions & 3 deletions ibis/expr/types/maps.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
if TYPE_CHECKING:
from collections.abc import Iterable, Mapping

from ibis.expr.types import dt
from ibis.expr import datatypes as dt


@public
Expand Down Expand Up @@ -492,6 +492,8 @@ def map(
│ {'b': 3.0} │
└───────────────────────────────────────────────┘
"""
from ibis.expr import datatypes as dt

if keys is None:
if type is None:
raise TypeError("Must specify a type when keys is None")
Expand All @@ -502,8 +504,11 @@ def map(
else:
if values is None:
keys, values = tuple(keys.keys()), tuple(keys.values())
if not isinstance(keys, Value) and len(keys) == 0 and type is None:
raise TypeError("Must specify a type when keys is empty")
type = dt.dtype(type) if type is not None else None
key_type = dt.Array(value_type=type.key_type) if type is not None else None
value_type = dt.Array(value_type=type.value_type) if type is not None else None
keys = ibis.array(keys, type=key_type)
values = ibis.array(values, type=value_type)
result = ops.Map(keys, values).to_expr()

if type is not None:
Expand Down

0 comments on commit 9207925

Please sign in to comment.