Extend tests for transformXxx Presto lambda functions (facebookincuba…

…tor#9289) Summary: Pull Request resolved: facebookincubator#9289 Extend tests for transform, transform_keys and transform_values lambda functions to cover scenario described in facebookincubator#9288 Reviewed By: bikramSingh91 Differential Revision: D55474277 fbshipit-source-id: 95bf9322430f99924ac0ee41e47aafda9e6959cd
PHILO-HE · Mar 29, 2024 · 90fc393 · 90fc393
1 parent cb58cba
commit 90fc393
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 20 deletions.
diff --git a/velox/functions/prestosql/tests/TransformKeysTest.cpp b/velox/functions/prestosql/tests/TransformKeysTest.cpp
@@ -69,19 +69,33 @@ TEST_F(TransformKeysTest, evaluateSubsetOfRows) {
           nullEvery(13)),
   });
 
-  SelectivityVector inputRows(size, false);
-  inputRows.setValidRange(0, size / 3, true);
-  inputRows.updateBounds();
+  // Test using 2 selectivity vectors. One of size 100 with 33 first rows
+  // selected. Another of size 33 with all rows selected. Both should produce
+  // the same result.
 
-  auto result = evaluate<MapVector>(
-      "transform_keys(c0, (k, v) -> k + 5)", input, inputRows);
   auto expectedResult = makeMapVector<int64_t, int32_t>(
       size / 3,
       [](auto row) { return row % 5; },
       [](auto row) { return row % 7 + 5; },
       [](auto row) { return row % 11; },
       nullEvery(13));
-  assertEqualVectors(expectedResult, result, inputRows);
+
+  {
+    SelectivityVector inputRows(size, false);
+    inputRows.setValidRange(0, size / 3, true);
+    inputRows.updateBounds();
+
+    auto result = evaluate<MapVector>(
+        "transform_keys(c0, (k, v) -> k + 5)", input, inputRows);
+    assertEqualVectors(expectedResult, result, inputRows);
+  }
+
+  {
+    SelectivityVector inputRows(size / 3);
+    auto result = evaluate<MapVector>(
+        "transform_keys(c0, (k, v) -> k + 5)", input, inputRows);
+    assertEqualVectors(expectedResult, result, inputRows);
+  }
 }
 
 TEST_F(TransformKeysTest, duplicateKeys) {

diff --git a/velox/functions/prestosql/tests/TransformTest.cpp b/velox/functions/prestosql/tests/TransformTest.cpp
@@ -45,19 +45,32 @@ TEST_F(TransformTest, evaluateSubsetOfRows) {
   auto input = makeRowVector(
       {makeArrayVector<int64_t>(size, modN(5), modN(7), nullEvery(11))});
 
-  SelectivityVector inputRows(size, false);
-  inputRows.setValidRange(0, size / 3, true);
-  inputRows.updateBounds();
-
-  auto result =
-      evaluate<ArrayVector>("transform(c0, x -> x + 5)", input, inputRows);
+  // Test using 2 selectivity vectors. One of size 100 with 33 first rows
+  // selected. Another of size 33 with all rows selected. Both should produce
+  // the same result.
 
   auto expectedResult = makeArrayVector<int64_t>(
       size / 3,
       modN(5),
       [](vector_size_t row) { return row % 7 + 5; },
       nullEvery(11));
-  assertEqualVectors(expectedResult, result, inputRows);
+  {
+    SelectivityVector inputRows(size, false);
+    inputRows.setValidRange(0, size / 3, true);
+    inputRows.updateBounds();
+
+    auto result =
+        evaluate<ArrayVector>("transform(c0, x -> x + 5)", input, inputRows);
+    assertEqualVectors(expectedResult, result, inputRows);
+  }
+
+  {
+    SelectivityVector inputRows(size / 3);
+
+    auto result =
+        evaluate<ArrayVector>("transform(c0, x -> x + 5)", input, inputRows);
+    assertEqualVectors(expectedResult, result, inputRows);
+  }
 }
 
 TEST_F(TransformTest, differentResultType) {

diff --git a/velox/functions/prestosql/tests/TransformValuesTest.cpp b/velox/functions/prestosql/tests/TransformValuesTest.cpp
@@ -69,19 +69,33 @@ TEST_F(TransformValuesTest, evaluateSubsetOfRows) {
           nullEvery(13)),
   });
 
-  SelectivityVector inputRows(size, false);
-  inputRows.setValidRange(0, size / 3, true);
-  inputRows.updateBounds();
-
-  auto result = evaluate<MapVector>(
-      "transform_values(c0, (k, v) -> v + 5)", input, inputRows);
+  // Test using 2 selectivity vectors. One of size 100 with 33 first rows
+  // selected. Another of size 33 with all rows selected. Both should produce
+  // the same result.
   auto expectedResult = makeMapVector<int32_t, int64_t>(
       size / 3,
       [](auto row) { return row % 5; },
       [](auto row) { return row % 7; },
       [](auto row) { return row % 11 + 5; },
       nullEvery(13));
-  assertEqualVectors(expectedResult, result, inputRows);
+
+  {
+    SelectivityVector inputRows(size, false);
+    inputRows.setValidRange(0, size / 3, true);
+    inputRows.updateBounds();
+
+    auto result = evaluate<MapVector>(
+        "transform_values(c0, (k, v) -> v + 5)", input, inputRows);
+    assertEqualVectors(expectedResult, result, inputRows);
+  }
+
+  {
+    SelectivityVector inputRows(size / 3);
+
+    auto result = evaluate<MapVector>(
+        "transform_values(c0, (k, v) -> v + 5)", input, inputRows);
+    assertEqualVectors(expectedResult, result, inputRows);
+  }
 }
 
 TEST_F(TransformValuesTest, differentResultType) {