feat(native): Add endpoint for expression optimization in sidecar

presto-native-execution/pom.xml

@@ -326,6 +326,31 @@
                 </exclusion>
             </exclusions>
         </dependency>
+
+        <dependency>
+            <groupId>com.facebook.airlift</groupId>
+            <artifactId>bootstrap</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>com.facebook.airlift</groupId>
+            <artifactId>configuration</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>com.facebook.airlift</groupId>
+            <artifactId>http-client</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>com.facebook.airlift</groupId>
+            <artifactId>json</artifactId>
+            <scope>test</scope>
+        </dependency>
+
         <dependency>
             <groupId>com.facebook.airlift</groupId>
             <artifactId>log</artifactId>
@@ -372,6 +397,18 @@
             <version>${project.version}</version>
             <scope>test</scope>
         </dependency>
+
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-core</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-databind</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

presto-native-execution/presto_cpp/main/CMakeLists.txt

@@ -60,6 +60,7 @@ target_link_libraries(
   $<TARGET_OBJECTS:presto_protocol>
   presto_common
   presto_exception
+  presto_expression_optimizer
   presto_function_metadata
   presto_connectors
   presto_http

presto-native-execution/presto_cpp/main/PrestoServer.cpp

@@ -41,6 +41,7 @@
 #include "presto_cpp/main/operators/PartitionAndSerialize.h"
 #include "presto_cpp/main/operators/ShuffleExchangeSource.h"
 #include "presto_cpp/main/operators/ShuffleRead.h"
+#include "presto_cpp/main/types/ExpressionOptimizer.h"
 #include "presto_cpp/main/types/PrestoToVeloxQueryPlan.h"
 #include "presto_cpp/main/types/VeloxPlanConversion.h"
 #include "velox/common/base/Counters.h"
@@ -473,6 +474,8 @@ void PrestoServer::run() {
   createTaskManager();
 
   if (systemConfig->prestoNativeSidecar()) {
+    expression::registerInRewrite();
+
     registerSidecarEndpoints();
   }
 
@@ -1715,6 +1718,21 @@ void PrestoServer::registerSidecarEndpoints() {
               http::sendOkResponse(downstream, getFunctionsMetadata(catalog));
             });
       });
+  httpServer_->registerPost(
+      "/v1/expressions",
+      [&](proxygen::HTTPMessage* message,
+          const std::vector<std::unique_ptr<folly::IOBuf>>& body,
+          proxygen::ResponseHandler* downstream) {
+        const json::array_t inputRowExpressions =
+            json::parse(util::extractMessageBody(body));
+        expression::optimizeExpressions(
+            message->getHeaders(),
+            inputRowExpressions,
+            downstream,
+            driverExecutor_.get(),
+            pool_.get());
+      });
+
   httpServer_->registerPost(
       "/v1/velox/plan",
       [server = this](
@@ -1832,4 +1850,5 @@ void PrestoServer::registerTraceNodeFactories() {
         return nullptr;
       });
 }
+
 } // namespace facebook::presto

presto-native-execution/presto_cpp/main/types/CMakeLists.txt

@@ -33,6 +33,20 @@ set_property(TARGET presto_types PROPERTY JOB_POOL_LINK presto_link_job_pool)
 add_library(presto_velox_plan_conversion OBJECT VeloxPlanConversion.cpp)
 target_link_libraries(presto_velox_plan_conversion velox_type)
 
+add_library(velox_to_presto_expr VeloxToPrestoExpr.cpp)
+
+target_link_libraries(
+  velox_to_presto_expr
+  presto_exception
+  presto_type_converter
+  presto_types
+  presto_protocol
+)
+
+add_library(presto_expression_optimizer ExpressionOptimizer.cpp)
+
+target_link_libraries(presto_expression_optimizer presto_types velox_to_presto_expr)
+
 if(PRESTO_ENABLE_TESTING)
   add_subdirectory(tests)
 endif()

presto-native-execution/presto_cpp/main/types/ExpressionOptimizer.cpp

