From 86cd6b39f9206448fa7f48d92c6c9e69b4cc5233 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 15 Jul 2026 11:40:08 -0700 Subject: [PATCH] Optimize percentile aggregations --- .../DeterministicConcurrentIndexedTable.java | 18 +- .../PercentileAggregationFunction.java | 114 +- .../PercentileTDigestAccumulator.java | 1049 +++++++++++++++++ .../PercentileTDigestAggregationFunction.java | 130 +- .../groupby/DefaultGroupByExecutor.java | 7 +- .../DictionaryBasedGroupKeyGenerator.java | 43 +- .../core/common/ObjectSerDeUtilsTest.java | 60 + .../core/data/table/IndexedTableTest.java | 106 ++ .../PercentileAggregationFunctionTest.java | 129 ++ ...centileTDigestAggregationFunctionTest.java | 1010 +++++++++++++++- .../DictionaryBasedGroupKeyGeneratorTest.java | 76 +- .../queries/FilteredAggregationsTest.java | 14 + .../PercentileTDigestMVQueriesTest.java | 14 +- .../queries/PercentileTDigestQueriesTest.java | 66 +- pinot-perf/README.md | 6 + .../percentile-aggregation-100m-results.md | 316 +++++ .../BenchmarkPercentileAggregation.java | 163 +++ ...BenchmarkPercentileTDigestAggregation.java | 114 ++ .../BenchmarkPercentileTDigestCombine.java | 978 +++++++++++++++ ...kPercentileTDigestStarTreeAggregation.java | 195 +++ ...hmarkPercentileTDigestValueAggregator.java | 104 ++ .../PercentileTDigestValueAggregator.java | 40 +- .../PercentileTDigestValueAggregatorTest.java | 154 +++ 23 files changed, 4810 insertions(+), 96 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAccumulator.java create mode 100644 pinot-perf/benchmark-results/percentile-aggregation-100m-results.md create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileAggregation.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestAggregation.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestCombine.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestStarTreeAggregation.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestValueAggregator.java create mode 100644 pinot-segment-local/src/test/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregatorTest.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/table/DeterministicConcurrentIndexedTable.java b/pinot-core/src/main/java/org/apache/pinot/core/data/table/DeterministicConcurrentIndexedTable.java index 6028100f2334..d4838da40e1c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/table/DeterministicConcurrentIndexedTable.java +++ b/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 map = (ConcurrentSkipListMap) _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 + } } } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunction.java index e87d860f93f5..5ae76bd477d9 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunction.java @@ -36,6 +36,14 @@ public class PercentileAggregationFunction extends NullableSingleInputAggregationFunction { 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; } /** diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAccumulator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAccumulator.java new file mode 100644 index 000000000000..569cc269167a --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAccumulator.java @@ -0,0 +1,1049 @@ +/** + * 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.query.aggregation.function; + +import com.tdunning.math.stats.Centroid; +import com.tdunning.math.stats.MergingDigest; +import com.tdunning.math.stats.TDigest; +import java.nio.BufferUnderflowException; +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Collection; +import java.util.List; + +/// Accumulates raw values and serialized TDigests into centroids without repeatedly sorting existing centroids. +/// +/// The accumulator is used while a percentile TDigest result, or one of its group-by results, is being built. Raw +/// values are sorted in small batches, while serialized TDigest centroids are decoded in their existing sorted order. +/// Both are linearly merged with the accumulated centroids and compressed using the same weight-limit rule as +/// [MergingDigest]. [#toTDigest()] normalizes the result to a standard [MergingDigest], so the public intermediate +/// and wire representations remain unchanged. Intermediate serialization retains the standard small encoding when +/// its explicit capacity is required to decode an oversized compact digest. +/// +/// The accumulator implements [TDigest] so process-local combine and reduction can retain the primitive state. +/// Quantile queries and CDFs read the primitive state directly. Centroid iteration and serialization use the +/// canonical library implementation returned from [#toTDigest()]. +/// +/// Serialized group state keeps its first digest pending, and allocates primitive centroid buffers only when another +/// input must be merged. The raw-value buffer remains unallocated for serialized state until a raw value is added. +/// +/// Instances are externally serialized per aggregation or group key and are not safe for concurrent mutation. +final class PercentileTDigestAccumulator extends TDigest { + private static final int MIN_RAW_BUFFER_SIZE = 256; + private static final int MAX_RAW_BUFFER_SIZE = 10_000; + private static final int CENTROID_CAPACITY_PADDING = 10; + private static final int VERBOSE_ENCODING = 1; + private static final int SMALL_ENCODING = 2; + private static final int VERBOSE_HEADER_SIZE = 32; + private static final int VERBOSE_CENTROID_SIZE = 16; + private static final int SMALL_HEADER_SIZE = 30; + private static final int SMALL_CENTROID_SIZE = 8; + private static final int DEFAULT_MERGE_BUFFER_MULTIPLIER = 5; + + private final double _compression; + private final int _centroidCapacity; + private double[] _rawValues; + private double[] _centroidMeans; + private double[] _centroidWeights; + private double[] _outputMeans; + private double[] _outputWeights; + private double[] _incomingMeans; + private double[] _incomingWeights; + private byte[] _pendingSerializedTDigest; + private double _pendingSerializedTotalWeight = Double.NaN; + private boolean _hasSerializedInput; + + private int _numRawValues; + private int _numCentroids; + private int _serializedMainCapacity; + private int _serializedBufferCapacity; + private double _totalWeight; + private double _min = Double.POSITIVE_INFINITY; + private double _max = Double.NEGATIVE_INFINITY; + + PercentileTDigestAccumulator(int compression) { + this((double) compression, true, true); + } + + private PercentileTDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers) { + if (!(compression > 0.0) || !Double.isFinite(compression)) { + throw new IllegalArgumentException("TDigest compression must be positive: " + compression); + } + _compression = compression; + long roundedCompression = (long) Math.ceil(compression); + _centroidCapacity = getCentroidCapacity(compression); + if (allocateRawBuffer) { + _rawValues = new double[getRawBufferSize(roundedCompression)]; + } + if (allocateMergeBuffers) { + _centroidMeans = new double[_centroidCapacity]; + _centroidWeights = new double[_centroidCapacity]; + _outputMeans = new double[_centroidCapacity]; + _outputWeights = new double[_centroidCapacity]; + } + } + + static PercentileTDigestAccumulator forSerializedTDigest(byte[] bytes) { + return new PercentileTDigestAccumulator(readCompression(bytes), false, false); + } + + static PercentileTDigestAccumulator forSerializedTDigest(SerializedTDigestInput input) { + return new PercentileTDigestAccumulator(input._compression, false, false); + } + + static PercentileTDigestAccumulator forSerializedTDigestWithMergeBuffers(byte[] bytes) { + return new PercentileTDigestAccumulator(readCompression(bytes), false, true); + } + + static PercentileTDigestAccumulator forSerializedTDigest(ByteBuffer input) { + byte[] bytes = new byte[input.remaining()]; + input.get(bytes); + PercentileTDigestAccumulator accumulator = forSerializedTDigest(bytes); + accumulator.addSerializedTDigest(bytes); + return accumulator; + } + + static PercentileTDigestAccumulator forReduction(double compression) { + return new PercentileTDigestAccumulator(compression, false, false); + } + + void add(double[] values, int from, int toExclusive) { + if (from >= toExclusive) { + return; + } + materializePendingSerializedTDigest(); + ensureRawBuffer(); + while (from < toExclusive) { + int numValues = Math.min(toExclusive - from, _rawValues.length - _numRawValues); + int rawOffset = _numRawValues; + for (int i = 0; i < numValues; i++) { + double value = values[from + i]; + if (Double.isNaN(value)) { + throw new IllegalArgumentException("Cannot add NaN to t-digest"); + } + _rawValues[rawOffset + i] = value; + } + from += numValues; + _numRawValues += numValues; + if (_numRawValues == _rawValues.length) { + flush(); + } + } + } + + @Override + public void add(double value) { + if (Double.isNaN(value)) { + throw new IllegalArgumentException("Cannot add NaN to t-digest"); + } + materializePendingSerializedTDigest(); + ensureRawBuffer(); + if (_numRawValues == _rawValues.length) { + flush(); + } + _rawValues[_numRawValues++] = value; + } + + @Override + public void add(double value, int weight) { + if (weight == 1) { + add(value); + return; + } + if (Double.isNaN(value)) { + throw new IllegalArgumentException("Cannot add NaN to t-digest"); + } + materializePendingSerializedTDigest(); + flush(); + ensureIncomingCapacity(1); + _incomingMeans[0] = value; + _incomingWeights[0] = weight; + if (mergeSorted(_incomingMeans, _incomingWeights, 1, weight)) { + _min = Math.min(_min, value); + _max = Math.max(_max, value); + } + } + + @Override + public void add(TDigest other) { + if (other instanceof PercentileTDigestAccumulator) { + addAccumulator((PercentileTDigestAccumulator) other); + return; + } + + other.compress(); + int numCentroids = other.centroidCount(); + materializePendingSerializedTDigest(); + flush(); + ensureIncomingCapacity(numCentroids); + int index = 0; + double incomingWeight = 0.0; + for (Centroid centroid : other.centroids()) { + _incomingMeans[index] = centroid.mean(); + double weight = centroid.count(); + _incomingWeights[index] = weight; + incomingWeight += weight; + index++; + } + if (mergeSorted(_incomingMeans, _incomingWeights, index, incomingWeight)) { + _min = Math.min(_min, other.getMin()); + _max = Math.max(_max, other.getMax()); + } + } + + @Override + public void add(List others) { + for (TDigest other : others) { + add(other); + } + } + + private void addAccumulator(PercentileTDigestAccumulator other) { + if (other._pendingSerializedTDigest != null) { + addSerializedTDigest(other._pendingSerializedTDigest); + return; + } + + other.compress(); + materializePendingSerializedTDigest(); + flush(); + preserveSerializedCapacity(other._serializedMainCapacity, other._serializedBufferCapacity); + if (mergeSorted(other._centroidMeans, other._centroidWeights, other._numCentroids, other._totalWeight)) { + _min = Math.min(_min, other._min); + _max = Math.max(_max, other._max); + } + } + + void addSerializedTDigest(byte[] bytes) { + flush(); + if (!_hasSerializedInput && _numCentroids == 0 && _totalWeight == 0.0) { + _pendingSerializedTotalWeight = getSerializedTotalWeight(bytes); + _pendingSerializedTDigest = bytes; + _hasSerializedInput = true; + return; + } + materializePendingSerializedTDigest(); + mergeSerializedTDigest(bytes); + _hasSerializedInput = true; + } + + void addSerializedTDigest(SerializedTDigestInput input) { + flush(); + if (!_hasSerializedInput && _numCentroids == 0 && _totalWeight == 0.0) { + _pendingSerializedTDigest = input._bytes; + _pendingSerializedTotalWeight = Double.NaN; + _hasSerializedInput = true; + return; + } + materializePendingSerializedTDigest(); + input.decode(); + mergeSerializedTDigest(input); + _hasSerializedInput = true; + } + + private void mergeSerializedTDigest(byte[] bytes) { + ByteBuffer input = ByteBuffer.wrap(bytes); + int encoding = input.getInt(); + if (encoding != VERBOSE_ENCODING && encoding != SMALL_ENCODING) { + throw new IllegalStateException("Invalid format for serialized histogram"); + } + double min = input.getDouble(); + double max = input.getDouble(); + int numCentroids; + int mainCapacity = 0; + int bufferCapacity = 0; + boolean initialize = _numCentroids == 0 && _totalWeight == 0.0; + double[] incomingMeans; + double[] incomingWeights; + if (encoding == VERBOSE_ENCODING) { + double serializedCompression = input.getDouble(); + numCentroids = input.getInt(); + checkCentroidCount(numCentroids); + checkCentroidsAvailable(input, numCentroids, VERBOSE_CENTROID_SIZE); + checkVerboseCentroidCapacity(serializedCompression, numCentroids); + if (initialize) { + ensureCentroidCapacity(numCentroids); + incomingMeans = _centroidMeans; + incomingWeights = _centroidWeights; + } else { + ensureIncomingCapacity(numCentroids); + incomingMeans = _incomingMeans; + incomingWeights = _incomingWeights; + } + for (int i = 0; i < numCentroids; i++) { + incomingWeights[i] = input.getDouble(); + incomingMeans[i] = input.getDouble(); + } + } else { + double serializedCompression = input.getFloat(); + mainCapacity = input.getShort(); + bufferCapacity = input.getShort(); + checkSmallArrayCapacities(mainCapacity, bufferCapacity); + mainCapacity = resolveSmallMainCapacity(serializedCompression, mainCapacity); + numCentroids = input.getShort(); + checkCentroidCount(numCentroids); + checkCentroidsAvailable(input, numCentroids, SMALL_CENTROID_SIZE); + checkSmallCentroidCapacity(mainCapacity, numCentroids); + if (initialize) { + ensureCentroidCapacity(numCentroids); + incomingMeans = _centroidMeans; + incomingWeights = _centroidWeights; + } else { + ensureIncomingCapacity(numCentroids); + incomingMeans = _incomingMeans; + incomingWeights = _incomingWeights; + } + for (int i = 0; i < numCentroids; i++) { + incomingWeights[i] = input.getFloat(); + incomingMeans[i] = input.getFloat(); + } + } + preserveSerializedCapacity(mainCapacity, bufferCapacity); + + if (initialize) { + double totalWeight = getTotalWeight(incomingWeights, numCentroids); + if (totalWeight != 0.0) { + _numCentroids = numCentroids; + _totalWeight = totalWeight; + _min = min; + _max = max; + } + } else if (mergeSorted(incomingMeans, incomingWeights, numCentroids)) { + _min = Math.min(_min, min); + _max = Math.max(_max, max); + } + } + + private void mergeSerializedTDigest(SerializedTDigestInput input) { + preserveSerializedCapacity(input._mainCapacity, input._bufferCapacity); + boolean initialize = _numCentroids == 0 && _totalWeight == 0.0; + if (initialize) { + if (input._totalWeight != 0.0) { + ensureCentroidCapacity(input._numCentroids); + System.arraycopy(input._means, 0, _centroidMeans, 0, input._numCentroids); + System.arraycopy(input._weights, 0, _centroidWeights, 0, input._numCentroids); + _numCentroids = input._numCentroids; + _totalWeight = input._totalWeight; + _min = input._min; + _max = input._max; + } + } else if (mergeSorted(input._means, input._weights, input._numCentroids, input._totalWeight)) { + _min = Math.min(_min, input._min); + _max = Math.max(_max, input._max); + } + } + + @Override + public void compress() { + materializePendingSerializedTDigest(); + flush(); + } + + @Override + public long size() { + if (_pendingSerializedTDigest != null) { + if (Double.isNaN(_pendingSerializedTotalWeight)) { + _pendingSerializedTotalWeight = getSerializedTotalWeight(_pendingSerializedTDigest); + } + return (long) _pendingSerializedTotalWeight; + } + return (long) (_totalWeight + _numRawValues); + } + + @Override + public double cdf(double value) { + compress(); + if (_numCentroids == 0) { + return Double.NaN; + } + if (_numCentroids == 1) { + double width = _max - _min; + if (value < _min) { + return 0.0; + } else if (value > _max) { + return 1.0; + } else if (value - _min <= width) { + return 0.5; + } else { + return (value - _min) / width; + } + } + if (value <= _min) { + return 0.0; + } + if (value >= _max) { + return 1.0; + } + if (value <= _centroidMeans[0]) { + double width = _centroidMeans[0] - _min; + return width > 0.0 ? (value - _min) / width * _centroidWeights[0] / _totalWeight / 2.0 : 0.0; + } + int lastIndex = _numCentroids - 1; + if (value >= _centroidMeans[lastIndex]) { + double width = _max - _centroidMeans[lastIndex]; + return width > 0.0 + ? 1.0 - (_max - value) / width * _centroidWeights[lastIndex] / _totalWeight / 2.0 : 1.0; + } + + double weightSoFar = _centroidWeights[0] / 2.0; + for (int i = 0; i < lastIndex; i++) { + if (_centroidMeans[i] == value) { + double startWeight = weightSoFar; + int equalIndex = i; + while (equalIndex < lastIndex && _centroidMeans[equalIndex + 1] == value) { + weightSoFar += _centroidWeights[equalIndex] + _centroidWeights[equalIndex + 1]; + equalIndex++; + } + return (startWeight + weightSoFar) / 2.0 / _totalWeight; + } + if (_centroidMeans[i] <= value && _centroidMeans[i + 1] > value) { + double width = _centroidMeans[i + 1] - _centroidMeans[i]; + double halfWeight = (_centroidWeights[i] + _centroidWeights[i + 1]) / 2.0; + return width > 0.0 + ? (weightSoFar + halfWeight * (value - _centroidMeans[i]) / width) / _totalWeight + : weightSoFar + halfWeight / _totalWeight; + } + weightSoFar += (_centroidWeights[i] + _centroidWeights[i + 1]) / 2.0; + } + throw new IllegalStateException("Unable to compute TDigest CDF"); + } + + @Override + public double quantile(double quantile) { + if (quantile < 0.0 || quantile > 1.0) { + throw new IllegalArgumentException("q should be in [0,1], got " + quantile); + } + compress(); + if (_numCentroids == 0) { + return Double.NaN; + } + if (_numCentroids == 1) { + return _centroidMeans[0]; + } + + double index = quantile * _totalWeight; + if (index < _centroidWeights[0] / 2.0) { + return _min + 2.0 * index / _centroidWeights[0] * (_centroidMeans[0] - _min); + } + double weightSoFar = _centroidWeights[0] / 2.0; + int lastIndex = _numCentroids - 1; + for (int i = 0; i < lastIndex; i++) { + double halfWeight = (_centroidWeights[i] + _centroidWeights[i + 1]) / 2.0; + if (weightSoFar + halfWeight > index) { + double leftWeight = index - weightSoFar; + double rightWeight = weightSoFar + halfWeight - index; + return weightedAverage(_centroidMeans[i], rightWeight, _centroidMeans[i + 1], leftWeight); + } + weightSoFar += halfWeight; + } + double leftWeight = index - _totalWeight - _centroidWeights[lastIndex] / 2.0; + double rightWeight = _centroidWeights[lastIndex] / 2.0 - leftWeight; + return weightedAverage(_centroidMeans[lastIndex], leftWeight, _max, rightWeight); + } + + private static double weightedAverage(double firstValue, double firstWeight, double secondValue, + double secondWeight) { + if (firstValue > secondValue) { + return weightedAverage(secondValue, secondWeight, firstValue, firstWeight); + } + double average = (firstValue * firstWeight + secondValue * secondWeight) / (firstWeight + secondWeight); + return Math.max(firstValue, Math.min(average, secondValue)); + } + + @Override + public Collection centroids() { + return toTDigest().centroids(); + } + + @Override + public double compression() { + return _compression; + } + + @Override + public int byteSize() { + return toTDigest().byteSize(); + } + + @Override + public int smallByteSize() { + return toTDigest().smallByteSize(); + } + + @Override + public void asBytes(ByteBuffer buffer) { + toTDigest().asBytes(buffer); + } + + @Override + public void asSmallBytes(ByteBuffer buffer) { + toTDigest().asSmallBytes(buffer); + } + + @Override + public TDigest recordAllData() { + return toTDigest().recordAllData(); + } + + @Override + public boolean isRecording() { + return false; + } + + @Override + public int centroidCount() { + if (_pendingSerializedTDigest != null) { + return getSerializedCentroidCount(_pendingSerializedTDigest); + } + flush(); + normalizeCentroidCapacity(); + return _numCentroids; + } + + @Override + public double getMin() { + if (_pendingSerializedTDigest != null) { + return ByteBuffer.wrap(_pendingSerializedTDigest).getDouble(Integer.BYTES); + } + flush(); + return _min; + } + + @Override + public double getMax() { + if (_pendingSerializedTDigest != null) { + return ByteBuffer.wrap(_pendingSerializedTDigest).getDouble(Integer.BYTES + Double.BYTES); + } + flush(); + return _max; + } + + byte[] serialize() { + if (_pendingSerializedTDigest != null) { + getSerializedTotalWeight(_pendingSerializedTDigest); + return _pendingSerializedTDigest.clone(); + } + flush(); + if (requiresCapacityPreservingEncoding()) { + return toCapacityPreservingBytes(); + } + TDigest tDigest = toTDigest(); + byte[] bytes = new byte[tDigest.byteSize()]; + tDigest.asBytes(ByteBuffer.wrap(bytes)); + return bytes; + } + + TDigest toTDigest() { + if (_pendingSerializedTDigest != null) { + return MergingDigest.fromBytes(ByteBuffer.wrap(_pendingSerializedTDigest)); + } + flush(); + if (_numCentroids == 0) { + return TDigest.createMergingDigest(_compression); + } + + if (requiresCapacityPreservingEncoding()) { + return toCapacityPreservingTDigest(); + } + normalizeCentroidCapacity(); + ByteBuffer buffer = ByteBuffer.allocate(VERBOSE_HEADER_SIZE + VERBOSE_CENTROID_SIZE * _numCentroids); + buffer.putInt(VERBOSE_ENCODING); + buffer.putDouble(_min); + buffer.putDouble(_max); + buffer.putDouble(_compression); + buffer.putInt(_numCentroids); + for (int i = 0; i < _numCentroids; i++) { + buffer.putDouble(_centroidWeights[i]); + buffer.putDouble(_centroidMeans[i]); + } + buffer.flip(); + return MergingDigest.fromBytes(buffer); + } + + private boolean requiresCapacityPreservingEncoding() { + return _numCentroids > _centroidCapacity && _serializedMainCapacity >= _numCentroids + && (double) (float) _compression == _compression; + } + + private void recompressCentroids() { + double[] means = _centroidMeans; + double[] weights = _centroidWeights; + int numCentroids = _numCentroids; + double totalWeight = _totalWeight; + _numCentroids = 0; + _totalWeight = 0.0; + if (!mergeSorted(means, weights, numCentroids, totalWeight)) { + throw new IllegalStateException("Cannot recompress an empty TDigest"); + } + } + + private void normalizeCentroidCapacity() { + if (_numCentroids <= _centroidCapacity || _serializedMainCapacity < _numCentroids + || (double) (float) _compression == _compression) { + return; + } + recompressCentroids(); + if (_numCentroids > _centroidCapacity) { + throw new IllegalStateException("Unable to recompress TDigest to its default centroid capacity: " + + _numCentroids); + } + } + + private TDigest toCapacityPreservingTDigest() { + return MergingDigest.fromBytes(ByteBuffer.wrap(toCapacityPreservingBytes())); + } + + private byte[] toCapacityPreservingBytes() { + if (_numCentroids > Short.MAX_VALUE) { + throw new IllegalStateException("TDigest has too many centroids for capacity-preserving encoding: " + + _numCentroids); + } + int mainCapacity = Math.min(Short.MAX_VALUE, Math.max(_numCentroids, _serializedMainCapacity)); + long defaultBufferCapacity = Math.multiplyExact(DEFAULT_MERGE_BUFFER_MULTIPLIER, (long) Math.ceil(_compression)); + int bufferCapacity = Math.toIntExact(Math.min(Short.MAX_VALUE, + Math.max(Math.max((long) _serializedBufferCapacity, mainCapacity + 1L), defaultBufferCapacity))); + ByteBuffer buffer = ByteBuffer.allocate(SMALL_HEADER_SIZE + SMALL_CENTROID_SIZE * _numCentroids); + buffer.putInt(SMALL_ENCODING); + buffer.putDouble(_min); + buffer.putDouble(_max); + buffer.putFloat((float) _compression); + buffer.putShort((short) mainCapacity); + buffer.putShort((short) bufferCapacity); + buffer.putShort((short) _numCentroids); + for (int i = 0; i < _numCentroids; i++) { + buffer.putFloat((float) _centroidWeights[i]); + buffer.putFloat((float) _centroidMeans[i]); + } + return buffer.array(); + } + + private void flush() { + if (_numRawValues == 0) { + return; + } + + Arrays.sort(_rawValues, 0, _numRawValues); + int preferredOutputCapacity = getPreferredOutputCapacity(_numRawValues); + ensureOutputCapacity(1); + double newTotalWeight = _totalWeight + _numRawValues; + double normalizer = _compression / (Math.PI * newTotalWeight); + int rawIndex = 0; + int centroidIndex = 0; + int numOutputCentroids = 0; + double weightSoFar = 0.0; + + while (rawIndex < _numRawValues || centroidIndex < _numCentroids) { + double mean; + double weight; + if (centroidIndex == _numCentroids || (rawIndex < _numRawValues + && _rawValues[rawIndex] <= _centroidMeans[centroidIndex])) { + mean = _rawValues[rawIndex++]; + weight = 1.0; + } else { + mean = _centroidMeans[centroidIndex]; + weight = _centroidWeights[centroidIndex++]; + } + if (numOutputCentroids == 0) { + _outputMeans[0] = mean; + _outputWeights[0] = weight; + numOutputCentroids = 1; + } else { + int currentIndex = numOutputCentroids - 1; + double proposedWeight = _outputWeights[currentIndex] + weight; + double z = proposedWeight * normalizer; + double q0 = weightSoFar / newTotalWeight; + double q2 = (weightSoFar + proposedWeight) / newTotalWeight; + if (z * z <= q0 * (1.0 - q0) && z * z <= q2 * (1.0 - q2)) { + _outputWeights[currentIndex] = proposedWeight; + _outputMeans[currentIndex] += (mean - _outputMeans[currentIndex]) * weight / proposedWeight; + } else { + weightSoFar += _outputWeights[currentIndex]; + ensureOutputCapacity(Math.max(preferredOutputCapacity, numOutputCentroids + 1)); + _outputMeans[numOutputCentroids] = mean; + _outputWeights[numOutputCentroids] = weight; + numOutputCentroids++; + } + } + } + finishMerge(numOutputCentroids, newTotalWeight); + _numRawValues = 0; + _min = Math.min(_min, _centroidMeans[0]); + _max = Math.max(_max, _centroidMeans[_numCentroids - 1]); + } + + private boolean mergeSorted(double[] incomingMeans, double[] incomingWeights, int incomingCount) { + return mergeSorted(incomingMeans, incomingWeights, incomingCount, + getTotalWeight(incomingWeights, incomingCount)); + } + + private boolean mergeSorted(double[] incomingMeans, double[] incomingWeights, int incomingCount, + double incomingWeight) { + if (incomingWeight == 0.0) { + return false; + } + + int preferredOutputCapacity = getPreferredOutputCapacity(incomingCount); + ensureOutputCapacity(1); + double newTotalWeight = _totalWeight + incomingWeight; + double normalizer = _compression / (Math.PI * newTotalWeight); + int incomingIndex = 0; + int centroidIndex = 0; + int numOutputCentroids = 0; + double weightSoFar = 0.0; + while (incomingIndex < incomingCount || centroidIndex < _numCentroids) { + double mean; + double weight; + if (centroidIndex == _numCentroids || (incomingIndex < incomingCount + && incomingMeans[incomingIndex] <= _centroidMeans[centroidIndex])) { + mean = incomingMeans[incomingIndex]; + weight = incomingWeights[incomingIndex++]; + } else { + mean = _centroidMeans[centroidIndex]; + weight = _centroidWeights[centroidIndex++]; + } + if (numOutputCentroids == 0) { + _outputMeans[0] = mean; + _outputWeights[0] = weight; + numOutputCentroids = 1; + } else { + int currentIndex = numOutputCentroids - 1; + double proposedWeight = _outputWeights[currentIndex] + weight; + double z = proposedWeight * normalizer; + double q0 = weightSoFar / newTotalWeight; + double q2 = (weightSoFar + proposedWeight) / newTotalWeight; + if (z * z <= q0 * (1.0 - q0) && z * z <= q2 * (1.0 - q2)) { + _outputWeights[currentIndex] = proposedWeight; + _outputMeans[currentIndex] += (mean - _outputMeans[currentIndex]) * weight / proposedWeight; + } else { + weightSoFar += _outputWeights[currentIndex]; + ensureOutputCapacity(Math.max(preferredOutputCapacity, numOutputCentroids + 1)); + _outputMeans[numOutputCentroids] = mean; + _outputWeights[numOutputCentroids] = weight; + numOutputCentroids++; + } + } + } + finishMerge(numOutputCentroids, newTotalWeight); + return true; + } + + private void finishMerge(int numOutputCentroids, double totalWeight) { + double[] temporary = _centroidMeans; + _centroidMeans = _outputMeans; + _outputMeans = temporary; + temporary = _centroidWeights; + _centroidWeights = _outputWeights; + _outputWeights = temporary; + _numCentroids = numOutputCentroids; + _totalWeight = totalWeight; + } + + private void ensureIncomingCapacity(int capacity) { + if (_incomingMeans == null || capacity > _incomingMeans.length) { + int newCapacity = Math.max(capacity, _centroidCapacity); + _incomingMeans = new double[newCapacity]; + _incomingWeights = new double[newCapacity]; + } + } + + private void materializePendingSerializedTDigest() { + if (_pendingSerializedTDigest != null) { + byte[] bytes = _pendingSerializedTDigest; + _pendingSerializedTDigest = null; + _pendingSerializedTotalWeight = Double.NaN; + mergeSerializedTDigest(bytes); + } + } + + private void ensureCentroidCapacity(int capacity) { + if (capacity > 0 && (_centroidMeans == null || capacity > _centroidMeans.length)) { + _centroidMeans = new double[capacity]; + _centroidWeights = new double[capacity]; + } + } + + private void ensureRawBuffer() { + if (_rawValues == null) { + _rawValues = new double[getRawBufferSize((long) Math.ceil(_compression))]; + } + } + + private static int getRawBufferSize(long roundedCompression) { + return (int) Math.min(MAX_RAW_BUFFER_SIZE, Math.max(MIN_RAW_BUFFER_SIZE, 2L * roundedCompression)); + } + + private static int getCentroidCapacity(double compression) { + long roundedCompression = (long) Math.ceil(compression); + return Math.toIntExact( + Math.addExact(Math.multiplyExact(2L, roundedCompression), CENTROID_CAPACITY_PADDING)); + } + + private void preserveSerializedCapacity(int mainCapacity, int bufferCapacity) { + if (mainCapacity > 0) { + _serializedMainCapacity = Math.max(_serializedMainCapacity, mainCapacity); + } + if (bufferCapacity > 0) { + _serializedBufferCapacity = Math.max(_serializedBufferCapacity, bufferCapacity); + } + } + + private int getPreferredOutputCapacity(int incomingCount) { + return Math.min(_centroidCapacity, Math.addExact(_numCentroids, incomingCount)); + } + + private void ensureOutputCapacity(int capacity) { + if (_outputMeans == null) { + int newCapacity = Math.max(capacity, _centroidCapacity); + _outputMeans = new double[newCapacity]; + _outputWeights = new double[newCapacity]; + } else if (capacity > _outputMeans.length) { + int newCapacity = Math.max(capacity, Math.multiplyExact(_outputMeans.length, 2)); + _outputMeans = Arrays.copyOf(_outputMeans, newCapacity); + _outputWeights = Arrays.copyOf(_outputWeights, newCapacity); + } + } + + private static double getTotalWeight(double[] weights, int count) { + double totalWeight = 0.0; + for (int i = 0; i < count; i++) { + totalWeight += weights[i]; + } + return totalWeight; + } + + private static double readCompression(byte[] bytes) { + ByteBuffer input = ByteBuffer.wrap(bytes); + int encoding = input.getInt(); + if (encoding == VERBOSE_ENCODING) { + input.getDouble(); + input.getDouble(); + return input.getDouble(); + } + if (encoding == SMALL_ENCODING) { + input.getDouble(); + input.getDouble(); + return input.getFloat(); + } + throw new IllegalStateException("Invalid format for serialized histogram"); + } + + private static int getSerializedCentroidCount(byte[] bytes) { + ByteBuffer input = ByteBuffer.wrap(bytes); + int encoding = input.getInt(); + input.getDouble(); + input.getDouble(); + int numCentroids; + if (encoding == VERBOSE_ENCODING) { + double compression = input.getDouble(); + numCentroids = input.getInt(); + checkCentroidCount(numCentroids); + checkCentroidsAvailable(input, numCentroids, VERBOSE_CENTROID_SIZE); + checkVerboseCentroidCapacity(compression, numCentroids); + } else if (encoding == SMALL_ENCODING) { + double compression = input.getFloat(); + int mainCapacity = input.getShort(); + int bufferCapacity = input.getShort(); + checkSmallArrayCapacities(mainCapacity, bufferCapacity); + mainCapacity = resolveSmallMainCapacity(compression, mainCapacity); + numCentroids = input.getShort(); + checkCentroidCount(numCentroids); + checkCentroidsAvailable(input, numCentroids, SMALL_CENTROID_SIZE); + checkSmallCentroidCapacity(mainCapacity, numCentroids); + } else { + throw new IllegalStateException("Invalid format for serialized histogram"); + } + return numCentroids; + } + + private static double getSerializedTotalWeight(byte[] bytes) { + ByteBuffer input = ByteBuffer.wrap(bytes); + int encoding = input.getInt(); + input.getDouble(); + input.getDouble(); + int numCentroids; + int centroidSize; + if (encoding == VERBOSE_ENCODING) { + double compression = input.getDouble(); + numCentroids = input.getInt(); + centroidSize = VERBOSE_CENTROID_SIZE; + checkCentroidCount(numCentroids); + checkCentroidsAvailable(input, numCentroids, centroidSize); + checkVerboseCentroidCapacity(compression, numCentroids); + } else if (encoding == SMALL_ENCODING) { + double compression = input.getFloat(); + int mainCapacity = input.getShort(); + int bufferCapacity = input.getShort(); + checkSmallArrayCapacities(mainCapacity, bufferCapacity); + mainCapacity = resolveSmallMainCapacity(compression, mainCapacity); + numCentroids = input.getShort(); + centroidSize = SMALL_CENTROID_SIZE; + checkCentroidCount(numCentroids); + checkCentroidsAvailable(input, numCentroids, centroidSize); + checkSmallCentroidCapacity(mainCapacity, numCentroids); + } else { + throw new IllegalStateException("Invalid format for serialized histogram"); + } + + double totalWeight = 0.0; + for (int i = 0; i < numCentroids; i++) { + totalWeight += centroidSize == VERBOSE_CENTROID_SIZE ? input.getDouble() : input.getFloat(); + input.position(input.position() + centroidSize / 2); + } + return totalWeight; + } + + /// Reusable, invocation-local view of one serialized TDigest for group-by-MV fanout. + /// + /// [#reset(byte\[\])] parses and validates the header once per input row. Centroid arrays are decoded lazily and + /// reused across rows, so all groups for a row merge the same primitive input without sharing mutable accumulator + /// state. The input must remain thread-confined and must not outlive the aggregation call that owns it. + static final class SerializedTDigestInput { + private byte[] _bytes; + private double _min; + private double _max; + private double _compression; + private int _numCentroids; + private int _mainCapacity; + private int _bufferCapacity; + private int _centroidOffset; + private int _centroidSize; + private double[] _means; + private double[] _weights; + private double _totalWeight; + private boolean _decoded; + + void reset(byte[] bytes) { + ByteBuffer input = ByteBuffer.wrap(bytes); + int encoding = input.getInt(); + if (encoding != VERBOSE_ENCODING && encoding != SMALL_ENCODING) { + throw new IllegalStateException("Invalid format for serialized histogram"); + } + _bytes = bytes; + _min = input.getDouble(); + _max = input.getDouble(); + if (encoding == VERBOSE_ENCODING) { + _compression = input.getDouble(); + _numCentroids = input.getInt(); + _mainCapacity = 0; + _bufferCapacity = 0; + _centroidSize = VERBOSE_CENTROID_SIZE; + } else { + _compression = input.getFloat(); + _mainCapacity = input.getShort(); + _bufferCapacity = input.getShort(); + checkSmallArrayCapacities(_mainCapacity, _bufferCapacity); + _mainCapacity = resolveSmallMainCapacity(_compression, _mainCapacity); + _numCentroids = input.getShort(); + _centroidSize = SMALL_CENTROID_SIZE; + } + checkCentroidCount(_numCentroids); + checkCentroidsAvailable(input, _numCentroids, _centroidSize); + if (_centroidSize == VERBOSE_CENTROID_SIZE) { + checkVerboseCentroidCapacity(_compression, _numCentroids); + } else { + checkSmallCentroidCapacity(_mainCapacity, _numCentroids); + } + _centroidOffset = input.position(); + _decoded = false; + } + + private void decode() { + if (_decoded) { + return; + } + ensureCapacity(_numCentroids); + ByteBuffer input = ByteBuffer.wrap(_bytes); + input.position(_centroidOffset); + double totalWeight = 0.0; + if (_centroidSize == VERBOSE_CENTROID_SIZE) { + for (int i = 0; i < _numCentroids; i++) { + double weight = input.getDouble(); + _weights[i] = weight; + _means[i] = input.getDouble(); + totalWeight += weight; + } + } else { + for (int i = 0; i < _numCentroids; i++) { + double weight = input.getFloat(); + _weights[i] = weight; + _means[i] = input.getFloat(); + totalWeight += weight; + } + } + _totalWeight = totalWeight; + _decoded = true; + } + + private void ensureCapacity(int capacity) { + if (capacity > 0 && (_means == null || capacity > _means.length)) { + int newCapacity = _means == null ? capacity : Math.max(capacity, Math.multiplyExact(_means.length, 2)); + _means = new double[newCapacity]; + _weights = new double[newCapacity]; + } + } + } + + private static void checkCentroidCount(int numCentroids) { + if (numCentroids < 0) { + throw new IllegalArgumentException("Invalid negative TDigest centroid count: " + numCentroids); + } + } + + private static void checkCentroidsAvailable(ByteBuffer input, int numCentroids, int centroidSize) { + if (numCentroids > input.remaining() / centroidSize) { + throw new BufferUnderflowException(); + } + } + + private static void checkSmallArrayCapacities(int mainCapacity, int bufferCapacity) { + if (mainCapacity < -1) { + throw new NegativeArraySizeException(Integer.toString(mainCapacity)); + } + if (bufferCapacity < -1) { + throw new NegativeArraySizeException(Integer.toString(bufferCapacity)); + } + } + + private static int resolveSmallMainCapacity(double compression, int mainCapacity) { + if (mainCapacity != -1) { + return mainCapacity; + } + return getSerializedDefaultCentroidCapacity(compression); + } + + private static void checkVerboseCentroidCapacity(double compression, int numCentroids) { + checkSmallCentroidCapacity(getSerializedDefaultCentroidCapacity(compression), numCentroids); + } + + private static int getSerializedDefaultCentroidCapacity(double compression) { + int capacity = (int) (2.0 * Math.ceil(compression)); + capacity += CENTROID_CAPACITY_PADDING; + if (capacity < 0) { + throw new NegativeArraySizeException(Integer.toString(capacity)); + } + return capacity; + } + + private static void checkSmallCentroidCapacity(int mainCapacity, int numCentroids) { + if (numCentroids > mainCapacity) { + throw new ArrayIndexOutOfBoundsException( + "Index " + mainCapacity + " out of bounds for length " + mainCapacity); + } + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAggregationFunction.java index 799fdc750169..8fd195c6c89f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAggregationFunction.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. * diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java index 529f2b15f706..78d7610b6f6e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DefaultGroupByExecutor.java @@ -102,7 +102,7 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a int maxInitialResultHolderCapacity = queryContext.getMaxInitialResultHolderCapacity(); Map groupByExpressionSizesFromPredicates = null; if (queryContext.isOptimizeMaxInitialResultHolderCapacity()) { - groupByExpressionSizesFromPredicates = getGroupByExpressionSizesFromPredicates(queryContext); + groupByExpressionSizesFromPredicates = getGroupByExpressionSizesFromPredicates(queryContext, projectOperator); } if (groupKeyGenerator != null) { _groupKeyGenerator = groupKeyGenerator; @@ -152,7 +152,8 @@ public DefaultGroupByExecutor(QueryContext queryContext, AggregationFunction[] a * 2. Ensure the top-level filter context consists solely of AND-type filters; other types for example OR we cannot * guarantee deterministic sizes for GroupBy expressions. */ - private Map getGroupByExpressionSizesFromPredicates(QueryContext queryContext) { + private Map getGroupByExpressionSizesFromPredicates(QueryContext queryContext, + BaseProjectOperator projectOperator) { FilterContext filterContext = queryContext.getFilter(); if (filterContext == null || queryContext.getGroupByExpressions() == null) { return null; @@ -188,6 +189,8 @@ private Map getGroupByExpressionSizesFromPredicates( // Populate the group-by expressions with sizes from the predicate map return queryContext.getGroupByExpressions().stream() .filter(predicateSizeMap::containsKey) + // An MV predicate filters rows, but each matching row can still produce other group values. + .filter(expression -> projectOperator.getResultColumnContext(expression).isSingleValue()) .collect(Collectors.toMap( expression -> expression, expression -> predicateSizeMap.getOrDefault(expression, null) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java index cfe70b785bba..b4a734c86e6c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGenerator.java @@ -118,8 +118,8 @@ public DictionaryBasedGroupKeyGenerator(BaseProjectOperator projectOperator, // only one call will be made to the dictionary to extract each raw value. _internedDictionaryValues = _numGroupByExpressions > 1 ? new Object[_numGroupByExpressions][] : null; Map cardinalityMap = new HashMap<>(_numGroupByExpressions); - long cardinalityProduct = 1L; - boolean longOverflow = false; + long rawKeyCardinalityProduct = 1L; + boolean rawKeyLongOverflow = false; for (int i = 0; i < _numGroupByExpressions; i++) { ExpressionContext groupByExpression = groupByExpressions[i]; ColumnContext columnContext = projectOperator.getResultColumnContext(groupByExpression); @@ -131,44 +131,49 @@ public DictionaryBasedGroupKeyGenerator(BaseProjectOperator projectOperator, if (_internedDictionaryValues != null && cardinality < MAX_DICTIONARY_INTERN_TABLE_SIZE) { _internedDictionaryValues[i] = new Object[cardinality]; } - if (!longOverflow) { - if (cardinalityProduct > Long.MAX_VALUE / cardinality) { - longOverflow = true; + if (!rawKeyLongOverflow) { + if (rawKeyCardinalityProduct > Long.MAX_VALUE / cardinality) { + rawKeyLongOverflow = true; } else { - cardinalityProduct *= cardinality; + rawKeyCardinalityProduct *= cardinality; } } _isSingleValueColumn[i] = columnContext.isSingleValue(); } + Long optimizedCardinality = null; if (groupByExpressionSizesFromPredicates != null) { - Pair optimizedCardinality = getOptimizedGroupByCardinality(groupByExpressionSizesFromPredicates, - cardinalityMap); - if (optimizedCardinality.getLeft() && optimizedCardinality.getRight() != null) { - longOverflow = false; - cardinalityProduct = Math.min(optimizedCardinality.getRight(), cardinalityProduct); + Pair optimizedCardinalityResult = + getOptimizedGroupByCardinality(groupByExpressionSizesFromPredicates, cardinalityMap); + if (optimizedCardinalityResult.getLeft()) { + optimizedCardinality = optimizedCardinalityResult.getRight(); } } + if (optimizedCardinality != null) { + _globalGroupIdUpperBound = (int) Math.min(optimizedCardinality, numGroupsLimit); + } else if (rawKeyLongOverflow || rawKeyCardinalityProduct > Integer.MAX_VALUE) { + _globalGroupIdUpperBound = numGroupsLimit; + } else { + _globalGroupIdUpperBound = Math.min((int) rawKeyCardinalityProduct, numGroupsLimit); + } // NOTE: We need to clean up the thread-local map before using it in case RawKeyHolder.close() is not called // for the previous segment // TODO: Ensure RawKeyHolder.close() - if (longOverflow) { + if (rawKeyLongOverflow) { // ArrayMapBasedHolder - _globalGroupIdUpperBound = numGroupsLimit; Object2IntOpenHashMap groupIdMap = THREAD_LOCAL_INT_ARRAY_MAP.get(); clearAndTrim(groupIdMap); _rawKeyHolder = new ArrayMapBasedHolder(groupIdMap); } else { - if (cardinalityProduct > Integer.MAX_VALUE) { + if (rawKeyCardinalityProduct > Integer.MAX_VALUE) { // LongMapBasedHolder - _globalGroupIdUpperBound = numGroupsLimit; Long2IntOpenHashMap groupIdMap = THREAD_LOCAL_LONG_MAP.get(); clearAndTrim(groupIdMap); _rawKeyHolder = new LongMapBasedHolder(groupIdMap); } else { - _globalGroupIdUpperBound = Math.min((int) cardinalityProduct, numGroupsLimit); - // arrayBaseHolder fails with ArrayIndexOutOfBoundsException if numGroupsLimit < cardinalityProduct - // because array doesn't fit all (potentially unsorted) values - if (cardinalityProduct > arrayBasedThreshold || numGroupsLimit < cardinalityProduct) { + // ArrayBasedHolder uses raw dictionary ids as group ids, so it requires an array covering the full raw-key + // domain. Predicate cardinality can reduce the number of groups without reducing the raw dictionary ids. + if (rawKeyCardinalityProduct > arrayBasedThreshold + || _globalGroupIdUpperBound < rawKeyCardinalityProduct) { // IntMapBasedHolder IntGroupIdMap groupIdMap = THREAD_LOCAL_INT_MAP.get(); groupIdMap.clearAndTrim(); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/common/ObjectSerDeUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/common/ObjectSerDeUtilsTest.java index d3d13e1680a3..5dd7c023ab84 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/common/ObjectSerDeUtilsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/common/ObjectSerDeUtilsTest.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++) { diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/table/IndexedTableTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/table/IndexedTableTest.java index 17dea94d32bb..f464999e022e 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/table/IndexedTableTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/table/IndexedTableTest.java @@ -18,19 +18,30 @@ */ package org.apache.pinot.core.data.table; +import com.tdunning.math.stats.Centroid; +import com.tdunning.math.stats.TDigest; +import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import org.apache.pinot.common.CustomObject; +import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.common.utils.DataSchema; import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.SyntheticBlockValSets; +import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.function.AggregationFunction.SerializedIntermediateResult; +import org.apache.pinot.core.query.aggregation.function.PercentileTDigestAggregationFunction; import org.apache.pinot.core.query.request.context.QueryContext; import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; import org.apache.pinot.spi.utils.CommonConstants.Server; @@ -120,6 +131,101 @@ public void testConcurrentIndexedTable() } } + @DataProvider(name = "concurrentPercentileTDigestTables") + public static Object[][] concurrentPercentileTDigestTables() { + return new Object[][]{{false, false}, {false, true}, {true, false}, {true, true}}; + } + + @Test(dataProvider = "concurrentPercentileTDigestTables", timeOut = 30_000) + public void testConcurrentPercentileTDigestUpdatesForSameKey(boolean deterministic, boolean serialized) + throws InterruptedException, ExecutionException, TimeoutException { + int numThreads = 8; + int segmentsPerThread = 32; + int valuesPerSegment = 64; + long expectedSize = (long) numThreads * segmentsPerThread * valuesPerSegment; + ExpressionContext metricExpression = ExpressionContext.forIdentifier("m1"); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext( + "SELECT PERCENTILETDIGEST(m1, 75) FROM testTable GROUP BY d1"); + DataSchema dataSchema = new DataSchema(new String[]{"d1", "percentileTDigest(m1, 75)"}, + new ColumnDataType[]{ColumnDataType.STRING, ColumnDataType.OBJECT}); + ExecutorService tableExecutor = Executors.newCachedThreadPool(); + ExecutorService writers = Executors.newFixedThreadPool(numThreads); + CountDownLatch ready = new CountDownLatch(numThreads); + CountDownLatch start = new CountDownLatch(1); + try { + IndexedTable indexedTable = deterministic + ? new DeterministicConcurrentIndexedTable(dataSchema, false, queryContext, 1, + Integer.MAX_VALUE, Integer.MAX_VALUE, INITIAL_CAPACITY, tableExecutor) + : new ConcurrentIndexedTable(dataSchema, false, queryContext, 1, + Integer.MAX_VALUE, Integer.MAX_VALUE, INITIAL_CAPACITY, tableExecutor); + List> futures = new ArrayList<>(numThreads); + for (int threadIndex = 0; threadIndex < numThreads; threadIndex++) { + int currentThread = threadIndex; + futures.add(writers.submit(() -> { + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(metricExpression, 75.0, false); + ready.countDown(); + start.await(); + for (int segmentIndex = 0; segmentIndex < segmentsPerThread; segmentIndex++) { + double[] values = new double[valuesPerSegment]; + long firstValue = ((long) currentThread * segmentsPerThread + segmentIndex) * valuesPerSegment; + for (int valueIndex = 0; valueIndex < valuesPerSegment; valueIndex++) { + values[valueIndex] = firstValue + valueIndex; + } + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(valuesPerSegment, resultHolder, + Map.of(metricExpression, SyntheticBlockValSets.Double.create(null, values))); + TDigest digest = function.extractAggregationResult(resultHolder); + if (serialized) { + SerializedIntermediateResult serializedResult = function.serializeIntermediateResult(digest); + digest = function.deserializeIntermediateResult(new CustomObject(serializedResult.getType(), + ByteBuffer.wrap(serializedResult.getBytes()))); + } + indexedTable.upsert(new Key(new Object[]{"same-key"}), + new Record(new Object[]{"same-key", digest})); + } + return null; + })); + } + Assert.assertTrue(ready.await(10, TimeUnit.SECONDS), "Writer threads did not become ready"); + start.countDown(); + for (Future future : futures) { + future.get(10, TimeUnit.SECONDS); + } + + indexedTable.finish(false); + Assert.assertEquals(indexedTable.size(), 1); + TDigest result = (TDigest) indexedTable.iterator().next().getValues()[1]; + Assert.assertEquals(result.size(), expectedSize); + Assert.assertEquals(result.getMin(), 0.0); + Assert.assertEquals(result.getMax(), expectedSize - 1.0); + + long centroidWeight = 0L; + double previousMean = Double.NEGATIVE_INFINITY; + for (Centroid centroid : result.centroids()) { + Assert.assertTrue(Double.isFinite(centroid.mean())); + Assert.assertTrue(centroid.count() > 0); + Assert.assertTrue(centroid.mean() + 1e-12 >= previousMean); + centroidWeight += centroid.count(); + previousMean = centroid.mean(); + } + Assert.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 = result.quantile(quantile); + Assert.assertTrue(Double.isFinite(value)); + Assert.assertTrue(value >= previousQuantile); + Assert.assertEquals(value, quantile * (expectedSize - 1), expectedSize * 0.02); + previousQuantile = value; + } + } finally { + start.countDown(); + writers.shutdownNow(); + tableExecutor.shutdownNow(); + } + } + @Test(dataProvider = "initDataProvider") public void testNonConcurrentIndexedTable(String orderBy, List survivors) { QueryContext queryContext = QueryContextConverterUtils.getQueryContext( diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunctionTest.java index 3c2ecdde0112..b5257bbb5ede 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileAggregationFunctionTest.java @@ -19,9 +19,138 @@ package org.apache.pinot.core.query.aggregation.function; +import it.unimi.dsi.fastutil.doubles.DoubleArrayList; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.SplittableRandom; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.testng.Assert; +import org.testng.annotations.Test; + public class PercentileAggregationFunctionTest extends AbstractPercentileAggregationFunctionTest { + private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("col"); + private static final double[] PERCENTILES = {0.0, 0.01, 1.0, 25.0, 50.0, 75.0, 99.0, 99.999, 100.0}; + @Override public String callStr(String column, int percent) { return "PERCENTILE(" + column + ", " + percent + ")"; } + + @Test + public void testSelectionMatchesFullSortForRandomValues() { + SplittableRandom random = new SplittableRandom(42); + int[] sizes = {1, 2, 3, 15, 31, 32, 33, 63, 127, 128, 129, 1024, 10_001}; + for (int size : sizes) { + double[] values = new double[size]; + for (int i = 0; i < size; i++) { + values[i] = random.nextDouble(-1_000_000.0, 1_000_000.0); + } + assertPercentiles(values, PERCENTILES); + } + } + + @Test + public void testSelectionMatchesFullSortForAdversarialValues() { + int size = 4097; + double[] ascending = new double[size]; + double[] descending = new double[size]; + double[] duplicates = new double[size]; + double[] equal = new double[size]; + double[] organPipe = new double[size]; + Arrays.fill(equal, 7.0); + for (int i = 0; i < size; i++) { + ascending[i] = i; + descending[i] = size - i; + duplicates[i] = i % 7; + organPipe[i] = Math.min(i, size - 1 - i); + } + + List inputs = new ArrayList<>(); + inputs.add(ascending); + inputs.add(descending); + inputs.add(duplicates); + inputs.add(equal); + inputs.add(organPipe); + double[] specialValues = {Double.NEGATIVE_INFINITY, -Double.MAX_VALUE, -1.0, -0.0, 0.0, Double.MIN_VALUE, + Double.MAX_VALUE, Double.POSITIVE_INFINITY, Double.NaN, Double.longBitsToDouble(0x7ff0000000000001L), + Double.NaN}; + double[] repeatedSpecialValues = new double[257]; + for (int i = 0; i < repeatedSpecialValues.length; i++) { + repeatedSpecialValues[i] = specialValues[i % specialValues.length]; + } + inputs.add(repeatedSpecialValues); + for (double[] values : inputs) { + assertPercentiles(values, PERCENTILES); + } + } + + @Test + public void testSelectionUsesLogicalListSize() { + double[] logicalValues = new double[257]; + for (int i = 0; i < logicalValues.length; i++) { + logicalValues[i] = (i * 37) % logicalValues.length; + } + DoubleArrayList values = new DoubleArrayList(512); + for (double value : logicalValues) { + values.add(value); + } + Arrays.fill(values.elements(), values.size(), values.elements().length, Double.NaN); + + assertResult(values, logicalValues, 0.0); + assertResult(values, logicalValues, 75.0); + assertResult(values, logicalValues, 100.0); + } + + @Test + public void testRepeatedExtraction() { + SplittableRandom random = new SplittableRandom(42); + double[] source = new double[10_001]; + for (int i = 0; i < source.length; i++) { + source[i] = random.nextDouble(); + } + DoubleArrayList values = DoubleArrayList.wrap(source.clone()); + for (double percentile : new double[]{75.0, 25.0, 99.0, 50.0, 75.0}) { + assertResult(values, source, percentile); + } + } + + @Test + public void testMergeAfterExtraction() { + SplittableRandom random = new SplittableRandom(42); + double[] first = new double[10_001]; + double[] second = new double[4097]; + for (int i = 0; i < first.length; i++) { + first[i] = random.nextDouble(); + } + for (int i = 0; i < second.length; i++) { + second[i] = random.nextDouble(); + } + + PercentileAggregationFunction function = new PercentileAggregationFunction(EXPRESSION, 75.0, false); + DoubleArrayList values = DoubleArrayList.wrap(first.clone()); + assertResult(values, first, 75.0); + function.merge(values, DoubleArrayList.wrap(second.clone())); + double[] combined = Arrays.copyOf(first, first.length + second.length); + System.arraycopy(second, 0, combined, first.length, second.length); + assertResult(values, combined, 75.0); + } + + private static void assertPercentiles(double[] values, double[] percentiles) { + for (double percentile : percentiles) { + assertResult(DoubleArrayList.wrap(values.clone()), values, percentile); + } + } + + private static void assertResult(DoubleArrayList actualValues, double[] sourceValues, double percentile) { + double[] sortedValues = sourceValues.clone(); + Arrays.sort(sortedValues); + int index = percentile == 100.0 ? sortedValues.length - 1 + : (int) ((long) sortedValues.length * percentile / 100); + double expected = sortedValues[index]; + PercentileAggregationFunction function = new PercentileAggregationFunction(EXPRESSION, percentile, false); + double actual = function.extractFinalResult(actualValues); + Assert.assertEquals(Double.doubleToLongBits(actual), Double.doubleToLongBits(expected), + "Unexpected result for percentile " + percentile + " and size " + sourceValues.length); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAggregationFunctionTest.java index 5df07f7ccc50..d94ec545cf85 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAggregationFunctionTest.java @@ -18,26 +18,786 @@ */ package org.apache.pinot.core.query.aggregation.function; +import com.tdunning.math.stats.Centroid; +import com.tdunning.math.stats.MergingDigest; +import com.tdunning.math.stats.TDigest; +import java.lang.reflect.Field; +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; import java.util.Map; +import java.util.SplittableRandom; +import java.util.function.IntPredicate; +import org.apache.pinot.common.CustomObject; import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.ObjectSerDeUtils; +import org.apache.pinot.core.common.SyntheticBlockValSets; +import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.function.AggregationFunction.SerializedIntermediateResult; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; import org.apache.pinot.segment.spi.Constants; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.roaringbitmap.RoaringBitmap; import org.testng.Assert; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; public class PercentileTDigestAggregationFunctionTest { + private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("col"); + + @DataProvider(name = "rowCounts") + public static Object[][] rowCounts() { + return new Object[][]{{0}, {1}, {255}, {256}, {257}, {10_003}}; + } + + @Test(dataProvider = "rowCounts") + public void testRawNumericAggregation(int numRows) { + SplittableRandom random = new SplittableRandom(42); + double[] values = new double[numRows]; + for (int i = 0; i < numRows; i++) { + values[i] = random.nextDouble(); + } + + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(numRows, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, values))); + TDigest result = function.extractAggregationResult(resultHolder); + + Assert.assertTrue(result instanceof PercentileTDigestAccumulator); + Assert.assertEquals(result.size(), numRows); + Assert.assertEquals(result.compression(), PercentileTDigestAggregationFunction.DEFAULT_TDIGEST_COMPRESSION); + if (numRows == 0) { + Assert.assertTrue(Double.isNaN(function.extractFinalResult(result))); + } else { + double[] sortedValues = values.clone(); + Arrays.sort(sortedValues); + double expected = sortedValues[(int) ((long) numRows * 75 / 100)]; + Assert.assertEquals(function.extractFinalResult(result), expected, 0.01); + } + + TDigest deserialized = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize( + ObjectSerDeUtils.TDIGEST_SER_DE.serialize(result)); + Assert.assertTrue(deserialized instanceof MergingDigest); + Assert.assertEquals(deserialized.size(), numRows); + Assert.assertEquals(deserialized.quantile(0.75), result.quantile(0.75)); + } + + @Test + public void testRawNumericAggregationWithNulls() { + int numRows = 1_024; + double[] values = new double[numRows]; + for (int i = 0; i < numRows; i++) { + values[i] = i; + } + RoaringBitmap nullBitmap = RoaringBitmap.bitmapOf(0, 1, 254, 255, 256, 511, 512, 1_023); + + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, true); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(numRows, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(nullBitmap, values))); + TDigest result = function.extractAggregationResult(resultHolder); + + Assert.assertEquals(result.size(), numRows - nullBitmap.getCardinality()); + double[] nonNullValues = new double[numRows - nullBitmap.getCardinality()]; + int index = 0; + for (int i = 0; i < numRows; i++) { + if (!nullBitmap.contains(i)) { + nonNullValues[index++] = values[i]; + } + } + Assert.assertEquals(function.extractFinalResult(result), nonNullValues[nonNullValues.length * 75 / 100], 2.0); + + RoaringBitmap allNull = new RoaringBitmap(); + allNull.add(0L, numRows); + resultHolder = function.createAggregationResultHolder(); + function.aggregate(numRows, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(allNull, values))); + Assert.assertEquals(function.extractAggregationResult(resultHolder).size(), 0L); + } + + @Test + public void testRawMultiValueAggregation() { + double[][] values = {{0.1, 0.2}, {}, {0.3}, {0.7, 0.8, 0.9}, {1.0}}; + RoaringBitmap nullBitmap = RoaringBitmap.bitmapOf(2); + BlockValSet blockValSet = new SyntheticBlockValSets.Base() { + @Override + public RoaringBitmap getNullBitmap() { + return nullBitmap; + } + + @Override + public DataType getValueType() { + return DataType.DOUBLE; + } + + @Override + public boolean isSingleValue() { + return false; + } + + @Override + public double[][] getDoubleValuesMV() { + return values; + } + }; + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, true); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(values.length, resultHolder, Map.of(EXPRESSION, blockValSet)); + TDigest result = function.extractAggregationResult(resultHolder); + + Assert.assertEquals(result.size(), 6L); + Assert.assertEquals(function.extractFinalResult(result), 0.8, 0.1); + } + + @Test + public void testDuplicateAndSpecialValues() { + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 50.0, false); + double[] duplicateValues = new double[10_003]; + Arrays.fill(duplicateValues, 42.0); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(duplicateValues.length, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, duplicateValues))); + TDigest result = function.extractAggregationResult(resultHolder); + Assert.assertEquals(result.size(), duplicateValues.length); + Assert.assertEquals(function.extractFinalResult(result), 42.0); + Assert.assertEquals(result.getMin(), 42.0); + Assert.assertEquals(result.getMax(), 42.0); + + double[] specialValues = {Double.NEGATIVE_INFINITY, -1.0, -0.0, 0.0, 1.0, Double.POSITIVE_INFINITY}; + resultHolder = function.createAggregationResultHolder(); + function.aggregate(specialValues.length, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, specialValues))); + result = function.extractAggregationResult(resultHolder); + Assert.assertEquals(result.size(), specialValues.length); + Assert.assertEquals(result.getMin(), Double.NEGATIVE_INFINITY); + Assert.assertEquals(result.getMax(), Double.POSITIVE_INFINITY); + Assert.assertTrue(Double.isNaN(result.quantile(0.0))); + Assert.assertTrue(Double.isNaN(result.quantile(1.0))); + Assert.assertEquals(result.quantile(0.5), 0.0); + } + + @Test + public void testSerializedAggregationVerboseAndSmallEncodings() { + int numDigests = 32; + int valuesPerDigest = 200; + byte[][] verboseValues = createSerializedDigests(numDigests, valuesPerDigest, 100.0, ignored -> false); + byte[][] smallValues = createSerializedDigests(numDigests, valuesPerDigest, 100.0, ignored -> true); + + TDigest verboseResult = aggregateSerialized(verboseValues, null, false); + TDigest smallResult = aggregateSerialized(smallValues, null, false); + long expectedSize = (long) numDigests * valuesPerDigest; + double expectedPercentile = 0.75; + for (TDigest result : new TDigest[]{verboseResult, smallResult}) { + Assert.assertTrue(result instanceof PercentileTDigestAccumulator); + Assert.assertEquals(result.size(), expectedSize); + Assert.assertEquals(result.compression(), 100.0); + Assert.assertEquals(result.quantile(0.75), expectedPercentile, 0.005); + Assert.assertEquals(result.getMin(), 0.0); + Assert.assertEquals(result.getMax(), (expectedSize - 1.0) / expectedSize, 0.000_001); + Assert.assertTrue(((PercentileTDigestAccumulator) result).toTDigest() instanceof MergingDigest); + } + Assert.assertEquals(smallResult.quantile(0.75), verboseResult.quantile(0.75), 0.001); + } + + @Test + public void testSerializedAggregationWithFirstAndNonLeadingNulls() { + int numDigests = 6; + int valuesPerDigest = 20; + byte[][] values = createSerializedDigests(numDigests, valuesPerDigest, 100.0, index -> (index & 1) == 0); + RoaringBitmap nullBitmap = RoaringBitmap.bitmapOf(0, 2, 5); + for (int index : nullBitmap) { + values[index] = new byte[0]; + } + + TDigest result = aggregateSerialized(values, nullBitmap, true); + int expectedSize = (numDigests - nullBitmap.getCardinality()) * valuesPerDigest; + Assert.assertEquals(result.size(), expectedSize); + double[] expectedValues = new double[expectedSize]; + int outputIndex = 0; + for (int digestIndex = 0; digestIndex < numDigests; digestIndex++) { + if (!nullBitmap.contains(digestIndex)) { + for (int valueIndex = 0; valueIndex < valuesPerDigest; valueIndex++) { + expectedValues[outputIndex++] = serializedValue(digestIndex, valueIndex, numDigests, valuesPerDigest); + } + } + } + Arrays.sort(expectedValues); + Assert.assertEquals(result.getMin(), expectedValues[0]); + Assert.assertEquals(result.getMax(), expectedValues[expectedSize - 1]); + Assert.assertEquals(result.quantile(0.75), expectedValues[expectedSize * 75 / 100], 0.02); + } + + @Test + public void testSerializedAggregationAllNull() { + byte[][] values = {new byte[0], new byte[0], new byte[0]}; + RoaringBitmap nullBitmap = new RoaringBitmap(); + nullBitmap.add(0L, values.length); + + TDigest result = aggregateSerialized(values, nullBitmap, true); + Assert.assertEquals(result.size(), 0L); + Assert.assertEquals(result.compression(), PercentileTDigestAggregationFunction.DEFAULT_TDIGEST_COMPRESSION); + Assert.assertTrue(Double.isNaN(result.quantile(0.75))); + } + + @Test + public void testSerializedAggregationAcrossCallsAndBatchBoundaries() { + int[] batchSizes = {1, 63, 64, 65, 257}; + int numDigests = Arrays.stream(batchSizes).sum(); + int valuesPerDigest = 10; + byte[][] values = createSerializedDigests(numDigests, valuesPerDigest, 100.0, index -> (index & 1) != 0); + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + int offset = 0; + for (int batchSize : batchSizes) { + byte[][] batch = Arrays.copyOfRange(values, offset, offset + batchSize); + function.aggregate(batchSize, resultHolder, Map.of(EXPRESSION, bytesBlockValSet(batch, null))); + offset += batchSize; + } + + TDigest result = function.extractAggregationResult(resultHolder); + long expectedSize = (long) numDigests * valuesPerDigest; + Assert.assertEquals(result.size(), expectedSize); + Assert.assertEquals(result.quantile(0.75), 0.75, 0.005); + } + + @Test + public void testSerializedAggregationPreservesFirstDigestDoubleCompression() { + double compression = 42.125; + byte[][] values = createSerializedDigests(8, 100, compression, ignored -> false); + + TDigest result = aggregateSerialized(values, null, false); + Assert.assertEquals(result.size(), 800L); + Assert.assertEquals(result.compression(), compression); + Assert.assertEquals(result.quantile(0.75), 0.75, 0.01); + } + + @Test + public void testSerializedGroupBySV() { + int numDigests = 12; + int valuesPerDigest = 100; + double compression = 42.125; + byte[][] values = createSerializedDigests(numDigests, valuesPerDigest, compression, index -> (index & 1) == 0); + int[] groupKeys = {0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2}; + RoaringBitmap nullBitmap = RoaringBitmap.bitmapOf(0, 7); + for (int index : nullBitmap) { + values[index] = new byte[0]; + } + + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, true); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(4, 4); + function.aggregateGroupBySV(numDigests, groupKeys, resultHolder, + Map.of(EXPRESSION, bytesBlockValSet(values, nullBitmap))); + function.aggregateGroupBySV(numDigests, groupKeys, resultHolder, + Map.of(EXPRESSION, bytesBlockValSet(values, nullBitmap))); + + assertSerializedAccumulatorState(resultHolder, 0); + assertSerializedAccumulatorState(resultHolder, 1); + assertSerializedAccumulatorState(resultHolder, 2); + double denominator = numDigests * (double) valuesPerDigest; + assertGroupResult(function, resultHolder, 0, 6L * valuesPerDigest, compression, 3.0 / denominator, + (99.0 * numDigests + 9.0) / denominator); + assertGroupResult(function, resultHolder, 1, 6L * valuesPerDigest, compression, 1.0 / denominator, + (99.0 * numDigests + 10.0) / denominator); + assertGroupResult(function, resultHolder, 2, 8L * valuesPerDigest, compression, 2.0 / denominator, + (99.0 * numDigests + 11.0) / denominator); + Assert.assertEquals(function.extractGroupByResult(resultHolder, 3).size(), 0L); + } + + @Test + public void testSerializedGroupByMV() { + int numDigests = 6; + int valuesPerDigest = 100; + byte[][] values = createSerializedDigests(numDigests, valuesPerDigest, 100.0, index -> (index & 1) != 0); + int[][] groupKeys = {{0, 1}, {1}, {0, 2}, {2, 3}, {0, 3}, {1, 2, 3}}; + RoaringBitmap nullBitmap = RoaringBitmap.bitmapOf(3); + values[3] = new byte[0]; + + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, true); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(4, 4); + function.aggregateGroupByMV(numDigests, groupKeys, resultHolder, + Map.of(EXPRESSION, bytesBlockValSet(values, nullBitmap))); + + for (int groupKey = 0; groupKey < 4; groupKey++) { + assertSerializedAccumulatorState(resultHolder, groupKey); + } + double denominator = numDigests * (double) valuesPerDigest; + assertGroupResult(function, resultHolder, 0, 3L * valuesPerDigest, 100.0, 0.0, + (99.0 * numDigests + 4.0) / denominator); + assertGroupResult(function, resultHolder, 1, 3L * valuesPerDigest, 100.0, 0.0, + (99.0 * numDigests + 5.0) / denominator); + assertGroupResult(function, resultHolder, 2, 2L * valuesPerDigest, 100.0, 2.0 / denominator, + (99.0 * numDigests + 5.0) / denominator); + assertGroupResult(function, resultHolder, 3, 2L * valuesPerDigest, 100.0, 4.0 / denominator, + (99.0 * numDigests + 5.0) / denominator); + } + + @Test + public void testSerializedGroupByMVSharedInputEdgeCases() + throws ReflectiveOperationException { + TDigest empty = TDigest.createMergingDigest(20.0); + TDigest one = TDigest.createMergingDigest(100.0); + one.add(1.0); + ByteBuffer smallBuffer = ByteBuffer.allocate(one.smallByteSize()); + one.asSmallBytes(smallBuffer); + TDigest two = TDigest.createMergingDigest(100.0); + two.add(2.0); + byte[][] values = { + new byte[0], ObjectSerDeUtils.TDIGEST_SER_DE.serialize(empty), smallBuffer.array(), + ObjectSerDeUtils.TDIGEST_SER_DE.serialize(two) + }; + int[][] groupKeys = {{}, {0, 1}, {0, 1, 2, 2}, {0, 1, 2, 3}}; + + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(4, 4); + function.aggregateGroupByMV(values.length, groupKeys, resultHolder, + Map.of(EXPRESSION, bytesBlockValSet(values, null))); + + PercentileTDigestAccumulator pendingAccumulator = resultHolder.getResult(3); + for (String fieldName + : new String[]{"_rawValues", "_centroidMeans", "_outputMeans", "_incomingMeans"}) { + Field field = PercentileTDigestAccumulator.class.getDeclaredField(fieldName); + field.setAccessible(true); + Assert.assertNull(field.get(pendingAccumulator)); + } + + long[] expectedSizes = {2L, 2L, 3L, 1L}; + double[] expectedCompressions = {20.0, 20.0, 100.0, 100.0}; + double[] expectedMinimums = {1.0, 1.0, 1.0, 2.0}; + for (int groupKey = 0; groupKey < 4; groupKey++) { + TDigest result = function.extractGroupByResult(resultHolder, groupKey); + Assert.assertEquals(result.size(), expectedSizes[groupKey]); + Assert.assertEquals(result.compression(), expectedCompressions[groupKey]); + Assert.assertEquals(result.getMin(), expectedMinimums[groupKey]); + Assert.assertEquals(result.getMax(), 2.0); + } + } + + @Test + public void testSerializedGroupByMVPreservesOversizedSmallCapacity() { + int numCentroids = 80; + ByteBuffer oversizedSmall = ByteBuffer.allocate(Integer.BYTES + 2 * Double.BYTES + Float.BYTES + + 3 * Short.BYTES + 2 * Float.BYTES * numCentroids); + oversizedSmall.putInt(2); + oversizedSmall.putDouble(0.0); + oversizedSmall.putDouble(numCentroids - 1.0); + oversizedSmall.putFloat(20.0F); + oversizedSmall.putShort((short) numCentroids); + oversizedSmall.putShort((short) 100); + oversizedSmall.putShort((short) numCentroids); + for (int i = 0; i < numCentroids; i++) { + oversizedSmall.putFloat(1.0F); + oversizedSmall.putFloat(i); + } + TDigest empty = TDigest.createMergingDigest(20.0); + + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + TDigest lazyResult = function.deserializeIntermediateResult( + new CustomObject(ObjectSerDeUtils.ObjectType.TDigest.getValue(), ByteBuffer.wrap(oversizedSmall.array()))); + assertOversizedSmallIntermediateRoundTrip(function, lazyResult, numCentroids); + + GroupByResultHolder resultHolder = function.createGroupByResultHolder(1, 1); + function.aggregateGroupByMV(2, new int[][]{{0}, {0}}, resultHolder, + Map.of(EXPRESSION, bytesBlockValSet( + new byte[][]{oversizedSmall.array(), ObjectSerDeUtils.TDIGEST_SER_DE.serialize(empty)}, null))); + + TDigest result = function.extractGroupByResult(resultHolder, 0); + Assert.assertEquals(result.size(), numCentroids); + Assert.assertEquals(result.compression(), 20.0); + Assert.assertEquals(result.getMin(), 0.0); + Assert.assertEquals(result.getMax(), numCentroids - 1.0); + Assert.assertEquals(result.quantile(0.75), 59.5, 1.0); + assertOversizedSmallIntermediateRoundTrip(function, result, numCentroids); + result.add(numCentroids); + Assert.assertEquals(result.size(), numCentroids + 1L); + + double preciseCompression = 20.0000001; + TDigest preciseEmpty = TDigest.createMergingDigest(preciseCompression); + GroupByResultHolder preciseResultHolder = function.createGroupByResultHolder(1, 1); + function.aggregateGroupByMV(2, new int[][]{{0}, {0}}, preciseResultHolder, + Map.of(EXPRESSION, bytesBlockValSet( + new byte[][]{ObjectSerDeUtils.TDIGEST_SER_DE.serialize(preciseEmpty), oversizedSmall.array()}, null))); + TDigest preciseResult = function.extractGroupByResult(preciseResultHolder, 0); + Assert.assertEquals(preciseResult.compression(), preciseCompression); + Assert.assertEquals(preciseResult.size(), numCentroids); + Assert.assertTrue(preciseResult.centroidCount() <= 2 * Math.ceil(preciseCompression) + 10); + } + + private static void assertOversizedSmallIntermediateRoundTrip(PercentileTDigestAggregationFunction function, + TDigest result, int numCentroids) { + SerializedIntermediateResult serializedResult = function.serializeIntermediateResult(result); + Assert.assertEquals(serializedResult.getType(), ObjectSerDeUtils.ObjectType.TDigest.getValue()); + byte[] serializedBytes = serializedResult.getBytes(); + Assert.assertEquals(ByteBuffer.wrap(serializedBytes).getInt(), 2); + TDigest roundTripped = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(serializedBytes); + Assert.assertEquals(roundTripped.size(), numCentroids); + Assert.assertEquals(roundTripped.centroidCount(), numCentroids); + Assert.assertEquals(roundTripped.compression(), 20.0); + Assert.assertEquals(roundTripped.getMin(), 0.0); + Assert.assertEquals(roundTripped.getMax(), numCentroids - 1.0); + double previousValue = Double.NEGATIVE_INFINITY; + for (double quantile : new double[]{0.0, 0.5, 0.75, 0.95, 0.99, 1.0}) { + double value = roundTripped.quantile(quantile); + Assert.assertTrue(Double.isFinite(value)); + Assert.assertTrue(value >= previousValue); + Assert.assertEquals(value, result.quantile(quantile)); + previousValue = value; + } + } + + @Test + public void testSerializedGroupByMVSharedInputUsesCentroidCountForInitialCapacity() + throws ReflectiveOperationException { + ByteBuffer oneCentroid = ByteBuffer.allocate(30 + 2 * Float.BYTES); + oneCentroid.putInt(2); + oneCentroid.putDouble(1.0); + oneCentroid.putDouble(1.0); + oneCentroid.putFloat(1_000_000.0F); + oneCentroid.putShort((short) 1); + oneCentroid.putShort((short) 1); + oneCentroid.putShort((short) 1); + oneCentroid.putFloat(1.0F); + oneCentroid.putFloat(1.0F); + TDigest empty = TDigest.createMergingDigest(20.0); + byte[] emptyBytes = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(empty); + PercentileTDigestAccumulator accumulator = PercentileTDigestAccumulator.forSerializedTDigest(emptyBytes); + accumulator.addSerializedTDigest(emptyBytes); + PercentileTDigestAccumulator.SerializedTDigestInput input = + new PercentileTDigestAccumulator.SerializedTDigestInput(); + input.reset(oneCentroid.array()); + accumulator.addSerializedTDigest(input); + + Field meansField = PercentileTDigestAccumulator.SerializedTDigestInput.class.getDeclaredField("_means"); + meansField.setAccessible(true); + Assert.assertEquals(((double[]) meansField.get(input)).length, 1); + Assert.assertEquals(accumulator.toTDigest().size(), 1L); + } + + @Test + public void testSerializedAccumulatorAllocatesRawBufferLazily() + throws ReflectiveOperationException { + byte[] serializedDigest = createSerializedDigests(1, 100, 100.0, ignored -> false)[0]; + Field rawValuesField = PercentileTDigestAccumulator.class.getDeclaredField("_rawValues"); + rawValuesField.setAccessible(true); + + PercentileTDigestAccumulator accumulator = PercentileTDigestAccumulator.forSerializedTDigest(serializedDigest); + Assert.assertNull(rawValuesField.get(accumulator)); + accumulator.add(new double[0], 0, 0); + Assert.assertNull(rawValuesField.get(accumulator)); + accumulator.addSerializedTDigest(serializedDigest); + accumulator.toTDigest(); + Assert.assertNull(rawValuesField.get(accumulator)); + + accumulator.add(new double[]{1.0}, 0, 1); + Assert.assertNotNull(rawValuesField.get(accumulator)); + accumulator.addSerializedTDigest(serializedDigest); + TDigest result = accumulator.toTDigest(); + Assert.assertEquals(result.size(), 201L); + Assert.assertEquals(result.getMin(), 0.0); + Assert.assertEquals(result.getMax(), 1.0); + Assert.assertNull(rawValuesField.get( + PercentileTDigestAccumulator.forSerializedTDigestWithMergeBuffers(serializedDigest))); + Assert.assertNotNull(rawValuesField.get(new PercentileTDigestAccumulator(100))); + } + + @Test + public void testSerializedAccumulatorGrowsBuffersForRawValues() { + TDigest serializedDigest = TDigest.createMergingDigest(100.0); + serializedDigest.add(-1.0); + byte[] serializedBytes = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(serializedDigest); + PercentileTDigestAccumulator accumulator = + PercentileTDigestAccumulator.forSerializedTDigest(serializedBytes); + accumulator.addSerializedTDigest(serializedBytes); + + double[] rawValues = new double[768]; + for (int i = 0; i < rawValues.length; i++) { + rawValues[i] = i; + } + accumulator.add(rawValues, 0, rawValues.length); + + TDigest result = accumulator.toTDigest(); + Assert.assertEquals(result.size(), 769L); + Assert.assertEquals(result.getMin(), -1.0); + Assert.assertEquals(result.getMax(), 767.0); + } + + @Test + public void testSerializedAccumulatorPreservesCompressionAfterEmptyDigests() { + TDigest first = TDigest.createMergingDigest(20.0); + TDigest second = TDigest.createMergingDigest(30.0); + TDigest third = TDigest.createMergingDigest(40.0); + third.add(1.0); + byte[] firstBytes = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(first); + PercentileTDigestAccumulator accumulator = PercentileTDigestAccumulator.forSerializedTDigest(firstBytes); + accumulator.addSerializedTDigest(firstBytes); + accumulator.addSerializedTDigest(ObjectSerDeUtils.TDIGEST_SER_DE.serialize(second)); + accumulator.addSerializedTDigest(ObjectSerDeUtils.TDIGEST_SER_DE.serialize(third)); + + TDigest result = accumulator.toTDigest(); + Assert.assertEquals(result.compression(), 20.0); + Assert.assertEquals(result.size(), 1L); + Assert.assertEquals(result.quantile(0.5), 1.0); + } + + @Test + public void testRawNumericGroupByStillUsesTDigestState() { + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(4, new int[]{0, 0, 1, 1}, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, new double[]{1.0, 2.0, 3.0, 4.0}))); + + Assert.assertTrue(resultHolder.getResult(0) instanceof TDigest); + Assert.assertEquals(function.extractGroupByResult(resultHolder, 0).size(), 2L); + Assert.assertEquals(function.extractGroupByResult(resultHolder, 1).size(), 2L); + } + + @DataProvider(name = "malformedSerializedDigests") + public static Object[][] malformedSerializedDigests() { + ByteBuffer invalidEncoding = ByteBuffer.allocate(Integer.BYTES); + invalidEncoding.putInt(99); + ByteBuffer truncatedVerbose = ByteBuffer.allocate(4 * Integer.BYTES + 2 * Double.BYTES); + truncatedVerbose.putInt(1); + truncatedVerbose.putDouble(0.0); + truncatedVerbose.putDouble(1.0); + truncatedVerbose.putDouble(100.0); + truncatedVerbose.putInt(1); + ByteBuffer truncatedSmall = ByteBuffer.allocate(3 * Integer.BYTES + 2 * Double.BYTES + Short.BYTES); + truncatedSmall.putInt(2); + truncatedSmall.putDouble(0.0); + truncatedSmall.putDouble(1.0); + truncatedSmall.putFloat(100.0F); + truncatedSmall.putShort((short) 210); + truncatedSmall.putShort((short) 500); + truncatedSmall.putShort((short) 1); + ByteBuffer oversizedVerbose = ByteBuffer.allocate(4 * Integer.BYTES + 2 * Double.BYTES); + oversizedVerbose.putInt(1); + oversizedVerbose.putDouble(0.0); + oversizedVerbose.putDouble(1.0); + oversizedVerbose.putDouble(100.0); + oversizedVerbose.putInt(10_000_000); + int numOversizedVerboseCentroids = 51; + ByteBuffer populatedOversizedVerbose = ByteBuffer.allocate(4 * Integer.BYTES + 2 * Double.BYTES + + 2 * Double.BYTES * numOversizedVerboseCentroids); + populatedOversizedVerbose.putInt(1); + populatedOversizedVerbose.putDouble(0.0); + populatedOversizedVerbose.putDouble(numOversizedVerboseCentroids - 1.0); + populatedOversizedVerbose.putDouble(20.0); + populatedOversizedVerbose.putInt(numOversizedVerboseCentroids); + for (int i = 0; i < numOversizedVerboseCentroids; i++) { + populatedOversizedVerbose.putDouble(1.0); + populatedOversizedVerbose.putDouble(i); + } + ByteBuffer oversizedSmall = ByteBuffer.allocate(3 * Integer.BYTES + 2 * Double.BYTES + Short.BYTES); + oversizedSmall.putInt(2); + oversizedSmall.putDouble(0.0); + oversizedSmall.putDouble(1.0); + oversizedSmall.putFloat(100.0F); + oversizedSmall.putShort((short) 210); + oversizedSmall.putShort((short) 500); + oversizedSmall.putShort(Short.MAX_VALUE); + ByteBuffer negativeMainCapacity = ByteBuffer.allocate(30); + negativeMainCapacity.putInt(2); + negativeMainCapacity.putDouble(0.0); + negativeMainCapacity.putDouble(1.0); + negativeMainCapacity.putFloat(100.0F); + negativeMainCapacity.putShort((short) -2); + negativeMainCapacity.putShort((short) 1); + negativeMainCapacity.putShort((short) 0); + ByteBuffer negativeBufferCapacity = ByteBuffer.allocate(30); + negativeBufferCapacity.putInt(2); + negativeBufferCapacity.putDouble(0.0); + negativeBufferCapacity.putDouble(1.0); + negativeBufferCapacity.putFloat(100.0F); + negativeBufferCapacity.putShort((short) 1); + negativeBufferCapacity.putShort((short) -2); + negativeBufferCapacity.putShort((short) 0); + ByteBuffer centroidsExceedMainCapacity = ByteBuffer.allocate(30 + 4 * Float.BYTES); + centroidsExceedMainCapacity.putInt(2); + centroidsExceedMainCapacity.putDouble(0.0); + centroidsExceedMainCapacity.putDouble(1.0); + centroidsExceedMainCapacity.putFloat(100.0F); + centroidsExceedMainCapacity.putShort((short) 1); + centroidsExceedMainCapacity.putShort((short) 2); + centroidsExceedMainCapacity.putShort((short) 2); + centroidsExceedMainCapacity.putFloat(1.0F); + centroidsExceedMainCapacity.putFloat(0.0F); + centroidsExceedMainCapacity.putFloat(1.0F); + centroidsExceedMainCapacity.putFloat(1.0F); + return new Object[][]{ + {invalidEncoding.array()}, {truncatedVerbose.array()}, {truncatedSmall.array()}, {oversizedVerbose.array()}, + {populatedOversizedVerbose.array()}, {oversizedSmall.array()}, {negativeMainCapacity.array()}, + {negativeBufferCapacity.array()}, {centroidsExceedMainCapacity.array()} + }; + } + + @Test(dataProvider = "malformedSerializedDigests") + public void testMalformedSerializedDigestMatchesMergingDigest(byte[] bytes) { + RuntimeException expected = Assert.expectThrows(RuntimeException.class, + () -> ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytes)); + RuntimeException actual = Assert.expectThrows(RuntimeException.class, + () -> aggregateSerialized(new byte[][]{bytes}, null, false)); + Assert.assertEquals(actual.getClass(), expected.getClass()); + Assert.assertEquals(actual.getMessage(), expected.getMessage()); + + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(1, 1); + RuntimeException groupByActual = Assert.expectThrows(RuntimeException.class, + () -> function.aggregateGroupBySV(1, new int[]{0}, resultHolder, + Map.of(EXPRESSION, bytesBlockValSet(new byte[][]{bytes}, null)))); + Assert.assertEquals(groupByActual.getClass(), expected.getClass()); + Assert.assertEquals(groupByActual.getMessage(), expected.getMessage()); + + TDigest empty = TDigest.createMergingDigest(100.0); + byte[] emptyBytes = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(empty); + RuntimeException materializedActual = Assert.expectThrows(RuntimeException.class, + () -> aggregateSerialized(new byte[][]{emptyBytes, bytes}, null, false)); + Assert.assertEquals(materializedActual.getClass(), expected.getClass()); + Assert.assertEquals(materializedActual.getMessage(), expected.getMessage()); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(1, 1); + RuntimeException groupByMVActual = Assert.expectThrows(RuntimeException.class, + () -> function.aggregateGroupByMV(2, new int[][]{{0}, {0}}, groupByMVResultHolder, + Map.of(EXPRESSION, bytesBlockValSet(new byte[][]{emptyBytes, bytes}, null)))); + Assert.assertEquals(groupByMVActual.getClass(), expected.getClass()); + Assert.assertEquals(groupByMVActual.getMessage(), expected.getMessage()); + } + + @DataProvider(name = "compressionFactors") + public static Object[][] compressionFactors() { + return new Object[][]{{10}, {50}, {100}, {200}, {1_000}, {10_000}}; + } + + @DataProvider(name = "reducerMergeCases") + public static Object[][] reducerMergeCases() { + List cases = new ArrayList<>(); + for (int fanIn : new int[]{8, 32, 128}) { + for (int compression : new int[]{20, 100, 1_000}) { + for (ReducerDistribution distribution : ReducerDistribution.values()) { + for (MergeOrder mergeOrder : MergeOrder.values()) { + for (ReducerInput reducerInput : ReducerInput.values()) { + cases.add(new Object[]{fanIn, compression, distribution, mergeOrder, reducerInput}); + } + } + } + } + } + return cases.toArray(new Object[0][]); + } + + @Test(dataProvider = "reducerMergeCases") + public void testReducerMergeStructureAndAccuracy(int fanIn, int compression, ReducerDistribution distribution, + MergeOrder mergeOrder, ReducerInput reducerInput) { + int valuesPerDigest = 64; + int numValues = fanIn * valuesPerDigest; + double[] rawValues = new double[numValues]; + TDigest[] sources = new TDigest[fanIn]; + for (int digestIndex = 0; digestIndex < fanIn; digestIndex++) { + SplittableRandom random = new SplittableRandom(0x5EEDL + digestIndex); + double[] sourceValues = new double[valuesPerDigest]; + for (int valueIndex = 0; valueIndex < valuesPerDigest; valueIndex++) { + double value = distribution.value(random.nextDouble()); + rawValues[digestIndex * valuesPerDigest + valueIndex] = value; + sourceValues[valueIndex] = value; + } + sources[digestIndex] = createReducerSource(sourceValues, compression, reducerInput); + } + + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, compression, false); + TDigest result = TDigest.createMergingDigest(compression); + for (int sourceIndex : mergeOrder.sourceIndexes(fanIn)) { + result = function.merge(result, sources[sourceIndex]); + } + + Arrays.sort(rawValues); + String caseDescription = + fanIn + "/" + compression + "/" + distribution + "/" + mergeOrder + "/" + reducerInput; + double maxRankError = compression == 20 ? 0.06 : 0.02; + assertValidReducerResult(result, rawValues, caseDescription, maxRankError); + + SerializedIntermediateResult serializedResult = function.serializeIntermediateResult(result); + Assert.assertEquals(serializedResult.getType(), ObjectSerDeUtils.ObjectType.TDigest.getValue()); + TDigest roundTripped = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(serializedResult.getBytes()); + Assert.assertTrue(roundTripped instanceof MergingDigest); + assertValidReducerResult(roundTripped, rawValues, caseDescription + "/standard-wire", maxRankError); + for (double quantile : new double[]{0.0, 0.5, 0.75, 0.95, 0.99, 1.0}) { + Assert.assertEquals(result.quantile(quantile), roundTripped.quantile(quantile), + "Wire round trip changed p" + quantile * 100 + " for " + caseDescription); + } + } + + @Test + public void testReducerMergeEmptyAndSingleton() { + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + for (ReducerInput reducerInput : ReducerInput.values()) { + TDigest empty = createReducerSource(new double[0], 100, reducerInput); + TDigest singleton = createReducerSource(new double[]{42.0}, 100, reducerInput); + + TDigest result = function.merge(empty, singleton); + Assert.assertSame(result, singleton); + Assert.assertSame(function.merge(result, createReducerSource(new double[0], 100, reducerInput)), result); + assertValidReducerResult(result, new double[]{42.0}, "singleton/" + reducerInput); + } + } + + @Test(dataProvider = "compressionFactors") + public void testCustomCompression(int compression) { + int numRows = 20_003; + SplittableRandom random = new SplittableRandom(42); + double[] values = new double[numRows]; + for (int i = 0; i < numRows; i++) { + values[i] = random.nextDouble(); + } + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, compression, false); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(numRows, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, values))); + TDigest result = function.extractAggregationResult(resultHolder); + + Assert.assertEquals(result.size(), numRows); + Assert.assertEquals(result.compression(), compression); + double[] sortedValues = values.clone(); + Arrays.sort(sortedValues); + Assert.assertEquals(function.extractFinalResult(result), sortedValues[numRows * 75 / 100], 0.02); + } + + @Test(expectedExceptions = IllegalArgumentException.class, + expectedExceptionsMessageRegExp = "Cannot add NaN to t-digest") + public void testNaNRejected() { + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(3, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, new double[]{1.0, Double.NaN, 2.0}))); + } @Test public void testCanUseStarTreeDefaultCompression() { PercentileTDigestAggregationFunction function = new PercentileTDigestAggregationFunction( - ExpressionContext.forIdentifier("col"), 95, false); + EXPRESSION, 95, false); Assert.assertTrue(function.canUseStarTree(Map.of())); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, "100"))); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, 100))); Assert.assertFalse(function.canUseStarTree(Map.of(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, "200"))); - function = new PercentileTDigestAggregationFunction(ExpressionContext.forIdentifier("col"), 99, 100, true); + function = new PercentileTDigestAggregationFunction(EXPRESSION, 99, 100, true); Assert.assertTrue(function.canUseStarTree(Map.of())); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, "100"))); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, 100))); @@ -47,11 +807,255 @@ public void testCanUseStarTreeDefaultCompression() { @Test public void testCanUseStarTreeCustomCompression() { PercentileTDigestAggregationFunction function = new PercentileTDigestAggregationFunction( - ExpressionContext.forIdentifier("col"), 95, 200, false); + EXPRESSION, 95, 200, false); Assert.assertFalse(function.canUseStarTree(Map.of())); Assert.assertFalse(function.canUseStarTree(Map.of(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, "100"))); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, "200"))); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.PERCENTILETDIGEST_COMPRESSION_FACTOR_KEY, 200))); } + + private static byte[][] createSerializedDigests(int numDigests, int valuesPerDigest, double compression, + IntPredicate useSmallEncoding) { + byte[][] serializedDigests = new byte[numDigests][]; + for (int digestIndex = 0; digestIndex < numDigests; digestIndex++) { + TDigest digest = TDigest.createMergingDigest(compression); + for (int valueIndex = 0; valueIndex < valuesPerDigest; valueIndex++) { + digest.add(serializedValue(digestIndex, valueIndex, numDigests, valuesPerDigest)); + } + if (useSmallEncoding.test(digestIndex)) { + ByteBuffer buffer = ByteBuffer.allocate(digest.smallByteSize()); + digest.asSmallBytes(buffer); + serializedDigests[digestIndex] = buffer.array(); + } else { + serializedDigests[digestIndex] = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(digest); + } + } + return serializedDigests; + } + + private static double serializedValue(int digestIndex, int valueIndex, int numDigests, int valuesPerDigest) { + return (valueIndex * (double) numDigests + digestIndex) / (numDigests * (double) valuesPerDigest); + } + + private static TDigest aggregateSerialized(byte[][] values, RoaringBitmap nullBitmap, boolean nullHandlingEnabled) { + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, nullHandlingEnabled); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(values.length, resultHolder, Map.of(EXPRESSION, bytesBlockValSet(values, nullBitmap))); + return function.extractAggregationResult(resultHolder); + } + + private static void assertSerializedAccumulatorState(GroupByResultHolder resultHolder, int groupKey) { + Assert.assertTrue(resultHolder.getResult(groupKey) instanceof PercentileTDigestAccumulator); + } + + private static void assertGroupResult(PercentileTDigestAggregationFunction function, + GroupByResultHolder resultHolder, int groupKey, long expectedSize, double expectedCompression, + double expectedMin, double expectedMax) { + TDigest result = function.extractGroupByResult(resultHolder, groupKey); + Assert.assertTrue(result instanceof PercentileTDigestAccumulator); + Assert.assertEquals(result.size(), expectedSize); + Assert.assertEquals(result.compression(), expectedCompression); + Assert.assertEquals(result.getMin(), expectedMin); + Assert.assertEquals(result.getMax(), expectedMax); + Assert.assertEquals(result.quantile(0.75), 0.75, 0.02); + Assert.assertTrue(((PercentileTDigestAccumulator) result).toTDigest() instanceof MergingDigest); + } + + private static void assertValidReducerResult(TDigest result, double[] sortedValues, String caseDescription) { + assertValidReducerResult(result, sortedValues, caseDescription, 0.02); + } + + private static void assertValidReducerResult(TDigest result, double[] sortedValues, String caseDescription, + double maxRankError) { + Assert.assertEquals(result.size(), sortedValues.length, "Total weight differs for " + caseDescription); + long centroidWeight = 0L; + double previousMean = Double.NEGATIVE_INFINITY; + for (Centroid centroid : result.centroids()) { + Assert.assertTrue(Double.isFinite(centroid.mean()), "Non-finite centroid mean for " + caseDescription); + Assert.assertTrue(centroid.count() > 0, "Non-positive centroid weight for " + caseDescription); + Assert.assertTrue(centroid.mean() + 1e-12 >= previousMean, + "Centroids are not sorted for " + caseDescription); + centroidWeight += centroid.count(); + previousMean = centroid.mean(); + } + Assert.assertEquals(centroidWeight, sortedValues.length, + "Centroid weights do not sum to the input size for " + caseDescription); + + double previousQuantile = Double.NEGATIVE_INFINITY; + for (double quantile : new double[]{0.0, 0.5, 0.75, 0.95, 0.99, 1.0}) { + double value = result.quantile(quantile); + Assert.assertTrue(Double.isFinite(value), "Non-finite p" + quantile * 100 + " for " + caseDescription); + Assert.assertTrue(value >= previousQuantile, "Quantiles are not monotonic for " + caseDescription); + previousQuantile = value; + if (quantile == 0.0) { + Assert.assertEquals(value, sortedValues[0], "Incorrect minimum for " + caseDescription); + } else if (quantile == 1.0) { + Assert.assertEquals(value, sortedValues[sortedValues.length - 1], + "Incorrect maximum for " + caseDescription); + } else { + int lowerRank = lowerBound(sortedValues, value); + int upperRank = upperBound(sortedValues, value); + double rankError = Math.max(0.0, Math.max(lowerRank / (double) sortedValues.length - quantile, + quantile - upperRank / (double) sortedValues.length)); + Assert.assertTrue(rankError <= maxRankError, + "Rank error " + rankError + " exceeds the frozen envelope for p" + quantile * 100 + " in " + + caseDescription); + } + } + } + + private static TDigest createReducerSource(double[] values, int compression, ReducerInput reducerInput) { + PercentileTDigestAggregationFunction function = + new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, compression, false); + if (reducerInput == ReducerInput.SERVER_LOCAL) { + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(values.length, resultHolder, + Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, values))); + return function.extractAggregationResult(resultHolder); + } + + TDigest digest = TDigest.createMergingDigest(compression); + for (double value : values) { + digest.add(value); + } + byte[] bytes = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(digest); + return function.deserializeIntermediateResult(new CustomObject(ObjectSerDeUtils.ObjectType.TDigest.getValue(), + ByteBuffer.wrap(bytes))); + } + + private static int lowerBound(double[] sortedValues, double value) { + int low = 0; + int high = sortedValues.length; + while (low < high) { + int middle = (low + high) >>> 1; + if (sortedValues[middle] < value) { + low = middle + 1; + } else { + high = middle; + } + } + return low; + } + + private static int upperBound(double[] sortedValues, double value) { + int low = 0; + int high = sortedValues.length; + while (low < high) { + int middle = (low + high) >>> 1; + if (sortedValues[middle] <= value) { + low = middle + 1; + } else { + high = middle; + } + } + return low; + } + + private enum ReducerDistribution { + UNIFORM { + @Override + double value(double unitValue) { + return unitValue; + } + }, + SKEWED { + @Override + double value(double unitValue) { + return Math.pow(unitValue, 8.0); + } + }, + HEAVY_TAIL { + @Override + double value(double unitValue) { + return Math.pow(1.0 - 0.999 * unitValue, -1.5) - 1.0; + } + }, + BIMODAL { + @Override + double value(double unitValue) { + return unitValue < 0.5 ? 0.1 + 0.2 * unitValue : 0.7 + 0.2 * unitValue; + } + }, + DUPLICATE_HEAVY { + @Override + double value(double unitValue) { + return unitValue < 0.7 ? 0.1 : unitValue < 0.9 ? 0.5 : 0.9; + } + }; + + abstract double value(double unitValue); + } + + private enum ReducerInput { + SERVER_LOCAL, + DISTRIBUTED + } + + private enum MergeOrder { + ORIGINAL { + @Override + int[] sourceIndexes(int fanIn) { + int[] indexes = new int[fanIn]; + for (int i = 0; i < fanIn; i++) { + indexes[i] = i; + } + return indexes; + } + }, + REVERSED { + @Override + int[] sourceIndexes(int fanIn) { + int[] indexes = ORIGINAL.sourceIndexes(fanIn); + for (int i = 0; i < fanIn / 2; i++) { + int otherIndex = fanIn - i - 1; + int value = indexes[i]; + indexes[i] = indexes[otherIndex]; + indexes[otherIndex] = value; + } + return indexes; + } + }, + RANDOMIZED { + @Override + int[] sourceIndexes(int fanIn) { + int[] indexes = ORIGINAL.sourceIndexes(fanIn); + SplittableRandom random = new SplittableRandom(0xBADC0FFEE0DDF00DL + fanIn); + for (int i = indexes.length - 1; i > 0; i--) { + int otherIndex = random.nextInt(i + 1); + int value = indexes[i]; + indexes[i] = indexes[otherIndex]; + indexes[otherIndex] = value; + } + return indexes; + } + }; + + abstract int[] sourceIndexes(int fanIn); + } + + private static BlockValSet bytesBlockValSet(byte[][] values, RoaringBitmap nullBitmap) { + return new SyntheticBlockValSets.Base() { + @Override + public RoaringBitmap getNullBitmap() { + return nullBitmap; + } + + @Override + public DataType getValueType() { + return DataType.BYTES; + } + + @Override + public boolean isSingleValue() { + return true; + } + + @Override + public byte[][] getBytesValuesSV() { + return values; + } + }; + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGeneratorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGeneratorTest.java index 7d78f3f914e0..26d3f5c10bfb 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGeneratorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/DictionaryBasedGroupKeyGeneratorTest.java @@ -21,9 +21,11 @@ import java.io.File; import java.util.ArrayList; import java.util.Arrays; +import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Random; import java.util.Set; import org.apache.commons.io.FileUtils; @@ -98,7 +100,8 @@ private void setup() value += 1 + _random.nextInt(MAX_STEP_LENGTH); } for (String mvColumn : MV_COLUMNS) { - int numMultiValues = 1 + _random.nextInt(MAX_NUM_MULTI_VALUES); + int numMultiValues = i == UNIQUE_ROWS / 2 && mvColumn.equals(MV_COLUMNS[0]) ? 2 + : 1 + _random.nextInt(MAX_NUM_MULTI_VALUES); Integer[] values = new Integer[numMultiValues]; for (int k = 0; k < numMultiValues; k++) { values[k] = value; @@ -133,9 +136,10 @@ private void setup() driver.build(); IndexSegment indexSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR_PATH, SEGMENT_NAME), ReadMode.heap); - // Generate a random query to filter out 2 unique rows - int docId1 = _random.nextInt(50); - int docId2 = docId1 + 1 + _random.nextInt(50); + // Filter 2 unique rows with sparse dictionary ids so optimized-capacity tests cannot accidentally use raw ids as + // dense group ids. + int docId1 = UNIQUE_ROWS / 2; + int docId2 = UNIQUE_ROWS - 1; // NOTE: put all columns into group-by so that transform operator has expressions for all columns String query = String.format("SELECT COUNT(*) FROM testTable WHERE %s IN (%d, %d) GROUP BY %s, %s", FILTER_COLUMN, docId1, @@ -250,6 +254,70 @@ public void testArrayMapBasedSingleValue() { assertEquals(DictionaryBasedGroupKeyGenerator.THREAD_LOCAL_INT_ARRAY_MAP.get().size(), 0); } + @Test(dataProvider = "optimizedCardinalityDataProvider") + public void testOptimizedCardinalityWithSparseDictionaryIds(String[] groupByColumns, int expectedUpperBound) { + ExpressionContext[] expressions = getExpressions(groupByColumns); + Map groupByExpressionSizesFromPredicates = new HashMap<>(); + for (ExpressionContext expression : expressions) { + groupByExpressionSizesFromPredicates.put(expression, 2); + } + DictionaryBasedGroupKeyGenerator dictionaryBasedGroupKeyGenerator = + new DictionaryBasedGroupKeyGenerator(_projectOperator, expressions, + Server.DEFAULT_QUERY_EXECUTOR_NUM_GROUPS_LIMIT, + Server.DEFAULT_QUERY_EXECUTOR_MAX_INITIAL_RESULT_HOLDER_CAPACITY, + groupByExpressionSizesFromPredicates); + assertEquals(dictionaryBasedGroupKeyGenerator.getGlobalGroupKeyUpperBound(), expectedUpperBound, _errorMessage); + + dictionaryBasedGroupKeyGenerator.generateKeysForBlock(_valueBlock, SV_GROUP_KEY_BUFFER); + + assertEquals(dictionaryBasedGroupKeyGenerator.getCurrentGroupKeyUpperBound(), 2, _errorMessage); + compareSingleValueBuffer(); + testGetGroupKeys(dictionaryBasedGroupKeyGenerator.getGroupKeys(), 2); + dictionaryBasedGroupKeyGenerator.close(); + } + + @DataProvider(name = "optimizedCardinalityDataProvider") + public Object[][] optimizedCardinalityDataProvider() { + return new Object[][]{ + // Raw key cardinality fits array-based holder, but sparse dictionary ids require remapping. + {new String[]{"s1"}, 2}, + // Raw key cardinality fits int, long and int-array map holders respectively. + {new String[]{"s1", "s2", "s3"}, 8}, + {new String[]{"s1", "s2", "s3", "s4", "s5"}, 32}, + {SV_COLUMNS, 1 << SV_COLUMNS.length} + }; + } + + @Test + public void testOptimizedCardinalityDoesNotUnderestimateMultiValueGroups() { + ExpressionContext expression = ExpressionContext.forIdentifier("m1"); + int predicateDictId = _valueBlock.getBlockValueSet(expression).getDictionaryIdsMV()[0][0]; + Object predicateValue = _projectOperator.getResultColumnContext(expression).getDictionary() + .getInternal(predicateDictId); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(String.format( + "SELECT COUNT(*), m1 FROM testTable WHERE m1 IN (%s) GROUP BY m1 LIMIT 1000", predicateValue)); + queryContext.setOptimizeMaxInitialResultHolderCapacity(true); + DefaultGroupByExecutor defaultGroupByExecutor = + new DefaultGroupByExecutor(queryContext, new ExpressionContext[]{expression}, _projectOperator); + GroupKeyGenerator groupKeyGenerator = defaultGroupByExecutor.getGroupKeyGenerator(); + int dictionaryCardinality = _projectOperator.getResultColumnContext(expression).getDictionary().length(); + assertEquals(groupKeyGenerator.getGlobalGroupKeyUpperBound(), dictionaryCardinality, _errorMessage); + + groupKeyGenerator.generateKeysForBlock(_valueBlock, MV_GROUP_KEY_BUFFER); + + assertTrue(MV_GROUP_KEY_BUFFER[0].length > 1, _errorMessage); + Set groupIds = new HashSet<>(); + for (int[] rowGroupIds : MV_GROUP_KEY_BUFFER) { + for (int groupId : rowGroupIds) { + assertTrue(groupId != GroupKeyGenerator.INVALID_ID, _errorMessage); + groupIds.add(groupId); + } + } + assertTrue(groupIds.size() > 1, _errorMessage); + testGetGroupKeys(groupKeyGenerator.getGroupKeys(), groupIds.size()); + groupKeyGenerator.close(); + } + /** * Helper method to compare the values inside the single value group key buffer. * diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/FilteredAggregationsTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/FilteredAggregationsTest.java index b615da21f78b..71932764d904 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/FilteredAggregationsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/FilteredAggregationsTest.java @@ -403,6 +403,20 @@ public void testGroupByMultipleColumns() { testQuery(filterQuery, nonFilterQuery); } + @Test + public void testGroupBySparseDictionaryIdsWithOptimizedInitialCapacity() { + String filteredAggregationsOptions = + "SET " + CommonConstants.Broker.Request.QueryOptionKey.FILTERED_AGGREGATIONS_SKIP_EMPTY_GROUPS + "=true; "; + String query = "SELECT NO_INDEX_COL, " + + "SUM(INT_COL) FILTER(WHERE INT_COL = 29998) AS sum1, " + + "SUM(INT_COL) FILTER(WHERE INT_COL = 29999) AS sum2 " + + "FROM MyTable WHERE NO_INDEX_COL IN (29998, 29999) " + + "GROUP BY NO_INDEX_COL ORDER BY NO_INDEX_COL"; + String optimizedQuery = filteredAggregationsOptions + "SET " + + CommonConstants.Broker.Request.QueryOptionKey.OPTIMIZE_MAX_INITIAL_RESULT_HOLDER_CAPACITY + "=true; " + query; + testQuery(optimizedQuery, filteredAggregationsOptions + query); + } + @Test public void testGroupByCaseAlternative() { String filterQuery = "SELECT SUM(INT_COL), SUM(INT_COL) FILTER(WHERE INT_COL > 25000) AS total_sum FROM MyTable " diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/PercentileTDigestMVQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/PercentileTDigestMVQueriesTest.java index c2b9fc48c7ba..1860cdbbc830 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/PercentileTDigestMVQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/PercentileTDigestMVQueriesTest.java @@ -26,6 +26,7 @@ import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.StarTreeIndexConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.DimensionFieldSpec; @@ -54,15 +55,18 @@ */ public class PercentileTDigestMVQueriesTest extends PercentileTDigestQueriesTest { private static final int MAX_NUM_MULTI_VALUES = 10; + private long _numValues; @Override protected void buildSegment() throws Exception { + _numValues = 0L; List rows = new ArrayList<>(NUM_ROWS); for (int i = 0; i < NUM_ROWS; i++) { GenericRow row = new GenericRow(); int numMultiValues = RANDOM.nextInt(MAX_NUM_MULTI_VALUES) + 1; + _numValues += numMultiValues; Double[] values = new Double[numMultiValues]; TDigest tDigest = TDigest.createMergingDigest(PercentileTDigestAggregationFunction.DEFAULT_TDIGEST_COMPRESSION); TDigest tDigestCustom = TDigest.createMergingDigest(CUSTOM_COMPRESSION); @@ -93,8 +97,11 @@ protected void buildSegment() schema.addField(new MetricFieldSpec(TDIGEST_COLUMN, FieldSpec.DataType.BYTES)); schema.addField(new MetricFieldSpec(TDIGEST_CUSTOM_COMPRESSION_COLUMN, FieldSpec.DataType.BYTES)); schema.addField(new DimensionFieldSpec(GROUP_BY_COLUMN, FieldSpec.DataType.STRING, true)); + StarTreeIndexConfig starTreeIndexConfig = new StarTreeIndexConfig(List.of(GROUP_BY_COLUMN), null, + List.of(STAR_TREE_TDIGEST_COLUMN), null, 1); TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME) - .setNoDictionaryColumns(List.of(TDIGEST_COLUMN, TDIGEST_CUSTOM_COMPRESSION_COLUMN)).build(); + .setNoDictionaryColumns(List.of(TDIGEST_COLUMN, TDIGEST_CUSTOM_COMPRESSION_COLUMN)) + .setStarTreeIndexConfigs(List.of(starTreeIndexConfig)).build(); SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, schema); config.setOutDir(INDEX_DIR.getPath()); @@ -108,6 +115,11 @@ protected void buildSegment() } } + @Override + protected long getExpectedTotalDigestSize() { + return _numValues; + } + @Override protected String getAggregationQuery(int percentile) { return String.format("SELECT PERCENTILE%1$dMV(%2$s), PERCENTILETDIGEST%1$dMV(%2$s), PERCENTILETDIGEST%1$d(%3$s), " diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/PercentileTDigestQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/PercentileTDigestQueriesTest.java index 66edad32b228..ac7c2cb301f1 100644 --- a/pinot-core/src/test/java/org/apache/pinot/queries/PercentileTDigestQueriesTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/queries/PercentileTDigestQueriesTest.java @@ -25,10 +25,13 @@ import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; +import java.util.HashMap; import java.util.Iterator; import java.util.List; +import java.util.Map; import java.util.Random; import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.request.PinotQuery; import org.apache.pinot.common.response.broker.BrokerResponseNative; import org.apache.pinot.core.operator.blocks.results.AggregationResultsBlock; import org.apache.pinot.core.operator.blocks.results.GroupByResultsBlock; @@ -37,12 +40,16 @@ import org.apache.pinot.core.query.aggregation.function.PercentileTDigestAggregationFunction; import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult; import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; import org.apache.pinot.segment.spi.ImmutableSegment; import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.SegmentContext; import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.StarTreeIndexConfig; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.data.DimensionFieldSpec; @@ -53,12 +60,14 @@ import org.apache.pinot.spi.data.readers.RecordReader; import org.apache.pinot.spi.utils.ReadMode; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.apache.pinot.sql.parsers.CalciteSqlParser; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; /** @@ -89,6 +98,7 @@ public class PercentileTDigestQueriesTest extends BaseQueriesTest { protected static final String TDIGEST_COLUMN = "tDigestColumn"; protected static final String TDIGEST_CUSTOM_COMPRESSION_COLUMN = "tDigestColumnCustom"; protected static final String GROUP_BY_COLUMN = "groupByColumn"; + protected static final String STAR_TREE_TDIGEST_COLUMN = "percentileTDigest__" + DOUBLE_COLUMN; protected static final String[] GROUPS = new String[]{"G1", "G2", "G3"}; protected static final long RANDOM_SEED = System.nanoTime(); protected static final Random RANDOM = new Random(RANDOM_SEED); @@ -155,8 +165,11 @@ protected void buildSegment() schema.addField(new MetricFieldSpec(TDIGEST_COLUMN, FieldSpec.DataType.BYTES)); schema.addField(new MetricFieldSpec(TDIGEST_CUSTOM_COMPRESSION_COLUMN, FieldSpec.DataType.BYTES)); schema.addField(new DimensionFieldSpec(GROUP_BY_COLUMN, FieldSpec.DataType.STRING, true)); + StarTreeIndexConfig starTreeIndexConfig = new StarTreeIndexConfig(List.of(GROUP_BY_COLUMN), null, + List.of(STAR_TREE_TDIGEST_COLUMN), null, 1); TableConfig tableConfig = new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME) - .setNoDictionaryColumns(List.of(TDIGEST_COLUMN, TDIGEST_CUSTOM_COMPRESSION_COLUMN)).build(); + .setNoDictionaryColumns(List.of(TDIGEST_COLUMN, TDIGEST_CUSTOM_COMPRESSION_COLUMN)) + .setStarTreeIndexConfigs(List.of(starTreeIndexConfig)).build(); SegmentGeneratorConfig config = new SegmentGeneratorConfig(tableConfig, schema); config.setOutDir(INDEX_DIR.getPath()); @@ -246,6 +259,38 @@ public void testInterSegmentGroupBy() { } } + @Test + public void testStarTreeGroupBy() { + String query = String.format("SELECT PERCENTILETDIGEST(%s, 75) FROM %s GROUP BY %s", DOUBLE_COLUMN, TABLE_NAME, + GROUP_BY_COLUMN); + + GroupByOperator starTreeOperator = getOperator(query); + Map starTreeResults = getGroupByResults(starTreeOperator.nextBlock()); + long starTreeDocsScanned = starTreeOperator.getExecutionStatistics().getNumDocsScanned(); + + GroupByOperator nonStarTreeOperator = getGroupByOperator(query, false); + Map nonStarTreeResults = getGroupByResults(nonStarTreeOperator.nextBlock()); + long nonStarTreeDocsScanned = nonStarTreeOperator.getExecutionStatistics().getNumDocsScanned(); + + assertEquals(nonStarTreeDocsScanned, NUM_ROWS); + assertTrue(starTreeDocsScanned < nonStarTreeDocsScanned); + assertEquals(starTreeResults.keySet(), nonStarTreeResults.keySet()); + + long totalDigestSize = 0; + for (Map.Entry entry : starTreeResults.entrySet()) { + TDigest starTreeDigest = entry.getValue(); + TDigest nonStarTreeDigest = nonStarTreeResults.get(entry.getKey()); + assertEquals(starTreeDigest.size(), nonStarTreeDigest.size()); + assertEquals(starTreeDigest.quantile(0.75), nonStarTreeDigest.quantile(0.75), DELTA, ERROR_MESSAGE); + totalDigestSize += starTreeDigest.size(); + } + assertEquals(totalDigestSize, getExpectedTotalDigestSize()); + } + + protected long getExpectedTotalDigestSize() { + return NUM_ROWS; + } + protected String getAggregationQuery(int percentile) { return String.format("SELECT PERCENTILE%1$d(%2$s), PERCENTILETDIGEST%1$d(%2$s), PERCENTILETDIGEST%1$d(%3$s), " + "PERCENTILE(%2$s, %1$d), PERCENTILETDIGEST(%2$s, %1$d), PERCENTILETDIGEST(%3$s, %1$d), " @@ -258,6 +303,25 @@ private String getGroupByQuery(int percentile) { return String.format("%s GROUP BY %s", getAggregationQuery(percentile), GROUP_BY_COLUMN); } + private GroupByOperator getGroupByOperator(String query, boolean useStarTree) { + PinotQuery pinotQuery = CalciteSqlParser.compileToPinotQuery(query); + QueryContext queryContext = QueryContextConverterUtils.getQueryContext(pinotQuery); + queryContext.setSkipStarTree(!useStarTree); + return (GroupByOperator) PLAN_MAKER.makeSegmentPlanNode(new SegmentContext(getIndexSegment()), queryContext).run(); + } + + private Map getGroupByResults(GroupByResultsBlock resultsBlock) { + AggregationGroupByResult groupByResult = resultsBlock.getAggregationGroupByResult(); + assertNotNull(groupByResult); + Map results = new HashMap<>(); + Iterator groupKeyIterator = groupByResult.getGroupKeyIterator(); + while (groupKeyIterator.hasNext()) { + GroupKeyGenerator.GroupKey groupKey = groupKeyIterator.next(); + results.put((String) groupKey._keys[0], (TDigest) groupByResult.getResultForGroupId(0, groupKey._groupId)); + } + return results; + } + private void assertTDigest(TDigest tDigest, DoubleList doubleList) { for (int percentile = 0; percentile <= 100; percentile++) { double expected; diff --git a/pinot-perf/README.md b/pinot-perf/README.md index cb4d47850a0d..dd39de2d31ac 100644 --- a/pinot-perf/README.md +++ b/pinot-perf/README.md @@ -40,6 +40,12 @@ $ cd target/pinot-perf-pkg/bin $ ./pinot-BenchmarkDictionary.sh ``` +# Percentile aggregation benchmarks + +See the [100-million-row percentile aggregation benchmark](benchmark-results/percentile-aggregation-100m-results.md) +for sample `PERCENTILE` and `PERCENTILETDIGEST` queries, a sample star-tree index config, guidance on choosing between +the functions, reproducible JMH commands, and baseline-versus-optimized raw-value and star-tree results. + # Vector benchmark suite `org.apache.pinot.perf.BenchmarkVectorIndex` is the canonical entry point for Pinot's vector diff --git a/pinot-perf/benchmark-results/percentile-aggregation-100m-results.md b/pinot-perf/benchmark-results/percentile-aggregation-100m-results.md new file mode 100644 index 000000000000..b0bfe9dfa83d --- /dev/null +++ b/pinot-perf/benchmark-results/percentile-aggregation-100m-results.md @@ -0,0 +1,316 @@ + + +# Percentile Aggregation: 100-Million-Row Benchmark + +- Date: 2026-07-15 +- Baseline: `068b458883d6` +- Reducer baseline: `633aab66c6` +- Host: Apple M4 Pro (14 cores, 24 GiB), macOS 26.5.2 aarch64, OpenJDK 21.0.10 + +These benchmarks measure single-threaded aggregation-function performance over deterministic pseudo-random `DOUBLE` +values. The raw-value workloads exercise the production block aggregation API over 100 million rows. The star-tree +query workload merges 10,000 stored TDigest entries that represent 100 million source rows, and the construction +workload isolates the raw-value aggregator used to build those entries. These numbers describe aggregation kernels on +one machine, not distributed end-to-end query latency or complete segment-generation time. + +## User guide + +The optimizations are automatic and require no query changes. Use exact `PERCENTILE` when the result must be an exact +order statistic. Use `PERCENTILETDIGEST` when bounded aggregation state and mergeable intermediate results are more +important than exactness. A star-tree can pre-aggregate `PERCENTILETDIGEST`; exact `PERCENTILE` cannot use a stored +star-tree percentile because its aggregation state retains every input value. + +| Function | Result | Aggregation state | +|---|---|---| +| `PERCENTILE(metric, 75)` | Exact P75 order statistic | Retains all input values | +| `PERCENTILETDIGEST(metric, 75)` | Approximate P75, default compression `100` | Bounded TDigest | +| `PERCENTILETDIGEST(metric, 75, 200)` | Approximate P75 with higher compression | Larger TDigest with potentially better accuracy | + +### Sample queries + +```sql +-- Exact P75. +SELECT PERCENTILE(latencyMs, 75) AS p75Exact +FROM myTable; + +-- Approximate P75 with the default TDigest compression of 100. +SELECT PERCENTILETDIGEST(latencyMs, 75) AS p75Approx +FROM myTable; + +-- Approximate P75 with an explicit compression factor. +SELECT PERCENTILETDIGEST(latencyMs, 75, 200) AS p75Approx +FROM myTable; + +-- Approximate P75 per region; a matching star-tree can serve the stored TDigest metric. +SELECT region, PERCENTILETDIGEST(latencyMs, 75) AS p75Approx +FROM myTable +GROUP BY region; +``` + +### Sample star-tree table config + +Add the TDigest function-column pair to a star-tree index when common queries group or filter on a stable set of +dimensions. This sample uses the default compression `100` and matches the two-argument `PERCENTILETDIGEST` query +shown above; a query with an explicit compression must match the compression configured for the index. + +```json +{ + "tableIndexConfig": { + "starTreeIndexConfigs": [ + { + "dimensionsSplitOrder": ["region", "deviceType"], + "skipStarNodeCreationForDimensions": [], + "functionColumnPairs": ["percentileTDigest__latencyMs"], + "maxLeafRecords": 10000 + } + ] + } +} +``` + +Choose the split order and leaf size for the table's query patterns and cardinalities. The benchmark below isolates +the stored-metric aggregation path after star-tree traversal; it does not measure planner, traversal, or forward-index +I/O. The optimized serialized merge applies to both non-grouped and group-by aggregation. Each serialized group keeps +its first digest pending and allocates primitive merge buffers only if another stored digest reaches that group; its +raw-value buffer remains lazy unless a raw value is added. + +## Results + +JMH 1.37 ran two forks, five one-second measurement iterations, one thread, an 8 GiB maximum heap, and the GC profiler +for each matched comparison. The raw-value and construction comparisons used two one-second warmup iterations; the +serialized star-tree query comparison used three. Each reported comparison contains ten measured iterations. + +### Exact percentile + +| Benchmark | Baseline (ms/op) | Optimized (ms/op) | Latency reduction | Speedup | +|---|---:|---:|---:|---:| +| Fresh result holder, aggregation plus P75 | 7,710.804 ± 287.081 | 764.626 ± 47.073 | 90.084% | 10.084x | +| Reused result holder, aggregation plus P75 | 7,499.571 ± 185.532 | 707.555 ± 16.087 | 90.565% | 10.599x | +| P75 extraction from 100 million values | 7,448.796 ± 75.498 | 618.253 ± 12.825 | 91.700% | 12.048x | + +The fresh-holder workload improved from 12.969 million to 130.783 million rows per second. Exact results were checked +bit-for-bit against a fully sorted oracle on every invocation. + +### TDigest percentile + +| Benchmark | Baseline (ms/op) | Optimized (ms/op) | Latency reduction | Speedup | +|---|---:|---:|---:|---:| +| Aggregation plus TDigest P75 | 5,677.240 ± 149.058 | 2,555.127 ± 49.673 | 54.994% | 2.222x | + +TDigest throughput improved from 17.614 million to 39.137 million rows per second. The optimized result was +`0.7499280847`, compared with the exact result `0.7499315464`, for an absolute error of `0.0000034617`. Neither TDigest +run triggered garbage collection; normalized allocation changed from 20,520 to 31,872 bytes per operation. + +### Star-tree TDigest query path + +| Benchmark | Baseline (ms/op) | Optimized (ms/op) | Latency reduction | Speedup | +|---|---:|---:|---:|---:| +| Merge 10,000 stored digests representing 100 million rows, then P75 | 85.658 ± 1.211 | 7.100 ± 0.186 | 91.711% | 12.065x | + +Normalized allocation fell from 198,223,964 bytes to 26,361 bytes per operation, and measured GC events fell from ten +to zero. The optimized digest retained an exact size of 100 million and produced P75 `0.7498860479`, versus the fully +sorted oracle `0.7499315464` (absolute error `0.0000454985`). Profiling showed the baseline repeatedly materializing, +shuffling, and sorting centroids as each stored digest was added. The optimized path decodes already-sorted serialized +centroids and linearly merges them into the accumulator. + +### Star-tree TDigest group-by query path + +This incremental comparison uses commit `4a950bde9455`, the preceding version of this change with the optimized +non-grouped path but the original group-by implementation, as its baseline. Every workload aggregates 10,000 stored +digests representing 100 million source rows. Entries are distributed round-robin, and the timed operation extracts +and validates every group result. + +| Groups | Stored digests per group | Baseline (ms/op) | Optimized (ms/op) | Latency reduction | Speedup | Baseline allocation (B/op) | Optimized allocation (B/op) | +|---:|---:|---:|---:|---:|---:|---:|---:| +| 100 | 100 | 119.244 ± 66.230 | 9.275 ± 0.161 | 92.222% | 12.856x | 197,636,620 | 2,894,952 | +| 1,000 | 10 | 96.039 ± 11.979 | 11.371 ± 0.257 | 88.159% | 8.446x | 192,278,646 | 28,821,662 | +| 10,000 | 1 | 12.107 ± 6.227 | 9.078 ± 0.412 | Not distinguishable | Parity | 136,520,147 | 137,480,126 | + +The 100- and 1,000-group workloads exercise centroid merging and reduce normalized allocation by 98.535% and 85.010%, +respectively. With 10,000 groups, each group receives exactly one stored digest, so there is no source digest merge to +optimize; extraction dominates and allocation is effectively unchanged (+0.703%). Its baseline and optimized latency +confidence intervals overlap, so the nominal difference between their means is not claimed as a speedup. That +high-cardinality case keeps the first +serialized digest pending, constructs only the final result digest, and does not allocate the accumulator's raw-value +buffer. The baseline had isolated latency outliers, reflected in its wider confidence intervals. + +### Server-local and distributed TDigest reduction + +The reducer benchmark runs on the same host with OpenJDK 25. It prebuilds segment digests and exact raw-value oracles +outside the timed region, creates fresh targets per invocation, and measures both the merge kernel and the complete +`ConcurrentIndexedTable` update path. `PROMOTED_LOCAL` models raw group-by segment results that arrive as ordinary +`MergingDigest` objects. `ACCUMULATOR_LOCAL` models materialized accumulator results from non-grouped and StarTree +segment execution. `ACCUMULATOR_WIRE` models lazy broker-side deserialization of standard TDigest bytes. + +The representative comparison uses 32 segment results, compression 100, uniform input, two forks, and five measured +iterations per fork. + +| Benchmark | Pairwise 3.2 (ms/op) | Promoted local | Accumulator local | Serialized/wire | +|---|---:|---:|---:|---:| +| Merge kernel | 0.2573 | 0.0325 (7.92x) | 0.0186 (13.82x) | 0.0220 (11.68x) | +| Full `IndexedTable` combine | 0.2683 | 0.0340 (7.90x) | 0.0201 (13.33x) | 0.0233 (11.52x) | +| Full combine plus final percentile extraction | 0.2983 | 0.0353 (8.46x) | 0.0201 (14.83x) | 0.0239 (12.47x) | + +Normalized kernel allocation fell from 205,162 B/op to 24,000 B/op for promoted local state (-88.3%), 11,936 B/op +for accumulator-local state (-94.2%), and 15,328 B/op for wire state (-92.5%). Direct percentile extraction reads the +primitive centroid buffers and does not construct a temporary source digest. + +The production-shaped workload uses 1,440 groups, seven TDigest metrics, fan-in 32, and distinct prebuilt source +objects for every group and metric. This removes the cache-locality advantage of the smaller shared-corpus kernel. + +| Path | Full combine (ms/op) | Speedup | Allocation (B/op) | Allocation reduction | +|---|---:|---:|---:|---:| +| TDigest 3.2 pairwise | 4,574.358 ± 309.333 | 1.00x | 2,074,920,845 | - | +| Promoted raw group-by state | 856.308 ± 176.240 | 5.34x | 248,823,888 | 88.01% | +| Materialized accumulator / StarTree-local | 222.986 ± 2.308 | 20.51x | 127,215,599 | 93.87% | +| Lazy distributed wire state | 283.015 ± 6.168 | 16.16x | 161,407,415 | 92.22% | + +A separate non-forked `/usr/bin/time -l` diagnostic over the same unique-source shape measured maximum RSS of +5,701 MB for pairwise, 5,773 MB for promoted local (+1.26%), 3,042 MB for accumulator-local (-46.65%), and 3,017 MB +for wire state (-47.09%). It is a peak-memory check only; the two-fork JMH results above are the latency measurements. + +The improvement holds across configured compression factors: + +| Compression | Promoted local | Accumulator local | Serialized/wire | +|---:|---:|---:|---:| +| 50 | 7.17x / -84.8% allocation | 12.14x / -90.5% | 10.98x / -88.8% | +| 100 | 7.67x / -86.7% allocation | 13.63x / -92.5% | 11.80x / -90.9% | +| 200 | 8.62x / -87.9% allocation | 15.27x / -93.6% | 12.75x / -92.1% | + +A 36-shape fan-in/distribution/order sweep covered fan-in 8/32/128; uniform, skewed, bimodal, and duplicate-heavy +inputs; and original, reversed, and seeded-random order across all four implementations (144 implementation/shape +cases). Every result retained exact total weight, finite positive centroids, and monotonic p0/p50/p75/p95/p99/p100. +Mean P75 error improved from 0.000175 for pairwise 3.2 to 0.000144 for every primitive-accumulator mode; the maximum +sample error was 0.001353 for pairwise and 0.000709 for the candidates. The focused correctness suite additionally +covers compression 20, 100, and 1,000 for the full parameter cross-product and round-trips every reduced result +through standard TDigest 3.2 bytes. + +#### TDigest 3.3 dependency-only experiment + +A separate build changed only the TDigest dependency from 3.2 to 3.3. The control harness supplies the same 128 +verbose centroids to both versions. Across 36 workload shapes, the geometric mean speedup was 4.47x in the kernel and +4.03x through `IndexedTable`, with 70.7% and 65.6% lower allocation. It is not used by this change: mean P75 error +increased 18.9x, the maximum sampled P75 error reached 0.01138, and mean resulting centroids changed from 123.1 to +42.2. A high-compression small-encoding mixed-version boundary also failed. In a separate 3.3 representative run, +`add(List.of(source))` was 2.27x slower than pairwise and bounded batches were 1.31-1.52x slower while allocating more. +The accepted reducer therefore stays on TDigest 3.2 and preserves standard TDigest wire bytes. + +### Star-tree TDigest construction kernel + +| Workload | Baseline (ms/op) | Optimized (ms/op) | Latency reduction | Speedup | +|---|---:|---:|---:|---:| +| 1 million rows, 1,000 rows per dimension group | 1,840.391 ± 47.275 | 53.770 ± 1.470 | 97.078% | 34.227x | + +The previous size-accounting call compressed the TDigest after every input value, defeating the digest's batching. +The optimized path maintains a safe serialized-size bound without compressing and defers compression until +serialization. A separate optimized 100-million-row run with 1,000 rows per group completed in 5,455.170 ms, or +18.331 million rows per second. The 100-million-row number is an actual optimized measurement; no unmeasured baseline +is extrapolated for it. +This construction benchmark includes one serialization per group but excludes dimension sorting and forward-index +I/O. + +## Reproduce + +Build the benchmark package from the repository root: + +```bash +./mvnw -pl pinot-perf -am clean package -DskipTests +``` + +Run the exact percentile benchmark: + +```bash +java -Xms4g -Xmx8g \ + -cp 'pinot-perf/target/pinot-perf-pkg/lib/*' \ + org.openjdk.jmh.Main \ + 'org.apache.pinot.perf.aggregation.BenchmarkPercentileAggregation.*' \ + -wi 2 -i 5 -f 2 -w 1s -r 1s -t 1 -to 30m -gc true -prof gc +``` + +Run the TDigest benchmark: + +```bash +java -Xms4g -Xmx8g \ + -cp 'pinot-perf/target/pinot-perf-pkg/lib/*' \ + org.openjdk.jmh.Main \ + 'org.apache.pinot.perf.aggregation.BenchmarkPercentileTDigestAggregation.*' \ + -wi 2 -i 5 -f 2 -w 1s -r 1s -t 1 -to 30m -gc true -prof gc +``` + +Run the stored star-tree TDigest query benchmark. Its 10,000 input entries represent 100 million source rows: + +```bash +java -Xms4g -Xmx8g \ + -cp 'pinot-perf/target/pinot-perf-pkg/lib/*' \ + org.openjdk.jmh.Main \ + 'org.apache.pinot.perf.aggregation.BenchmarkPercentileTDigestStarTreeAggregation.*' \ + -wi 3 -i 5 -f 2 -w 1s -r 1s -t 1 -to 30m -gc true -prof gc +``` + +Run the star-tree construction kernel over 100 million raw rows with 1,000 rows per dimension group: + +```bash +java -Xms4g -Xmx8g \ + -cp 'pinot-perf/target/pinot-perf-pkg/lib/*' \ + org.openjdk.jmh.Main \ + 'org.apache.pinot.perf.aggregation.BenchmarkPercentileTDigestValueAggregator.*' \ + -p _numRows=100000000 -p _rowsPerGroup=1000 \ + -wi 1 -i 3 -f 1 -w 1s -r 1s -t 1 -to 30m -gc true -prof gc +``` + +Run the representative reducer comparison: + +```bash +java -Xms4g -Xmx8g \ + -cp 'pinot-perf/target/pinot-perf-pkg/lib/*' \ + org.openjdk.jmh.Main \ + 'org.apache.pinot.perf.aggregation.BenchmarkPercentileTDigestCombine.*' \ + -p _numGroups=1 -p _numMetrics=1 -p _fanIn=32 -p _compression=100 \ + -p _distribution=UNIFORM -p _mergeOrder=ORIGINAL \ + -p _implementation=PAIRWISE,PROMOTED_LOCAL,ACCUMULATOR_LOCAL,ACCUMULATOR_WIRE \ + -p _sourceLayout=NATIVE -p _sourceReuse=SHARED \ + -wi 2 -i 5 -f 2 -w 1s -r 1s -t 1 -prof gc +``` + +Run the production-shaped combine with unique prebuilt source state: + +```bash +java -Xms4g -Xmx8g \ + -cp 'pinot-perf/target/pinot-perf-pkg/lib/*' \ + org.openjdk.jmh.Main \ + 'org.apache.pinot.perf.aggregation.BenchmarkPercentileTDigestCombine.combineIndexedTable$' \ + -p _numGroups=1440 -p _numMetrics=7 -p _fanIn=32 -p _compression=100 \ + -p _distribution=UNIFORM -p _mergeOrder=ORIGINAL \ + -p _implementation=PAIRWISE,PROMOTED_LOCAL,ACCUMULATOR_LOCAL,ACCUMULATOR_WIRE \ + -p _sourceLayout=NATIVE -p _sourceReuse=UNIQUE \ + -wi 2 -i 5 -f 2 -w 1s -r 1s -t 1 -prof gc +``` + +For the dependency-only experiment, change only `` in the root `pom.xml` from `3.2` to `3.3`, +rebuild `pinot-core` and `pinot-perf`, and run the same command for each build with `_implementation=PAIRWISE` and +`_sourceLayout=FIXED_VERBOSE`. Write each result with `-rf json -rff ` and construct the runtime +classpath with exactly one `t-digest` JAR; a package directory reused across builds can otherwise retain both JARs. + +The raw query benchmarks generate the same 100 million values from `SplittableRandom(42)` during trial setup. The +timed methods include result extraction and validation. Dataset generation and exact-oracle construction are outside +the measured region. The star-tree query benchmark similarly generates its stored digests during trial setup. The +construction benchmark generates a reusable deterministic value block before measurement. diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileAggregation.java b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileAggregation.java new file mode 100644 index 000000000000..327cf61b3879 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileAggregation.java @@ -0,0 +1,163 @@ +/** + * 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.perf.aggregation; + +import it.unimi.dsi.fastutil.doubles.DoubleArrayList; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.SplittableRandom; +import java.util.concurrent.TimeUnit; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.SyntheticBlockValSets; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.function.PercentileAggregationFunction; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; + +/// Measures exact percentile aggregation over 100 million rows on one machine. +/// +/// The fresh aggregation benchmark includes result-list allocation and growth while feeding 100 million deterministic +/// pseudo-random values through the production aggregation function in Pinot-sized 10,000-row blocks. The reused +/// aggregation benchmark isolates the steady-state compute kernel, and the extraction benchmark isolates the final +/// order-statistic calculation. Every invocation checks its result against an exact precomputed oracle. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Warmup(iterations = 2, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork(value = 2, jvmArgsAppend = {"-Xms4g", "-Xmx8g"}) +@Threads(1) +public class BenchmarkPercentileAggregation { + private static final int NUM_ROWS = 100_000_000; + private static final int BLOCK_SIZE = DocIdSetPlanNode.MAX_DOC_PER_CALL; + private static final int NUM_BLOCKS = NUM_ROWS / BLOCK_SIZE; + private static final int PERCENTILE_INDEX = (int) ((long) NUM_ROWS * 75 / 100); + private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("col"); + + /// State for the end-to-end aggregation-function benchmark. + @State(Scope.Benchmark) + public static class AggregationState { + private PercentileAggregationFunction _function; + private AggregationResultHolder _resultHolder; + private List> _blockValSetMaps; + private double _expectedResult; + + @Setup(Level.Trial) + public void setUp() { + if (NUM_ROWS % BLOCK_SIZE != 0) { + throw new IllegalStateException("NUM_ROWS must be divisible by BLOCK_SIZE"); + } + SplittableRandom random = new SplittableRandom(42); + _blockValSetMaps = new ArrayList<>(NUM_BLOCKS); + double[] expectedValues = new double[NUM_ROWS]; + for (int blockId = 0; blockId < NUM_BLOCKS; blockId++) { + double[] values = new double[BLOCK_SIZE]; + for (int i = 0; i < BLOCK_SIZE; i++) { + values[i] = random.nextDouble(); + } + _blockValSetMaps.add(Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, values))); + System.arraycopy(values, 0, expectedValues, blockId * BLOCK_SIZE, BLOCK_SIZE); + } + Arrays.sort(expectedValues); + _expectedResult = expectedValues[PERCENTILE_INDEX]; + _function = new PercentileAggregationFunction(EXPRESSION, 75.0, false); + _resultHolder = _function.createAggregationResultHolder(); + } + + @Setup(Level.Invocation) + public void resetResult() { + DoubleArrayList values = _resultHolder.getResult(); + if (values != null) { + values.clear(); + } + } + } + + /// State for the final percentile extraction benchmark. + @State(Scope.Benchmark) + public static class ExtractionState { + private PercentileAggregationFunction _function; + private double[] _sourceValues; + private double[] _workingValues; + private DoubleArrayList _valueList; + private double _expectedResult; + + @Setup(Level.Trial) + public void setUp() { + SplittableRandom random = new SplittableRandom(42); + _sourceValues = new double[NUM_ROWS]; + for (int i = 0; i < NUM_ROWS; i++) { + _sourceValues[i] = random.nextDouble(); + } + _workingValues = _sourceValues.clone(); + Arrays.sort(_workingValues); + _expectedResult = _workingValues[PERCENTILE_INDEX]; + _valueList = DoubleArrayList.wrap(_workingValues); + _function = new PercentileAggregationFunction(EXPRESSION, 75.0, false); + } + + @Setup(Level.Invocation) + public void resetValues() { + System.arraycopy(_sourceValues, 0, _workingValues, 0, NUM_ROWS); + } + } + + @Benchmark + public double aggregatePercentile75Fresh(AggregationState state) { + return aggregate(state, state._function.createAggregationResultHolder()); + } + + @Benchmark + public double aggregatePercentile75Reused(AggregationState state) { + return aggregate(state, state._resultHolder); + } + + private static double aggregate(AggregationState state, AggregationResultHolder resultHolder) { + for (int i = 0; i < NUM_BLOCKS; i++) { + state._function.aggregate(BLOCK_SIZE, resultHolder, state._blockValSetMaps.get(i)); + } + DoubleArrayList result = state._function.extractAggregationResult(resultHolder); + return verify(state._function.extractFinalResult(result), state._expectedResult); + } + + @Benchmark + public double extractPercentile75(ExtractionState state) { + return verify(state._function.extractFinalResult(state._valueList), state._expectedResult); + } + + private static double verify(double actual, double expected) { + if (Double.compare(actual, expected) != 0) { + throw new IllegalStateException("Unexpected percentile: " + actual + ", expected: " + expected); + } + return actual; + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestAggregation.java b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestAggregation.java new file mode 100644 index 000000000000..fc23656f79c3 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestAggregation.java @@ -0,0 +1,114 @@ +/** + * 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.perf.aggregation; + +import com.tdunning.math.stats.TDigest; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Map; +import java.util.SplittableRandom; +import java.util.concurrent.TimeUnit; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.SyntheticBlockValSets; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.function.PercentileTDigestAggregationFunction; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; + +/// Measures percentile TDigest aggregation over 100 million rows on one machine. +/// +/// Each invocation creates a fresh result holder and feeds 100 million deterministic pseudo-random values through the +/// production aggregation function in Pinot-sized 10,000-row blocks. The result is checked for its row count and +/// accuracy against an exact precomputed 75th percentile. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Warmup(iterations = 2, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork(value = 2, jvmArgsAppend = {"-Xms4g", "-Xmx8g"}) +@Threads(1) +public class BenchmarkPercentileTDigestAggregation { + private static final int NUM_ROWS = 100_000_000; + private static final int BLOCK_SIZE = DocIdSetPlanNode.MAX_DOC_PER_CALL; + private static final int NUM_BLOCKS = NUM_ROWS / BLOCK_SIZE; + private static final int PERCENTILE_INDEX = (int) ((long) NUM_ROWS * 75 / 100); + private static final double MAX_ABSOLUTE_ERROR = 0.001; + private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("col"); + + /// State for the end-to-end aggregation-function benchmark. + @State(Scope.Benchmark) + public static class AggregationState { + private PercentileTDigestAggregationFunction _function; + private List> _blockValSetMaps; + private double _expectedResult; + + @Setup(Level.Trial) + public void setUp() { + if (NUM_ROWS % BLOCK_SIZE != 0) { + throw new IllegalStateException("NUM_ROWS must be divisible by BLOCK_SIZE"); + } + SplittableRandom random = new SplittableRandom(42); + _blockValSetMaps = new ArrayList<>(NUM_BLOCKS); + double[] expectedValues = new double[NUM_ROWS]; + for (int blockId = 0; blockId < NUM_BLOCKS; blockId++) { + double[] values = new double[BLOCK_SIZE]; + for (int i = 0; i < BLOCK_SIZE; i++) { + values[i] = random.nextDouble(); + } + _blockValSetMaps.add(Map.of(EXPRESSION, SyntheticBlockValSets.Double.create(null, values))); + System.arraycopy(values, 0, expectedValues, blockId * BLOCK_SIZE, BLOCK_SIZE); + } + Arrays.sort(expectedValues); + _expectedResult = expectedValues[PERCENTILE_INDEX]; + _function = new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + } + } + + @Benchmark + public double aggregatePercentileTDigest75(AggregationState state) { + AggregationResultHolder resultHolder = state._function.createAggregationResultHolder(); + for (int i = 0; i < NUM_BLOCKS; i++) { + state._function.aggregate(BLOCK_SIZE, resultHolder, state._blockValSetMaps.get(i)); + } + TDigest result = state._function.extractAggregationResult(resultHolder); + if (result.size() != NUM_ROWS) { + throw new IllegalStateException("Unexpected TDigest size: " + result.size() + ", expected: " + NUM_ROWS); + } + return verify(state._function.extractFinalResult(result), state._expectedResult); + } + + private static double verify(double actual, double expected) { + if (!Double.isFinite(actual) || Math.abs(actual - expected) > MAX_ABSOLUTE_ERROR) { + throw new IllegalStateException("Unexpected percentile: " + actual + ", expected: " + expected); + } + return actual; + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestCombine.java b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestCombine.java new file mode 100644 index 000000000000..89aad4208063 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestCombine.java @@ -0,0 +1,978 @@ +/** + * 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.perf.aggregation; + +import com.tdunning.math.stats.Centroid; +import com.tdunning.math.stats.MergingDigest; +import com.tdunning.math.stats.TDigest; +import java.nio.ByteBuffer; +import java.security.CodeSource; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.IdentityHashMap; +import java.util.Iterator; +import java.util.List; +import java.util.Map; +import java.util.SplittableRandom; +import java.util.StringJoiner; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import org.apache.pinot.common.CustomObject; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.core.common.ObjectSerDeUtils; +import org.apache.pinot.core.data.table.ConcurrentIndexedTable; +import org.apache.pinot.core.data.table.IndexedTable; +import org.apache.pinot.core.data.table.Key; +import org.apache.pinot.core.data.table.Record; +import org.apache.pinot.core.query.aggregation.function.AggregationFunction; +import org.apache.pinot.core.query.aggregation.function.AggregationFunction.SerializedIntermediateResult; +import org.apache.pinot.core.query.aggregation.function.PercentileTDigestAggregationFunction; +import org.apache.pinot.core.query.request.context.QueryContext; +import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils; +import org.openjdk.jmh.annotations.AuxCounters; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.infra.BenchmarkParams; + +/// Measures the TDigest reducer kernel and the complete [ConcurrentIndexedTable] group-by combine path. +/// +/// Source digests and the exact raw-value oracle are built once per trial. Each invocation receives fresh merge +/// targets so a previous invocation cannot change its input. Accuracy and centroid statistics come from an untimed +/// 32-merge quality sweep. Run with JMH's GC profiler to report allocation and GC metrics, for example `-prof gc`. +/// Error counters report absolute error in billionths because all generated distributions are bounded to `[0, 1]`. +/// +/// The TDigest dependency version is deliberately a build-level dimension. Compare identical `PAIRWISE` runs from a +/// 3.2 build and a 3.3 build. Use `_sourceLayout=FIXED_VERBOSE` to hold the serialized centroid shape constant for the +/// attributable pairwise experiment; the default `NATIVE` mode includes version-specific source compression. +/// `SINGLETON_LIST` is accepted as a command-line JMH parameter for the 3.3 experiment, but is excluded from the +/// default matrix and rejected on TDigest 3.2 because that version corrupts non-empty targets in `add(List)`. +/// `PROMOTED_LOCAL` measures promotion of raw group-by [MergingDigest] results. `ACCUMULATOR_LOCAL` models +/// materialized accumulator results from non-grouped and StarTree segment execution. `ACCUMULATOR_WIRE` keeps +/// deserialized sources lazy to model broker reduction. Use `_sourceReuse=UNIQUE` for production-shaped runs so +/// every group and metric has distinct prebuilt source state; the default shared corpus keeps the full matrix small. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Warmup(iterations = 2, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork(value = 2, jvmArgsAppend = {"-Xms4g", "-Xmx8g"}) +@Threads(1) +@State(Scope.Thread) +public class BenchmarkPercentileTDigestCombine { + private static final int VALUES_PER_DIGEST = 10_000; + private static final int QUALITY_SAMPLES = 32; + private static final double ERROR_SCALE = 1_000_000_000.0; + private static final Pattern TDIGEST_JAR_VERSION = Pattern.compile("t-digest-(\\d+)\\.(\\d+)(?:\\.[^/]*)?\\.jar"); + + public enum Distribution { + UNIFORM, + SKEWED, + BIMODAL, + DUPLICATE_HEAVY + } + + public enum MergeOrder { + ORIGINAL, + REVERSED, + RANDOMIZED + } + + public enum SourceLayout { + NATIVE, + FIXED_VERBOSE + } + + public enum SourceReuse { + SHARED, + UNIQUE + } + + public enum Implementation { + PAIRWISE(0), + PROMOTED_LOCAL(0), + ACCUMULATOR_LOCAL(0), + ACCUMULATOR_WIRE(0), + SINGLETON_LIST(1), + BATCH_8(8), + BATCH_16(16), + BATCH_32(32); + + private final int _batchSize; + + Implementation(int batchSize) { + _batchSize = batchSize; + } + + private boolean isBatching() { + return _batchSize >= 8; + } + } + + @Param({"1", "1440"}) + private int _numGroups; + + @Param({"1", "7"}) + private int _numMetrics; + + @Param({"8", "32", "128"}) + private int _fanIn; + + @Param({"50", "100", "200"}) + private int _compression; + + @Param + private Distribution _distribution; + + @Param + private MergeOrder _mergeOrder; + + @Param({"NATIVE"}) + private SourceLayout _sourceLayout; + + @Param({"SHARED"}) + private SourceReuse _sourceReuse; + + @Param({"PAIRWISE", "PROMOTED_LOCAL", "ACCUMULATOR_LOCAL", "ACCUMULATOR_WIRE"}) + private Implementation _implementation; + + private ReducerAggregationFunction[] _functions; + private TDigest[] _sources; + private byte[][] _serializedSources; + private TDigest[] _firstSources; + private int[] _sourceOrder; + private double _averageInputCentroidCount; + private double _expectedP75; + private double _expectedP95; + private double _expectedP99; + private QueryContext _queryContext; + private DataSchema _dataSchema; + private ExecutorService _executorService; + private QualityStats _qualityStats; + + @Setup(Level.Trial) + public void setUpTrial() { + rejectUnsafeSingletonListMerge(); + buildSourcesAndOracle(); + buildMergeOrder(); + buildQuery(); + _executorService = Executors.newSingleThreadExecutor(); + buildQualityStats(); + } + + @Setup(Level.Invocation) + public void setUpInvocation() { + prepareFirstSources(); + } + + @TearDown(Level.Trial) + public void tearDownTrial() { + if (_executorService != null) { + _executorService.shutdownNow(); + } + } + + @Benchmark + public TDigest[] mergeKernel(QualityCounters counters) { + resetPending(); + TDigest[] targets = new TDigest[_numGroups * _numMetrics]; + for (int targetIndex = 0; targetIndex < targets.length; targetIndex++) { + targets[targetIndex] = _firstSources[targetIndex]; + } + for (int position = 1; position < _fanIn; position++) { + for (int groupId = 0; groupId < _numGroups; groupId++) { + int targetOffset = groupId * _numMetrics; + for (int metricId = 0; metricId < _numMetrics; metricId++) { + int targetIndex = targetOffset + metricId; + targets[targetIndex] = _functions[metricId].merge(targets[targetIndex], + _sources[sourceSlot(position, targetIndex)]); + } + } + } + flushPending(); + counters.record(_qualityStats); + return targets; + } + + @Benchmark + public IndexedTable combineIndexedTable(QualityCounters counters) { + IndexedTable table = buildIndexedTable(false); + counters.record(_qualityStats); + return table; + } + + @Benchmark + public IndexedTable combineIndexedTableAndExtract(QualityCounters counters) { + IndexedTable table = buildIndexedTable(true); + counters.record(_qualityStats); + return table; + } + + private IndexedTable buildIndexedTable(boolean extractFinalResult) { + resetPending(); + DataSchema dataSchema = extractFinalResult + ? new DataSchema(_dataSchema.getColumnNames().clone(), _dataSchema.getColumnDataTypes().clone()) : _dataSchema; + IndexedTable table = new ConcurrentIndexedTable(dataSchema, false, _queryContext, _numGroups, + Integer.MAX_VALUE, Integer.MAX_VALUE, _numGroups, _executorService); + for (int groupId = 0; groupId < _numGroups; groupId++) { + Object[] values = new Object[_numMetrics + 1]; + values[0] = groupId; + for (int metricId = 0; metricId < _numMetrics; metricId++) { + int targetIndex = groupId * _numMetrics + metricId; + values[metricId + 1] = _firstSources[targetIndex]; + } + table.upsert(new Key(new Object[]{groupId}), new Record(values)); + } + for (int position = 1; position < _fanIn; position++) { + for (int groupId = 0; groupId < _numGroups; groupId++) { + Object[] values = new Object[_numMetrics + 1]; + values[0] = groupId; + for (int metricId = 0; metricId < _numMetrics; metricId++) { + int targetIndex = groupId * _numMetrics + metricId; + values[metricId + 1] = _sources[sourceSlot(position, targetIndex)]; + } + table.upsert(new Key(new Object[]{groupId}), new Record(values)); + } + } + flushPending(); + table.finish(false, extractFinalResult); + return table; + } + + private void buildSourcesAndOracle() { + _sources = new TDigest[_fanIn]; + _serializedSources = new byte[_fanIn][]; + double[] rawValues = new double[_fanIn * VALUES_PER_DIGEST]; + int totalInputCentroids = 0; + for (int sourceId = 0; sourceId < _fanIn; sourceId++) { + MergingDigest nativeDigest = null; + if (_sourceLayout == SourceLayout.NATIVE) { + nativeDigest = _implementation == Implementation.PAIRWISE ? new MergingDigest(_compression) + : new ImmutableMergingDigest(_compression); + } + double[] sourceValues = new double[VALUES_PER_DIGEST]; + SplittableRandom random = new SplittableRandom(0x6A09E667F3BCC909L + sourceId); + int valueOffset = sourceId * VALUES_PER_DIGEST; + for (int valueId = 0; valueId < VALUES_PER_DIGEST; valueId++) { + double value = nextValue(random, valueId); + if (nativeDigest != null) { + nativeDigest.add(value); + } + sourceValues[valueId] = value; + rawValues[valueOffset + valueId] = value; + } + TDigest digest; + byte[] fixedBytes = null; + if (nativeDigest != null) { + if (nativeDigest instanceof ImmutableMergingDigest immutableDigest) { + immutableDigest.freeze(); + } else { + nativeDigest.compress(); + } + digest = nativeDigest; + } else { + Arrays.sort(sourceValues); + FixedTDigest fixedDigest = new FixedTDigest(sourceValues, _compression); + digest = fixedDigest; + fixedBytes = fixedDigest.bytes(); + } + if (isAccumulatorSource()) { + PercentileTDigestAggregationFunction function = new PercentileTDigestAggregationFunction( + ExpressionContext.forIdentifier("metric"), 75.0, _compression, false); + SerializedIntermediateResult serialized = function.serializeIntermediateResult(digest); + byte[] bytes = serialized.getBytes(); + _serializedSources[sourceId] = bytes; + digest = function.deserializeIntermediateResult( + new CustomObject(serialized.getType(), ByteBuffer.wrap(bytes))); + if (_implementation == Implementation.ACCUMULATOR_LOCAL) { + digest.compress(); + } + } else { + _serializedSources[sourceId] = fixedBytes != null ? fixedBytes + : ObjectSerDeUtils.TDIGEST_SER_DE.serialize(digest); + } + _sources[sourceId] = digest; + totalInputCentroids += digest.centroidCount(); + } + _averageInputCentroidCount = (double) totalInputCentroids / _fanIn; + Arrays.sort(rawValues); + _expectedP75 = exactQuantile(rawValues, 0.75); + _expectedP95 = exactQuantile(rawValues, 0.95); + _expectedP99 = exactQuantile(rawValues, 0.99); + expandSources(); + } + + private void expandSources() { + if (_sourceReuse == SourceReuse.SHARED) { + return; + } + int numTargets = _numGroups * _numMetrics; + byte[][] baseSerializedSources = _serializedSources; + _sources = new TDigest[_fanIn * numTargets]; + _serializedSources = new byte[_sources.length][]; + PercentileTDigestAggregationFunction function = new PercentileTDigestAggregationFunction( + ExpressionContext.forIdentifier("metric"), 75.0, _compression, false); + for (int sourceId = 0; sourceId < _fanIn; sourceId++) { + byte[] sourceBytes = baseSerializedSources[sourceId]; + for (int targetIndex = 0; targetIndex < numTargets; targetIndex++) { + int expandedIndex = sourceId * numTargets + targetIndex; + byte[] bytes = sourceBytes.clone(); + if (isAccumulatorSource()) { + _serializedSources[expandedIndex] = bytes; + TDigest source = function.deserializeIntermediateResult( + new CustomObject(ObjectSerDeUtils.ObjectType.TDigest.getValue(), ByteBuffer.wrap(bytes))); + if (_implementation == Implementation.ACCUMULATOR_LOCAL) { + source.compress(); + } + _sources[expandedIndex] = source; + } else { + _serializedSources[expandedIndex] = bytes; + _sources[expandedIndex] = _sourceLayout == SourceLayout.FIXED_VERBOSE ? new FixedTDigest(bytes) + : ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytes); + } + } + } + } + + private static byte[] createFixedVerboseBytes(double[] sortedValues, double compression) { + int numCentroids = Math.min(sortedValues.length, (int) Math.ceil(compression * 1.28)); + ByteBuffer buffer = ByteBuffer.allocate(32 + numCentroids * 16); + buffer.putInt(1); + buffer.putDouble(sortedValues[0]); + buffer.putDouble(sortedValues[sortedValues.length - 1]); + buffer.putDouble(compression); + buffer.putInt(numCentroids); + for (int centroidId = 0; centroidId < numCentroids; centroidId++) { + int start = centroidId * sortedValues.length / numCentroids; + int end = (centroidId + 1) * sortedValues.length / numCentroids; + double sum = 0.0; + for (int valueId = start; valueId < end; valueId++) { + sum += sortedValues[valueId]; + } + double weight = end - start; + buffer.putDouble(weight); + buffer.putDouble(sum / weight); + } + return buffer.array(); + } + + private void buildMergeOrder() { + _sourceOrder = new int[_fanIn]; + for (int i = 0; i < _fanIn; i++) { + _sourceOrder[i] = i; + } + if (_mergeOrder == MergeOrder.REVERSED) { + for (int left = 0, right = _fanIn - 1; left < right; left++, right--) { + int value = _sourceOrder[left]; + _sourceOrder[left] = _sourceOrder[right]; + _sourceOrder[right] = value; + } + } else if (_mergeOrder == MergeOrder.RANDOMIZED) { + SplittableRandom random = new SplittableRandom(0xBB67AE8584CAA73BL); + for (int i = _fanIn - 1; i > 0; i--) { + int other = random.nextInt(i + 1); + int value = _sourceOrder[i]; + _sourceOrder[i] = _sourceOrder[other]; + _sourceOrder[other] = value; + } + } + } + + private void buildQuery() { + StringJoiner selectExpressions = new StringJoiner(", "); + String[] columnNames = new String[_numMetrics + 1]; + ColumnDataType[] columnDataTypes = new ColumnDataType[_numMetrics + 1]; + columnNames[0] = "groupKey"; + columnDataTypes[0] = ColumnDataType.INT; + for (int metricId = 0; metricId < _numMetrics; metricId++) { + String metricName = "metric" + metricId; + selectExpressions.add("PERCENTILETDIGEST(" + metricName + ", 75, " + _compression + ")"); + columnNames[metricId + 1] = metricName; + columnDataTypes[metricId + 1] = ColumnDataType.OBJECT; + } + _queryContext = QueryContextConverterUtils.getQueryContext( + "SELECT " + selectExpressions + " FROM testTable GROUP BY groupKey"); + AggregationFunction[] aggregationFunctions = _queryContext.getAggregationFunctions(); + if (aggregationFunctions == null || aggregationFunctions.length != _numMetrics) { + throw new IllegalStateException("Unexpected aggregation functions in benchmark query"); + } + _functions = new ReducerAggregationFunction[_numMetrics]; + for (int metricId = 0; metricId < _numMetrics; metricId++) { + ReducerAggregationFunction function = new ReducerAggregationFunction( + ExpressionContext.forIdentifier("metric" + metricId), _compression, _implementation); + _functions[metricId] = function; + aggregationFunctions[metricId] = function; + } + _dataSchema = new DataSchema(columnNames, columnDataTypes); + } + + private void buildQualityStats() { + double totalResultingCentroids = 0.0; + double maxResultingCentroids = 0.0; + double totalP75Error = 0.0; + double maxP75Error = 0.0; + double totalP95Error = 0.0; + double maxP95Error = 0.0; + double totalP99Error = 0.0; + double maxP99Error = 0.0; + for (int sampleId = 0; sampleId < QUALITY_SAMPLES; sampleId++) { + resetPending(); + TDigest digest = newFirstSource(sourceSlot(0, 0), 0); + for (int position = 1; position < _fanIn; position++) { + digest = _functions[0].merge(digest, _sources[sourceSlot(position, 0)]); + } + flushPending(); + digest.compress(); + validateDigest(digest); + double p75Error = Math.abs(digest.quantile(0.75) - _expectedP75) * ERROR_SCALE; + double p95Error = Math.abs(digest.quantile(0.95) - _expectedP95) * ERROR_SCALE; + double p99Error = Math.abs(digest.quantile(0.99) - _expectedP99) * ERROR_SCALE; + double resultingCentroids = digest.centroidCount(); + totalResultingCentroids += resultingCentroids; + maxResultingCentroids = Math.max(maxResultingCentroids, resultingCentroids); + totalP75Error += p75Error; + maxP75Error = Math.max(maxP75Error, p75Error); + totalP95Error += p95Error; + maxP95Error = Math.max(maxP95Error, p95Error); + totalP99Error += p99Error; + maxP99Error = Math.max(maxP99Error, p99Error); + } + _qualityStats = new QualityStats(_averageInputCentroidCount, totalResultingCentroids / QUALITY_SAMPLES, + maxResultingCentroids, totalP75Error / QUALITY_SAMPLES, maxP75Error, totalP95Error / QUALITY_SAMPLES, + maxP95Error, totalP99Error / QUALITY_SAMPLES, maxP99Error); + prepareFirstSources(); + validateIndexedTable(buildIndexedTable(false)); + } + + private void validateIndexedTable(IndexedTable table) { + if (table.size() != _numGroups) { + throw new IllegalStateException("Unexpected IndexedTable size: " + table.size() + ", expected: " + _numGroups); + } + Iterator iterator = table.iterator(); + while (iterator.hasNext()) { + Record record = iterator.next(); + Object[] values = record.getValues(); + for (int metricId = 0; metricId < _numMetrics; metricId++) { + validateDigest((TDigest) values[metricId + 1]); + } + } + } + + private void prepareFirstSources() { + int numTargets = _numGroups * _numMetrics; + _firstSources = new TDigest[numTargets]; + for (int targetIndex = 0; targetIndex < numTargets; targetIndex++) { + _firstSources[targetIndex] = newFirstSource(sourceSlot(0, targetIndex), targetIndex % _numMetrics); + } + } + + private TDigest newFirstSource(int sourceIndex, int metricId) { + if (isAccumulatorSource()) { + byte[] bytes = _serializedSources[sourceIndex]; + TDigest target = _functions[metricId].deserializeIntermediateResult( + new CustomObject(ObjectSerDeUtils.ObjectType.TDigest.getValue(), ByteBuffer.wrap(bytes))); + if (_implementation == Implementation.ACCUMULATOR_LOCAL) { + target.compress(); + } + return target; + } + return ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(_serializedSources[sourceIndex]); + } + + private boolean isAccumulatorSource() { + return _implementation == Implementation.ACCUMULATOR_LOCAL + || _implementation == Implementation.ACCUMULATOR_WIRE; + } + + private int sourceSlot(int position, int targetIndex) { + int sourceId = _sourceOrder[position]; + return _sourceReuse == SourceReuse.SHARED ? sourceId : sourceId * _numGroups * _numMetrics + targetIndex; + } + + private void validateDigest(TDigest digest) { + long expectedWeight = expectedWeight(); + if (digest.size() != expectedWeight) { + throw new IllegalStateException( + "Unexpected TDigest weight: " + digest.size() + ", expected: " + expectedWeight); + } + long centroidWeight = 0L; + double previousMean = Double.NEGATIVE_INFINITY; + for (Centroid centroid : digest.centroids()) { + if (!Double.isFinite(centroid.mean()) || centroid.count() <= 0 || centroid.mean() + 1e-12 < previousMean) { + throw new IllegalStateException("Invalid centroid: previousMean=" + previousMean + ", mean=" + centroid.mean() + + ", weight=" + centroid.count()); + } + centroidWeight += centroid.count(); + previousMean = centroid.mean(); + } + if (centroidWeight != expectedWeight) { + throw new IllegalStateException( + "Unexpected centroid weight: " + centroidWeight + ", expected: " + expectedWeight); + } + double[] quantiles = {digest.quantile(0.0), digest.quantile(0.5), digest.quantile(0.75), + digest.quantile(0.95), digest.quantile(0.99), digest.quantile(1.0)}; + for (int i = 0; i < quantiles.length; i++) { + if (!Double.isFinite(quantiles[i]) || (i > 0 && quantiles[i] + 1e-12 < quantiles[i - 1])) { + throw new IllegalStateException("Invalid quantiles: " + Arrays.toString(quantiles)); + } + } + } + + private void flushPending() { + for (ReducerAggregationFunction function : _functions) { + function.flushAll(); + } + } + + private void resetPending() { + for (ReducerAggregationFunction function : _functions) { + function.resetPending(); + } + } + + private long expectedWeight() { + return (long) _fanIn * VALUES_PER_DIGEST; + } + + private double nextValue(SplittableRandom random, int valueId) { + switch (_distribution) { + case UNIFORM: + return random.nextDouble(); + case SKEWED: + return Math.pow(random.nextDouble(), 8.0); + case BIMODAL: + double mode = (valueId & 1) == 0 ? 0.2 : 0.8; + double radius = Math.sqrt(-2.0 * Math.log(Math.max(Double.MIN_NORMAL, random.nextDouble()))); + double normal = radius * Math.cos(2.0 * Math.PI * random.nextDouble()); + return Math.max(0.0, Math.min(1.0, mode + normal * 0.04)); + case DUPLICATE_HEAVY: + int bucket = random.nextInt(100); + if (bucket < 80) { + return 0.5; + } else if (bucket < 90) { + return 0.1; + } else if (bucket < 98) { + return 0.9; + } + return random.nextDouble(); + default: + throw new IllegalStateException("Unsupported distribution: " + _distribution); + } + } + + private static double exactQuantile(double[] sortedValues, double quantile) { + double index = quantile * (sortedValues.length - 1); + int lower = (int) index; + int upper = Math.min(lower + 1, sortedValues.length - 1); + double fraction = index - lower; + return sortedValues[lower] + fraction * (sortedValues[upper] - sortedValues[lower]); + } + + private void rejectUnsafeSingletonListMerge() { + if (_sourceLayout == SourceLayout.FIXED_VERBOSE + && (_implementation == Implementation.SINGLETON_LIST || _implementation.isBatching())) { + throw new IllegalArgumentException("FIXED_VERBOSE sources do not support list or batch merging"); + } + if (_sourceReuse == SourceReuse.UNIQUE + && (_implementation == Implementation.SINGLETON_LIST || _implementation.isBatching())) { + throw new IllegalArgumentException("UNIQUE sources do not support list or batch merging"); + } + if (_implementation != Implementation.SINGLETON_LIST) { + return; + } + CodeSource codeSource = TDigest.class.getProtectionDomain().getCodeSource(); + String location = codeSource != null ? codeSource.getLocation().toString() : ""; + Matcher matcher = TDIGEST_JAR_VERSION.matcher(location); + if (!matcher.find()) { + throw new IllegalStateException("Cannot determine TDigest version from: " + location); + } + int major = Integer.parseInt(matcher.group(1)); + int minor = Integer.parseInt(matcher.group(2)); + if (major < 3 || (major == 3 && minor < 3)) { + throw new IllegalArgumentException( + "SINGLETON_LIST is unsafe on TDigest " + major + "." + minor + "; use TDigest 3.3 or newer"); + } + } + + private static final class QualityStats { + private final double _inputCentroidCount; + private final double _resultingCentroidCount; + private final double _maxResultingCentroidCount; + private final double _p75ErrorPpb; + private final double _maxP75ErrorPpb; + private final double _p95ErrorPpb; + private final double _maxP95ErrorPpb; + private final double _p99ErrorPpb; + private final double _maxP99ErrorPpb; + + private QualityStats(double inputCentroidCount, double resultingCentroidCount, double maxResultingCentroidCount, + double p75ErrorPpb, double maxP75ErrorPpb, double p95ErrorPpb, double maxP95ErrorPpb, double p99ErrorPpb, + double maxP99ErrorPpb) { + _inputCentroidCount = inputCentroidCount; + _resultingCentroidCount = resultingCentroidCount; + _maxResultingCentroidCount = maxResultingCentroidCount; + _p75ErrorPpb = p75ErrorPpb; + _maxP75ErrorPpb = maxP75ErrorPpb; + _p95ErrorPpb = p95ErrorPpb; + _maxP95ErrorPpb = maxP95ErrorPpb; + _p99ErrorPpb = p99ErrorPpb; + _maxP99ErrorPpb = maxP99ErrorPpb; + } + } + + @State(Scope.Thread) + @AuxCounters(AuxCounters.Type.EVENTS) + public static class QualityCounters { + private double _aggregationDivisor; + private long _samples; + private double _inputCentroidCount; + private double _resultingCentroidCount; + private double _maxResultingCentroidCount; + private double _p75ErrorPpb; + private double _maxP75ErrorPpb; + private double _p95ErrorPpb; + private double _maxP95ErrorPpb; + private double _p99ErrorPpb; + private double _maxP99ErrorPpb; + + @Setup(Level.Trial) + public void setUp(BenchmarkParams benchmarkParams) { + _aggregationDivisor = (double) benchmarkParams.getMeasurement().getCount() + * Math.max(1, benchmarkParams.getForks()); + } + + @Setup(Level.Iteration) + public void reset() { + _samples = 0L; + _inputCentroidCount = 0.0; + _resultingCentroidCount = 0.0; + _maxResultingCentroidCount = 0.0; + _p75ErrorPpb = 0.0; + _maxP75ErrorPpb = 0.0; + _p95ErrorPpb = 0.0; + _maxP95ErrorPpb = 0.0; + _p99ErrorPpb = 0.0; + _maxP99ErrorPpb = 0.0; + } + + public double inputCentroidCount() { + return perResult(_inputCentroidCount); + } + + public double resultingCentroidCount() { + return perResult(_resultingCentroidCount); + } + + public double maxResultingCentroidCount() { + return perResult(_maxResultingCentroidCount); + } + + public double p75ErrorPpb() { + return perResult(_p75ErrorPpb); + } + + public double maxP75ErrorPpb() { + return perResult(_maxP75ErrorPpb); + } + + public double p95ErrorPpb() { + return perResult(_p95ErrorPpb); + } + + public double maxP95ErrorPpb() { + return perResult(_maxP95ErrorPpb); + } + + public double p99ErrorPpb() { + return perResult(_p99ErrorPpb); + } + + public double maxP99ErrorPpb() { + return perResult(_maxP99ErrorPpb); + } + + private void record(QualityStats stats) { + _samples++; + _inputCentroidCount += stats._inputCentroidCount; + _resultingCentroidCount += stats._resultingCentroidCount; + _maxResultingCentroidCount += stats._maxResultingCentroidCount; + _p75ErrorPpb += stats._p75ErrorPpb; + _maxP75ErrorPpb += stats._maxP75ErrorPpb; + _p95ErrorPpb += stats._p95ErrorPpb; + _maxP95ErrorPpb += stats._maxP95ErrorPpb; + _p99ErrorPpb += stats._p99ErrorPpb; + _maxP99ErrorPpb += stats._maxP99ErrorPpb; + } + + private double perResult(double sum) { + return _samples == 0L ? 0.0 : sum / _samples / _aggregationDivisor; + } + } + + /// Immutable verbose TDigest input used to keep the exact source centroid layout identical across dependency + /// versions. TDigest 3.3 force-compresses a [MergingDigest] when its centroids are read, so a library digest cannot + /// represent this control input without changing the state under test. + private static final class FixedTDigest extends TDigest { + private final byte[] _bytes; + private final double _compression; + private final double _min; + private final double _max; + private final long _size; + private final List _centroids; + + private FixedTDigest(double[] sortedValues, double compression) { + this(createFixedVerboseBytes(sortedValues, compression)); + } + + private FixedTDigest(byte[] bytes) { + _bytes = bytes; + ByteBuffer input = ByteBuffer.wrap(bytes); + if (input.getInt() != 1) { + throw new IllegalArgumentException("Expected verbose TDigest encoding"); + } + _min = input.getDouble(); + _max = input.getDouble(); + _compression = input.getDouble(); + int numCentroids = input.getInt(); + List centroids = new ArrayList<>(numCentroids); + long size = 0L; + for (int i = 0; i < numCentroids; i++) { + double weight = input.getDouble(); + double mean = input.getDouble(); + int count = Math.toIntExact((long) weight); + if (count != weight) { + throw new IllegalArgumentException("Non-integral fixed TDigest weight: " + weight); + } + centroids.add(new Centroid(mean, count)); + size += count; + } + _centroids = List.copyOf(centroids); + _size = size; + } + + private byte[] bytes() { + return _bytes.clone(); + } + + @Override + public void add(double value) { + throw new UnsupportedOperationException("Fixed TDigest is immutable"); + } + + @Override + public void add(double value, int weight) { + throw new UnsupportedOperationException("Fixed TDigest is immutable"); + } + + @Override + public void add(List others) { + throw new UnsupportedOperationException("Fixed TDigest is immutable"); + } + + @Override + public void add(TDigest other) { + throw new UnsupportedOperationException("Fixed TDigest is immutable"); + } + + @Override + public void compress() { + } + + @Override + public long size() { + return _size; + } + + @Override + public double cdf(double value) { + return toTDigest().cdf(value); + } + + @Override + public double quantile(double quantile) { + return toTDigest().quantile(quantile); + } + + @Override + public int centroidCount() { + return _centroids.size(); + } + + @Override + public List centroids() { + return _centroids; + } + + @Override + public double compression() { + return _compression; + } + + @Override + public int byteSize() { + return _bytes.length; + } + + @Override + public int smallByteSize() { + return toTDigest().smallByteSize(); + } + + @Override + public void asBytes(ByteBuffer buffer) { + buffer.put(_bytes); + } + + @Override + public void asSmallBytes(ByteBuffer buffer) { + toTDigest().asSmallBytes(buffer); + } + + @Override + public TDigest recordAllData() { + return toTDigest().recordAllData(); + } + + @Override + public boolean isRecording() { + return false; + } + + @Override + public double getMin() { + return _min; + } + + @Override + public double getMax() { + return _max; + } + + private TDigest toTDigest() { + return MergingDigest.fromBytes(ByteBuffer.wrap(_bytes)); + } + } + + private static final class ImmutableMergingDigest extends MergingDigest { + private boolean _frozen; + + private ImmutableMergingDigest(double compression) { + super(compression); + } + + private void freeze() { + super.compress(); + _frozen = true; + } + + @Override + public void compress() { + if (!_frozen) { + super.compress(); + } + } + } + + private static final class ReducerAggregationFunction extends PercentileTDigestAggregationFunction { + private final Implementation _implementation; + private final Map> _pending = new IdentityHashMap<>(); + + private ReducerAggregationFunction(ExpressionContext expression, int compression, + Implementation implementation) { + super(expression, 75.0, compression, false); + _implementation = implementation; + } + + @Override + public TDigest merge(TDigest intermediateResult1, TDigest intermediateResult2) { + if (!_implementation.isBatching()) { + if (_implementation == Implementation.PAIRWISE) { + if (intermediateResult1.size() == 0L) { + return intermediateResult2; + } + if (intermediateResult2.size() == 0L) { + return intermediateResult1; + } + intermediateResult1.add(intermediateResult2); + return intermediateResult1; + } + if (_implementation == Implementation.SINGLETON_LIST && intermediateResult1.size() != 0L + && intermediateResult2.size() != 0L) { + intermediateResult1.add(List.of(intermediateResult2)); + return intermediateResult1; + } + return super.merge(intermediateResult1, intermediateResult2); + } + if (intermediateResult1.size() == 0L) { + return intermediateResult2; + } + if (intermediateResult2.size() == 0L) { + return intermediateResult1; + } + List pending = _pending.computeIfAbsent(intermediateResult1, + ignored -> new ArrayList<>(_implementation._batchSize)); + pending.add(intermediateResult2); + if (pending.size() == _implementation._batchSize) { + flush(intermediateResult1, pending); + } + return intermediateResult1; + } + + private void flushAll() { + if (_pending.isEmpty()) { + return; + } + for (Map.Entry> entry : _pending.entrySet()) { + flush(entry.getKey(), entry.getValue()); + } + _pending.clear(); + } + + private void resetPending() { + if (!_pending.isEmpty()) { + throw new IllegalStateException("Pending TDigest batches were not flushed"); + } + } + + private void flush(TDigest target, List pending) { + if (pending.isEmpty()) { + return; + } + TDigest batch = TDigest.createMergingDigest(_compressionFactor); + batch.add(pending); + pending.clear(); + target.add(batch); + } + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestStarTreeAggregation.java b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestStarTreeAggregation.java new file mode 100644 index 000000000000..ff97ca0999b3 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestStarTreeAggregation.java @@ -0,0 +1,195 @@ +/** + * 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.perf.aggregation; + +import com.tdunning.math.stats.TDigest; +import java.util.Map; +import java.util.SplittableRandom; +import java.util.concurrent.TimeUnit; +import javax.annotation.Nullable; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.ObjectSerDeUtils; +import org.apache.pinot.core.common.SyntheticBlockValSets; +import org.apache.pinot.core.plan.DocIdSetPlanNode; +import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.function.PercentileTDigestAggregationFunction; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Level; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; +import org.roaringbitmap.RoaringBitmap; + +/// Measures percentile TDigest aggregation over serialized star-tree entries representing 100 million source rows. +/// +/// The benchmark creates 10,000 deterministic serialized digests, each summarizing 10,000 source values, and feeds one +/// Pinot-sized block through the production `BYTES` aggregation path. The group-by workload distributes those entries +/// round-robin across 100, 1,000, or 10,000 result groups. It measures query-time deserialization, merging, and result +/// extraction, but deliberately excludes segment generation and forward-index I/O. +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Warmup(iterations = 2, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork(value = 2, jvmArgsAppend = {"-Xms4g", "-Xmx8g"}) +@Threads(1) +public class BenchmarkPercentileTDigestStarTreeAggregation { + private static final int NUM_SOURCE_ROWS = 100_000_000; + private static final int NUM_STAR_TREE_ROWS = DocIdSetPlanNode.MAX_DOC_PER_CALL; + private static final int SOURCE_ROWS_PER_DIGEST = NUM_SOURCE_ROWS / NUM_STAR_TREE_ROWS; + private static final double EXPECTED_PERCENTILE_75 = 0.7499315463846341; + private static final double MAX_ABSOLUTE_ERROR = 0.001; + private static final ExpressionContext EXPRESSION = ExpressionContext.forIdentifier("col"); + + /// State for the serialized star-tree aggregation benchmark. + @State(Scope.Benchmark) + public static class AggregationState { + private PercentileTDigestAggregationFunction _function; + private Map _blockValSetMap; + + @Setup(Level.Trial) + public void setUp() { + if (NUM_SOURCE_ROWS % NUM_STAR_TREE_ROWS != 0) { + throw new IllegalStateException("NUM_SOURCE_ROWS must be divisible by NUM_STAR_TREE_ROWS"); + } + _blockValSetMap = Map.of(EXPRESSION, new BytesBlockValSet(createSerializedDigests())); + _function = new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + } + } + + /// State for the serialized star-tree group-by aggregation benchmark. + @State(Scope.Benchmark) + public static class GroupByAggregationState { + @Param({"100", "1000", "10000"}) + private int _numGroups; + + private PercentileTDigestAggregationFunction _function; + private Map _blockValSetMap; + private int[] _groupKeys; + + @Setup(Level.Trial) + public void setUp() { + if (NUM_STAR_TREE_ROWS % _numGroups != 0) { + throw new IllegalStateException("NUM_STAR_TREE_ROWS must be divisible by the number of groups"); + } + _blockValSetMap = Map.of(EXPRESSION, new BytesBlockValSet(createSerializedDigests())); + _function = new PercentileTDigestAggregationFunction(EXPRESSION, 75.0, false); + _groupKeys = new int[NUM_STAR_TREE_ROWS]; + for (int rowId = 0; rowId < NUM_STAR_TREE_ROWS; rowId++) { + _groupKeys[rowId] = rowId % _numGroups; + } + } + } + + @Benchmark + public double aggregatePercentileTDigest75(AggregationState state) { + AggregationResultHolder resultHolder = state._function.createAggregationResultHolder(); + state._function.aggregate(NUM_STAR_TREE_ROWS, resultHolder, state._blockValSetMap); + TDigest result = state._function.extractAggregationResult(resultHolder); + if (result.size() != NUM_SOURCE_ROWS) { + throw new IllegalStateException("Unexpected TDigest size: " + result.size() + ", expected: " + NUM_SOURCE_ROWS); + } + double percentile = state._function.extractFinalResult(result); + if (!Double.isFinite(percentile) || Math.abs(percentile - EXPECTED_PERCENTILE_75) > MAX_ABSOLUTE_ERROR) { + throw new IllegalStateException( + "Unexpected percentile: " + percentile + ", expected: " + EXPECTED_PERCENTILE_75); + } + return percentile; + } + + @Benchmark + public double aggregateGroupByPercentileTDigest75(GroupByAggregationState state) { + GroupByResultHolder resultHolder = state._function.createGroupByResultHolder(state._numGroups, state._numGroups); + state._function.aggregateGroupBySV(NUM_STAR_TREE_ROWS, state._groupKeys, resultHolder, state._blockValSetMap); + + long totalSize = 0L; + double percentileChecksum = 0.0; + long expectedGroupSize = NUM_SOURCE_ROWS / state._numGroups; + for (int groupKey = 0; groupKey < state._numGroups; groupKey++) { + TDigest result = state._function.extractGroupByResult(resultHolder, groupKey); + if (result.size() != expectedGroupSize) { + throw new IllegalStateException( + "Unexpected TDigest size: " + result.size() + ", expected: " + expectedGroupSize); + } + double percentile = state._function.extractFinalResult(result); + if (!Double.isFinite(percentile) || Math.abs(percentile - 0.75) > 0.05) { + throw new IllegalStateException("Unexpected percentile for group " + groupKey + ": " + percentile); + } + totalSize += result.size(); + percentileChecksum += percentile; + } + if (totalSize != NUM_SOURCE_ROWS) { + throw new IllegalStateException("Unexpected total TDigest size: " + totalSize + ", expected: " + NUM_SOURCE_ROWS); + } + return percentileChecksum; + } + + private static byte[][] createSerializedDigests() { + SplittableRandom random = new SplittableRandom(42); + byte[][] serializedDigests = new byte[NUM_STAR_TREE_ROWS][]; + for (int rowId = 0; rowId < NUM_STAR_TREE_ROWS; rowId++) { + TDigest digest = + TDigest.createMergingDigest(PercentileTDigestAggregationFunction.DEFAULT_TDIGEST_COMPRESSION); + for (int i = 0; i < SOURCE_ROWS_PER_DIGEST; i++) { + digest.add(random.nextDouble()); + } + serializedDigests[rowId] = ObjectSerDeUtils.TDIGEST_SER_DE.serialize(digest); + } + return serializedDigests; + } + + private static final class BytesBlockValSet extends SyntheticBlockValSets.Base { + private final byte[][] _values; + + private BytesBlockValSet(byte[][] values) { + _values = values; + } + + @Override + public DataType getValueType() { + return DataType.BYTES; + } + + @Override + public boolean isSingleValue() { + return true; + } + + @Override + public byte[][] getBytesValuesSV() { + return _values; + } + + @Nullable + @Override + public RoaringBitmap getNullBitmap() { + return null; + } + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestValueAggregator.java b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestValueAggregator.java new file mode 100644 index 000000000000..da92f80e4e0c --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/aggregation/BenchmarkPercentileTDigestValueAggregator.java @@ -0,0 +1,104 @@ +/** + * 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.perf.aggregation; + +import com.tdunning.math.stats.TDigest; +import java.util.List; +import java.util.SplittableRandom; +import java.util.concurrent.TimeUnit; +import org.apache.pinot.segment.local.aggregator.PercentileTDigestValueAggregator; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Threads; +import org.openjdk.jmh.annotations.Warmup; + +/// Measures the raw-value TDigest aggregation kernel used while constructing a star-tree index. +/// +/// The default models one million source rows split into groups of 1,000 rows with identical dimension keys. Each +/// group is serialized because every aggregate eventually becomes a star-tree forward-index entry. Use +/// `-p _numRows=100000000` for a 100-million-row run and vary `_rowsPerGroup` to model the table's dimension +/// cardinality. This deliberately excludes dimension sorting and forward-index I/O so it isolates +/// [PercentileTDigestValueAggregator]. +@State(Scope.Benchmark) +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Warmup(iterations = 2, time = 1) +@Measurement(iterations = 5, time = 1) +@Fork(value = 2, jvmArgsAppend = {"-Xms4g", "-Xmx8g"}) +@Threads(1) +public class BenchmarkPercentileTDigestValueAggregator { + private static final int VALUE_BLOCK_SIZE = 1 << 16; + + @Param({"1000000"}) + int _numRows; + + @Param({"1000"}) + int _rowsPerGroup; + + private double[] _values; + + @Setup + public void setUp() { + if (_numRows <= 0) { + throw new IllegalArgumentException("numRows must be positive"); + } + if (_rowsPerGroup <= 0) { + throw new IllegalArgumentException("rowsPerGroup must be positive"); + } + _values = new double[VALUE_BLOCK_SIZE]; + SplittableRandom random = new SplittableRandom(42); + for (int i = 0; i < VALUE_BLOCK_SIZE; i++) { + _values[i] = random.nextDouble(); + } + } + + @Benchmark + public long aggregateAndSerializeRawValues() { + PercentileTDigestValueAggregator aggregator = new PercentileTDigestValueAggregator(List.of()); + long totalAggregatedRows = 0; + long totalSerializedBytes = 0; + for (int from = 0; from < _numRows; from += _rowsPerGroup) { + int to = Math.min(_numRows, from + _rowsPerGroup); + TDigest digest = aggregator.getInitialAggregatedValue(_values[from & (VALUE_BLOCK_SIZE - 1)]); + for (int i = from + 1; i < to; i++) { + aggregator.applyRawValue(digest, _values[i & (VALUE_BLOCK_SIZE - 1)]); + } + int maxByteSize = aggregator.getMaxAggregatedValueByteSize(); + byte[] serialized = aggregator.serializeAggregatedValue(digest); + if (serialized.length > maxByteSize) { + throw new IllegalStateException( + "Serialized TDigest exceeds registered maximum: " + serialized.length + " > " + maxByteSize); + } + totalAggregatedRows += digest.size(); + totalSerializedBytes += serialized.length; + } + if (totalAggregatedRows != _numRows) { + throw new IllegalStateException("Unexpected aggregated row count: " + totalAggregatedRows); + } + return totalSerializedBytes; + } +} diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregator.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregator.java index 772b6a528916..0f0d5b1818b8 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregator.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregator.java @@ -19,6 +19,7 @@ package org.apache.pinot.segment.local.aggregator; import com.tdunning.math.stats.TDigest; +import java.nio.ByteBuffer; import java.util.List; import org.apache.pinot.common.request.context.ExpressionContext; import org.apache.pinot.segment.local.utils.CustomSerDeUtils; @@ -31,6 +32,9 @@ public class PercentileTDigestValueAggregator implements ValueAggregator getDefaultCentroidCapacity(value.compression())) { + // Verbose decoding recreates the default centroid capacity, while compact decoding preserves the stored one. + ByteBuffer buffer = ByteBuffer.allocate(value.smallByteSize()); + value.asSmallBytes(buffer); + bytes = buffer.array(); + } else { + bytes = CustomSerDeUtils.TDIGEST_SER_DE.serialize(value); + } + _maxByteSize = Math.max(_maxByteSize, bytes.length); + return bytes; } @Override public TDigest deserializeAggregatedValue(byte[] bytes) { return CustomSerDeUtils.TDIGEST_SER_DE.deserialize(bytes); } + + private void updateMaxByteSize(TDigest value) { + long defaultCapacity = getDefaultCentroidCapacity(value.compression()); + long maxCentroids = Math.max(value.centroidCount(), Math.min(value.size(), defaultCapacity)); + _maxByteSize = Math.max(_maxByteSize, getMaxVerboseByteSize(maxCentroids)); + } + + private static int getMaxVerboseByteSize(long centroidCount) { + long maxCentroids = Math.max(centroidCount, 0L); + return Math.toIntExact(Math.addExact(VERBOSE_HEADER_SIZE, + Math.multiplyExact(VERBOSE_CENTROID_SIZE, maxCentroids))); + } + + private static long getDefaultCentroidCapacity(double compression) { + return Math.addExact(Math.multiplyExact(2L, (long) Math.ceil(compression)), CENTROID_CAPACITY_PADDING); + } } diff --git a/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregatorTest.java b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregatorTest.java new file mode 100644 index 000000000000..709311872fae --- /dev/null +++ b/pinot-segment-local/src/test/java/org/apache/pinot/segment/local/aggregator/PercentileTDigestValueAggregatorTest.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.segment.local.aggregator; + +import com.tdunning.math.stats.TDigest; +import java.nio.ByteBuffer; +import java.util.List; +import org.apache.pinot.common.request.Literal; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.segment.local.utils.CustomSerDeUtils; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +/// Tests the serialization-size contract and deferred compression used by [PercentileTDigestValueAggregator]. +public class PercentileTDigestValueAggregatorTest { + + @Test(dataProvider = "compressionBounds") + public void testRawValuesTrackMaxVerboseSizeWithoutCompression(int compression, int numValues, + int expectedMaxByteSize) { + PercentileTDigestValueAggregator aggregator = newAggregator(compression); + assertEquals(aggregator.getMaxAggregatedValueByteSize(), 0); + + TDigest digest = aggregator.getInitialAggregatedValue(0.0); + for (int i = 1; i < numValues; i++) { + aggregator.applyRawValue(digest, i); + } + assertEquals(digest.centroidCount(), 0); + assertEquals(aggregator.getMaxAggregatedValueByteSize(), expectedMaxByteSize); + byte[] serialized = aggregator.serializeAggregatedValue(digest); + assertTrue(serialized.length <= expectedMaxByteSize); + } + + @DataProvider + public static Object[][] compressionBounds() { + return new Object[][]{ + {10, 30, 512}, + {100, 210, 3_392}, + {1_000, 2_010, 32_192} + }; + } + + @Test + public void testRawUpdatesDoNotForceCompression() { + PercentileTDigestValueAggregator aggregator = newAggregator(100); + TDigest digest = aggregator.getInitialAggregatedValue(0.0); + for (int i = 1; i < 256; i++) { + aggregator.applyRawValue(digest, i); + } + + assertEquals(digest.size(), 256L); + assertEquals(digest.centroidCount(), 0); + int maxByteSize = aggregator.getMaxAggregatedValueByteSize(); + byte[] serialized = aggregator.serializeAggregatedValue(digest); + assertTrue(digest.centroidCount() > 0); + assertTrue(serialized.length <= maxByteSize); + } + + @Test + public void testPreAggregatedCompressionExpandsRegisteredBound() { + TDigest input = TDigest.createMergingDigest(200); + input.add(42.0); + byte[] inputBytes = CustomSerDeUtils.TDIGEST_SER_DE.serialize(input); + + PercentileTDigestValueAggregator aggregator = newAggregator(10); + TDigest result = aggregator.getInitialAggregatedValue(inputBytes); + for (int i = 0; i < 409; i++) { + aggregator.applyRawValue(result, i); + } + + assertEquals(result.compression(), 200.0); + assertEquals(aggregator.getMaxAggregatedValueByteSize(), 6_592); + int maxByteSize = aggregator.getMaxAggregatedValueByteSize(); + assertTrue(aggregator.serializeAggregatedValue(result).length <= maxByteSize); + } + + @Test + public void testOversizedSmallEncodedDigestRemainsReadable() { + int centroidCount = 300; + byte[] smallEncoding = createSmallEncoding(10, 400, 500, centroidCount); + + PercentileTDigestValueAggregator aggregator = newAggregator(10); + TDigest result = aggregator.getInitialAggregatedValue(smallEncoding); + int maxByteSize = aggregator.getMaxAggregatedValueByteSize(); + byte[] serialized = aggregator.serializeAggregatedValue(result); + TDigest clone = aggregator.cloneAggregatedValue(result); + + assertEquals(result.centroidCount(), centroidCount); + assertEquals(maxByteSize, 4_832); + assertEquals(ByteBuffer.wrap(serialized).getInt(), 2); + assertEquals(serialized.length, smallEncoding.length); + assertTrue(serialized.length <= maxByteSize); + assertEquals(clone.size(), result.size()); + assertEquals(clone.centroidCount(), result.centroidCount()); + assertEquals(clone.quantile(0.75), result.quantile(0.75)); + } + + @Test + public void testAggregatedUpdateDoesNotForceDestinationCompression() { + PercentileTDigestValueAggregator aggregator = newAggregator(100); + TDigest destination = aggregator.getInitialAggregatedValue(1.0); + TDigest source = TDigest.createMergingDigest(100); + source.add(2.0); + source.add(3.0); + + aggregator.applyAggregatedValue(destination, source); + + assertEquals(destination.size(), 3L); + assertEquals(destination.centroidCount(), 0); + int maxByteSize = aggregator.getMaxAggregatedValueByteSize(); + byte[] serialized = aggregator.serializeAggregatedValue(destination); + assertTrue(serialized.length <= maxByteSize); + } + + private static PercentileTDigestValueAggregator newAggregator(int compression) { + return new PercentileTDigestValueAggregator( + List.of(ExpressionContext.forLiteral(Literal.intValue(compression)))); + } + + private static byte[] createSmallEncoding(int compression, int centroidCapacity, int bufferSize, + int centroidCount) { + ByteBuffer buffer = ByteBuffer.allocate(30 + centroidCount * 2 * Float.BYTES); + buffer.putInt(2); + buffer.putDouble(0.0); + buffer.putDouble(centroidCount - 1.0); + buffer.putFloat(compression); + buffer.putShort((short) centroidCapacity); + buffer.putShort((short) bufferSize); + buffer.putShort((short) centroidCount); + for (int i = 0; i < centroidCount; i++) { + buffer.putFloat(1.0F); + buffer.putFloat(i); + } + return buffer.array(); + } +}