[Groupby Query Metrics] Add merge buffer tracking

docs/operations/metrics.md

@@ -89,9 +89,14 @@ Most metric values reset each emission period, as specified in `druid.monitoring
 |`mergeBuffer/used`|Number of merge buffers used from the merge buffer pool.|This metric is only available if the `GroupByStatsMonitor` module is included.|Depends on the number of groupBy queries needing merge buffers.|
 |`mergeBuffer/queries`|Number of groupBy queries that acquired a batch of buffers from the merge buffer pool.|This metric is only available if the `GroupByStatsMonitor` module is included.|Depends on the number of groupBy queries needing merge buffers.|
 |`mergeBuffer/acquisitionTimeNs`|Total time in nanoseconds to acquire merge buffer for groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`mergeBuffer/maxAcquisitionTimeNs`|Maximum time in nanoseconds to acquire merge buffer for any single groupBy query within the emission period.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`mergeBuffer/bytesUsed`|Number of bytes used by merge buffers to process groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`mergeBuffer/maxBytesUsed`|Maximum number of bytes used by merge buffers for any single groupBy query within the emission period.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
 |`groupBy/spilledQueries`|Number of groupBy queries that have spilled onto the disk.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
 |`groupBy/spilledBytes`|Number of bytes spilled on the disk by the groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`groupBy/maxSpilledBytes`|Maximum number of bytes spilled to disk by any single groupBy query within the emission period.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
 |`groupBy/mergeDictionarySize`|Size of on-heap merge dictionary in bytes.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`groupBy/maxMergeDictionarySize`|Maximum size of the on-heap merge dictionary in bytes observed for any single groupBy query within the emission period.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
 
 ### Historical
 
@@ -113,9 +118,14 @@ Most metric values reset each emission period, as specified in `druid.monitoring
 |`mergeBuffer/used`|Number of merge buffers used from the merge buffer pool.|This metric is only available if the `GroupByStatsMonitor` module is included.|Depends on the number of groupBy queries needing merge buffers.|
 |`mergeBuffer/queries`|Number of groupBy queries that acquired a batch of buffers from the merge buffer pool.|This metric is only available if the `GroupByStatsMonitor` module is included.|Depends on the number of groupBy queries needing merge buffers.|
 |`mergeBuffer/acquisitionTimeNs`|Total time in nanoseconds to acquire merge buffer for groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`mergeBuffer/maxAcquisitionTimeNs`|Maximum time in nanoseconds to acquire merge buffer for any single groupBy query within the emission period.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`mergeBuffer/bytesUsed`|Number of bytes used by merge buffers to process groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`mergeBuffer/maxBytesUsed`|Maximum number of bytes used by merge buffers for any single groupBy query within the emission period.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
 |`groupBy/spilledQueries`|Number of groupBy queries that have spilled onto the disk.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
 |`groupBy/spilledBytes`|Number of bytes spilled on the disk by the groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`groupBy/maxSpilledBytes`|Maximum number of bytes spilled to disk by any single groupBy query within the emission period.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
 |`groupBy/mergeDictionarySize`|Size of on-heap merge dictionary in bytes.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
+|`groupBy/maxMergeDictionarySize`|Maximum size of the on-heap merge dictionary in bytes observed for any single groupBy query within the emission period.|This metric is only available if the `GroupByStatsMonitor` module is included.|Varies|
 
 ### Real-time
 
@@ -140,9 +150,14 @@ to represent the task ID are deprecated and will be removed in a future release.
 |`mergeBuffer/used`|Number of merge buffers used from the merge buffer pool. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Depends on the number of groupBy queries needing merge buffers.|
 |`mergeBuffer/queries`|Number of groupBy queries that acquired a batch of buffers from the merge buffer pool. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Depends on the number of groupBy queries needing merge buffers.|
 |`mergeBuffer/acquisitionTimeNs`|Total time in nanoseconds to acquire merge buffer for groupBy queries. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
+|`mergeBuffer/maxAcquisitionTimeNs`|Maximum time in nanoseconds to acquire merge buffer for any single groupBy query within the emission period. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
+|`mergeBuffer/bytesUsed`|Number of bytes used by merge buffers to process groupBy queries.|This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
+|`mergeBuffer/maxBytesUsed`|Maximum number of bytes used by merge buffers for any single groupBy query within the emission period. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
 |`groupBy/spilledQueries`|Number of groupBy queries that have spilled onto the disk. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
 |`groupBy/spilledBytes`|Number of bytes spilled on the disk by the groupBy queries. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