@@ -0,0 +1,161 @@
+/*
+ * 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 "presto_cpp/main/types/ExpressionOptimizer.h"
+#include "presto_cpp/main/common/Configs.h"
+#include "presto_cpp/main/common/Utils.h"
+#include "presto_cpp/main/http/HttpServer.h"
+#include "presto_cpp/main/types/PrestoToVeloxExpr.h"
+#include "presto_cpp/main/types/TypeParser.h"
+#include "presto_cpp/main/types/VeloxToPrestoExpr.h"
+#include "presto_cpp/presto_protocol/core/presto_protocol_core.h"
+#include "velox/expression/Expr.h"
+#include "velox/expression/ExprOptimizer.h"
+#include "velox/expression/ExprRewriteRegistry.h"
+
+namespace facebook::presto::expression {
+
+namespace {
+
+constexpr char const* kIn = "in";
+constexpr char const* kEvaluated = "EVALUATED";
+constexpr char const* kTimezoneHeader = "X-Presto-Time-Zone";
+constexpr char const* kOptimizerLevelHeader =
+    "X-Presto-Expression-Optimizer-Level";
+
+velox::core::TypedExprPtr rewriteInExpression(
+    const velox::core::TypedExprPtr& expr) {
+  if (!expr->isCallKind()) {
+    return nullptr;
+  }
+
+  const auto* callExpr = expr->asUnchecked<velox::core::CallTypedExpr>();
+  if (callExpr->name() != kIn) {
+    return nullptr;
+  }
+
+  const auto& valueExpr = callExpr->inputs().at(0);
+  if (!valueExpr->isConstantKind()) {
+    return nullptr;
+  }
+
+  const auto* constValueExpr =
+      valueExpr->asUnchecked<velox::core::ConstantTypedExpr>();
+  const auto inputs = callExpr->inputs();
+  const auto numInputs = inputs.size();
+  std::vector<velox::core::TypedExprPtr> optimizedInputs;
+  optimizedInputs.emplace_back(valueExpr);
+
+  for (auto i = 1; i < numInputs; i++) {
+    if (inputs[i]->isConstantKind()) {
+      const auto* constInput =
+          inputs[i]->asUnchecked<velox::core::ConstantTypedExpr>();
+      if (*constInput == *constValueExpr) {
+        return std::make_shared<velox::core::ConstantTypedExpr>(
+            velox::BOOLEAN(), true);
+      }
+    } else {
+      optimizedInputs.push_back(inputs[i]);
+    }
+  }
+
+  return std::make_shared<velox::core::CallTypedExpr>(
+      expr->type(), std::move(optimizedInputs), callExpr->name());
+}
+} // namespace
+
+void registerInRewrite() {
+  velox::expression::ExprRewriteRegistry::instance().registerRewrite(
+      rewriteInExpression);
+}
+
+void optimizeExpressions(
+    const proxygen::HTTPHeaders& httpHeaders,
+    const json::array_t inputRowExpressions,
+    proxygen::ResponseHandler* downstream,
+    folly::Executor* driverExecutor,
+    velox::memory::MemoryPool* pool) {
+  static const velox::expression::MakeFailExpr kMakeFailExpr =
+      [](const std::string& error,
+         const velox::TypePtr& type) -> velox::core::TypedExprPtr {
+    return std::make_shared<velox::core::CastTypedExpr>(
+        type,
+        std::vector<velox::core::TypedExprPtr>{
+            std::make_shared<velox::core::CallTypedExpr>(
+                velox::UNKNOWN(),
+                std::vector<velox::core::TypedExprPtr>{
+                    std::make_shared<velox::core::ConstantTypedExpr>(
+                        velox::VARCHAR(), error)},
+                fmt::format(
+                    "{}fail",
+                    SystemConfig::instance()->prestoDefaultNamespacePrefix()))},
+        false);
+  };
+
+  const auto& timezone = httpHeaders.getSingleOrEmpty(kTimezoneHeader);
+  const auto& optimizerLevel =
+      httpHeaders.getSingleOrEmpty(kOptimizerLevelHeader);
+  std::unordered_map<std::string, std::string> config(
+      {{velox::core::QueryConfig::kSessionTimezone, timezone},
+       {velox::core::QueryConfig::kAdjustTimestampToTimezone, "true"}});
+  auto queryConfig = velox::core::QueryConfig{std::move(config)};
+  auto queryCtx =
+      velox::core::QueryCtx::create(driverExecutor, std::move(queryConfig));
+
+  TypeParser typeParser;
+  const VeloxExprConverter veloxExprConverter(pool, &typeParser);
+  expression::VeloxToPrestoExprConverter veloxToPrestoExprConverter(pool);
+  json j;
+  json result = json::array();
+
+  for (const RowExpressionPtr& inputRowExpr : inputRowExpressions) {
+    protocol::to_json(j, inputRowExpr);
+    VLOG(1) << "Input RowExpression: " << j.dump();
+    auto expr = veloxExprConverter.toVeloxExpr(inputRowExpr);
+    VLOG(1) << "Converted TypedExpr: " << expr->toString();
+    auto optimized =
+        velox::expression::optimize(expr, queryCtx.get(), pool, kMakeFailExpr);
+    VLOG(1) << "Optimized TypedExpr: " << optimized->toString();
+
+    try {
+      if (optimizerLevel == kEvaluated) {
+        const auto evalResult = velox::exec::tryEvaluateConstantExpression(
+            optimized, pool, queryCtx, /* suppressEvaluationFailures */ false);
+        optimized =
+            std::make_shared<velox::core::ConstantTypedExpr>(evalResult);
+      }
+    } catch (const velox::VeloxUserError& e) {
+      LOG(ERROR) << fmt::format(
+          "Expression evaluation failed with error: {}.", e.message());
+      http::sendErrorResponse(downstream, e.what());
+      return;
+    }
+
+    auto resultExpression =
+        veloxToPrestoExprConverter.getRowExpression(optimized);
+    if (resultExpression == nullptr) {
+      // Log Velox to Presto expression conversion error and suppress it to
+      // return the unoptimized RowExpression.
+      resultExpression = inputRowExpr;
+    }
+    protocol::to_json(j, resultExpression);
+    VLOG(1) << "Optimized RowExpression: " << j.dump();
+
+    result.push_back(resultExpression);
+  }
+
+  http::sendOkResponse(downstream, result);
+}
+
+} // namespace facebook::presto::expression

