reserve

Signed-off-by: Seaven <[email protected]>
StarRocks · Jan 2, 2025 · 86f65d8 · 86f65d8
1 parent 3099c16
commit 86f65d8
Show file tree

Hide file tree

Showing 6 changed files with 14 additions and 8 deletions.
diff --git a/be/src/common/config.h b/be/src/common/config.h
@@ -1535,5 +1535,5 @@ CONF_mBool(enable_hk_shuffle_sort, "false");
 
 CONF_mBool(enable_hk_string_select, "false");
 
-CONF_mBool(enable_aggregate_reserve, "false");
+CONF_mBool(enable_hk_aggregate_reserve, "false");
 } // namespace starrocks::config
diff --git a/be/src/exec/aggregate/agg_hash_map.h b/be/src/exec/aggregate/agg_hash_map.h
@@ -106,6 +106,8 @@ struct AggHashMapWithKey {
     HashMap hash_map;
     AggStatistics* agg_stat;
 
+    void reserve(size_t capacity) { hash_map.reserve(capacity); }
+
     ////// Common Methods ////////
     template <typename Func>
     void build_hash_map(size_t chunk_size, const Columns& key_columns, MemPool* pool, Func&& allocate_func,

diff --git a/be/src/exec/aggregate/agg_hash_set.h b/be/src/exec/aggregate/agg_hash_set.h
@@ -105,6 +105,8 @@ struct AggHashSet {
         return x.rehash_number;
     }
 
+    void reserve(size_t n) { hash_set.reserve(n); }
+
     ////// Common Methods ////////
     void build_hash_set(size_t chunk_size, const Columns& key_columns, MemPool* pool) {
         return static_cast<Impl*>(this)->template build_set<true>(chunk_size, key_columns, pool, nullptr);

diff --git a/be/src/exec/aggregate/agg_hash_variant.h b/be/src/exec/aggregate/agg_hash_variant.h
@@ -488,7 +488,9 @@ struct AggHashMapVariant {
 
     size_t allocated_memory_usage(const MemPool* pool) const;
 
-    void reserve(int capacity) {}
+    void reserve(size_t capacity) {
+        visit([&](const auto& hash_map_with_key) { return hash_map_with_key->reserve(capacity); });
+    }
 
 private:
     Type _type = Type::phase1_slice;
@@ -600,7 +602,7 @@ struct AggHashSetVariant {
         return visit([](auto& hash_set_with_key) { return hash_set_with_key->rehash_count(); });
     }
 
-    void reserve(int capacity) {
+    void reserve(size_t capacity) {
         visit([&](auto& hash_set_with_key) { hash_set_with_key->reserve(capacity); });
     }
 

diff --git a/be/src/exec/aggregator.cpp b/be/src/exec/aggregator.cpp
@@ -1367,10 +1367,6 @@ void Aggregator::_init_agg_hash_variant(HashVariantType& hash_variant) {
         }
     }
 
-    if (config::enable_aggregate_reserve) {
-        hash_variant.reserve(_params->predicted_rows / 2);
-    }
-
     VLOG_ROW << "hash type is "
              << static_cast<typename std::underlying_type<typename HashVariantType::Type>::type>(type);
     hash_variant.init(_state, type, _agg_stat);
@@ -1381,6 +1377,10 @@ void Aggregator::_init_agg_hash_variant(HashVariantType& hash_variant) {
             variant->fixed_byte_size = fixed_byte_size;
         }
     });
+
+    if (config::enable_hk_aggregate_reserve) {
+        hash_variant.reserve(_params->predicted_rows / 2);
+    }
 }
 
 void Aggregator::build_hash_map(size_t chunk_size, bool agg_group_by_with_limit) {

diff --git a/fe/fe-core/src/main/java/com/starrocks/planner/PlanNode.java b/fe/fe-core/src/main/java/com/starrocks/planner/PlanNode.java
@@ -521,7 +521,7 @@ private void treeToThriftHelper(TPlan container) {
         msg.node_id = id.asInt();
         msg.num_children = children.size();
         msg.limit = limit;
-        msg.predicted_rows = predictedRows;
+        msg.setPredicted_rows(predictedRows);
 
         for (TupleId tid : tupleIds) {
             msg.addToRow_tuples(tid.asInt());