From f37adf5c905fb51258c92843a5c6c71983be769e Mon Sep 17 00:00:00 2001 From: PerumalsamyR Date: Tue, 14 Jul 2026 23:58:18 -0700 Subject: [PATCH] Prevent NaN centroid means in TDigestDouble Merging centroids used the incremental mean update mean + (value - mean) * w / weight, where the intermediate difference or product can overflow to infinity for finite values of large magnitude, eventually turning stored centroid means into NaN. Fall back to an overflow-safe weighted average when that happens, and normalize weights in weightedAverage so quantile interpolation cannot overflow either. Also ignore infinite values in update(), consistent with NaN handling, and validate deserialized state (finite min/max, centroid means and buffered values, positive centroid weights), since serialized input cannot be assumed to be well-formed. Addresses #702 --- .../datasketches/tdigest/TDigestDouble.java | 56 ++++++++++-- .../tdigest/TDigestDoubleTest.java | 85 +++++++++++++++++++ 2 files changed, 136 insertions(+), 5 deletions(-) diff --git a/src/main/java/org/apache/datasketches/tdigest/TDigestDouble.java b/src/main/java/org/apache/datasketches/tdigest/TDigestDouble.java index 0f059937d..09ad6dba0 100644 --- a/src/main/java/org/apache/datasketches/tdigest/TDigestDouble.java +++ b/src/main/java/org/apache/datasketches/tdigest/TDigestDouble.java @@ -95,11 +95,12 @@ public short getK() { } /** - * Update this TDigest with the given value + * Update this TDigest with the given value. + * NaN and infinity are ignored. * @param value to update the TDigest with */ public void update(final double value) { - if (Double.isNaN(value)) { return; } + if (!Double.isFinite(value)) { return; } if (numBuffered_ == (centroidsCapacity_ * BUFFER_MULTIPLIER)) { compress(); } bufferValues_[numBuffered_] = value; numBuffered_++; @@ -423,6 +424,7 @@ public static TDigestDouble heapify(final MemorySegment seg, final boolean isFlo } else { value = posSeg.getDouble(); } + checkDeserializedValue(value, "value"); return new TDigestDouble(reverseMerge, k, value, value, new double[] {value}, new long[] {1}, 1, null); } final int numCentroids = posSeg.getInt(); @@ -436,17 +438,22 @@ public static TDigestDouble heapify(final MemorySegment seg, final boolean isFlo min = posSeg.getDouble(); max = posSeg.getDouble(); } + checkDeserializedValue(min, "min"); + checkDeserializedValue(max, "max"); final double[] means = new double[numCentroids]; final long[] weights = new long[numCentroids]; long totalWeight = 0; for (int i = 0; i < numCentroids; i++) { means[i] = isFloat ? posSeg.getFloat() : posSeg.getDouble(); weights[i] = isFloat ? posSeg.getInt() : posSeg.getLong(); + checkDeserializedValue(means[i], "centroid mean"); + checkDeserializedWeight(weights[i]); totalWeight += weights[i]; } final double[] buffered = new double[numBuffered]; for (int i = 0; i < numBuffered; i++) { buffered[i] = isFloat ? posSeg.getFloat() : posSeg.getDouble(); + checkDeserializedValue(buffered[i], "buffered value"); } return new TDigestDouble(reverseMerge, k, min, max, means, weights, totalWeight, buffered); } @@ -464,12 +471,16 @@ private static TDigestDouble heapifyCompat(final MemorySegment seg) { final double max = seg.get(JAVA_DOUBLE_UNALIGNED_BIG_ENDIAN, offset); offset += Double.BYTES; final short k = (short) seg.get(JAVA_DOUBLE_UNALIGNED_BIG_ENDIAN, offset); offset += Double.BYTES; final int numCentroids = seg.get(JAVA_INT_UNALIGNED_BIG_ENDIAN, offset); offset += Integer.BYTES; + checkDeserializedValue(min, "min"); + checkDeserializedValue(max, "max"); final double[] means = new double[numCentroids]; final long[] weights = new long[numCentroids]; long totalWeight = 0; for (int i = 0; i < numCentroids; i++) { weights[i] = (long) seg.get(JAVA_DOUBLE_UNALIGNED_BIG_ENDIAN, offset); offset += Double.BYTES; means[i] = seg.get(JAVA_DOUBLE_UNALIGNED_BIG_ENDIAN, offset); offset += Double.BYTES; + checkDeserializedValue(means[i], "centroid mean"); + checkDeserializedWeight(weights[i]); totalWeight += weights[i]; } return new TDigestDouble(false, k, min, max, means, weights, totalWeight, null); @@ -482,17 +493,33 @@ private static TDigestDouble heapifyCompat(final MemorySegment seg) { // they can be derived from k in the constructor seg.get(JAVA_INT_UNALIGNED_BIG_ENDIAN, offset); offset += Integer.BYTES; final int numCentroids = seg.get(JAVA_SHORT_UNALIGNED_BIG_ENDIAN, offset); offset += Short.BYTES; + checkDeserializedValue(min, "min"); + checkDeserializedValue(max, "max"); final double[] means = new double[numCentroids]; final long[] weights = new long[numCentroids]; long totalWeight = 0; for (int i = 0; i < numCentroids; i++) { weights[i] = (long) seg.get(JAVA_FLOAT_UNALIGNED_BIG_ENDIAN, offset); offset += Float.BYTES; means[i] = seg.get(JAVA_FLOAT_UNALIGNED_BIG_ENDIAN, offset); offset += Float.BYTES; + checkDeserializedValue(means[i], "centroid mean"); + checkDeserializedWeight(weights[i]); totalWeight += weights[i]; } return new TDigestDouble(false, k, min, max, means, weights, totalWeight, null); } + private static void checkDeserializedValue(final double value, final String description) { + if (!Double.isFinite(value)) { + throw new SketchesArgumentException("Deserialized " + description + " must be finite, actual: " + value); + } + } + + private static void checkDeserializedWeight(final long weight) { + if (weight < 1) { + throw new SketchesArgumentException("Deserialized centroid weight must be positive, actual: " + weight); + } + } + /** * Human-readable summary of this TDigest as a string * @return summary of this TDigest @@ -594,8 +621,8 @@ private void merge(final double[] values, final long[] weights, final long weigh } if (addThis) { // merge into existing centroid centroidWeights_[numCentroids_ - 1] += weights[current]; - centroidMeans_[numCentroids_ - 1] += ((values[current] - centroidMeans_[numCentroids_ - 1]) - * weights[current]) / centroidWeights_[numCentroids_ - 1]; + centroidMeans_[numCentroids_ - 1] = mergedMean(centroidMeans_[numCentroids_ - 1], + values[current], weights[current], centroidWeights_[numCentroids_ - 1]); } else { // copy to a new centroid weightSoFar += centroidWeights_[numCentroids_ - 1]; centroidMeans_[numCentroids_] = values[current]; @@ -641,7 +668,26 @@ static double z(final double compression, final double n) { } } + /* + * The weights are normalized before multiplying so that each term is bounded by the magnitude + * of its input, otherwise the products can overflow to infinity for values of large magnitude. + */ private static double weightedAverage(final double x1, final double w1, final double x2, final double w2) { - return ((x1 * w1) + (x2 * w2)) / (w1 + w2); + final double weight = w1 + w2; + return (x1 * (w1 / weight)) + (x2 * (w2 / weight)); + } + + /* + * Computes the mean of a centroid after merging in the given value with weight w, + * where weight is the total weight of the centroid including w. + * The intermediate (value - mean) or its product with w can overflow to infinity + * even when both inputs are finite (e.g. means near opposite ends of the double range), + * which eventually turns the stored mean into NaN. In that case fall back to + * the overflow-safe weighted average, which stays finite. + */ + private static double mergedMean(final double mean, final double value, final long w, final long weight) { + final double newMean = mean + (((value - mean) * w) / weight); + if (Double.isFinite(newMean)) { return newMean; } + return weightedAverage(mean, weight - w, value, w); } } diff --git a/src/test/java/org/apache/datasketches/tdigest/TDigestDoubleTest.java b/src/test/java/org/apache/datasketches/tdigest/TDigestDoubleTest.java index c940648be..18e4103cf 100644 --- a/src/test/java/org/apache/datasketches/tdigest/TDigestDoubleTest.java +++ b/src/test/java/org/apache/datasketches/tdigest/TDigestDoubleTest.java @@ -26,7 +26,9 @@ import static org.testng.Assert.assertTrue; import java.lang.foreign.MemorySegment; +import java.lang.foreign.ValueLayout; +import org.apache.datasketches.common.SketchesArgumentException; import org.apache.datasketches.common.SketchesStateException; import org.apache.datasketches.common.TestUtil; import org.testng.annotations.Test; @@ -151,6 +153,89 @@ public void serializeDeserializeNonEmpty() { assertEquals(td2.getQuantile(0.5), td1.getQuantile(0.5)); } + @Test + public void updateIgnoresNaNAndInfinity() { + final TDigestDouble td = new TDigestDouble(); + td.update(Double.NaN); + td.update(Double.POSITIVE_INFINITY); + td.update(Double.NEGATIVE_INFINITY); + assertTrue(td.isEmpty()); + td.update(1); + td.update(Double.POSITIVE_INFINITY); + td.update(Double.NEGATIVE_INFINITY); + assertEquals(td.getTotalWeight(), 1); + assertEquals(td.getMinValue(), 1.0); + assertEquals(td.getMaxValue(), 1.0); + } + + // issue #702 + @Test + public void extremeValuesDoNotProduceNaN() { + final TDigestDouble td = new TDigestDouble(); + final int n = 10000; + for (int i = 0; i < n; i++) { + td.update(((i & 1) == 0) ? Double.MAX_VALUE : -Double.MAX_VALUE); + } + final byte[] bytes = td.toByteArray(); // compresses as a side effect + assertEquals(td.getTotalWeight(), n); + assertTrue(Double.isFinite(td.getQuantile(0.25))); + assertTrue(Double.isFinite(td.getQuantile(0.5))); + assertTrue(Double.isFinite(td.getQuantile(0.75))); + // all serialized centroid means must be finite, otherwise heapify throws + final TDigestDouble td2 = TDigestDouble.heapify(MemorySegment.ofArray(bytes)); + assertEquals(td2.getTotalWeight(), n); + assertEquals(td2.getMinValue(), -Double.MAX_VALUE); + assertEquals(td2.getMaxValue(), Double.MAX_VALUE); + } + + // serialized layout: preamble 16 bytes, min 8 bytes, max 8 bytes, + // then (mean 8 bytes, weight 8 bytes) per centroid + private static byte[] serializeNonEmpty() { + final TDigestDouble td = new TDigestDouble(); + for (int i = 0; i < 1000; i++) { + td.update(i); + } + return td.toByteArray(); + } + + @Test + public void deserializeNaNCentroidMean() { + final byte[] bytes = serializeNonEmpty(); + MemorySegment.ofArray(bytes).set(ValueLayout.JAVA_DOUBLE_UNALIGNED, 32, Double.NaN); + assertThrows(SketchesArgumentException.class, () -> TDigestDouble.heapify(MemorySegment.ofArray(bytes))); + } + + @Test + public void deserializeInfiniteCentroidMean() { + final byte[] bytes = serializeNonEmpty(); + MemorySegment.ofArray(bytes).set(ValueLayout.JAVA_DOUBLE_UNALIGNED, 32, Double.NEGATIVE_INFINITY); + assertThrows(SketchesArgumentException.class, () -> TDigestDouble.heapify(MemorySegment.ofArray(bytes))); + } + + @Test + public void deserializeZeroCentroidWeight() { + final byte[] bytes = serializeNonEmpty(); + MemorySegment.ofArray(bytes).set(ValueLayout.JAVA_LONG_UNALIGNED, 40, 0L); + assertThrows(SketchesArgumentException.class, () -> TDigestDouble.heapify(MemorySegment.ofArray(bytes))); + } + + @Test + public void deserializeNaNMinValue() { + final byte[] bytes = serializeNonEmpty(); + MemorySegment.ofArray(bytes).set(ValueLayout.JAVA_DOUBLE_UNALIGNED, 16, Double.NaN); + assertThrows(SketchesArgumentException.class, () -> TDigestDouble.heapify(MemorySegment.ofArray(bytes))); + } + + @Test + public void deserializeNaNSingleValue() { + final TDigestDouble td = new TDigestDouble(); + td.update(1); + final byte[] bytes = td.toByteArray(); + // single-value layout: preamble 8 bytes, then the value + MemorySegment.ofArray(bytes).set(ValueLayout.JAVA_DOUBLE_UNALIGNED, 8, Double.NaN); + assertThrows(SketchesArgumentException.class, () -> TDigestDouble.heapify(MemorySegment.ofArray(bytes))); + } + @Test public void deserializeFromReferenceImplementationDouble() { final byte[] bytes = TestUtil.getFileBytes(resPath, "tdigest_ref_k100_n10000_double.sk");