feat: Add API to return function signatures by name (#12289)

Summary:
Pull Request resolved: #12289

Add velox::getFunctionSignatures(name) API to get a list of function signatures by name.

Use this API to optimize lambda function resolution logic.

Reviewed By: pedroerp

Differential Revision: D69368657

fbshipit-source-id: e8713b2e5791c1186495a0a6423b1ab374ea8a6b

pyvelox/signatures.cpp

@@ -106,7 +106,7 @@ void addSignatureBindings(py::module& m, bool asModuleLocalDefinitions) {
 
   m.def(
       "get_function_signatures",
-      &getFunctionSignatures,
+      []() { return getFunctionSignatures(); },
       py::return_value_policy::reference,
       "Returns a dictionary of the current signatures.");
 

velox/functions/FunctionRegistry.cpp

@@ -64,6 +64,27 @@ FunctionSignatureMap getFunctionSignatures() {
   return result;
 }
 
+std::vector<const exec::FunctionSignature*> getFunctionSignatures(
+    const std::string& functionName) {
+  // Some functions have both simple and vector implementations (for different
+  // signatures). Collect all signatures.
+  // Check simple functions first.
+  auto signatures = exec::simpleFunctions().getFunctionSignatures(functionName);
+
+  // Check vector functions.
+  auto& vectorFunctions = exec::vectorFunctionFactories();
+  vectorFunctions.withRLock([&](const auto& functions) {
+    auto it = functions.find(functionName);
+    if (it != functions.end()) {
+      for (const auto& signature : it->second.signatures) {
+        signatures.push_back(signature.get());
+      }
+    }
+  });
+
+  return signatures;
+}
+
 FunctionSignatureMap getVectorFunctionSignatures() {
   FunctionSignatureMap result;
   populateVectorFunctionSignatures(result);

velox/functions/FunctionRegistry.h

@@ -31,6 +31,11 @@ using FunctionSignatureMap = std::
 /// The mapping is function name -> list of function signatures
 FunctionSignatureMap getFunctionSignatures();
 
+/// Returns a list of function signatures for a given function name. Returns
+/// empty list if function with specified name not found.
+std::vector<const exec::FunctionSignature*> getFunctionSignatures(
+    const std::string& functionName);
+
 /// Returns a mapping of all Vector functions registered in Velox
 /// The mapping is function name -> list of function signatures
 FunctionSignatureMap getVectorFunctionSignatures();

velox/functions/tests/FunctionRegistryTest.cpp

@@ -107,6 +107,34 @@ class FunctionRegistryTest : public testing::Test {
   }
 };
 
+TEST_F(FunctionRegistryTest, getFunctionSignaturesByName) {
+  {
+    auto signatures = getFunctionSignatures("func_one");
+    ASSERT_EQ(signatures.size(), 1);
+    ASSERT_EQ(
+        signatures.at(0)->toString(),
+        exec::FunctionSignatureBuilder()
+            .returnType("varchar")
+            .argumentType("varchar")
+            .build()
+            ->toString());
+  }
+
+  {
+    auto signatures = getFunctionSignatures("vector_func_one");
+    ASSERT_EQ(signatures.size(), 1);
+    ASSERT_EQ(
+        signatures.at(0)->toString(),
+        exec::FunctionSignatureBuilder()
+            .returnType("bigint")
+            .argumentType("varchar")
+            .build()
+            ->toString());
+  }
+
+  ASSERT_TRUE(getFunctionSignatures("non-existent-function").empty());
+}
+
 TEST_F(FunctionRegistryTest, getFunctionSignatures) {
   auto functionSignatures = getFunctionSignatures();
   ASSERT_EQ(functionSignatures.size(), 14);

velox/parse/Expressions.cpp

@@ -384,11 +384,9 @@ const exec::FunctionSignature* findLambdaSignature(
 const exec::FunctionSignature* findLambdaSignature(
     const std::shared_ptr<const CallExpr>& callExpr) {
   // Look for a scalar lambda function.
-  auto allSignatures = getFunctionSignatures();
-  auto it = allSignatures.find(callExpr->getFunctionName());
-
-  if (it != allSignatures.end()) {
-    return findLambdaSignature(it->second, callExpr);
+  auto scalarSignatures = getFunctionSignatures(callExpr->getFunctionName());
+  if (!scalarSignatures.empty()) {
+    return findLambdaSignature(scalarSignatures, callExpr);
   }
 
   // Look for an aggregate lambda function.