prefetch agg data

Signed-off-by: guo-shaoge <[email protected]>
pingcap · Dec 20, 2024 · 2e94829 · 2e94829
1 parent 6750737
commit 2e94829
Showing 1 changed file with 28 additions and 6 deletions.
diff --git a/dbms/src/AggregateFunctions/IAggregateFunction.h b/dbms/src/AggregateFunctions/IAggregateFunction.h
@@ -235,20 +235,42 @@ class IAggregateFunctionHelper : public IAggregateFunction
         Arena * arena,
         ssize_t if_argument_pos = -1) const override
     {
+        const auto end = start_offset + batch_size;
+        static constexpr size_t prefetch_step = 16;
         if (if_argument_pos >= 0)
         {
             const auto & flags = assert_cast<const ColumnUInt8 &>(*columns[if_argument_pos]).getData();
-            for (size_t i = start_offset; i < start_offset + batch_size; ++i)
+            for (size_t i = start_offset; i < end; ++i)
             {
-                if (flags[i] && places[i - start_offset])
-                    static_cast<const Derived *>(this)->add(places[i - start_offset] + place_offset, columns, i, arena);
+                const auto place_idx = i - start_offset;
+                const auto prefetch_idx = place_idx + prefetch_step;
+
+                if (flags[i] && places[place_idx])
+                {
+                    // Only prefetch first agg function data for simplicity.
+                    // This may not work when there are too many agg functions.
+                    if likely (prefetch_idx < end)
+                        __builtin_prefetch(places[prefetch_idx]);
+
+                    static_cast<const Derived *>(this)->add(places[place_idx] + place_offset, columns, i, arena);
+                }
             }
         }
         else
         {
-            for (size_t i = start_offset; i < start_offset + batch_size; ++i)
-                if (places[i - start_offset])
-                    static_cast<const Derived *>(this)->add(places[i - start_offset] + place_offset, columns, i, arena);
+            for (size_t i = start_offset; i < end; ++i)
+            {
+                const auto place_idx = i - start_offset;
+                const auto prefetch_idx = place_idx + prefetch_step;
+
+                if (places[place_idx])
+                {
+                    if likely (prefetch_idx < end)
+                        __builtin_prefetch(places[prefetch_idx]);
+
+                    static_cast<const Derived *>(this)->add(places[place_idx] + place_offset, columns, i, arena);
+                }
+            }
         }
     }