From c1f20cdbc89deac980f646a967485ce7f51fab8e Mon Sep 17 00:00:00 2001
From: duanmeng <macduan@tencent.com>
Date: Fri, 28 Jun 2024 23:49:10 +0800
Subject: [PATCH] Add max spill level in RowNumber fuzzer

---
 velox/exec/fuzzer/RowNumberFuzzer.cpp | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/velox/exec/fuzzer/RowNumberFuzzer.cpp b/velox/exec/fuzzer/RowNumberFuzzer.cpp
index aba81e918025..c117b2dcde4b 100644
--- a/velox/exec/fuzzer/RowNumberFuzzer.cpp
+++ b/velox/exec/fuzzer/RowNumberFuzzer.cpp
@@ -50,6 +50,11 @@ DEFINE_double(
 
 DEFINE_bool(enable_spill, true, "Whether to test plans with spilling enabled.");
 
+DEFINE_int32(
+    max_spill_level,
+    -1,
+    "Max spill level, -1 means random [0, 7], otherwise the actual level.");
+
 DEFINE_bool(
     enable_oom_injection,
     false,
@@ -280,7 +285,10 @@ RowVectorPtr RowNumberFuzzer::execute(
   int32_t spillPct{0};
   if (injectSpill) {
     spillDirectory = exec::test::TempDirectoryPath::create();
+    const auto maxSpillLevel =
+        FLAGS_max_spill_level == -1 ? randInt(0, 7) : FLAGS_max_spill_level;
     builder.config(core::QueryConfig::kSpillEnabled, true)
+        .config(core::QueryConfig::kMaxSpillLevel, maxSpillLevel)
         .config(core::QueryConfig::kRowNumberSpillEnabled, true)
         .spillDirectory(spillDirectory->getPath());
     spillPct = 10;