Fix flakiness in partitionedOutputWithLargeInput (facebookincubator#8761

)

Summary:

This change ensures that the test checks against the output metrics
of the partitionedOutput operator instead of the RemoteExchange
operator where input vectors can be merged based on their memory
footprint (facebookincubator#7404)
which can vary as the input is randomly generated.

Reviewed By: kevinwilfong

Differential Revision: D53824959

velox/exec/tests/MultiFragmentTest.cpp

@@ -574,10 +574,9 @@ TEST_F(MultiFragmentTest, partitionedOutput) {
 TEST_F(MultiFragmentTest, partitionedOutputWithLargeInput) {
   // Verify that partitionedOutput operator is able to split a single input
   // vector if it hits memory or row limits.
-  // We create a large vector that hits the row limit (70% - 120% of 10,000)
-  // which would hit a task level memory limit of 1MB unless its split up.
+  // We create a large vector that hits the row limit (70% - 120% of 10,000).
   // This test exercises splitting up the input both from the edges and the
-  // middle as it ends up splitting it in ~ 3 splits.
+  // middle as it ends up splitting it into at least 3.
   setupSources(1, 30'000);
   const int64_t kRootMemoryLimit = 1 << 20; // 1MB
   // Single Partition
@@ -595,10 +594,10 @@ TEST_F(MultiFragmentTest, partitionedOutputWithLargeInput) {
 
     auto task =
         assertQuery(op, {leafTaskId}, "SELECT c0, c1, c2, c3, c4 FROM tmp");
-    auto taskStats = toPlanStats(task->taskStats());
-    ASSERT_GT(taskStats.at("0").inputVectors, 2);
     ASSERT_TRUE(waitForTaskCompletion(leafTask.get()))
         << leafTask->taskId() << "state: " << leafTask->state();
+    auto taskStats = toPlanStats(leafTask->taskStats());
+    ASSERT_GT(taskStats.at("1").outputVectors, 2);
   }
 
   // Multiple partitions but round-robin.
@@ -636,11 +635,10 @@ TEST_F(MultiFragmentTest, partitionedOutputWithLargeInput) {
 
     auto task = assertQuery(
         op, intermediateTaskIds, "SELECT c0, c1, c2, c3, c4 FROM tmp");
-    auto taskStats = toPlanStats(task->taskStats());
-    ASSERT_GT(taskStats.at("0").inputVectors, 2);
-
     ASSERT_TRUE(waitForTaskCompletion(leafTask.get()))
         << "state: " << leafTask->state();
+    auto taskStats = toPlanStats(leafTask->taskStats());
+    ASSERT_GT(taskStats.at("1").outputVectors, 2);
   }
 }