feat(functions): Add support for REST based remote functions

Co-authored-by: Wills Feng <[email protected]>

CMakeLists.txt

@@ -138,7 +138,7 @@ option(VELOX_ENABLE_ABFS "Build Abfs Connector" OFF)
 option(VELOX_ENABLE_HDFS "Build Hdfs Connector" OFF)
 option(VELOX_ENABLE_PARQUET "Enable Parquet support" OFF)
 option(VELOX_ENABLE_ARROW "Enable Arrow support" OFF)
-option(VELOX_ENABLE_REMOTE_FUNCTIONS "Enable remote function support" OFF)
+option(VELOX_ENABLE_REMOTE_FUNCTIONS "Enable remote function support" ON)
 option(VELOX_ENABLE_CCACHE "Use ccache if installed." ON)
 
 option(VELOX_BUILD_TEST_UTILS "Builds Velox test utilities" OFF)
@@ -534,6 +534,8 @@ if(VELOX_ENABLE_REMOTE_FUNCTIONS)
   find_package(fizz CONFIG REQUIRED)
   find_package(wangle CONFIG REQUIRED)
   find_package(FBThrift CONFIG REQUIRED)
+  velox_set_source(curl)
+  velox_resolve_dependency(curl)
 endif()
 
 if(VELOX_ENABLE_GCS)

velox/functions/remote/client/CMakeLists.txt

@@ -16,11 +16,20 @@ velox_add_library(velox_functions_remote_thrift_client ThriftClient.cpp)
 velox_link_libraries(velox_functions_remote_thrift_client
                      PUBLIC remote_function_thrift FBThrift::thriftcpp2)
 
+velox_add_library(velox_functions_remote_rest_client RestClient.cpp)
+velox_link_libraries(velox_functions_remote_rest_client Folly::folly
+                     CURL::libcurl)
+
 velox_add_library(velox_functions_remote Remote.cpp)
 velox_link_libraries(
   velox_functions_remote
   PUBLIC velox_expression
+         velox_memory
+         velox_exec
+         velox_vector
+         velox_presto_serializer
          velox_functions_remote_thrift_client
+         velox_functions_remote_rest_client
          velox_functions_remote_get_serde
          velox_type_fbhive
          Folly::folly)

velox/functions/remote/client/Remote.cpp

@@ -16,34 +16,70 @@
 
 #include "velox/functions/remote/client/Remote.h"
 
+#include <fmt/format.h>
 #include <folly/io/async/EventBase.h>
+#include <sstream>
+#include <string>
+
+#include "velox/common/memory/ByteStream.h"
 #include "velox/expression/Expr.h"
 #include "velox/expression/VectorFunction.h"
+#include "velox/functions/remote/client/RestClient.h"
 #include "velox/functions/remote/client/ThriftClient.h"
 #include "velox/functions/remote/if/GetSerde.h"
 #include "velox/functions/remote/if/gen-cpp2/RemoteFunctionServiceAsyncClient.h"
+#include "velox/serializers/PrestoSerializer.h"
 #include "velox/type/fbhive/HiveTypeSerializer.h"
 #include "velox/vector/VectorStream.h"
 
