From dbff1edcd34ff05ecf657b50f3164e06ded73a81 Mon Sep 17 00:00:00 2001
From: Alexey Milovidov <milovidov@yandex-team.ru>
Date: Mon, 14 Jun 2021 05:26:05 +0300
Subject: [PATCH] Fix memory tracking of aggregate function topK

---
 src/AggregateFunctions/AggregateFunctionTopK.h               | 3 ++-
 src/Common/SpaceSaving.h                                     | 5 +++--
 .../queries/0_stateless/01910_memory_tracking_topk.reference | 0
 tests/queries/0_stateless/01910_memory_tracking_topk.sql     | 4 ++++
 4 files changed, 9 insertions(+), 3 deletions(-)
 create mode 100644 tests/queries/0_stateless/01910_memory_tracking_topk.reference
 create mode 100644 tests/queries/0_stateless/01910_memory_tracking_topk.sql

diff --git a/src/AggregateFunctions/AggregateFunctionTopK.h b/src/AggregateFunctions/AggregateFunctionTopK.h
index 30d69b8ca7b..1988b1f8b2a 100644
--- 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
index b7353d803b7..d1e6d079d17 100644
--- 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 @@ private:
     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
new file mode 100644
index 00000000000..e69de29bb2d
diff --git a/tests/queries/0_stateless/01910_memory_tracking_topk.sql b/tests/queries/0_stateless/01910_memory_tracking_topk.sql
new file mode 100644
index 00000000000..e6309c1eeb8
--- /dev/null
+++ 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 }