Fix backward compatible issue in DistinctCountThetaSketchAggregationF…

…unction (apache#12347)
ankitsultana · Feb 5, 2024 · 5c81c0a · 5c81c0a
1 parent 71b72e1
commit 5c81c0a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 16 deletions.
diff --git a/...pinot/core/query/aggregation/function/DistinctCountRawThetaSketchAggregationFunction.java b/...pinot/core/query/aggregation/function/DistinctCountRawThetaSketchAggregationFunction.java
@@ -53,7 +53,8 @@ public String extractFinalResult(List<ThetaSketchAccumulator> accumulators) {
     int numAccumulators = accumulators.size();
     List<Sketch> mergedSketches = new ArrayList<>(numAccumulators);
 
-    for (ThetaSketchAccumulator accumulator : accumulators) {
+    for (Object object : accumulators) {
+      ThetaSketchAccumulator accumulator = convertSketchAccumulator(object);
       accumulator.setThreshold(_accumulatorThreshold);
       accumulator.setSetOperationBuilder(_setOperationBuilder);
       mergedSketches.add(accumulator.getResult());

diff --git a/...he/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java b/...he/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java
@@ -947,12 +947,9 @@ public List<ThetaSketchAccumulator> extractAggregationResult(AggregationResultHo
     }
 
     if (result.get(0) instanceof Sketch) {
-      int numSketches = result.size();
-      ArrayList<ThetaSketchAccumulator> thetaSketchAccumulators = new ArrayList<>(numSketches);
+      ArrayList<ThetaSketchAccumulator> thetaSketchAccumulators = new ArrayList<>(result.size());
       for (Object o : result) {
-        ThetaSketchAccumulator thetaSketchAccumulator =
-            new ThetaSketchAccumulator(_setOperationBuilder, _accumulatorThreshold);
-        thetaSketchAccumulator.apply((Sketch) o);
+        ThetaSketchAccumulator thetaSketchAccumulator = convertSketchAccumulator(o);
         thetaSketchAccumulators.add(thetaSketchAccumulator);
       }
       return thetaSketchAccumulators;
@@ -974,12 +971,9 @@ public List<ThetaSketchAccumulator> extractGroupByResult(GroupByResultHolder gro
     }
 
     if (result.get(0) instanceof Sketch) {
-      int numSketches = result.size();
-      ArrayList<ThetaSketchAccumulator> thetaSketchAccumulators = new ArrayList<>(numSketches);
+      ArrayList<ThetaSketchAccumulator> thetaSketchAccumulators = new ArrayList<>(result.size());
       for (Object o : result) {
-        ThetaSketchAccumulator thetaSketchAccumulator =
-            new ThetaSketchAccumulator(_setOperationBuilder, _accumulatorThreshold);
-        thetaSketchAccumulator.apply((Sketch) o);
+        ThetaSketchAccumulator thetaSketchAccumulator = convertSketchAccumulator(o);
         thetaSketchAccumulators.add(thetaSketchAccumulator);
       }
       return thetaSketchAccumulators;
@@ -1024,7 +1018,8 @@ public Comparable extractFinalResult(List<ThetaSketchAccumulator> accumulators)
     int numAccumulators = accumulators.size();
     List<Sketch> mergedSketches = new ArrayList<>(numAccumulators);
 
-    for (ThetaSketchAccumulator accumulator : accumulators) {
+    for (Object accumulatorObject : accumulators) {
+      ThetaSketchAccumulator accumulator = convertSketchAccumulator(accumulatorObject);
       accumulator.setThreshold(_accumulatorThreshold);
       accumulator.setSetOperationBuilder(_setOperationBuilder);
       mergedSketches.add(accumulator.getResult());
@@ -1037,14 +1032,14 @@ public Comparable extractFinalResult(List<ThetaSketchAccumulator> accumulators)
   // The AggregationDataTableReducer casts intermediate results to Objects and although the code compiles,
   // types might still be incompatible at runtime due to type erasure.
   // Due to performance overheads of redundant casts, this should be removed at some future point.
-  private ThetaSketchAccumulator convertSketchAccumulator(Object mergeResult) {
-    if (mergeResult instanceof Sketch) {
-      Sketch sketch = (Sketch) mergeResult;
+  protected ThetaSketchAccumulator convertSketchAccumulator(Object result) {
+    if (result instanceof Sketch) {
+      Sketch sketch = (Sketch) result;
       ThetaSketchAccumulator accumulator = new ThetaSketchAccumulator(_setOperationBuilder, _accumulatorThreshold);
       accumulator.apply(sketch);
       return accumulator;
     }
-    return (ThetaSketchAccumulator) mergeResult;
+    return (ThetaSketchAccumulator) result;
   }
 
   /**