test: Add custom function tests (#25480)

pramodsatya · mohsaka · soumiiow · web-flow · commit 491f8978b58f · 2025-10-03T19:38:04.000-04:00
## Description Adds e2e tests for Presto C++ custom functions. Testcases borrowed from [AbstractTestQueries](https://github.com/prestodb/presto/blob/master/presto-tests/src/main/java/com/facebook/presto/tests/AbstractTestQueries.java) in `presto-tests`. ``` == NO RELEASE NOTE == ``` Co-authored-by: mohsaka <michael.ohsaka@ibm.com> Co-authored-by: soumiiow <soumyaduriseti@gmail.com>
diff --git a/.github/workflows/prestocpp-linux-build-and-unit-test.yml b/.github/workflows/prestocpp-linux-build-and-unit-test.yml
@@ -79,6 +79,8 @@ jobs:
           github.event_name == 'schedule' || needs.changes.outputs.codechange == 'true'
         run: |
           source /opt/rh/gcc-toolset-12/enable
+          cp -r presto-native-tests/presto_cpp/tests presto-native-execution/presto_cpp/main/functions/dynamic_registry
+          echo "add_subdirectory(tests)" >> presto-native-execution/presto_cpp/main/functions/dynamic_registry/CMakeLists.txt
           cd presto-native-execution
           cmake \
             -B _build/release \
@@ -126,6 +128,7 @@ jobs:
           path: |
             presto-native-execution/_build/release/presto_cpp/main/presto_server
             presto-native-execution/_build/release/velox/velox/functions/remote/server/velox_functions_remote_server_main
+            presto-native-execution/_build/release/presto_cpp/main/functions/dynamic_registry/tests/
 
   prestocpp-linux-presto-e2e-tests:
     needs: [changes, prestocpp-linux-build-for-test]
diff --git a/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java b/presto-native-execution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
@@ -120,6 +120,7 @@ public static class HiveQueryRunnerBuilder
         private Map<String, String> extraCoordinatorProperties = new HashMap<>();
         private Map<String, String> hiveProperties = new HashMap<>();
         private Map<String, String> tpcdsProperties = new HashMap<>();
+        private Optional<String> pluginDirectory = Optional.empty();
         private String security;
         private boolean addStorageFormatToPath;
         private boolean coordinatorSidecarEnabled;
@@ -171,6 +172,12 @@ public HiveQueryRunnerBuilder setRemoteFunctionServerUds(Optional<String> remote
             return this;
         }
 
+        public HiveQueryRunnerBuilder setPluginDirectory(Optional<String> pluginDirectory)
+        {
+            this.pluginDirectory = pluginDirectory;
+            return this;
+        }
+
         public HiveQueryRunnerBuilder setFailOnNestedLoopJoin(boolean failOnNestedLoopJoin)
         {
             this.failOnNestedLoopJoin = failOnNestedLoopJoin;
@@ -281,7 +288,7 @@ public QueryRunner build()
             Optional<BiFunction<Integer, URI, Process>> externalWorkerLauncher = Optional.empty();
             if (this.useExternalWorkerLauncher) {
                 externalWorkerLauncher = getExternalWorkerLauncher("hive", serverBinary, cacheMaxSize, remoteFunctionServerUds,
-                        failOnNestedLoopJoin, coordinatorSidecarEnabled, builtInWorkerFunctionsEnabled, enableRuntimeMetricsCollection, enableSsdCache, implicitCastCharNToVarchar);
+                        pluginDirectory, failOnNestedLoopJoin, coordinatorSidecarEnabled, builtInWorkerFunctionsEnabled, enableRuntimeMetricsCollection, enableSsdCache, implicitCastCharNToVarchar);
             }
             return HiveQueryRunner.createQueryRunner(
                     ImmutableList.of(),
@@ -370,7 +377,7 @@ public QueryRunner build()
             Optional<BiFunction<Integer, URI, Process>> externalWorkerLauncher = Optional.empty();
             if (this.useExternalWorkerLauncher) {
                 externalWorkerLauncher = getExternalWorkerLauncher("iceberg", serverBinary, cacheMaxSize, remoteFunctionServerUds,
-                        false, false, false, false, false, false);
+                        Optional.empty(), false, false, false, false, false, false);
             }
             return IcebergQueryRunner.builder()
                     .setExtraProperties(extraProperties)
@@ -464,6 +471,7 @@ public static Optional<BiFunction<Integer, URI, Process>> getExternalWorkerLaunc
             String prestoServerPath,
             int cacheMaxSize,
             Optional<String> remoteFunctionServerUds,
+            Optional<String> pluginDirectory,
             Boolean failOnNestedLoopJoin,
             boolean isCoordinatorSidecarEnabled,
             boolean isBuiltInWorkerFunctionsEnabled,
@@ -517,6 +525,10 @@ else if (isBuiltInWorkerFunctionsEnabled) {
                                     "remote-function-server.signature.files.directory.path=%s%n", configProperties, REMOTE_FUNCTION_CATALOG_NAME, remoteFunctionServerUds.get(), jsonSignaturesPath);
                         }
 
