Optimizing Aggregate Holder Array using ThreadLocal arrays

pinot-core/src/main/java/org/apache/pinot/core/operator/query/GroupByOperator.java

@@ -143,6 +143,9 @@ protected GroupByResultsBlock getNextBlock() {
       if (groupByExecutor.getNumGroups() > trimSize) {
         TableResizer tableResizer = new TableResizer(_dataSchema, _queryContext);
         Collection<IntermediateRecord> intermediateRecords = groupByExecutor.trimGroupByResult(trimSize, tableResizer);
+        // intermediateRecords no longer references the holders; safe to return their arrays
+        // to the ThreadLocal cache for the next query.
+        groupByExecutor.releaseHolders();
         GroupByResultsBlock resultsBlock = new GroupByResultsBlock(_dataSchema, intermediateRecords, _queryContext);
         resultsBlock.setNumGroupsLimitReached(numGroupsLimitReached);
         resultsBlock.setNumGroupsWarningLimitReached(numGroupsWarningLimitReached);

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java

@@ -241,4 +241,11 @@ public GroupKeyGenerator getGroupKeyGenerator() {
   public GroupByResultHolder[] getGroupByResultHolders() {
     return _groupByResultHolders;
   }
+
+  @Override
+  public void releaseHolders() {
+    for (GroupByResultHolder holder : _groupByResultHolders) {
+      holder.release();
+    }
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DoubleGroupByResultHolder.java

@@ -19,13 +19,19 @@
 package org.apache.pinot.core.query.aggregation.groupby;
 
 import com.google.common.base.Preconditions;
+import java.lang.ref.SoftReference;
 import java.util.Arrays;
 
 
 /**
  * Result Holder implemented using DoubleArray.
  */
 public class DoubleGroupByResultHolder implements GroupByResultHolder {
+  // Don't cache arrays above this size; outlier queries let them GC normally
+  // instead of pinning per-worker-thread allocations forever.
+  private static final int MAX_CACHEABLE_CAPACITY = 1_000_000;
+  private static final ThreadLocal<SoftReference<double[]>> THREAD_LOCAL_ARRAY = new ThreadLocal<>();
+
   private final int _maxCapacity;
   private final double _defaultValue;
 
@@ -55,22 +61,51 @@ public void ensureCapacity(int capacity) {
     Preconditions.checkArgument(capacity <= _maxCapacity);
 
     if (capacity > _resultHolderCapacity) {
+      SoftReference<double[]> ref = THREAD_LOCAL_ARRAY.get();
+      double[] cached = (ref != null) ? ref.get() : null;
+      // Don't use the ThreadLocal array if it's smaller than the capacity
+      if (cached != null && cached.length < capacity) {
+        cached = null;
+      }
       int copyLength = _resultHolderCapacity;
-      _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity);
+      if (cached != null) {
+        _resultHolderCapacity = cached.length;
+      } else {
+        _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity);
 
-      // Cap the growth to maximum possible number of group keys
-      _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity);
+        // Cap the growth to maximum possible number of group keys
+        _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity);
+      }
 
       double[] current = _resultArray;
-      _resultArray = new double[_resultHolderCapacity];
+      if (cached != null) {
+        _resultArray = cached;
+      } else {
+        _resultArray = new double[_resultHolderCapacity];
+      }
+
       System.arraycopy(current, 0, _resultArray, 0, copyLength);
 
-      if (_defaultValue != 0.0) {
+      // Reused arrays may contain stale data beyond copyLength; fresh arrays only need
+      // filling when the default value is non-zero.
+      if (cached != null || _defaultValue != 0.0) {
         Arrays.fill(_resultArray, copyLength, _resultHolderCapacity, _defaultValue);
       }
     }
   }
 
+  @Override
+  public void release() {
+    if (_resultArray.length <= MAX_CACHEABLE_CAPACITY) {
+      SoftReference<double[]> ref = THREAD_LOCAL_ARRAY.get();
+      double[] cached = (ref != null) ? ref.get() : null;
+      if (cached == null || _resultArray.length > cached.length) {
+        THREAD_LOCAL_ARRAY.set(new SoftReference<>(_resultArray));
+      }
+    }
+    _resultArray = null;
+  }
+
   @Override
   public double getDoubleResult(int groupKey) {
     if (groupKey == GroupKeyGenerator.INVALID_ID) {

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DummyGroupByResultHolder.java

@@ -53,4 +53,8 @@ public <T> T getResult(int groupKey) {
   @Override
   public void ensureCapacity(int capacity) {
   }
+
+  @Override
+  public void release() {
+  }
 }

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByExecutor.java

@@ -60,4 +60,14 @@ public interface GroupByExecutor {
   GroupKeyGenerator getGroupKeyGenerator();
 
   GroupByResultHolder[] getGroupByResultHolders();
+
+  /**
+   * Releases the per-aggregation result holders, allowing each holder to return its backing
+   * array to the ThreadLocal cache for reuse by the next query.
+   * <p>Must only be called once the holders' data is no longer needed (i.e. after
+   * {@link #trimGroupByResult} or after the {@link AggregationGroupByResult} returned by
+   * {@link #getResult} has been fully consumed). Accessing the holders after this call is
+   * undefined.
+   */
+  void releaseHolders();
 }

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/GroupByResultHolder.java

@@ -86,4 +86,9 @@ public interface GroupByResultHolder {
    * @param capacity
    */
   void ensureCapacity(int capacity);
+
+  /**
+   * Reset the threadlocal array.
+   */
+  void release();
 }

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/IntGroupByResultHolder.java

@@ -20,17 +20,22 @@
 package org.apache.pinot.core.query.aggregation.groupby;
 
 import com.google.common.base.Preconditions;
+import java.lang.ref.SoftReference;
 import java.util.Arrays;
 
 
 public class IntGroupByResultHolder implements GroupByResultHolder {
 
+  // Don't cache arrays above this size; outlier queries let them GC normally
+  // instead of pinning per-worker-thread allocations forever.
+  private static final int MAX_CACHEABLE_CAPACITY = 1_000_000;
+  private static final ThreadLocal<SoftReference<int[]>> THREAD_LOCAL_ARRAY = new ThreadLocal<>();
+
   private final int _maxCapacity;
   private final int _defaultValue;
 
   private int _resultHolderCapacity;
   private int[] _resultArray;
-
   /**
    * Constructor for the class.
    *
@@ -54,22 +59,51 @@ public void ensureCapacity(int capacity) {
     Preconditions.checkArgument(capacity <= _maxCapacity);
 
     if (capacity > _resultHolderCapacity) {
+      SoftReference<int[]> ref = THREAD_LOCAL_ARRAY.get();
+      int[] cached = (ref != null) ? ref.get() : null;
+      // Don't use the ThreadLocal array if it's smaller than the capacity
+      if (cached != null && cached.length < capacity) {
+        cached = null;
+      }
+
       int copyLength = _resultHolderCapacity;
-      _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity);
+      if (cached != null) {
+        _resultHolderCapacity = cached.length;
+      } else {
+        _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity);
 
-      // Cap the growth to maximum possible number of group keys
-      _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity);
+        // Cap the growth to maximum possible number of group keys
+        _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity);
+      }
 
       int[] current = _resultArray;
-      _resultArray = new int[_resultHolderCapacity];
+      if (cached != null) {
+        _resultArray = cached;
+      } else {
+        _resultArray = new int[_resultHolderCapacity];
+      }
       System.arraycopy(current, 0, _resultArray, 0, copyLength);
 
-      if (_defaultValue != 0) {
+      // Reused arrays may contain stale data beyond copyLength; fresh arrays only need
+      // filling when the default value is non-zero.
+      if (cached != null || _defaultValue != 0) {
         Arrays.fill(_resultArray, copyLength, _resultHolderCapacity, _defaultValue);
       }
     }
   }
 
+  @Override
+  public void release() {
+    if (_resultArray.length <= MAX_CACHEABLE_CAPACITY) {
+      SoftReference<int[]> ref = THREAD_LOCAL_ARRAY.get();
+      int[] cached = (ref != null) ? ref.get() : null;
+      if (cached == null || _resultArray.length > cached.length) {
+        THREAD_LOCAL_ARRAY.set(new SoftReference<>(_resultArray));
+      }
+    }
+    _resultArray = null;
+  }
+
   @Override
   public double getDoubleResult(int groupKey) {
     throw new UnsupportedOperationException();

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/ObjectGroupByResultHolder.java

@@ -19,12 +19,19 @@
 package org.apache.pinot.core.query.aggregation.groupby;
 
 import com.google.common.base.Preconditions;
+import java.lang.ref.SoftReference;
+import java.util.Arrays;
 
 
 /**
  * Result Holder implemented using ObjectArray.
  */
 public class ObjectGroupByResultHolder implements GroupByResultHolder {
+  // Don't cache arrays above this size; outlier queries let them GC normally
+  // instead of pinning per-worker-thread allocations forever.
+  private static final int MAX_CACHEABLE_CAPACITY = 1_000_000;
+  private static final ThreadLocal<SoftReference<Object[]>> THREAD_LOCAL_ARRAY = new ThreadLocal<>();
+
   private final int _maxCapacity;
 
   private int _resultHolderCapacity;
@@ -43,21 +50,51 @@ public ObjectGroupByResultHolder(int initialCapacity, int maxCapacity) {
     _resultArray = new Object[initialCapacity];
   }
 
+
   @Override
   public void ensureCapacity(int capacity) {
     Preconditions.checkArgument(capacity <= _maxCapacity);
 
     if (capacity > _resultHolderCapacity) {
+      SoftReference<Object[]> ref = THREAD_LOCAL_ARRAY.get();
+      Object[] cached = (ref != null) ? ref.get() : null;
+      if (cached != null && cached.length < capacity) {
+        cached = null;
+      }
       int copyLength = _resultHolderCapacity;
-      _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity);
+      if (cached != null) {
+        _resultHolderCapacity = cached.length;
+      } else {
+        _resultHolderCapacity = Math.max(_resultHolderCapacity * 2, capacity);
 
-      // Cap the growth to maximum possible number of group keys
-      _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity);
+        // Cap the growth to maximum possible number of group keys
+        _resultHolderCapacity = Math.min(_resultHolderCapacity, _maxCapacity);
+      }
 
       Object[] current = _resultArray;
-      _resultArray = new Object[_resultHolderCapacity];
+      if (cached != null) {
+        _resultArray = cached;
+      } else {
+        _resultArray = new Object[_resultHolderCapacity];
+      }
       System.arraycopy(current, 0, _resultArray, 0, copyLength);
+      // release() nulls slots before caching, so reused arrays start clean past copyLength.
+    }
+  }
+
+  @Override
+  public void release() {
+    // Null slots so per-query aggregation results can be GC'd, regardless of whether
+    // we end up caching the array for reuse.
+    Arrays.fill(_resultArray, 0, _resultHolderCapacity, null);
+    if (_resultArray.length <= MAX_CACHEABLE_CAPACITY) {
+      SoftReference<Object[]> ref = THREAD_LOCAL_ARRAY.get();
+      Object[] cached = (ref != null) ? ref.get() : null;
+      if (cached == null || _resultArray.length > cached.length) {
+        THREAD_LOCAL_ARRAY.set(new SoftReference<>(_resultArray));
+      }
     }
+    _resultArray = null;
   }
 
   @Override