Address review: return RoaringBitmap/BitSet directly from getDictIdBitmap/getDictIdBitSet

- getDictIdBitmap now returns wrapper._dictIdBitmap (RoaringBitmap) directly
- getDictIdBitSet now returns dictIdsWrapper._bitSet (BitSet) directly
- Removed now-unused helper methods from DictIdsWrapper and GroupByDictIdsWrapper
- Updated all call sites to use the raw types directly
- Test: assert reference equality in batch-reuse test to directly verify wrapper reuse

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: 2 people

pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java

@@ -119,7 +119,10 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult
     Dictionary dictionary = blockValSet.getDictionary();
     if (dictionary != null) {
       int[] dictIds = blockValSet.getDictionaryIdsSV();
-      getDictIdBitSet(aggregationResultHolder, dictionary).addDictIds(dictIds, length);
+      BitSet bitSet = getDictIdBitSet(aggregationResultHolder, dictionary);
+      for (int i = 0; i < length; i++) {
+        bitSet.set(dictIds[i]);
+      }
       return;
     }
 
@@ -167,9 +170,11 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult
     Dictionary dictionary = blockValSet.getDictionary();
     if (dictionary != null) {
       int[][] dictIds = blockValSet.getDictionaryIdsMV();
-      DictIdsWrapper dictIdsWrapper = getDictIdBitSet(aggregationResultHolder, dictionary);
+      BitSet bitSet = getDictIdBitSet(aggregationResultHolder, dictionary);
       for (int i = 0; i < length; i++) {
-        dictIdsWrapper.addDictIds(dictIds[i]);
+        for (int dictId : dictIds[i]) {
+          bitSet.set(dictId);
+        }
       }
       return;
     }
@@ -638,30 +643,30 @@ public boolean canUseStarTree(Map<String, Object> functionParameters) {
   }
 
   /**
-   * Returns the {@link GroupByDictIdsWrapper} for the given group key, creating a new one if absent.
-   * Uses a {@link RoaringBitmap} (sparse) so memory scales with distinct values per group, not dictionary size.
+   * Returns the {@link RoaringBitmap} for the given group key, creating a new {@link GroupByDictIdsWrapper} if absent.
+   * Uses a sparse bitmap so memory scales with distinct values per group, not dictionary size.
    */
