fix(native): Enable Velox to Presto lambda expression conversion

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

@@ -91,7 +91,7 @@ protocol::RowExpressionOptimizationResult optimizeExpression(
   }
 
   result.optimizedExpression =
-      veloxToPrestoConverter.getRowExpression(optimized, input);
+      veloxToPrestoConverter.getRowExpression(optimized);
   return result;
 }
 

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

@@ -228,6 +228,30 @@ SpecialFormExpressionPtr VeloxToPrestoExprConverter::getDereferenceExpression(
   return result;
 }
 
+LambdaDefinitionExpressionPtr VeloxToPrestoExprConverter::getLambdaExpression(
+    const velox::core::LambdaTypedExpr* lambdaExpr) const {
+  static constexpr char const* kLambda = "lambda";
+
+  json result;
+  result["@type"] = kLambda;
+  const auto& signature = lambdaExpr->signature();
+  std::vector<protocol::TypeSignature> argumentTypes;
+  argumentTypes.reserve(signature->children().size());
+  for (const auto& type : signature->children()) {
+    argumentTypes.emplace_back(getTypeSignature(type));
+  }
+  result["argumentTypes"] = argumentTypes;
+
+  std::vector<std::string> arguments;
+  arguments.reserve(signature->names().size());
+  for (const auto& name : signature->names()) {
+    arguments.emplace_back(name);
+  }
+  result["arguments"] = arguments;
+  result["body"] = getRowExpression(lambdaExpr->body());
+  return result;
+}
+
 CallExpressionPtr VeloxToPrestoExprConverter::getCallExpression(
     const velox::core::CallTypedExpr* expr) const {
   static constexpr char const* kCall = "call";
@@ -271,8 +295,7 @@ CallExpressionPtr VeloxToPrestoExprConverter::getCallExpression(
 }
 
 RowExpressionPtr VeloxToPrestoExprConverter::getRowExpression(
-    const velox::core::TypedExprPtr& expr,
-    const RowExpressionPtr& inputRowExpr) const {
+    const velox::core::TypedExprPtr& expr) const {
   switch (expr->kind()) {
     case velox::core::ExprKind::kConstant: {
       const auto* constantExpr =
@@ -312,24 +335,23 @@ RowExpressionPtr VeloxToPrestoExprConverter::getRowExpression(
       }
       return getCallExpression(callTypedExpr);
     }
+    case velox::core::ExprKind::kLambda: {
+      const auto* lambdaExpr =
+          expr->asUnchecked<velox::core::LambdaTypedExpr>();
+      return getLambdaExpression(lambdaExpr);
+    }
     // Presto does not have a RowExpression type for kConcat and kInput Velox
-    // expressions. Presto's expression optimizer does not support optimization
-    // of lambda expressions.
+    // expressions. Presto to Velox expression conversion should not generate
+    // Velox expressions of these types.
     case velox::core::ExprKind::kConcat:
       [[fallthrough]];
     case velox::core::ExprKind::kInput:
       [[fallthrough]];
-    case velox::core::ExprKind::kLambda:
-      [[fallthrough]];
-    default: {
-      // Log Velox to Presto expression conversion error and return the
-      // unoptimized input RowExpression.
-      LOG(ERROR) << fmt::format(
+    default:
+      VELOX_FAIL(
           "Unable to convert Velox expression: {} of kind: {} to Presto RowExpression.",
           expr->toString(),
           velox::core::ExprKindName::toName(expr->kind()));
-      return inputRowExpr;
-    }
   }
 }
 

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

@@ -24,6 +24,8 @@ using ConstantExpressionPtr =
     std::shared_ptr<facebook::presto::protocol::ConstantExpression>;
 using CallExpressionPtr =
     std::shared_ptr<facebook::presto::protocol::CallExpression>;
+using LambdaDefinitionExpressionPtr =
+    std::shared_ptr<facebook::presto::protocol::LambdaDefinitionExpression>;
 using SpecialFormExpressionPtr =
     std::shared_ptr<facebook::presto::protocol::SpecialFormExpression>;
 
@@ -59,11 +61,9 @@ class VeloxToPrestoExprConverter {
       : pool_(pool) {}
 
   /// Converts a Velox expression `expr` to a Presto protocol RowExpression.
-  /// The input Presto RowExpression is returned as the `defaultResult` in case
-  /// the Velox to Presto expression conversion fails.
+  /// Throws an exception if the expression conversion fails.
   RowExpressionPtr getRowExpression(
-      const velox::core::TypedExprPtr& expr,
-      const RowExpressionPtr& defaultResult = nullptr) const;
+      const velox::core::TypedExprPtr& expr) const;
 
  private:
   /// This function is used to serialize a constant velox vector to
@@ -99,6 +99,11 @@ class VeloxToPrestoExprConverter {
   SpecialFormExpressionPtr getDereferenceExpression(
       const velox::core::DereferenceTypedExpr* dereferenceExpr) const;
 
+  /// Helper function to construct a Presto
+  /// `protocol::LambdaDefinitionExpression` from a Velox lambda expression.
+  LambdaDefinitionExpressionPtr getLambdaExpression(
+      const velox::core::LambdaTypedExpr* lambdaExpr) const;
+
   /// Helper function to construct a Presto `protocol::CallExpression` from a
   /// Velox call expression.
   CallExpressionPtr getCallExpression(

presto-native-sidecar-plugin/src/test/java/com/facebook/presto/sidecar/expressions/TestNativeExpressionInterpreter.java

@@ -98,6 +98,13 @@ public void tearDown()
         closeAllRuntimeException(queryRunner);
     }
 
+    @Test
+    public void testLambda()
+    {
+        optimize("array_sort(ARRAY['apple', 'banana', 'pear'], x -> IF(x = 'banana', NULL, length(x)))");
+        optimize("array_sort(ARRAY[ROW('a', 3), ROW('b', 1), ROW('c', 2)], x -> x[2])");
+    }
+
     ///  Velox permits Bigint to Varchar cast but Presto does not.
     @Override
     @Test