xiangfu0 commented on code in PR #19795: URL: https://github.com/apache/pinot/pull/19795#discussion_r4239363031
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/NonFiniteAwareTDigest.java: ########## @@ -0,0 +1,500 @@ +/** + * 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.customobject.tdigest; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.segment.local.customobject.tdigest.TDigestCodec.SerializedTDigestMetadata; + + +/// Stores non-finite values as exact tail masses while delegating finite values to an owned native digest. +/// Emits legacy t-digest bytes and preserves historical fractional endpoint provenance across merges and copies. +/// Instances are mutable and require external synchronization. +public final class NonFiniteAwareTDigest extends TDigest { + private final PercentileTDigestAccumulator _finiteDigest; + private double _negativeInfinityWeight; + private double _positiveInfinityWeight; + private List<Centroid> _finiteCentroids; + private byte[] _originalFractionalBytes; + + public static NonFiniteAwareTDigest forLegacyAggregation(double compression) { + return new NonFiniteAwareTDigest(PercentileTDigestAccumulator.forLegacyAggregation(compression), 0.0, 0.0); + } + + private NonFiniteAwareTDigest(PercentileTDigestAccumulator finiteDigest, double negativeInfinityWeight, + double positiveInfinityWeight) { + _finiteDigest = finiteDigest; + _negativeInfinityWeight = negativeInfinityWeight; + _positiveInfinityWeight = positiveInfinityWeight; + } + + public NonFiniteAwareTDigest copy() { + if (_originalFractionalBytes != null) { + return fromBytes(_originalFractionalBytes); + } + PercentileTDigestAccumulator source = _finiteDigest; + PercentileTDigestAccumulator finiteDigest; + if (!source.hasValidStatistics() || source.hasOriginalFractionalPayload()) { + finiteDigest = PercentileTDigestAccumulator.fromBytes(source.serialize()); + } else { + finiteDigest = PercentileTDigestAccumulator.forLegacyAggregation(compression()); + if (!source.isEmpty()) { + finiteDigest.add(source); Review Comment: Copy now takes a native buffered-state snapshot without flushing or compressing the source. Mutable buffers are independent, and compression mode, merge direction, retained bytes and fractional provenance are preserved. The seeded regression verifies unchanged source bytes and quantiles after further aggregation, including independent clone mutation and infinity tails. Addressed in [277a170ffe](https://github.com/apache/pinot/commit/277a170ffe224d680ee00c6f56efd422bb673990). All 526 affected tests, actual 3.2/3.3 reader checks and the unchanged 270-case accuracy oracle pass. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/PercentileTDigestAccumulator.java: ########## @@ -0,0 +1,1498 @@ +/** + * 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.customobject.tdigest; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.List; +import org.apache.pinot.segment.local.customobject.tdigest.TDigestCodec.SerializedTDigestMetadata; + + + +/// Accumulates raw values and serialized TDigests into primitive centroids using Pinot's K1 weight-limit rule. +/// +/// Raw values are sorted in small batches; serialized centroids retain their existing sorted order. Both are +/// linearly merged with accumulated centroids. K1 preserves Pinot's existing middle-quantile behavior, and the +/// legacy verbose and small encodings preserve stored data and mixed-version wire compatibility. +/// +/// Serialized state keeps its first digest pending and allocates centroid buffers only when another input must be +/// merged or the digest is queried. Raw-value buffers remain unallocated until a raw value is added. +/// +/// Historical corrupted payloads remain opaque and byte-exact, with NaN statistics. Their reported historical mass +/// may be invalid; they cannot be mutated or mixed with another distribution. Such an operation fails before +/// changing aggregate mass, because neither invented statistics nor synthetic NaN bytes are safe for old readers. +/// Unchanged historical fractional boundaries retain their original bytes, including when configured compression +/// differs. After a merge, inherited fractional endpoints retain precise verbose weights and finite statistics when +/// unit endpoints cannot preserve their mass. These keep the historical t-digest 3.3 assertion-enabled reader and +/// recompression limitation. Newly introduced unsupported fractional endpoints are rejected. +/// +/// Instances are externally serialized per aggregation or group key and are not safe for concurrent mutation. +public 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 MIN_COMPRESSION = 10; + private static final int TWO_LEVEL_COMPRESSION_MULTIPLIER = 2; + private static final int MIN_INCREMENTAL_COMPRESSION = 50; + + private final double _compression; + private final boolean _useTwoLevelRawCompression; + private final double _workingCompression; + private final double _publicMaxWeightScale; + private final double _workingMaxWeightScale; + private final int _centroidCapacity; + private final int _pendingCentroidCapacity; + private double[] _rawValues; + private double[] _centroidMeans; + private double[] _centroidWeights; + private double[] _outputMeans; + private double[] _outputWeights; + private double[] _incomingMeans; + private double[] _incomingWeights; + private double[] _sortedIncomingMeans; + private double[] _sortedIncomingWeights; + private SerializedTDigestInput _serializedTDigestInput; + private byte[] _pendingSerializedTDigest; + private byte[] _originalFractionalBytes; + private double _historicalFractionalMinMean = Double.NaN; + private double _historicalFractionalMaxMean = Double.NaN; + private SerializedTDigestMetadata _pendingSerializedMetadata; + private boolean _legacyDegraded; + private boolean _hasFractionalWeights; + + private int _numRawValues; + private int _numCentroids; + private int _numIncomingCentroids; + private int _mergeCount; + private boolean _publiclyCompressed; + private boolean _incomingCentroidsSorted = true; + private double _totalWeight; + private double _incomingWeight; + private double _min = Double.POSITIVE_INFINITY; + private double _max = Double.NEGATIVE_INFINITY; + + public PercentileTDigestAccumulator(int compression) { + this(positiveCompression(compression), true, true); + } + + public PercentileTDigestAccumulator(double compression) { + this(compression, true, true); + } + + private PercentileTDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers) { + this(compression, allocateRawBuffer, allocateMergeBuffers, false); + } + + private PercentileTDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers, + boolean useTwoLevelRawCompression) { + TDigestCodec.validateCompression(compression); + _useTwoLevelRawCompression = useTwoLevelRawCompression; + _compression = Math.max(MIN_COMPRESSION, compression); + _workingCompression = TWO_LEVEL_COMPRESSION_MULTIPLIER * _compression; + _publicMaxWeightScale = calculateK1MaxWeightScale(_compression); + _workingMaxWeightScale = calculateK1MaxWeightScale(_workingCompression); + long roundedCompression = (long) Math.ceil(_compression); + _centroidCapacity = TDigestCodec.getDefaultCentroidCapacity(_compression); + _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity, useTwoLevelRawCompression); + if (allocateRawBuffer) { + _rawValues = new double[getInitialRawBufferSize(roundedCompression)]; + } + if (allocateMergeBuffers) { + int initialCapacity = Math.min(_centroidCapacity, MIN_RAW_BUFFER_SIZE); + _centroidMeans = new double[initialCapacity]; + _centroidWeights = new double[initialCapacity]; + _outputMeans = new double[initialCapacity]; + _outputWeights = new double[initialCapacity]; + } + } + + /// Creates a lazy digest with the legacy two-level compression used by Star-tree and general digest aggregation. + public static PercentileTDigestAccumulator forLegacyAggregation(double compression) { + return new PercentileTDigestAccumulator(compression, false, false, true); + } + + /// Reads and retains one legacy digest, preserving the general aggregation compression policy. + public static PercentileTDigestAccumulator fromBytes(byte[] bytes) { + return fromBytes(ByteBuffer.wrap(bytes)); + } + + /// Reads one payload and advances the buffer past it. + public static PercentileTDigestAccumulator fromBytes(ByteBuffer bytes) { + SerializedTDigestInput input = new SerializedTDigestInput(); + input.reset(bytes); + return fromBytes(input); + } + + /// Retains an already validated input without rescanning its centroid bytes. + public static PercentileTDigestAccumulator fromBytes(SerializedTDigestInput input) { + PercentileTDigestAccumulator digest = forLegacyAggregation(input.getCompression()); + digest.addSerializedTDigest(input); + return digest; + } + + public static PercentileTDigestAccumulator forReduction(double compression) { + return new PercentileTDigestAccumulator(positiveCompression(compression), false, false); + } + + /// Copies primitive centroids without encoding and decoding them. Zero weights are ignored, and unordered + /// means are sorted when initializing a stored, already compressed distribution; caller arrays are not mutated. + /// Subsequent inputs use the normal buffered merge. Positive weights and consistent extrema are required. + public void addCentroids(double[] means, double[] weights, int count, double min, double max, + boolean alreadyCompressed) { + requireMutable(); + if (count < 0 || count > means.length || count > weights.length) { + throw new IllegalArgumentException("Invalid TDigest centroid count: " + count); + } + double totalWeight = 0.0; + int positiveCount = 0; + boolean sorted = true; + boolean fractionalWeights = false; + double previousMean = Double.NEGATIVE_INFINITY; + for (int i = 0; i < count; i++) { + double weight = weights[i]; + if (weight == 0.0) { + continue; + } + double mean = means[i]; + if (!(weight > 0.0) || !Double.isFinite(weight) || !(mean >= min && mean <= max)) { + throw new IllegalArgumentException("Invalid TDigest centroid"); + } + totalWeight += weight; + sorted &= mean >= previousMean; + previousMean = mean; + fractionalWeights |= weight != Math.rint(weight); + positiveCount++; + } + checkTotalWeight(getTotalWeight() + totalWeight); + if (positiveCount == 0) { + return; + } + materializePendingSerializedTDigest(); + _originalFractionalBytes = null; + _hasFractionalWeights |= fractionalWeights; + if (hasNoInputs() && alreadyCompressed) { + ensureCentroidCapacity(Math.addExact(positiveCount, 2)); + for (int i = 0; i < count; i++) { + if (weights[i] > 0.0) { + _centroidMeans[_numCentroids] = means[i]; + _centroidWeights[_numCentroids++] = weights[i]; + } + } + if (!sorted) { + it.unimi.dsi.fastutil.Arrays.mergeSort(0, _numCentroids, + (first, second) -> Double.compare(_centroidMeans[first], _centroidMeans[second]), + (first, second) -> { + double mean = _centroidMeans[first]; + _centroidMeans[first] = _centroidMeans[second]; + _centroidMeans[second] = mean; + double weight = _centroidWeights[first]; + _centroidWeights[first] = _centroidWeights[second]; + _centroidWeights[second] = weight; + }); + } + _totalWeight = totalWeight; + _min = min; + _max = max; + normalizeBoundaryCentroids(); + _publiclyCompressed = true; + } else { + if (sorted && positiveCount == count) { + bufferIncomingCentroids(means, weights, count, totalWeight, min, max); + } else { + // Caller-owned arrays are copied only when filtering or sorting is necessary. + double[] sortedMeans = new double[positiveCount]; + double[] sortedWeights = new double[positiveCount]; + int index = 0; + for (int i = 0; i < count; i++) { + if (weights[i] > 0.0) { + sortedMeans[index] = means[i]; + sortedWeights[index++] = weights[i]; + } + } + if (!sorted) { + sortCentroids(sortedMeans, sortedWeights, positiveCount); + } + bufferIncomingCentroids(sortedMeans, sortedWeights, positiveCount, totalWeight, min, max); + } + _min = Math.min(_min, min); + _max = Math.max(_max, max); + } + } + + public void add(double[] values, int from, int toExclusive) { + if (from >= toExclusive) { + return; + } + requireMutable(); + materializePendingSerializedTDigest(); + while (from < toExclusive) { + ensureRawBuffer(); + if (_numRawValues == _rawValues.length + || _numRawValues + _numIncomingCentroids == getPendingInputLimit()) { + flush(); + } + int numValues = Math.min(toExclusive - from, Math.min(_rawValues.length - _numRawValues, + getPendingInputLimit() - _numRawValues - _numIncomingCentroids)); + 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; + } + _originalFractionalBytes = null; + from += numValues; + _numRawValues += numValues; + if (_numRawValues == _rawValues.length + && (!_useTwoLevelRawCompression || _rawValues.length == _pendingCentroidCapacity) + || _numRawValues + _numIncomingCentroids == getPendingInputLimit()) { + flush(); + } + } + } + + @Override + public void add(double value) { + if (Double.isNaN(value)) { + throw new IllegalArgumentException("Cannot add NaN to t-digest"); + } + requireMutable(); + materializePendingSerializedTDigest(); + ensureRawBuffer(); + if (_numRawValues == _rawValues.length + || _numRawValues + _numIncomingCentroids == getPendingInputLimit()) { + flush(); + } + _originalFractionalBytes = null; + _rawValues[_numRawValues++] = value; + } + + @Override + public void add(double value, double weight) { + if (weight == 1) { + add(value); + return; + } + if (Double.isNaN(value)) { + throw new IllegalArgumentException("Cannot add NaN to t-digest"); + } + if (!(weight > 0.0) || !Double.isFinite(weight)) { + throw new IllegalArgumentException("TDigest weight must be positive: " + weight); + } + requireMutable(); + materializePendingSerializedTDigest(); + bufferIncomingCentroid(value, weight); + _originalFractionalBytes = null; + _hasFractionalWeights |= weight != Math.rint(weight); + _min = Math.min(_min, value); + _max = Math.max(_max, value); + } + + @Override + public void add(TDigest other) { + if (other.hasValidStatistics() && other.getTotalWeight() == 0.0) { + return; + } + requireMutable(); + if (other instanceof PercentileTDigestAccumulator) { + addAccumulator((PercentileTDigestAccumulator) other); + return; + } + if (!other.hasValidStatistics()) { + materializePendingSerializedTDigest(); + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(ByteBuffer.wrap(other.serialize()), false); + addSerializedTDigest(input); + return; + } + double otherWeight = other.getTotalWeight(); + checkTotalWeight(getTotalWeight() + otherWeight); + Collection<Centroid> centroids = other.centroids(); + double[] means = new double[centroids.size()]; + double[] weights = new double[means.length]; + int count = 0; + for (Centroid centroid : centroids) { + double weight = centroid.weight(); + if (weight == 0.0) { + continue; + } + if (!(weight > 0.0) || !Double.isFinite(weight) || Double.isNaN(centroid.mean())) { + throw new IllegalArgumentException("Invalid TDigest centroid"); + } + means[count] = centroid.mean(); + weights[count++] = weight; + } + // Validate all generic centroids before changing the target, then use the same whole-source primitive path. + addCentroids(means, weights, count, other.getMin(), other.getMax(), false); + inheritHistoricalFractionalBoundaries(other); + } + + private void addAccumulator(PercentileTDigestAccumulator other) { + if (other == this) { + compress(); + if (_numCentroids == 0) { + return; + } + // Snapshot only self-merges; fresh fractional mass need not be representable by a legacy wire payload. + addCentroids(Arrays.copyOf(_centroidMeans, _numCentroids), Arrays.copyOf(_centroidWeights, _numCentroids), + _numCentroids, _min, _max, false); + return; + } + if (other._originalFractionalBytes != null && getTotalWeight() == 0.0) { + materializePendingSerializedTDigest(); + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(other._originalFractionalBytes); + input._retainedBytes = other._originalFractionalBytes; + addSerializedTDigest(input); + return; + } + if (other._pendingSerializedTDigest != null) { + // Materialize before resetting the reusable reader: materialization uses that same reader. + materializePendingSerializedTDigest(); + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(other._pendingSerializedTDigest, other._pendingSerializedMetadata); + input._retainedBytes = other._pendingSerializedTDigest; + addSerializedTDigest(input); + return; + } + if (other._legacyDegraded) { + addSerializedTDigest(other.serialize()); + return; + } + + // MergingDigest compressed merge sources in place. Aggregation owns these mutable inputs exclusively. + other.compress(); + materializePendingSerializedTDigest(); + checkTotalWeight(_totalWeight + _incomingWeight + _numRawValues + other._totalWeight); + _hasFractionalWeights |= other._hasFractionalWeights; + inheritHistoricalFractionalBoundaries(other._historicalFractionalMinMean, other._historicalFractionalMaxMean); + bufferIncomingCentroids(other._centroidMeans, other._centroidWeights, other._numCentroids, other._totalWeight, + other._min, other._max); + if (other._totalWeight > 0.0) { + _originalFractionalBytes = null; + } + _min = Math.min(_min, other._min); + _max = Math.max(_max, other._max); + } + + public void addSerializedTDigest(byte[] bytes) { + if (_legacyDegraded) { + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(bytes); + addSerializedTDigest(input); + return; + } + requireMutable(); + materializePendingSerializedTDigest(); + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(bytes); + addSerializedTDigest(input); + } + + private SerializedTDigestInput getSerializedTDigestInput() { + if (_serializedTDigestInput == null) { + _serializedTDigestInput = new SerializedTDigestInput(); + } + return _serializedTDigestInput; + } + + public void addSerializedTDigest(SerializedTDigestInput input) { + if (ignoreEmptyInputForDegradedState(input)) { + return; + } + requireMutable(); + if (prepareSerializedTDigest(input)) { + mergeSerializedTDigest(input, false); + } + } + + public void addSerializedTDigestDirect(SerializedTDigestInput input) { + if (ignoreEmptyInputForDegradedState(input)) { + return; + } + requireMutable(); + if (prepareSerializedTDigest(input)) { + mergeSerializedTDigest(input, true); + } + } + + private boolean ignoreEmptyInputForDegradedState(SerializedTDigestInput input) { + if (!_legacyDegraded) { + return false; + } + input.inspectMetadata(); + return !input._metadata.needsLegacyFallback() && input._metadata.totalWeight() == 0.0; + } + + private boolean prepareSerializedTDigest(SerializedTDigestInput input) { + if (_pendingSerializedTDigest == null && hasNoInputs()) { + input.inspectMetadata(); + _hasFractionalWeights |= input._metadata.fractionalWeights(); + _pendingSerializedTDigest = input.retainBytes(); + _pendingSerializedMetadata = input._metadata; + if (input.recordHistoricalFractionalBoundaries(this)) { + // Alias the retained snapshot; historical fractional bytes remain byte-exact after read-only materialization. + _originalFractionalBytes = _pendingSerializedTDigest; + } + if (input._metadata.needsLegacyFallback()) { + _legacyDegraded = true; + } + return false; + } + materializePendingSerializedTDigest(); + input.decode(); + if (input._metadata.needsLegacyFallback()) { + if (hasNoInputs()) { + _pendingSerializedTDigest = input.retainBytes(); + _pendingSerializedMetadata = input._metadata; + _legacyDegraded = true; + return false; + } + throw corruptedMutation(); + } + checkTotalWeight(_totalWeight + _incomingWeight + _numRawValues + input._metadata.totalWeight()); + input.recordHistoricalFractionalBoundaries(this); + if (input._metadata.totalWeight() > 0.0) { + _originalFractionalBytes = null; + } + _hasFractionalWeights |= input._metadata.fractionalWeights(); + return true; + } + + private boolean hasNoInputs() { + return _numCentroids == 0 && _totalWeight == 0.0 && _numRawValues == 0 && _numIncomingCentroids == 0; + } + + private void mergeSerializedTDigest(SerializedTDigestInput input, boolean direct) { + if (input._metadata.totalWeight() == 0.0) { + return; + } + if (hasNoInputs()) { + if (input._metadata.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._metadata.totalWeight(); + _min = input._metadata.min(); + _max = input._metadata.max(); + // Stored centroids already went through the writer's public compression. Reading or repairing their + // endpoints needs another K1 pass only when the output compression differs from the stored compression. + _publiclyCompressed = _compression == input._metadata.compression(); + } + return; + } + if (direct) { + flush(); + boolean runBackwards = (_mergeCount++ & 1) != 0; + if (!mergeSorted(input._means, input._weights, input._numCentroids, input._metadata.totalWeight(), + getIncrementalCompression(), runBackwards)) { + return; + } + } else { + bufferIncomingCentroids(input._means, input._weights, input._numCentroids, input._metadata.totalWeight(), + input._metadata.min(), input._metadata.max()); + } + _min = Math.min(_min, input._metadata.min()); + _max = Math.max(_max, input._metadata.max()); + } + + @Override + public void compress() { + if (_legacyDegraded) { + return; + } + materializePendingSerializedTDigest(); + if (_numRawValues > 0 || _numIncomingCentroids > 0) { + // SQL reduction combines pending values and existing centroids directly at public compression, matching + // MergingDigest 3.3. Its writer avoids an extra working-compression pass that can increase reducer rank error. + flush(_compression); Review Comment: Public compress() now owns the same working-then-public compression policy as serialize() for legacy aggregation. SQL reduction retains its single final pass. Seed-42 regressions compare exact bytes and quantiles for 500 and 2,000 inputs; removing the fix reproduces the 960-versus-944-byte divergence. Addressed in [277a170ffe](https://github.com/apache/pinot/commit/277a170ffe224d680ee00c6f56efd422bb673990). All 526 affected tests, actual 3.2/3.3 reader checks and the unchanged 270-case accuracy oracle pass. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/PercentileTDigestAccumulator.java: ########## @@ -0,0 +1,1498 @@ +/** + * 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.customobject.tdigest; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collection; +import java.util.List; +import org.apache.pinot.segment.local.customobject.tdigest.TDigestCodec.SerializedTDigestMetadata; + + + +/// Accumulates raw values and serialized TDigests into primitive centroids using Pinot's K1 weight-limit rule. +/// +/// Raw values are sorted in small batches; serialized centroids retain their existing sorted order. Both are +/// linearly merged with accumulated centroids. K1 preserves Pinot's existing middle-quantile behavior, and the +/// legacy verbose and small encodings preserve stored data and mixed-version wire compatibility. +/// +/// Serialized state keeps its first digest pending and allocates centroid buffers only when another input must be +/// merged or the digest is queried. Raw-value buffers remain unallocated until a raw value is added. +/// +/// Historical corrupted payloads remain opaque and byte-exact, with NaN statistics. Their reported historical mass +/// may be invalid; they cannot be mutated or mixed with another distribution. Such an operation fails before +/// changing aggregate mass, because neither invented statistics nor synthetic NaN bytes are safe for old readers. +/// Unchanged historical fractional boundaries retain their original bytes, including when configured compression +/// differs. After a merge, inherited fractional endpoints retain precise verbose weights and finite statistics when +/// unit endpoints cannot preserve their mass. These keep the historical t-digest 3.3 assertion-enabled reader and +/// recompression limitation. Newly introduced unsupported fractional endpoints are rejected. +/// +/// Instances are externally serialized per aggregation or group key and are not safe for concurrent mutation. +public 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 MIN_COMPRESSION = 10; + private static final int TWO_LEVEL_COMPRESSION_MULTIPLIER = 2; + private static final int MIN_INCREMENTAL_COMPRESSION = 50; + + private final double _compression; + private final boolean _useTwoLevelRawCompression; + private final double _workingCompression; + private final double _publicMaxWeightScale; + private final double _workingMaxWeightScale; + private final int _centroidCapacity; + private final int _pendingCentroidCapacity; + private double[] _rawValues; + private double[] _centroidMeans; + private double[] _centroidWeights; + private double[] _outputMeans; + private double[] _outputWeights; + private double[] _incomingMeans; + private double[] _incomingWeights; + private double[] _sortedIncomingMeans; + private double[] _sortedIncomingWeights; + private SerializedTDigestInput _serializedTDigestInput; + private byte[] _pendingSerializedTDigest; + private byte[] _originalFractionalBytes; + private double _historicalFractionalMinMean = Double.NaN; + private double _historicalFractionalMaxMean = Double.NaN; + private SerializedTDigestMetadata _pendingSerializedMetadata; + private boolean _legacyDegraded; + private boolean _hasFractionalWeights; + + private int _numRawValues; + private int _numCentroids; + private int _numIncomingCentroids; + private int _mergeCount; + private boolean _publiclyCompressed; + private boolean _incomingCentroidsSorted = true; + private double _totalWeight; + private double _incomingWeight; + private double _min = Double.POSITIVE_INFINITY; + private double _max = Double.NEGATIVE_INFINITY; + + public PercentileTDigestAccumulator(int compression) { + this(positiveCompression(compression), true, true); + } + + public PercentileTDigestAccumulator(double compression) { + this(compression, true, true); + } + + private PercentileTDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers) { + this(compression, allocateRawBuffer, allocateMergeBuffers, false); + } + + private PercentileTDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers, + boolean useTwoLevelRawCompression) { + TDigestCodec.validateCompression(compression); + _useTwoLevelRawCompression = useTwoLevelRawCompression; + _compression = Math.max(MIN_COMPRESSION, compression); + _workingCompression = TWO_LEVEL_COMPRESSION_MULTIPLIER * _compression; + _publicMaxWeightScale = calculateK1MaxWeightScale(_compression); + _workingMaxWeightScale = calculateK1MaxWeightScale(_workingCompression); + long roundedCompression = (long) Math.ceil(_compression); + _centroidCapacity = TDigestCodec.getDefaultCentroidCapacity(_compression); + _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity, useTwoLevelRawCompression); + if (allocateRawBuffer) { + _rawValues = new double[getInitialRawBufferSize(roundedCompression)]; + } + if (allocateMergeBuffers) { + int initialCapacity = Math.min(_centroidCapacity, MIN_RAW_BUFFER_SIZE); + _centroidMeans = new double[initialCapacity]; + _centroidWeights = new double[initialCapacity]; + _outputMeans = new double[initialCapacity]; + _outputWeights = new double[initialCapacity]; + } + } + + /// Creates a lazy digest with the legacy two-level compression used by Star-tree and general digest aggregation. + public static PercentileTDigestAccumulator forLegacyAggregation(double compression) { + return new PercentileTDigestAccumulator(compression, false, false, true); + } + + /// Reads and retains one legacy digest, preserving the general aggregation compression policy. + public static PercentileTDigestAccumulator fromBytes(byte[] bytes) { + return fromBytes(ByteBuffer.wrap(bytes)); + } + + /// Reads one payload and advances the buffer past it. + public static PercentileTDigestAccumulator fromBytes(ByteBuffer bytes) { + SerializedTDigestInput input = new SerializedTDigestInput(); + input.reset(bytes); + return fromBytes(input); + } + + /// Retains an already validated input without rescanning its centroid bytes. + public static PercentileTDigestAccumulator fromBytes(SerializedTDigestInput input) { + PercentileTDigestAccumulator digest = forLegacyAggregation(input.getCompression()); + digest.addSerializedTDigest(input); + return digest; + } + + public static PercentileTDigestAccumulator forReduction(double compression) { + return new PercentileTDigestAccumulator(positiveCompression(compression), false, false); + } + + /// Copies primitive centroids without encoding and decoding them. Zero weights are ignored, and unordered + /// means are sorted when initializing a stored, already compressed distribution; caller arrays are not mutated. + /// Subsequent inputs use the normal buffered merge. Positive weights and consistent extrema are required. + public void addCentroids(double[] means, double[] weights, int count, double min, double max, + boolean alreadyCompressed) { + requireMutable(); + if (count < 0 || count > means.length || count > weights.length) { + throw new IllegalArgumentException("Invalid TDigest centroid count: " + count); + } + double totalWeight = 0.0; + int positiveCount = 0; + boolean sorted = true; + boolean fractionalWeights = false; + double previousMean = Double.NEGATIVE_INFINITY; + for (int i = 0; i < count; i++) { + double weight = weights[i]; + if (weight == 0.0) { + continue; + } + double mean = means[i]; + if (!(weight > 0.0) || !Double.isFinite(weight) || !(mean >= min && mean <= max)) { + throw new IllegalArgumentException("Invalid TDigest centroid"); + } + totalWeight += weight; + sorted &= mean >= previousMean; + previousMean = mean; + fractionalWeights |= weight != Math.rint(weight); + positiveCount++; + } + checkTotalWeight(getTotalWeight() + totalWeight); + if (positiveCount == 0) { + return; + } + materializePendingSerializedTDigest(); + _originalFractionalBytes = null; + _hasFractionalWeights |= fractionalWeights; + if (hasNoInputs() && alreadyCompressed) { + ensureCentroidCapacity(Math.addExact(positiveCount, 2)); + for (int i = 0; i < count; i++) { + if (weights[i] > 0.0) { + _centroidMeans[_numCentroids] = means[i]; + _centroidWeights[_numCentroids++] = weights[i]; + } + } + if (!sorted) { + it.unimi.dsi.fastutil.Arrays.mergeSort(0, _numCentroids, + (first, second) -> Double.compare(_centroidMeans[first], _centroidMeans[second]), + (first, second) -> { + double mean = _centroidMeans[first]; + _centroidMeans[first] = _centroidMeans[second]; + _centroidMeans[second] = mean; + double weight = _centroidWeights[first]; + _centroidWeights[first] = _centroidWeights[second]; + _centroidWeights[second] = weight; + }); + } + _totalWeight = totalWeight; + _min = min; + _max = max; + normalizeBoundaryCentroids(); + _publiclyCompressed = true; + } else { + if (sorted && positiveCount == count) { + bufferIncomingCentroids(means, weights, count, totalWeight, min, max); + } else { + // Caller-owned arrays are copied only when filtering or sorting is necessary. + double[] sortedMeans = new double[positiveCount]; + double[] sortedWeights = new double[positiveCount]; + int index = 0; + for (int i = 0; i < count; i++) { + if (weights[i] > 0.0) { + sortedMeans[index] = means[i]; + sortedWeights[index++] = weights[i]; + } + } + if (!sorted) { + sortCentroids(sortedMeans, sortedWeights, positiveCount); + } + bufferIncomingCentroids(sortedMeans, sortedWeights, positiveCount, totalWeight, min, max); + } + _min = Math.min(_min, min); + _max = Math.max(_max, max); + } + } + + public void add(double[] values, int from, int toExclusive) { + if (from >= toExclusive) { + return; + } + requireMutable(); + materializePendingSerializedTDigest(); + while (from < toExclusive) { + ensureRawBuffer(); + if (_numRawValues == _rawValues.length + || _numRawValues + _numIncomingCentroids == getPendingInputLimit()) { + flush(); + } + int numValues = Math.min(toExclusive - from, Math.min(_rawValues.length - _numRawValues, + getPendingInputLimit() - _numRawValues - _numIncomingCentroids)); + 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; + } + _originalFractionalBytes = null; + from += numValues; + _numRawValues += numValues; + if (_numRawValues == _rawValues.length + && (!_useTwoLevelRawCompression || _rawValues.length == _pendingCentroidCapacity) + || _numRawValues + _numIncomingCentroids == getPendingInputLimit()) { + flush(); + } + } + } + + @Override + public void add(double value) { + if (Double.isNaN(value)) { + throw new IllegalArgumentException("Cannot add NaN to t-digest"); + } + requireMutable(); + materializePendingSerializedTDigest(); + ensureRawBuffer(); + if (_numRawValues == _rawValues.length + || _numRawValues + _numIncomingCentroids == getPendingInputLimit()) { + flush(); + } + _originalFractionalBytes = null; + _rawValues[_numRawValues++] = value; + } + + @Override + public void add(double value, double weight) { + if (weight == 1) { + add(value); + return; + } + if (Double.isNaN(value)) { + throw new IllegalArgumentException("Cannot add NaN to t-digest"); + } + if (!(weight > 0.0) || !Double.isFinite(weight)) { + throw new IllegalArgumentException("TDigest weight must be positive: " + weight); + } + requireMutable(); + materializePendingSerializedTDigest(); + bufferIncomingCentroid(value, weight); + _originalFractionalBytes = null; + _hasFractionalWeights |= weight != Math.rint(weight); + _min = Math.min(_min, value); + _max = Math.max(_max, value); + } + + @Override + public void add(TDigest other) { + if (other.hasValidStatistics() && other.getTotalWeight() == 0.0) { + return; + } + requireMutable(); + if (other instanceof PercentileTDigestAccumulator) { + addAccumulator((PercentileTDigestAccumulator) other); + return; + } + if (!other.hasValidStatistics()) { + materializePendingSerializedTDigest(); + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(ByteBuffer.wrap(other.serialize()), false); + addSerializedTDigest(input); + return; + } + double otherWeight = other.getTotalWeight(); + checkTotalWeight(getTotalWeight() + otherWeight); + Collection<Centroid> centroids = other.centroids(); + double[] means = new double[centroids.size()]; + double[] weights = new double[means.length]; + int count = 0; + for (Centroid centroid : centroids) { + double weight = centroid.weight(); + if (weight == 0.0) { + continue; + } + if (!(weight > 0.0) || !Double.isFinite(weight) || Double.isNaN(centroid.mean())) { + throw new IllegalArgumentException("Invalid TDigest centroid"); + } + means[count] = centroid.mean(); + weights[count++] = weight; + } + // Validate all generic centroids before changing the target, then use the same whole-source primitive path. + addCentroids(means, weights, count, other.getMin(), other.getMax(), false); + inheritHistoricalFractionalBoundaries(other); + } + + private void addAccumulator(PercentileTDigestAccumulator other) { + if (other == this) { + compress(); + if (_numCentroids == 0) { + return; + } + // Snapshot only self-merges; fresh fractional mass need not be representable by a legacy wire payload. + addCentroids(Arrays.copyOf(_centroidMeans, _numCentroids), Arrays.copyOf(_centroidWeights, _numCentroids), + _numCentroids, _min, _max, false); + return; + } + if (other._originalFractionalBytes != null && getTotalWeight() == 0.0) { + materializePendingSerializedTDigest(); + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(other._originalFractionalBytes); + input._retainedBytes = other._originalFractionalBytes; + addSerializedTDigest(input); + return; + } + if (other._pendingSerializedTDigest != null) { + // Materialize before resetting the reusable reader: materialization uses that same reader. + materializePendingSerializedTDigest(); + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(other._pendingSerializedTDigest, other._pendingSerializedMetadata); + input._retainedBytes = other._pendingSerializedTDigest; + addSerializedTDigest(input); + return; + } + if (other._legacyDegraded) { + addSerializedTDigest(other.serialize()); + return; + } + + // MergingDigest compressed merge sources in place. Aggregation owns these mutable inputs exclusively. + other.compress(); + materializePendingSerializedTDigest(); + checkTotalWeight(_totalWeight + _incomingWeight + _numRawValues + other._totalWeight); + _hasFractionalWeights |= other._hasFractionalWeights; + inheritHistoricalFractionalBoundaries(other._historicalFractionalMinMean, other._historicalFractionalMaxMean); + bufferIncomingCentroids(other._centroidMeans, other._centroidWeights, other._numCentroids, other._totalWeight, + other._min, other._max); + if (other._totalWeight > 0.0) { + _originalFractionalBytes = null; + } + _min = Math.min(_min, other._min); + _max = Math.max(_max, other._max); + } + + public void addSerializedTDigest(byte[] bytes) { + if (_legacyDegraded) { + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(bytes); + addSerializedTDigest(input); + return; + } + requireMutable(); + materializePendingSerializedTDigest(); + SerializedTDigestInput input = getSerializedTDigestInput(); + input.reset(bytes); + addSerializedTDigest(input); + } + + private SerializedTDigestInput getSerializedTDigestInput() { + if (_serializedTDigestInput == null) { + _serializedTDigestInput = new SerializedTDigestInput(); + } + return _serializedTDigestInput; + } + + public void addSerializedTDigest(SerializedTDigestInput input) { + if (ignoreEmptyInputForDegradedState(input)) { + return; + } + requireMutable(); + if (prepareSerializedTDigest(input)) { + mergeSerializedTDigest(input, false); + } + } + + public void addSerializedTDigestDirect(SerializedTDigestInput input) { + if (ignoreEmptyInputForDegradedState(input)) { + return; + } + requireMutable(); + if (prepareSerializedTDigest(input)) { + mergeSerializedTDigest(input, true); + } + } + + private boolean ignoreEmptyInputForDegradedState(SerializedTDigestInput input) { + if (!_legacyDegraded) { + return false; + } + input.inspectMetadata(); + return !input._metadata.needsLegacyFallback() && input._metadata.totalWeight() == 0.0; + } + + private boolean prepareSerializedTDigest(SerializedTDigestInput input) { + if (_pendingSerializedTDigest == null && hasNoInputs()) { + input.inspectMetadata(); + _hasFractionalWeights |= input._metadata.fractionalWeights(); + _pendingSerializedTDigest = input.retainBytes(); + _pendingSerializedMetadata = input._metadata; + if (input.recordHistoricalFractionalBoundaries(this)) { + // Alias the retained snapshot; historical fractional bytes remain byte-exact after read-only materialization. + _originalFractionalBytes = _pendingSerializedTDigest; + } + if (input._metadata.needsLegacyFallback()) { + _legacyDegraded = true; + } + return false; + } + materializePendingSerializedTDigest(); + input.decode(); + if (input._metadata.needsLegacyFallback()) { + if (hasNoInputs()) { + _pendingSerializedTDigest = input.retainBytes(); + _pendingSerializedMetadata = input._metadata; + _legacyDegraded = true; + return false; + } + throw corruptedMutation(); + } + checkTotalWeight(_totalWeight + _incomingWeight + _numRawValues + input._metadata.totalWeight()); + input.recordHistoricalFractionalBoundaries(this); + if (input._metadata.totalWeight() > 0.0) { + _originalFractionalBytes = null; + } + _hasFractionalWeights |= input._metadata.fractionalWeights(); + return true; + } + + private boolean hasNoInputs() { + return _numCentroids == 0 && _totalWeight == 0.0 && _numRawValues == 0 && _numIncomingCentroids == 0; + } + + private void mergeSerializedTDigest(SerializedTDigestInput input, boolean direct) { + if (input._metadata.totalWeight() == 0.0) { + return; + } + if (hasNoInputs()) { + if (input._metadata.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._metadata.totalWeight(); + _min = input._metadata.min(); + _max = input._metadata.max(); + // Stored centroids already went through the writer's public compression. Reading or repairing their + // endpoints needs another K1 pass only when the output compression differs from the stored compression. + _publiclyCompressed = _compression == input._metadata.compression(); + } + return; + } + if (direct) { + flush(); + boolean runBackwards = (_mergeCount++ & 1) != 0; + if (!mergeSorted(input._means, input._weights, input._numCentroids, input._metadata.totalWeight(), + getIncrementalCompression(), runBackwards)) { + return; + } + } else { + bufferIncomingCentroids(input._means, input._weights, input._numCentroids, input._metadata.totalWeight(), + input._metadata.min(), input._metadata.max()); + } + _min = Math.min(_min, input._metadata.min()); + _max = Math.max(_max, input._metadata.max()); + } + + @Override + public void compress() { + if (_legacyDegraded) { + return; + } + materializePendingSerializedTDigest(); + if (_numRawValues > 0 || _numIncomingCentroids > 0) { + // SQL reduction combines pending values and existing centroids directly at public compression, matching + // MergingDigest 3.3. Its writer avoids an extra working-compression pass that can increase reducer rank error. + flush(_compression); + _publiclyCompressed = true; + } else if (!_publiclyCompressed && _numCentroids > 0) { + recompressCentroids(_compression); + _publiclyCompressed = true; + } + } + + @Override + public boolean hasValidStatistics() { + return !_legacyDegraded; + } + + @Override + public double getTotalWeight() { + return _pendingSerializedTDigest != null ? _pendingSerializedMetadata.totalWeight() + : _totalWeight + _incomingWeight + _numRawValues; + } + + /// Bounds serialized output centroids from buffered inputs, including up to two boundary splits before a merge + /// and two more before writing, without flushing or using the potentially huge input mass. + public long getCentroidCountUpperBound() { + return _pendingSerializedTDigest != null ? (long) _pendingSerializedMetadata.centroidCount() + 4L + : (long) _numCentroids + _numRawValues + _numIncomingCentroids + 4L; + } + + /// Whether unchanged historical fractional boundaries still require their original byte representation. + public boolean hasOriginalFractionalPayload() { + return _originalFractionalBytes != null; + } + + @Override + public double cdf(double value) { + if (Double.isNaN(value) || Double.isInfinite(value)) { + throw new IllegalArgumentException(String.format("Invalid value: %f", value)); + } + if (_legacyDegraded) { + return Double.NaN; + } + materializePendingSerializedTDigest(); + flush(); + 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; + if (width > 0.0) { + if (_centroidWeights[0] < 1.0) { + return (_centroidWeights[0] / 2.0 * (value - _min) / width) / _totalWeight; + } + return value == _min ? 0.5 / _totalWeight + : (1.0 + (value - _min) / width * (_centroidWeights[0] / 2.0 - 1.0)) / _totalWeight; + } + return 0.0; + } + int lastIndex = _numCentroids - 1; + if (value > _centroidMeans[lastIndex]) { + double width = _max - _centroidMeans[lastIndex]; + if (width > 0.0) { + if (_centroidWeights[lastIndex] < 1.0) { + return 1.0 - (_centroidWeights[lastIndex] / 2.0 * (_max - value) / width) / _totalWeight; + } + return value == _max ? 1.0 - 0.5 / _totalWeight + : 1.0 - (1.0 + (_max - value) / width * (_centroidWeights[lastIndex] / 2.0 - 1.0)) + / _totalWeight; + } + return 1.0; + } + + double weightSoFar = 0.0; + for (int i = 0; i < lastIndex; i++) { + if (_centroidMeans[i] == value) { + double equalWeight = 0.0; + int equalIndex = i; + while (equalIndex < _numCentroids && _centroidMeans[equalIndex] == value) { + equalWeight += _centroidWeights[equalIndex]; + equalIndex++; + } + return (weightSoFar + equalWeight / 2.0) / _totalWeight; + } + if (_centroidMeans[i] <= value && _centroidMeans[i + 1] > value) { + if (!Double.isFinite(_centroidMeans[i]) || !Double.isFinite(_centroidMeans[i + 1])) { + return (weightSoFar + _centroidWeights[i]) / _totalWeight; + } + double width = _centroidMeans[i + 1] - _centroidMeans[i]; + if (width > 0.0) { + double leftExcludedWeight = getPointMassHalfWeight(i); + double rightExcludedWeight = getPointMassHalfWeight(i + 1); + double halfWeight = (_centroidWeights[i] + _centroidWeights[i + 1]) / 2.0; + double interpolatedWeight = halfWeight - leftExcludedWeight - rightExcludedWeight; + double base = weightSoFar + _centroidWeights[i] / 2.0 + leftExcludedWeight; + return (base + interpolatedWeight * (value - _centroidMeans[i]) / width) / _totalWeight; + } + return (weightSoFar + (_centroidWeights[i] + _centroidWeights[i + 1]) / 2.0) / _totalWeight; + } + weightSoFar += _centroidWeights[i]; + } + if (value == _centroidMeans[lastIndex]) { + return 1.0 - Math.min(1.0, _centroidWeights[lastIndex]) / 2.0 / _totalWeight; + } + throw new IllegalStateException("Unable to compute TDigest CDF"); + } + + @Override + public double quantile(double quantile) { + if (Double.isNaN(quantile) || quantile < 0.0 || quantile > 1.0) { + throw new IllegalArgumentException("q should be in [0,1], got " + quantile); + } + if (_legacyDegraded) { + return Double.NaN; + } + materializePendingSerializedTDigest(); + flush(); + if (_numCentroids == 0) { + return Double.NaN; + } + if (quantile == 0.0) { + return _min; + } + if (quantile == 1.0) { + return _max; + } + if (_numCentroids == 1) { + return _centroidMeans[0]; + } + + double index = quantile * _totalWeight; + double firstWeight = _centroidWeights[0]; + if (firstWeight < 1.0 && index < firstWeight / 2.0) { + return weightedAverage(_min, firstWeight / 2.0 - index, _centroidMeans[0], index); + } + if (firstWeight >= 1.0 && index < 1.0) { + return _min; + } + if (_centroidWeights[0] > 1.0 && index < _centroidWeights[0] / 2.0) { + return _min + (index - 1.0) / (_centroidWeights[0] / 2.0 - 1.0) * (_centroidMeans[0] - _min); + } + int lastIndex = _numCentroids - 1; + double lastWeight = _centroidWeights[lastIndex]; + if (lastWeight < 1.0 && index >= _totalWeight - lastWeight / 2.0) { + return weightedAverage(_centroidMeans[lastIndex], _totalWeight - index, _max, + index - (_totalWeight - lastWeight / 2.0)); + } + if (lastWeight >= 1.0 && index > _totalWeight - 1.0) { + return _max; + } + if (_centroidWeights[lastIndex] > 1.0 + && _totalWeight - index <= _centroidWeights[lastIndex] / 2.0) { + return _max - (_totalWeight - index - 1.0) / (_centroidWeights[lastIndex] / 2.0 - 1.0) + * (_max - _centroidMeans[lastIndex]); + } + + double weightSoFar = _centroidWeights[0] / 2.0; + for (int i = 0; i < lastIndex; i++) { + double halfWeight = (_centroidWeights[i] + _centroidWeights[i + 1]) / 2.0; + if (weightSoFar + halfWeight > index) { + double leftExcludedWeight = getPointMassHalfWeight(i); + if (index - weightSoFar < leftExcludedWeight) { + return _centroidMeans[i]; + } + double rightExcludedWeight = getPointMassHalfWeight(i + 1); + if (weightSoFar + halfWeight - index <= rightExcludedWeight) { + return _centroidMeans[i + 1]; + } + double leftWeight = index - weightSoFar - leftExcludedWeight; + double rightWeight = weightSoFar + halfWeight - index - rightExcludedWeight; + return weightedAverage(_centroidMeans[i], rightWeight, _centroidMeans[i + 1], leftWeight); + } + weightSoFar += halfWeight; + } + double weightToMax = index - _totalWeight - _centroidWeights[lastIndex] / 2.0; + double weightFromLastCentroid = _centroidWeights[lastIndex] / 2.0 - weightToMax; + return weightedAverage(_centroidMeans[lastIndex], weightToMax, _max, weightFromLastCentroid); + } + + private double getPointMassHalfWeight(int index) { Review Comment: Documented the intentional duplicate-plateau interpolation improvement and possible cross-upgrade result changes in the migration guide and PR description. Added the raw 70/20/10 duplicate-heavy fixture alongside the existing plateau/CDF tests. Removing plateau exclusion makes its q(0.69) shift from 0.1 to about 0.172, so the regression detects the old interpolation. Addressed in [277a170ffe](https://github.com/apache/pinot/commit/277a170ffe224d680ee00c6f56efd422bb673990). All 526 affected tests, actual 3.2/3.3 reader checks and the unchanged 270-case accuracy oracle pass. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/NonFiniteAwareTDigest.java: ########## @@ -0,0 +1,500 @@ +/** + * 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.customobject.tdigest; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.Collection; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.segment.local.customobject.tdigest.TDigestCodec.SerializedTDigestMetadata; + + +/// Stores non-finite values as exact tail masses while delegating finite values to an owned native digest. +/// Emits legacy t-digest bytes and preserves historical fractional endpoint provenance across merges and copies. +/// Instances are mutable and require external synchronization. +public final class NonFiniteAwareTDigest extends TDigest { + private final PercentileTDigestAccumulator _finiteDigest; + private double _negativeInfinityWeight; + private double _positiveInfinityWeight; + private List<Centroid> _finiteCentroids; + private byte[] _originalFractionalBytes; + + public static NonFiniteAwareTDigest forLegacyAggregation(double compression) { + return new NonFiniteAwareTDigest(PercentileTDigestAccumulator.forLegacyAggregation(compression), 0.0, 0.0); + } + + private NonFiniteAwareTDigest(PercentileTDigestAccumulator finiteDigest, double negativeInfinityWeight, + double positiveInfinityWeight) { + _finiteDigest = finiteDigest; + _negativeInfinityWeight = negativeInfinityWeight; + _positiveInfinityWeight = positiveInfinityWeight; + } + + public NonFiniteAwareTDigest copy() { + if (_originalFractionalBytes != null) { + return fromBytes(_originalFractionalBytes); + } + PercentileTDigestAccumulator source = _finiteDigest; + PercentileTDigestAccumulator finiteDigest; + if (!source.hasValidStatistics() || source.hasOriginalFractionalPayload()) { + finiteDigest = PercentileTDigestAccumulator.fromBytes(source.serialize()); + } else { + finiteDigest = PercentileTDigestAccumulator.forLegacyAggregation(compression()); + if (!source.isEmpty()) { + finiteDigest.add(source); + invalidateDerivedCaches(); + } + } + if (source.hasValidStatistics()) { + finiteDigest.inheritHistoricalFractionalBoundaries(source); + } + return new NonFiniteAwareTDigest(finiteDigest, _negativeInfinityWeight, _positiveInfinityWeight); + } + + public static NonFiniteAwareTDigest fromBytes(byte[] bytes) { + SerializedTDigestInput validated = new SerializedTDigestInput(); + validated.reset(bytes); + SerializedTDigestMetadata metadata = validated.getMetadata(); + if (!metadata.hasNonFiniteMeans() || metadata.needsLegacyFallback()) { + return new NonFiniteAwareTDigest(PercentileTDigestAccumulator.fromBytes(validated), 0.0, 0.0); + } + double encodedMin = metadata.min(); + double encodedMax = metadata.max(); + double compression = metadata.compression(); + int centroidCount = metadata.centroidCount(); + double negativeInfinityWeight = 0.0; + double positiveInfinityWeight = 0.0; + double[] finiteMeans = new double[centroidCount]; + double[] finiteWeights = new double[centroidCount]; + TDigestCodec.inspectSerialized(ByteBuffer.wrap(bytes), metadata, finiteMeans, finiteWeights); + int finiteCount = 0; + double finiteMin = Double.POSITIVE_INFINITY; + double finiteMax = Double.NEGATIVE_INFINITY; + for (int i = 0; i < centroidCount; i++) { + double weight = finiteWeights[i]; + double mean = finiteMeans[i]; + if (weight == 0.0) { + continue; + } + if (mean == Double.NEGATIVE_INFINITY) { + negativeInfinityWeight += weight; + } else if (mean == Double.POSITIVE_INFINITY) { + positiveInfinityWeight += weight; + } else { + finiteWeights[finiteCount] = weight; + finiteMeans[finiteCount] = mean; + finiteMin = Math.min(finiteMin, mean); + finiteMax = Math.max(finiteMax, mean); + finiteCount++; + } + } + + PercentileTDigestAccumulator finiteDigest = + PercentileTDigestAccumulator.forLegacyAggregation(compression); + if (finiteCount > 0) { + finiteMin = Double.isFinite(encodedMin) ? encodedMin : finiteMin; + finiteMax = Double.isFinite(encodedMax) ? encodedMax : finiteMax; + finiteDigest.addCentroids(finiteMeans, finiteWeights, finiteCount, finiteMin, finiteMax, true); + } + validated.recordHistoricalFractionalBoundaries(finiteDigest); + NonFiniteAwareTDigest wrapped = + new NonFiniteAwareTDigest(finiteDigest, negativeInfinityWeight, positiveInfinityWeight); + if (metadata.fractionalWeights()) { + wrapped._originalFractionalBytes = bytes.clone(); + } + return wrapped; + } + + @Override + public void add(double value) { + add(value, 1); + } + + @Override + public void add(double value, double weight) { + if (Double.isNaN(value)) { + throw new IllegalArgumentException("Cannot add NaN to t-digest"); + } + if (!(weight > 0.0) || !Double.isFinite(weight)) { + throw new IllegalArgumentException("TDigest weight must be positive: " + weight); + } + requireMutable(); + checkTotalWeight(getTotalWeight() + weight); + if (value == Double.NEGATIVE_INFINITY) { + _negativeInfinityWeight += weight; + } else if (value == Double.POSITIVE_INFINITY) { + _positiveInfinityWeight += weight; + } else { + _finiteDigest.add(value, weight); + } + mutated(); + } + + @Override + public void add(TDigest other) { + if (other.isEmpty()) { + return; + } + requireMutable(); + if (!other.hasValidStatistics() && !isEmpty()) { + throw corruptedMutation(); + } + if (other.hasValidStatistics()) { + checkTotalWeight(getTotalWeight() + other.getTotalWeight()); + } + if (other instanceof NonFiniteAwareTDigest) { + NonFiniteAwareTDigest wrapped = (NonFiniteAwareTDigest) other; + double negativeInfinityWeight = wrapped._negativeInfinityWeight; + double positiveInfinityWeight = wrapped._positiveInfinityWeight; + if (!wrapped._finiteDigest.isEmpty()) { + _finiteDigest.add(wrapped._finiteDigest); + wrapped.invalidateDerivedCaches(); + } + if (wrapped.hasValidStatistics()) { + _finiteDigest.inheritHistoricalFractionalBoundaries(wrapped); + } + _negativeInfinityWeight += negativeInfinityWeight; + _positiveInfinityWeight += positiveInfinityWeight; + mutated(); + return; + } + if (!other.hasValidStatistics() || Double.isFinite(other.getMin()) && Double.isFinite(other.getMax())) { + _finiteDigest.add(other); + mutated(); + return; + } + // Read precise centroid weights directly; generic sources need no full digest serialization for their tails. + Collection<Centroid> centroids = other.centroids(); + double[] finiteMeans = new double[centroids.size()]; + double[] finiteWeights = new double[centroids.size()]; + int finiteCount = 0; + double finiteMin = Double.POSITIVE_INFINITY; + double finiteMax = Double.NEGATIVE_INFINITY; + double negativeInfinityWeight = 0.0; + double positiveInfinityWeight = 0.0; + double totalWeight = getTotalWeight(); + for (Centroid centroid : centroids) { + double mean = centroid.mean(); + double weight = centroid.weight(); + if (weight == 0.0) { + continue; + } + if (!(weight > 0.0) || !Double.isFinite(weight) || Double.isNaN(mean)) { + throw new IllegalArgumentException("Invalid TDigest centroid"); + } + totalWeight += weight; + if (mean == Double.NEGATIVE_INFINITY) { + negativeInfinityWeight += weight; + } else if (mean == Double.POSITIVE_INFINITY) { + positiveInfinityWeight += weight; + } else { + finiteMeans[finiteCount] = mean; + finiteWeights[finiteCount++] = weight; + finiteMin = Math.min(finiteMin, mean); + finiteMax = Math.max(finiteMax, mean); + } + } + checkTotalWeight(totalWeight); + if (finiteCount > 0) { + finiteMin = Double.isFinite(other.getMin()) ? other.getMin() : finiteMin; + finiteMax = Double.isFinite(other.getMax()) ? other.getMax() : finiteMax; + _finiteDigest.addCentroids( + finiteMeans, finiteWeights, finiteCount, finiteMin, finiteMax, false); + } + _finiteDigest.inheritHistoricalFractionalBoundaries(other); + _negativeInfinityWeight += negativeInfinityWeight; + _positiveInfinityWeight += positiveInfinityWeight; + mutated(); + } + + private void requireMutable() { + if (!hasValidStatistics()) { + throw corruptedMutation(); + } + } + + private static IllegalArgumentException corruptedMutation() { + return new IllegalArgumentException("Cannot merge or mutate a historically corrupted TDigest; " + + "rebuild stored digests from source data before merging"); + } + + private void mutated() { + _originalFractionalBytes = null; + invalidateDerivedCaches(); + } + + @Override + public void compress() { + _finiteDigest.compress(); + invalidateDerivedCaches(); + } + + @Override + public double getTotalWeight() { + return _finiteDigest.getTotalWeight() + _negativeInfinityWeight + _positiveInfinityWeight; + } + + @Override + double getHistoricalFractionalBoundaryMean(boolean lowerBoundary) { + return _finiteDigest.getHistoricalFractionalBoundaryMean(lowerBoundary); + } + + @Override + public boolean hasValidStatistics() { + return _finiteDigest.hasValidStatistics(); + } + + @Override + public double cdf(double value) { + if (Double.isNaN(value) || Double.isInfinite(value)) { + throw new IllegalArgumentException(String.format("Invalid value: %f", value)); + } + double size = getTotalWeight(); + if (size == 0.0 || !hasValidStatistics()) { + return Double.NaN; + } + double finiteSize = _finiteDigest.getTotalWeight(); + double finiteWeight = finiteSize == 0L ? 0.0 : _finiteDigest.cdf(value) * finiteSize; + return (_negativeInfinityWeight + finiteWeight) / size; + } + + @Override + public double quantile(double quantile) { + if (Double.isNaN(quantile) || quantile < 0.0 || quantile > 1.0) { + throw new IllegalArgumentException("q should be in [0,1], got " + quantile); + } + double size = getTotalWeight(); + if (size == 0.0 || !hasValidStatistics()) { + return Double.NaN; + } + double index = quantile * size; + if (_negativeInfinityWeight > 0L && index < _negativeInfinityWeight) { + return Double.NEGATIVE_INFINITY; + } + if (_positiveInfinityWeight > 0L && index >= size - _positiveInfinityWeight) { + return Double.POSITIVE_INFINITY; + } + double finiteSize = _finiteDigest.getTotalWeight(); + if (finiteSize == 0L) { + return _negativeInfinityWeight > 0L ? Double.NEGATIVE_INFINITY : Double.POSITIVE_INFINITY; + } + double finiteQuantile = (index - _negativeInfinityWeight) / finiteSize; + return _finiteDigest.quantile(Math.max(0.0, Math.min(finiteQuantile, 1.0))); + } + + @Override + public Collection<Centroid> centroids() { + return getAllCentroids(); + } + + @Override + public double compression() { + return _finiteDigest.compression(); + } + + @Override + public int maxSerializedByteSize() { + if (_originalFractionalBytes != null) { + return _originalFractionalBytes.length; + } + if (_negativeInfinityWeight == 0.0 && _positiveInfinityWeight == 0.0) { + return _finiteDigest.maxSerializedByteSize(); + } + int finiteByteSize = _finiteDigest.maxSerializedByteSize(); + PercentileTDigestAccumulator finiteDigest = _finiteDigest; + if (finiteDigest.hasOriginalFractionalPayload()) { + // Infinite tails can make retained fractional boundaries interior; their finite view may split endpoints. + finiteByteSize = Math.max(finiteByteSize, + TDigestCodec.getMaxLegacyCompatibleByteSize(compression(), finiteDigest.getCentroidCountUpperBound())); + } + boolean hasFiniteValues = _finiteDigest.getTotalWeight() > 0.0; + int negativeInfinityCount = getInfinityCentroidCount(_negativeInfinityWeight, true, + !hasFiniteValues && _positiveInfinityWeight == 0.0); + int positiveInfinityCount = getInfinityCentroidCount(_positiveInfinityWeight, + !hasFiniteValues && _negativeInfinityWeight == 0.0, true); + int tailCount = Math.addExact(negativeInfinityCount, positiveInfinityCount); + int tailByteSize = Math.multiplyExact(TDigestCodec.VERBOSE_CENTROID_SIZE, tailCount); + long finiteCount = hasFiniteValues ? finiteDigest.getCentroidCountUpperBound() : 0L; + int compatibleBound = + TDigestCodec.getMaxLegacyCompatibleByteSize(compression(), Math.addExact(finiteCount, tailCount)); + return Math.min(Math.addExact(finiteByteSize, tailByteSize), compatibleBound); + } + + @Override + public int centroidCount() { + return getAllCentroids().size(); + } + + @Override + public double getMin() { + if (_negativeInfinityWeight > 0L) { + return Double.NEGATIVE_INFINITY; + } + if (!_finiteDigest.isEmpty()) { + return _finiteDigest.getMin(); + } + return Double.POSITIVE_INFINITY; + } + + @Override + public double getMax() { + if (_positiveInfinityWeight > 0L) { + return Double.POSITIVE_INFINITY; + } + if (!_finiteDigest.isEmpty()) { + return _finiteDigest.getMax(); + } + return Double.NEGATIVE_INFINITY; + } + + private List<Centroid> getFiniteCentroids() { + if (_finiteCentroids == null) { + // Match the wire view without serializing finite-only fractional boundaries surrounded by infinite tails. + PercentileTDigestAccumulator finiteDigest = _finiteDigest; + finiteDigest.prepareCentroidsForSerialization(); + invalidateDerivedCaches(); + _finiteCentroids = new ArrayList<>(_finiteDigest.centroids()); + } + return _finiteCentroids; + } + + private List<Centroid> getAllCentroids() { + List<Centroid> finiteCentroids = getFiniteCentroids(); + boolean hasFiniteValues = _finiteDigest.getTotalWeight() > 0.0; + List<Centroid> centroids = new ArrayList<>(finiteCentroids.size() + 6); + appendInfinityCentroids(centroids, Double.NEGATIVE_INFINITY, _negativeInfinityWeight, true, + !hasFiniteValues && _positiveInfinityWeight == 0L); + centroids.addAll(finiteCentroids); + appendInfinityCentroids(centroids, Double.POSITIVE_INFINITY, _positiveInfinityWeight, + !hasFiniteValues && _negativeInfinityWeight == 0L, true); + if (centroids.size() > 1 + && (centroids.getFirst().weight() > 1.0 || centroids.getLast().weight() > 1.0)) { + int count = centroids.size(); + double[] means = new double[Math.addExact(count, 2)]; + double[] weights = new double[means.length]; + for (int i = 0; i < count; i++) { + Centroid centroid = centroids.get(i); + means[i] = centroid.mean(); + weights[i] = centroid.weight(); + } + count = PercentileTDigestAccumulator.normalizeSerializedBoundaries( Review Comment: Degraded wrapper centroid views now skip boundary normalization and preserve the encoded means and weights. Tests cover NaN/inverted extrema, infinite weights and overflowing total mass, including cloned views and byte-exact serialization. Addressed in [277a170ffe](https://github.com/apache/pinot/commit/277a170ffe224d680ee00c6f56efd422bb673990). All 526 affected tests, actual 3.2/3.3 reader checks and the unchanged 270-case accuracy oracle pass. -- 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]
