diff --git a/docs/source/api.rst b/docs/source/api.rst index eb80edfc40cb..977e15f12098 100644 --- a/docs/source/api.rst +++ b/docs/source/api.rst @@ -17,6 +17,9 @@ These methods are available directly in the ``ibis`` module namespace. case literal + array + struct + map schema table timestamp diff --git a/ibis/expr/api.py b/ibis/expr/api.py index ffd5dc92fcca..6e46c8b4de6b 100644 --- a/ibis/expr/api.py +++ b/ibis/expr/api.py @@ -101,7 +101,9 @@ ValueExpr, array, literal, + map, null, + struct, ) from ibis.expr.window import ( cumulative_window, @@ -183,6 +185,7 @@ 'join', 'least', 'literal', + 'map', 'NA', 'negate', 'now', @@ -197,6 +200,7 @@ 'schema', 'Schema', 'sequence', + 'struct', 'table', 'time', 'timestamp', diff --git a/ibis/expr/datatypes.py b/ibis/expr/datatypes.py index 98f48efac2c2..fbcfd69fed83 100644 --- a/ibis/expr/datatypes.py +++ b/ibis/expr/datatypes.py @@ -535,7 +535,7 @@ def from_dict( @property def pairs(self) -> Mapping: - return collections.OrderedDict(zip(self.names, self.types)) + return dict(zip(self.names, self.types)) def __getitem__(self, key: str) -> DataType: return self.pairs[key] @@ -1414,6 +1414,24 @@ def can_cast_string_to_temporal( Collection = TypeVar('Collection', Array, Set) +@castable.register(Map, Map) +def can_cast_map(source, target, **kwargs): + return castable(source.key_type, target.key_type) and castable( + source.value_type, target.value_type + ) + + +@castable.register(Struct, Struct) +def can_cast_struct(source, target, **kwargs): + source_pairs = source.pairs + target_pairs = target.pairs + for name in {*source.names, *target.names}: + if name in target_pairs: + if not castable(source_pairs[name], target_pairs[name]): + return False + return True + + @castable.register(Array, Array) @castable.register(Set, Set) def can_cast_variadic( @@ -1460,14 +1478,6 @@ def can_cast_special_string(source, target, **kwargs): return True -# @castable.register(Map, Map) -# def can_cast_maps(source, target): -# return (source.equals(target) or -# source.equals(Map(null, null)) or -# source.equals(Map(any, any))) -# TODO cast category - - def cast(source: DataType | str, target: DataType | str, **kwargs) -> DataType: """Attempts to implicitly cast from source dtype to target dtype""" source, result_target = dtype(source), dtype(target) diff --git a/ibis/expr/types.py b/ibis/expr/types.py index 3a309a6c224f..2c1d5a1db56f 100644 --- a/ibis/expr/types.py +++ b/ibis/expr/types.py @@ -1,3 +1,4 @@ +import collections import itertools import os import webbrowser @@ -518,7 +519,7 @@ def __dir__(self): def _resolve(self, exprs): exprs = util.promote_list(exprs) - return list(map(self._ensure_expr, exprs)) + return list(self._ensure_expr(x) for x in exprs) def _ensure_expr(self, expr): if isinstance(expr, str): @@ -1247,7 +1248,7 @@ def literal(value, type=None): return ops.Literal(value, dtype=dtype).to_expr() -def array(values): +def array(values, type=None): """Create an array expression. If the input expressions are all column expressions, then the output will @@ -1270,6 +1271,9 @@ def array(values): array_value : ArrayValue An array column (if the inputs are column expressions), or an array scalar (if the inputs are Python literals) + type : ibis type or string, optional + An instance of :class:`ibis.expr.datatypes.DataType` or a string + indicating the ibis type of `value`. Examples -------- @@ -1293,7 +1297,7 @@ def array(values): ) else: try: - return literal(list(values)) + return literal(list(values), type=type) except com.IbisTypeError as e: raise com.IbisTypeError( 'Could not create an array scalar from the values provided ' @@ -1302,6 +1306,65 @@ def array(values): ) from e +def struct(value, type=None): + '''Create a struct literal from a dict or other mapping. + + Parameters + ---------- + value + the literal struct value + type + An instance of :class:`ibis.expr.datatypes.DataType` or a string + indicating the ibis type of `value`. + + Returns + ------- + StructValue + An expression representing a literal struct (compound type with fields + of fixed types) + + Examples + -------- + Create struct literal from dict with inferred type: + >>> import ibis + >>> t = ibis.struct(dict(a=1, b='foo')) + + Create struct literal from dict with specified type: + >>> t = ibis.struct(dict(a=1, b='foo'), type='struct') + ''' + return literal(collections.OrderedDict(value), type=type) + + +def map(value, type=None): + '''Create a map literal from a dict or other mapping. + + Parameters + ---------- + value + the literal map value + type + An instance of :class:`ibis.expr.datatypes.DataType` or a string + indicating the ibis type of `value`. + + Returns + ------- + MapValue + An expression representing a literal map (associative array with + key/value pairs of fixed types) + + Examples + -------- + Create map literal from dict with inferred type: + >>> import ibis + >>> t = ibis.map(dict(a=1, b=2)) + + Create map literal from dict with specified type: + >>> import ibis + >>> t = ibis.map(dict(a=1, b=2), type='map') + ''' + return literal(dict(value), type=type) + + class UnnamedMarker: pass diff --git a/ibis/tests/expr/test_array.py b/ibis/tests/expr/test_array.py index 3626dcdb1240..ee976671f29e 100644 --- a/ibis/tests/expr/test_array.py +++ b/ibis/tests/expr/test_array.py @@ -6,15 +6,16 @@ @pytest.mark.parametrize( - ['arg', 'type'], + ['arg', 'typestr', 'type'], [ - ([1, 2, 3], dt.Array(dt.int8)), - ([1, 2, 3.0], dt.Array(dt.double)), - (['a', 'b', 'c'], dt.Array(dt.string)), + ([1, 2, 3], None, dt.Array(dt.int8)), + ([1, 2, 3], 'array', dt.Array(dt.int16)), + ([1, 2, 3.0], None, dt.Array(dt.double)), + (['a', 'b', 'c'], None, dt.Array(dt.string)), ], ) -def test_array_literal(arg, type): - x = ibis.literal(arg) +def test_array_literal(arg, typestr, type): + x = ibis.literal(arg, type=typestr) assert x._arg.value == arg assert x.type() == type diff --git a/ibis/tests/expr/test_literal.py b/ibis/tests/expr/test_literal.py index a72d5a071743..551c33b2278d 100644 --- a/ibis/tests/expr/test_literal.py +++ b/ibis/tests/expr/test_literal.py @@ -92,3 +92,60 @@ def test_normalized_underlying_value(userinput, literal_type, expected_type): a = ibis.literal(userinput, type=literal_type) assert isinstance(a.op().value, expected_type) + + +@pytest.mark.parametrize( + 'value', + [ + dict(field1='value1', field2=3.14), + dict(field1='value1', field2=1), # coerceable type + dict(field2=2.72, field1='value1'), # wrong field order + dict(field1='value1', field2=3.14, field3='extra'), # extra field + ], +) +def test_struct_literal(value): + typestr = "struct" + a = ibis.struct(value, type=typestr) + assert a.op().value == value + assert a.type() == datatypes.dtype(typestr) + + +@pytest.mark.parametrize( + 'value', + [ + dict(field1='value1', field2='3.14'), # non-coerceable type + dict(field1='value1', field3=3.14), # wrong field name + dict(field1='value1'), # missing field + ], +) +def test_struct_literal_non_castable(value): + typestr = "struct" + with pytest.raises( + (KeyError, TypeError, ibis.common.exceptions.IbisTypeError) + ): + ibis.struct(value, type=typestr) + + +@pytest.mark.parametrize( + 'value', + [ + dict(key1='value1', key2='value2'), + ], +) +def test_map_literal(value): + typestr = "map" + a = ibis.map(value, type=typestr) + assert a.op().value == value + assert a.type() == datatypes.dtype(typestr) + + +@pytest.mark.parametrize( + 'value', + [ + dict(key1='value1', key2=6.25), # heterogenous map values + ], +) +def test_map_literal_non_castable(value): + typestr = "map" + with pytest.raises(ibis.common.exceptions.IbisTypeError): + ibis.map(value, type=typestr)