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

docs/data/sql_functions.yml

@@ -925,6 +925,21 @@ 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. A NULL key matches a NULL key in the map. 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.
+      eg.
+        -- 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))
   - 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,21 @@ 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. A NULL key matches a NULL key in the map. 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.
+      eg.
+        -- 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))
   - 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_keys
     Expression.map_union

flink-python/pyflink/table/expression.py

@@ -1966,6 +1966,23 @@ def map_entries(self) -> 'Expression':
         """
         return _unary_op("mapEntries")(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.
+
+        A None key matches a None key in the map. 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
+        """
+        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

@@ -160,6 +160,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_KEYS;
 import static org.apache.flink.table.functions.BuiltInFunctionDefinitions.MAP_UNION;
@@ -1967,6 +1968,26 @@ 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>A {@code NULL} key matches a {@code NULL} key in the map. 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
+     * }</pre>
+     */
+    public OutType mapContainsKey(InType key) {
+        return toApiSpecificExpression(
+                unresolvedCall(MAP_CONTAINS_KEY, toExpr(), objectToExpression(key)));
+    }
+
     /**
      * Returns a map created by merging at least one map. These maps should have a common map type.
      * If there are overlapping keys, the value from 'map2' will overwrite the value from 'map1',

flink-table/flink-table-common/src/main/java/org/apache/flink/table/functions/BuiltInFunctionDefinitions.java

@@ -113,6 +113,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(
+                                    Arrays.asList("map", "key"),
+                                    Arrays.asList(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,65 @@
+/*
+ * 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.Optional;
+
+import static org.apache.flink.table.types.logical.utils.LogicalTypeCasts.supportsImplicitCast;
+
+/**
+ * Specific {@link ArgumentTypeStrategy} for {@link BuiltInFunctionDefinitions#MAP_CONTAINS_KEY}.
+ */
+@Internal
+class MapKeyArgumentTypeStrategy implements ArgumentTypeStrategy {
+
+    @Override
+    public Optional<DataType> inferArgumentType(
+            CallContext callContext, int argumentPos, boolean throwOnFailure) {
+        final MapType mapType =
+                (MapType) callContext.getArgumentDataTypes().get(0).getLogicalType();
+        final LogicalType actualKeyType =
+                callContext.getArgumentDataTypes().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 Optional.empty();
+    }
+
+    @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

@@ -104,6 +104,9 @@ public static InputTypeStrategy windowTimeIndicator() {
     public static final ArgumentTypeStrategy ARRAY_FULLY_COMPARABLE =
             new ArrayComparableElementArgumentTypeStrategy(StructuredComparison.FULL);
 
+    /** 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(
+                                "Invalid input arguments. Expected signatures are:\n"
+                                        + "f(<MAP>, <MAP KEY>)"),
                 TestSpec.forStrategy(sequence(SpecificInputTypeStrategies.ARRAY_FULLY_COMPARABLE))
                         .expectSignature("f(<ARRAY<COMPARABLE>>)")
                         .calledWithArgumentTypes(DataTypes.ARRAY(DataTypes.ROW()))

flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/functions/MapFunctionITCase.java

@@ -46,10 +46,12 @@
 import static org.apache.flink.table.api.DataTypes.TIME;
 import static org.apache.flink.table.api.DataTypes.TIMESTAMP;
 import static org.apache.flink.table.api.Expressions.$;
+import static org.apache.flink.table.api.Expressions.array;
 import static org.apache.flink.table.api.Expressions.call;
 import static org.apache.flink.table.api.Expressions.lit;
 import static org.apache.flink.table.api.Expressions.map;
 import static org.apache.flink.table.api.Expressions.mapFromArrays;
+import static org.apache.flink.table.api.Expressions.nullOf;
 import static org.apache.flink.util.CollectionUtil.entry;
 
 /** Test {@link BuiltInFunctionDefinitions#MAP} and its return type. */
@@ -73,6 +75,7 @@ Stream<TestSetSpec> getTestSetSpecs() {
                         mapValuesTestCases(),
                         mapEntriesTestCases(),
                         mapFromArraysTestCases(),
+                        mapContainsKeyTestCases(),
                         mapUnionTestCases())
                 .flatMap(s -> s);
     }
@@ -406,6 +409,102 @@ private Stream<TestSetSpec> mapFromArraysTestCases() {
                                         DataTypes.STRING(), DataTypes.ARRAY(DataTypes.INT()))));
     }
 
+    private Stream<TestSetSpec> 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),
+                                "Invalid input arguments. Expected signatures are:\n"
+                                        + "MAP_CONTAINS_KEY(map <MAP>, key <MAP KEY>)")
+                        .testSqlValidationError(
+                                "MAP_CONTAINS_KEY(f0, TRUE)",
+                                "Invalid input arguments. Expected signatures are:\n"
+                                        + "MAP_CONTAINS_KEY(map <MAP>, key <MAP KEY>)"),
+                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()));
+    }
+
     private Stream<TestSetSpec> mapUnionTestCases() {
         return Stream.of(
                 TestSetSpec.forFunction(BuiltInFunctionDefinitions.MAP_UNION)