enable spill manager

apache · Feb 11, 2025 · be95884 · be95884
1 parent c90f170
commit be95884
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 0 deletions.
diff --git a/cpp-ch/local-engine/Operator/GraceAggregatingTransform.cpp b/cpp-ch/local-engine/Operator/GraceAggregatingTransform.cpp
@@ -22,6 +22,7 @@
 #include <Common/GlutenConfig.h>
 #include <Common/QueryContext.h>
 #include <Common/formatReadable.h>
+#include "Processors/IProcessor.h"
 
 namespace DB::ErrorCodes
 {
@@ -60,6 +61,9 @@ GraceAggregatingTransform::GraceAggregatingTransform(
     if (enable_spill_test)
         buckets.emplace(1, BufferFileStream());
     current_data_variants = std::make_shared<DB::AggregatedDataVariants>();
+
+    // IProcessor::spillable, MemorySpillScheduler will trigger the spill by enable this flag.
+    spillable = true;
 }
 
 GraceAggregatingTransform::~GraceAggregatingTransform()
@@ -476,6 +480,7 @@ void GraceAggregatingTransform::mergeOneBlock(const DB::Block & block, bool is_o
     {
         rehashDataVariants();
     }
+    force_spill = false;
 
     LOG_DEBUG(
         logger,
@@ -534,6 +539,13 @@ void GraceAggregatingTransform::mergeOneBlock(const DB::Block & block, bool is_o
 
 bool GraceAggregatingTransform::isMemoryOverflow()
 {
+    if (force_spill)
+    {
+        auto stats = getMemoryStats();
+        if (stats.spillable_memory_bytes > force_spill_on_bytes * 0.8)
+            return true;
+    }
+
     /// More greedy memory usage strategy.
     if (!current_data_variants)
         return false;
@@ -580,4 +592,41 @@ bool GraceAggregatingTransform::isMemoryOverflow()
     return false;
 }
 
+DB::ProcessorMemoryStats GraceAggregatingTransform::getMemoryStats()
+{
+    DB::ProcessorMemoryStats stats;
+    if (!current_data_variants)
+        return stats;
+
+    stats.need_reserved_memory_bytes = current_data_variants->aggregates_pool->allocatedBytes();
+    for (size_t i = current_bucket_index + 1; i < getBucketsNum(); ++i)
+    {
+        auto & file_stream = buckets[i];
+        stats.spillable_memory_bytes += file_stream.pending_bytes;
+    }
+
+    if (per_key_memory_usage > 0)
+    {
+        auto current_result_rows = current_data_variants->size();
+        stats.need_reserved_memory_bytes += current_result_rows * per_key_memory_usage;
+        stats.spillable_memory_bytes += current_result_rows * per_key_memory_usage;
+    }
+    else
+    {
+        // This is a rough estimation, we don't know the exact memory usage for each key.
+        stats.spillable_memory_bytes += current_data_variants->aggregates_pool->allocatedBytes();
+    }
+    return stats;
+}
+
+bool GraceAggregatingTransform::spillOnSize(size_t bytes)
+{
+    auto stats = getMemoryStats();
+    if (stats.spillable_memory_bytes < bytes * 0.8)
+        return false;
+    force_spill = true;
+    force_spill_on_bytes = bytes;
+    return true;
+}
+
 }
diff --git a/cpp-ch/local-engine/Operator/GraceAggregatingTransform.h b/cpp-ch/local-engine/Operator/GraceAggregatingTransform.h
@@ -23,6 +23,7 @@
 #include <Processors/IProcessor.h>
 #include <Processors/Transforms/AggregatingTransform.h>
 #include <Poco/Logger.h>
+#include "Common/MemorySpillScheduler.h"
 #include <Common/AggregateUtil.h>
 
 
@@ -106,7 +107,13 @@ class GraceAggregatingTransform : public DB::IProcessor
     void checkAndSetupCurrentDataVariants();
     /// Merge one block into current_data_variants.
     void mergeOneBlock(const DB::Block & block, bool is_original_block);
+
+    // spill control
     bool isMemoryOverflow();
+    DB::ProcessorMemoryStats getMemoryStats() override;
+    bool spillOnSize(size_t bytes) override;
+    bool force_spill = false; // a force flag to trigger spill
+    bool force_spill_on_bytes = 0;
 
     bool input_finished = false;
     bool has_input = false;