Fix memory tracking of aggregate function topK

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;

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;

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 }