Merge pull request #25260 from ClickHouse/fix-topk-memory-tracking

Fix memory tracking of aggregate function `topK`
ClickHouse · Jun 14, 2021 · 2370c8b · 2370c8b
2 parents 8a26777 + dbff1ed
commit 2370c8b
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 3 deletions.
diff --git a/src/AggregateFunctions/AggregateFunctionTopK.h b/src/AggregateFunctions/AggregateFunctionTopK.h
@@ -118,7 +118,8 @@ struct AggregateFunctionTopKGenericData
  *  For such columns topK() can be implemented more efficiently (especially for small numeric arrays).
  */
 template <bool is_plain_column, bool is_weighted>
-class AggregateFunctionTopKGeneric : public IAggregateFunctionDataHelper<AggregateFunctionTopKGenericData, AggregateFunctionTopKGeneric<is_plain_column, is_weighted>>
+class AggregateFunctionTopKGeneric
+    : public IAggregateFunctionDataHelper<AggregateFunctionTopKGenericData, AggregateFunctionTopKGeneric<is_plain_column, is_weighted>>
 {
 private:
     using State = AggregateFunctionTopKGenericData;

diff --git a/src/Common/SpaceSaving.h b/src/Common/SpaceSaving.h
@@ -5,6 +5,7 @@
 
 #include <boost/range/adaptor/reversed.hpp>
 
+#include <Common/AllocatorWithMemoryTracking.h>
 #include <Common/ArenaWithFreeLists.h>
 #include <Common/HashTable/Hash.h>
 #include <Common/HashTable/HashMap.h>
@@ -382,8 +383,8 @@ class SpaceSaving
     using CounterMap = HashMapWithStackMemory<TKey, Counter *, Hash, 4>;
 
     CounterMap counter_map;
-    std::vector<Counter *> counter_list;
-    std::vector<UInt64> alpha_map;
+    std::vector<Counter *, AllocatorWithMemoryTracking<Counter *>> counter_list;
+    std::vector<UInt64, AllocatorWithMemoryTracking<UInt64>> alpha_map;
     SpaceSavingArena<TKey> arena;
     size_t m_capacity;
     size_t removed_keys = 0;

diff --git a/tests/queries/0_stateless/01910_memory_tracking_topk.reference b/tests/queries/0_stateless/01910_memory_tracking_topk.reference
diff --git a/tests/queries/0_stateless/01910_memory_tracking_topk.sql b/tests/queries/0_stateless/01910_memory_tracking_topk.sql
@@ -0,0 +1,4 @@
+-- Memory limit must correctly apply, triggering an exception:
+
+SET max_memory_usage = '100M';
+SELECT length(topK(5592405)(tuple(number))) FROM numbers(10) GROUP BY number; -- { serverError 241 }