+|`groupBy/maxSpilledBytes`|Maximum number of bytes spilled to disk by any single groupBy query within the emission period. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
 |`groupBy/mergeDictionarySize`|Size of on-heap merge dictionary in bytes. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
+|`groupBy/maxMergeDictionarySize`|Maximum size of the on-heap merge dictionary in bytes observed for any single groupBy query within the emission period. This metric is only available if the `GroupByStatsMonitor` module is included.|`dataSource`, `taskId`|Varies|
 
 ### Jetty
 

processing/src/main/java/org/apache/druid/query/groupby/GroupByStatsProvider.java

@@ -27,7 +27,8 @@
 import java.util.concurrent.atomic.AtomicLong;
 
 /**
- * Metrics collector for groupBy queries like spilled bytes, merge buffer acquistion time, dictionary size.
+ * Collects groupBy query metrics (spilled bytes, merge buffer usage, dictionary size) per-query, then
+ * aggregates them when queries complete. Stats are retrieved and reset periodically via {@link #getStatsSince()}.
  */
 @LazySingleton
 public class GroupByStatsProvider
@@ -60,34 +61,67 @@ public synchronized void closeQuery(QueryResourceId resourceId)
 
   public synchronized AggregateStats getStatsSince()
   {
-    return aggregateStatsContainer.reset();
+    AggregateStats aggregateStats = new AggregateStats(aggregateStatsContainer);
+    aggregateStatsContainer.reset();
+    return aggregateStats;
   }
 
   public static class AggregateStats
   {
     private long mergeBufferQueries = 0;
     private long mergeBufferAcquisitionTimeNs = 0;
+    private long totalMergeBufferUsedBytes = 0;
+    private long maxMergeBufferAcquisitionTimeNs = 0;
+    private long maxMergeBufferUsedBytes = 0;
     private long spilledQueries = 0;
     private long spilledBytes = 0;
+    private long maxSpilledBytes = 0;
     private long mergeDictionarySize = 0;
+    private long maxMergeDictionarySize = 0;
 
     public AggregateStats()
     {
     }
 
+    public AggregateStats(AggregateStats aggregateStats)
+    {
+      this(
+          aggregateStats.mergeBufferQueries,
+          aggregateStats.mergeBufferAcquisitionTimeNs,
+          aggregateStats.totalMergeBufferUsedBytes,
+          aggregateStats.maxMergeBufferAcquisitionTimeNs,
+          aggregateStats.maxMergeBufferUsedBytes,
+          aggregateStats.spilledQueries,
+          aggregateStats.spilledBytes,
+          aggregateStats.maxSpilledBytes,
+          aggregateStats.mergeDictionarySize,
+          aggregateStats.maxMergeDictionarySize
+      );
+    }
+
     public AggregateStats(
         long mergeBufferQueries,
         long mergeBufferAcquisitionTimeNs,
+        long totalMergeBufferUsedBytes,
+        long maxMergeBufferAcquisitionTimeNs,
+        long maxMergeBufferUsedBytes,
         long spilledQueries,
         long spilledBytes,
-        long mergeDictionarySize
+        long maxSpilledBytes,
+        long mergeDictionarySize,
+        long maxMergeDictionarySize
     )
     {
       this.mergeBufferQueries = mergeBufferQueries;
       this.mergeBufferAcquisitionTimeNs = mergeBufferAcquisitionTimeNs;
+      this.totalMergeBufferUsedBytes = totalMergeBufferUsedBytes;
+      this.maxMergeBufferAcquisitionTimeNs = maxMergeBufferAcquisitionTimeNs;
+      this.maxMergeBufferUsedBytes = maxMergeBufferUsedBytes;
       this.spilledQueries = spilledQueries;
       this.spilledBytes = spilledBytes;
+      this.maxSpilledBytes = maxSpilledBytes;
       this.mergeDictionarySize = mergeDictionarySize;
+      this.maxMergeDictionarySize = maxMergeDictionarySize;
     }
 
     public long getMergeBufferQueries()
@@ -100,6 +134,21 @@ public long getMergeBufferAcquisitionTimeNs()
       return mergeBufferAcquisitionTimeNs;
     }
 
