fix(executor): DDSketch interpolated quantiles and top-k membership on sketchlib planner-pin/0.4.0-plus

Cargo.toml

@@ -23,4 +23,8 @@ resolver = "2"
 arrow = { version = "58.4", default-features = false }
 # The hasher Arrow already uses, for the executor's hot interning maps.
 ahash = { version = "0.8", default-features = false, features = ["std", "runtime-rng"] }
-asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "eb04e7f243c1d5e5c346cd19842f2402d796b620" }
+# `planner-pin/0.4.0-plus`: release 0.4.0 (eb04e7f) plus the open sketchlib
+# PRs #245 (checked f64 updates, heap rescoring), #246 (KLL lazy-buffer
+# follow-ups), #247 (Count Sketch stack medians) and #248
+# (`DdSketch::quantile_interpolated`). Moves to a release once they merge.
+asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "63e565506426020c505ef0b84e0c91f450d09e7c" }

crates/executor/src/summary_kernels/dd_sketch.rs

@@ -45,11 +45,16 @@ impl AggregateCore for DDSketchAccumulator {
     /// Quantiles, and the total sample count as a bare `PointCount`.
     fn estimate(&self, query: &SketchStatistic) -> Result<f64, KernelError> {
         match query {
-            SketchStatistic::Quantile { q } if (0.0..=1.0).contains(q) => self
+            // PromQL `quantile_over_time` and SQL `percentile_cont`
+            // interpolate between the order statistics at ranks ⌊q (n − 1)⌋
+            // and ⌈q (n − 1)⌉. The Planner never sends a q outside [0, 1]
+            // (Pass 1 prunes the sketch candidate,
+            // `QuantileOutsideUnitInterval`; the exact plan answers ±Inf),
+            // so sketchlib's error for one, NaN included, is passed on.
+            SketchStatistic::Quantile { q } => self
                 .inner
-                .get_value_at_quantile(*q)
+                .quantile_interpolated(*q)?
                 .ok_or_else(|| "DDSketch quantile of an empty population".into()),
-            SketchStatistic::Quantile { .. } => Err("quantile must be in [0, 1]".into()),
             SketchStatistic::PointCount { value: None, .. } => Ok(self.inner.get_count() as f64),
             other => Err(format!("DDSketch does not answer {other:?}").into()),
         }
@@ -95,18 +100,45 @@ mod tests {
         assert!(dd.estimate(&SketchStatistic::Cardinality).is_err());
     }
 
-    // The quantile at q reads the value of rank ⌈q·n⌉; q = 0 and q = 1 read
-    // the exact minimum and maximum.
+    // The quantile at q interpolates linearly between the order statistics
+    // at ranks ⌊q (n − 1)⌋ and ⌈q (n − 1)⌉, as PromQL `quantile_over_time`
+    // and SQL `percentile_cont` do: [20, 40] at q = 0.9 is 38 (nearest rank
+    // read 40), within α of the interpolated value; q = 0 and q = 1 read the
+    // exact minimum and maximum.
     #[test]
-    fn quantile_reads_rank_ceil_q_n() {
-        let mut dd = DDSketchAccumulator::new(0.01);
-        for v in [1.0, 2.0, 3.0, 4.0] {
+    fn quantile_interpolates_between_adjacent_order_statistics() {
+        let alpha = 0.01;
+        let mut dd = DDSketchAccumulator::new(alpha);
+        for v in [20.0, 40.0] {
             dd.inner.add(&v);
         }
         let at = |q| dd.estimate(&SketchStatistic::Quantile { q }).unwrap();
-        assert_eq!(at(0.0), 1.0);
-        assert!((at(0.5) - 2.0).abs() <= 0.01 * 2.0, "{}", at(0.5));
-        assert_eq!(at(0.9), 4.0);
-        assert_eq!(at(1.0), 4.0);
+        assert!((at(0.9) - 38.0).abs() <= alpha * 38.0, "{}", at(0.9));
+        assert_eq!(at(0.0), 20.0);
+        assert_eq!(at(1.0), 40.0);
+        let mut four = DDSketchAccumulator::new(alpha);
+        for v in [1.0, 2.0, 3.0, 4.0] {
+            four.inner.add(&v);
+        }
+        // Rank 0.5 · 3 = 1.5: halfway between 2 and 3.
+        let median = four
+            .estimate(&SketchStatistic::Quantile { q: 0.5 })
+            .unwrap();
+        assert!((median - 2.5).abs() <= alpha * 2.5, "{median}");
+    }
+
+    // The Planner never sends a φ outside [0, 1] (Pass 1 prunes the sketch
+    // candidate, `QuantileOutsideUnitInterval`, and the exact plan answers
+    // ±Inf), so the kernel refuses one; NaN likewise.
+    #[test]
+    fn quantile_outside_the_unit_interval_is_an_error() {
+        let mut dd = DDSketchAccumulator::new(0.01);
+        dd.inner.add(&1.0);
+        for q in [-0.1, 1.1, f64::NAN] {
+            assert!(
+                dd.estimate(&SketchStatistic::Quantile { q }).is_err(),
+                "{q}"
+            );
+        }
     }
 }

crates/executor/src/summary_kernels/factory.rs

@@ -809,17 +809,19 @@ mod ddsketch_domain {
     }
 
     /// Quantiles over 0, 0, 3, 5 and over signed samples are within alpha of
-    /// the sample at rank ⌈q·n⌉.
+    /// PromQL's interpolated quantile, at rank q (n − 1) between adjacent
+    /// order statistics: here each pair of neighbours shares a sign or
+    /// includes zero.
     #[test]
     fn quantiles_over_zero_and_signed_samples() {
         for (samples, expected) in [
             (
                 vec![0.0, 0.0, 3.0, 5.0],
-                [(0.25, 0.0), (0.5, 0.0), (0.75, 3.0), (1.0, 5.0)],
+                [(0.25, 0.0), (0.5, 1.5), (0.75, 3.5), (1.0, 5.0)],
             ),
             (
                 vec![-7.0, -3.0, 0.0, 4.0],
-                [(0.25, -7.0), (0.5, -3.0), (0.75, 0.0), (1.0, 4.0)],
+                [(0.25, -4.0), (0.5, -1.5), (0.75, 1.0), (1.0, 4.0)],
             ),
         ] {
             let mut updater = DDSketchAccumulatorUpdater::new(0.01);

crates/executor/src/summary_kernels/weighted_frequency.rs

@@ -83,17 +83,13 @@ pub struct WeightedFrequency {
     width: usize,
     depth: usize,
     capacity: usize,
-    /// The sum of every weight's magnitude, which bounds every counter's:
-    /// while it is finite no counter can overflow.
-    total_magnitude: f64,
 }
 
 impl std::fmt::Debug for WeightedFrequency {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("WeightedFrequency")
             .field("algorithm", &self.algorithm())
             .field("shape", &self.shape())
-            .field("total_magnitude", &self.total_magnitude)
             .finish()
     }
 }
@@ -176,7 +172,6 @@ impl WeightedFrequency {
             width,
             depth,
             capacity,
-            total_magnitude: 0.0,
         })
     }
 
@@ -193,12 +188,12 @@ impl WeightedFrequency {
 
     /// The weight rule: finite, and nonnegative for Count-Min.
     pub(crate) fn check_weight(&self, weight: f64) -> Result<(), String> {
-        if !weight.is_finite() || (self.algorithm() == FrequencyAlgorithm::Cms && weight < 0.0) {
-            return Err(
-                "weighted frequency requires finite weights; CMS additionally requires nonnegative weights".into(),
-            );
+        match self.algorithm() {
+            FrequencyAlgorithm::Cms => CountMinSketch::<Matrix, RegularPath>::check_weight(weight),
+            FrequencyAlgorithm::CountSketch => {
+                CountSketch::<Matrix, RegularPath>::check_weight(weight)
+            }
         }
-        Ok(())
     }
 
     /// Empties the state in place, keeping its algorithm and shape.
@@ -209,32 +204,48 @@ impl WeightedFrequency {
             Sketch::CmsHeap(s) => s.clear(),
             Sketch::CountSketchHeap(s) => s.clear(),
         }
-        self.total_magnitude = 0.0;
     }
 
     /// Adds `weight` to the identity tuple `values`. An error, and the state
     /// unchanged, for a weight [`Self::check_weight`] refuses, an identity
-    /// that is not Null, Bool, Int64, finite Float64 or Utf8, or a total
-    /// weight past `f64::MAX`.
+    /// that is not Null, Bool, Int64, finite Float64 or Utf8, or an update
+    /// that would carry a counter past `f64::MAX`.
     pub fn update(&mut self, values: &[Value], weight: f64) -> Result<(), Error> {
         let key = identity_key(values)?;
         self.update_key(&key, weight)
     }
 
+    /// Adds `weight` to `key`. A heap state keeps the old kernel's
+    /// membership: every resident is rescored after the counters change and
+    /// before `key` is offered, so the heap admits `key` against the
+    /// residents' current estimates, not the counts they had when last
+    /// touched.
     fn update_key(&mut self, key: &[u8], weight: f64) -> Result<(), Error> {
         self.check_weight(weight).map_err(Error::Operator)?;
-        let total = self.total_magnitude + weight.abs();
-        if !total.is_finite() {
-            return Err(Error::Operator("weighted frequency sum overflow".into()));
-        }
         let key = DataInput::Bytes(key);
+        // The weight passed `check_weight`, so sketchlib refuses only a
+        // counter carried past `f64::MAX`, leaving the sketch unchanged.
+        let overflow = |e: String| Error::Operator(format!("weighted frequency sum overflow: {e}"));
         match &mut self.sketch {
-            Sketch::Cms(s) => s.update_weighted(&key, weight),
-            Sketch::CountSketch(s) => s.update_weighted(&key, weight),
-            Sketch::CmsHeap(s) => s.update_weighted(&key, weight),
-            Sketch::CountSketchHeap(s) => s.update_weighted(&key, weight),
+            Sketch::Cms(s) => s.try_update_weighted(&key, weight).map_err(overflow)?,
+            Sketch::CountSketch(s) => s.try_update_weighted(&key, weight).map_err(overflow)?,
+            Sketch::CmsHeap(s) => {
+                s.count_min_sketch_mut()
+                    .try_update_weighted(&key, weight)
+                    .map_err(overflow)?;
+                s.rescore_heap();
+                let estimate = s.estimate(&key);
+                s.heap_mut().update(&key, estimate);
+            }
+            Sketch::CountSketchHeap(s) => {
+                s.count_sketch_mut()
+                    .try_update_weighted(&key, weight)
+                    .map_err(overflow)?;
+                s.rescore_heap();
+                let estimate = s.estimate(&key);
+                s.heap_mut().update(&key, estimate);
+            }
         }
-        self.total_magnitude = total;
         Ok(())
     }
 
@@ -306,10 +317,6 @@ impl WeightedFrequency {
         if self.algorithm() != other.algorithm() || self.shape() != other.shape() {
             return Err(Error::Invalid("weighted frequency shape mismatch".into()));
         }
-        let total = self.total_magnitude + other.total_magnitude;
-        if !total.is_finite() {
-            return Err(Error::Operator("weighted frequency merge overflow".into()));
-        }
         let mut merged = self.clone();
         match (&mut merged.sketch, &other.sketch) {
             (Sketch::Cms(s), Sketch::Cms(o)) => s.merge(o),
@@ -318,7 +325,11 @@ impl WeightedFrequency {
             (Sketch::CountSketchHeap(s), Sketch::CountSketchHeap(o)) => s.merge(o),
             _ => unreachable!("one algorithm and capacity is one kind of sketch"),
         }
-        merged.total_magnitude = total;
+        // Both inputs are left unchanged, so a merge past `f64::MAX` is
+        // refused by inspecting the copy.
+        if merged.counters().iter().any(|c| !c.is_finite()) {
+            return Err(Error::Operator("weighted frequency merge overflow".into()));
+        }
         Ok(merged)
     }
 }
@@ -360,6 +371,63 @@ impl AggregateCore for WeightedFrequency {
 mod tests {
     use super::*;
 
+    // Top-k membership equals the old kernel's, which rescored every
+    // candidate after every update and kept the `capacity` best: on small,
+    // colliding sketches, an update also moves residents it did not touch,
+    // and after every update the retained keys are the reference model's
+    // (up to exact ties).
+    #[test]
+    fn membership_matches_rescoring_every_candidate_after_every_update() {
+        for algorithm in [FrequencyAlgorithm::Cms, FrequencyAlgorithm::CountSketch] {
+            let (width, depth, capacity) = (8, 3, 4);
+            let mut state = WeightedFrequency::new(algorithm, width, depth, capacity).unwrap();
+            // The reference: candidate keys, rescored from the same counters.
+            let mut reference: Vec<i64> = Vec::new();
+            let mut seed = 0x2545_f491_4f6c_dd1du64;
+            for step in 0..400 {
+                seed ^= seed << 13;
+                seed ^= seed >> 7;
+                seed ^= seed << 17;
+                let key = (seed % 24) as i64;
+                // Distinct fractional weights, so no two estimates tie.
+                let weight = 1.0 + (seed >> 40) as f64 / (1u64 << 24) as f64;
+                let weight = if algorithm == FrequencyAlgorithm::CountSketch && step % 3 == 0 {
+                    -weight
+                } else {
+                    weight
+                };
+                state.update(&[Value::Int64(key)], weight).unwrap();
+                if !reference.contains(&key) {
+                    reference.push(key);
+                }
+                let score = |k: i64| state.point_estimate(&[Value::Int64(k)]).unwrap();
+                reference.sort_by(|a, b| score(*b).total_cmp(&score(*a)));
+                reference.truncate(capacity);
+                // Keys that share every cell tie exactly, and the two
+                // models break ties differently, so membership is compared by
+                // the retained keys' current estimates.
+                let scores = |keys: &[i64]| {
+                    let mut scores: Vec<f64> = keys.iter().map(|&k| score(k)).collect();
+                    scores.sort_by(f64::total_cmp);
+                    scores
+                };
+                let got: Vec<i64> = state
+                    .rows(capacity)
+                    .into_iter()
+                    .map(|row| match row[0] {
+                        Value::Int64(k) => k,
+                        _ => unreachable!(),
+                    })
+                    .collect();
+                assert_eq!(
+                    scores(&got),
+                    scores(&reference),
+                    "{algorithm:?} after update {step}: {got:?} vs {reference:?}"
+                );
+            }
+        }
+    }
+
     // The charge is sketchlib's count of the sketch plus this adapter's own
     // fields, each struct counted once.
     #[test]
@@ -405,18 +473,21 @@ mod tests {
         assert!(matches!(rows[2][..], [Value::Int64(1), Value::Float64(score)] if score == -10.5));
     }
 
-    // Weights outside the rule, non-finite or unsupported identities and a
-    // total weight past f64::MAX are errors that leave the state unchanged.
+    // Weights outside the rule, non-finite or unsupported identities and an
+    // update carrying a counter past f64::MAX are errors that leave the state
+    // unchanged.
     #[test]
     fn invalid_updates_leave_the_state_unchanged() {
         let mut cms = WeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 3, 4).unwrap();
         cms.update(&[Value::Int64(1)], f64::MAX).unwrap();
         let snapshot =
             |s: &WeightedFrequency| (s.counters().to_vec(), format!("{:?}", s.rows(usize::MAX)));
         let before = snapshot(&cms);
-        for weight in [-1.0, f64::INFINITY, f64::NAN, f64::MAX] {
+        for weight in [-1.0, f64::INFINITY, f64::NAN] {
             assert!(cms.update(&[Value::Null], weight).is_err());
         }
+        // A counter carried past f64::MAX: the same key's cells again.
+        assert!(cms.update(&[Value::Int64(1)], f64::MAX).is_err());
         for identity in [
             Value::Float64(f64::NAN),
             Value::Float64(f64::INFINITY),

crates/integration-tests/tests/promql_numeric_regressions.rs

@@ -108,17 +108,7 @@ fn sum_rate_and_increase_have_real_exact_accumulator_nodes() {
 /// The PromQL quantile over a DDSketch: the linear interpolation at rank
 /// `q (n - 1)` between the sketch's two neighbouring order statistics.
 fn interpolated_quantile(sketch: &asap_sketchlib::DdSketch, q: f64) -> f64 {
-    let n = sketch.get_count();
-    // The `j`-th smallest (from 0): the sketch's value at rank `ceil(p n) = j + 1`.
-    let order_statistic = |j: u64| {
-        sketch
-            .get_value_at_quantile((j as f64 + 0.5) / n as f64)
-            .unwrap()
-    };
-    let rank = q * (n - 1) as f64;
-    let (lower, upper) = (rank.floor() as u64, rank.ceil() as u64);
-    let fraction = rank - lower as f64;
-    order_statistic(lower) * (1. - fraction) + order_statistic(upper) * fraction
+    sketch.quantile_interpolated(q).unwrap().unwrap()
 }
 
 /// A finite, nonzero approximate denominator does not certify a ratio bound.

crates/integration-tests/tests/promql_to_post_asap.rs

@@ -839,18 +839,8 @@ fn ddsketch_ratio_bound_holds_for_signed_pinned_sketch_evaluations() {
                 sketch.try_add(&v).unwrap();
             }
             // PromQL's quantile: interpolated between neighbouring order
-            // statistics, the `j`-th smallest read at rank `j + 1`.
-            let n = sketch.get_count();
-            let interpolated = |q: f64| {
-                let order_statistic = |j: u64| {
-                    sketch
-                        .get_value_at_quantile((j as f64 + 0.5) / n as f64)
-                        .unwrap()
-                };
-                let r = q * (n - 1) as f64;
-                order_statistic(r.floor() as u64) * (1. - r.fract())
-                    + order_statistic(r.ceil() as u64) * r.fract()
-            };
+            // statistics.
+            let interpolated = |q: f64| sketch.quantile_interpolated(q).unwrap().unwrap();
             let want = exact(0.9) / exact(0.5);
             let got = interpolated(0.9) / interpolated(0.5);
             assert!((got - want).abs() / want.abs() <= bound + 1e-12);