Add FunctionRegistry APIs to retrieve function metadata

facebookincubator · Mar 29, 2024 · 8a28ec9 · 8a28ec9
1 parent cb58cba
commit 8a28ec9
Show file tree

Hide file tree

Showing 7 changed files with 152 additions and 82 deletions.
diff --git a/velox/expression/VectorFunction.cpp b/velox/expression/VectorFunction.cpp
@@ -70,6 +70,13 @@ std::optional<std::vector<FunctionSignaturePtr>> getVectorFunctionSignatures(
       });
 }
 
+std::optional<VectorFunctionMetadata> getVectorFunctionMetadata(
+    const std::string& name) {
+  return applyToVectorFunctionEntry<VectorFunctionMetadata>(
+      name,
+      [&](const auto& /*name*/, const auto& entry) { return entry.metadata; });
+}
+
 TypePtr resolveVectorFunction(
     const std::string& functionName,
     const std::vector<TypePtr>& argTypes) {

diff --git a/velox/expression/VectorFunction.h b/velox/expression/VectorFunction.h
@@ -175,6 +175,11 @@ class SimpleFunctionAdapterFactory {
 std::optional<std::vector<FunctionSignaturePtr>> getVectorFunctionSignatures(
     const std::string& name);
 
+/// Returns metadata corresponding to function with the specified name. Returns
+/// std::nullopt if there is no function with the specified name.
+std::optional<VectorFunctionMetadata> getVectorFunctionMetadata(
+    const std::string& name);
+
 /// Given name of vector function and argument types, returns
 /// the return type if function exists and have a signature that binds to the
 /// input types otherwise returns nullptr.

diff --git a/velox/functions/CMakeLists.txt b/velox/functions/CMakeLists.txt
@@ -15,7 +15,7 @@ add_library(velox_function_registry FunctionRegistry.cpp)
 add_library(velox_coverage_util CoverageUtil.cpp)
 
 target_link_libraries(velox_function_registry velox_expression velox_type
-                      velox_core velox_exception)
+                      velox_core velox_exception velox_exec)
 target_link_libraries(velox_coverage_util velox_function_registry)
 add_subdirectory(lib)
 if(${VELOX_ENABLE_PRESTO_FUNCTIONS})

diff --git a/velox/functions/CoverageUtil.cpp b/velox/functions/CoverageUtil.cpp
@@ -274,86 +274,6 @@ void printCoverageMap(
   std::cout << out.str() << std::endl;
 }
 
-// A function name is a companion function's if the name is an existing
-// aggregation functio name followed by a specific suffixes.
-bool isCompanionFunctionName(
-    const std::string& name,
-    const std::unordered_map<std::string, exec::AggregateFunctionEntry>&
-        aggregateFunctions) {
-  auto suffixOffset = name.rfind("_partial");
-  if (suffixOffset == std::string::npos) {
-    suffixOffset = name.rfind("_merge_extract");
-  }
-  if (suffixOffset == std::string::npos) {
-    suffixOffset = name.rfind("_merge");
-  }
-  if (suffixOffset == std::string::npos) {
-    suffixOffset = name.rfind("_extract");
-  }
-  if (suffixOffset == std::string::npos) {
-    return false;
-  }
-  return aggregateFunctions.count(name.substr(0, suffixOffset)) > 0;
-}
-
-/// Returns alphabetically sorted list of scalar functions available in Velox,
-/// excluding companion functions.
-std::vector<std::string> getSortedScalarNames() {
-  // Do not print "internal" functions.
-  static const std::unordered_set<std::string> kBlockList = {"row_constructor"};
-
-  auto functions = getFunctionSignatures();
-
-  std::vector<std::string> names;
-  names.reserve(functions.size());
-  exec::aggregateFunctions().withRLock([&](const auto& aggregateFunctions) {
-    for (const auto& func : functions) {
-      const auto& name = func.first;
-      if (!isCompanionFunctionName(name, aggregateFunctions) &&
-          kBlockList.count(name) == 0) {
-        names.emplace_back(name);
-      }
-    }
-  });
-  std::sort(names.begin(), names.end());
-  return names;
-}
-
-/// Returns alphabetically sorted list of aggregate functions available in
-/// Velox, excluding compaion functions.
-std::vector<std::string> getSortedAggregateNames() {
-  std::vector<std::string> names;
-  exec::aggregateFunctions().withRLock([&](const auto& functions) {
-    names.reserve(functions.size());
-    for (const auto& entry : functions) {
-      if (!isCompanionFunctionName(entry.first, functions)) {
-        names.push_back(entry.first);
-      }
-    }
-  });
-  std::sort(names.begin(), names.end());
-  return names;
-}
-
-/// Returns alphabetically sorted list of window functions available in Velox,
-/// excluding companion functions.
-std::vector<std::string> getSortedWindowNames() {
-  const auto& functions = exec::windowFunctions();
-
-  std::vector<std::string> names;
-  names.reserve(functions.size());
-  exec::aggregateFunctions().withRLock([&](const auto& aggregateFunctions) {
-    for (const auto& entry : functions) {
-      if (!isCompanionFunctionName(entry.first, aggregateFunctions) &&
-          aggregateFunctions.count(entry.first) == 0) {
-        names.emplace_back(entry.first);
-      }
-    }
-  });
-  std::sort(names.begin(), names.end());
-  return names;
-}
-
 /// Takes a super-set of simple, vector and aggregate function names and prints
 /// coverage map showing which of these functions are available in Velox.
 /// Companion functions are excluded.

diff --git a/velox/functions/FunctionRegistry.cpp b/velox/functions/FunctionRegistry.cpp
@@ -135,4 +135,96 @@ resolveVectorFunctionWithMetadata(
   return exec::resolveVectorFunctionWithMetadata(functionName, argTypes);
 }
 
+bool isCompanionFunctionName(
+    const std::string& name,
+    const std::unordered_map<std::string, exec::AggregateFunctionEntry>&
+        aggregateFunctions) {
+  auto suffixOffset = name.rfind("_partial");
+  if (suffixOffset == std::string::npos) {
+    suffixOffset = name.rfind("_merge_extract");
+  }
+  if (suffixOffset == std::string::npos) {
+    suffixOffset = name.rfind("_merge");
+  }
+  if (suffixOffset == std::string::npos) {
+    suffixOffset = name.rfind("_extract");
+  }
+  if (suffixOffset == std::string::npos) {
+    return false;
+  }
+  return aggregateFunctions.count(name.substr(0, suffixOffset)) > 0;
+}
+
+std::vector<std::string> getSortedScalarNames() {
+  // Do not print "internal" functions.
+  static const std::unordered_set<std::string> kBlockList = {
+      "row_constructor", "in", "is_null"};
+
+  auto functions = getFunctionSignatures();
+
+  std::vector<std::string> names;
+  names.reserve(functions.size());
+  exec::aggregateFunctions().withRLock([&](const auto& aggregateFunctions) {
+    for (const auto& func : functions) {
+      const auto& name = func.first;
+      if (!isCompanionFunctionName(name, aggregateFunctions) &&
+          kBlockList.count(name) == 0) {
+        names.emplace_back(name);
+      }
+    }
+  });
+  std::sort(names.begin(), names.end());
+  return names;
+}
+
+std::vector<std::string> getSortedAggregateNames() {
+  std::vector<std::string> names;
+  exec::aggregateFunctions().withRLock([&](const auto& functions) {
+    names.reserve(functions.size());
+    for (const auto& entry : functions) {
+      if (!isCompanionFunctionName(entry.first, functions)) {
+        names.push_back(entry.first);
+      }
+    }
+  });
+  std::sort(names.begin(), names.end());
+  return names;
+}
+
+std::vector<std::string> getSortedWindowNames() {
+  const auto& functions = exec::windowFunctions();
+
+  std::vector<std::string> names;
+  names.reserve(functions.size());
+  exec::aggregateFunctions().withRLock([&](const auto& aggregateFunctions) {
+    for (const auto& entry : functions) {
+      if (!isCompanionFunctionName(entry.first, aggregateFunctions) &&
+          aggregateFunctions.count(entry.first) == 0) {
+        names.emplace_back(entry.first);
+      }
+    }
+  });
+  std::sort(names.begin(), names.end());
+  return names;
+}
+
+std::optional<exec::VectorFunctionMetadata> getFunctionMetadata(
+    const std::string& functionName) {
+  auto simpleFunctionMetadata =
+      exec::simpleFunctions().getFunctionSignaturesAndMetadata(functionName);
+  if (simpleFunctionMetadata.size()) {
+    // Functions like abs are registered as simple functions for primitive
+    // types, and as a vector function for complex types like DECIMAL. So do not
+    // throw an error if function metadata is not found in simple function
+    // signature map.
+    return simpleFunctionMetadata.back().first;
+  }
+
+  auto vectorFunctionMetadata = exec::getVectorFunctionMetadata(functionName);
+  if (vectorFunctionMetadata.has_value()) {
+    return vectorFunctionMetadata.value();
+  }
+  return std::nullopt;
+}
+
 } // namespace facebook::velox
