Add PercentileTDigest support for MergeAndRollup aggregation

pinot-core/src/main/java/org/apache/pinot/core/common/MinionConstants.java

@@ -189,7 +189,7 @@ public static class RealtimeToOfflineSegmentsTask extends MergeTask {
             DISTINCTCOUNTRAWTHETASKETCH, DISTINCTCOUNTTUPLESKETCH, DISTINCTCOUNTRAWINTEGERSUMTUPLESKETCH,
             SUMVALUESINTEGERSUMTUPLESKETCH, AVGVALUEINTEGERSUMTUPLESKETCH, DISTINCTCOUNTHLLPLUS,
             DISTINCTCOUNTRAWHLLPLUS, DISTINCTCOUNTCPCSKETCH, DISTINCTCOUNTRAWCPCSKETCH, DISTINCTCOUNTULL,
-            DISTINCTCOUNTRAWULL, PERCENTILEKLL, PERCENTILERAWKLL);
+            DISTINCTCOUNTRAWULL, PERCENTILEKLL, PERCENTILERAWKLL, PERCENTILETDIGEST, PERCENTILERAWTDIGEST);
   }
 
   // Generate segment and push to controller based on batch ingestion configs

pinot-core/src/main/java/org/apache/pinot/core/segment/processing/aggregator/PercentileTDigestAggregator.java

