misc: Add ability to remove a function from all registries (facebooki…

…ncubator#12537)

Summary:

This addresses an internal use-case where, during initialization,
all Presto functions are registered before some are overwritten
with custom implementations. Since vector functions are prioritized
over simple functions, this approach enables replacing a previously
registered vector function with a simple function. To simplify the
API and prevent confusion between removing specific signatures versus
all, this API is designed to remove all signatures, providing a clean
slate for the user.

Reviewed By: kgpai

Differential Revision: D70594709

velox/expression/SimpleFunctionRegistry.cpp

@@ -68,6 +68,11 @@ bool SimpleFunctionRegistry::registerFunctionInternal(
   });
 }
 
+void SimpleFunctionRegistry::removeFunction(const std::string& name) {
+  const auto sanitizedName = sanitizeName(name);
+  registeredFunctions_.withWLock([&](auto& map) { map.erase(sanitizedName); });
+}
+
 namespace {
 // This function is not thread safe. It should be called only from within a
 // synchronized read region of registeredFunctions_.

velox/expression/SimpleFunctionRegistry.h

@@ -91,6 +91,8 @@ class SimpleFunctionRegistry {
     }
   }
 
+  void removeFunction(const std::string& name);
+
   std::vector<std::string> getFunctionNames() const {
     std::vector<std::string> result;
     registeredFunctions_.withRLock([&](const auto& map) {

velox/functions/FunctionRegistry.cpp

@@ -181,4 +181,10 @@ resolveVectorFunctionWithMetadata(
   return exec::resolveVectorFunctionWithMetadata(functionName, argTypes);
 }
 
+void removeFunction(const std::string& functionName) {
+  exec::mutableSimpleFunctions().removeFunction(functionName);
+  exec::vectorFunctionFactories().withWLock(
+      [&](auto& functionMap) { functionMap.erase(functionName); });
+}
+
 } // namespace facebook::velox

velox/functions/FunctionRegistry.h

@@ -97,6 +97,10 @@ resolveVectorFunctionWithMetadata(
     const std::string& functionName,
     const std::vector<TypePtr>& argTypes);
 
+/// Given name of a function, removes it from both the simple and vector
+/// function registries (including all signatures).
+void removeFunction(const std::string& functionName);
+
 /// Clears the function registry.
 void clearFunctionRegistry();
 

velox/functions/tests/FunctionRegistryTest.cpp

@@ -81,6 +81,10 @@ inline void registerTestFunctions() {
   VELOX_REGISTER_VECTOR_FUNCTION(udf_vector_func_three, "vector_func_three");
   VELOX_REGISTER_VECTOR_FUNCTION(udf_vector_func_four, "vector_func_four");
 }
+
+inline void registerTestVectorFunctionOne(const std::string& functionName) {
+  VELOX_REGISTER_VECTOR_FUNCTION(udf_vector_func_one, functionName);
+}
 } // namespace
 
 class FunctionRegistryTest : public testing::Test {
@@ -107,6 +111,44 @@ class FunctionRegistryTest : public testing::Test {
   }
 };
 
+TEST_F(FunctionRegistryTest, removeFunction) {
+  const std::string functionName = "func_to_remove";
+  auto checkFunctionExists = [&](const std::string& name,
+                                 bool vectorFuncSignatures,
+                                 bool simpleFuncSignatures) {
+    EXPECT_EQ(
+        getFunctionSignatures(name).size(),
+        vectorFuncSignatures + simpleFuncSignatures);
+    EXPECT_EQ(getVectorFunctionSignatures().count(name), vectorFuncSignatures);
+    EXPECT_EQ(
+        exec::simpleFunctions().getFunctionSignatures(name).size(),
+        simpleFuncSignatures);
+  };
+
+  checkFunctionExists(functionName, 0, 0);
+
+  // Only vector function registered
+  registerTestVectorFunctionOne(functionName);
+  checkFunctionExists(functionName, 1, 0);
+  removeFunction(functionName);
+  checkFunctionExists(functionName, 0, 0);
+
+  // Only simple function registered
+  registerFunction<FuncOne, Varchar, Varchar>(
+      std::vector<std::string>{functionName});
+  checkFunctionExists(functionName, 0, 1);
+  removeFunction(functionName);
+  checkFunctionExists(functionName, 0, 0);
+
+  // Both vector and simple function registered
+  registerTestVectorFunctionOne(functionName);
+  registerFunction<FuncOne, Varchar, Varchar>(
+      std::vector<std::string>{functionName});
+  checkFunctionExists(functionName, 1, 1);
+  removeFunction(functionName);
+  checkFunctionExists(functionName, 0, 0);
+}
+
 TEST_F(FunctionRegistryTest, getFunctionSignaturesByName) {
   {
     auto signatures = getFunctionSignatures("func_one");