Jackie-Jiang commented on code in PR #19795: URL: https://github.com/apache/pinot/pull/19795#discussion_r4252796665
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigest.java: ########## @@ -0,0 +1,146 @@ +/** + * 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.Collection; +import java.util.List; + +/// Mutable percentile digest base for Pinot aggregation and legacy t-digest byte encodings. +/// Implementations require external synchronization. Historical corrupted payloads can be retained unchanged, +/// but cannot be mutated or mixed with healthy state. +public abstract class TDigest { + public abstract void add(double value); + + public void add(double value, int weight) { + add(value, (double) weight); + } + + /// Adds positive finite mass; rejects total-weight overflow. Fractional interior mass is supported, but + /// serialization rejects fresh endpoint mass below one or singleton mass between one and two: legacy readers + /// require unit endpoints. Thus `add(5, 0.5)` is accepted but cannot be serialized alone. + public abstract void add(double value, double weight); + + /// Merges an exclusively owned source; implementations may compress that mutable source in place. + public abstract void add(TDigest other); + + public void add(List<? extends TDigest> others) { + for (TDigest other : others) { + add(other); + } + } + + public abstract void compress(); + + /// Returns the legacy long size view, saturating when the precise mass exceeds its range. + public long size() { + return (long) getTotalWeight(); + } + + /// Returns mass without narrowing fractional or large weights. Add/merge reject overflow of a valid total. + /// Invalid historical payloads may report their original negative or non-finite total, but cannot be merged. + public abstract double getTotalWeight(); + + /// Returns false when historical numerical corruption leaves the distribution unknown. + public boolean hasValidStatistics() { + return true; + } + + /// An unknown distribution remains present even when its historical weights sum to zero. + public boolean isEmpty() { + return getTotalWeight() == 0.0 && hasValidStatistics(); + } + + public abstract double cdf(double value); + + public abstract double quantile(double quantile); + + /// Returns immutable centroid values with their precise double weights. + /// Reading this view may flush buffered inputs without applying the final public compression. + public abstract Collection<Centroid> centroids(); + + public abstract double compression(); + + /// Serializes one independently owned legacy-compatible payload. Implementations own their compression policy. + /// Unchanged historical encodings may be retained byte-exactly; fresh unrepresentable boundaries are rejected. + public abstract byte[] serialize(); + + /// Finalizes an independently owned result without changing a live aggregation that may receive more inputs. + protected abstract byte[] serializeSnapshot(); + + /// Bounds independently owned serialization without requiring a two-step size/write protocol. + /// Mutable implementations can override this to avoid flushing buffered inputs. + public int maxSerializedByteSize() { + return serialize().length; + } + + public abstract int centroidCount(); + + public abstract double getMin(); + + public abstract double getMax(); + + /// Returns the inherited unsupported fractional endpoint mean, or NaN. Mergers propagate this provenance + /// from validated legacy bytes without granting permission to new fractional extrema. + double getHistoricalFractionalBoundaryMean(boolean lowerBoundary) { + return Double.NaN; + } + + /// Extremal fresh fractional input means, retained across merges to prevent reusing historical permission for + /// new mass at the same endpoint. NaN means this digest has received no fresh fractional inputs. + double getFreshFractionalBoundaryMean(boolean lowerBoundary) { + return Double.NaN; + } + + /// Immutable centroid value preserving fractional and large mass. + public static final class Centroid { + private final double _mean; + private final double _weight; Review Comment: Optional cleanup: replace the `Centroid` boilerplate with a record. `Centroid` is an immutable pair with a constructor, accessors, and handwritten value equality/hashCode. `public record Centroid(double mean, double weight) {}` would preserve those semantics, including NaN and signed-zero equality, while removing roughly 35 lines. This package already uses records. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/NonFiniteAwareTDigest.java: ########## @@ -0,0 +1,749 @@ +/** + * 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 TDigestAccumulator _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(TDigestAccumulator.forLegacyAggregation(compression), 0.0, 0.0); + } + + private NonFiniteAwareTDigest(TDigestAccumulator 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(); Review Comment: Optional cleanup: share immutable retained bytes internally. Both `copy()` and the empty-receiver adoption path clone `_originalBytes` unnecessarily. Input bytes are already cloned on entry, internal mutations only clear the reference, and serialization clones outgoing bytes. Sharing the retained array internally would remove one payload allocation per copy/adoption, matching `TDigestAccumulator`'s existing approach. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigestAccumulator.java: ########## @@ -104,64 +110,201 @@ public TDigestAccumulator(int compression) { this((double) compression, true, true); } + public TDigestAccumulator(double compression) { + this(compression, true, true); + } + private TDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers) { - if (!(compression > 0.0) || !Double.isFinite(compression)) { - throw new IllegalArgumentException("TDigest compression must be positive: " + compression); - } - super.setScaleFunction(ScaleFunction.K_1); - _compression = Math.max(MIN_COMPRESSION, compression); + this(compression, allocateRawBuffer, allocateMergeBuffers, false); + } + + private TDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers, + boolean useTwoLevelRawCompression) { + _compression = TDigestCodec.normalizeCompression(compression); + _useTwoLevelRawCompression = useTwoLevelRawCompression; _workingCompression = TWO_LEVEL_COMPRESSION_MULTIPLIER * _compression; _publicMaxWeightScale = calculateK1MaxWeightScale(_compression); _workingMaxWeightScale = calculateK1MaxWeightScale(_workingCompression); long roundedCompression = (long) Math.ceil(_compression); - _centroidCapacity = getCentroidCapacity(_compression); - _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity); + _centroidCapacity = TDigestCodec.getDefaultCentroidCapacity(_compression); + _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity, useTwoLevelRawCompression); if (allocateRawBuffer) { - _rawValues = new double[getRawBufferSize(roundedCompression)]; + _rawValues = new double[getInitialRawBufferSize(roundedCompression)]; } if (allocateMergeBuffers) { - _centroidMeans = new double[_centroidCapacity]; - _centroidWeights = new double[_centroidCapacity]; - _outputMeans = new double[_centroidCapacity]; - _outputWeights = new double[_centroidCapacity]; + 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]; } } - public static TDigestAccumulator forSerializedTDigest(byte[] bytes) { - return new TDigestAccumulator(readCompression(bytes), false, false); + /// Creates a lazy digest with the legacy two-level compression used by Star-tree and general digest aggregation. + public static TDigestAccumulator forLegacyAggregation(double compression) { + return new TDigestAccumulator(compression, false, false, true); } - public static TDigestAccumulator forSerializedTDigest(SerializedTDigestInput input) { - return new TDigestAccumulator(input._compression, false, false); + /// Reads and retains one legacy digest, preserving the general aggregation compression policy. + public static TDigestAccumulator fromBytes(byte[] bytes) { + return fromBytes(ByteBuffer.wrap(bytes)); } - public static TDigestAccumulator forSerializedTDigestWithMergeBuffers(byte[] bytes) { - return new TDigestAccumulator(readCompression(bytes), false, true); + /// Reads one payload and advances the buffer past it. + public static TDigestAccumulator fromBytes(ByteBuffer bytes) { + SerializedTDigestInput input = new SerializedTDigestInput(); + input.reset(bytes); + return fromBytes(input); } - public static TDigestAccumulator forSerializedTDigest(ByteBuffer input) { - byte[] bytes = new byte[input.remaining()]; - input.get(bytes); - TDigestAccumulator accumulator = forSerializedTDigest(bytes); - accumulator.addSerializedTDigest(bytes); - return accumulator; + /// Retains an already validated input without rescanning its centroid bytes. + public static TDigestAccumulator fromBytes(SerializedTDigestInput input) { + TDigestAccumulator digest = forLegacyAggregation(input.getCompression()); + digest.addSerializedTDigest(input); + return digest; } public static TDigestAccumulator forReduction(double compression) { return new TDigestAccumulator(compression, false, false); } + /// Snapshots buffered inputs and centroid state without materializing or compressing the source. Mutable buffers + /// are copied; immutable retained bytes and metadata are shared. Merge and sort scratch buffers remain lazy. + TDigestAccumulator copy() { + TDigestAccumulator copy = + new TDigestAccumulator(_compression, false, false, _useTwoLevelRawCompression); + copy._rawValues = _rawValues != null ? _rawValues.clone() : null; + copy._centroidMeans = _numCentroids > 0 ? Arrays.copyOf(_centroidMeans, _numCentroids) : null; + copy._centroidWeights = _numCentroids > 0 ? Arrays.copyOf(_centroidWeights, _numCentroids) : null; + copy._incomingMeans = _numIncomingCentroids > 0 ? Arrays.copyOf(_incomingMeans, _numIncomingCentroids) : null; + copy._incomingWeights = _numIncomingCentroids > 0 ? Arrays.copyOf(_incomingWeights, _numIncomingCentroids) : null; + copy._pendingSerializedTDigest = _pendingSerializedTDigest; + copy._pendingSerializedMetadata = _pendingSerializedMetadata; + copy._originalFractionalBytes = _originalFractionalBytes; + copy._originalBoundaryBytes = _originalBoundaryBytes; + copy._historicalFractionalMinMean = _historicalFractionalMinMean; + copy._historicalFractionalMaxMean = _historicalFractionalMaxMean; + copy._freshFractionalMinMean = _freshFractionalMinMean; + copy._freshFractionalMaxMean = _freshFractionalMaxMean; + copy._legacyDegraded = _legacyDegraded; + copy._hasFractionalWeights = _hasFractionalWeights; + copy._deferBoundaryNormalization = _deferBoundaryNormalization; + copy._numRawValues = _numRawValues; + copy._numCentroids = _numCentroids; + copy._numIncomingCentroids = _numIncomingCentroids; + copy._mergeCount = _mergeCount; + copy._publiclyCompressed = _publiclyCompressed; + copy._incomingCentroidsSorted = _incomingCentroidsSorted; + copy._totalWeight = _totalWeight; + copy._incomingWeight = _incomingWeight; + copy._min = _min; + copy._max = _max; + return copy; + } + + /// 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) { + addCentroids(means, weights, count, min, max, alreadyCompressed, false); + } + + /// Adds validated legacy centroids without marking their existing fractional mass as fresh input. + void addHistoricalCentroids(double[] means, double[] weights, int count, double min, double max, + boolean alreadyCompressed) { + addCentroids(means, weights, count, min, max, alreadyCompressed, true); + } + + private void addCentroids(double[] means, double[] weights, int count, double min, double max, + boolean alreadyCompressed, boolean historical) { + 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; + } + prepareForMutation(); + _originalFractionalBytes = null; + _hasFractionalWeights |= fractionalWeights; + if (!historical && fractionalWeights) { + for (int i = 0; i < count; i++) { + recordFreshFractionalMean(means[i], weights[i]); + } + } + 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) { + sortCentroids(_centroidMeans, _centroidWeights, _numCentroids); + } + _totalWeight = totalWeight; + _min = min; + _max = max; + if (!historical) { + 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); Review Comment: Optional cleanup: remove duplicate extrema updates. `bufferIncomingCentroids()` now updates min/max on both its direct and buffered paths, making these outer assignments redundant. The same duplication occurs at lines 467–468. Removing both pairs would give the helper sole ownership of those updates. Keep the separate serialized-merge updates at lines 605–606, which are still needed. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigestCodec.java: ########## @@ -0,0 +1,872 @@ +/** + * 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; + +/// Reads and writes the legacy t-digest 3.2/3.3 formats, validating lengths before allocating centroid arrays. +/// Fresh output respects both readers' capacities without narrowing doubles. Historical fractional endpoints and +/// opaque corrupted payloads retain their original limitations; byte preservation does not ensure reverse upgrades. +/// Endpoint repair, capacity fallback and point-mass interpolation can change quantiles on historical data. +/// This utility has no shared mutable state; callers must exclusively own input buffers and parser state. +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 + || count == 1 && firstWeight != 1.0 && firstWeight < 2.0)) { + throw new IllegalArgumentException("Cannot serialize fractional TDigest boundary mass for legacy readers"); + } + } + + /// Allows unsupported endpoints only at exact means inherited from validated legacy bytes. A tolerance could + /// grant the exception to fresh extrema. Fresh fractional input at or beyond an endpoint permanently excludes it, + /// even if its historical mass was merged away or the new weight equals the historical weight. + static boolean hasInheritedFractionalBoundaryEncoding(TDigest source, int count, double firstMean, + double firstWeight, double lastMean, double lastWeight) { + if (count <= 0 || !source.hasValidStatistics() || !(firstWeight > 0.0) || !(lastWeight > 0.0) + || !Double.isFinite(firstWeight) || !Double.isFinite(lastWeight)) { + return false; + } + boolean singleton = count == 1 && firstWeight != 1.0 && firstWeight < 2.0; + boolean firstUnsupported = firstWeight < 1.0 || singleton; + boolean lastUnsupported = lastWeight < 1.0 || singleton; + return (firstUnsupported || lastUnsupported) + && (!firstUnsupported || firstMean == source.getHistoricalFractionalBoundaryMean(true) + && !(source.getFreshFractionalBoundaryMean(true) <= firstMean)) + && (!lastUnsupported || lastMean == source.getHistoricalFractionalBoundaryMean(false) + && !(source.getFreshFractionalBoundaryMean(false) >= lastMean)); + } + + private static void checkSerializedBoundaryWeights(byte[] bytes, SerializedTDigestMetadata metadata) { + int count = metadata.centroidCount(); + if (count > 0) { + ByteBuffer encoded = ByteBuffer.wrap(bytes); + int firstOffset = metadata.centroidOffset(); + int lastOffset = firstOffset + (count - 1) * metadata.centroidSize(); + double firstWeight = metadata.encoding() == VERBOSE_ENCODING + ? encoded.getDouble(firstOffset) : encoded.getFloat(firstOffset); + double lastWeight = metadata.encoding() == VERBOSE_ENCODING + ? encoded.getDouble(lastOffset) : encoded.getFloat(lastOffset); + checkLegacyBoundaryWeights(count, firstWeight, lastWeight); + } + } + + /// Converts a verbose `MergingDigest` payload to a representation readable by t-digest 3.2 and 3.3. + /// + /// The input must use the standard verbose `MergingDigest` encoding. The returned payload stays verbose when it + /// already fits both legacy default capacities, uses compact encoding only when every narrowed field is exact, and + /// otherwise reduces interior centroids without narrowing doubles. Compression below ten is normalized in the + /// output header because the 3.2 verbose reader sizes its arrays from the raw header before applying that minimum. + /// Opaque degraded state is retained without reducing its unknown distribution, including over-capacity payloads. + /// Fresh fractional endpoints that legacy unit-boundary repair cannot represent are rejected without reweighting. + public static byte[] makeLegacyCompatible(byte[] verboseBytes) { + return makeLegacyCompatible(verboseBytes, inspectSerialized(ByteBuffer.wrap(verboseBytes), false)); + } + + /// Uses metadata validated for these exact retained bytes, avoiding another centroid validation walk. + public static byte[] makeLegacyCompatible(byte[] verboseBytes, SerializedTDigestMetadata metadata) { + return makeLegacyCompatible(verboseBytes, metadata, false); + } + + /// Retains exact verbose weights for inherited historical fractional endpoints that cannot become unit endpoints. + /// All other compatibility checks apply. These payloads keep their historical assertion-enabled 3.3 read/merge + /// limitation; callers must establish inherited endpoint provenance before enabling this exception. + static byte[] makeLegacyCompatible(byte[] verboseBytes, SerializedTDigestMetadata metadata, + boolean inheritedFractionalBoundaries) { + if (metadata.encoding() != VERBOSE_ENCODING) { + throw new IllegalArgumentException("Expected verbose TDigest encoding"); + } + if (metadata.needsLegacyFallback()) { + return verboseBytes; + } + if (!Double.isNaN(metadata.recoveredInfinityMean())) { + double[] means = new double[metadata.centroidCount()]; + double[] weights = new double[means.length]; + decodeSerializedCentroids(ByteBuffer.wrap(verboseBytes), metadata, means, weights); + if (!inheritedFractionalBoundaries) { + return serializeRecoveredCentroids(metadata, means, weights); + } + verboseBytes = serializeCentroids(metadata.compression(), metadata.min(), metadata.max(), means, weights, + means.length); + metadata = inspectSerialized(ByteBuffer.wrap(verboseBytes), false); + } + if (metadata.hasZeroWeightCentroids()) { + verboseBytes = removeZeroWeightCentroids(verboseBytes, metadata); + metadata = inspectSerialized(ByteBuffer.wrap(verboseBytes), false); + } + if (!inheritedFractionalBoundaries) { + checkSerializedBoundaryWeights(verboseBytes, metadata); + } + double compression = metadata.compression(); + if (metadata.encodedCompression() != compression) { + verboseBytes = verboseBytes.clone(); + ByteBuffer.wrap(verboseBytes).putDouble(20, compression); + } + int centroidCount = metadata.centroidCount(); + int legacyCapacity = + Math.min(getLegacyDefaultCentroidCapacity(compression), getDefaultCentroidCapacity(compression)); + if (centroidCount <= legacyCapacity) { + return verboseBytes; + } + int mainCapacity = getDefaultCentroidCapacity(compression); + long bufferCapacity = DEFAULT_MERGE_BUFFER_MULTIPLIER * (long) mainCapacity; + if (centroidCount > mainCapacity || inheritedFractionalBoundaries || (double) (float) compression != compression + || centroidCount > Short.MAX_VALUE + || mainCapacity > Short.MAX_VALUE || bufferCapacity > Short.MAX_VALUE) { + return reduceVerboseCentroids(verboseBytes, legacyCapacity); + } + ByteBuffer verbose = ByteBuffer.wrap(verboseBytes); + ByteBuffer centroids = verbose.duplicate(); + centroids.position(metadata.centroidOffset()); + for (int i = 0; i < centroidCount; i++) { + double weight = centroids.getDouble(); + double mean = centroids.getDouble(); + if ((double) (float) weight != weight || (double) (float) mean != mean) { + return reduceVerboseCentroids(verboseBytes, legacyCapacity); + } + } + + ByteBuffer small = ByteBuffer.allocate(SMALL_HEADER_SIZE + SMALL_CENTROID_SIZE * centroidCount); + verbose.rewind(); + verbose.getInt(); + small.putInt(SMALL_ENCODING); + small.putDouble(verbose.getDouble()); + small.putDouble(verbose.getDouble()); + verbose.getDouble(); + verbose.getInt(); + small.putFloat((float) compression); + small.putShort((short) mainCapacity); + small.putShort((short) bufferCapacity); + small.putShort((short) centroidCount); + for (int i = 0; i < centroidCount; i++) { + small.putFloat((float) verbose.getDouble()); + small.putFloat((float) verbose.getDouble()); + } + return small.array(); + } + + private static byte[] serializeRecoveredCentroids(SerializedTDigestMetadata metadata, double[] means, + double[] weights) { + TDigestAccumulator digest = TDigestAccumulator.forLegacyAggregation(metadata.compression()); + digest.addHistoricalCentroids(means, weights, means.length, metadata.min(), metadata.max(), true); + digest.inheritHistoricalFractionalBoundaries(metadata.historicalFractionalMinMean(), + metadata.historicalFractionalMaxMean()); + return digest.serialize(); + } + + private static byte[] removeZeroWeightCentroids(byte[] verboseBytes, SerializedTDigestMetadata metadata) { + ByteBuffer source = ByteBuffer.wrap(verboseBytes); + source.position(metadata.centroidOffset()); + ByteBuffer filtered = ByteBuffer.allocate(metadata.encodedLength()); + filtered.put(verboseBytes, 0, VERBOSE_HEADER_SIZE); + int count = 0; + for (int i = 0; i < metadata.centroidCount(); i++) { + double weight = source.getDouble(); + double mean = source.getDouble(); + if (weight != 0.0) { + filtered.putDouble(weight).putDouble(mean); + count++; + } + } + filtered.putInt(28, count); + return Arrays.copyOf(filtered.array(), filtered.position()); + } + + /// Validates serialized lengths and centroid values without allocating centroid arrays. + public static double validateSerialized(byte[] bytes) { + return validateSerialized(ByteBuffer.wrap(bytes)); + } + + /// Validates one legacy digest without changing `input` and returns its total centroid weight. + /// + /// Trailing bytes remain available to the caller for intermediate values containing several objects. + public static double validateSerialized(ByteBuffer input) { + return validateSerialized(input, true); + } + + private static double validateSerialized(ByteBuffer input, boolean checkCapacity) { + return inspectSerialized(input, checkCapacity).totalWeight(); + } + + /// Header and numerical metadata for one legacy byte payload. The centroid bytes remain owned by the caller. + /// Header-only inspection leaves total weight as NaN until centroid inspection validates the numerical fields. + /// + /// Legacy writers can emit poisoned numerical state after overflow or broken fractional boundary repair. + /// Such state is retained with degraded quantiles, rather than being passed to the sorted K1 kernel. + public record SerializedTDigestMetadata(int encoding, double min, double max, double compression, + int centroidCount, int centroidOffset, int centroidSize, + int encodedLength, double totalWeight, boolean hasNonFiniteMeans, boolean needsLegacyFallback, + boolean unorderedMeans, boolean fractionalWeights, boolean weightedBoundaries, boolean hasZeroWeightCentroids, + double encodedCompression, double recoveredInfinityMean, double historicalFractionalMinMean, + double historicalFractionalMaxMean) { + /// Whether the encoded centroid view can be used unchanged by healthy pending-state readers and writers. + public boolean isCanonicalPayload() { + return !needsLegacyFallback && !weightedBoundaries && !unorderedMeans && !hasZeroWeightCentroids + && Double.isNaN(recoveredInfinityMean); + } + } + + /// Invocation-local parser state. Primitive byte-array reads avoid per-row buffer and metadata allocations. + /// Public metadata and retained pending payloads receive immutable snapshots instead of sharing this state. + static final class SerializedTDigestState { + private static final VarHandle SHORT_VIEW = + MethodHandles.byteArrayViewVarHandle(short[].class, ByteOrder.BIG_ENDIAN); + private static final VarHandle INT_VIEW = MethodHandles.byteArrayViewVarHandle(int[].class, ByteOrder.BIG_ENDIAN); + private static final VarHandle LONG_VIEW = MethodHandles.byteArrayViewVarHandle(long[].class, ByteOrder.BIG_ENDIAN); + + int _encoding; + double _min; + double _max; + double _compression; + int _centroidCount; + int _centroidOffset; + int _centroidSize; + int _encodedLength; + double _totalWeight; + boolean _hasNonFiniteMeans; + boolean _needsLegacyFallback; + boolean _unorderedMeans; + boolean _fractionalWeights; + boolean _weightedBoundaries; + boolean _hasZeroWeightCentroids; + double _encodedCompression; + double _recoveredInfinityMean; + double _historicalFractionalMinMean; + double _historicalFractionalMaxMean; + private byte[] _bytes; + private ByteBuffer _buffer; + private int _position; + private int _limit; + + void reset(byte[] bytes) { + _bytes = bytes; + _buffer = null; + _position = 0; + _limit = bytes.length; + } + + void reset(ByteBuffer input) { + _bytes = null; + _buffer = input.slice().order(ByteOrder.BIG_ENDIAN); + _position = 0; + _limit = _buffer.limit(); + } + + void copyFrom(SerializedTDigestMetadata metadata) { + _encoding = metadata.encoding(); + _min = metadata.min(); + _max = metadata.max(); + _compression = metadata.compression(); + _centroidCount = metadata.centroidCount(); + _centroidOffset = metadata.centroidOffset(); + _centroidSize = metadata.centroidSize(); + _encodedLength = metadata.encodedLength(); + _totalWeight = metadata.totalWeight(); + _hasNonFiniteMeans = metadata.hasNonFiniteMeans(); + _needsLegacyFallback = metadata.needsLegacyFallback(); + _unorderedMeans = metadata.unorderedMeans(); + _fractionalWeights = metadata.fractionalWeights(); + _weightedBoundaries = metadata.weightedBoundaries(); + _hasZeroWeightCentroids = metadata.hasZeroWeightCentroids(); + _encodedCompression = metadata.encodedCompression(); + _recoveredInfinityMean = metadata.recoveredInfinityMean(); + _historicalFractionalMinMean = metadata.historicalFractionalMinMean(); + _historicalFractionalMaxMean = metadata.historicalFractionalMaxMean(); + } + + SerializedTDigestMetadata snapshot() { + return new SerializedTDigestMetadata(_encoding, _min, _max, _compression, _centroidCount, _centroidOffset, + _centroidSize, _encodedLength, _totalWeight, _hasNonFiniteMeans, _needsLegacyFallback, _unorderedMeans, + _fractionalWeights, _weightedBoundaries, _hasZeroWeightCentroids, _encodedCompression, _recoveredInfinityMean, + _historicalFractionalMinMean, _historicalFractionalMaxMean); + } + + double weight(int index) { + _position = _centroidOffset + index * _centroidSize; + return _centroidSize == VERBOSE_CENTROID_SIZE ? readDouble() : readFloat(); + } + + double mean(int index) { + _position = _centroidOffset + index * _centroidSize + _centroidSize / 2; + double mean = _centroidSize == VERBOSE_CENTROID_SIZE ? readDouble() : readFloat(); + return Double.isNaN(mean) ? _recoveredInfinityMean : Math.max(_min, Math.min(mean, _max)); + } + + private int readPosition(int length) { + if (length > _limit - _position) { + throw new BufferUnderflowException(); + } + int position = _position; + _position += length; + return position; + } + + private int readInt() { + int position = readPosition(Integer.BYTES); + return _bytes != null ? (int) INT_VIEW.get(_bytes, position) : _buffer.getInt(position); + } + + private short readShort() { + int position = readPosition(Short.BYTES); + return _bytes != null ? (short) SHORT_VIEW.get(_bytes, position) : _buffer.getShort(position); + } + + private double readDouble() { + int position = readPosition(Double.BYTES); + return _bytes != null + ? Double.longBitsToDouble((long) LONG_VIEW.get(_bytes, position)) + : _buffer.getDouble(position); + } + + private float readFloat() { + int position = readPosition(Float.BYTES); + return _bytes != null ? Float.intBitsToFloat((int) INT_VIEW.get(_bytes, position)) : _buffer.getFloat(position); + } + } + + /// Reads a finalized payload's quantile without cloning bytes or constructing a mutable digest. + static double quantile(byte[] bytes, double quantile) { + if (Double.isNaN(quantile) || quantile < 0.0 || quantile > 1.0) { + throw new IllegalArgumentException("q should be in [0,1], got " + quantile); + } + SerializedTDigestState state = new SerializedTDigestState(); + state.reset(bytes); + readSerializedHeader(state, true); + if (Double.isNaN(state._totalWeight)) { + inspectSerialized(state, null, null); + } + if (state._needsLegacyFallback || state._totalWeight == 0.0) { + return Double.NaN; + } + if (state._unorderedMeans || state._hasZeroWeightCentroids) { + // Rare retained historical payloads need the reader's filtering and stable sort, but no recompression. + SerializedTDigestInput input = new SerializedTDigestInput(); + input.reset(bytes, state.snapshot()); + input.decode(); + return TDigestAccumulator.quantile(quantile, input._numCentroids, state._totalWeight, state._min, state._max, + input._means, input._weights, null); + } + return TDigestAccumulator.quantile(quantile, state._centroidCount, state._totalWeight, state._min, state._max, + null, null, state); + } + + /// Inspects one payload without mutating `input`, scanning the centroid values once. + public static SerializedTDigestMetadata inspectSerialized(ByteBuffer input) { + return inspectSerialized(input, true); + } + + /// Inspects a payload, optionally enforcing the legacy decoder's declared centroid capacity. + public static SerializedTDigestMetadata inspectSerialized(ByteBuffer input, boolean checkCapacity) { + SerializedTDigestState state = new SerializedTDigestState(); + state.reset(input); + readSerializedHeader(state, checkCapacity); + if (Double.isNaN(state._totalWeight)) { + inspectSerialized(state, null, null); + } + return state.snapshot(); + } + + /// Reads the fixed-size header without allocating centroid arrays; oversized historical state is also inspected. + public static SerializedTDigestMetadata readSerializedHeader(ByteBuffer input) { + return readSerializedHeader(input, true); + } + + /// Reads a fixed-size header, optionally enforcing the declared centroid capacity. + /// + /// NaN or inverted nonempty extrema mark legacy degraded state. Oversized payloads require numerical inspection: + /// historical degraded state remains readable even if an old decoder's arrays could not hold its centroids. + public static SerializedTDigestMetadata readSerializedHeader(ByteBuffer input, boolean checkCapacity) { + SerializedTDigestState state = new SerializedTDigestState(); + state.reset(input); + readSerializedHeader(state, checkCapacity); + return state.snapshot(); + } + + static void readSerializedHeader(SerializedTDigestState state, boolean checkCapacity) { + int encoding = state.readInt(); + if (encoding != VERBOSE_ENCODING && encoding != SMALL_ENCODING) { + throw new IllegalStateException("Invalid format for serialized histogram"); + } + double min = state.readDouble(); + double max = state.readDouble(); + double compression; + double encodedCompression; + int mainCapacity; + int bufferCapacity = 0; + int centroidCount; + int centroidSize; + if (encoding == VERBOSE_ENCODING) { + encodedCompression = state.readDouble(); + compression = normalizeCompression(encodedCompression); + mainCapacity = Math.max(getDefaultCentroidCapacity(compression), getLegacyDefaultCentroidCapacity(compression)); + centroidCount = state.readInt(); + centroidSize = VERBOSE_CENTROID_SIZE; + } else { + encodedCompression = state.readFloat(); + compression = normalizeCompression(encodedCompression); + mainCapacity = state.readShort(); + bufferCapacity = state.readShort(); + checkSmallArrayCapacity(mainCapacity); + checkSmallArrayCapacity(bufferCapacity); + if (mainCapacity == -1) { + mainCapacity = getDefaultCentroidCapacity(compression); + } + centroidCount = state.readShort(); + centroidSize = SMALL_CENTROID_SIZE; + } + if (centroidCount < 0) { + throw new IllegalArgumentException("Invalid negative TDigest centroid count: " + centroidCount); + } + if (centroidCount > (state._limit - state._position) / centroidSize) { + throw new BufferUnderflowException(); + } + int centroidOffset = state._position; + int encodedLength = centroidOffset + centroidCount * centroidSize; + state._encoding = encoding; + state._min = min; + state._max = max; + state._compression = compression; + state._centroidCount = centroidCount; + state._centroidOffset = centroidOffset; + state._centroidSize = centroidSize; + state._encodedLength = encodedLength; + state._totalWeight = Double.NaN; + state._hasNonFiniteMeans = false; + state._needsLegacyFallback = Double.isNaN(min) || Double.isNaN(max) || centroidCount > 0 && min > max; + state._unorderedMeans = false; + state._fractionalWeights = false; + state._weightedBoundaries = false; + state._hasZeroWeightCentroids = false; + state._encodedCompression = encodedCompression; + state._recoveredInfinityMean = Double.NaN; + state._historicalFractionalMinMean = Double.NaN; + state._historicalFractionalMaxMean = Double.NaN; + if (checkCapacity && centroidCount > mainCapacity) { + inspectSerialized(state, null, null); + if (!state._needsLegacyFallback) { + throw new IllegalArgumentException("TDigest centroid count exceeds capacity: " + centroidCount); + } + } + } + + /// Validates centroids in one pass and optionally decodes their means and weights into reusable arrays. + /// + /// The header must describe this exact input, and each supplied array must hold all encoded centroids. Endpoint + /// rounding inherited from compact re-encoding or old merge arithmetic is restored to the double header bounds. + public static SerializedTDigestMetadata inspectSerialized(ByteBuffer input, SerializedTDigestMetadata header, + double[] means, double[] weights) { + SerializedTDigestState state = new SerializedTDigestState(); + state.reset(input); + state.copyFrom(header); + inspectSerialized(state, means, weights); + return state.snapshot(); Review Comment: Optional cleanup: reduce the exposed internal API. This four-argument `inspectSerialized` overload has no production or test callers; the reusable parser uses the state overload. Delete it rather than retain an unnecessary public buffer/header contract. Now that these classes share a sub-package, `TDigestAccumulator`'s `getCentroidCountUpperBound()`, `hasOriginalFractionalPayload()`, and `copyCentroids()` can also become package-private: all current callers are in that package. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigestAccumulator.java: ########## @@ -186,243 +327,299 @@ public void add(double value) { if (Double.isNaN(value)) { throw new IllegalArgumentException("Cannot add NaN to t-digest"); } - materializePendingSerializedTDigest(); + requireMutable(); + prepareForMutation(); ensureRawBuffer(); - if (_numRawValues == _rawValues.length - || _numRawValues + _numIncomingCentroids == getPendingInputLimit()) { - flush(); - } + flushIfBuffersFull(true); + _originalFractionalBytes = null; _rawValues[_numRawValues++] = value; } @Override - public void add(double value, int weight) { + 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"); } - materializePendingSerializedTDigest(); + if (!(weight > 0.0) || !Double.isFinite(weight)) { + throw new IllegalArgumentException("TDigest weight must be positive: " + weight); + } + requireMutable(); + prepareForMutation(); bufferIncomingCentroid(value, weight); + _originalFractionalBytes = null; + _hasFractionalWeights |= weight != Math.rint(weight); + recordFreshFractionalMean(value, 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 TDigestAccumulator) { addAccumulator((TDigestAccumulator) other); return; } - if (other instanceof MergingDigest && other.size() > Integer.MAX_VALUE) { - // Centroid.count() exposes an int even though MergingDigest stores double weights internally. Preserve large - // pre-aggregated weights by ingesting the serialized double-precision representation instead. - addSerializedTDigest(TDigestUtils.serialize(other)); + if (other instanceof NonFiniteAwareTDigest) { + TDigestAccumulator finiteDigest = ((NonFiniteAwareTDigest) other).finiteDigestIfNoTails(); + if (finiteDigest != null) { + addAccumulator(finiteDigest); + 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(); - materializePendingSerializedTDigest(); + boolean wrapped = other instanceof NonFiniteAwareTDigest; + int capacity = wrapped ? Math.addExact(centroids.size(), 2) : centroids.size(); + double[] means = new double[capacity]; + double[] weights = new double[capacity]; + int count = 0; for (Centroid centroid : centroids) { - bufferIncomingCentroid(centroid.mean(), centroid.count()); + 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; } - _min = Math.min(_min, other.getMin()); - _max = Math.max(_max, other.getMax()); - } - - @Override - public void add(List<? extends TDigest> others) { - for (TDigest other : others) { - add(other); + if (wrapped) { + count = normalizeSerializedBoundaries(means, weights, count, other.getMin(), other.getMax()); } + // 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, wrapped); + inheritHistoricalFractionalBoundaries(other); } private void addAccumulator(TDigestAccumulator other) { + if (other == this) { + compress(); + prepareForMutation(); + if (_numCentroids == 0) { + return; + } + // Snapshot only self-merges; fresh fractional mass need not be representable by a legacy wire payload. + addHistoricalCentroids(Arrays.copyOf(_centroidMeans, _numCentroids), + Arrays.copyOf(_centroidWeights, _numCentroids), _numCentroids, _min, _max, false); + return; + } + if (_deferBoundaryNormalization && !other._deferBoundaryNormalization) { + other.prepareForMutation(); + } + if (other._originalBoundaryBytes != null) { + if (getTotalWeight() != 0.0) { + throw unrepresentableBoundaryMutation(); + } + addSerializedTDigest(other._originalBoundaryBytes); + 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) { - addSerializedTDigest(other._pendingSerializedTDigest); + // 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; } Review Comment: Optional cleanup: remove the unreachable degraded-source branch. Every transition to `_legacyDegraded` retains `_pendingSerializedTDigest`, and materialization explicitly skips degraded state. Consequently, the preceding pending-input branch always handles degraded sources and returns before this branch. Removing it would simplify the state handling without removing legacy recovery behavior. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigestAccumulator.java: ########## @@ -104,64 +110,201 @@ public TDigestAccumulator(int compression) { this((double) compression, true, true); } + public TDigestAccumulator(double compression) { + this(compression, true, true); + } + private TDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers) { - if (!(compression > 0.0) || !Double.isFinite(compression)) { - throw new IllegalArgumentException("TDigest compression must be positive: " + compression); - } - super.setScaleFunction(ScaleFunction.K_1); - _compression = Math.max(MIN_COMPRESSION, compression); + this(compression, allocateRawBuffer, allocateMergeBuffers, false); + } + + private TDigestAccumulator(double compression, boolean allocateRawBuffer, boolean allocateMergeBuffers, + boolean useTwoLevelRawCompression) { + _compression = TDigestCodec.normalizeCompression(compression); + _useTwoLevelRawCompression = useTwoLevelRawCompression; _workingCompression = TWO_LEVEL_COMPRESSION_MULTIPLIER * _compression; _publicMaxWeightScale = calculateK1MaxWeightScale(_compression); _workingMaxWeightScale = calculateK1MaxWeightScale(_workingCompression); long roundedCompression = (long) Math.ceil(_compression); - _centroidCapacity = getCentroidCapacity(_compression); - _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity); + _centroidCapacity = TDigestCodec.getDefaultCentroidCapacity(_compression); + _pendingCentroidCapacity = getPendingCentroidCapacity(_centroidCapacity, useTwoLevelRawCompression); if (allocateRawBuffer) { - _rawValues = new double[getRawBufferSize(roundedCompression)]; + _rawValues = new double[getInitialRawBufferSize(roundedCompression)]; } if (allocateMergeBuffers) { - _centroidMeans = new double[_centroidCapacity]; - _centroidWeights = new double[_centroidCapacity]; - _outputMeans = new double[_centroidCapacity]; - _outputWeights = new double[_centroidCapacity]; + 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]; } } - public static TDigestAccumulator forSerializedTDigest(byte[] bytes) { - return new TDigestAccumulator(readCompression(bytes), false, false); + /// Creates a lazy digest with the legacy two-level compression used by Star-tree and general digest aggregation. + public static TDigestAccumulator forLegacyAggregation(double compression) { + return new TDigestAccumulator(compression, false, false, true); } - public static TDigestAccumulator forSerializedTDigest(SerializedTDigestInput input) { - return new TDigestAccumulator(input._compression, false, false); + /// Reads and retains one legacy digest, preserving the general aggregation compression policy. + public static TDigestAccumulator fromBytes(byte[] bytes) { + return fromBytes(ByteBuffer.wrap(bytes)); } - public static TDigestAccumulator forSerializedTDigestWithMergeBuffers(byte[] bytes) { - return new TDigestAccumulator(readCompression(bytes), false, true); + /// Reads one payload and advances the buffer past it. + public static TDigestAccumulator fromBytes(ByteBuffer bytes) { + SerializedTDigestInput input = new SerializedTDigestInput(); + input.reset(bytes); + return fromBytes(input); } - public static TDigestAccumulator forSerializedTDigest(ByteBuffer input) { - byte[] bytes = new byte[input.remaining()]; - input.get(bytes); - TDigestAccumulator accumulator = forSerializedTDigest(bytes); - accumulator.addSerializedTDigest(bytes); - return accumulator; + /// Retains an already validated input without rescanning its centroid bytes. + public static TDigestAccumulator fromBytes(SerializedTDigestInput input) { + TDigestAccumulator digest = forLegacyAggregation(input.getCompression()); + digest.addSerializedTDigest(input); + return digest; } public static TDigestAccumulator forReduction(double compression) { return new TDigestAccumulator(compression, false, false); } + /// Snapshots buffered inputs and centroid state without materializing or compressing the source. Mutable buffers + /// are copied; immutable retained bytes and metadata are shared. Merge and sort scratch buffers remain lazy. + TDigestAccumulator copy() { + TDigestAccumulator copy = + new TDigestAccumulator(_compression, false, false, _useTwoLevelRawCompression); + copy._rawValues = _rawValues != null ? _rawValues.clone() : null; Review Comment: Optional allocation improvement: copy only active raw values for serialization snapshots. `serializeSnapshot()` calls `copy()`, which clones the full raw buffer, including unused capacity. A serialization-specific copy could retain only `_numRawValues`, while sharing the rest of the copy implementation. An isolated comparison saved approximately 1.2–2 KB per RAW result in three buffered cases; all eight checked output hashes matched. Preserve general `copy()` capacity because it controls future flush timing. This demonstrates allocation savings, not an end-to-end latency improvement. -- 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]
