Merge pull request #38748 from azat/Map-combinator-fix

Fix use-after-free for Map combinator that leads to incorrect result

src/AggregateFunctions/AggregateFunctionMap.h

@@ -169,12 +169,21 @@ public:
         {
             const auto & it = merged_maps.find(elem.first);
 
-            if (it != merged_maps.end())
+            AggregateDataPtr nested_place;
+            if (it == merged_maps.end())
             {
-                nested_func->merge(it->second, elem.second, arena);
+                // elem.second cannot be copied since this it will be destroyed after merging,
+                // and lead to use-after-free.
+                nested_place = arena->alignedAlloc(nested_func->sizeOfData(), nested_func->alignOfData());
+                nested_func->create(nested_place);
+                merged_maps.emplace(elem.first, nested_place);
             }
             else
-                merged_maps[elem.first] = elem.second;
+            {
+                nested_place = it->second;
+            }
+
+            nested_func->merge(nested_place, elem.second, arena);
         }
     }
 

tests/queries/0_stateless/02351_Map_combinator_dist.reference

@@ -0,0 +1,4 @@
+1
+1
+1
+1

tests/queries/0_stateless/02351_Map_combinator_dist.sql

@@ -0,0 +1,83 @@
+-- Tags: no-backward-compatibility-check:22.6
+
+-- https://github.com/ClickHouse/ClickHouse/issues/35359
+
+-- sumMap
+SELECT x[67]
+FROM
+(
+    SELECT
+        A,
+        sumMap(CAST(arrayMap(x -> (x, 1), r), 'Map(UInt8,Int64)')) AS x
+    FROM remote('127.{1,1}', view(
+        SELECT
+            number AS A,
+            range(150) AS r
+        FROM numbers(60)
+        WHERE (A % 2) = shardNum()
+    ))
+    GROUP BY A
+    LIMIT 100000000
+)
+WHERE A = 53
+SETTINGS prefer_localhost_replica = 0, distributed_aggregation_memory_efficient = 1, group_by_two_level_threshold = 0, group_by_two_level_threshold_bytes = 0;
+
+-- minMap
+SELECT x[0]
+FROM
+(
+    SELECT
+        A,
+        minMap(CAST(arrayMap(x -> (x, 1), r), 'Map(UInt8,Int64)')) AS x
+    FROM remote('127.{1,1}', view(
+        SELECT
+            number AS A,
+            range(150) AS r
+        FROM numbers(60)
+        WHERE (A % 2) = shardNum()
+    ))
+    GROUP BY A
+    LIMIT 100000000
+)
+WHERE A = 41
+SETTINGS prefer_localhost_replica = 0, distributed_aggregation_memory_efficient = 1, group_by_two_level_threshold = 0, group_by_two_level_threshold_bytes = 0;
+
+-- maxMap
+SELECT x[0]
+FROM
+(
+    SELECT
+        A,
+        maxMap(CAST(arrayMap(x -> (x, 1), r), 'Map(UInt8,Int64)')) AS x
+    FROM remote('127.{1,1}', view(
+        SELECT
+            number AS A,
+            range(150) AS r
+        FROM numbers(60)
+        WHERE (A % 2) = shardNum()
+    ))
+    GROUP BY A
+    LIMIT 100000000
+)
+WHERE A = 41
+SETTINGS prefer_localhost_replica = 0, distributed_aggregation_memory_efficient = 1, group_by_two_level_threshold = 0, group_by_two_level_threshold_bytes = 0;
+
+-- avgMap
+SELECT x[0]
+FROM
+(
+    SELECT
+        A,
+        avgMap(CAST(arrayMap(x -> (x, 1), r), 'Map(UInt8,Int64)')) AS x
+    FROM remote('127.{1,1}', view(
+        SELECT
+            number AS A,
+            range(150) AS r
+        FROM numbers(60)
+        WHERE (A % 2) = shardNum()
+    ))
+    GROUP BY A
+    LIMIT 100000000
+)
+WHERE A = 41
+SETTINGS prefer_localhost_replica = 0, distributed_aggregation_memory_efficient = 1, group_by_two_level_threshold = 0, group_by_two_level_threshold_bytes = 0;