diff --git a/docs/data/sql_functions.yml b/docs/data/sql_functions.yml index e6ce79e47f772..d3af2b96bae8d 100644 --- a/docs/data/sql_functions.yml +++ b/docs/data/sql_functions.yml @@ -925,6 +925,25 @@ collection: - sql: MAP_ENTRIES(map) table: MAP.mapEntries() description: Returns an array of all entries in the given map. No order guaranteed. + - sql: MAP_CONTAINS_KEY(map, key) + table: MAP.mapContainsKey(key) + description: | + Returns TRUE if the given key exists in the map, FALSE otherwise. Returns NULL if the map is + NULL. + + If the search key is NULL, the function returns TRUE when the map contains a NULL key. + The given key is cast implicitly to the map's key type where Flink's implicit casting rules + allow it; otherwise the call fails validation. + + Examples + -- TRUE + MAP_CONTAINS_KEY(MAP['a', 1, 'b', 2], 'a') + -- FALSE + MAP_CONTAINS_KEY(MAP['a', 1, 'b', 2], 'z') + -- TRUE + MAP_CONTAINS_KEY(MAP[CAST(NULL AS STRING), 1], CAST(NULL AS STRING)) + -- TRUE, the TINYINT key is cast to the map's INT key type + MAP_CONTAINS_KEY(MAP[1, 'a'], CAST(1 AS TINYINT)) - sql: MAP_FROM_ARRAYS(array_of_keys, array_of_values) table: mapFromArrays(array_of_keys, array_of_values) description: Returns a map created from an arrays of keys and values. Note that the lengths of two arrays should be the same. diff --git a/docs/data/sql_functions_zh.yml b/docs/data/sql_functions_zh.yml index b5819bbede25a..34912265281d1 100644 --- a/docs/data/sql_functions_zh.yml +++ b/docs/data/sql_functions_zh.yml @@ -1052,6 +1052,25 @@ collection: - sql: MAP_ENTRIES(map) table: MAP.mapEntries() description: 以数组形式返回 map 中的所有 entry，不保证顺序。 + - sql: MAP_CONTAINS_KEY(map, key) + table: MAP.mapContainsKey(key) + description: | + Returns TRUE if the given key exists in the map, FALSE otherwise. Returns NULL if the map is + NULL. + + If the search key is NULL, the function returns TRUE when the map contains a NULL key. + The given key is cast implicitly to the map's key type where Flink's implicit casting rules + allow it; otherwise the call fails validation. + + Examples + -- TRUE + MAP_CONTAINS_KEY(MAP['a', 1, 'b', 2], 'a') + -- FALSE + MAP_CONTAINS_KEY(MAP['a', 1, 'b', 2], 'z') + -- TRUE + MAP_CONTAINS_KEY(MAP[CAST(NULL AS STRING), 1], CAST(NULL AS STRING)) + -- TRUE, the TINYINT key is cast to the map's INT key type + MAP_CONTAINS_KEY(MAP[1, 'a'], CAST(1 AS TINYINT)) - sql: MAP_FROM_ARRAYS(array_of_keys, array_of_values) table: mapFromArrays(array_of_keys, array_of_values) description: 返回由 key 的数组 keys 和 value 的数组 values 创建的 map。请注意两个数组的长度应该相等。 diff --git a/flink-python/docs/reference/pyflink.table/expressions.rst b/flink-python/docs/reference/pyflink.table/expressions.rst index 2e2cee6cb93d4..a76320245a214 100644 --- a/flink-python/docs/reference/pyflink.table/expressions.rst +++ b/flink-python/docs/reference/pyflink.table/expressions.rst @@ -256,6 +256,7 @@ advanced type helper functions Expression.array_min Expression.array_sort Expression.array_union + Expression.map_contains_key Expression.map_entries Expression.map_from_entries Expression.map_keys diff --git a/flink-python/pyflink/table/expression.py b/flink-python/pyflink/table/expression.py index e82e3bc02d2f2..fe5fb4c960130 100644 --- a/flink-python/pyflink/table/expression.py +++ b/flink-python/pyflink/table/expression.py @@ -1984,6 +1984,24 @@ def map_from_entries(self) -> 'Expression': """ return _unary_op("mapFromEntries")(self) + def map_contains_key(self, key) -> 'Expression': + """ + Returns True if the given key exists in the map, False otherwise. Returns None if the map + is None. + + If the search key is None, the function returns True when the map contains a None key. + The given key is cast implicitly to the map's key type where Flink's implicit casting + rules allow it; otherwise the call fails validation. + + Examples: + :: + + >>> map_("a", 1, "b", 2).map_contains_key("a") # True + >>> map_("a", 1, "b", 2).map_contains_key("z") # False + >>> map_(1, "a").map_contains_key(lit(1, DataTypes.TINYINT())) # True + """ + return _binary_op("mapContainsKey")(self, key) + # ---------------------------- time definition functions ----------------------------- @property diff --git a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/BaseExpressions.java b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/BaseExpressions.java index bd95cb9a73f4f..3f8cf41be96a3 100644 --- a/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/BaseExpressions.java +++ b/flink-table/flink-table-api-java/src/main/java/org/apache/flink/table/api/internal/BaseExpressions.java @@ -161,6 +161,7 @@ import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.LPAD; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.LTRIM; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.MAKE_VALID_UTF8; +import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.MAP_CONTAINS_KEY; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.MAP_ENTRIES; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.MAP_FROM_ENTRIES; import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.MAP_KEYS; @@ -1969,6 +1970,27 @@ public OutType mapEntries() { return toApiSpecificExpression(unresolvedCall(MAP_ENTRIES, toExpr())); } + /** + * Returns {@code TRUE} if the given key exists in the map, {@code FALSE} otherwise. Returns + * {@code NULL} if the map is {@code NULL}. + * + *

