From fc74211f430c563d9316742e9022f602acde73b7 Mon Sep 17 00:00:00 2001
From: Chengcheng Jin <chengcheng.jin@intel.com>
Date: Wed, 11 Dec 2024 14:27:22 +0000
Subject: [PATCH] add test

---
 velox/exec/tests/SpillTest.cpp | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/velox/exec/tests/SpillTest.cpp b/velox/exec/tests/SpillTest.cpp
index db67c5c3c3745..4a8d75a1d6141 100644
--- 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());
   }
 }