xiangfu0 commented on code in PR #19795: URL: https://github.com/apache/pinot/pull/19795#discussion_r4244257238
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigestCodec.java: ########## @@ -0,0 +1,850 @@ +/** + * 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.lang.invoke.MethodHandles; +import java.lang.invoke.VarHandle; +import java.nio.BufferUnderflowException; +import java.nio.ByteBuffer; +import java.nio.ByteOrder; +import java.util.Arrays; + +/// Compatibility helpers for the serialized t-digest format shared by t-digest 3.2 and 3.3. +/// +/// Version 3.3 requires the first and last centroid to have unit weight whenever a digest is recompressed. Version +/// 3.2 did not maintain that invariant, so an otherwise valid 3.2 digest can fail an assertion after being read by +/// 3.3. Pure reads retain the legacy weighted boundary centroids. Adjacent equal means use Pinot's point-mass +/// interpolation, so quantile/CDF values can change on unchanged bytes, including intervals next to a plateau. +/// Mutation and serialization repair their boundaries only when unit endpoints can preserve their total weight and +/// first moment; otherwise unchanged historical bytes remain readable and mutation fails explicitly. +/// Serialization uses the compact form when 3.3's larger low-compression centroid buffer would not fit in a 3.2 +/// default buffer and every centroid is exactly representable in that format. Otherwise, it merges the nearest +/// interior centroids until the verbose representation fits the 3.2 capacity. This keeps intermediate values readable +/// in both directions during a rolling upgrade without silently narrowing double-precision values to floats. +/// +/// Pinot owns the in-memory K1 implementation and retains these legacy bytes for stored columns and mixed-version +/// intermediate results. Decoding checks lengths and numerical state before allocating centroid arrays. +/// NaN means at an identifiable infinite tail are restored to that infinite bound without changing their weights. +/// Other poisoned states remain byte-readable with NaN statistics because their distribution cannot be recovered. +/// They retain their original bytes; attempting to mutate them or mix them with another distribution fails explicitly. +/// Opaque poisoned payloads retain their original centroids even if an old reader's default array cannot hold them; +/// these historical exceptions remain readable by Pinot, but are not guaranteed to survive a reverse upgrade. +/// Historical fractional boundary masses below one, and singleton masses between one and two, retain their original +/// bytes. Fresh serialization rejects these boundaries: t-digest 3.3's unit-endpoint assertion cannot be satisfied +/// without changing their mass. Fractional interior weights remain supported. +public final class TDigestCodec { + public static final int VERBOSE_ENCODING = 1; + public static final int SMALL_ENCODING = 2; + private static final int LEGACY_CAPACITY_PADDING = 10; + public static final int VERBOSE_HEADER_SIZE = 32; + public static final int VERBOSE_CENTROID_SIZE = 16; + public static final int SMALL_HEADER_SIZE = 30; + public static final int SMALL_CENTROID_SIZE = 8; + public static final int DEFAULT_MERGE_BUFFER_MULTIPLIER = 5; + + private TDigestCodec() { + } + + /// Reads the compression from a validated header without decoding centroid arrays. + public static double readCompression(byte[] bytes) { + return readSerializedHeader(ByteBuffer.wrap(bytes)).compression(); + } + + /// Rejects non-finite compression while retaining the legacy clamping of finite settings below ten. + public static void validateCompression(double compression) { + if (!Double.isFinite(compression)) { + throw new IllegalArgumentException("Invalid TDigest compression: " + compression); + } + } + + /// Applies the legacy minimum compression to query, ingestion and serialized-header settings. + public static double normalizeCompression(double compression) { + validateCompression(compression); + return Math.max(compression, 10.0); + } + + /// Writes validated centroid arrays in the standard verbose encoding without sorting or recompression. + /// The caller retains the arrays; only the first `count` entries are written, with normalized compression. + public static byte[] serializeCentroids(double compression, double min, double max, double[] means, + double[] weights, int count) { + if (count < 0 || count > means.length || count > weights.length) { + throw new IllegalArgumentException("Invalid TDigest centroid count: " + count); + } + ByteBuffer encoded = ByteBuffer.allocate(Math.addExact(VERBOSE_HEADER_SIZE, + Math.multiplyExact(VERBOSE_CENTROID_SIZE, count))); + encoded.putInt(VERBOSE_ENCODING).putDouble(min).putDouble(max).putDouble(normalizeCompression(compression)); + encoded.putInt(count); + for (int i = 0; i < count; i++) { + encoded.putDouble(weights[i]).putDouble(means[i]); + } + return encoded.array(); + } + + /// Encodes validated, sorted native centroids and selects a legacy-compatible representation. + static byte[] serializeCompatibleCentroids(double compression, double min, double max, double[] means, + double[] weights, int count, boolean inheritedFractionalBoundaries) { + if (count < 0 || count > 0 && (count > means.length || count > weights.length || !(min <= max))) { + throw new IllegalArgumentException("Invalid TDigest centroid count or extrema"); + } + ByteBuffer buffer = ByteBuffer.allocate( + Math.addExact(VERBOSE_HEADER_SIZE, + Math.multiplyExact(VERBOSE_CENTROID_SIZE, count))); + buffer.putInt(VERBOSE_ENCODING); + buffer.putDouble(min); + buffer.putDouble(max); + buffer.putDouble(compression); + buffer.putInt(count); + double encodedWeight = 0.0; + boolean hasNonFiniteMeans = false; + boolean fractionalWeights = false; + for (int i = 0; i < count; i++) { + if (!(weights[i] > 0.0) || !Double.isFinite(weights[i]) || !(means[i] >= min && means[i] <= max)) { + throw new IllegalArgumentException("Invalid TDigest centroid"); + } + buffer.putDouble(weights[i]); + buffer.putDouble(means[i]); + encodedWeight += weights[i]; + hasNonFiniteMeans |= !Double.isFinite(means[i]); + fractionalWeights |= weights[i] != Math.rint(weights[i]); + } + if (!Double.isFinite(encodedWeight)) { + throw new IllegalArgumentException("TDigest total weight exceeds the supported range"); + } + boolean weightedBoundaries = count == 1 ? weights[0] >= 2.0 + : count > 1 && (weights[0] > 1.0 || weights[count - 1] > 1.0); + // Collect exact wire metadata during the write rather than decoding and validating our fresh bytes again. + SerializedTDigestMetadata metadata = new SerializedTDigestMetadata(VERBOSE_ENCODING, min, max, + compression, count, + VERBOSE_HEADER_SIZE, VERBOSE_CENTROID_SIZE, buffer.capacity(), encodedWeight, + hasNonFiniteMeans, false, false, fractionalWeights, weightedBoundaries, false, compression, Double.NaN, + Double.NaN, Double.NaN); + return makeLegacyCompatible(buffer.array(), metadata, inheritedFractionalBoundaries); + } + + /// Rejects fresh boundary masses that legacy readers cannot recompress without inventing mass. + /// Call after boundary normalization and removing zero-weight centroids; fractional interior mass is supported. + static void checkLegacyBoundaryWeights(int count, double firstWeight, double lastWeight) { + if (count > 0 && (firstWeight < 1.0 || lastWeight < 1.0 Review Comment: Keeping validation at serialization is intentional: add(5, 0.5) followed by add(4) and add(6) produces supported fractional interior mass and serializes successfully (total 2.5), so rejecting the first add would prohibit that valid use. Verbose encoding preserves fractions but does not repair the old compressor's unit-endpoint assumptions. The existing API/compatibility notes and tests explicitly cover the fresh singleton rejection and inherited-byte exception. Also, genuine 3.3's public API takes add(double, int), not add(double, double); fractional historical encodings are a separate case from that public call. Rechecked both outcomes without changing this policy. Validated at `3821623d46933dd233b8110e123ecb4f86ac1e96`: 554 affected tests, genuine 3.2/3.3 bidirectional byte checks, the 270-case oracle, applicable precommit checks and independent Pinot review pass. New-head CI and maintainer approval remain separate gates. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/NonFiniteAwareTDigest.java: ########## @@ -0,0 +1,726 @@ +/** + * 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.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 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 PercentileTDigestAccumulator _finiteDigest; + private double _negativeInfinityWeight; + private double _positiveInfinityWeight; + private List<Centroid> _finiteCentroids; + private byte[] _originalBytes; + private boolean _unrepresentableInfinityTail; + private boolean _pendingHistoricalBoundaryRepair; + private double _historicalMin; + private double _historicalMax; + + 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; + } + + /// Snapshots the aggregation without flushing or compressing the source's buffered inputs. + public NonFiniteAwareTDigest copy() { + NonFiniteAwareTDigest copy = + new NonFiniteAwareTDigest(_finiteDigest.copy(), _negativeInfinityWeight, _positiveInfinityWeight); + copy._originalBytes = _originalBytes == null ? null : _originalBytes.clone(); + copy._unrepresentableInfinityTail = _unrepresentableInfinityTail; + copy._pendingHistoricalBoundaryRepair = _pendingHistoricalBoundaryRepair; + copy._historicalMin = _historicalMin; + copy._historicalMax = _historicalMax; + return copy; + } + + 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); + } + validated.decode(); + if (!PercentileTDigestAccumulator.canNormalizeBoundaries(validated._means, validated._weights, + validated._numCentroids, metadata.min(), metadata.max())) { + return new NonFiniteAwareTDigest(PercentileTDigestAccumulator.fromBytes(validated), 0.0, 0.0); + } + double encodedMin = metadata.min(); + double encodedMax = metadata.max(); + double compression = metadata.compression(); + int centroidCount = validated._numCentroids; + double negativeInfinityWeight = 0.0; + double positiveInfinityWeight = 0.0; + double[] finiteMeans = validated._means; + double[] finiteWeights = validated._weights; + 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 (negativeInfinityWeight > 0.0 || positiveInfinityWeight > 0.0) { + finiteDigest.deferBoundaryNormalization(); + } + if (finiteCount > 0) { + finiteMin = Double.isFinite(encodedMin) ? encodedMin : finiteMin; + finiteMax = Double.isFinite(encodedMax) ? encodedMax : finiteMax; + finiteDigest.addHistoricalCentroids(finiteMeans, finiteWeights, finiteCount, finiteMin, finiteMax, true); + } + validated.recordHistoricalFractionalBoundaries(finiteDigest); + NonFiniteAwareTDigest wrapped = + new NonFiniteAwareTDigest(finiteDigest, negativeInfinityWeight, positiveInfinityWeight); + wrapped._pendingHistoricalBoundaryRepair = metadata.weightedBoundaries(); + wrapped._historicalMin = encodedMin; + wrapped._historicalMax = encodedMax; + wrapped._unrepresentableInfinityTail = !canSplitInfinityWeight(negativeInfinityWeight) + || !canSplitInfinityWeight(positiveInfinityWeight); + if (metadata.fractionalWeights() || wrapped._unrepresentableInfinityTail) { + wrapped._originalBytes = 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); + } + PercentileTDigestAccumulator.checkTotalWeight(getTotalWeight() + weight); + double infinityWeight = 0.0; + if (value == Double.NEGATIVE_INFINITY) { + addInfinityWeight(_negativeInfinityWeight, weight); + infinityWeight = addInfinityWeight(boundaryPreparedView()._negativeInfinityWeight, weight); + } else if (value == Double.POSITIVE_INFINITY) { + addInfinityWeight(_positiveInfinityWeight, weight); + infinityWeight = addInfinityWeight(boundaryPreparedView()._positiveInfinityWeight, weight); + } + checkSplittableInfinityWeight(infinityWeight); + requireMutable(); + if (value == Double.NEGATIVE_INFINITY) { + _finiteDigest.deferBoundaryNormalization(); + _finiteDigest.recordFreshFractionalMean(value, weight); + _negativeInfinityWeight = infinityWeight; + } else if (value == Double.POSITIVE_INFINITY) { + _finiteDigest.deferBoundaryNormalization(); + _finiteDigest.recordFreshFractionalMean(value, weight); + _positiveInfinityWeight = infinityWeight; + } else { + _finiteDigest.add(value, weight); + } + mutated(); + } + + @Override + public void add(TDigest other) { + if (other.isEmpty()) { + return; + } + if (!hasValidStatistics()) { + throw PercentileTDigestAccumulator.corruptedMutation(); + } + if (!other.hasValidStatistics() && !isEmpty()) { + throw PercentileTDigestAccumulator.corruptedMutation(); + } + if (other.hasValidStatistics()) { + PercentileTDigestAccumulator.checkTotalWeight(getTotalWeight() + other.getTotalWeight()); + } + if (other instanceof NonFiniteAwareTDigest) { + NonFiniteAwareTDigest wrapped = (NonFiniteAwareTDigest) other; + if (isEmpty() && wrapped._unrepresentableInfinityTail) { + _finiteDigest = wrapped._finiteDigest.copy(); + _negativeInfinityWeight = wrapped._negativeInfinityWeight; + _positiveInfinityWeight = wrapped._positiveInfinityWeight; + _originalBytes = wrapped._originalBytes.clone(); + _unrepresentableInfinityTail = true; + _pendingHistoricalBoundaryRepair = wrapped._pendingHistoricalBoundaryRepair; + _historicalMin = wrapped._historicalMin; + _historicalMax = wrapped._historicalMax; + invalidateDerivedCaches(); + return; + } + addInfinityWeight(_negativeInfinityWeight, wrapped._negativeInfinityWeight); + addInfinityWeight(_positiveInfinityWeight, wrapped._positiveInfinityWeight); + NonFiniteAwareTDigest preparedReceiver = boundaryPreparedView(); + NonFiniteAwareTDigest preparedSource = wrapped.boundaryPreparedView(); + double negativeInfinityWeight = + addInfinityWeight(preparedReceiver._negativeInfinityWeight, preparedSource._negativeInfinityWeight); + double positiveInfinityWeight = + addInfinityWeight(preparedReceiver._positiveInfinityWeight, preparedSource._positiveInfinityWeight); + checkSplittableInfinityWeight(negativeInfinityWeight); + checkSplittableInfinityWeight(positiveInfinityWeight); + requireMutable(); + wrapped.prepareHistoricalBoundaries(); + if (negativeInfinityWeight > 0.0 || positiveInfinityWeight > 0.0) { + _finiteDigest.deferBoundaryNormalization(); + } + try { + if (!wrapped._finiteDigest.isEmpty()) { + _finiteDigest.add(wrapped._finiteDigest); + } + } catch (RuntimeException | Error e) { + restoreBoundaryNormalizationAfterFailedMerge(); + throw e; + } finally { + // Native merging can compress the source before a later validation failure. + 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())) { + requireMutable(); + _finiteDigest.add(other); + mutated(); + return; + } + // Read precise centroid weights directly; generic sources need no full digest serialization for their tails. + if (other instanceof PercentileTDigestAccumulator) { + PercentileTDigestAccumulator nativeDigest = (PercentileTDigestAccumulator) other; + if (isEmpty() && nativeDigest.hasUnrepresentableBoundaryPayload()) { + _finiteDigest = nativeDigest.copy(); + mutated(); + return; + } + } + Collection<Centroid> centroids = other.centroids(); + double[] finiteMeans = new double[Math.addExact(centroids.size(), 2)]; + double[] finiteWeights = new double[finiteMeans.length]; + int count = 0; + double sourceMin = other.getMin(); + double sourceMax = other.getMax(); + double negativeInfinityWeight = 0.0; + double positiveInfinityWeight = 0.0; + for (Centroid centroid : centroids) { + double mean = centroid.mean(); + double weight = centroid.weight(); + if (weight == 0.0) { + continue; + } + if (!(weight > 0.0) || !Double.isFinite(weight) || !(mean >= sourceMin && mean <= sourceMax)) { + throw new IllegalArgumentException("Invalid TDigest centroid"); + } Review Comment: Fixed the accumulator-source branch to use the same addInfinityWeight guard. The existing precision-loss regression now covers both signs and both orders of duplicated infinity weights 1e-20 and 1.0. These cases previously accepted a merge that absorbed the small mass; they now reject it while preserving receiver bytes and the source historical bytes. The expanded regression fails at the prior head and passes with the fix. Validated at `3821623d46933dd233b8110e123ecb4f86ac1e96`: 554 affected tests, genuine 3.2/3.3 bidirectional byte checks, the 270-case oracle, applicable precommit checks and independent Pinot review pass. New-head CI and maintainer approval remain separate gates. -- 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]
