From 37ba7cb4ca58543da21c40f22192365323d90cf3 Mon Sep 17 00:00:00 2001
From: Kevin Wilfong <kevinwilfong@meta.com>
Date: Wed, 31 Jan 2024 14:55:01 -0800
Subject: [PATCH] Break up AllTypes.nonSortedSpillFunctions test (#8621)

Summary:
Pull Request resolved: https://github.com/facebookincubator/velox/pull/8621

Instances of this test can run for upwards of 10 minutes (depending on the host and available resources).

It looks like we can easily break this test up further using test parameters.  Making maxSpillRunRows a parameter reduces the
run time of individual tests by about a third.

Reviewed By: amitkdutta

Differential Revision: D53283194

fbshipit-source-id: 73368ecedd8450932dceeb369e439afdafeea2ce
---
 velox/exec/tests/SpillerTest.cpp | 102 +++++++++++++++++--------------
 1 file changed, 56 insertions(+), 46 deletions(-)

diff --git a/velox/exec/tests/SpillerTest.cpp b/velox/exec/tests/SpillerTest.cpp
index c03e3c37ebc2..c1bad6e6bf06 100644
--- a/velox/exec/tests/SpillerTest.cpp
+++ b/velox/exec/tests/SpillerTest.cpp
@@ -1013,63 +1013,73 @@ class SpillerTest : public exec::test::RowContainerTestBase {
   std::unique_ptr<Spiller> spiller_;
 };
 
+struct AllTypesTestParam {
+  TestParam param;
+  uint64_t maxSpillRunRows;
+};
+
 class AllTypes : public SpillerTest,
-                 public testing::WithParamInterface<TestParam> {
+                 public testing::WithParamInterface<AllTypesTestParam> {
  public:
-  AllTypes() : SpillerTest(GetParam()) {}
+  AllTypes()
+      : SpillerTest(GetParam().param),
+        maxSpillRunRows_(GetParam().maxSpillRunRows) {}
+
+  static std::vector<AllTypesTestParam> getTestParams() {
+    auto testParams = TestParamsBuilder().getTestParams();
+
+    std::vector<AllTypesTestParam> allTypesTestParams;
+    for (const auto& testParam : testParams) {
+      for (const auto& maxSpillRunRows :
+           std::vector<uint64_t>{0, 101, 1'000'000}) {
+        allTypesTestParams.push_back({testParam, maxSpillRunRows});
+      }
+    }
 
-  static std::vector<TestParam> getTestParams() {
-    return TestParamsBuilder().getTestParams();
+    return allTypesTestParams;
   }
+
+ protected:
+  uint64_t maxSpillRunRows_;
 };
 
 TEST_P(AllTypes, nonSortedSpillFunctions) {
-  struct {
-    uint64_t maxSpillRunRows;
-
-    std::string debugString() const {
-      return fmt::format("maxSpillRunRows {}", maxSpillRunRows);
+  if (type_ == Spiller::Type::kOrderByInput ||
+      type_ == Spiller::Type::kOrderByOutput ||
+      type_ == Spiller::Type::kAggregateInput ||
+      type_ == Spiller::Type::kAggregateOutput) {
+    setupSpillData(rowType_, numKeys_, 5'000, 1, nullptr, {});
+    sortSpillData();
+    setupSpiller(100'000, 0, false, maxSpillRunRows_);
+    {
+      RowVectorPtr dummyVector;
+      VELOX_ASSERT_THROW(
+          spiller_->spill(0, dummyVector), "Unexpected spiller type");
     }
-  } testSettings[] = {{0}, {101}, {1'000'000}};
 
-  for (const auto& testData : testSettings) {
-    if (type_ == Spiller::Type::kOrderByInput ||
-        type_ == Spiller::Type::kOrderByOutput ||
-        type_ == Spiller::Type::kAggregateInput ||
-        type_ == Spiller::Type::kAggregateOutput) {
-      setupSpillData(rowType_, numKeys_, 5'000, 1, nullptr, {});
-      sortSpillData();
-      setupSpiller(100'000, 0, false, testData.maxSpillRunRows);
-      {
-        RowVectorPtr dummyVector;
-        VELOX_ASSERT_THROW(
-            spiller_->spill(0, dummyVector), "Unexpected spiller type");
-      }
-
-      if (type_ == Spiller::Type::kOrderByOutput) {
-        RowContainerIterator rowIter;
-        std::vector<char*> rows(5'000);
-        int numListedRows{0};
-        numListedRows = rowContainer_->listRows(&rowIter, 5000, rows.data());
-        ASSERT_EQ(numListedRows, 5000);
-        spiller_->spill(rows);
-      } else {
-        spiller_->spill();
-      }
-
-      ASSERT_FALSE(spiller_->finalized());
-      SpillPartitionSet spillPartitionSet;
-      spiller_->finishSpill(spillPartitionSet);
-      ASSERT_TRUE(spiller_->finalized());
-      ASSERT_EQ(spillPartitionSet.size(), 1);
-      verifySortedSpillData(spillPartitionSet.begin()->second.get());
-      continue;
+    if (type_ == Spiller::Type::kOrderByOutput) {
+      RowContainerIterator rowIter;
+      std::vector<char*> rows(5'000);
+      int numListedRows{0};
+      numListedRows = rowContainer_->listRows(&rowIter, 5000, rows.data());
+      ASSERT_EQ(numListedRows, 5000);
+      spiller_->spill(rows);
+    } else {
+      spiller_->spill();
     }
-    testNonSortedSpill(2, 5'000, 3, 1, testData.maxSpillRunRows);
-    testNonSortedSpill(2, 5'000, 3, 1'000'000'000, testData.maxSpillRunRows);
-    // Empty case.
-    testNonSortedSpill(1, 5'000, 0, 1, testData.maxSpillRunRows);
+
+    ASSERT_FALSE(spiller_->finalized());
+    SpillPartitionSet spillPartitionSet;
+    spiller_->finishSpill(spillPartitionSet);
+    ASSERT_TRUE(spiller_->finalized());
+    ASSERT_EQ(spillPartitionSet.size(), 1);
+    verifySortedSpillData(spillPartitionSet.begin()->second.get());
+    return;
   }
+  testNonSortedSpill(2, 5'000, 3, 1, maxSpillRunRows_);
+  testNonSortedSpill(2, 5'000, 3, 1'000'000'000, maxSpillRunRows_);
+  // Empty case.
+  testNonSortedSpill(1, 5'000, 0, 1, maxSpillRunRows_);
 }
 
 class NoHashJoin : public SpillerTest,