+    public long getMaxMergeBufferAcquisitionTimeNs()
+    {
+      return maxMergeBufferAcquisitionTimeNs;
+    }
+
+    public long getTotalMergeBufferUsedBytes()
+    {
+      return totalMergeBufferUsedBytes;
+    }
+
+    public long getMaxMergeBufferUsedBytes()
+    {
+      return maxMergeBufferUsedBytes;
+    }
+
     public long getSpilledQueries()
     {
       return spilledQueries;
@@ -110,50 +159,63 @@ public long getSpilledBytes()
       return spilledBytes;
     }
 
+    public long getMaxSpilledBytes()
+    {
+      return maxSpilledBytes;
+    }
+
     public long getMergeDictionarySize()
     {
       return mergeDictionarySize;
     }
 
+    public long getMaxMergeDictionarySize()
+    {
+      return maxMergeDictionarySize;
+    }
+
     public void addQueryStats(PerQueryStats perQueryStats)
     {
       if (perQueryStats.getMergeBufferAcquisitionTimeNs() > 0) {
         mergeBufferQueries++;
         mergeBufferAcquisitionTimeNs += perQueryStats.getMergeBufferAcquisitionTimeNs();
+        maxMergeBufferAcquisitionTimeNs = Math.max(
+            maxMergeBufferAcquisitionTimeNs,
+            perQueryStats.getMergeBufferAcquisitionTimeNs()
+        );
+        totalMergeBufferUsedBytes += perQueryStats.getMergeBufferTotalUsedBytes();
+        maxMergeBufferUsedBytes = Math.max(maxMergeBufferUsedBytes, perQueryStats.getMergeBufferTotalUsedBytes());
       }
 
       if (perQueryStats.getSpilledBytes() > 0) {
         spilledQueries++;
         spilledBytes += perQueryStats.getSpilledBytes();
+        maxSpilledBytes = Math.max(maxSpilledBytes, perQueryStats.getSpilledBytes());
       }
 
       mergeDictionarySize += perQueryStats.getMergeDictionarySize();
+      maxMergeDictionarySize = Math.max(maxMergeDictionarySize, perQueryStats.getMergeDictionarySize());
     }
 
-    public AggregateStats reset()
+    public void reset()
     {
-      AggregateStats aggregateStats =
-          new AggregateStats(
-              mergeBufferQueries,
-              mergeBufferAcquisitionTimeNs,
-              spilledQueries,
-              spilledBytes,
-              mergeDictionarySize
-          );
-
       this.mergeBufferQueries = 0;
       this.mergeBufferAcquisitionTimeNs = 0;
+      this.maxMergeBufferAcquisitionTimeNs = 0;
+      this.totalMergeBufferUsedBytes = 0;
+      this.maxMergeBufferUsedBytes = 0;
       this.spilledQueries = 0;
       this.spilledBytes = 0;
+      this.maxSpilledBytes = 0;
       this.mergeDictionarySize = 0;
-
-      return aggregateStats;
+      this.maxMergeDictionarySize = 0;
     }
   }
 
   public static class PerQueryStats
   {
     private final AtomicLong mergeBufferAcquisitionTimeNs = new AtomicLong(0);
+    private final AtomicLong mergeBufferTotalUsedBytes = new AtomicLong(0);
     private final AtomicLong spilledBytes = new AtomicLong(0);
     private final AtomicLong mergeDictionarySize = new AtomicLong(0);
 
@@ -162,6 +224,11 @@ public void mergeBufferAcquisitionTime(long delay)
       mergeBufferAcquisitionTimeNs.addAndGet(delay);
     }
 
+    public void mergeBufferTotalUsedBytes(long bytes)
+    {
+      mergeBufferTotalUsedBytes.addAndGet(bytes);
+    }
+
     public void spilledBytes(long bytes)
     {
       spilledBytes.addAndGet(bytes);
@@ -177,6 +244,11 @@ public long getMergeBufferAcquisitionTimeNs()
       return mergeBufferAcquisitionTimeNs.get();
     }
 
+    public long getMergeBufferTotalUsedBytes()
+    {
+      return mergeBufferTotalUsedBytes.get();
+    }
+
     public long getSpilledBytes()
     {
       return spilledBytes.get();

processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/AbstractBufferHashGrouper.java

@@ -173,6 +173,15 @@ public void close()
     aggregators.reset();
   }
 
