[Enhancement] perform prefetch for AggHashMapWithSerializedKey (#57175)

Signed-off-by: zihe.liu <[email protected]>

Author: ZiheLiu

be/src/exec/aggregate/agg_hash_map.h

@@ -591,6 +591,12 @@ struct AggHashMapWithSerializedKey : public AggHashMapWithKey<HashMap, AggHashMa
     using Iterator = typename HashMap::iterator;
     using ResultVector = Buffer<Slice>;
 
+    struct CacheEntry {
+        KeyType key;
+        size_t hashval;
+    };
+    std::vector<CacheEntry> caches;
+
     template <class... Args>
     AggHashMapWithSerializedKey(int chunk_size, Args&&... args)
             : Base(chunk_size, std::forward<Args>(args)...),
@@ -686,7 +692,23 @@ struct AggHashMapWithSerializedKey : public AggHashMapWithKey<HashMap, AggHashMa
         for (const auto& key_column : key_columns) {
             key_column->serialize_batch(buffer, slice_sizes, chunk_size, max_one_row_size);
         }
+        if (this->hash_map.bucket_count() < prefetch_threhold) {
+            this->template compute_agg_states_by_cols_non_prefetch<Func, allocate_and_compute_state,
+                                                                   compute_not_founds>(
+                    chunk_size, key_columns, pool, std::move(allocate_func), agg_states, not_founds,
+                    max_serialize_each_row);
+        } else {
+            this->template compute_agg_states_by_cols_prefetch<Func, allocate_and_compute_state, compute_not_founds>(
+                    chunk_size, key_columns, pool, std::move(allocate_func), agg_states, not_founds,
+                    max_serialize_each_row);
+        }
+    }
 
+    template <typename Func, bool allocate_and_compute_state, bool compute_not_founds>
+    ALWAYS_NOINLINE void compute_agg_states_by_cols_non_prefetch(size_t chunk_size, const Columns& key_columns,
+                                                                 MemPool* pool, Func&& allocate_func,
+                                                                 Buffer<AggDataPtr>* agg_states, Filter* not_founds,
+                                                                 size_t max_serialize_each_row) {
         for (size_t i = 0; i < chunk_size; ++i) {
             Slice key = {buffer + i * max_one_row_size, slice_sizes[i]};
             if constexpr (allocate_and_compute_state) {
@@ -714,6 +736,50 @@ struct AggHashMapWithSerializedKey : public AggHashMapWithKey<HashMap, AggHashMa
         }
     }
 
+    template <typename Func, bool allocate_and_compute_state, bool compute_not_founds>
+    ALWAYS_NOINLINE void compute_agg_states_by_cols_prefetch(size_t chunk_size, const Columns& key_columns,
+                                                             MemPool* pool, Func&& allocate_func,
+                                                             Buffer<AggDataPtr>* agg_states, Filter* not_founds,
+                                                             size_t max_serialize_each_row) {
+        caches.resize(chunk_size);
+        for (size_t i = 0; i < chunk_size; ++i) {
+            caches[i].key = KeyType(Slice(buffer + i * max_one_row_size, slice_sizes[i]));
+        }
+        for (size_t i = 0; i < chunk_size; ++i) {
+            caches[i].hashval = this->hash_map.hash_function()(caches[i].key);
+        }
+
+        for (size_t i = 0; i < chunk_size; ++i) {
+            if (i + AGG_HASH_MAP_DEFAULT_PREFETCH_DIST < chunk_size) {
+                this->hash_map.prefetch_hash(caches[i + AGG_HASH_MAP_DEFAULT_PREFETCH_DIST].hashval);
+            }
+
+            const auto& key = caches[i].key;
+            if constexpr (allocate_and_compute_state) {
+                auto iter = this->hash_map.lazy_emplace_with_hash(key, caches[i].hashval, [&](const auto& ctor) {
+                    if constexpr (compute_not_founds) {
+                        DCHECK(not_founds);
+                        (*not_founds)[i] = 1;
+                    }
+                    // we must persist the slice before insert
+                    uint8_t* pos = pool->allocate(key.size);
+                    strings::memcpy_inlined(pos, key.data, key.size);
+                    Slice pk{pos, key.size};
+                    AggDataPtr pv = allocate_func(pk);
+                    ctor(pk, pv);
+                });
+                (*agg_states)[i] = iter->second;
+            } else if constexpr (compute_not_founds) {
+                DCHECK(not_founds);
+                if (auto iter = this->hash_map.find(key, caches[i].hashval); iter != this->hash_map.end()) {
+                    (*agg_states)[i] = iter->second;
+                } else {
+                    (*not_founds)[i] = 1;
+                }
+            }
+        }
+    }
+
     uint32_t get_max_serialize_size(const Columns& key_columns) {
         uint32_t max_size = 0;
         for (const auto& key_column : key_columns) {