Skip to content

Commit

Permalink
Add API to return function signatures by name
Browse files Browse the repository at this point in the history
Summary:
Add velox::getFunctionSignatures(name) API to get a list of function signatures by name.

Use this API to optimize lambda function resolution logic.

Differential Revision: D69368657
  • Loading branch information
mbasmanova authored and facebook-github-bot committed Feb 9, 2025
1 parent 710d449 commit f801149
Show file tree
Hide file tree
Showing 4 changed files with 58 additions and 5 deletions.
22 changes: 22 additions & 0 deletions velox/functions/FunctionRegistry.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,28 @@ FunctionSignatureMap getFunctionSignatures() {
return result;
}

std::vector<const exec::FunctionSignature*> getFunctionSignatures(
const std::string& functionName) {
// Check simple functions first.
auto signatures = exec::simpleFunctions().getFunctionSignatures(functionName);
if (!signatures.empty()) {
return signatures;
}

// 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);
Expand Down
5 changes: 5 additions & 0 deletions velox/functions/FunctionRegistry.h
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
28 changes: 28 additions & 0 deletions velox/functions/tests/FunctionRegistryTest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
8 changes: 3 additions & 5 deletions velox/parse/Expressions.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down

0 comments on commit f801149

Please sign in to comment.