Fix early partial aggregation handle issue (#5063)

Summary: We shall only check early partial aggregation abandon for partial aggregation. Pull Request resolved: #5063 Reviewed By: spershin Differential Revision: D46258008 Pulled By: xiaoxmeng fbshipit-source-id: 956aa4b2f8a805b63ff8045210e2e493f41219a1
facebookincubator · May 30, 2023 · f72d6c6 · f72d6c6
1 parent 576c2f4
commit f72d6c6
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/velox/common/memory/tests/MockSharedArbitratorTest.cpp b/velox/common/memory/tests/MockSharedArbitratorTest.cpp
@@ -1327,7 +1327,6 @@ DEBUG_ONLY_TEST_F(
     MockSharedArbitrationTest,
     freeUnusedCapacityWhenReclaimMemoryPool) {
   const int allocationSize = kMemoryCapacity / 4;
-  const int reclaimedMemoryCapacity = kMemoryCapacity;
   std::shared_ptr<MockQuery> reclaimedQuery = addQuery(kMemoryCapacity);
   MockMemoryOperator* reclaimedQueryOp = addMemoryOp(reclaimedQuery);
   // The buffer to free later.

diff --git a/velox/exec/HashAggregation.cpp b/velox/exec/HashAggregation.cpp
@@ -162,6 +162,7 @@ HashAggregation::HashAggregation(
 }
 
 bool HashAggregation::abandonPartialAggregationEarly(int64_t numOutput) const {
+  VELOX_CHECK(isPartialOutput_ && !isGlobal_);
   return numInputRows_ > abandonPartialAggregationMinRows_ &&
       100 * numOutput / numInputRows_ >= abandonPartialAggregationMinPct_;
 }
@@ -194,7 +195,7 @@ void HashAggregation::addInput(RowVectorPtr input) {
   // NOTE: we should not trigger partial output flush in case of global
   // aggregation as the final aggregator will handle it the same way as the
   // partial aggregator. Hence, we have to use more memory anyway.
-  const bool abandonPartialEarly =
+  const bool abandonPartialEarly = isPartialOutput_ && !isGlobal_ &&
       abandonPartialAggregationEarly(groupingSet_->numDistinct());
   if (isPartialOutput_ && !isGlobal_ &&
       (abandonPartialEarly ||