Add support for signature bindings in PyVelox (facebookincubator#4333)

Summary:
Adding support for signature bindings will allow us to track changes in signatures across changes. This is required so we can  diff changes in signatures during CI time.

Pull Request resolved: facebookincubator#4333

Reviewed By: pedroerp

Differential Revision: D44179392

Pulled By: kgpai

fbshipit-source-id: eec405cc745222e18aa06e0ddc2912e880ee68f6

CMakeLists.txt

@@ -136,7 +136,7 @@ if(${VELOX_BUILD_PYTHON_PACKAGE})
   set(VELOX_ENABLE_AGGREGATES OFF)
   set(VELOX_ENABLE_HIVE_CONNECTOR OFF)
   set(VELOX_ENABLE_TPCH_CONNECTOR OFF)
-  set(VELOX_ENABLE_SPARK_FUNCTIONS OFF)
+  set(VELOX_ENABLE_SPARK_FUNCTIONS ON)
   set(VELOX_ENABLE_EXAMPLES OFF)
   set(VELOX_ENABLE_S3 OFF)
   set(VELOX_ENABLE_SUBSTRAIT OFF)

pyvelox/CMakeLists.txt

@@ -17,12 +17,18 @@ if(VELOX_BUILD_PYTHON_PACKAGE)
   include_directories(SYSTEM ${CMAKE_SOURCE_DIR})
   add_definitions(-DCREATE_PYVELOX_MODULE -DVELOX_DISABLE_GOOGLETEST)
   # Define our Python module:
-  pybind11_add_module(pyvelox MODULE pyvelox.cpp pyvelox.h)
-
+  pybind11_add_module(pyvelox MODULE pyvelox.cpp signatures.cpp)
   # Link with Velox:
   target_link_libraries(
-    pyvelox PRIVATE velox_type velox_vector velox_core velox_exec
-                    velox_functions_prestosql velox_parse_parser)
+    pyvelox
+    PRIVATE velox_type
+            velox_vector
+            velox_core
+            velox_exec
+            velox_functions_prestosql
+            velox_parse_parser
+            velox_functions_prestosql
+            velox_functions_spark)
 
   install(TARGETS pyvelox LIBRARY DESTINATION .)
 else()

pyvelox/pyvelox.cpp

@@ -14,7 +14,8 @@
  * limitations under the License.
  */
 
-#include "pyvelox.h" // @manual
+#include "pyvelox.h"
+#include "signatures.h"
 
 namespace facebook::velox::py {
 using namespace velox;
@@ -133,7 +134,7 @@ PYBIND11_MODULE(pyvelox, m) {
   )pbdoc";
 
   addVeloxBindings(m);
-
+  addSignatureBindings(m);
   m.attr("__version__") = "dev";
 }
 #endif

pyvelox/signatures.cpp

@@ -0,0 +1,79 @@
+/*
+ * 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.
+ */
+
+#include "signatures.h" // @manual
+#include "velox/functions/FunctionRegistry.h"
+#include "velox/functions/prestosql/registration/RegistrationFunctions.h"
+#include "velox/functions/sparksql/Register.h"
+
+namespace facebook::velox::py {
+
+namespace py = pybind11;
+
+void registerPrestoFunctions(const std::string& prefix) {
+  facebook::velox::functions::prestosql::registerAllScalarFunctions(prefix);
+}
+
+void registerSparkFunctions(const std::string& prefix) {
+  facebook::velox::functions::sparksql::registerFunctions(prefix);
+}
+
+void addSignatureBindings(py::module& m, bool asModuleLocalDefinitions) {
+  // TypeSignature
+  py::class_<exec::TypeSignature> typeSignature(
+      m, "TypeSignature", py::module_local(asModuleLocalDefinitions));
+  typeSignature.def("__str__", &exec::TypeSignature::toString);
+  typeSignature.def("base_name", &exec::TypeSignature::baseName);
+  typeSignature.def("parameters", &exec::TypeSignature::parameters);
+
+  // FunctionSignature
+  py::class_<exec::FunctionSignature> functionSignature(
+      m, "FunctionSignature", py::module_local(asModuleLocalDefinitions));
+
+  functionSignature.def("__str__", &exec::FunctionSignature::toString);
+  functionSignature.def("return_type", &exec::FunctionSignature::returnType);
+  functionSignature.def(
+      "argument_types", &exec::FunctionSignature::argumentTypes);
+  functionSignature.def(
+      "variable_arity", &exec::FunctionSignature::variableArity);
+  functionSignature.def("variables", &exec::FunctionSignature::variables);
+  functionSignature.def(
+      "constant_arguments", &exec::FunctionSignature::constantArguments);
+
+  m.def(
+      "clear_signatures",
+      &clearFunctionRegistry,
+      "Clears the function registry.");
+
+  m.def(
+      "register_spark_signatures",
+      &registerSparkFunctions,
+      "Adds Spark signatures to the function registry.",
+      py::arg("prefix") = "");
+
+  m.def(
+      "register_presto_signatures",
+      &registerPrestoFunctions,
+      "Adds Presto signatures to the function registry.",
+      py::arg("prefix") = "");
+
+  m.def(
+      "get_function_signatures",
+      &getFunctionSignatures,
+      py::return_value_policy::reference,
+      "Returns a dictionary of the current signatures.");
+}
+} // namespace facebook::velox::py

