[core] Optimize TDigest percentile aggregation (#18996)

Author: xiangfu0

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

@@ -108,20 +108,18 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde
     if (blockValSet.getValueType() == DataType.BYTES) {
       // Serialized TDigest
       byte[][] bytesValues = blockValSet.getBytesValuesSV();
-      foldNotNull(length, blockValSet, (TDigest) aggregationResultHolder.getResult(), (tDigest, from, toEx) -> {
-        if (tDigest != null) {
-          for (int i = from; i < toEx; i++) {
-            tDigest.add(ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytesValues[i]));
-          }
-        } else {
-          tDigest = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytesValues[0]);
-          aggregationResultHolder.setValue(tDigest);
-          for (int i = 1; i < length; i++) {
-            tDigest.add(ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytesValues[i]));
+      foldNotNull(length, blockValSet,
+          (PercentileTDigestAccumulator) aggregationResultHolder.getResult(), (current, from, toExclusive) -> {
+            if (current == null) {
+              current = PercentileTDigestAccumulator.forSerializedTDigestWithMergeBuffers(bytesValues[from]);
+              aggregationResultHolder.setValue(current);
+            }
+            for (int i = from; i < toExclusive; i++) {
+              current.addSerializedTDigest(bytesValues[i]);
+            }
+            return current;
           }
-        }
-        return tDigest;
-      });
+      );
       return;
     }
 
@@ -134,22 +132,17 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde
 
   protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet) {
     double[] doubleValues = blockValSet.getDoubleValuesSV();
-    TDigest tDigest = getDefaultTDigest(aggregationResultHolder, _compressionFactor);
-    forEachNotNull(length, blockValSet, (from, to) -> {
-      for (int i = from; i < to; i++) {
-        tDigest.add(doubleValues[i]);
-      }
-    });
+    PercentileTDigestAccumulator accumulator = getAccumulator(aggregationResultHolder, _compressionFactor);
+    forEachNotNull(length, blockValSet, (from, to) -> accumulator.add(doubleValues, from, to));
   }
 
   protected void aggregateMV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet) {
     double[][] valuesArray = blockValSet.getDoubleValuesMV();
-    TDigest tDigest = getDefaultTDigest(aggregationResultHolder, _compressionFactor);
+    PercentileTDigestAccumulator accumulator = getAccumulator(aggregationResultHolder, _compressionFactor);
     forEachNotNull(length, blockValSet, (from, to) -> {
       for (int i = from; i < to; i++) {
-        for (double value : valuesArray[i]) {
-          tDigest.add(value);
-        }
+        double[] values = valuesArray[i];
+        accumulator.add(values, 0, values.length);
       }
     });
   }
@@ -163,14 +156,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol
       byte[][] bytesValues = blockValSet.getBytesValuesSV();
       forEachNotNull(length, blockValSet, (from, to) -> {
         for (int i = from; i < to; i++) {
-          TDigest value = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytesValues[i]);
           int groupKey = groupKeyArray[i];
-          TDigest tDigest = groupByResultHolder.getResult(groupKey);
-          if (tDigest != null) {
-            tDigest.add(value);
-          } else {
-            groupByResultHolder.setValueForKey(groupKey, value);
-          }
+          getSerializedAccumulator(groupByResultHolder, groupKey, bytesValues[i])
+              .addSerializedTDigest(bytesValues[i]);
         }
       });
       return;
@@ -213,17 +201,17 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult
     if (blockValSet.getValueType() == DataType.BYTES) {
       // Serialized TDigest
       byte[][] bytesValues = blockValSet.getBytesValuesSV();
+      PercentileTDigestAccumulator.SerializedTDigestInput input =
+          new PercentileTDigestAccumulator.SerializedTDigestInput();
       forEachNotNull(length, blockValSet, (from, to) -> {
         for (int i = from; i < to; i++) {
-          TDigest value = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytesValues[i]);
-          for (int groupKey : groupKeysArray[i]) {
-            TDigest tDigest = groupByResultHolder.getResult(groupKey);
-            if (tDigest != null) {
-              tDigest.add(value);
-            } else {
-              // Create a new TDigest for the group
-              groupByResultHolder.setValueForKey(groupKey, ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytesValues[i]));
-            }
+          int[] groupKeys = groupKeysArray[i];
+          if (groupKeys.length == 0) {
+            continue;
+          }
+          input.reset(bytesValues[i]);
+          for (int groupKey : groupKeys) {
+            getSerializedAccumulator(groupByResultHolder, groupKey, input).addSerializedTDigest(input);
           }
         }
       });