-  protected static GroupByDictIdsWrapper getDictIdBitmap(GroupByResultHolder groupByResultHolder, int groupKey,
+  protected static RoaringBitmap getDictIdBitmap(GroupByResultHolder groupByResultHolder, int groupKey,
       Dictionary dictionary) {
     GroupByDictIdsWrapper wrapper = groupByResultHolder.getResult(groupKey);
     if (wrapper == null) {
       wrapper = new GroupByDictIdsWrapper(dictionary);
       groupByResultHolder.setValueForKey(groupKey, wrapper);
     }
-    return wrapper;
+    return wrapper._dictIdBitmap;
   }
 
   /**
-   * Returns the {@link DictIdsWrapper} from the result holder, creating a new one if absent.
+   * Returns the {@link BitSet} from the result holder, creating a new {@link DictIdsWrapper} if absent.
    */
-  protected static DictIdsWrapper getDictIdBitSet(AggregationResultHolder aggregationResultHolder,
+  protected static BitSet getDictIdBitSet(AggregationResultHolder aggregationResultHolder,
       Dictionary dictionary) {
     DictIdsWrapper dictIdsWrapper = aggregationResultHolder.getResult();
     if (dictIdsWrapper == null) {
       dictIdsWrapper = new DictIdsWrapper(dictionary);
       aggregationResultHolder.setValue(dictIdsWrapper);
     }
-    return dictIdsWrapper;
+    return dictIdsWrapper._bitSet;
   }
 
   /**
@@ -735,22 +740,6 @@ protected static final class DictIdsWrapper {
       _dictionary = dictionary;
       _bitSet = new BitSet(dictionary.length());
     }
-
-    void set(int dictId) {
-      _bitSet.set(dictId);
-    }
-
-    void addDictIds(int[] dictIds) {
-      for (int dictId : dictIds) {
-        _bitSet.set(dictId);
-      }
-    }
-
-    void addDictIds(int[] dictIds, int length) {
-      for (int i = 0; i < length; i++) {
-        _bitSet.set(dictIds[i]);
-      }
-    }
   }
 
   /**
@@ -768,13 +757,5 @@ protected static final class GroupByDictIdsWrapper {
       _dictionary = dictionary;
       _dictIdBitmap = new RoaringBitmap();
     }
-
-    void add(int dictId) {
-      _dictIdBitmap.add(dictId);
-    }
-
-    void add(int[] dictIds) {
-      _dictIdBitmap.add(dictIds);
-    }
   }
 }

pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java

@@ -19,6 +19,7 @@
 package org.apache.pinot.core.query.aggregation.function;
 
 import com.clearspring.analytics.stream.cardinality.HyperLogLog;
+import java.util.BitSet;
 import java.util.List;
 import java.util.Map;
 import org.apache.pinot.common.request.Literal;
@@ -89,7 +90,10 @@ public void testBitSetDeduplicationProducesCorrectCardinality() {
         List.of(ExpressionContext.forIdentifier("col"),
             ExpressionContext.forLiteral(Literal.intValue(12))));
 
-    DistinctCountHLLAggregationFunction.getDictIdBitSet(holder, dictionary).addDictIds(dictIds, dictIds.length);
+    BitSet bitSet = DistinctCountHLLAggregationFunction.getDictIdBitSet(holder, dictionary);
+    for (int dictId : dictIds) {
+      bitSet.set(dictId);
+    }
 
     HyperLogLog result = function.extractAggregationResult(holder);
     // With log2m=12 the expected error is ~1.6%; allow 5% to be safe
@@ -121,7 +125,10 @@ public void testBitSetLargeCardinalityDictionary() {
         List.of(ExpressionContext.forIdentifier("col"),
             ExpressionContext.forLiteral(Literal.intValue(14))));
 
-    DistinctCountHLLAggregationFunction.getDictIdBitSet(holder, dictionary).addDictIds(dictIds, dictIds.length);
+    BitSet bitSet = DistinctCountHLLAggregationFunction.getDictIdBitSet(holder, dictionary);
+    for (int dictId : dictIds) {
+      bitSet.set(dictId);
+    }
 
     HyperLogLog result = function.extractAggregationResult(holder);
     // log2m=14 gives ~0.8% error; allow 5%
@@ -152,14 +159,21 @@ public void testDictIdBitSetIsReusedAcrossBatches() {
     for (int i = 0; i < 100; i++) {
       batch1[i] = i;
     }
-    DistinctCountHLLAggregationFunction.getDictIdBitSet(holder, dictionary).addDictIds(batch1, batch1.length);
+    BitSet bitSet = DistinctCountHLLAggregationFunction.getDictIdBitSet(holder, dictionary);
+    for (int dictId : batch1) {
+      bitSet.set(dictId);
+    }
 
-    // Batch 2: dict IDs 100–199 (using the same holder — wrapper must be reused)
+    // Batch 2: dict IDs 100–199 (using the same holder — wrapper must be reused, same BitSet instance)
     int[] batch2 = new int[100];
     for (int i = 0; i < 100; i++) {
       batch2[i] = 100 + i;
     }
-    DistinctCountHLLAggregationFunction.getDictIdBitSet(holder, dictionary).addDictIds(batch2, batch2.length);
+    BitSet bitSet2 = DistinctCountHLLAggregationFunction.getDictIdBitSet(holder, dictionary);
+    Assert.assertSame(bitSet, bitSet2, "getDictIdBitSet must return the same BitSet on subsequent calls");
+    for (int dictId : batch2) {
+      bitSet2.set(dictId);
+    }
 
     HyperLogLog result = function.extractAggregationResult(holder);
     Assert.assertEquals(result.cardinality(), numDistinct, numDistinct * 0.05,