feat: Add Spark transform_keys function

facebookincubator · Jan 9, 2025 · 8c98dcf · 8c98dcf
1 parent ff943fd
commit 8c98dcf
Show file tree

Hide file tree

Showing 9 changed files with 36 additions and 3 deletions.
diff --git a/velox/functions/lib/CMakeLists.txt b/velox/functions/lib/CMakeLists.txt
@@ -39,6 +39,7 @@ velox_add_library(
   StringEncodingUtils.cpp
   SubscriptUtil.cpp
   TimeUtils.cpp
+  TransformKeys.cpp
   Utf8Utils.cpp)
 
 velox_link_libraries(

diff --git a/velox/functions/prestosql/TransformKeys.cpp → velox/functions/lib/TransformKeys.cpp b/velox/functions/prestosql/TransformKeys.cpp → velox/functions/lib/TransformKeys.cpp
@@ -13,6 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+#include "velox/functions/lib/TransformKeys.h"
 #include "velox/expression/Expr.h"
 #include "velox/expression/VectorFunction.h"
 #include "velox/functions/lib/CheckDuplicateKeys.h"
@@ -142,4 +143,8 @@ VELOX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
     exec::VectorFunctionMetadataBuilder().defaultNullBehavior(false).build(),
     std::make_unique<TransformKeysFunction>());
 
+void registerTransformKeysFunction(const std::string& name) {
+  VELOX_REGISTER_VECTOR_FUNCTION(udf_transform_keys, name);
+}
+
 } // namespace facebook::velox::functions
diff --git a/velox/functions/lib/TransformKeys.h b/velox/functions/lib/TransformKeys.h
@@ -0,0 +1,25 @@
+/*
+ * 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 <string>
+
+namespace facebook::velox::functions {
+
+void registerTransformKeysFunction(const std::string& name);
+
+} // namespace facebook::velox::functions
diff --git a/velox/functions/lib/tests/CMakeLists.txt b/velox/functions/lib/tests/CMakeLists.txt
@@ -24,6 +24,7 @@ add_executable(
   Re2FunctionsTest.cpp
   RepeatTest.cpp
   TDigestTest.cpp
+  TransformKeysTest.cpp
   Utf8Test.cpp
   ZetaDistributionTest.cpp)
 

diff --git a/...ons/prestosql/tests/TransformKeysTest.cpp → ...functions/lib/tests/TransformKeysTest.cpp b/...ons/prestosql/tests/TransformKeysTest.cpp → ...functions/lib/tests/TransformKeysTest.cpp
diff --git a/velox/functions/prestosql/CMakeLists.txt b/velox/functions/prestosql/CMakeLists.txt
@@ -51,7 +51,6 @@ velox_add_library(
   Subscript.cpp
   ToUtf8.cpp
   Transform.cpp
-  TransformKeys.cpp
   TransformValues.cpp
   TypeOf.cpp
   URIParser.cpp

diff --git a/velox/functions/prestosql/registration/MapFunctionsRegistration.cpp b/velox/functions/prestosql/registration/MapFunctionsRegistration.cpp
@@ -17,6 +17,7 @@
 #include "velox/expression/VectorFunction.h"
 #include "velox/functions/Registerer.h"
 #include "velox/functions/lib/MapConcat.h"
+#include "velox/functions/lib/TransformKeys.h"
 #include "velox/functions/prestosql/Map.h"
 #include "velox/functions/prestosql/MapFunctions.h"
 #include "velox/functions/prestosql/MapNormalize.h"
@@ -82,7 +83,6 @@ void registerMapKeyExists(const std::string& prefix) {
 
 void registerMapFunctions(const std::string& prefix) {
   VELOX_REGISTER_VECTOR_FUNCTION(udf_map_filter, prefix + "map_filter");
-  VELOX_REGISTER_VECTOR_FUNCTION(udf_transform_keys, prefix + "transform_keys");
   VELOX_REGISTER_VECTOR_FUNCTION(
       udf_transform_values, prefix + "transform_values");
   registerMapFunction(prefix + "map", false /*allowDuplicateKeys*/);
@@ -102,6 +102,7 @@ void registerMapFunctions(const std::string& prefix) {
   VELOX_REGISTER_VECTOR_FUNCTION(
       udf_no_values_match, prefix + "no_values_match");
 
+  registerTransformKeysFunction(prefix + "transform_keys");
   registerMapConcatFunction(prefix + "map_concat");
 
   registerFunction<

diff --git a/velox/functions/prestosql/tests/CMakeLists.txt b/velox/functions/prestosql/tests/CMakeLists.txt
@@ -95,7 +95,6 @@ add_executable(
   SplitToMapTest.cpp
   StringFunctionsTest.cpp
   TimestampWithTimeZoneCastTest.cpp
-  TransformKeysTest.cpp
   TransformTest.cpp
   TransformValuesTest.cpp
   TrimFunctionsTest.cpp

diff --git a/velox/functions/sparksql/registration/RegisterMap.cpp b/velox/functions/sparksql/registration/RegisterMap.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 #include "velox/functions/lib/RegistrationHelpers.h"
+#include "velox/functions/lib/TransformKeys.h"
 #include "velox/functions/sparksql/Size.h"
 
 namespace facebook::velox::functions {
@@ -40,6 +41,7 @@ void registerMapFunctions(const std::string& prefix) {
   // This is the semantics of spark.sql.ansi.enabled = false.
   registerElementAtFunction(prefix + "element_at", true);
   registerSize(prefix + "size");
+  registerTransformKeysFunction(prefix + "transform_keys");
 }
 } // namespace sparksql
 } // namespace facebook::velox::functions