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");