diff --git a/velox/functions/FunctionRegistry.h b/velox/functions/FunctionRegistry.h
@@ -18,6 +18,8 @@
 #include <string>
 #include <vector>
 
+#include "velox/exec/Aggregate.h"
+#include "velox/exec/WindowFunction.h"
 #include "velox/expression/FunctionMetadata.h"
 #include "velox/expression/FunctionSignature.h"
 #include "velox/type/Type.h"
@@ -85,4 +87,24 @@ resolveVectorFunctionWithMetadata(
 /// Clears the function registry.
 void clearFunctionRegistry();
 
+/// A function name is a companion function's if the name is an existing
+/// aggregation function name followed by a specific suffixes.
+bool isCompanionFunctionName(
+    const std::string& name,
+    const std::unordered_map<std::string, exec::AggregateFunctionEntry>&
+        aggregateFunctions);
+
+/// Returns sorted list of scalar function names available in Velox.
+std::vector<std::string> getSortedScalarNames();
+
+/// Returns sorted list of aggregate function names available in Velox.
+std::vector<std::string> getSortedAggregateNames();
+
+/// Returns sorted list of window function names available in Velox.
+std::vector<std::string> getSortedWindowNames();
+
+/// Get the function metadata corresponding to functionName.
+std::optional<exec::VectorFunctionMetadata> getFunctionMetadata(
+    const std::string& functionName);
+
 } // namespace facebook::velox
