diff --git a/docs/source/release.rst b/docs/source/release.rst index a52a938e1263..4bbf28894f69 100644 --- a/docs/source/release.rst +++ b/docs/source/release.rst @@ -34,6 +34,7 @@ New Features * ``ibis.pandas.from_dataframe`` convenience function (:issue:`1155`) * Remove the restriction on ``ROW_NUMBER()`` requiring it to have an ``ORDER BY`` clause (:issue:`1371`) +* Add ``.get()`` operation on a Map type (:issue:`1376`) Bug Fixes ~~~~~~~~~ diff --git a/ibis/expr/api.py b/ibis/expr/api.py index 4fee58656e3d..c13306f22197 100644 --- a/ibis/expr/api.py +++ b/ibis/expr/api.py @@ -1992,13 +1992,27 @@ def _array_slice(array, index): _add_methods(ArrayValue, _array_column_methods) + # --------------------------------------------------------------------- # Map API +def get(expr, key, default): + """ + Return the mapped value for this key, or the default + if the key does not exist + + Parameters + ---------- + key : any + default : any + """ + return _ops.MapValueOrDefaultForKey(expr, key, default).to_expr() + _map_column_methods = dict( length=_unary_op('length', _ops.MapLength), __getitem__=_binop_expr('__getitem__', _ops.MapValueForKey), + get=get, keys=_unary_op('keys', _ops.MapKeys), values=_unary_op('values', _ops.MapValues), __add__=_binop_expr('__add__', _ops.MapConcat), diff --git a/ibis/expr/operations.py b/ibis/expr/operations.py index d5c556bb6b6b..ffadcca0ca23 100644 --- a/ibis/expr/operations.py +++ b/ibis/expr/operations.py @@ -2746,6 +2746,26 @@ def output_type(self): return rules.shape_like(self.args[0], map_type.value_type) +class MapValueOrDefaultForKey(ValueOp): + + input_type = [ + rules.map(dt.any, dt.any), + rules.one_of((dt.string, dt.int_), name='key'), + rules.value(name='default') + ] + + def output_type(self): + map_type = self.args[0].type() + value_type = map_type.value_type + default_type = self.default.type() + + if default_type is not dt.null and value_type != default_type: + raise ValueError("default type: {} must be the same " + "as the map value_type {}".format( + default_type, value_type)) + return rules.shape_like(self.args[0], map_type.value_type) + + class MapKeys(ValueOp): input_type = [rules.map(dt.any, dt.any)] diff --git a/ibis/expr/tests/test_value_exprs.py b/ibis/expr/tests/test_value_exprs.py index fe6a4632860a..76226b4dc00c 100644 --- a/ibis/expr/tests/test_value_exprs.py +++ b/ibis/expr/tests/test_value_exprs.py @@ -150,13 +150,23 @@ def test_simple_map_operations(): expr = ibis.literal(value) expr2 = ibis.literal(value2) assert isinstance(expr, ir.MapValue) - assert isinstance(expr['b'].op(), ops.MapValueForKey) assert isinstance(expr.length().op(), ops.MapLength) - assert isinstance(expr.keys().op(), ops.MapKeys) - assert isinstance(expr.values().op(), ops.MapValues) assert isinstance((expr + expr2).op(), ops.MapConcat) assert isinstance((expr2 + expr).op(), ops.MapConcat) + default = ibis.literal([0.0]) + assert isinstance(expr.get('d', default).op(), ops.MapValueOrDefaultForKey) + + # test for an invalid default type, nulls are ok + with pytest.raises(ValueError): + expr.get('d', ibis.literal('foo')) + assert isinstance(expr.get('d', ibis.literal(None)).op(), + ops.MapValueOrDefaultForKey) + + assert isinstance(expr['b'].op(), ops.MapValueForKey) + assert isinstance(expr.keys().op(), ops.MapKeys) + assert isinstance(expr.values().op(), ops.MapValues) + @pytest.mark.parametrize( ['value', 'expected_type'],