Fix tdigest

Author: LiaCastaneda

datafusion/functions-aggregate-common/src/tdigest.rs

@@ -436,6 +436,13 @@ impl TDigest {
 
         let rank = q * self.count;
 
+        // Matches airlift's boundary: rank >= totalWeight - 1 → return max.
+        // Without this, p90 on sparse data interpolates into a value that was
+        // never observed (e.g. returns 550 instead of 1000 for a 10-value set).
+        if rank >= self.count - 1.0 {
+            return self.max();
+        }
+
         let mut pos: usize;
         let mut t;
         if q > 0.5 {
@@ -491,8 +498,28 @@ impl TDigest {
             }
         }
 
-        let value = self.centroids[pos].mean()
-            + ((rank - t) / self.centroids[pos].weight() - 0.5) * delta;
+        // Matches airlift's single-sample guard: when the rank falls between two
+        // weight-1 centroids, treat them as point masses and snap to the nearer
+        // one rather than interpolating an invented value between them.
+        let offset = (rank - t) / self.centroids[pos].weight() - 0.5;
+        if offset < 0.0 {
+            // rank is in the left half of pos — the straddling pair is (pos-1, pos)
+            let left = pos.saturating_sub(1);
+            if self.centroids[left].weight() == 1.0 && self.centroids[pos].weight() == 1.0
+            {
+                return self.centroids[pos].mean();
+            }
+        } else {
+            // rank is in the right half of pos — the straddling pair is (pos, pos+1)
+            let right = (pos + 1).min(self.centroids.len() - 1);
+            if self.centroids[pos].weight() == 1.0
+                && self.centroids[right].weight() == 1.0
+            {
+                return self.centroids[pos].mean();
+            }
+        }
+
+        let value = self.centroids[pos].mean() + offset * delta;
 
         // In `merge_digests()`: `min` is initialized to Inf, `max` is initialized to -Inf
         // and gets updated according to different `TDigest`s
@@ -735,6 +762,21 @@ mod tests {
         assert_state_roundtrip!(t);
     }
 
+    // Verify that sparse datasets produce the same percentile results as Trino.
+    // Previously DataFusion would interpolate between single-sample centroids,
+    // returning invented values like 550 (not in the dataset).
+    // Trino returns: p50=10, p90=1000, p99=1000.
+    #[test]
+    fn test_sparse_dataset_matches_trino() {
+        let values = vec![1.0, 2.0, 3.0, 4.0, 5.0, 10.0, 20.0, 50.0, 100.0, 1000.0];
+        let t = TDigest::new(100);
+        let t = t.merge_unsorted_f64(values);
+
+        assert_eq!(t.estimate_quantile(0.5), 10.0);
+        assert_eq!(t.estimate_quantile(0.9), 1000.0);
+        assert_eq!(t.estimate_quantile(0.99), 1000.0);
+    }
+
     #[test]
     fn test_size() {
         let t = TDigest::new(10);