@@ -268,21 +256,21 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR
 
   @Override
   public TDigest extractAggregationResult(AggregationResultHolder aggregationResultHolder) {
-    TDigest tDigest = aggregationResultHolder.getResult();
-    if (tDigest == null) {
+    PercentileTDigestAccumulator accumulator = aggregationResultHolder.getResult();
+    if (accumulator == null) {
       return TDigest.createMergingDigest(_compressionFactor);
     } else {
-      return tDigest;
+      return accumulator;
     }
   }
 
   @Override
   public TDigest extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) {
-    TDigest tDigest = groupByResultHolder.getResult(groupKey);
-    if (tDigest == null) {
+    Object result = groupByResultHolder.getResult(groupKey);
+    if (result == null) {
       return TDigest.createMergingDigest(_compressionFactor);
     } else {
-      return tDigest;
+      return (TDigest) result;
     }
   }
 
@@ -294,8 +282,15 @@ public TDigest merge(TDigest intermediateResult1, TDigest intermediateResult2) {
     if (intermediateResult2.size() == 0L) {
       return intermediateResult1;
     }
-    intermediateResult1.add(intermediateResult2);
-    return intermediateResult1;
+    if (intermediateResult1 instanceof PercentileTDigestAccumulator) {
+      intermediateResult1.add(intermediateResult2);
+      return intermediateResult1;
+    }
+    PercentileTDigestAccumulator accumulator =
+        PercentileTDigestAccumulator.forReduction(intermediateResult1.compression());
+    accumulator.add(intermediateResult1);
+    accumulator.add(intermediateResult2);
+    return accumulator;
   }
 
   @Override
@@ -305,13 +300,18 @@ public ColumnDataType getIntermediateResultColumnType() {
 
   @Override
   public SerializedIntermediateResult serializeIntermediateResult(TDigest tDigest) {
-    return new SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.TDigest.getValue(),
-        ObjectSerDeUtils.TDIGEST_SER_DE.serialize(tDigest));
+    byte[] bytes;
+    if (tDigest instanceof PercentileTDigestAccumulator) {
+      bytes = ((PercentileTDigestAccumulator) tDigest).serialize();
+    } else {
+      bytes = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(tDigest);
+    }
+    return new SerializedIntermediateResult(ObjectSerDeUtils.ObjectType.TDigest.getValue(), bytes);
   }
 
   @Override
   public TDigest deserializeIntermediateResult(CustomObject customObject) {
-    return ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(customObject.getBuffer());
+    return PercentileTDigestAccumulator.forSerializedTDigest(customObject.getBuffer());
   }
 
   @Override
@@ -353,6 +353,36 @@ protected static TDigest getDefaultTDigest(AggregationResultHolder aggregationRe
     return tDigest;
   }
 
+  private static PercentileTDigestAccumulator getAccumulator(AggregationResultHolder aggregationResultHolder,
+      int compressionFactor) {
+    PercentileTDigestAccumulator accumulator = aggregationResultHolder.getResult();
+    if (accumulator == null) {
+      accumulator = new PercentileTDigestAccumulator(compressionFactor);
+      aggregationResultHolder.setValue(accumulator);
+    }
+    return accumulator;
+  }
+
+  private static PercentileTDigestAccumulator getSerializedAccumulator(GroupByResultHolder groupByResultHolder,
+      int groupKey, byte[] bytes) {
+    PercentileTDigestAccumulator accumulator = groupByResultHolder.getResult(groupKey);
+    if (accumulator == null) {
+      accumulator = PercentileTDigestAccumulator.forSerializedTDigest(bytes);
+      groupByResultHolder.setValueForKey(groupKey, accumulator);
+    }
+    return accumulator;
+  }
+
+  private static PercentileTDigestAccumulator getSerializedAccumulator(GroupByResultHolder groupByResultHolder,
+      int groupKey, PercentileTDigestAccumulator.SerializedTDigestInput input) {
+    PercentileTDigestAccumulator accumulator = groupByResultHolder.getResult(groupKey);
+    if (accumulator == null) {
+      accumulator = PercentileTDigestAccumulator.forSerializedTDigest(input);
+      groupByResultHolder.setValueForKey(groupKey, accumulator);
+    }
+    return accumulator;
+  }
+
   /**
    * Returns the TDigest for the given group key if exists, or creates a new one with default compression.
    *

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

@@ -20,6 +20,7 @@
 
 import com.clearspring.analytics.stream.cardinality.HyperLogLog;
 import com.dynatrace.hash4j.distinctcount.UltraLogLog;
+import com.tdunning.math.stats.Centroid;
 import com.tdunning.math.stats.TDigest;
 import it.unimi.dsi.fastutil.doubles.Double2LongOpenHashMap;
 import it.unimi.dsi.fastutil.doubles.DoubleArrayList;
@@ -33,7 +34,9 @@
 import it.unimi.dsi.fastutil.longs.LongArrayList;
 import it.unimi.dsi.fastutil.objects.ObjectArrayList;
 import it.unimi.dsi.fastutil.objects.ObjectLinkedOpenHashSet;
+import java.nio.ByteBuffer;
 import java.util.Arrays;
+import java.util.Base64;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -64,6 +67,7 @@
 import org.apache.pinot.segment.local.customobject.TupleIntSketchAccumulator;
 import org.apache.pinot.segment.local.customobject.ValueLongPair;
 import org.apache.pinot.segment.local.utils.UltraLogLogUtils;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
 
 import static org.testng.Assert.assertEquals;
@@ -77,6 +81,23 @@ public class ObjectSerDeUtilsTest {
   private static final String ERROR_MESSAGE = "Random seed: " + RANDOM_SEED;
 
   private static final int NUM_ITERATIONS = 100;
+  private static final String TDIGEST_3_2_VERBOSE_COMPRESSION_20 =
+      "AAAAAT+5mZmZmZmaP+zMzMzMzM1ANAAAAAAAAAAAABtAKgAAAAAAAD+5mZmZmZmaQBgAAAAAAAA/uZmZmZmZmkAcAAAAAAAAP7mZ"
+          + "mZmZmZpAHAAAAAAAAD+5mZmZmZmaQBwAAAAAAAA/uZmZmZmZmkAmAAAAAAAAP7mZmZmZmZpAKgAAAAAAAD+5mZmZmZmaQCYAAAAA"
+          + "AAA/uZmZmZmZmkAmAAAAAAAAP7mZmZmZmZpAJAAAAAAAAD+5mZmZmZmaQC4AAAAAAAA/uZmZmZmZmkAmAAAAAAAAP7mZmZmZmZpA"
+          + "LAAAAAAAAD+5mZmZmZmaQCIAAAAAAAA/uZmZmZmZmkAyAAAAAAAAP8VVVVVVVVZAMQAAAAAAAD/gAAAAAAAAQCwAAAAAAAA/4AAA"
+          + "AAAAAEAsAAAAAAAAP+AAAAAAAABAKAAAAAAAAD/gAAAAAAAAQCQAAAAAAAA/564UeuFHrkAUAAAAAAAAP+zMzMzMzM1AFAAAAAAA"
+          + "AD/szMzMzMzNQBwAAAAAAAA/7MzMzMzMzUAQAAAAAAAAP+zMzMzMzM1ACAAAAAAAAD/szMzMzMzNP/AAAAAAAAA/7MzMzMzMzT/w"
+          + "AAAAAAAAP+zMzMzMzM0=";
+  private static final String TDIGEST_3_2_SMALL_COMPRESSION_1000 =
+      "AAAAAj+INkk/eVmAQJLlmhJu+WhEegAAB9oTiABAP4AAADxBsko/gAAAPRQvNz+AAAA9fAMgP4AAAD20Elw/gAAAPexqgT+AAAA+"
+          + "E5aAP4AAAD4yP88/gAAAPlJFjz+AAAA+c73OP4AAAD6LYDU/gAAAPp2zoz+AAAA+sOdBP4AAAD7FClg/gAAAPtotkD+AAAA+8GMX"
+          + "P4AAAD8D32Y/gAAAPxArOT+AAAA/HSDzP4AAAD8qzbU/gAAAPzk/8j+AAAA/SIegP4AAAD9YtmU/gAAAP2nf1D+AAAA/fBmxP4AA"
+          + "AD+Hvh4/gAAAP5IRRj+AAAA/nRWBP4AAAD+o294/gAAAP7V3mj+AAAA/wv55P4AAAD/RiSw/gAAAP+Ez2T+AAAA/8h6sP4AAAEAC"
+          + "N00/gAAAQAwnIz+AAABAFveKP4AAAEAixUk/gAAAQC+yDj+AAABAPeWGP4AAAEBNjrY/gAAAQF7lsz+AAABAci3iP4AAAECD3G8/"
+          + "gAAAQI/1LD+AAABAnZ6aP4AAAECtJUM/gAAAQL7pkj+AAABA02aDP4AAAEDrOzE/gAAAQQOcVz+AAABBFDtQP4AAAEEoOWA/gAAA"
+          + "QUChLz+AAABBXvGHP4AAAEGCsOA/gAAAQZutuD+AAABBvS/yP4AAAEHr7EA/gAAAQhhGxT+AAABCTlsKP4AAAEKWasI/gAAAQvfH"
+          + "yz+AAABDgpJvP4AAAESXLNE=";
 
   @Test
   public void testString() {
@@ -332,6 +353,45 @@ public void testTDigest() {
     }
   }
 
+  @DataProvider(name = "tdigest32Fixtures")
+  public static Object[][] tdigest32Fixtures() {
+    return new Object[][]{
+        {TDIGEST_3_2_VERBOSE_COMPRESSION_20, 1, 20.0, 256L, 0.1, 0.9},
+        {TDIGEST_3_2_SMALL_COMPRESSION_1000, 2, 1_000.0, 64L, 0.011822292565892623, 1209.4004609431959}
+    };
+  }
+
+  @Test(dataProvider = "tdigest32Fixtures")
+  public void testTDigest32Fixtures(String base64Bytes, int expectedEncoding, double expectedCompression,
+      long expectedSize, double expectedMin, double expectedMax) {
+    byte[] bytes = Base64.getDecoder().decode(base64Bytes);
+    assertEquals(ByteBuffer.wrap(bytes).getInt(), expectedEncoding);
+    TDigest digest = ObjectSerDeUtils.deserialize(bytes, ObjectSerDeUtils.ObjectType.TDigest);
+    assertEquals(digest.compression(), expectedCompression);
+    assertEquals(digest.size(), expectedSize);
+    assertEquals(digest.getMin(), expectedMin);
+    assertEquals(digest.getMax(), expectedMax);
+
+    long centroidWeight = 0L;
+    double previousMean = Double.NEGATIVE_INFINITY;
+    for (Centroid centroid : digest.centroids()) {
+      assertTrue(Double.isFinite(centroid.mean()));
+      assertTrue(centroid.count() > 0);
+      assertTrue(centroid.mean() + 1e-12 >= previousMean);
+      centroidWeight += centroid.count();
+      previousMean = centroid.mean();
+    }
+    assertEquals(centroidWeight, expectedSize);
+
+    double previousQuantile = Double.NEGATIVE_INFINITY;
+    for (double quantile : new double[]{0.0, 0.5, 0.75, 0.95, 0.99, 1.0}) {
+      double value = digest.quantile(quantile);
+      assertTrue(Double.isFinite(value));
+      assertTrue(value >= previousQuantile);
+      previousQuantile = value;
+    }
+  }
+
   @Test
   public void testInt2LongMap() {
     for (int i = 0; i < NUM_ITERATIONS; i++) {