presto-native-execution/presto_cpp/main/types/ExpressionOptimizer.h

@@ -0,0 +1,53 @@
+/*
+ * 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 <folly/io/IOBuf.h>
+#include <proxygen/httpserver/ResponseHandler.h>
+#include <proxygen/lib/http/HTTPHeaders.h>
+#include "presto_cpp/external/json/nlohmann/json.hpp"
+#include "velox/common/memory/MemoryPool.h"
+#include "velox/core/Expressions.h"
+
+namespace facebook::presto::expression {
+
+/// Registers rewrite for IN expression when the `value` being searched for is
+/// constant. The rewrite returns `true` if `value` is in the IN-list along with
+/// other non-constant expressions; IN expressions with a constant IN-list will
+/// be constant folded by the ExprOptimizer in Velox. Constant expressions from
+/// IN-list that do not match with `value` are also removed.
+/// This rewrite should be in Presto only and not in Velox, since IN is not a
+/// FunctionCallToSpecialForm in Velox and it is evaluated with a Filter instead
+/// for cases where the IN-list is not constant.
+void registerInRewrite();
+
+/// Optimizes RowExpressions received in a http request and returns a http
+/// response containing the result of expression optimization.
+/// @param httpHeaders Headers from the http request, contains the timezone
+/// from Presto coordinator and the expression optimizer level.
+/// @param inputRowExpressions List of RowExpressions to be optimized.
+/// @param downstream Returns the result of expression optimization as a http
+/// response. If expression optimizer level is `EVALUATED` and the evaluation of
+/// any expression from the input fails, the http response contains the error
+/// message encountered during evaluation with a 500 response code. Otherwise,
+/// the http response contains the list of optimized RowExpressions, serialized
+/// as an array of JSON objects, with 200 response code.
+/// @param driverExecutor Driver CPU executor.
+/// @param pool Memory pool.
+void optimizeExpressions(
+    const proxygen::HTTPHeaders& httpHeaders,
+    const nlohmann::json::array_t inputRowExpressions,
+    proxygen::ResponseHandler* downstream,
+    folly::Executor* driverExecutor,
+    velox::memory::MemoryPool* pool);
+} // namespace facebook::presto::expression