[branch-1.2](bug) Fix exception-unsafe in aggregation node (#31253)

cherry-pick #28483
apache · Feb 22, 2024 · 1a4da2b · 1a4da2b
1 parent d6bb15e
commit 1a4da2b
Showing 1 changed file with 20 additions and 6 deletions.
diff --git a/be/src/vec/exec/vaggregation_node.h b/be/src/vec/exec/vaggregation_node.h
@@ -737,12 +737,26 @@ struct AggregateDataContainer {
 private:
     void _expand() {
         _index_in_sub_container = 0;
-        _current_keys = _arena_pool.alloc(_size_of_key * SUB_CONTAINER_CAPACITY);
-        _key_containers.emplace_back(_current_keys);
-
-        _current_agg_data = (AggregateDataPtr)_arena_pool.alloc(_size_of_aggregate_states *
-                                                                SUB_CONTAINER_CAPACITY);
-        _value_containers.emplace_back(_current_agg_data);
+        _current_keys = nullptr;
+        _current_agg_data = nullptr;
+        try {
+            _current_keys = _arena_pool.alloc(_size_of_key * SUB_CONTAINER_CAPACITY);
+            _key_containers.emplace_back(_current_keys);
+
+            _current_agg_data = (AggregateDataPtr)_arena_pool.alloc(_size_of_aggregate_states *
+                                                                    SUB_CONTAINER_CAPACITY);
+            _value_containers.emplace_back(_current_agg_data);
+        } catch (...) {
+            if (_current_keys) {
+                _key_containers.pop_back();
+                _current_keys = nullptr;
+            }
+            if (_current_agg_data) {
+                _value_containers.pop_back();
+                _current_agg_data = nullptr;
+            }
+            throw;
+        }
     }
 
 private: