spill agg blocks using BlockInputStream

dbms/src/Core/tests/gtest_spiller.cpp

@@ -512,6 +512,46 @@ try
 }
 CATCH
 
+TEST_F(SpillerTest, SpillFileNumberUsingBlockInputStream)
+try
+{
+    auto blocks = generateBlocks(50);
+    size_t total_block_size = 0;
+    size_t total_block_rows = 0;
+    for (const auto & block : blocks)
+    {
+        total_block_size += block.bytes();
+        total_block_rows += block.rows();
+    }
+
+    auto spiller_config_with_small_max_spill_size = *spill_config_ptr;
+    spiller_config_with_small_max_spill_size.max_cached_data_bytes_in_spiller = total_block_size / 50;
+    spiller_config_with_small_max_spill_size.max_spilled_bytes_per_file
+        = spiller_config_with_small_max_spill_size.max_cached_data_bytes_in_spiller;
+    spiller_config_with_small_max_spill_size.max_spilled_rows_per_file = total_block_rows / 50;
+
+    /// case 1, sorted spiller, only 1 file per spill
+    Spiller sorted_spiller(spiller_config_with_small_max_spill_size, true, 1, spiller_test_header, logger);
+    BlocksList block_list;
+    block_list.insert(block_list.end(), blocks.begin(), blocks.end());
+    auto block_input_stream = std::make_shared<BlocksListBlockInputStream>(std::move(block_list));
+    sorted_spiller.spillBlocksUsingBlockInputStream(block_input_stream, 0, []() { return false; });
+    sorted_spiller.finishSpill();
+    auto restore_streams = sorted_spiller.restoreBlocks(0);
+    ASSERT_TRUE(restore_streams.size() == 1);
+
+    /// case 2, non-sorted spiller, multiple file
+    Spiller non_sorted_spiller(spiller_config_with_small_max_spill_size, false, 1, spiller_test_header, logger);
+    block_list.clear();
+    block_list.insert(block_list.end(), blocks.begin(), blocks.end());
+    block_input_stream = std::make_shared<BlocksListBlockInputStream>(std::move(block_list));
+    non_sorted_spiller.spillBlocksUsingBlockInputStream(block_input_stream, 0, []() { return false; });
+    non_sorted_spiller.finishSpill();
+    restore_streams = non_sorted_spiller.restoreBlocks(0);
+    ASSERT_TRUE(restore_streams.size() > 1);
+}
+CATCH
+
 TEST_F(SpillerTest, SpillAndMeetCancelled)
 try
 {

dbms/src/DataStreams/AggHashTableToBlocksBlockInputStream.h

@@ -0,0 +1,77 @@
+// Copyright 2023 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <DataStreams/IProfilingBlockInputStream.h>
+#include <Interpreters/Aggregator.h>
+
+namespace DB
+{
+template <typename Method>
+class AggHashTableToBlocksBlockInputStream : public IProfilingBlockInputStream
+{
+    static constexpr auto NAME = "AggHashTableToBlocks";
+
+public:
+    AggHashTableToBlocksBlockInputStream(
+        Aggregator & aggregator_,
+        AggregatedDataVariants & aggregated_data_variants_,
+        Method & method_)
+        : aggregator(aggregator_)
+        , aggregated_data_variants(aggregated_data_variants_)
+        , method(method_)
+        , max_block_rows(0)
+        , max_block_bytes(0)
+        , current_bucket(0)
+        , total_bucket(Method::Data::NUM_BUCKETS)
+    {}
+
+    Block getHeader() const override { return aggregator.getHeader(false); }
+    String getName() const override { return NAME; }
+    std::pair<size_t, size_t> maxBlockRowAndBytes() const { return {max_block_rows, max_block_bytes}; }
+
+protected:
+    Block readImpl() override
+    {
+        if (current_bucket < total_bucket)
+        {
+            auto block = aggregator.convertOneBucketToBlock(
+                aggregated_data_variants,
+                method,
+                aggregated_data_variants.aggregates_pool,
+                false,
+                current_bucket++);
+            size_t block_rows = block.rows();
+            size_t block_bytes = block.bytes();
+
+            if (block_rows > max_block_rows)
+                max_block_rows = block_rows;
+            if (block_bytes > max_block_bytes)
+                max_block_bytes = block_bytes;
+            return block;
+        }
+        return {};
+    }
+
+private:
+    Aggregator & aggregator;
+    AggregatedDataVariants & aggregated_data_variants;
+    Method & method;
+    size_t max_block_rows;
+    size_t max_block_bytes;
+    size_t current_bucket;
+    size_t total_bucket;
+};
+} // namespace DB

dbms/src/Interpreters/Aggregator.cpp

@@ -18,6 +18,7 @@
 #include <Common/Stopwatch.h>
 #include <Common/ThresholdUtils.h>
 #include <Common/typeid_cast.h>
+#include <DataStreams/AggHashTableToBlocksBlockInputStream.h>
 #include <DataStreams/materializeBlock.h>
 #include <DataTypes/DataTypeAggregateFunction.h>
 #include <DataTypes/DataTypeNullable.h>
@@ -1019,40 +1020,22 @@ void Aggregator::spillImpl(AggregatedDataVariants & data_variants, Method & meth
     RUNTIME_ASSERT(
         agg_spill_context->getSpiller() != nullptr,
         "spiller must not be nullptr in Aggregator when spilling");
-    size_t max_temporary_block_size_rows = 0;
-    size_t max_temporary_block_size_bytes = 0;
 
-    auto update_max_sizes = [&](const Block & block) {
-        size_t block_size_rows = block.rows();
-        size_t block_size_bytes = block.bytes();
-
-        if (block_size_rows > max_temporary_block_size_rows)
-            max_temporary_block_size_rows = block_size_rows;
-        if (block_size_bytes > max_temporary_block_size_bytes)
-            max_temporary_block_size_bytes = block_size_bytes;
-    };
-
-    Blocks blocks;
-    for (size_t bucket = 0; bucket < Method::Data::NUM_BUCKETS; ++bucket)
-    {
-        /// memory in hash table is released after `convertOneBucketToBlock`,
-        /// so the peak memory usage is not increased although we save all
-        /// the blocks before the actual spill
-        blocks.push_back(convertOneBucketToBlock(data_variants, method, data_variants.aggregates_pool, false, bucket));
-        update_max_sizes(blocks.back());
-    }
-    agg_spill_context->getSpiller()->spillBlocks(std::move(blocks), 0);
+    auto block_input_stream
+        = std::make_shared<AggHashTableToBlocksBlockInputStream<Method>>(*this, data_variants, method);
+    agg_spill_context->getSpiller()->spillBlocksUsingBlockInputStream(block_input_stream, 0, is_cancelled);
     agg_spill_context->finishOneSpill(thread_num);
 
     /// Pass ownership of the aggregate functions states:
     /// `data_variants` will not destroy them in the destructor, they are now owned by ColumnAggregateFunction objects.
     data_variants.aggregator = nullptr;
 
+    auto [max_block_rows, max_block_bytes] = block_input_stream->maxBlockRowAndBytes();
     LOG_TRACE(
         log,
         "Max size of temporary bucket blocks: {} rows, {:.3f} MiB.",
-        max_temporary_block_size_rows,
-        (max_temporary_block_size_bytes / 1048576.0));
+        max_block_rows,
+        (max_block_bytes / 1048576.0));
 }
 
 

dbms/src/Interpreters/Aggregator.h

@@ -48,6 +48,8 @@ extern const int UNKNOWN_AGGREGATED_DATA_VARIANT;
 }
 
 class IBlockOutputStream;
+template <typename Method>
+class AggHashTableToBlocksBlockInputStream;
 
 
 /** Different data structures that can be used for aggregation
@@ -1374,6 +1376,9 @@ class Aggregator
     void destroyImpl(Table & table) const;
 
     void destroyWithoutKey(AggregatedDataVariants & result) const;
+
+    template <typename Method>
+    friend class AggHashTableToBlocksBlockInputStream;
 };
 
 /** Get the aggregation variant by its type. */