[core] Optimize exact and TDigest percentile aggregation, including StarTree group-by

pinot-core/src/main/java/org/apache/pinot/core/data/table/DeterministicConcurrentIndexedTable.java

@@ -26,13 +26,13 @@
 
 
 public class DeterministicConcurrentIndexedTable extends IndexedTable {
-
   public DeterministicConcurrentIndexedTable(DataSchema dataSchema, boolean hasFinalInput,
       QueryContext queryContext, int resultSize,
       int trimSize, int trimThreshold, int initialCapacity, ExecutorService executorService) {
     super(dataSchema, hasFinalInput, queryContext, resultSize, trimSize, trimThreshold,
         new ConcurrentSkipListMap<>(), executorService);
   }
+
   /**
    * Thread safe implementation of upsert for inserting {@link Record} into {@link Table}
    */
@@ -44,12 +44,16 @@ public boolean upsert(Key key, Record record) {
 
   protected void upsertWithoutOrderBy(Key key, Record record) {
     ConcurrentSkipListMap<Key, Record> map = (ConcurrentSkipListMap<Key, Record>) _lookupMap;
-
-    if (map.size() < _resultSize || map.containsKey(key)) {
-      addOrUpdateRecord(key, record);
-    } else if (!map.isEmpty() && key.compareTo(map.lastKey()) < 0) {
-      addOrUpdateRecord(key, record);
-      map.pollLastEntry(); // evict the largest key after insertion
+    // ConcurrentSkipListMap.compute() can invoke a remapping function concurrently and retry it. Aggregation merge
+    // functions mutate the existing record, so protect the entire admission, update, and eviction operation from
+    // duplicate mutation and cross-key eviction races.
+    synchronized (map) {
+      if (map.size() < _resultSize || map.containsKey(key)) {
+        addOrUpdateRecord(key, record);
+      } else if (!map.isEmpty() && key.compareTo(map.lastKey()) < 0) {
+        addOrUpdateRecord(key, record);
+        map.pollLastEntry(); // evict the largest key after insertion
+      }
     }
   }
 }

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

@@ -36,6 +36,14 @@
 public class PercentileAggregationFunction extends NullableSingleInputAggregationFunction<DoubleArrayList, Double> {
   private static final double DEFAULT_FINAL_RESULT = Double.NEGATIVE_INFINITY;
 
+  /// Maximum candidate range length sorted directly instead of partitioned during selection. Direct sorting avoids
+  /// partition overhead once no more than `32` values remain.
+  private static final int DIRECT_SORT_MAX_LENGTH = 32;
+
+  /// Minimum candidate range length that uses Tukey's ninther for pivot selection. Smaller ranges use the median of
+  /// the first, middle, and last values; larger ranges amortize the extra comparisons with better pivot selection.
+  private static final int NINTHER_MIN_LENGTH = 128;
+
   //version 0 functions specified in the of form PERCENTILE<2-digits>(column)
   //version 1 functions of form PERCENTILE(column, <2-digits>.<16-digits>)
   protected final int _version;
@@ -242,13 +250,111 @@ public Double extractFinalResult(DoubleArrayList intermediateResult) {
       }
     } else {
       double[] values = intermediateResult.elements();
-      Arrays.sort(values, 0, size);
-      if (_percentile == 100) {
-        return values[size - 1];
+      int index = _percentile == 100 ? size - 1 : (int) ((long) size * _percentile / 100);
+      return select(values, size, index);
+    }
+  }
+
+  /// Selects the requested rank in place using the same total ordering as a full sort.
+  ///
+  /// Three-way partitioning handles duplicate-heavy inputs, while the depth limit falls back to sorting the remaining
+  /// range to avoid quadratic behavior on adversarial inputs.
+  private static double select(double[] values, int size, int index) {
+    if (index == 0) {
+      double minimum = values[0];
+      for (int i = 1; i < size; i++) {
+        if (compare(values[i], minimum) < 0) {
+          minimum = values[i];
+        }
+      }
+      return minimum;
+    }
+    if (index == size - 1) {
+      double maximum = values[0];
+      for (int i = 1; i < size; i++) {
+        if (compare(values[i], maximum) > 0) {
+          maximum = values[i];
+        }
+      }
+      return maximum;
+    }
+
+    int left = 0;
+    int right = size - 1;
+    int depthLimit = 2 * (31 - Integer.numberOfLeadingZeros(size));
+    while (right - left >= DIRECT_SORT_MAX_LENGTH) {
+      if (depthLimit-- == 0) {
+        Arrays.sort(values, left, right + 1);
+        return values[index];
+      }
+
+      double pivot = choosePivot(values, left, right);
+      int lower = left;
+      int current = left;
+      int upper = right;
+      while (current <= upper) {
+        int comparison = compare(values[current], pivot);
+        if (comparison < 0) {
+          swap(values, lower++, current++);
+        } else if (comparison > 0) {
+          swap(values, current, upper--);
+        } else {
+          current++;
+        }
+      }
+
+      if (index < lower) {
+        right = lower - 1;
+      } else if (index > upper) {
+        left = upper + 1;
       } else {
-        return values[(int) ((long) size * _percentile / 100)];
+        return values[index];
+      }
+    }
+    Arrays.sort(values, left, right + 1);
+    return values[index];
+  }
+
+  private static double choosePivot(double[] values, int left, int right) {
+    int length = right - left + 1;
+    int middle = left + (length >>> 1);
+    if (length >= NINTHER_MIN_LENGTH) {
+      int step = length >>> 3;
+      return medianOfThree(
+          medianOfThree(values[left], values[left + step], values[left + (step << 1)]),
+          medianOfThree(values[middle - step], values[middle], values[middle + step]),
+          medianOfThree(values[right - (step << 1)], values[right - step], values[right]));
+    }
+    return medianOfThree(values[left], values[middle], values[right]);
+  }
+
+  private static double medianOfThree(double first, double second, double third) {
+    if (compare(first, second) < 0) {
+      if (compare(second, third) < 0) {
+        return second;
       }
+      return compare(first, third) < 0 ? third : first;
     }
+    if (compare(first, third) < 0) {
+      return first;
+    }
+    return compare(second, third) < 0 ? third : second;
+  }
+
+  private static int compare(double first, double second) {
+    if (first < second) {
+      return -1;
+    }
+    if (first > second) {
+      return 1;
+    }
+    return Double.compare(first, second);
+  }
+
+  private static void swap(double[] values, int first, int second) {
+    double value = values[first];
+    values[first] = values[second];
+    values[second] = value;
   }
 
   /**