diff --git a/presto-docs/src/main/sphinx/functions/map.rst b/presto-docs/src/main/sphinx/functions/map.rst index e2cd0f5da5a55..e055cef83a34a 100644 --- a/presto-docs/src/main/sphinx/functions/map.rst +++ b/presto-docs/src/main/sphinx/functions/map.rst @@ -99,6 +99,12 @@ Map Functions SELECT map_subset(MAP(ARRAY[1,2], ARRAY['a','b']), ARRAY[]); -- {} SELECT map_subset(MAP(ARRAY[], ARRAY[]), ARRAY[1,2]); -- {} +.. function:: map_key_exists(x(K, V), k) -> boolean + + Returns whether the given key exists in the map. Returns ``true`` if key is present in the input map, returns ``false`` if not present. + + SELECT map_key_exists(MAP(ARRAY['x','y'], ARRAY[100,200]), 'x'); -- TRUE + .. function:: map_keys(x(K,V)) -> array(K) Returns all the keys in the map ``x``. diff --git a/presto-main/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java b/presto-main/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java index 929436e7dd626..dba6659963d92 100644 --- a/presto-main/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java +++ b/presto-main/src/main/java/com/facebook/presto/operator/scalar/sql/MapSqlFunctions.java @@ -35,6 +35,17 @@ public static String mapKeysByTopNValues() return "RETURN IF(n < 0, fail('n must be greater than or equal to 0'), map_keys(map_top_n(input, n)))"; } + @SqlInvokedScalarFunction(value = "map_key_exists", deterministic = true, calledOnNullInput = false) + @Description("Returns whether a given key exists in a map.") + @TypeParameter("K") + @TypeParameter("V") + @SqlParameters({@SqlParameter(name = "input", type = "map(K, V)"), @SqlParameter(name = "k", type = "K")}) + @SqlType("boolean") + public static String mapKeysExists() + { + return "RETURN CONTAINS(MAP_KEYS(input), k)"; + } + @SqlInvokedScalarFunction(value = "map_top_n", deterministic = true, calledOnNullInput = true) @Description("Truncates map items. Keeps only the top N elements by value.") @TypeParameter("K") diff --git a/presto-main/src/test/java/com/facebook/presto/operator/scalar/sql/TestMapKeyExists.java b/presto-main/src/test/java/com/facebook/presto/operator/scalar/sql/TestMapKeyExists.java new file mode 100644 index 0000000000000..d0895c7e7f007 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/operator/scalar/sql/TestMapKeyExists.java @@ -0,0 +1,85 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.operator.scalar.sql; + +import com.facebook.presto.operator.scalar.AbstractTestFunctions; +import org.testng.annotations.Test; + +import static com.facebook.presto.common.type.BooleanType.BOOLEAN; + +public class TestMapKeyExists + extends AbstractTestFunctions +{ + @Test + public void testBasic() + { + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY[1, 2, 3], ARRAY[4, 5, 6]), 2)", + BOOLEAN, + true); + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY[-1, -2, -3], ARRAY[4, 5, 6]), 2)", + BOOLEAN, + false); + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY['ab', 'bc', 'cd'], ARRAY['x', 'y', 'z']), 'abc')", + BOOLEAN, + false); + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY[123.0, 99.5, 1000.99], ARRAY['x', 'y', 'z']), 123.0)", + BOOLEAN, + true); + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY['x', 'y'], ARRAY[NULL, 1]), 'x')", + BOOLEAN, + true); + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY['x', 'y'], ARRAY[NULL, NULL]), 'z')", + BOOLEAN, + false); + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY[NAN(), 123.21], ARRAY['val1', 'val2']), 123.21)", + BOOLEAN, + true); + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY[NAN(), 123.21], ARRAY['val1', 'val2']), NAN())", + BOOLEAN, + false); + } + + @Test + public void testEmpty() + { + assertFunction("MAP_KEY_EXISTS(MAP(ARRAY[], ARRAY[]), 5)", BOOLEAN, false); + } + + @Test + public void testNull() + { + assertFunction("MAP_KEY_EXISTS(NULL, 1)", BOOLEAN, null); + } + + @Test + public void testComplexKeys() + { + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY[ROW('x', 1), ROW('y', 2)], ARRAY[1, 2]), ROW('x', 1))", + BOOLEAN, + true); + assertFunction( + "MAP_KEY_EXISTS(MAP(ARRAY[ROW('x', 1), ROW('x', -2)], ARRAY[2, 1]), ROW('y', 1))", + BOOLEAN, + false); + } +}