xiangfu0 commented on code in PR #18996: URL: https://github.com/apache/pinot/pull/18996#discussion_r3591611690
########## pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/PercentileTDigestAccumulator.java: ########## @@ -0,0 +1,448 @@ +/** + * 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.MergingDigest; +import com.tdunning.math.stats.TDigest; +import java.nio.BufferUnderflowException; +import java.nio.ByteBuffer; +import java.util.Arrays; + +/// 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 intermediate-result +/// serialization and distributed merging remain unchanged. +/// +/// 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 thread-confined to one aggregation result or group-by result. +final class PercentileTDigestAccumulator { + 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_CENTROID_SIZE = 8; + + 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 boolean _hasSerializedInput; + + private int _numRawValues; + private int _numCentroids; + 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 = Math.toIntExact( + Math.addExact(Math.multiplyExact(2L, roundedCompression), CENTROID_CAPACITY_PADDING)); + 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 forSerializedTDigestWithMergeBuffers(byte[] bytes) { + return new PercentileTDigestAccumulator(readCompression(bytes), false, true); + } + + 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(); + } + } + } + + void addSerializedTDigest(byte[] bytes) { + flush(); + if (!_hasSerializedInput && _numCentroids == 0 && _totalWeight == 0.0) { + validateSerializedTDigest(bytes); + _pendingSerializedTDigest = bytes; + _hasSerializedInput = true; + return; + } + materializePendingSerializedTDigest(); + mergeSerializedTDigest(bytes); + _hasSerializedInput = true; + } + + private void mergeSerializedTDigest(byte[] bytes) { + + ByteBuffer input = ByteBuffer.wrap(bytes); + int encoding = input.getInt(); + double min = input.getDouble(); + double max = input.getDouble(); + int numCentroids; + boolean initialize = _numCentroids == 0 && _totalWeight == 0.0; + double[] incomingMeans; + double[] incomingWeights; + if (encoding == VERBOSE_ENCODING) { + input.getDouble(); + numCentroids = input.getInt(); + checkCentroidCount(numCentroids); + checkCentroidsAvailable(input, numCentroids, VERBOSE_CENTROID_SIZE); + 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 if (encoding == SMALL_ENCODING) { + input.getFloat(); + input.getShort(); + input.getShort(); + numCentroids = input.getShort(); + checkCentroidCount(numCentroids); + checkCentroidsAvailable(input, numCentroids, SMALL_CENTROID_SIZE); + 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(); + } + } else { + throw new IllegalStateException("Invalid format for serialized histogram"); + } + + 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); + } + } + + TDigest toTDigest() { + if (_pendingSerializedTDigest != null) { + return MergingDigest.fromBytes(ByteBuffer.wrap(_pendingSerializedTDigest)); + } + flush(); + if (_numCentroids == 0) { + return TDigest.createMergingDigest(_compression); + } + + ByteBuffer buffer = ByteBuffer.allocate(VERBOSE_HEADER_SIZE + VERBOSE_CENTROID_SIZE * _numCentroids); Review Comment: This materialized path always re-emits verbose TDigest bytes, but t-digest 3.2 verbose decoding constructs a digest with only the default centroid capacity. If the accumulator first materializes a valid small-encoded digest whose centroid count exceeds that default capacity, `MergingDigest.fromBytes(buffer)` can throw during query result extraction. Preserve the small-encoding capacity or build the result digest with enough centroid capacity before returning it. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
