[FLINK-40355][table] Add new MAP_CONTAINS_KEY function

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.

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。请注意两个数组的长度应该相等。

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

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

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}.
+     *
+     * <p>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.
+     *
+     * <p>Examples:
+     *
+     * <pre>{@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
+     * }</pre>
+     */
+    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.

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")

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<DataType> inferArgumentType(
+            CallContext callContext, int argumentPos, boolean throwOnFailure) {
+        List<DataType> 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("<MAP KEY>");
+    }
+}

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}.
      *

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(<MAP>, <MAP KEY>)")
+                        .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(<ARRAY<COMPARABLE>>)")
                         .calledWithArgumentTypes(DataTypes.ARRAY(DataTypes.ROW()))