Significantly reduce memory usage in AggregatingInOrderTransform

azat · azat · commit 2a2f85836516 · 2020-10-03T00:56:58.000+03:00
Clean the aggregates pools (Arena's objects) between flushes, this will
reduce memory usage significantly (since Arena is not intended for
memory reuse in the already full Chunk's)

Before this patch you cannot run SELECT FROM huge_table GROUP BY
primary_key SETTINGS optimize_aggregation_in_order=1 (and the whole
point of optimize_aggregation_in_order got lost), while after, this
should be possible.
diff --git a/src/Processors/Transforms/AggregatingInOrderTransform.cpp b/src/Processors/Transforms/AggregatingInOrderTransform.cpp
@@ -4,6 +4,12 @@
 namespace DB
 {
 
+namespace ErrorCodes
+{
+    extern const int LOGICAL_ERROR;
+}
+
+
 AggregatingInOrderTransform::AggregatingInOrderTransform(
     Block header, AggregatingTransformParamsPtr params_,
     const SortDescription & group_by_description_, size_t res_block_size_)
@@ -140,6 +146,24 @@ void AggregatingInOrderTransform::consume(Chunk chunk)
                 block_end_reached = true;
                 need_generate = true;
                 cur_block_size = 0;
+
+                /// Arenas cannot be destroyed here, since later, in FinalizingSimpleTransform
+                /// there will be finalizeChunk(), but even after
+                /// finalizeChunk() we cannot destroy arena, since some memory
+                /// from Arena still in use, so we attach it to the Chunk to
+                /// remove it once it will be consumed.
+                if (params->final)
+                {
+                    if (variants.aggregates_pools.size() != 1)
+                        throw Exception("Too much arenas", ErrorCodes::LOGICAL_ERROR);
+
+                    Arenas arenas(1, std::make_shared<Arena>());
+                    std::swap(variants.aggregates_pools, arenas);
+                    variants.aggregates_pool = variants.aggregates_pools.at(0).get();
+
+                    chunk.setChunkInfo(std::make_shared<AggregatedArenasChunkInfo>(std::move(arenas)));
+                }
+
                 return;
             }
 
diff --git a/src/Processors/Transforms/AggregatingTransform.h b/src/Processors/Transforms/AggregatingTransform.h
@@ -8,6 +8,15 @@
 namespace DB
 {
 
+class AggregatedArenasChunkInfo : public ChunkInfo
+{
+public:
+    Arenas arenas;
+    AggregatedArenasChunkInfo(Arenas arenas_)
+        : arenas(std::move(arenas_))
+    {}
+};
+
 class AggregatedChunkInfo : public ChunkInfo
 {
 public:
