xiangfu0 commented on code in PR #19795: URL: https://github.com/apache/pinot/pull/19795#discussion_r4243031926
########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/SerializedTDigestInput.java: ########## @@ -0,0 +1,169 @@ +/** + * 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 org.apache.pinot.segment.local.customobject.tdigest.TDigestCodec.SerializedTDigestMetadata; +import org.apache.pinot.segment.local.customobject.tdigest.TDigestCodec.SerializedTDigestState; + +/// Reusable, invocation-local view of one serialized TDigest for group-by-MV fanout. +/// +/// [#reset(byte\[\])] validates the header once per input row. The first decode combines numerical validation +/// with filling reusable centroid arrays, so all groups for a row merge the same primitive input without sharing +/// mutable accumulator state. A pending accumulator requests an immutable byte snapshot, created at most once per +/// reset +/// and shared by the row's group fanout. The input must remain thread-confined and must not outlive the aggregation +/// call that owns it. Pending-only inputs inspect their bytes without allocating centroid arrays; their first +/// later materialization reads the already validated values without another numerical validation pass. +public final class SerializedTDigestInput { + byte[] _bytes; + byte[] _retainedBytes; + int _numCentroids; + double[] _means; + double[] _weights; + private boolean _decoded; + private boolean _metadataInspected; + + final SerializedTDigestState _metadata = new SerializedTDigestState(); + private SerializedTDigestMetadata _metadataSnapshot; + + public void reset(byte[] bytes) { + _metadata.reset(bytes); + TDigestCodec.readSerializedHeader(_metadata, true); + resetState(bytes); + _metadataSnapshot = null; + } + + /// Owns one snapshot of the encoded digest and advances the source buffer past it. + public void reset(ByteBuffer bytes) { + reset(bytes, true); + } + + /// Reads a source digest, optionally enforcing the historical decoder's declared centroid capacity. + public void reset(ByteBuffer bytes, boolean checkCapacity) { + SerializedTDigestMetadata metadata = TDigestCodec.readSerializedHeader(bytes, checkCapacity); + byte[] snapshot = new byte[metadata.encodedLength()]; + bytes.get(snapshot); + reset(snapshot, metadata); + _retainedBytes = snapshot; + } + + void reset(byte[] bytes, SerializedTDigestMetadata metadata) { + _metadata.reset(bytes); + _metadata.copyFrom(metadata); + resetState(bytes); + _metadataSnapshot = _metadataInspected ? metadata : null; + } + + private void resetState(byte[] bytes) { + _bytes = bytes; + _retainedBytes = null; + _numCentroids = _metadata._centroidCount; + _decoded = false; + _metadataInspected = !Double.isNaN(_metadata._totalWeight); + } + + public double getCompression() { + return _metadata._compression; + } + + /// Returns validated metadata without allocating centroid arrays for a digest that will remain pending. + public SerializedTDigestMetadata getMetadata() { + inspectMetadata(); + if (_metadataSnapshot == null) { + _metadataSnapshot = _metadata.snapshot(); + } + return _metadataSnapshot; + } + + void inspectMetadata() { + if (!_metadataInspected) { + TDigestCodec.inspectSerialized(_metadata, null, null); + _metadataInspected = true; + } + } + + /// Propagates endpoint provenance cached during the one numerical inspection of this serialized distribution. + /// The encoded endpoint tuple remains valid when a previous fanout decode has split boundary centroids. + boolean recordHistoricalFractionalBoundaries(PercentileTDigestAccumulator target) { + inspectMetadata(); + if (!_metadata._fractionalWeights || _metadata._needsLegacyFallback) { + return false; + } + target.requireMutable(); + double minMean = _metadata._historicalFractionalMinMean; + double maxMean = _metadata._historicalFractionalMaxMean; + target.inheritHistoricalFractionalBoundaries(minMean, maxMean); + return !Double.isNaN(minMean) || !Double.isNaN(maxMean); + } + + byte[] retainBytes() { + if (_retainedBytes == null) { + _retainedBytes = _bytes.clone(); + } + return _retainedBytes; + } + + void decode() { + if (_decoded) { + return; + } + int encodedCentroidCount = _numCentroids; + ensureCapacity(Math.addExact(encodedCentroidCount, 2)); + if (!_metadataInspected) { + TDigestCodec.inspectSerialized(_metadata, _means, _weights); + _metadataInspected = true; + } else { + // Pending state was already inspected before retaining bytes. Its later first read only materializes + // those validated centroids; regular merged inputs compute flags and weights in the decode pass above. + TDigestCodec.decodeSerializedCentroids(_metadata, _means, _weights); + } + if (!_metadata._needsLegacyFallback) { + if (_metadata._hasZeroWeightCentroids) { + int nonZeroCount = 0; + for (int i = 0; i < encodedCentroidCount; i++) { + if (_weights[i] != 0.0) { + _means[nonZeroCount] = _means[i]; + _weights[nonZeroCount++] = _weights[i]; + } + } + encodedCentroidCount = nonZeroCount; + } + if (_metadata._unorderedMeans) { + PercentileTDigestAccumulator.sortCentroids(_means, _weights, encodedCentroidCount); + } + _numCentroids = encodedCentroidCount; + } + _decoded = true; + } + + /// Repairs endpoints only when this input is being combined with another distribution. + void normalizeBoundaries() { + _numCentroids = PercentileTDigestAccumulator.normalizeSerializedBoundaries(_means, _weights, _numCentroids, Review Comment: Fixed in `006cf74a28`. `SerializedTDigestInput` now records successful boundary normalization for the current row and clears that flag on every reset. Later MV groups reuse the normalized arrays without repeating the spare-capacity check; original bytes and metadata remain intact, with no extra array allocation. The new actual `aggregateGroupByMV` regression fails on `bfaebae6b8` with the reported exception and passes after the fix. It covers two pre-populated groups, repeated weighted-boundary rows, and a third group initialized after normalization. All 546 affected tests, real 3.2/3.3 bidirectional bytes checks and the 270-case oracle pass. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/PercentileTDigestAccumulator.java: ########## @@ -0,0 +1,1722 @@ +/** + * 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. + */ +/* + * Portions of the CDF, quantile, weighted-average and K1 weight-limit implementation are adapted from + * MergingDigest, AbstractTDigest and ScaleFunction in t-digest 3.3 (https://github.com/tdunning/t-digest). + * Modified for Pinot's primitive centroid storage, lazy serialized inputs and legacy boundary handling. + * Original license notice for these portions: + * + * Licensed to Ted Dunning 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 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 byte[] _originalBoundaryBytes; + private double _historicalFractionalMinMean = Double.NaN; + private double _historicalFractionalMaxMean = Double.NaN; + private double _freshFractionalMinMean = Double.NaN; + private double _freshFractionalMaxMean = Double.NaN; + private SerializedTDigestMetadata _pendingSerializedMetadata; + private boolean _legacyDegraded; + private boolean _hasFractionalWeights; + private boolean _deferBoundaryNormalization; + + 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((double) 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) { + _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 = 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(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. + PercentileTDigestAccumulator copy() { + PercentileTDigestAccumulator copy = + new PercentileTDigestAccumulator(_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); + } + } + + public void add(double[] values, int from, int toExclusive) { + if (from >= toExclusive) { + return; + } + requireMutable(); + prepareForMutation(); + 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(); + prepareForMutation(); + 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(); + 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 PercentileTDigestAccumulator) { + addAccumulator((PercentileTDigestAccumulator) other); + return; + } + if (other instanceof NonFiniteAwareTDigest) { + PercentileTDigestAccumulator 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(); + 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) { + 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; + } + 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(PercentileTDigestAccumulator 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) { + // 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(); + other.prepareForMutation(); + prepareForMutation(); + checkTotalWeight(_totalWeight + _incomingWeight + _numRawValues + other._totalWeight); + _hasFractionalWeights |= other._hasFractionalWeights; + inheritHistoricalFractionalBoundaries(other); + 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 || _originalBoundaryBytes != null) { + 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 && _originalBoundaryBytes == null + && (_pendingSerializedMetadata == null || !_pendingSerializedMetadata.weightedBoundaries())) { + 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.getMetadata(); + 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.getMetadata(); + _legacyDegraded = true; + return false; + } + throw corruptedMutation(); + } + if (!_deferBoundaryNormalization) { + input.normalizeBoundaries(); + } + prepareForMutation(); + 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 (_useTwoLevelRawCompression) { + // Public compression and serialization both retain the legacy working pass before final compression. + flush(); + } + 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])) { + if (_centroidMeans[i] == Double.NEGATIVE_INFINITY) { + return (weightSoFar + _centroidWeights[i]) / _totalWeight; + } + if (_centroidWeights[i] == 1.0 && _centroidWeights[i + 1] == 1.0) { + return (weightSoFar + 1.0) / _totalWeight; + } + double leftExcludedWeight = _centroidWeights[i] == 1.0 ? 0.5 : 0.0; + double rightExcludedWeight = _centroidWeights[i + 1] == 1.0 ? 0.5 : 0.0; + double interpolatedWeight = (_centroidWeights[i] + _centroidWeights[i + 1]) / 2.0 + - leftExcludedWeight - rightExcludedWeight; + double base = weightSoFar + _centroidWeights[i] / 2.0 + leftExcludedWeight; + return (base + interpolatedWeight * (value - _centroidMeans[i]) + / (_centroidMeans[i + 1] - _centroidMeans[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 (_numCentroids == 1) { + return _centroidMeans[0]; + } + if (quantile == 0.0) { + return _min; + } + if (quantile == 1.0) { + return _max; + } + + 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) { + // Adjacent equal means establish a repeated-value plateau. Keep its complete mass out of interpolation at + // either edge; an isolated weighted centroid still represents an unknown distribution around its mean. + double mean = _centroidMeans[index]; + boolean pointMass = _centroidWeights[index] == 1.0 || index > 0 && _centroidMeans[index - 1] == mean Review Comment: Confirmed the same-byte behavior change against the real 3.3 reader: for the synthetic four-centroid payload, q(0.2) changes from approximately 2 to 5 and cdf(4.999) from 0.3124625 to 0.125, while serialized bytes remain unchanged. The plateau rule is intentional and already mentioned in the migration note/PR description, but the codec and weighted-boundary paragraph overstated interpolation parity. `006cf74a28` corrects both claims and explicitly documents quantile AND CDF changes in intervals adjacent to a plateau. The PR description has also been clarified; these descriptions now distinguish preserved bytes from intentionally changed interpolation. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/PercentileTDigestAccumulator.java: ########## @@ -0,0 +1,1722 @@ +/** + * 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. + */ +/* + * Portions of the CDF, quantile, weighted-average and K1 weight-limit implementation are adapted from + * MergingDigest, AbstractTDigest and ScaleFunction in t-digest 3.3 (https://github.com/tdunning/t-digest). + * Modified for Pinot's primitive centroid storage, lazy serialized inputs and legacy boundary handling. + * Original license notice for these portions: + * + * Licensed to Ted Dunning 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 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 byte[] _originalBoundaryBytes; + private double _historicalFractionalMinMean = Double.NaN; + private double _historicalFractionalMaxMean = Double.NaN; + private double _freshFractionalMinMean = Double.NaN; + private double _freshFractionalMaxMean = Double.NaN; + private SerializedTDigestMetadata _pendingSerializedMetadata; + private boolean _legacyDegraded; + private boolean _hasFractionalWeights; + private boolean _deferBoundaryNormalization; + + 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((double) 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) { + _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 = 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(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. + PercentileTDigestAccumulator copy() { + PercentileTDigestAccumulator copy = + new PercentileTDigestAccumulator(_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); + } + } + + public void add(double[] values, int from, int toExclusive) { + if (from >= toExclusive) { + return; + } + requireMutable(); + prepareForMutation(); + 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(); + prepareForMutation(); + 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(); + 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 PercentileTDigestAccumulator) { + addAccumulator((PercentileTDigestAccumulator) other); + return; + } + if (other instanceof NonFiniteAwareTDigest) { + PercentileTDigestAccumulator 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(); + 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) { + 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; + } + 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(PercentileTDigestAccumulator 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) { + // 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(); + other.prepareForMutation(); + prepareForMutation(); + checkTotalWeight(_totalWeight + _incomingWeight + _numRawValues + other._totalWeight); + _hasFractionalWeights |= other._hasFractionalWeights; + inheritHistoricalFractionalBoundaries(other); + 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 || _originalBoundaryBytes != null) { + 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 && _originalBoundaryBytes == null + && (_pendingSerializedMetadata == null || !_pendingSerializedMetadata.weightedBoundaries())) { + 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.getMetadata(); + 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.getMetadata(); + _legacyDegraded = true; + return false; + } + throw corruptedMutation(); Review Comment: Confirmed; this is the deliberate fail-fast contract previously discussed in [the earlier corruption thread](https://github.com/apache/pinot/pull/19795#discussion_r4240344248). The [operator remediation section](https://github.com/xiangfu0/pinot/blob/006cf74a280b793524709a7193c8c1fc23a9352b/compatibility-verifier/README.md#t-digest-migration) explicitly covers a single corrupt row failing a query, merge-rollup task or TDigest star-tree build; retries cannot repair the bytes, and operators must identify affected digests with `hasValidStatistics()` and rebuild/reingest from a known-good source. We retain this policy: the unknown distribution cannot produce a valid merged digest, and silently dropping its mass or inventing poisoned merged bytes would weaken the byte-compatibility guarantees. No automatic quarantine, NaN merge fallback or configuration gate is added. An isolated probe confirms solo reads retain NaN/original bytes and a rejected merge leaves healthy mass/statistics unc hanged; the existing corruption regressions are included in the passing 546-test run. ########## pinot-segment-local/src/main/java/org/apache/pinot/segment/local/customobject/tdigest/TDigestCodec.java: ########## @@ -0,0 +1,853 @@ +/** + * 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 centroids and interpolation. 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)) { Review Comment: The old query accumulator did reject non-positive compression. This exact policy difference was acknowledged in [the earlier compression thread](https://github.com/apache/pinot/pull/19795#discussion_r4241326604), where the shared finite clamp was retained intentionally and documented in the [migration release note](https://github.com/xiangfu0/pinot/blob/006cf74a280b793524709a7193c8c1fc23a9352b/compatibility-verifier/README.md#t-digest-migration). The previous general/Star-tree `TDigestUtils` factories already delegated to the legacy clamping constructor; the strict helper guard was in serialized-byte compatibility processing. Also, the standard SQL factory still rejects a negative literal before construction, so the `PERCENTILETDIGEST(..., -50)` example still fails there. We keep the previously documented consistent finite clamp, including the intentionally changed zero/Smart/reducer cases, while rejecting non-finite values. Applying the suggested guard to the shared codec would a lso reject historical headers. Existing tests pin both paths and pass in this validation; no policy reversal is made. ########## compatibility-verifier/tdigest-compatibility/run.sh: ########## @@ -0,0 +1,53 @@ +#!/usr/bin/env bash +# +# 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. +# + +# Run after LegacyTDigestCompatibilityTest. Legacy jars are test-tool artifacts, not Pinot dependencies. +set -euo pipefail +# Match the unit-test workflow's Maven Central retry and connection settings, including standalone runs. +export MAVEN_OPTS="${MAVEN_OPTS:-} -Dmaven.wagon.httpconnectionManager.ttlSeconds=25 \ +-Dmaven.wagon.http.retryHandler.count=30 -Dhttp.keepAlive=false -Dmaven.wagon.http.pool=false" +TDIGEST_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +TDIGEST_FIXTURES="$TDIGEST_ROOT/pinot-segment-local/target/tdigest-compat-fixtures" +TDIGEST_WORK="$TDIGEST_ROOT/pinot-segment-local/target/tdigest-legacy-reader" +if [[ ! -s "$TDIGEST_FIXTURES/manifest.tsv" ]]; then Review Comment: Added the requested notes beside the workflow gate and test-split placement in `006cf74a28`, including Maven Central/cache requirements for the real 3.2/3.3 test jars. The current workflow is already gated on `hashFiles(.../manifest.tsv)`, rather than `matrix.testset == 2`, so moving the segment tests to another shard moves verification with their fixture output. The comments now explicitly require keeping `LegacyTDigestCompatibilityTest` included. Only comments changed in those CI files; executable lines are identical, shell syntax passes, and the real reader/oracle harness passes. -- 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]