@@ -0,0 +1,62 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.segment.processing.aggregator;
+
+import com.tdunning.math.stats.TDigest;
+import java.util.Map;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import org.apache.pinot.core.query.aggregation.function.PercentileTDigestAggregationFunction;
+import org.apache.pinot.segment.spi.Constants;
+
+
+/**
+ * Aggregator for merging serialized TDigest sketches during segment processing
+ * (e.g., MergeAndRollup). Handles both {@code PERCENTILETDIGEST} and
+ * {@code PERCENTILERAWTDIGEST} aggregation types
+ */
+public class PercentileTDigestAggregator implements ValueAggregator {
+
+  @Override
+  public Object aggregate(Object value1, Object value2, Map<String, String> functionParameters) {
+    byte[] bytes1 = (byte[]) value1;
+    byte[] bytes2 = (byte[]) value2;
+
+    // Treat empty byte arrays (default null value for BYTES columns) as missing values
+    if (bytes1.length == 0) {
+      return bytes2;
+    } else if (bytes2.length == 0) {
+      return bytes1;
+    }
+
+    int compression = getCompression(functionParameters);
+    TDigest first = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytes1);
+    TDigest second = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytes2);
+    TDigest merged = TDigest.createMergingDigest(compression);
+    merged.add(first);
+    merged.add(second);
+    return ObjectSerDeUtils.TDIGEST_SER_DE.serialize(merged);
+  }
+
+  private int getCompression(Map<String, String> functionParameters) {
+    String compressionParam = functionParameters.get(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY);
+    return compressionParam != null
+        ? Integer.parseInt(compressionParam)
+        : PercentileTDigestAggregationFunction.DEFAULT_TDIGEST_COMPRESSION;
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/segment/processing/aggregator/ValueAggregatorFactory.java

@@ -64,6 +64,9 @@ public static ValueAggregator getValueAggregator(AggregationFunctionType aggrega
       case PERCENTILEKLL:
       case PERCENTILERAWKLL:
         return new PercentileKLLSketchAggregator();
+      case PERCENTILETDIGEST:
+      case PERCENTILERAWTDIGEST:
+        return new PercentileTDigestAggregator();
       default:
         throw new IllegalStateException("Unsupported aggregation type: " + aggregationType);
     }

pinot-core/src/test/java/org/apache/pinot/core/segment/processing/aggregator/PercentileTDigestAggregatorTest.java

@@ -0,0 +1,154 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.segment.processing.aggregator;
+
+import com.tdunning.math.stats.TDigest;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import org.apache.pinot.segment.spi.AggregationFunctionType;
+import org.apache.pinot.segment.spi.Constants;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+
+
+public class PercentileTDigestAggregatorTest {
+
+  private PercentileTDigestAggregator _aggregator;
+
+  @BeforeMethod
+  public void setUp() {
+    _aggregator = new PercentileTDigestAggregator();
+  }
+
+  @Test
+  public void testAggregateWithDefaultCompression() {
+    TDigest first = TDigest.createMergingDigest(100);
+    for (int i = 0; i < 100; i++) {
+      first.add(i);
+    }
+    TDigest second = TDigest.createMergingDigest(100);
+    for (int i = 100; i < 200; i++) {
+      second.add(i);
+    }
+
+    byte[] value1 = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(first);
+    byte[] value2 = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(second);
+
+    Map<String, String> functionParameters = new HashMap<>();
+    byte[] result = (byte[]) _aggregator.aggregate(value1, value2, functionParameters);
+
+    TDigest resultDigest = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(result);
+    assertNotNull(resultDigest);
+    assertEquals(resultDigest.size(), 200);
+    assertEquals(resultDigest.quantile(0.5), 99.5, 1);
+  }
+
+  @Test
+  public void testAggregateWithCustomCompression() {
+    TDigest first = TDigest.createMergingDigest(100);
+    for (int i = 0; i < 50; i++) {
+      first.add(i);
+    }
+    TDigest second = TDigest.createMergingDigest(100);
+    for (int i = 50; i < 100; i++) {
+      second.add(i);
+    }
+
+    byte[] value1 = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(first);
+    byte[] value2 = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(second);
+
+    Map<String, String> functionParameters = new HashMap<>();
+    functionParameters.put(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, "200");
+
+    byte[] result = (byte[]) _aggregator.aggregate(value1, value2, functionParameters);
+
+    TDigest resultDigest = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(result);
+    assertNotNull(resultDigest);
+    assertEquals(resultDigest.size(), 100);
+    assertEquals(resultDigest.quantile(0.5), 49.5, 1);
+  }
+
+  @Test
+  public void testAggregateWithBothEmptyBytes() {
+    byte[] empty1 = new byte[0];
+    byte[] empty2 = new byte[0];
+
+    Map<String, String> functionParameters = new HashMap<>();
+    byte[] result = (byte[]) _aggregator.aggregate(empty1, empty2, functionParameters);
+
+    // Both empty — treat as missing, return empty bytes
+    assertEquals(result.length, 0);
+  }
+
+  @Test
+  public void testAggregateWithFirstEmptyBytes() {
+    TDigest second = TDigest.createMergingDigest(100);
+    for (int i = 0; i < 50; i++) {
+      second.add(i);
+    }
+    byte[] empty = new byte[0];
+    byte[] value2 = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(second);
+
+    Map<String, String> functionParameters = new HashMap<>();
+    byte[] result = (byte[]) _aggregator.aggregate(empty, value2, functionParameters);
+
+    // Should return the non-empty side as-is
+    assertEquals(result, value2);
+    TDigest resultDigest = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(result);
+    assertEquals(resultDigest.size(), 50);
+  }
+
+  @Test
+  public void testAggregateWithSecondEmptyBytes() {
+    TDigest first = TDigest.createMergingDigest(100);
+    for (int i = 0; i < 50; i++) {
+      first.add(i);
+    }
+    byte[] value1 = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(first);
+    byte[] empty = new byte[0];
+
+    Map<String, String> functionParameters = new HashMap<>();
+    byte[] result = (byte[]) _aggregator.aggregate(value1, empty, functionParameters);
+
+    // Should return the non-empty side as-is
+    assertEquals(result, value1);
+    TDigest resultDigest = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(result);
+    assertEquals(resultDigest.size(), 50);
+  }
+
+  @Test
+  public void testFactoryReturnsAggregatorForNonRawType() {
+    ValueAggregator aggregator = ValueAggregatorFactory.getValueAggregator(
+        AggregationFunctionType.PERCENTILETDIGEST, DataType.BYTES);
+    assertTrue(aggregator instanceof PercentileTDigestAggregator);
+  }
+
+  @Test
+  public void testFactoryReturnsAggregatorForRawType() {
+    ValueAggregator aggregator = ValueAggregatorFactory.getValueAggregator(
+        AggregationFunctionType.PERCENTILERAWTDIGEST, DataType.BYTES);
+    assertTrue(aggregator instanceof PercentileTDigestAggregator);
+  }
+}

pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java

@@ -504,7 +504,8 @@ public void validateTaskConfigs(TableConfig tableConfig, Schema schema, Map<Stri
     // check no mis-configured aggregation function parameters
     Set<String> allowedFunctionParameterNames = ImmutableSet.of(Constants.CPCSKETCH_LGK_KEY.toLowerCase(),
         Constants.THETA_TUPLE_SKETCH_SAMPLING_PROBABILITY.toLowerCase(),
-        Constants.THETA_TUPLE_SKETCH_NOMINAL_ENTRIES.toLowerCase());
+        Constants.THETA_TUPLE_SKETCH_NOMINAL_ENTRIES.toLowerCase(),
+        Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY.toLowerCase());
     Map<String, Map<String, String>> aggregationFunctionParameters =
         MergeRollupTaskUtils.getAggregationFunctionParameters(taskConfigs);
     for (String fieldName : aggregationFunctionParameters.keySet()) {
@@ -515,10 +516,11 @@ public void validateTaskConfigs(TableConfig tableConfig, Schema schema, Map<Stri
       for (String functionParameterName : functionParameters.keySet()) {
         // check that function parameter name is valid
         Preconditions.checkState(allowedFunctionParameterNames.contains(functionParameterName.toLowerCase()),
-            "Aggregation function parameter name must be one of [lgK, samplingProbability, nominalEntries]!");
+            "Aggregation function parameter name must be one of " + allowedFunctionParameterNames + "!");
         // check that function parameter value is valid for nominal entries
         if (functionParameterName.equalsIgnoreCase(Constants.CPCSKETCH_LGK_KEY)
-            || functionParameterName.equalsIgnoreCase(Constants.THETA_TUPLE_SKETCH_NOMINAL_ENTRIES)) {
+            || functionParameterName.equalsIgnoreCase(Constants.THETA_TUPLE_SKETCH_NOMINAL_ENTRIES)
+            || functionParameterName.equalsIgnoreCase(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY)) {
           String value = functionParameters.get(functionParameterName);
           String err = "Aggregation function parameter \"" + functionParameterName + "\" on column \"" + fieldName
               + "\" has invalid value: " + value;