diff --git a/velox/functions/tests/FunctionRegistryTest.cpp b/velox/functions/tests/FunctionRegistryTest.cpp
@@ -193,7 +193,10 @@ VELOX_DECLARE_VECTOR_FUNCTION(
 VELOX_DECLARE_VECTOR_FUNCTION_WITH_METADATA(
     udf_vector_func_four,
     VectorFuncFour::signatures(),
-    exec::VectorFunctionMetadataBuilder().deterministic(false).build(),
+    exec::VectorFunctionMetadataBuilder()
+        .deterministic(false)
+        .defaultNullBehavior(false)
+        .build(),
     std::make_unique<VectorFuncFour>());
 
 inline void registerTestFunctions() {
@@ -588,4 +591,25 @@ TEST_F(FunctionRegistryTest, resolveWithMetadata) {
   result = resolveFunctionWithMetadata("non-existent-function", {VARCHAR()});
   EXPECT_FALSE(result.has_value());
 }
+
+TEST_F(FunctionRegistryTest, functionMetadata) {
+  auto checkMetadata = [&](StringView functionName,
+                           bool expectedDeterminism,
+                           bool expectedDefaultNullBehavior) {
+    auto metadata = getFunctionMetadata(functionName);
+    EXPECT_TRUE(metadata.has_value());
+    ASSERT_EQ(metadata.value().deterministic, expectedDeterminism);
+    ASSERT_EQ(
+        metadata.value().defaultNullBehavior, expectedDefaultNullBehavior);
+  };
+
+  // Validate VectorFunctionMetadata for simple functions func_one and func_two.
+  checkMetadata("func_one", false, true);
+  checkMetadata("func_two", true, false);
+
+  // Validate VectorFunctionMetadata for vector functions vector_func_three and
+  // vector_func_four.
+  checkMetadata("vector_func_three", true, true);
+  checkMetadata("vector_func_four", false, false);
+}
 } // namespace facebook::velox