From 02d77dc53f6a4c6f49e8f596c41dbc3c70370de3 Mon Sep 17 00:00:00 2001 From: xiaoxmeng Date: Thu, 11 Apr 2024 16:19:16 -0700 Subject: [PATCH] Fix flaky spill stats check in aggregation test The flaky is due to spill check condition is fragile which is based on whether we have received more than one input at partial aggregation. We shall change to final aggregation and given we have four drivers so it is not guarantee one aggregation has received more than one input row. Since we have recorded the spill injection count, then we just rely on this to check spill stats. This PR also restrict the case that we trigger spill for output memory reservation by checking if table is null or empty --- velox/exec/GroupingSet.cpp | 3 ++- .../lib/aggregates/tests/utils/AggregationTestBase.cpp | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/velox/exec/GroupingSet.cpp b/velox/exec/GroupingSet.cpp index 357f2db1d566..df60c7568e8f 100644 --- a/velox/exec/GroupingSet.cpp +++ b/velox/exec/GroupingSet.cpp @@ -900,7 +900,8 @@ void GroupingSet::ensureOutputFits() { // to reserve memory for the output as we can't reclaim much memory from this // operator itself. The output processing can reclaim memory from the other // operator or query through memory arbitration. - if (isPartial_ || spillConfig_ == nullptr || hasSpilled()) { + if (isPartial_ || spillConfig_ == nullptr || hasSpilled() || + table_ == nullptr || table_->numDistinct() == 0) { return; } diff --git a/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp b/velox/functions/lib/aggregates/tests/utils/AggregationTestBase.cpp index 7a04de032494..45776cb5f1a6 100644 --- 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