Fix array_sort Presto function when evaluating a subset of rows (face…

…bookincubator#7800)

Summary:
array_sort used to generate invalid dictionary vectors when evaluating only a subset of rows:

```
VeloxRuntimeError: dictionaryIndices->size() >= length * sizeof(vector_size_t) (36 vs. 40) Malformed dictionary, index array is shorter than DictionaryVector
```

Pull Request resolved: facebookincubator#7800

Reviewed By: xiaoxmeng

Differential Revision: D51693822

Pulled By: mbasmanova

fbshipit-source-id: 43657fdc01a9eed7d17508758711988fd895dfcf

velox/functions/prestosql/ArraySort.cpp

@@ -340,7 +340,10 @@ class ArraySortLambdaFunction : public exec::VectorFunction {
         context,
         throwOnNestedNull_);
     auto sortedElements = BaseVector::wrapInDictionary(
-        nullptr, indices, newNumElements, flatArray->elements());
+        nullptr,
+        indices,
+        indices->size() / sizeof(vector_size_t),
+        flatArray->elements());
 
     // Set nulls for rows not present in 'rows'.
     BufferPtr newNulls = addNullsForUnselectedRows(flatArray, rows);
@@ -349,7 +352,7 @@ class ArraySortLambdaFunction : public exec::VectorFunction {
         flatArray->pool(),
         flatArray->type(),
         std::move(newNulls),
-        flatArray->size(),
+        rows.end(),
         flatArray->offsets(),
         flatArray->sizes(),
         sortedElements);

velox/functions/prestosql/tests/ArraySortTest.cpp

@@ -510,13 +510,23 @@ TEST_F(ArraySortTest, lambda) {
     SCOPED_TRACE(lambdaExpr);
     auto result = evaluate(fmt::format("{}(c0, {})", name, lambdaExpr), data);
     assertEqualVectors(sortedAsc, result);
+
+    SelectivityVector firstRow(1);
+    result =
+        evaluate(fmt::format("{}(c0, {})", name, lambdaExpr), data, firstRow);
+    assertEqualVectors(sortedAsc->slice(0, 1), result);
   };
 
   auto testDesc = [&](const std::string& name, const std::string& lambdaExpr) {
     SCOPED_TRACE(name);
     SCOPED_TRACE(lambdaExpr);
     auto result = evaluate(fmt::format("{}(c0, {})", name, lambdaExpr), data);
     assertEqualVectors(sortedDesc, result);
+
+    SelectivityVector firstRow(1);
+    result =
+        evaluate(fmt::format("{}(c0, {})", name, lambdaExpr), data, firstRow);
+    assertEqualVectors(sortedDesc->slice(0, 1), result);
   };
 
   // Different ways to sort by length ascending.