feat: Add map_key_exists function (facebookincubator#11735)

Summary: Pull Request resolved: facebookincubator#11735

Reviewed By: xiaoxmeng, kgpai

Differential Revision: D65014614

fbshipit-source-id: c28f8f44824cc6e60c686e9648104ffb6ad36f9a

velox/functions/prestosql/MapFunctions.h

@@ -0,0 +1,35 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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.
+ */
+#pragma once
+
+#include "velox/expression/ComplexViewTypes.h"
+#include "velox/functions/Udf.h"
+
+namespace facebook::velox::functions {
+
+template <typename TExec>
+struct MapKeyExists {
+  VELOX_DEFINE_FUNCTION_TYPES(TExec);
+
+  void call(
+      bool& out,
+      const arg_type<Map<Generic<T1>, Generic<T2>>>& inputMap,
+      const arg_type<Generic<T1>>& key) {
+    out = (inputMap.find(key) != inputMap.end());
+  }
+};
+
+} // namespace facebook::velox::functions

velox/functions/prestosql/MapRemoveNullValues.h

@@ -15,6 +15,9 @@
  */
 #pragma once
 
+#include "velox/expression/ComplexViewTypes.h"
+#include "velox/functions/Udf.h"
+
 namespace facebook::velox::functions {
 
 template <typename TExec>

velox/functions/prestosql/registration/MapFunctionsRegistration.cpp

@@ -18,6 +18,7 @@
 #include "velox/functions/Registerer.h"
 #include "velox/functions/lib/MapConcat.h"
 #include "velox/functions/prestosql/Map.h"
+#include "velox/functions/prestosql/MapFunctions.h"
 #include "velox/functions/prestosql/MapNormalize.h"
 #include "velox/functions/prestosql/MapRemoveNullValues.h"
 #include "velox/functions/prestosql/MapSubset.h"
@@ -69,6 +70,14 @@ void registerMapRemoveNullValues(const std::string& prefix) {
       Map<Generic<T1>, Generic<T2>>>({prefix + "map_remove_null_values"});
 }
 
+void registerMapKeyExists(const std::string& prefix) {
+  registerFunction<
+      MapKeyExists,
+      bool,
+      Map<Generic<T1>, Generic<T2>>,
+      Generic<T1>>({prefix + "map_key_exists"});
+}
+
 } // namespace
 
 void registerMapFunctions(const std::string& prefix) {
@@ -116,6 +125,8 @@ void registerMapFunctions(const std::string& prefix) {
 
   registerMapRemoveNullValues(prefix);
 
+  registerMapKeyExists(prefix);
+
   registerFunction<
       MapNormalizeFunction,
       Map<Varchar, double>,

velox/functions/prestosql/tests/MapKeyExistsTest.cpp

@@ -0,0 +1,66 @@
+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * 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.
+ */
+#include "velox/common/base/tests/GTestUtils.h"
+#include "velox/functions/prestosql/tests/utils/FunctionBaseTest.h"
+
+using namespace facebook::velox::test;
+
+namespace facebook::velox::functions {
+namespace {
+
+class MapKeyExistsTest : public test::FunctionBaseTest {};
+
+TEST_F(MapKeyExistsTest, general) {
+  // numeric key
+  auto data = makeRowVector({makeMapVectorFromJson<int64_t, int32_t>(
+      {"{1:4, 10:5, 3:6}", "{-2:5}", "{}"})});
+
+  auto result = evaluate("map_key_exists(c0, 1)", data);
+  auto expected = makeFlatVector<bool>({true, false, false});
+
+  assertEqualVectors(expected, result);
+
+  // string key
+  data = makeRowVector({makeMapVectorFromJson<std::string, std::string>({
+      "{\"ad\":null, \"bc\":null, \"cd\":null}",
+  })});
+
+  result = evaluate("map_key_exists(c0, 'bc')", data);
+  expected = makeFlatVector<bool>(true);
+
+  assertEqualVectors(expected, result);
+}
+
+TEST_F(MapKeyExistsTest, testFloats) {
+  auto data = makeRowVector({makeMapVectorFromJson<double, int32_t>({
+      "{1:10, NaN:20, 0:null, 4:40, 5:50, 6:60}",
+      "{NaN:20, -0:20}",
+  })});
+
+  auto result = evaluate("map_key_exists(c0, cast('NaN' as REAL))", data);
+  auto expected = makeFlatVector<bool>({true, true});
+
+  result = evaluate("map_key_exists(c0, cast('0' as REAL))", data);
+  expected = makeFlatVector<bool>({true, false});
+
+  result = evaluate("map_key_exists(c0, cast('-0' as REAL))", data);
+  expected = makeFlatVector<bool>({true, true});
+
+  assertEqualVectors(expected, result);
+}
+
+} // namespace
+} // namespace facebook::velox::functions