add test

facebookincubator · Dec 11, 2024 · fc74211 · fc74211
1 parent 475cfa4
commit fc74211
Showing 1 changed file with 17 additions and 12 deletions.
diff --git a/velox/exec/tests/SpillTest.cpp b/velox/exec/tests/SpillTest.cpp
@@ -988,6 +988,7 @@ TEST_P(SpillTest, multipleSortKeys) {
       vectors.push_back(makeArrayVector<int32_t>({{1, 2, 3}, {3, 4, 5}}));
     }
     state.appendToPartition(partitionIndex, makeRowVector({vectors}));
+    state.appendToPartition(partitionIndex, makeRowVector({vectors}));
     state.finishFile(partitionIndex);
     EXPECT_TRUE(
         state.testingNonEmptySpilledPartitionSet().contains(partitionIndex));
@@ -997,18 +998,22 @@ TEST_P(SpillTest, multipleSortKeys) {
         1 << 20, testData.mergePrefixComparatorEnabled, pool(), &spillStats_);
     ASSERT_TRUE(merge != nullptr);
 
-    for (auto i = 0; i < timeValues.size(); ++i) {
-      auto* stream = merge->next();
-      ASSERT_NE(stream, nullptr);
-      ASSERT_EQ(
-          timeValues[i],
-          stream->decoded(0).valueAt<Timestamp>(stream->currentIndex()));
-
-      ASSERT_EQ(
-          intValues[i],
-          stream->decoded(1).valueAt<int64_t>(stream->currentIndex()));
-      stream->pop();
-    }
+    const auto testMerge = [] {
+      for (auto i = 0; i < timeValues.size(); ++i) {
+        auto* stream = merge->next();
+        ASSERT_NE(stream, nullptr);
+        ASSERT_EQ(
+            timeValues[i],
+            stream->decoded(0).valueAt<Timestamp>(stream->currentIndex()));
+
+        ASSERT_EQ(
+            intValues[i],
+            stream->decoded(1).valueAt<int64_t>(stream->currentIndex()));
+        stream->pop();
+      }
+    };
+    testMerge();
+    testMerge();
     ASSERT_EQ(nullptr, merge->next());
   }
 }