/* * 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; import com.facebook.presto.type.ArrayType; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import org.testng.annotations.Test; import java.util.HashMap; import java.util.Map; import static com.facebook.presto.spi.type.BigintType.BIGINT; import static com.facebook.presto.spi.type.BooleanType.BOOLEAN; import static com.facebook.presto.spi.type.DoubleType.DOUBLE; import static com.facebook.presto.spi.type.IntegerType.INTEGER; import static com.facebook.presto.spi.type.VarcharType.VARCHAR; import static com.facebook.presto.spi.type.VarcharType.createVarcharType; import static com.facebook.presto.type.UnknownType.UNKNOWN; import static com.facebook.presto.util.StructuralTestUtil.mapType; public class TestMapTransformValueFunction extends AbstractTestFunctions { @Test public void testEmpty() throws Exception { assertFunction("transform_values(map(ARRAY[], ARRAY[]), (k, v) -> NULL)", mapType(UNKNOWN, UNKNOWN), ImmutableMap.of()); assertFunction("transform_values(map(ARRAY[], ARRAY[]), (k, v) -> k)", mapType(UNKNOWN, UNKNOWN), ImmutableMap.of()); assertFunction("transform_values(map(ARRAY[], ARRAY[]), (k, v) -> v)", mapType(UNKNOWN, UNKNOWN), ImmutableMap.of()); assertFunction("transform_values(map(ARRAY[], ARRAY[]), (k, v) -> 0)", mapType(UNKNOWN, INTEGER), ImmutableMap.of()); assertFunction("transform_values(map(ARRAY[], ARRAY[]), (k, v) -> true)", mapType(UNKNOWN, BOOLEAN), ImmutableMap.of()); assertFunction("transform_values(map(ARRAY[], ARRAY[]), (k, v) -> 'value')", mapType(UNKNOWN, createVarcharType(5)), ImmutableMap.of()); assertFunction("transform_values(CAST (map(ARRAY[], ARRAY[]) AS MAP(BIGINT,VARCHAR)), (k, v) -> k + CAST(v as BIGINT))", mapType(BIGINT, BIGINT), ImmutableMap.of()); assertFunction("transform_values(CAST (map(ARRAY[], ARRAY[]) AS MAP(BIGINT,VARCHAR)), (k, v) -> CAST(k AS VARCHAR) || v)", mapType(BIGINT, VARCHAR), ImmutableMap.of()); } @Test public void testNullValue() throws Exception { Map<Integer, Void> sequenceToNullMap = new HashMap<>(); sequenceToNullMap.put(1, null); sequenceToNullMap.put(2, null); sequenceToNullMap.put(3, null); assertFunction("transform_values(map(ARRAY[1, 2, 3], ARRAY ['a', 'b', 'c']), (k, v) -> NULL)", mapType(INTEGER, UNKNOWN), sequenceToNullMap); Map<Integer, String> mapWithNullValue = new HashMap<>(); mapWithNullValue.put(1, "a"); mapWithNullValue.put(2, "b"); mapWithNullValue.put(3, null); assertFunction("transform_values(map(ARRAY[1, 2, 3], ARRAY ['a', 'b', NULL]), (k, v) -> v)", mapType(INTEGER, createVarcharType(1)), mapWithNullValue); assertFunction("transform_values(map(ARRAY[1, 2, 3], ARRAY [10, 11, NULL]), (k, v) -> to_base(v, 16))", mapType(INTEGER, createVarcharType(64)), mapWithNullValue); assertFunction("transform_values(map(ARRAY[1, 2, 3], ARRAY ['10', '11', 'Invalid']), (k, v) -> to_base(TRY_CAST(v as BIGINT), 16))", mapType(INTEGER, createVarcharType(64)), mapWithNullValue); assertFunction( "transform_values(map(ARRAY[1, 2, 3], ARRAY [0, 0, 0]), (k, v) -> element_at(map(ARRAY[1, 2], ARRAY['a', 'b']), k + v))", mapType(INTEGER, createVarcharType(1)), mapWithNullValue); assertFunction( "transform_values(map(ARRAY[1, 2, 3], ARRAY ['a', 'b', NULL]), (k, v) -> IF(v IS NULL, k + 1.0, k + 0.5))", mapType(INTEGER, DOUBLE), ImmutableMap.of(1, 1.5, 2, 2.5, 3, 4.0)); } @Test public void testBasic() throws Exception { assertFunction( "transform_values(map(ARRAY [1, 2, 3, 4], ARRAY [10, 20, 30, 40]), (k, v) -> k + v)", mapType(INTEGER, INTEGER), ImmutableMap.of(1, 11, 2, 22, 3, 33, 4, 44)); assertFunction( "transform_values(map(ARRAY ['a', 'b', 'c', 'd'], ARRAY [1, 2, 3, 4]), (k, v) -> v * v)", mapType(createVarcharType(1), INTEGER), ImmutableMap.of("a", 1, "b", 4, "c", 9, "d", 16)); assertFunction( "transform_values(map(ARRAY ['a', 'b', 'c', 'd'], ARRAY [1, 2, 3, 4]), (k, v) -> k || CAST(v as VARCHAR))", mapType(createVarcharType(1), VARCHAR), ImmutableMap.of("a", "a1", "b", "b2", "c", "c3", "d", "d4")); assertFunction( "transform_values(map(ARRAY[1, 2, 3], ARRAY [1.0, 1.4, 1.7]), (k, v) -> map(ARRAY[1, 2, 3], ARRAY['one', 'two', 'three'])[k] || '_' || CAST(v AS VARCHAR))", mapType(INTEGER, VARCHAR), ImmutableMap.of(1, "one_1.0", 2, "two_1.4", 3, "three_1.7")); } @Test public void testTypeCombinations() throws Exception { assertFunction( "transform_values(map(ARRAY [25, 26, 27], ARRAY [25, 26, 27]), (k, v) -> k + v)", mapType(INTEGER, INTEGER), ImmutableMap.of(25, 50, 26, 52, 27, 54)); assertFunction( "transform_values(map(ARRAY [25, 26, 27], ARRAY [26.1, 31.2, 37.1]), (k, v) -> CAST(v - k AS BIGINT))", mapType(INTEGER, BIGINT), ImmutableMap.of(25, 1L, 26, 5L, 27, 10L)); assertFunction( "transform_values(map(ARRAY [25, 27], ARRAY [false, true]), (k, v) -> if(v, k + 1, k + 2))", mapType(INTEGER, INTEGER), ImmutableMap.of(25, 27, 27, 28)); assertFunction( "transform_values(map(ARRAY [25, 26, 27], ARRAY ['abc', 'd', 'xy']), (k, v) -> k + length(v))", mapType(INTEGER, BIGINT), ImmutableMap.of(25, 28L, 26, 27L, 27, 29L)); assertFunction( "transform_values(map(ARRAY [25, 26, 27], ARRAY [ARRAY ['a'], ARRAY ['a', 'c'], ARRAY ['a', 'b', 'c']]), (k, v) -> k + cardinality(v))", mapType(INTEGER, BIGINT), ImmutableMap.of(25, 26L, 26, 28L, 27, 30L)); assertFunction( "transform_values(map(ARRAY [25.5, 26.75, 27.875], ARRAY [25, 26, 27]), (k, v) -> k - v)", mapType(DOUBLE, DOUBLE), ImmutableMap.of(25.5, 0.5, 26.75, 0.75, 27.875, 0.875)); assertFunction( "transform_values(map(ARRAY [25.5, 26.75, 27.875], ARRAY [25.0, 26.0, 27.0]), (k, v) -> k - v)", mapType(DOUBLE, DOUBLE), ImmutableMap.of(25.5, 0.5, 26.75, 0.75, 27.875, 0.875)); assertFunction( "transform_values(map(ARRAY [25.5, 27.5], ARRAY [false, true]), (k, v) -> if(v, k + 0.1, k + 0.2))", mapType(DOUBLE, DOUBLE), ImmutableMap.of(25.5, 25.7, 27.5, 27.6)); assertFunction( "transform_values(map(ARRAY [25.5, 26.5, 27.5], ARRAY ['a', 'def', 'xy']), (k, v) -> k + length(v))", mapType(DOUBLE, DOUBLE), ImmutableMap.of(25.5, 26.5, 26.5, 29.5, 27.5, 29.5)); assertFunction( "transform_values(map(ARRAY [25.5, 26.5, 27.5], ARRAY [ARRAY ['a'], ARRAY ['a', 'c'], ARRAY ['a', 'b', 'c']]), (k, v) -> k + cardinality(v))", mapType(DOUBLE, DOUBLE), ImmutableMap.of(25.5, 26.5, 26.5, 28.5, 27.5, 30.5)); assertFunction( "transform_values(map(ARRAY [true, false], ARRAY [25, 26]), (k, v) -> k AND v = 25)", mapType(BOOLEAN, BOOLEAN), ImmutableMap.of(true, true, false, false)); assertFunction( "transform_values(map(ARRAY [false, true], ARRAY [25.5, 26.5]), (k, v) -> k OR v > 100)", mapType(BOOLEAN, BOOLEAN), ImmutableMap.of(false, false, true, true)); assertFunction( "transform_values(map(ARRAY [true, false], ARRAY [false, null]), (k, v) -> NOT k OR v)", mapType(BOOLEAN, BOOLEAN), ImmutableMap.of(false, true, true, false)); assertFunction( "transform_values(map(ARRAY [false, true], ARRAY ['abc', 'def']), (k, v) -> NOT k AND v = 'abc')", mapType(BOOLEAN, BOOLEAN), ImmutableMap.of(false, true, true, false)); assertFunction( "transform_values(map(ARRAY [true, false], ARRAY [ARRAY ['a', 'b'], ARRAY ['a', 'b', 'c']]), (k, v) -> k OR cardinality(v) = 3)", mapType(BOOLEAN, BOOLEAN), ImmutableMap.of(false, true, true, true)); assertFunction( "transform_values(map(ARRAY ['s0', 's1', 's2'], ARRAY [25, 26, 27]), (k, v) -> k || ':' || CAST(v as VARCHAR))", mapType(createVarcharType(2), VARCHAR), ImmutableMap.of("s0", "s0:25", "s1", "s1:26", "s2", "s2:27")); assertFunction( "transform_values(map(ARRAY ['s0', 's1', 's2'], ARRAY [25.5, 26.5, 27.5]), (k, v) -> k || ':' || CAST(v as VARCHAR))", mapType(createVarcharType(2), VARCHAR), ImmutableMap.of("s0", "s0:25.5", "s1", "s1:26.5", "s2", "s2:27.5")); assertFunction( "transform_values(map(ARRAY ['s0', 's2'], ARRAY [false, true]), (k, v) -> if(v, k, CAST(v AS VARCHAR)))", mapType(createVarcharType(2), VARCHAR), ImmutableMap.of("s0", "false", "s2", "s2")); assertFunction( "transform_values(map(ARRAY ['s0', 's1', 's2'], ARRAY ['abc', 'def', 'xyz']), (k, v) -> k || ':' || v)", mapType(createVarcharType(2), VARCHAR), ImmutableMap.of("s0", "s0:abc", "s1", "s1:def", "s2", "s2:xyz")); assertFunction( "transform_values(map(ARRAY ['s0', 's1', 's2'], ARRAY [ARRAY ['a', 'b'], ARRAY ['a', 'c'], ARRAY ['a', 'b', 'c']]), (k, v) -> k || ':' || array_max(v))", mapType(createVarcharType(2), VARCHAR), ImmutableMap.of("s0", "s0:b", "s1", "s1:c", "s2", "s2:c")); assertFunction( "transform_values(map(ARRAY [ARRAY [1, 2], ARRAY [3, 4]], ARRAY [25, 26]), (k, v) -> if(v % 2 = 0, reverse(k), k))", mapType(new ArrayType(INTEGER), new ArrayType(INTEGER)), ImmutableMap.of(ImmutableList.of(1, 2), ImmutableList.of(1, 2), ImmutableList.of(3, 4), ImmutableList.of(4, 3))); assertFunction( "transform_values(map(ARRAY [ARRAY [1, 2], ARRAY [3, 4]], ARRAY [25.5, 26.5]), (k, v) -> CAST(k AS ARRAY(DOUBLE)) || v)", mapType(new ArrayType(INTEGER), new ArrayType(DOUBLE)), ImmutableMap.of(ImmutableList.of(1, 2), ImmutableList.of(1., 2., 25.5), ImmutableList.of(3, 4), ImmutableList.of(3., 4., 26.5))); assertFunction( "transform_values(map(ARRAY [ARRAY [1, 2], ARRAY [3, 4]], ARRAY [false, true]), (k, v) -> if(v, reverse(k), k))", mapType(new ArrayType(INTEGER), new ArrayType(INTEGER)), ImmutableMap.of(ImmutableList.of(1, 2), ImmutableList.of(1, 2), ImmutableList.of(3, 4), ImmutableList.of(4, 3))); assertFunction( "transform_values(map(ARRAY [ARRAY [1, 2], ARRAY []], ARRAY ['a', 'ff']), (k, v) -> k || from_base(v, 16))", mapType(new ArrayType(INTEGER), new ArrayType(BIGINT)), ImmutableMap.of(ImmutableList.of(1, 2), ImmutableList.of(1L, 2L, 10L), ImmutableList.of(), ImmutableList.of(255L))); assertFunction( "transform_values(map(ARRAY [ARRAY [3, 4], ARRAY []], ARRAY [ARRAY ['a', 'b', 'c'], ARRAY ['a', 'c']]), (k, v) -> transform(k, x -> CAST(x AS VARCHAR)) || v)", mapType(new ArrayType(INTEGER), new ArrayType(VARCHAR)), ImmutableMap.of(ImmutableList.of(3, 4), ImmutableList.of("3", "4", "a", "b", "c"), ImmutableList.of(), ImmutableList.of("a", "c"))); } }