diff --git a/presto-docs/src/main/sphinx/functions/map.rst b/presto-docs/src/main/sphinx/functions/map.rst index bc9a62bc5cd2f..bc7412fe0fcc5 100644 --- a/presto-docs/src/main/sphinx/functions/map.rst +++ b/presto-docs/src/main/sphinx/functions/map.rst @@ -214,3 +214,14 @@ Map Functions SELECT transform_values(MAP(ARRAY ['a', 'b'], ARRAY [1, 2]), (k, v) -> k || CAST(v as VARCHAR)); -- {a -> a1, b -> b2} SELECT transform_values(MAP(ARRAY [1, 2], ARRAY [1.0, 1.4]), -- {1 -> one_1.0, 2 -> two_1.4} (k, v) -> MAP(ARRAY[1, 2], ARRAY['one', 'two'])[k] || '_' || CAST(v AS VARCHAR)); + +.. function:: map_int_keys_to_array(map(int,V)) -> array(V) + Returns an ``array`` of values from the ``map`` with value at indexed by the original keys from ``map``:: + SELECT MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[3, 5, 6, 9], ARRAY['a', 'b', 'c', 'd'])) -> ARRAY[null, null, 'a', null, 'b', 'c', null, null, 'd'] + SELECT MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[3, 5, 6, 9], ARRAY['a', null, 'c', 'd'])) -> ARRAY[null, null, 'a', null, null, 'c', 'd'] + +.. function:: array_to_map_int_keys(array(v)) -> map(int, v) + Returns an ``map`` with indices of all non-null values from the ``array`` as keys and element at the specified index as the value:: + SELECT ARRAY_TO_MAP_INT_KEYS(CAST(ARRAY[3, 5, 6, 9] AS ARRAY)) -> MAP(ARRAY[1, 2, 3,4], ARRAY[3, 5, 6, 9]) + SELECT ARRAY_TO_MAP_INT_KEYS(CAST(ARRAY[3, 5, null, 6, 9] AS ARRAY)) -> MAP(ARRAY[1, 2, 4, 5], ARRAY[3, 5, 6, 9]) + SELECT ARRAY_TO_MAP_INT_KEYS(CAST(ARRAY[3, 5, null, 6, 9, null, null, 1] AS ARRAY)) -> MAP(ARRAY[1, 2, 4, 5, 8], ARRAY[3, 5, 6, 9, 1]) \ No newline at end of file diff --git a/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/sql/TestMapSqlFunctions.java b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/sql/TestMapSqlFunctions.java new file mode 100644 index 0000000000000..73996c37ed690 --- /dev/null +++ b/presto-main-base/src/test/java/com/facebook/presto/operator/scalar/sql/TestMapSqlFunctions.java @@ -0,0 +1,105 @@ +/* + * 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.common.type.ArrayType; +import com.facebook.presto.operator.scalar.AbstractTestFunctions; +import com.facebook.presto.spi.StandardErrorCode; +import com.google.common.collect.ImmutableMap; +import org.testng.annotations.Test; + +import static com.facebook.presto.common.type.IntegerType.INTEGER; +import static com.facebook.presto.common.type.VarcharType.createVarcharType; +import static com.facebook.presto.util.StructuralTestUtil.mapType; +import static java.util.Arrays.asList; + +public class TestMapSqlFunctions + extends AbstractTestFunctions +{ + @Test + public void testMapIntKeysToArray() + { + assertFunction("MAP_INT_KEYS_TO_ARRAY(CAST(MAP() AS MAP))", + new ArrayType(INTEGER), + null); + + assertFunction("MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[1, 3], ARRAY['a', 'b']))", + new ArrayType(createVarcharType(1)), + asList("a", null, "b")); + + assertFunction("MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[3, 5, 6, 9], ARRAY['a', 'b', 'c', 'd']))", + new ArrayType(createVarcharType(1)), + asList(null, null, "a", null, "b", "c", null, null, "d")); + + assertFunction("MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[3, 5, 6, 9], ARRAY['a', null, 'c', 'd']))", + new ArrayType(createVarcharType(1)), + asList(null, null, "a", null, null, "c", null, null, "d")); + + assertFunction("MAP_INT_KEYS_TO_ARRAY(CAST(MAP() AS MAP))", + new ArrayType(INTEGER), + null); + + assertFunction("MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[2], ARRAY[MAP(ARRAY[3, 5, 6, 9], ARRAY['a', null, 'c', 'd'])]))", + new ArrayType(mapType(INTEGER, createVarcharType(1))), + asList(null, asMap(asList(3, 5, 6, 9), asList("a", null, "c", "d")))); + + assertFunction("MAP_INT_KEYS_TO_ARRAY(CAST(MAP() AS MAP))", + new ArrayType(INTEGER), + null); + + assertInvalidFunction( + "MAP_INT_KEYS_TO_ARRAY(MAP(CAST(SEQUENCE(1,10000)||ARRAY[10001] AS ARRAY),SEQUENCE(1,10000)||ARRAY[10001]))", + StandardErrorCode.GENERIC_USER_ERROR, + "Max key value must be <= 10k for map_int_keys_to_array function"); + + assertInvalidFunction( + "MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[0],ARRAY[1]))", + StandardErrorCode.GENERIC_USER_ERROR, + "Only positive keys allowed in map_int_keys_to_array function, but got: 0"); + + assertInvalidFunction( + "MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[-1, 2], ARRAY['a', 'b']))", + StandardErrorCode.GENERIC_USER_ERROR, + "Only positive keys allowed in map_int_keys_to_array function, but got: -1"); + + assertInvalidFunction("MAP_INT_KEYS_TO_ARRAY(MAP(ARRAY[0, 2], ARRAY['a', 'b']))", + StandardErrorCode.GENERIC_USER_ERROR, + "Only positive keys allowed in map_int_keys_to_array function, but got: 0"); + } + + @Test + public void testArrayToMapIntKeys() + { + assertFunction("ARRAY_TO_MAP_INT_KEYS(CAST(ARRAY[3, 5, 6, 9] AS ARRAY))", + mapType(INTEGER, INTEGER), + ImmutableMap.of(1, 3, 2, 5, 3, 6, 4, 9)); + + assertFunction("ARRAY_TO_MAP_INT_KEYS(CAST(ARRAY[3, 5, null, 6, 9] AS ARRAY))", + mapType(INTEGER, INTEGER), + asMap(asList(1, 2, 4, 5), asList(3, 5, 6, 9))); + + assertFunction("ARRAY_TO_MAP_INT_KEYS(CAST(ARRAY[3, 5, null, 6, 9, null, null, 1] AS ARRAY))", + mapType(INTEGER, INTEGER), + asMap(asList(1, 2, 4, 5, 8), asList(3, 5, 6, 9, 1))); + + assertFunction("ARRAY_TO_MAP_INT_KEYS(CAST(NULL AS ARRAY))", + mapType(INTEGER, INTEGER), + null); + + assertInvalidFunction( + "ARRAY_TO_MAP_INT_KEYS(SEQUENCE(1,10000)||ARRAY[10001])", + StandardErrorCode.GENERIC_USER_ERROR, + "Max number of elements must be <= 10k for array_to_map_int_keys function"); + } +} diff --git a/presto-sql-helpers/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java b/presto-sql-helpers/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java index 498c068dedf39..8063b5975bca7 100644 --- a/presto-sql-helpers/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java +++ b/presto-sql-helpers/presto-sql-invoked-functions-plugin/src/main/java/com/facebook/presto/scalar/sql/MapSqlFunctions.java @@ -167,4 +167,31 @@ public static String noValuesMatch() { return "RETURN NONE_MATCH(MAP_VALUES(input), f)"; } + + @SqlInvokedScalarFunction(value = "map_int_keys_to_array", deterministic = true, calledOnNullInput = true) + @Description("Convert a map with (a small number of) int keys to an array 1..max_key (upto 10K) with nulls for missing keys") + @TypeParameter("V") + @SqlParameters({@SqlParameter(name = "input", type = "map(INTEGER, V)")}) + @SqlType("array(V)") + public static String mapIntKeysToArray() + { + return "RETURN IF(ARRAY_MAX(MAP_KEYS(input)) > 10000, " + + " FAIL('Max key value must be <= 10k for map_int_keys_to_array function'), " + + " IF(ARRAY_MIN(MAP_KEYS(input)) <= 0, " + + " FAIL('Only positive keys allowed in map_int_keys_to_array function, but got: ' || cast(ARRAY_MIN(MAP_KEYS(input)) as varchar)), " + + " TRANSFORM(SEQUENCE(1, ARRAY_MAX(MAP_KEYS(input))), " + + " k->element_at(input, k))))"; + } + + @SqlInvokedScalarFunction(value = "array_to_map_int_keys", deterministic = true, calledOnNullInput = true) + @Description("Convert an array to a map with array index->element value for all non-null element values") + @TypeParameter("V") + @SqlParameters({@SqlParameter(name = "input", type = "ARRAY(V)")}) + @SqlType("MAP(INTEGER, V)") + public static String arrayToMapIntKeys() + { + return "RETURN IF(CARDINALITY(input) > 10000," + + " FAIL('Max number of elements must be <= 10k for array_to_map_int_keys function'), " + + " MAP_FROM_ENTRIES(REMOVE_NULLS(TRANSFORM(CAST(SEQUENCE(1, CARDINALITY(input)) AS ARRAY), x->IF(input[x] IS NOT NULL, (x, input[x]))))))"; + } }