+                        if (pluginDirectory.isPresent()) {
+                            configProperties = format("%s%n" + "plugin.dir=%s%n", configProperties, pluginDirectory.get());
+                        }
+
                         if (failOnNestedLoopJoin) {
                             configProperties = format("%s%n" + "velox-plan-validator-fail-on-nested-loop-join=true%n", configProperties);
                         }
diff --git a/presto-native-tests/CMakeLists.txt b/presto-native-tests/CMakeLists.txt
@@ -0,0 +1,71 @@
+# 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.
+cmake_minimum_required(VERSION 3.10)
+
+# set the project name
+project(PrestoCpp)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED True)
+message("Appending CMAKE_CXX_FLAGS with ${SCRIPT_CXX_FLAGS}")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${SCRIPT_CXX_FLAGS}")
+
+# Known warnings that are benign can be disabled.
+set(DISABLED_WARNINGS
+    "-Wno-nullability-completeness -Wno-deprecated-declarations")
+
+# Important warnings that must be explicitly enabled.
+set(ENABLE_WARNINGS "-Wreorder")
+
+set(CMAKE_CXX_FLAGS
+    "${CMAKE_CXX_FLAGS} ${DISABLED_WARNINGS} ${ENABLE_WARNINGS}")
+
+set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
+
+if(DEFINED ENV{INSTALL_PREFIX})
+  message(STATUS "Dependency install directory set to: $ENV{INSTALL_PREFIX}")
+  list(APPEND CMAKE_PREFIX_PATH "$ENV{INSTALL_PREFIX}")
+endif()
+
+find_package(gflags REQUIRED)
+
+find_library(GLOG glog)
+
+find_package(fmt REQUIRED)
+
+if(NOT TARGET gflags::gflags)
+  # This is a bit convoluted, but we want to be able to use gflags::gflags as a
+  # target even when velox is built as a subproject which uses
+  # `find_package(gflags)` which does not create a globally imported target that
+  # we can ALIAS.
+  add_library(gflags_gflags INTERFACE)
+  target_link_libraries(gflags_gflags INTERFACE gflags)
+  add_library(gflags::gflags ALIAS gflags_gflags)
+endif()
+include_directories(.)
+
+include_directories(${CMAKE_BINARY_DIR})
+
+# Adding this down here prevents warnings in dependencies from stopping the
+# build.
+if("${TREAT_WARNINGS_AS_ERRORS}")
+  set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror")
+endif()
+
+if(DEFINED ENV{INSTALL_PREFIX})
+  # Allow installed package headers to be picked up before brew/system package
+  # headers. We set this after Velox since Velox handles INSTALL_PREFIX its own
+  # way.
+  include_directories(BEFORE "$ENV{INSTALL_PREFIX}/include")
+endif()
+
+add_subdirectory(presto_cpp)
diff --git a/presto-native-tests/Makefile b/presto-native-tests/Makefile
@@ -0,0 +1,67 @@
+# 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.
+.PHONY: build clean debug release submodules
+
+BUILD_BASE_DIR=_build
+BUILD_DIR=release
+BUILD_TYPE=Release
+TREAT_WARNINGS_AS_ERRORS ?= 1
+ENABLE_WALL ?= 1
+NUM_THREADS ?= $(shell getconf _NPROCESSORS_CONF 2>/dev/null || echo 1)
+CMAKE_PREFIX_PATH ?= "/usr/local"
+
+EXTRA_CMAKE_FLAGS ?= ""
+
+CMAKE_FLAGS := -DTREAT_WARNINGS_AS_ERRORS=${TREAT_WARNINGS_AS_ERRORS}
+CMAKE_FLAGS += -DENABLE_ALL_WARNINGS=${ENABLE_WALL}
+CMAKE_FLAGS += -DCMAKE_PREFIX_PATH=$(CMAKE_PREFIX_PATH)
+CMAKE_FLAGS += -DCMAKE_BUILD_TYPE=$(BUILD_TYPE)
+
+SHELL := /bin/bash
+
+# Use Ninja if available. If Ninja is used, pass through parallelism control flags.
+USE_NINJA ?= 1
+ifeq ($(USE_NINJA), 1)
+ifneq ($(shell which ninja), )
+CMAKE_FLAGS += -GNinja -DMAX_LINK_JOBS=$(MAX_LINK_JOBS) -DMAX_HIGH_MEM_JOBS=$(MAX_HIGH_MEM_JOBS)
+endif
+endif
+
+ifndef USE_CCACHE
+ifneq ($(shell which ccache), )
+CMAKE_FLAGS += -DCMAKE_CXX_COMPILER_LAUNCHER=ccache
+endif
+endif
+
+all: release			#: Build the release version
+
+clean:					#: Delete all build artifacts
+	rm -rf $(BUILD_BASE_DIR)
+
+cmake: submodules		#: Use CMake to create a Makefile build system
+	cmake -B "$(BUILD_BASE_DIR)/$(BUILD_DIR)" $(FORCE_COLOR) $(CMAKE_FLAGS) $(EXTRA_CMAKE_FLAGS)
+
+build:					#: Build the software based in BUILD_DIR and BUILD_TYPE variables
+	cmake --build $(BUILD_BASE_DIR)/$(BUILD_DIR) -j $(NUM_THREADS)
+
+debug:					#: Build with debugging symbols
+	$(MAKE) cmake BUILD_DIR=debug BUILD_TYPE=Debug
+	$(MAKE) build BUILD_DIR=debug
+
+release:				#: Build the release version
+	$(MAKE) cmake BUILD_DIR=release BUILD_TYPE=Release && \
+	$(MAKE) build BUILD_DIR=release
+
+help:					#: Show the help messages
+	@cat $(firstword $(MAKEFILE_LIST)) | \
+	awk '/^[-a-z]+:/' | \
+	awk -F: '{ printf("%-20s   %s\n", $$1, $$NF) }'
diff --git a/presto-native-tests/presto_cpp/CMakeLists.txt b/presto-native-tests/presto_cpp/CMakeLists.txt
@@ -0,0 +1,12 @@
+# 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.
+add_subdirectory(tests)
diff --git a/presto-native-tests/presto_cpp/tests/CMakeLists.txt b/presto-native-tests/presto_cpp/tests/CMakeLists.txt
@@ -0,0 +1,12 @@
+# 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.
+add_subdirectory(custom_functions)
diff --git a/presto-native-tests/presto_cpp/tests/custom_functions/CMakeLists.txt b/presto-native-tests/presto_cpp/tests/custom_functions/CMakeLists.txt
@@ -0,0 +1,39 @@
+# 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.
+set(PRESTO_AND_VELOX_INCLUDE_DIRS
+  ${CMAKE_SOURCE_DIR}/../presto-native-execution
+  ${CMAKE_SOURCE_DIR}/../presto-native-execution/velox
+)
+
+add_library(presto_cpp_custom_functions_test SHARED CustomFunctions.cpp)
+
+set(CMAKE_DYLIB_TEST_LINK_LIBRARIES fmt::fmt
+                                    xsimd)
+
+target_include_directories(presto_cpp_custom_functions_test
+  PRIVATE
+    ${PRESTO_AND_VELOX_INCLUDE_DIRS}
+    presto_dynamic_function_registrar
+)
+
+target_link_libraries(presto_cpp_custom_functions_test
+                      PRIVATE
+                      ${CMAKE_DYLIB_TEST_LINK_LIBRARIES})
+
+if(APPLE)
+  set(COMMON_LIBRARY_LINK_OPTIONS "-Wl,-undefined,dynamic_lookup")
+else()
+  set(COMMON_LIBRARY_LINK_OPTIONS "-Wl,--exclude-libs,ALL")
+endif()
+
+target_link_options(presto_cpp_custom_functions_test PRIVATE
+                    ${COMMON_LIBRARY_LINK_OPTIONS})
diff --git a/presto-native-tests/presto_cpp/tests/custom_functions/CustomFunctions.cpp b/presto-native-tests/presto_cpp/tests/custom_functions/CustomFunctions.cpp
@@ -0,0 +1,112 @@
+/*
+ * 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 <iostream>
+ #include "presto_cpp/main/functions/dynamic_registry/DynamicFunctionRegistrar.h"
+ #include "velox/expression/SimpleFunctionRegistry.h"
+ // This file defines a function that will be dynamically linked and
+ // registered. This function implements a custom function in the definition,
+ // and this library (.so) needs to provide a `void registerExtensions()`
+ // C function in the top-level namespace.
+ //
+ // (note the extern "C" directive to prevent the compiler from mangling the
+ // symbol name).
+
+namespace custom::functionRegistry {
+    template <typename T>
+    struct CustomAdd {
+      VELOX_DEFINE_FUNCTION_TYPES(T);
+      FOLLY_ALWAYS_INLINE bool call(
+        out_type<int64_t>& result,
+          const arg_type<int64_t>& x1,
+          const arg_type<int64_t>& x2) {
+        result = x1 + x2;
+        return true;
+      }
+    };
+
+    template <typename T>
+    struct SumArray {
+      VELOX_DEFINE_FUNCTION_TYPES(T);
+      FOLLY_ALWAYS_INLINE bool call(
+        out_type<int64_t>& result,
+          const arg_type<facebook::velox::Array<int64_t>>& arr) {
+        int64_t sum = 0;
+	for (auto val : arr) {
+	  if (val.has_value()) {
+	    sum += val.value();
+	  }
+	}
+        result = sum;
+        return true;
+      }
+    };
+
+    template <typename T>
+    struct MapSize {
+      VELOX_DEFINE_FUNCTION_TYPES(T);
+      FOLLY_ALWAYS_INLINE bool call(
+        out_type<int64_t>& result,
+          const arg_type<facebook::velox::Map<int64_t, int64_t>>& m) {
+        result = m.size();
+        return true;
+      }
+    };
+
+    template <typename T>
+    struct SumNestedArrayElements {
+      VELOX_DEFINE_FUNCTION_TYPES(T);
+      FOLLY_ALWAYS_INLINE bool call(
+        out_type<int64_t>& result,
+          const arg_type<facebook::velox::Array<facebook::velox::Array<int64_t>>>& arr) {
+        int64_t sum = 0;
+        for (auto innerOpt : arr) {
+          if (innerOpt.has_value()) {
+            for (auto val : innerOpt.value()) {
+              if (val.has_value()) {
+                sum += val.value();
+              }
+            }
+          }
+        }
+        result = sum;
+        return true;
+    }
+};
+} // namespace custom::functionRegistry
+
+extern "C" {
+    void registerExtensions() {
+      facebook::presto::registerPrestoFunction<
+          custom::functionRegistry::CustomAdd,
+          int64_t,
+          int64_t,
+          int64_t>("dynamic_custom_add");
+
+      facebook::presto::registerPrestoFunction<
+          custom::functionRegistry::SumArray,
+          int64_t,
+          facebook::velox::Array<int64_t>>("sum_array");
+
+      facebook::presto::registerPrestoFunction<
+          custom::functionRegistry::MapSize,
+          int64_t,
+          facebook::velox::Map<int64_t, int64_t>>("map_size");
+
+      facebook::presto::registerPrestoFunction<
+          custom::functionRegistry::SumNestedArrayElements,
+          int64_t,
+          facebook::velox::Array<facebook::velox::Array<int64_t>>>("sum_nested_array_elements");
+    }
+}
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/NativeTestsUtils.java
diff --git a/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestCustomFunctions.java b/presto-native-tests/src/test/java/com/facebook/presto/nativetests/TestCustomFunctions.java
