Reset buffer aggregators when resetting Groupers. (apache#16296)

Buffer aggregators can contain some cached objects within them, such as
Memory references or HLL Unions. Prior to this patch, various Grouper
implementations were not releasing this state when resetting their own
internal state, which could lead to excessive memory use.

This patch renames AggregatorAdapater#close to "reset", and updates
Grouper implementations to call this reset method whenever they reset
their internal state.

The base method on BufferAggregator and VectorAggregator remains named
"close", for compatibility with existing extensions, but the contract
is adjusted to say that the aggregator may be reused after the method
is called. All existing implementations in core already adhere to this
new contract, except for the ArrayOfDoubles build flavors, which are
updated in this patch to adhere.

Additionally, this patch harmonizes buffer sketch helpers to call their
clear method "clear" rather than a mix of "clear" and "close". (Others
were already using "clear".)

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchMergeBufferAggregator.java

@@ -75,7 +75,7 @@ public Object get(final ByteBuffer buf, final int position)
   @Override
   public void close()
   {
-    helper.close();
+    helper.clear();
   }
 
   @Override

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchMergeBufferAggregatorHelper.java

@@ -142,7 +142,7 @@ private void initializeEmptyUnion(ByteBuffer buf, int position)
     }
   }
 
-  public void close()
+  public void clear()
   {
     unions.clear();
     memCache.clear();

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/hll/HllSketchMergeVectorAggregator.java

@@ -102,7 +102,7 @@ public Object get(final ByteBuffer buf, final int position)
   @Override
   public void close()
   {
-    helper.close();
+    helper.clear();
   }
 
   @Override

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchBufferAggregator.java

@@ -85,7 +85,7 @@ public double getDouble(ByteBuffer buf, int position)
   @Override
   public void close()
   {
-    helper.close();
+    helper.clear();
   }
 
   @Override

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchBufferAggregatorHelper.java

@@ -95,7 +95,7 @@ public void relocate(int oldPosition, int newPosition, ByteBuffer oldBuffer, Byt
   /**
    * Returns a {@link Union} associated with a particular buffer location.
    *
-   * The Union object will be cached in this helper until {@link #close()} is called.
+   * The Union object will be cached in this helper until {@link #clear()} is called.
    */
   public Union getOrCreateUnion(ByteBuffer buf, int position)
   {
@@ -122,7 +122,7 @@ private Union createNewUnion(ByteBuffer buf, int position, boolean isWrapped)
     return union;
   }
 
-  public void close()
+  public void clear()
   {
     unions.clear();
     memCache.clear();

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/theta/SketchVectorAggregator.java

@@ -107,6 +107,6 @@ public void relocate(int oldPosition, int newPosition, ByteBuffer oldBuffer, Byt
   @Override
   public void close()
   {
-    helper.close();
+    helper.clear();
   }
 }

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchBuildAggregator.java

@@ -27,7 +27,6 @@
 import org.apache.druid.segment.data.IndexedInts;
 
 import javax.annotation.Nullable;
-
 import java.nio.ByteBuffer;
 import java.util.LinkedHashMap;
 import java.util.List;
@@ -48,6 +47,7 @@ public class ArrayOfDoublesSketchBuildAggregator implements Aggregator
   @Nullable
   private ArrayOfDoublesUpdatableSketch sketch;
 
+  private final int nominalEntries;
   private final boolean canLookupUtf8;
   private final boolean canCacheById;
   private final LinkedHashMap<Integer, Object> stringCache = new LinkedHashMap<Integer, Object>()
@@ -67,10 +67,7 @@ public ArrayOfDoublesSketchBuildAggregator(
   {
     this.keySelector = keySelector;
     this.valueSelectors = valueSelectors.toArray(new BaseDoubleColumnValueSelector[0]);
-    values = new double[valueSelectors.size()];
-    sketch = new ArrayOfDoublesUpdatableSketchBuilder().setNominalEntries(nominalEntries)
-                                                       .setNumberOfValues(valueSelectors.size()).build();
-
+    this.nominalEntries = nominalEntries;
     this.canCacheById = this.keySelector.nameLookupPossibleInAdvance();
     this.canLookupUtf8 = this.keySelector.supportsLookupNameUtf8();
   }
@@ -83,6 +80,15 @@ public ArrayOfDoublesSketchBuildAggregator(
   @Override
   public void aggregate()
   {
+    if (values == null) {
+      values = new double[valueSelectors.length];
+    }
+
+    if (sketch == null) {
+      sketch = new ArrayOfDoublesUpdatableSketchBuilder().setNominalEntries(nominalEntries)
+                                                         .setNumberOfValues(valueSelectors.length).build();
+    }
+
     final IndexedInts keys = keySelector.getRow();
     for (int i = 0; i < valueSelectors.length; i++) {
       if (valueSelectors[i].isNull()) {

extensions-core/datasketches/src/main/java/org/apache/druid/query/aggregation/datasketches/tuple/ArrayOfDoublesSketchBuildBufferAggregator.java

@@ -73,8 +73,6 @@ public ArrayOfDoublesSketchBuildBufferAggregator(
     this.valueSelectors = valueSelectors.toArray(new BaseDoubleColumnValueSelector[0]);
     this.nominalEntries = nominalEntries;
     this.maxIntermediateSize = maxIntermediateSize;
-    values = new double[valueSelectors.size()];
-
     this.canCacheById = this.keySelector.nameLookupPossibleInAdvance();
     this.canLookupUtf8 = this.keySelector.supportsLookupNameUtf8();
   }
@@ -92,6 +90,10 @@ public void init(final ByteBuffer buf, final int position)
   @Override
   public void aggregate(final ByteBuffer buf, final int position)
   {
+    if (values == null) {
+      values = new double[valueSelectors.length];
+    }
+
     for (int i = 0; i < valueSelectors.length; i++) {
       if (valueSelectors[i].isNull()) {
         return;

processing/src/main/java/org/apache/druid/query/aggregation/AggregatorAdapters.java

@@ -26,7 +26,6 @@
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 
 import javax.annotation.Nullable;
-import java.io.Closeable;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.List;
@@ -42,7 +41,7 @@
  * (2) Query engines are freed from the need to manage how much space each individual aggregator needs. They only
  * need to allocate a block of size "spaceNeeded".
  */
-public class AggregatorAdapters implements Closeable
+public class AggregatorAdapters
 {
   private static final Logger log = new Logger(AggregatorAdapters.class);
 
@@ -230,14 +229,14 @@ public void relocate(int oldPosition, int newPosition, ByteBuffer oldBuffer, Byt
   }
 
   /**
-   * Close all of our aggregators.
+   * Reset all of our aggregators, releasing resources held by them. After this, this instance may be reused or
+   * it may be discarded.
    */
-  @Override
-  public void close()
+  public void reset()
   {
     for (Adapter adapter : adapters) {
       try {
-        adapter.close();
+        adapter.reset();
       }
       catch (Exception e) {
         log.warn(e, "Could not close aggregator [%s], skipping.", adapter.getFactory().getName());
@@ -250,7 +249,7 @@ public void close()
    * BufferAggregator and VectorAggregator. Private, since it doesn't escape this class and the
    * only two implementations are private static classes below.
    */
-  private interface Adapter extends Closeable
+  private interface Adapter
   {
     void init(ByteBuffer buf, int position);
 
@@ -259,8 +258,7 @@ private interface Adapter extends Closeable
 
     void relocate(int oldPosition, int newPosition, ByteBuffer oldBuffer, ByteBuffer newBuffer);
 
-    @Override
-    void close();
+    void reset();
 
     AggregatorFactory getFactory();
 
@@ -293,7 +291,7 @@ public Object get(final ByteBuffer buf, final int position)
     }
 
     @Override
-    public void close()
+    public void reset()
     {
       aggregator.close();
     }
@@ -352,7 +350,7 @@ public Object get(final ByteBuffer buf, final int position)
     }
 
     @Override
-    public void close()
+    public void reset()
     {
       aggregator.close();
     }

processing/src/main/java/org/apache/druid/query/aggregation/BufferAggregator.java

@@ -158,7 +158,11 @@ default double getDouble(ByteBuffer buf, int position)
   }
 
   /**
-   * Release any resources used by the aggregator
+   * Release any resources used by the aggregator. The aggregator may be reused after this call, by calling
+   * {@link #init(ByteBuffer, int)} followed by other methods as normal.
+   *
+   * This call would be more properly named "reset", but we use the name "close" to improve compatibility with
+   * existing aggregator implementations in extensions.
    */
   void close();
 

processing/src/main/java/org/apache/druid/query/aggregation/VectorAggregator.java

@@ -83,7 +83,11 @@ default void relocate(int oldPosition, int newPosition, ByteBuffer oldBuffer, By
   }
 
   /**
-   * Release any resources used by the aggregator.
+   * Release any resources used by the aggregator. The aggregator may be reused after this call, by calling
+   * {@link #init(ByteBuffer, int)} followed by other methods as normal.
+   *
+   * This call would be more properly named "reset", but we use the name "close" to improve compatibility with
+   * existing aggregator implementations in extensions.
    */
   void close();
 }

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

@@ -170,7 +170,7 @@ public AggregateResult aggregate(KeyType key, int keyHash)
   public void close()
   {
     keySerde.reset();
-    aggregators.close();
+    aggregators.reset();
   }
 
   /**

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

@@ -269,6 +269,7 @@ public void reset()
   {
     // Clear the entire usedFlagBuffer
     usedFlagMemory.clear();
+    aggregators.reset();
   }
 
   @Override
@@ -280,7 +281,7 @@ public IntGrouperHashFunction hashFunction()
   @Override
   public void close()
   {
-    aggregators.close();
+    aggregators.reset();
   }
 
   @Override

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

@@ -158,6 +158,7 @@ public void reset()
     offsetList.reset();
     hashTable.reset();
     keySerde.reset();
+    aggregators.reset();
   }
 
   @Override

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

@@ -205,6 +205,7 @@ public void reset()
     }
 
     this.hashTable = createTable(buffer, tableStart, numBuckets);
+    this.aggregators.reset();
   }
 
   @Override
@@ -256,7 +257,7 @@ public void close()
   @Override
   public void close()
   {
-    aggregators.close();
+    aggregators.reset();
   }
 
   @VisibleForTesting

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

@@ -185,6 +185,7 @@ public void reset()
     hashTable.reset();
     keySerde.reset();
     offsetHeap.reset();
+    aggregators.reset();
     heapIndexUpdater.setHashTableBuffer(hashTable.getTableBuffer());
     hasIterated = false;
     offsetHeapIterableSize = 0;

processing/src/main/java/org/apache/druid/query/timeseries/TimeseriesQueryEngine.java

@@ -164,9 +164,9 @@ private Sequence<Result<TimeseriesResultValue>> processVectorized(
       }
 
       final VectorColumnSelectorFactory columnSelectorFactory = cursor.getColumnSelectorFactory();
-      final AggregatorAdapters aggregators = closer.register(
-          AggregatorAdapters.factorizeVector(columnSelectorFactory, query.getAggregatorSpecs())
-      );
+      final AggregatorAdapters aggregators =
+          AggregatorAdapters.factorizeVector(columnSelectorFactory, query.getAggregatorSpecs());
+      closer.register(aggregators::reset);
 
       final ResourceHolder<ByteBuffer> bufferHolder = closer.register(bufferPool.take());
 

processing/src/main/java/org/apache/druid/query/topn/BaseTopNAlgorithm.java

@@ -120,7 +120,7 @@ private void runWithCardinalityKnown(
 
       updateResults(params, theDimValSelector, aggregatesStore, resultBuilder);
 
-      closeAggregators(aggregatesStore);
+      resetAggregators(aggregatesStore);
 
       numProcessed += numToProcess;
       params.getCursor().reset();
@@ -151,7 +151,7 @@ private void runWithCardinalityUnknown(
     }
     long processedRows = scanAndAggregate(params, null, aggregatesStore);
     updateResults(params, null, aggregatesStore, resultBuilder);
-    closeAggregators(aggregatesStore);
+    resetAggregators(aggregatesStore);
     params.getCursor().reset();
     if (queryMetrics != null) {
       queryMetrics.addProcessedRows(processedRows);
@@ -199,7 +199,7 @@ protected abstract void updateResults(
       TopNResultBuilder resultBuilder
   );
 
-  protected abstract void closeAggregators(
+  protected abstract void resetAggregators(
       DimValAggregateStore dimValAggregateStore
   );
 

processing/src/main/java/org/apache/druid/query/topn/HeapBasedTopNAlgorithm.java

@@ -112,7 +112,7 @@ protected void updateResults(
   }
 
   @Override
-  protected void closeAggregators(TopNColumnAggregatesProcessor processor)
+  protected void resetAggregators(TopNColumnAggregatesProcessor processor)
   {
     processor.closeAggregators();
   }

processing/src/main/java/org/apache/druid/query/topn/PooledTopNAlgorithm.java

@@ -768,7 +768,7 @@ protected void updateResults(
   }
 
   @Override
-  protected void closeAggregators(BufferAggregator[] bufferAggregators)
+  protected void resetAggregators(BufferAggregator[] bufferAggregators)
   {
     for (BufferAggregator agg : bufferAggregators) {
       agg.close();

processing/src/main/java/org/apache/druid/query/topn/TimeExtractionTopNAlgorithm.java

@@ -135,7 +135,7 @@ protected void updateResults(
   }
 
   @Override
-  protected void closeAggregators(Map<Object, Aggregator[]> stringMap)
+  protected void resetAggregators(Map<Object, Aggregator[]> stringMap)
   {
     for (Aggregator[] aggregators : stringMap.values()) {
       for (Aggregator agg : aggregators) {

processing/src/test/java/org/apache/druid/query/groupby/epinephelinae/HashVectorGrouperTest.java

@@ -45,7 +45,7 @@ public void testCloseAggregatorAdaptorsShouldBeClosed()
     );
     grouper.initVectorized(512);
     grouper.close();
-    Mockito.verify(aggregatorAdapters, Mockito.times(1)).close();
+    Mockito.verify(aggregatorAdapters, Mockito.times(2)).reset();
   }
 
   @Test