Fix flaky spill stats check in aggregation test

facebookincubator · Apr 11, 2024 · 1fa1ca3 · 1fa1ca3
1 parent a3b4849
commit 1fa1ca3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp b/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp
@@ -415,7 +415,7 @@ void AggregationTestBase::testAggregationsWithCompanion(
         toPlanStats(task->taskStats()).at(partialNodeId).inputRows;
     const auto finalInputRows =
         toPlanStats(task->taskStats()).at(finalNodeId).inputRows;
-    if (partialInputRows > 1) {
+    if (exec::injectedSpillCount() > 0) {
       EXPECT_LT(0, spilledBytes(*task))
           << "partial inputRows: " << partialInputRows
           << " final inputRows: " << finalInputRows
@@ -854,7 +854,7 @@ void AggregationTestBase::testAggregationsImpl(
         toPlanStats(task->taskStats()).at(partialNodeId).inputRows;
     const auto finalInputRows =
         toPlanStats(task->taskStats()).at(finalNodeId).inputRows;
-    if (partialInputRows > 1) {
+    if (exec::injectedSpillCount() > 0) {
       EXPECT_LT(0, spilledBytes(*task))
           << "partial inputRows: " << partialInputRows
           << " final inputRows: " << finalInputRows