xx

velox/core/PlanNode.h

@@ -681,8 +681,25 @@ class AggregationNode : public PlanNode {
     // (https://github.com/facebookincubator/velox/issues/3263) and pre-grouped
     // aggregation (https://github.com/facebookincubator/velox/issues/3264). We
     // will add support later to re-enable.
-    return (isFinal() || isSingle()) && !(aggregates().empty()) &&
-        preGroupedKeys().empty() && queryConfig.aggregationSpillEnabled();
+    if (!queryConfig.aggregationSpillEnabled()) {
+      return false;
+    }
+
+    if (!isFinal() && !isSingle()) {
+      return false;
+    }
+
+    if (!preGroupedKeys().empty()) {
+      return false;
+    }
+
+    // aggregates().empty() means distinct aggregate
+    if (aggregates().empty() &&
+        !queryConfig.distinctAggregationSpillEnabled()) {
+      return false;
+    }
+
+    return true;
   }
 
   bool isFinal() const {

velox/core/QueryConfig.h

@@ -144,6 +144,10 @@ class QueryConfig {
   static constexpr const char* kAggregationSpillEnabled =
       "aggregation_spill_enabled";
 
+  /// Distinct aggregation spilling flag
+  static constexpr const char* kDistinctAggregationSpillEnabled =
+      "distinct_aggregation_spill_enabled";
+
   /// Join spilling flag, only applies if "spill_enabled" flag is set.
   static constexpr const char* kJoinSpillEnabled = "join_spill_enabled";
 
@@ -329,6 +333,11 @@ class QueryConfig {
     return get<bool>(kAggregationSpillEnabled, true);
   }
 
+  /// Returns 'is distinct aggregation spilling enabled' flag.
+  bool distinctAggregationSpillEnabled() const {
+    return get<bool>(kDistinctAggregationSpillEnabled, true);
+  }
+
   /// Returns 'is join spilling enabled' flag. Must also check the
   /// spillEnabled()!
   bool joinSpillEnabled() const {

velox/exec/HashAggregation.cpp

@@ -159,6 +159,9 @@ HashAggregation::HashAggregation(
       spillConfig_.has_value() ? &spillConfig_.value() : nullptr,
       &nonReclaimableSection_,
       operatorCtx_.get());
+
+  distinctAggregationSpillEnabled_ =
+      driverCtx->queryConfig().distinctAggregationSpillEnabled();
 }
 
 bool HashAggregation::abandonPartialAggregationEarly(int64_t numOutput) const {
@@ -198,7 +201,7 @@ void HashAggregation::addInput(RowVectorPtr input) {
 
   const bool abandonPartialEarly = isPartialOutput_ && !isGlobal_ &&
       abandonPartialAggregationEarly(groupingSet_->numDistinct());
-  if (isDistinct_) {
+  if (isDistinct_ && !distinctAggregationSpillEnabled_) {
     newDistincts_ = !groupingSet_->hashLookup().newGroups.empty();
 
     if (newDistincts_) {
@@ -349,13 +352,13 @@ RowVectorPtr HashAggregation::getOutput() {
   // - partial aggregation reached memory limit;
   // - distinct aggregation has new keys;
   // - running in partial streaming mode and have some output ready.
-  if (!noMoreInput_ && !partialFull_ && !newDistincts_ &&
+  if (!noMoreInput_ && !partialFull_ && !newDistincts_ && 
       !groupingSet_->hasOutput()) {
     input_ = nullptr;
     return nullptr;
   }
 
-  if (isDistinct_) {
+  if (isDistinct_ && !distinctAggregationSpillEnabled_) {
     if (!newDistincts_) {
       if (noMoreInput_) {
         finished_ = true;
@@ -372,8 +375,8 @@ RowVectorPtr HashAggregation::getOutput() {
     auto output = fillOutput(size, indices);
     numOutputRows_ += size;
 
-    // Drop reference to input_ to make it singly-referenced at the producer and
-    // allow for memory reuse.
+    // Drop reference to input_ to make it singly-referenced at the producer
+    // and allow for memory reuse.
     input_ = nullptr;
 
     resetPartialOutputIfNeed();

velox/exec/HashAggregation.h

@@ -80,6 +80,8 @@ class HashAggregation : public Operator {
   int64_t maxPartialAggregationMemoryUsage_;
   std::unique_ptr<GroupingSet> groupingSet_;
 
+  bool distinctAggregationSpillEnabled_{false};
+
   bool partialFull_ = false;
   bool newDistincts_ = false;
   bool finished_ = false;