+using namespace folly;
 namespace facebook::velox::functions {
 namespace {
 
 std::string serializeType(const TypePtr& type) {
-  // Use hive type serializer.
   return type::fbhive::HiveTypeSerializer::serialize(type);
 }
 
+std::string extractFunctionName(const std::string& input) {
+  size_t lastDot = input.find_last_of('.');
+  if (lastDot != std::string::npos) {
+    return input.substr(lastDot + 1);
+  }
+  return input;
+}
+
+std::string urlEncode(const std::string& value) {
+  std::ostringstream escaped;
+  escaped.fill('0');
+  escaped << std::hex;
+  for (char c : value) {
+    if (isalnum(static_cast<unsigned char>(c)) || c == '-' || c == '_' ||
+        c == '.' || c == '~') {
+      escaped << c;
+    } else {
+      escaped << '%' << std::setw(2) << int(static_cast<unsigned char>(c));
+    }
+  }
+  return escaped.str();
+}
+
 class RemoteFunction : public exec::VectorFunction {
  public:
   RemoteFunction(
       const std::string& functionName,
       const std::vector<exec::VectorFunctionArg>& inputArgs,
-      const RemoteVectorFunctionMetadata& metadata)
+      const RemoteVectorFunctionMetadata& metadata,
+      std::unique_ptr<HttpClient> httpClient = nullptr)
       : functionName_(functionName),
-        location_(metadata.location),
-        thriftClient_(getThriftClient(location_, &eventBase_)),
-        serdeFormat_(metadata.serdeFormat),
-        serde_(getSerde(serdeFormat_)) {
+        restClient_(httpClient ? std::move(httpClient) : getRestClient()),
+        metadata_(metadata) {
+    if (metadata.location.type() == typeid(SocketAddress)) {
+      location_ = boost::get<SocketAddress>(metadata.location);
+      thriftClient_ = getThriftClient(location_, &eventBase_);
+    } else if (metadata.location.type() == typeid(std::string)) {
+      url_ = boost::get<std::string>(metadata.location);
+    }
+
     std::vector<TypePtr> types;
     types.reserve(inputArgs.size());
     serializedInputTypes_.reserve(inputArgs.size());
@@ -62,7 +98,11 @@ class RemoteFunction : public exec::VectorFunction {
       exec::EvalCtx& context,
       VectorPtr& result) const override {
     try {
-      applyRemote(rows, args, outputType, context, result);
+      if ((metadata_.location.type() == typeid(SocketAddress))) {
+        applyRemote(rows, args, outputType, context, result);
+      } else if (metadata_.location.type() == typeid(std::string)) {
+        applyRestRemote(rows, args, outputType, context, result);
+      }
     } catch (const VeloxRuntimeError&) {
       throw;
     } catch (const std::exception&) {
@@ -71,6 +111,48 @@ class RemoteFunction : public exec::VectorFunction {
   }
 
  private:
+  void applyRestRemote(
+      const SelectivityVector& rows,
+      std::vector<VectorPtr>& args,
+      const TypePtr& outputType,
+      exec::EvalCtx& context,
+      VectorPtr& result) const {
+    try {
+      serializer::presto::PrestoVectorSerde serde;
+      auto remoteRowVector = std::make_shared<RowVector>(
+          context.pool(),
+          remoteInputType_,
+          BufferPtr{},
+          rows.end(),
+          std::move(args));
+
+      std::unique_ptr<IOBuf> requestBody =
+          std::make_unique<IOBuf>(rowVectorToIOBuf(
+              remoteRowVector, rows.end(), *context.pool(), &serde));
+
+      const std::string fullUrl = fmt::format(
+          "{}/v1/functions/{}/{}/{}/{}",
+          url_,
+          metadata_.schema.value_or("default"),
+          extractFunctionName(functionName_),
+          urlEncode(metadata_.functionId.value_or("default_function_id")),
+          metadata_.version.value_or("1"));
+
+      std::unique_ptr<IOBuf> responseBody =
+          restClient_->invokeFunction(fullUrl, std::move(requestBody));
+
+      auto outputRowVector = IOBufToRowVector(
+          *responseBody, ROW({outputType}), *context.pool(), &serde);
+
+      result = outputRowVector->childAt(0);
+    } catch (const std::exception& e) {
+      VELOX_FAIL(
+          "Error while executing remote function '{}': {}",
+          functionName_,
+          e.what());
+    }
+  }
+
   void applyRemote(
       const SelectivityVector& rows,
       std::vector<VectorPtr>& args,
@@ -97,11 +179,14 @@ class RemoteFunction : public exec::VectorFunction {
 
     auto requestInputs = request.inputs_ref();
     requestInputs->rowCount_ref() = remoteRowVector->size();
-    requestInputs->pageFormat_ref() = serdeFormat_;
+    requestInputs->pageFormat_ref() = metadata_.serdeFormat;
 
     // TODO: serialize only active rows.
     requestInputs->payload_ref() = rowVectorToIOBuf(
-        remoteRowVector, rows.end(), *context.pool(), serde_.get());
+        remoteRowVector,
+        rows.end(),
+        *context.pool(),
+        getSerde(metadata_.serdeFormat).get());
 
     try {
       thriftClient_->sync_invokeFunction(remoteResponse, request);
@@ -117,12 +202,15 @@ class RemoteFunction : public exec::VectorFunction {
         remoteResponse.get_result().get_payload(),
         ROW({outputType}),
         *context.pool(),
-        serde_.get());
+        getSerde(metadata_.serdeFormat).get());
     result = outputRowVector->childAt(0);
 
     if (auto errorPayload = remoteResponse.get_result().errorPayload()) {
       auto errorsRowVector = IOBufToRowVector(
-          *errorPayload, ROW({VARCHAR()}), *context.pool(), serde_.get());
+          *errorPayload,
+          ROW({VARCHAR()}),
+          *context.pool(),
+          getSerde(metadata_.serdeFormat).get());
       auto errorsVector =
           errorsRowVector->childAt(0)->asFlatVector<StringView>();
       VELOX_CHECK(errorsVector, "Should be convertible to flat vector");
@@ -142,16 +230,14 @@ class RemoteFunction : public exec::VectorFunction {
   }
 
   const std::string functionName_;
-  folly::SocketAddress location_;
-
-  folly::EventBase eventBase_;
+  EventBase eventBase_;
   std::unique_ptr<RemoteFunctionClient> thriftClient_;
-  remote::PageFormat serdeFormat_;
-  std::unique_ptr<VectorSerde> serde_;
-
-  // Structures we construct once to cache:
+  std::unique_ptr<HttpClient> restClient_;
+  SocketAddress location_;
+  std::string url_;
   RowTypePtr remoteInputType_;
   std::vector<std::string> serializedInputTypes_;
+  const RemoteVectorFunctionMetadata metadata_;
 };
 
 std::shared_ptr<exec::VectorFunction> createRemoteFunction(
@@ -169,7 +255,7 @@ void registerRemoteFunction(
     std::vector<exec::FunctionSignaturePtr> signatures,
     const RemoteVectorFunctionMetadata& metadata,
     bool overwrite) {
-  exec::registerStatefulVectorFunction(
+  registerStatefulVectorFunction(
       name,
       signatures,
       std::bind(

velox/functions/remote/client/Remote.h

@@ -16,20 +16,37 @@
 
 #pragma once
 
+#include <boost/variant.hpp>
 #include <folly/SocketAddress.h>
 #include "velox/expression/VectorFunction.h"
 #include "velox/functions/remote/if/gen-cpp2/RemoteFunction_types.h"
 
 namespace facebook::velox::functions {
 
 struct RemoteVectorFunctionMetadata : public exec::VectorFunctionMetadata {
-  /// Network address of the servr to communicate with. Note that this can hold
-  /// a network location (ip/port pair) or a unix domain socket path (see
+  /// URL of the HTTP/REST server for remote function.
+  /// Or Network address of the server to communicate with. Note that this can
+  /// hold a network location (ip/port pair) or a unix domain socket path (see
   /// SocketAddress::makeFromPath()).
-  folly::SocketAddress location;
+  boost::variant<folly::SocketAddress, std::string> location;
 
-  /// The serialization format to be used
+  /// The serialization format to be used when sending data to the remote.
   remote::PageFormat serdeFormat{remote::PageFormat::PRESTO_PAGE};
+
+  /// Optional schema defining the structure of the data or input/output types
+  /// involved in the remote function. This may include details such as column
+  /// names and data types.
+  std::optional<std::string> schema;
+
+  /// Optional identifier for the specific remote function to be invoked.
+  /// This can be useful when the same server hosts multiple functions,
+  /// and the client needs to specify which function to call.
+  std::optional<std::string> functionId;
+
+  /// Optional version information to be used when calling the remote function.
+  /// This can help in ensuring compatibility with a particular version of the
+  /// function if multiple versions are available on the server.
+  std::optional<std::string> version;
 };
 
 /// Registers a new remote function. It will use the meatadata defined in