+  /**
+   * This method is implemented to return the highest memory value claimed by the Grouper. This is only
+   * used for monitoring the size of the merge buffers used.
+   */
+  public long getMergeBufferUsedBytes()
+  {
+    return hashTable.getMaxTableBufferUsedBytes();
+  }
+
   /**
    * Populate a {@link ReusableEntry} with values from a particular bucket.
    */

processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/BufferHashGrouper.java

@@ -26,7 +26,6 @@
 import org.apache.druid.query.aggregation.AggregatorAdapters;
 import org.apache.druid.query.aggregation.AggregatorFactory;
 
-import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.util.AbstractList;
 import java.util.Collections;
@@ -50,7 +49,6 @@ public class BufferHashGrouper<KeyType> extends AbstractBufferHashGrouper<KeyTyp
   // to get a comparator that uses the ordering defined by the OrderByColumnSpec of a query.
   private final boolean useDefaultSorting;
 
-  @Nullable
   private ByteBufferIntList offsetList;
 
   public BufferHashGrouper(
@@ -154,6 +152,18 @@ public void reset()
     aggregators.reset();
   }
 
+  @Override
+  public long getMergeBufferUsedBytes()
+  {
+    if (!initialized) {
+      return 0L;
+    }
+
+    long hashTableUsage = hashTable.getMaxTableBufferUsedBytes();
+    long offSetListUsage = offsetList.getMaxMergeBufferUsedBytes();
+    return hashTableUsage + offSetListUsage;
+  }
+
   @Override
   public CloseableIterator<Entry<KeyType>> iterator(boolean sorted)
   {

processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/ByteBufferHashTable.java

@@ -26,6 +26,14 @@
 
 import java.nio.ByteBuffer;
 
+/**
+ * A fixed-width, open-addressing hash table that lives inside a caller-provided byte buffer.
+ * <p>
+ * The table uses a contiguous slice of the merge buffer as its backing store. Each bucket holds at most one entry,
+ * and occupies {@code bucketSizeWithHash} number of bytes. Collisions are resolved by continuously probing the
+ * next bucket to find an empty bucket to slot the new entry. The current table view is maintained as a
+ * {@link ByteBuffer} slice that moves and grows within the arena as the table expands.
+ */
 public class ByteBufferHashTable
 {
   public static int calculateTableArenaSizeWithPerBucketAdditionalSize(
@@ -79,6 +87,9 @@ public static int calculateTableArenaSizeWithFixedAdditionalSize(
   @Nullable
   protected BucketUpdateHandler bucketUpdateHandler;
 
+  // Keeps track on how many bytes is being used in the merge buffer.
+  protected long maxTableBufferUsedBytes;
+
   public ByteBufferHashTable(
       float maxLoadFactor,
       int initialBuckets,
@@ -97,6 +108,7 @@ public ByteBufferHashTable(
     this.maxSizeForTesting = maxSizeForTesting;
     this.tableArenaSize = buffer.capacity();
     this.bucketUpdateHandler = bucketUpdateHandler;
+    this.maxTableBufferUsedBytes = 0;
   }
 
   public void reset()
@@ -139,6 +151,7 @@ public void reset()
     bufferDup.position(tableStart);
     bufferDup.limit(tableStart + maxBuckets * bucketSizeWithHash);
     tableBuffer = bufferDup.slice();
+    updateMaxTableBufferUsedBytes();
 
     // Clear used bits of new table
     for (int i = 0; i < maxBuckets; i++) {
@@ -245,6 +258,7 @@ protected void initializeNewBucketKey(
     tableBuffer.putInt(Groupers.getUsedFlag(keyHash));
     tableBuffer.put(keyBuffer);
     size++;
+    updateMaxTableBufferUsedBytes();
 
     if (bucketUpdateHandler != null) {
       bucketUpdateHandler.handleNewBucket(offset);
@@ -381,6 +395,16 @@ public int getGrowthCount()
     return growthCount;
   }
 
+  protected void updateMaxTableBufferUsedBytes()
+  {
+    maxTableBufferUsedBytes = Math.max(maxTableBufferUsedBytes, (long) size * bucketSizeWithHash);
+  }
+
+  public long getMaxTableBufferUsedBytes()
+  {
+    return maxTableBufferUsedBytes;
+  }
+
   public interface BucketUpdateHandler
   {
     void handleNewBucket(int bucketOffset);