fix: Deserialized TDigest NaN checks (prestodb#25907)

Summary: Pull Request resolved: prestodb#25907

Reviewed By: Yuhta

Differential Revision:
D81173686

Privacy Context Container: L1202870

Author: natashasehgal

presto-main-base/src/main/java/com/facebook/presto/tdigest/TDigest.java

@@ -203,6 +203,13 @@ public static TDigest createTDigest(Slice slice)
             r.mean = new double[r.activeCentroids];
             sliceInput.readBytes(wrappedDoubleArray(r.weight), r.activeCentroids * SIZE_OF_DOUBLE);
             sliceInput.readBytes(wrappedDoubleArray(r.mean), r.activeCentroids * SIZE_OF_DOUBLE);
+
+            // Validate deserialized TDigest data
+            for (int i = 0; i < r.activeCentroids; i++) {
+                checkArgument(!isNaN(r.mean[i]), "Deserialized t-digest contains NaN mean value");
+                checkArgument(r.weight[i] > 0, "weight must be > 0");
+            }
+
             sliceInput.close();
             return r;
         }
@@ -712,6 +719,12 @@ public long estimatedInMemorySizeInBytes()
 
     public Slice serialize()
     {
+        // Validate data before serialization
+        for (int i = 0; i < activeCentroids; i++) {
+            checkArgument(!isNaN(mean[i]), "Cannot serialize t-digest with NaN mean value");
+            checkArgument(weight[i] > 0, "Cannot serialize t-digest with non-positive weight");
+        }
+
         SliceOutput sliceOutput = new DynamicSliceOutput(toIntExact(estimatedSerializedSizeInBytes()));
 
         sliceOutput.writeByte(1); // version 1 of T-Digest serialization

presto-main-base/src/test/java/com/facebook/presto/tdigest/TestTDigest.java

@@ -397,6 +397,16 @@ public void testGeometricDistribution()
         }
     }
 
+    @Test(expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*Cannot serialize t-digest with NaN mean value.*")
+    public void testSerializationWithNaNMean()
+    {
+        double[] means = {1.0, Double.NaN, 3.0};
+        double[] weights = {1.0, 1.0, 1.0};
+        TDigest tDigest = createTDigest(means, weights, STANDARD_COMPRESSION_FACTOR, 1.0, 3.0, 6.0, 3);
+
+        tDigest.serialize();
+    }
+
     @Test(enabled = false)
     public void testPoissonDistribution()
     {