Add support for multiarg lambdas in array_sort (facebookincubator#13387)

Summary:
Pull Request resolved: facebookincubator#13387

This change aims to add support for array_sort with lambda that can take multiple arguments. Currently Velox tries to convert any multiple argument lambda to a single argument one using SimpleMatcher. This works for a large majority of the cases and is pretty fast, however there are still significant number of cases where we still require support for true multiple arguments in the lambda.

This preliminary PR supports that use case ; We do this by evaluating the lambda for all the argument vectors in an iteration. In every iteration we keep one of the arguments constant for an array and compare it to the rest of its array. This gives us the position of the constant element wrt to the other elements in the array. We can do this N times to sort the array.

To illustrate, say lambda takes two arguments, say x, y.
We will map y to the elements of the array A and assign x as A[i] for iteration i.
Thus in the ith iteration we will compare x[i] against every y (i.e every element in A).
At the end of each iteration, based on the number of -1, 0, 1's we will know the position of the element at index i.

NOTE:
1. This implementation requires that the lambda always return -1, 0, or 1. This is enforced by Presto Java too.
2. This implementation has a further restriction that for the lambda f, f(a,a) => 0 when a = a.
3. This implementation aims to rewrite the lambda to the single lambda arg case failing which it uses the multi arg lambda.
4. This implementation also only supports 1 lambda entry and not multiple lambda entries.

Differential Revision: D74852355

Author: Krishna Pai

velox/core/QueryConfig.h

@@ -598,6 +598,13 @@ class QueryConfig {
   static constexpr const char* kFieldNamesInJsonCastEnabled =
       "field_names_in_json_cast_enabled";
 
+  /// Represents the maximum size of the array that can be sorted by the
+  /// array_sort() function, if the lambda takes more than one argument
+  /// and cannot be reduced by the SimpleMatcher. If the value is 0, then
+  /// this is disabled.
+  static constexpr const char* kArraySortMaxIterations =
+      "array_sort_max_iterations";
+
   bool selectiveNimbleReaderEnabled() const {
     return get<bool>(kSelectiveNimbleReaderEnabled, false);
   }
@@ -1089,6 +1096,10 @@ class QueryConfig {
     return get<bool>(kFieldNamesInJsonCastEnabled, false);
   }
 
+  uint32_t getArraySortMaxIterations() const {
+    return get<uint32_t>(kArraySortMaxIterations, 1000);
+  }
+
   template <typename T>
   T get(const std::string& key, const T& defaultValue) const {
     return config_->get<T>(key, defaultValue);

velox/expression/LambdaExpr.cpp

@@ -76,6 +76,10 @@ class ExprCallable : public Callable {
     lambdaCtx.swapErrors(elementErrors);
   }
 
+  const RowTypePtr getFunctionSignatures() override {
+    return signature_;
+  }
+
  private:
   void resetSharedExprs() {
     for (auto& expr : sharedExprsToReset_) {

velox/expression/tests/ExprCompilerTest.cpp

@@ -309,6 +309,11 @@ TEST_F(ExprCompilerTest, rewrites) {
       exprSet->expr(0)->toString(),
       "plus(1:BIGINT, array_sum_propagate_element_null(transform(c0, (x) -> multiply(x, 2:BIGINT))))");
 
+  exprSet = compile(makeTypedExpr(
+      "ARRAY_SORT(c0, (a, b) -> IF(POW(a, 3) < POW(b, 3), -1, IF(a = b, 0, 1)))",
+      ROW({"c0", "c1"}, {ARRAY(BIGINT()), BIGINT()})));
+  ASSERT_EQ(exprSet->size(), 1);
+
   exprSet = compile(makeTypedExpr(
       "reduce(c0, 1, (s, x) -> (s + 2) - x, s -> s)",
       ROW({"c0"}, {ARRAY(BIGINT())})));