[native] Add row expression optimizer

prestodb · Jan 8, 2025 · ad0a9e2 · ad0a9e2
1 parent 7913aa4
commit ad0a9e2
Show file tree

Hide file tree

Showing 16 changed files with 1,623 additions and 24 deletions.
diff --git a/presto-native-execution/presto_cpp/main/CMakeLists.txt b/presto-native-execution/presto_cpp/main/CMakeLists.txt
@@ -51,6 +51,7 @@ target_link_libraries(
   presto_http
   presto_operators
   presto_velox_conversion
+  presto_expression_optimizer
   velox_aggregates
   velox_caching
   velox_common_base

diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.cpp b/presto-native-execution/presto_cpp/main/PrestoServer.cpp
@@ -1546,6 +1546,26 @@ void PrestoServer::registerSidecarEndpoints() {
          proxygen::ResponseHandler* downstream) {
         http::sendOkResponse(downstream, getFunctionsMetadata());
       });
+  rowExpressionOptimizer_ =
+      std::make_unique<expression::RowExpressionOptimizer>();
+  httpServer_->registerPost(
+      "/v1/expressions",
+      [&](proxygen::HTTPMessage* message,
+          const std::vector<std::unique_ptr<folly::IOBuf>>& body,
+          proxygen::ResponseHandler* downstream) {
+        json::array_t inputRowExpressions =
+            json::parse(util::extractMessageBody(body));
+        auto result =
+            rowExpressionOptimizer_->optimize(message, inputRowExpressions);
+        if (result.second) {
+          VELOX_USER_CHECK(
+              result.first.is_array(),
+              "Output json should be an array of row expressions");
+          http::sendOkResponse(downstream, result.first);
+        } else {
+          http::sendErrorResponse(downstream, result.first);
+        }
+      });
   httpServer_->registerPost(
       "/v1/velox/plan",
       [server = this](

diff --git a/presto-native-execution/presto_cpp/main/PrestoServer.h b/presto-native-execution/presto_cpp/main/PrestoServer.h
@@ -25,6 +25,7 @@
 #include "presto_cpp/main/PeriodicHeartbeatManager.h"
 #include "presto_cpp/main/PrestoExchangeSource.h"
 #include "presto_cpp/main/PrestoServerOperations.h"
+#include "presto_cpp/main/types/RowExpressionOptimizer.h"
 #include "presto_cpp/main/types/VeloxPlanValidator.h"
 #include "velox/common/caching/AsyncDataCache.h"
 #include "velox/common/memory/MemoryAllocator.h"
@@ -291,6 +292,7 @@ class PrestoServer {
   std::string address_;
   std::string nodeLocation_;
   folly::SSLContextPtr sslContext_;
+  std::unique_ptr<expression::RowExpressionOptimizer> rowExpressionOptimizer_;
 };
 
 } // namespace facebook::presto
diff --git a/presto-native-execution/presto_cpp/main/types/CMakeLists.txt b/presto-native-execution/presto_cpp/main/types/CMakeLists.txt
@@ -34,6 +34,12 @@ add_library(presto_velox_conversion OBJECT VeloxPlanConversion.cpp)
 
 target_link_libraries(presto_velox_conversion velox_type)
 
+add_library(presto_expression_optimizer RowExpressionConverter.cpp
+                                        RowExpressionOptimizer.cpp)
+
+target_link_libraries(presto_expression_optimizer presto_type_converter
+                      presto_types presto_protocol)
+
 if(PRESTO_ENABLE_TESTING)
   add_subdirectory(tests)
 endif()
diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxExpr.cpp b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxExpr.cpp
@@ -33,32 +33,10 @@ std::string toJsonString(const T& value) {
 }
 
 std::string mapScalarFunction(const std::string& name) {
-  static const std::unordered_map<std::string, std::string> kFunctionNames = {
-      // Operator overrides: com.facebook.presto.common.function.OperatorType
-      {"presto.default.$operator$add", "presto.default.plus"},
-      {"presto.default.$operator$between", "presto.default.between"},
-      {"presto.default.$operator$divide", "presto.default.divide"},
-      {"presto.default.$operator$equal", "presto.default.eq"},
-      {"presto.default.$operator$greater_than", "presto.default.gt"},
-      {"presto.default.$operator$greater_than_or_equal", "presto.default.gte"},
-      {"presto.default.$operator$is_distinct_from",
-       "presto.default.distinct_from"},
-      {"presto.default.$operator$less_than", "presto.default.lt"},
-      {"presto.default.$operator$less_than_or_equal", "presto.default.lte"},
-      {"presto.default.$operator$modulus", "presto.default.mod"},
-      {"presto.default.$operator$multiply", "presto.default.multiply"},
-      {"presto.default.$operator$negation", "presto.default.negate"},
-      {"presto.default.$operator$not_equal", "presto.default.neq"},
-      {"presto.default.$operator$subtract", "presto.default.minus"},
-      {"presto.default.$operator$subscript", "presto.default.subscript"},
-      // Special form function overrides.
-      {"presto.default.in", "in"},
-  };
-
   std::string lowerCaseName = boost::to_lower_copy(name);
 
-  auto it = kFunctionNames.find(lowerCaseName);
-  if (it != kFunctionNames.end()) {
+  auto it = kPrestoOperatorMap.find(lowerCaseName);
+  if (it != kPrestoOperatorMap.end()) {
     return it->second;
   }
 
@@ -102,6 +80,15 @@ std::string getFunctionName(const protocol::SqlFunctionId& functionId) {
 
 } // namespace
 
+const std::unordered_map<std::string, std::string> veloxToPrestoOperatorMap() {
+  std::unordered_map<std::string, std::string> veloxToPrestoOperatorMap;
+  for (const auto& entry : kPrestoOperatorMap) {
+    veloxToPrestoOperatorMap[entry.second] = entry.first;
+  }
+  veloxToPrestoOperatorMap.insert({"cast", "presto.default.$operator$cast"});
+  return veloxToPrestoOperatorMap;
+}
+
 velox::variant VeloxExprConverter::getConstantValue(
     const velox::TypePtr& type,
     const protocol::Block& block) const {

diff --git a/presto-native-execution/presto_cpp/main/types/PrestoToVeloxExpr.h b/presto-native-execution/presto_cpp/main/types/PrestoToVeloxExpr.h
@@ -20,6 +20,30 @@
 
 namespace facebook::presto {
 
+static const std::unordered_map<std::string, std::string> kPrestoOperatorMap = {
+    // Operator overrides: com.facebook.presto.common.function.OperatorType
+    {"presto.default.$operator$add", "presto.default.plus"},
+    {"presto.default.$operator$between", "presto.default.between"},
+    {"presto.default.$operator$divide", "presto.default.divide"},
+    {"presto.default.$operator$equal", "presto.default.eq"},
+    {"presto.default.$operator$greater_than", "presto.default.gt"},
+    {"presto.default.$operator$greater_than_or_equal", "presto.default.gte"},
+    {"presto.default.$operator$is_distinct_from",
+     "presto.default.distinct_from"},
+    {"presto.default.$operator$less_than", "presto.default.lt"},
+    {"presto.default.$operator$less_than_or_equal", "presto.default.lte"},
+    {"presto.default.$operator$modulus", "presto.default.mod"},
+    {"presto.default.$operator$multiply", "presto.default.multiply"},
+    {"presto.default.$operator$negation", "presto.default.negate"},
+    {"presto.default.$operator$not_equal", "presto.default.neq"},
+    {"presto.default.$operator$subtract", "presto.default.minus"},
+    {"presto.default.$operator$subscript", "presto.default.subscript"},
+    // Special form function overrides.
+    {"presto.default.in", "in"},
+};
+
+const std::unordered_map<std::string, std::string> veloxToPrestoOperatorMap();
+
 class VeloxExprConverter {
  public:
   VeloxExprConverter(velox::memory::MemoryPool* pool, TypeParser* typeParser)