pyvelox/signatures.h

@@ -0,0 +1,36 @@
+/*
+ * 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 <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace facebook::velox::py {
+
+namespace py = pybind11;
+
+/// Adds Function signature bindings to module m.
+/// This adds bindings to select Presto and Spark function signatures.
+///
+/// @param m Module to add bindings to.
+/// @param asModuleLocalDefinitions If true then these bindings are only
+///  visible inside the module. Refer to
+///  https://pybind11.readthedocs.io/en/stable/advanced/classes.html#module-local-class-bindings
+///  for further details.
+void addSignatureBindings(py::module& m, bool asModuleLocalDefinitions = true);
+
+} // namespace facebook::velox::py

pyvelox/test/test_signatures.py

@@ -0,0 +1,59 @@
+# 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.
+
+import unittest
+
+import pyvelox.pyvelox as pv
+
+
+class TestFunctionSignatures(unittest.TestCase):
+    def test_clear_signatures(self):
+        pv.clear_signatures()
+        signatures = pv.get_function_signatures()
+        self.assertEqual(len(signatures), 0)
+
+    def test_get_signatures(self):
+        pv.register_presto_signatures()
+        presto_signatures = pv.get_function_signatures()
+        self.assertTrue(len(presto_signatures) > 0)
+
+        pv.clear_signatures()
+        pv.register_spark_signatures()
+        spark_signatures = pv.get_function_signatures()
+        self.assertTrue(len(spark_signatures) > 0)
+
+    def test_function_signature(self):
+        pv.clear_signatures()
+        pv.register_presto_signatures()
+        presto_signatures = pv.get_function_signatures()
+
+        concat_signatures = presto_signatures["concat"]
+        self.assertTrue(len(concat_signatures) > 0)
+        self.assertEqual(str(concat_signatures[0].return_type()), "varchar")
+        self.assertEqual(str(concat_signatures[0]), "(varchar,varchar...) -> varchar")
+
+    def test_function_prefix(self):
+        pv.clear_signatures()
+        pv.register_presto_signatures("foo")
+        presto_signatures = pv.get_function_signatures()
+
+        concat_signatures = presto_signatures["fooconcat"]
+        self.assertTrue(len(concat_signatures) > 0)
+
+        pv.clear_signatures()
+        pv.register_spark_signatures("bar")
+        spark_signatures = pv.get_function_signatures()
+
+        concat_signatures = spark_signatures["barconcat"]
+        self.assertTrue(len(concat_signatures) > 0)

setup.py

@@ -27,6 +27,7 @@
 import subprocess
 import sys
 from pathlib import Path
+
 from setuptools import Extension, find_packages, setup
 from setuptools.command.build_ext import build_ext
 
@@ -128,7 +129,6 @@ def build_extension(self, ext):
             "-DVELOX_BUILD_PYTHON_PACKAGE=ON",
             f"-DPYTHON_EXECUTABLE={exec_path} ",
             "-DVELOX_CODEGEN_SUPPORT=OFF",
-            "-DVELOX_BUILD_MINIMAL=ON",
         ]
         build_args = []
 

velox/expression/FunctionRegistry.h

@@ -152,6 +152,10 @@ class FunctionRegistry {
         : std::nullopt;
   }
 
+  void clearRegistry() {
+    registeredFunctions_.clear();
+  }
+
  private:
   template <typename T>
   static std::unique_ptr<T> CreateUdf() {

velox/functions/FunctionRegistry.cpp

@@ -77,6 +77,12 @@ FunctionSignatureMap getFunctionSignatures() {
   return result;
 }
 
+void clearFunctionRegistry() {
+  exec::SimpleFunctions().clearRegistry();
+  exec::vectorFunctionFactories().withWLock(
+      [](auto& functionMap) { functionMap.clear(); });
+}
+
 std::shared_ptr<const Type> resolveFunction(
     const std::string& functionName,
     const std::vector<TypePtr>& argTypes) {

velox/functions/FunctionRegistry.h

@@ -67,4 +67,6 @@ std::shared_ptr<const Type> resolveVectorFunction(
     const std::string& functionName,
     const std::vector<TypePtr>& argTypes);
 
+/// Clears the function registry.
+void clearFunctionRegistry();
 } // namespace facebook::velox