wills-feng · wills-feng · Jun 12, 2024
diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -1059,10 +1059,11 @@ void PrestoServer::registerRemoteFunctions() {
     } else {
       VELOX_FAIL(
           "To register remote functions using a json file path you need to "
-          "specify the remote server location using '{}', '{}' or '{}'.",
+          "specify the remote server location using '{}', '{}' or '{}' or {}.",
           SystemConfig::kRemoteFunctionServerThriftAddress,
           SystemConfig::kRemoteFunctionServerThriftPort,
-          SystemConfig::kRemoteFunctionServerThriftUdsPath);
+          SystemConfig::kRemoteFunctionServerThriftUdsPath,
+          SystemConfig::kRemoteFunctionServerRestURL);
     }
   }
 #endif

diff --git a/presto-native-execution/presto_cpp/main/RemoteFunctionRegisterer.cpp b/presto-native-execution/presto_cpp/main/RemoteFunctionRegisterer.cpp
@@ -18,6 +18,7 @@
 #include "velox/common/base/Exceptions.h"
 #include "velox/common/base/Fs.h"
 #include "velox/functions/remote/client/Remote.h"
+#include <boost/variant.hpp>
 
 namespace facebook::presto {
 namespace {
@@ -50,7 +51,7 @@ PageFormat fromSerdeString(const std::string_view& serdeName) {
 // registered.
 size_t processFile(
     const fs::path& filePath,
-    const folly::SocketAddress& location,
+    const boost::variant<folly::SocketAddress, proxygen::URL>& location,
     const std::string_view& prefix,
     const std::string_view& serde) {
   std::ifstream stream{filePath};
@@ -88,7 +89,7 @@ size_t processFile(
 
 size_t registerRemoteFunctions(
     const std::string& inputPath,
-    const folly::SocketAddress& location,
+    const boost::variant<folly::SocketAddress, proxygen::URL>& location,
     const std::string_view& prefix,
     const std::string_view& serde) {
   size_t signaturesCount = 0;

diff --git a/presto-native-execution/presto_cpp/main/RemoteFunctionRegisterer.h b/presto-native-execution/presto_cpp/main/RemoteFunctionRegisterer.h
@@ -14,6 +14,8 @@
 #pragma once
 
 #include <folly/SocketAddress.h>
+#include <boost/variant.hpp>
+#include <proxygen/lib/utils/URL.h>
 
 namespace facebook::presto {
 
@@ -36,7 +38,7 @@ namespace facebook::presto {
 /// Returns the number of signatures registered.
 size_t registerRemoteFunctions(
     const std::string& inputPath,
-    const folly::SocketAddress& location,
+    const boost::variant<folly::SocketAddress, proxygen::URL>& location,
     const std::string_view& prefix = "",
     const std::string_view& serde = "presto_page");
 

diff --git a/presto-native-execution/presto_cpp/main/common/Configs.cpp b/presto-native-execution/presto_cpp/main/common/Configs.cpp
@@ -20,6 +20,7 @@
 #include <boost/lexical_cast.hpp>
 #include <boost/uuid/uuid_generators.hpp>
 #include <boost/uuid/uuid_io.hpp>
+
 #if __has_include("filesystem")
 #include <filesystem>
 namespace fs = std::filesystem;
@@ -287,35 +288,42 @@ folly::Optional<std::string> SystemConfig::discoveryUri() const {
   return optionalProperty(kDiscoveryUri);
 }
 
-folly::Optional<folly::SocketAddress>
+folly::Optional<boost::variant<folly::SocketAddress, proxygen::URL>>
 SystemConfig::remoteFunctionServerLocation() const {
-  // First check if there is a UDS path registered. If there's one, use it.
-  auto remoteServerUdsPath =
-      optionalProperty(kRemoteFunctionServerThriftUdsPath);
-  if (remoteServerUdsPath.hasValue()) {
-    return folly::SocketAddress::makeFromPath(remoteServerUdsPath.value());
-  }
 
-  // Otherwise, check for address and port parameters.
-  auto remoteServerAddress =
-      optionalProperty(kRemoteFunctionServerThriftAddress);
-  auto remoteServerPort =
-      optionalProperty<uint16_t>(kRemoteFunctionServerThriftPort);
-
-  if (remoteServerPort.hasValue()) {
-    // Fallback to localhost if address is not specified.
-    return remoteServerAddress.hasValue()
-        ? folly::
-              SocketAddress{remoteServerAddress.value(), remoteServerPort.value()}
-        : folly::SocketAddress{"::1", remoteServerPort.value()};
-  } else if (remoteServerAddress.hasValue()) {
-    VELOX_FAIL(
-        "Remote function server port not provided using '{}'.",
-        kRemoteFunctionServerThriftPort);
-  }
+  boost::variant<folly::SocketAddress, proxygen::URL> ret;
+  // First check if there is a REST URL
+  auto remoteServerRestURL =
+      optionalProperty(kRemoteFunctionServerRestURL);
+  if (remoteServerRestURL.hasValue()) {
+    ret=proxygen::URL(remoteServerRestURL.value());
+  }else {
+    // check if there is a UDS path registered. If there's one, use it.
+    auto remoteServerUdsPath =
+        optionalProperty(kRemoteFunctionServerThriftUdsPath);
+    if (remoteServerUdsPath.hasValue()) {
+      ret = folly::SocketAddress::makeFromPath(remoteServerUdsPath.value());
+    }
 
-  // No remote function server configured.
-  return folly::none;
+    // Otherwise, check for address and port parameters.
+    auto remoteServerAddress =
+        optionalProperty(kRemoteFunctionServerThriftAddress);
+    auto remoteServerPort =
+        optionalProperty<uint16_t>(kRemoteFunctionServerThriftPort);
+
+    if (remoteServerPort.hasValue()) {
+      // Fallback to localhost if address is not specified.
+      ret = remoteServerAddress.hasValue()
+          ? folly::
+                SocketAddress{remoteServerAddress.value(), remoteServerPort.value()}
+          : folly::SocketAddress{"::1", remoteServerPort.value()};
+    } else if (remoteServerAddress.hasValue()) {
+      VELOX_FAIL(
+          "Remote function server port not provided using '{}'.",
+          kRemoteFunctionServerThriftPort);
+    }
+  }
+  return ret;
 }
 
 folly::Optional<std::string>

diff --git a/presto-native-execution/presto_cpp/main/common/Configs.h b/presto-native-execution/presto_cpp/main/common/Configs.h
@@ -17,6 +17,8 @@
 #include <memory>
 #include <string>
 #include <unordered_map>
+#include <boost/variant.hpp>
+#include <proxygen/lib/utils/URL.h>
 #include "velox/core/Config.h"
 
 namespace facebook::presto {
@@ -489,6 +491,10 @@ class SystemConfig : public ConfigBase {
   static constexpr std::string_view kRemoteFunctionServerThriftUdsPath{
       "remote-function-server.thrift.uds-path"};
 
+  /// HTTP URL used by the remote function rest server.
+  static constexpr std::string_view kRemoteFunctionServerRestURL{
+      "remote-function-server.rest.url"};
+
   /// Path where json files containing signatures for remote functions can be
   /// found.
   static constexpr std::string_view
@@ -570,7 +576,7 @@ class SystemConfig : public ConfigBase {
 
   folly::Optional<std::string> discoveryUri() const;
 
-  folly::Optional<folly::SocketAddress> remoteFunctionServerLocation() const;
+  folly::Optional<boost::variant<folly::SocketAddress, proxygen::URL>> remoteFunctionServerLocation() const;
 
   folly::Optional<std::string> remoteFunctionServerSignatureFilesDirectoryPath()
       const;

diff --git a/...ion/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeRemoteFunctions.java b/...ion/src/test/java/com/facebook/presto/nativeworker/AbstractTestNativeRemoteFunctions.java
@@ -24,6 +24,7 @@
 import static com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils.REMOTE_FUNCTION_JSON_SIGNATURES;
 import static com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils.setupJsonFunctionNamespaceManager;
 import static com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils.startRemoteFunctionServer;
+import static com.facebook.presto.nativeworker.PrestoNativeQueryRunnerUtils.startRemoteFunctionRestServer;
 
 @Test(groups = "remote-function")
 public abstract class AbstractTestNativeRemoteFunctions
@@ -51,7 +52,13 @@ protected void preInitQueryRunners()
         String propertyName = "REMOTE_FUNCTION_SERVER";
         remoteFunctionServerBinaryPath = getProperty(propertyName).orElseThrow(() -> new NullPointerException("Remote function server path missing. Add -DREMOTE_FUNCTION_SERVER=<path/to/binary>" +
                 " to your JVM arguments, or create an environment variable REMOTE_FUNCTION_SERVER with value <path/to/binary>."));
-        remoteFunctionServerUds = startRemoteFunctionServer(remoteFunctionServerBinaryPath);
+
+        String remoteFunctionServerType = "REMOTE_FUNCTION_SERVER_TYPE";
+        if(getProperty(remoteFunctionServerType).isPresent() && getProperty(remoteFunctionServerType).get().equals("REST")){
+            remoteFunctionServerUds = startRemoteFunctionRestServer(remoteFunctionServerBinaryPath);
+        }else {
+            remoteFunctionServerUds = startRemoteFunctionServer(remoteFunctionServerBinaryPath);
+        }
     }
 
     protected void postInitQueryRunners()

diff --git a/...xecution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java b/...xecution/src/test/java/com/facebook/presto/nativeworker/PrestoNativeQueryRunnerUtils.java
@@ -26,6 +26,7 @@
 
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.net.ServerSocket;
 import java.net.URI;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -35,6 +36,7 @@
 import java.util.UUID;
 import java.util.function.BiFunction;
 
+
 import static com.facebook.presto.hive.HiveTestUtils.getProperty;
 import static com.facebook.presto.iceberg.IcebergQueryRunner.createIcebergQueryRunner;
 import static com.facebook.presto.nativeworker.NativeQueryRunnerUtils.getNativeWorkerHiveProperties;
@@ -316,7 +318,7 @@ public static QueryRunner createNativeQueryRunner(boolean useThrift, String stor
     public static String startRemoteFunctionServer(String remoteFunctionServerBinaryPath)
     {
         try {
-            Path tempDirectoryPath = Files.createTempDirectory("RemoteFunctionServer");
+            Path tempDirectoryPath = Files.createTempDirectory(Paths.get("/tmp"),"RemoteFunctionServer");
             Path remoteFunctionServerUdsPath = tempDirectoryPath.resolve(REMOTE_FUNCTION_UDS);
             log.info("Temp directory for Remote Function Server: %s", tempDirectoryPath.toString());
 
@@ -333,6 +335,27 @@ public static String startRemoteFunctionServer(String remoteFunctionServerBinary
         }
     }
 
+    public static String startRemoteFunctionRestServer(String remoteFunctionServerBinaryPath)
+    {
+        try {
+            Path tempDirectoryPath = Files.createTempDirectory(Paths.get("/tmp"),"RemoteFunctionRestServer");
+            int port=new ServerSocket(0).getLocalPort();
+            log.info("Temp directory for Remote Function Rest Server: %s", tempDirectoryPath.toString());
+
+            Process p = new ProcessBuilder(Paths.get(remoteFunctionServerBinaryPath).toAbsolutePath().toString(), "--service_port", Integer.toString(port), "--function_prefix", REMOTE_FUNCTION_CATALOG_NAME + ".schema.")
+                    .directory(tempDirectoryPath.toFile())
+                    .redirectErrorStream(true)
+                    .redirectOutput(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("thrift_server.out").toFile()))
+                    .redirectError(ProcessBuilder.Redirect.to(tempDirectoryPath.resolve("thrift_server.err").toFile()))
+                    .start();
+
+            return String.format("http://127.0.0.1:%d/", port);
+        }
+        catch (IOException e) {
+            throw new UncheckedIOException(e);
+        }
+    }
+
     public static NativeQueryRunnerParameters getNativeQueryRunnerParameters()
     {
         Path prestoServerPath = Paths.get(getProperty("PRESTO_SERVER")
@@ -375,11 +398,20 @@ public static Optional<BiFunction<Integer, URI, Process>> getExternalWorkerLaunc
 
                         if (remoteFunctionServerUds.isPresent()) {
                             String jsonSignaturesPath = Resources.getResource(REMOTE_FUNCTION_JSON_SIGNATURES).getFile();
-                            configProperties = format("%s%n" +
-                                    "remote-function-server.catalog-name=%s%n" +
-                                    "remote-function-server.thrift.uds-path=%s%n" +
-                                    "remote-function-server.serde=presto_page%n" +
-                                    "remote-function-server.signature.files.directory.path=%s%n", configProperties, REMOTE_FUNCTION_CATALOG_NAME, remoteFunctionServerUds.get(), jsonSignaturesPath);
+                            if (remoteFunctionServerUds.get().startsWith("http")) {
+                                configProperties = format("%s%n" +
+                                        "remote-function-server.catalog-name=%s%n" +
+                                        "remote-function-server.rest.url=%s%n" +
+                                        "remote-function-server.serde=presto_page%n" +
+                                        "remote-function-server.signature.files.directory.path=%s%n", configProperties, REMOTE_FUNCTION_CATALOG_NAME, remoteFunctionServerUds.get(), jsonSignaturesPath);
+
+                            }else{
+                                configProperties = format("%s%n" +
+                                        "remote-function-server.catalog-name=%s%n" +
+                                        "remote-function-server.thrift.uds-path=%s%n" +
+                                        "remote-function-server.serde=presto_page%n" +
+                                        "remote-function-server.signature.files.directory.path=%s%n", configProperties, REMOTE_FUNCTION_CATALOG_NAME, remoteFunctionServerUds.get(), jsonSignaturesPath);
+                            }
                         }
                         Files.write(tempDirectoryPath.resolve("config.properties"), configProperties.getBytes());
                         Files.write(tempDirectoryPath.resolve("node.properties"),