If the search key is {@code NULL}, the function returns {@code TRUE} when the map contains + * a {@code NULL} key. The given key is cast implicitly to the map's key type where Flink's + * implicit casting rules allow it; otherwise the call fails validation. + * + *

Examples: + * + * 

{@code
+     * map("a", 1, "b", 2).mapContainsKey("a") // TRUE
+     * map("a", 1, "b", 2).mapContainsKey("z") // FALSE
+     * map(1, "a").mapContainsKey(lit(1).cast(DataTypes.TINYINT())) // TRUE
+     * }
+ */ + public OutType mapContainsKey(InType key) { + return toApiSpecificExpression( + unresolvedCall(MAP_CONTAINS_KEY, toExpr(), objectToExpression(key))); + } + /** * Returns a map created from the given array of entries. Each entry must be a row with exactly * two fields, where the first field becomes the key and the second one the value. diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java index a2a0cb5547f95..dcb2c6559ee1c 100644 --- a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java +++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java @@ -112,6 +112,7 @@ import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.INDEX; import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.JSON_ARGUMENT; import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.LATERAL_SNAPSHOT_INPUT_TYPE_STRATEGY; +import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.MAP_KEY_ARG; import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.ML_PREDICT_INPUT_TYPE_STRATEGY; import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.TO_CHANGELOG_INPUT_TYPE_STRATEGY; import static org.apache.flink.table.types.inference.strategies.SpecificInputTypeStrategies.TWO_EQUALS_COMPARABLE; @@ -211,6 +212,21 @@ ANY, and(logical(LogicalTypeRoot.BOOLEAN), LITERAL) "org.apache.flink.table.runtime.functions.scalar.MapEntriesFunction") .build(); + public static final BuiltInFunctionDefinition MAP_CONTAINS_KEY = + BuiltInFunctionDefinition.newBuilder() + .name("MAP_CONTAINS_KEY") + .kind(SCALAR) + .inputTypeStrategy( + sequence( + List.of("map", "key"), + List.of(logical(LogicalTypeRoot.MAP), MAP_KEY_ARG))) + .outputTypeStrategy( + nullableIfArgs( + ConstantArgumentCount.of(0), explicit(DataTypes.BOOLEAN()))) + .runtimeClass( + "org.apache.flink.table.runtime.functions.scalar.MapContainsKeyFunction") + .build(); + public static final BuiltInFunctionDefinition MAP_FROM_ARRAYS = BuiltInFunctionDefinition.newBuilder() .name("MAP_FROM_ARRAYS") diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/MapKeyArgumentTypeStrategy.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/MapKeyArgumentTypeStrategy.java new file mode 100644 index 0000000000000..da457b18db539 --- /dev/null +++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/MapKeyArgumentTypeStrategy.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.flink.table.types.inference.strategies; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.functions.BuiltInFunctionDefinitions; +import org.apache.flink.table.functions.FunctionDefinition; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.inference.ArgumentTypeStrategy; +import org.apache.flink.table.types.inference.CallContext; +import org.apache.flink.table.types.inference.Signature.Argument; +import org.apache.flink.table.types.logical.LogicalType; +import org.apache.flink.table.types.logical.MapType; + +import java.util.List; +import java.util.Optional; + +import static org.apache.flink.table.types.logical.utils.LogicalTypeCasts.supportsImplicitCast; + +/** + * Specific {@link ArgumentTypeStrategy} for the key argument of {@link + * BuiltInFunctionDefinitions#MAP_CONTAINS_KEY}. + */ +@Internal +class MapKeyArgumentTypeStrategy implements ArgumentTypeStrategy { + + @Override + public Optional inferArgumentType( + CallContext callContext, int argumentPos, boolean throwOnFailure) { + List argumentTypes = callContext.getArgumentDataTypes(); + final int mapArgumentPos = 0; + final MapType mapType = (MapType) argumentTypes.get(mapArgumentPos).getLogicalType(); + final LogicalType actualKeyType = argumentTypes.get(argumentPos).getLogicalType(); + LogicalType expectedKeyType = mapType.getKeyType(); + + if (!expectedKeyType.isNullable() && actualKeyType.isNullable()) { + expectedKeyType = expectedKeyType.copy(true); + } + + if (supportsImplicitCast(actualKeyType, expectedKeyType)) { + return Optional.of(DataTypes.of(expectedKeyType)); + } + return callContext.fail( + throwOnFailure, + "Unsupported argument type. Expected type '%s' but actual type was '%s'.", + expectedKeyType, + actualKeyType); + } + + @Override + public Argument getExpectedArgument(FunctionDefinition functionDefinition, int argumentPos) { + return Argument.of(""); + } +} diff --git a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/SpecificInputTypeStrategies.java b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/SpecificInputTypeStrategies.java index 5f64808afa203..e012f310953b3 100644 --- a/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/SpecificInputTypeStrategies.java +++ b/flink-table/flink-table-common/src/main/java/org/apache/flink/table/types/inference/strategies/SpecificInputTypeStrategies.java @@ -113,6 +113,9 @@ public static InputTypeStrategy plainJsonPath(final InputTypeStrategy signatures public static final ArgumentTypeStrategy ARRAY_OF_ENTRIES_ARG = new ArrayOfEntriesArgumentTypeStrategy(); + /** Argument type derived from the map key type. */ + public static final ArgumentTypeStrategy MAP_KEY_ARG = new MapKeyArgumentTypeStrategy(); + /** * Input strategy for {@link BuiltInFunctionDefinitions#JSON_OBJECT}. * diff --git a/flink-table/flink-table-common/src/test/java/org/apache/flink/table/types/inference/InputTypeStrategiesTest.java b/flink-table/flink-table-common/src/test/java/org/apache/flink/table/types/inference/InputTypeStrategiesTest.java index e6973c3e0d0a7..38a2c3b86c845 100644 --- a/flink-table/flink-table-common/src/test/java/org/apache/flink/table/types/inference/InputTypeStrategiesTest.java +++ b/flink-table/flink-table-common/src/test/java/org/apache/flink/table/types/inference/InputTypeStrategiesTest.java @@ -644,6 +644,43 @@ ANY, explicit(DataTypes.INT()) .expectArgumentTypes( DataTypes.ARRAY(DataTypes.INT().notNull()).notNull(), DataTypes.INT()), + TestSpec.forStrategy( + "MapKey argument type strategy implicitly casts the key", + sequence( + logical(LogicalTypeRoot.MAP), + SpecificInputTypeStrategies.MAP_KEY_ARG)) + .calledWithArgumentTypes( + DataTypes.MAP(DataTypes.BIGINT().notNull(), DataTypes.STRING()), + DataTypes.INT().notNull()) + .expectSignature("f(, )") + .expectArgumentTypes( + DataTypes.MAP(DataTypes.BIGINT().notNull(), DataTypes.STRING()), + DataTypes.BIGINT().notNull()), + TestSpec.forStrategy( + "MapKey argument type strategy widens a NOT NULL key type " + + "for a nullable argument", + sequence( + logical(LogicalTypeRoot.MAP), + SpecificInputTypeStrategies.MAP_KEY_ARG)) + .calledWithArgumentTypes( + DataTypes.MAP(DataTypes.BIGINT().notNull(), DataTypes.STRING()) + .notNull(), + DataTypes.BIGINT()) + .expectArgumentTypes( + DataTypes.MAP(DataTypes.BIGINT().notNull(), DataTypes.STRING()) + .notNull(), + DataTypes.BIGINT()), + TestSpec.forStrategy( + "MapKey argument type strategy rejects a key that cannot be cast", + sequence( + logical(LogicalTypeRoot.MAP), + SpecificInputTypeStrategies.MAP_KEY_ARG)) + .calledWithArgumentTypes( + DataTypes.MAP(DataTypes.INT(), DataTypes.STRING()), + DataTypes.BOOLEAN()) + .expectErrorMessage( + "Unsupported argument type. Expected type 'INT' but actual " + + "type was 'BOOLEAN'."), TestSpec.forStrategy(sequence(SpecificInputTypeStrategies.ARRAY_FULLY_COMPARABLE)) .expectSignature("f(>)") .calledWithArgumentTypes(DataTypes.ARRAY(DataTypes.ROW())) diff --git a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/MapFunctionITCase.java b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/MapFunctionITCase.java index 54b3b8cec05bc..7d24b7995fd07 100644 --- a/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/MapFunctionITCase.java +++ b/flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/MapFunctionITCase.java @@ -77,6 +77,7 @@ Stream getTestSetSpecs() { mapValuesTestCases(), mapEntriesTestCases(), mapFromArraysTestCases(), + mapContainsKeyTestCases(), mapFromEntriesTestCases(), mapUnionTestCases()) .flatMap(s -> s); @@ -591,6 +592,108 @@ private Stream mapFromEntriesTestCases() { .notNull())); } + private Stream mapContainsKeyTestCases() { + return Stream.of( + TestSetSpec.forFunction( + BuiltInFunctionDefinitions.MAP_CONTAINS_KEY, "Invalid input") + .onFieldsWithData(CollectionUtil.map(entry("a", 1))) + .andDataTypes(DataTypes.MAP(DataTypes.STRING(), DataTypes.INT())) + .testTableApiValidationError( + $("f0").mapContainsKey(true), + "Unsupported argument type. Expected type 'STRING' " + + "but actual type was 'BOOLEAN NOT NULL'.") + .testSqlValidationError( + "MAP_CONTAINS_KEY(f0, TRUE)", + "Unsupported argument type. Expected type 'STRING' " + + "but actual type was 'BOOLEAN NOT NULL'."), + TestSetSpec.forFunction(BuiltInFunctionDefinitions.MAP_CONTAINS_KEY) + .onFieldsWithData( + CollectionUtil.map(entry("a", 1), entry("b", 2)), + CollectionUtil.map(entry(1, 2), entry(null, 3)), + null, + CollectionUtil.map(entry(new Integer[] {1, 2}, "x"))) + .andDataTypes( + DataTypes.MAP(DataTypes.STRING(), DataTypes.INT()), + DataTypes.MAP(DataTypes.INT(), DataTypes.INT()), + DataTypes.MAP(DataTypes.STRING(), DataTypes.INT()), + DataTypes.MAP(DataTypes.ARRAY(DataTypes.INT()), DataTypes.STRING())) + .testResult( + $("f0").mapContainsKey("a"), + "MAP_CONTAINS_KEY(f0, 'a')", + true, + DataTypes.BOOLEAN()) + // a miss is FALSE, not NULL + .testResult( + $("f0").mapContainsKey("z"), + "MAP_CONTAINS_KEY(f0, 'z')", + false, + DataTypes.BOOLEAN()) + // only a NULL map yields NULL + .testResult( + $("f2").mapContainsKey("a"), + "MAP_CONTAINS_KEY(f2, 'a')", + null, + DataTypes.BOOLEAN()) + // a NULL probe finds a NULL key + .testResult( + $("f1").mapContainsKey(nullOf(DataTypes.INT())), + "MAP_CONTAINS_KEY(f1, CAST(NULL AS INT))", + true, + DataTypes.BOOLEAN()) + // an absent NULL key is FALSE + .testResult( + $("f0").mapContainsKey(nullOf(DataTypes.STRING())), + "MAP_CONTAINS_KEY(f0, CAST(NULL AS STRING))", + false, + DataTypes.BOOLEAN()) + // a miss past a NULL key is FALSE + .testResult( + $("f1").mapContainsKey(9), + "MAP_CONTAINS_KEY(f1, 9)", + false, + DataTypes.BOOLEAN()) + // complex keys compare structurally + .testResult( + $("f3").mapContainsKey(array(1, 2)), + "MAP_CONTAINS_KEY(f3, ARRAY[1, 2])", + true, + DataTypes.BOOLEAN()) + // there is no VARIANT literal, so the key is built with PARSE_JSON + .testResult( + map(call("PARSE_JSON", "1"), lit(1)) + .mapContainsKey(call("PARSE_JSON", "1")), + "MAP_CONTAINS_KEY(MAP[PARSE_JSON('1'), 1], PARSE_JSON('1'))", + true, + DataTypes.BOOLEAN().notNull()), + TestSetSpec.forFunction( + BuiltInFunctionDefinitions.MAP_CONTAINS_KEY, "Documented examples") + .onFieldsWithData(1) + .andDataTypes(DataTypes.INT().notNull()) + // a NOT NULL map yields a NOT NULL result + .testResult( + map(lit("a"), lit(1), lit("b"), lit(2)).mapContainsKey("a"), + "MAP_CONTAINS_KEY(MAP['a', 1, 'b', 2], 'a')", + true, + DataTypes.BOOLEAN().notNull()) + .testResult( + map(lit("a"), lit(1), lit("b"), lit(2)).mapContainsKey("z"), + "MAP_CONTAINS_KEY(MAP['a', 1, 'b', 2], 'z')", + false, + DataTypes.BOOLEAN().notNull()) + .testResult( + map(nullOf(DataTypes.STRING()), lit(1)) + .mapContainsKey(nullOf(DataTypes.STRING())), + "MAP_CONTAINS_KEY(MAP[CAST(NULL AS STRING), 1], CAST(NULL AS STRING))", + true, + DataTypes.BOOLEAN().notNull()) + .testResult( + map(lit(1), lit("a")) + .mapContainsKey(lit(1).cast(DataTypes.TINYINT())), + "MAP_CONTAINS_KEY(MAP[1, 'a'], CAST(1 AS TINYINT))", + true, + DataTypes.BOOLEAN().notNull())); + } + private Stream mapUnionTestCases() { return Stream.of( TestSetSpec.forFunction(BuiltInFunctionDefinitions.MAP_UNION) diff --git a/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/functions/scalar/MapContainsKeyFunction.java b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/functions/scalar/MapContainsKeyFunction.java new file mode 100644 index 0000000000000..61703b4af38d8 --- /dev/null +++ b/flink-table/flink-table-runtime/src/main/java/org/apache/flink/table/runtime/functions/scalar/MapContainsKeyFunction.java @@ -0,0 +1,102 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you 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 org.apache.flink.table.runtime.functions.scalar; + +import org.apache.flink.annotation.Internal; +import org.apache.flink.table.api.DataTypes; +import org.apache.flink.table.data.ArrayData; +import org.apache.flink.table.data.MapData; +import org.apache.flink.table.functions.BuiltInFunctionDefinitions; +import org.apache.flink.table.functions.FunctionContext; +import org.apache.flink.table.functions.SpecializedFunction.ExpressionEvaluator; +import org.apache.flink.table.functions.SpecializedFunction.SpecializedContext; +import org.apache.flink.table.types.DataType; +import org.apache.flink.table.types.KeyValueDataType; +import org.apache.flink.util.FlinkRuntimeException; + +import javax.annotation.Nullable; + +import java.lang.invoke.MethodHandle; + +import static org.apache.flink.table.api.Expressions.$; + +/** Implementation of {@link BuiltInFunctionDefinitions#MAP_CONTAINS_KEY}. */ +@Internal +public class MapContainsKeyFunction extends BuiltInScalarFunction { + + private final ArrayData.ElementGetter keyElementGetter; + private final ExpressionEvaluator equalityEvaluator; + private transient MethodHandle equalityHandle; + + public MapContainsKeyFunction(SpecializedContext context) { + super(BuiltInFunctionDefinitions.MAP_CONTAINS_KEY, context); + final DataType mapDataType = context.getCallContext().getArgumentDataTypes().get(0); + final DataType keyDataType = ((KeyValueDataType) mapDataType).getKeyDataType(); + + keyElementGetter = ArrayData.createElementGetter(keyDataType.getLogicalType()); + equalityEvaluator = + context.createEvaluator( + $("key").isEqual($("needle")), + DataTypes.BOOLEAN(), + DataTypes.FIELD("key", keyDataType.notNull().toInternal()), + DataTypes.FIELD("needle", keyDataType.notNull().toInternal())); + } + + @Override + public void open(FunctionContext context) throws Exception { + equalityHandle = equalityEvaluator.open(context); + } + + public @Nullable Boolean eval(@Nullable MapData map, @Nullable Object needle) { + if (map == null) { + return null; + } + final ArrayData keys = map.keyArray(); + final int size = map.size(); + if (needle == null) { + // A NULL needle matches a NULL key, unlike SQL `NULL = NULL` which yields UNKNOWN. + for (int pos = 0; pos < size; pos++) { + if (keyElementGetter.getElementOrNull(keys, pos) == null) { + return true; + } + } + } else { + for (int pos = 0; pos < size; pos++) { + final Object key = keyElementGetter.getElementOrNull(keys, pos); + if (key != null && isEqual(key, needle)) { + return true; + } + } + } + return false; + } + + private boolean isEqual(final Object key, final Object needle) { + try { + return (boolean) equalityHandle.invoke(key, needle); + } catch (Throwable t) { + throw new FlinkRuntimeException(t); + } + } + + @Override + public void close() throws Exception { + equalityEvaluator.close(); + } +}