Detect distribution drift against each series' recent windows

asap-precompute-rs/src/answers.rs

@@ -184,10 +184,151 @@ pub fn exact_answers(
     out
 }
 
+/// Thresholds for [`DriftTracker`].
+#[derive(Clone, Debug, serde::Deserialize)]
+#[serde(deny_unknown_fields)]
+pub struct DriftRules {
+    /// Windows merged into the baseline.
+    #[serde(default = "default_drift_baseline")]
+    pub baseline_windows: usize,
+    /// Kolmogorov-Smirnov distance above which a series has drifted.
+    #[serde(default = "default_ks")]
+    pub ks_threshold: f64,
+    /// Points needed in the current window before testing.
+    #[serde(default = "default_drift_points")]
+    pub min_points: u64,
+}
+
+fn default_drift_baseline() -> usize {
+    5
+}
+fn default_ks() -> f64 {
+    0.1
+}
+fn default_drift_points() -> u64 {
+    1000
+}
+
+impl Default for DriftRules {
+    fn default() -> Self {
+        Self {
+            baseline_windows: default_drift_baseline(),
+            ks_threshold: default_ks(),
+            min_points: default_drift_points(),
+        }
+    }
+}
+
+/// Largest gap between two CDFs over the points `xs` (Kolmogorov-Smirnov).
+pub fn ks_distance(current: impl Fn(f64) -> f64, baseline: impl Fn(f64) -> f64, xs: &[f64]) -> f64 {
+    xs.iter()
+        .map(|&x| (current(x) - baseline(x)).abs())
+        .fold(0.0, f64::max)
+}
+
+/// Turns each window's KS distance into a drift gauge and an event when a
+/// series starts or stops drifting. Sketch and exact use the same rules.
+#[derive(Default)]
+pub struct DriftTracker {
+    /// Thresholds.
+    pub rules: DriftRules,
+    drifting: HashMap<String, bool>,
+}
+
+impl DriftTracker {
+    /// Tracker with the given thresholds.
+    pub fn new(rules: DriftRules) -> Self {
+        Self {
+            rules,
+            drifting: HashMap::new(),
+        }
+    }
+
+    /// Records one window's distance from its baseline.
+    pub fn observe(
+        &mut self,
+        template: &SketchEnvelope,
+        method: &str,
+        ks: f64,
+        baseline_p50: f64,
+        current_p50: f64,
+    ) -> (Vec<SketchEnvelope>, Option<Event>) {
+        let drift = ks > self.rules.ks_threshold;
+        let out = vec![
+            gauge(template, "asap.drift.ks", method, &[], ks),
+            gauge(
+                template,
+                "asap.drift",
+                method,
+                &[],
+                if drift { 1.0 } else { 0.0 },
+            ),
+        ];
+        let key = format!("{}|{:?}", template.metric_name, template.labels);
+        let previous = self.drifting.insert(key, drift);
+        let series = template
+            .labels
+            .iter()
+            .map(|kv| format!("{}={}", kv.key, kv.value))
+            .collect::<Vec<_>>()
+            .join(",");
+        let event = (previous.unwrap_or(false) != drift).then(|| Event {
+            window_end_ms: template.window_end_ms,
+            method: method.into(),
+            kind: "drift".into(),
+            attributes: template
+                .labels
+                .iter()
+                .cloned()
+                .chain([
+                    KeyValue::new("drift", drift.to_string()),
+                    KeyValue::new("ks", format!("{ks:.3}")),
+                ])
+                .collect(),
+            body: if drift {
+                format!(
+                    "{} {{{series}}} drifted from its baseline: KS distance {ks:.3}, median {baseline_p50:.1} -> {current_p50:.1}",
+                    template.metric_name
+                )
+            } else {
+                format!(
+                    "{} {{{series}}} is back within its baseline: KS distance {ks:.3}",
+                    template.metric_name
+                )
+            },
+        });
+        (out, event)
+    }
+}
+
 #[cfg(test)]
 mod classifier_tests {
     use super::*;
 
+    #[test]
+    fn ks_and_drift_events() {
+        let uniform = |x: f64| x.clamp(0.0, 1.0);
+        let shifted = |x: f64| (x - 0.3).clamp(0.0, 1.0);
+        let xs: Vec<f64> = (0..=100).map(|i| i as f64 / 100.0).collect();
+        let ks = ks_distance(shifted, uniform, &xs);
+        assert!((ks - 0.3).abs() < 1e-9, "{ks}");
+        let mut t = DriftTracker::default();
+        let s = SketchEnvelope {
+            metric_name: "m".into(),
+            ..Default::default()
+        };
+        assert!(t.observe(&s, "sketch", 0.02, 1.0, 1.0).1.is_none());
+        let (out, event) = t.observe(&s, "sketch", 0.3, 1.0, 2.0);
+        assert_eq!(out[1].value, 1.0);
+        assert!(event.unwrap().body.contains("drifted"));
+        assert!(t
+            .observe(&s, "sketch", 0.03, 1.0, 1.0)
+            .1
+            .unwrap()
+            .body
+            .contains("back"));
+    }
+
     fn attr(key: &str) -> SketchEnvelope {
         SketchEnvelope {
             metric_name: "m".into(),

asap-precompute-rs/src/otap/answer_nodes.rs

@@ -54,10 +54,15 @@ use otel_arrow_dfe_pdata::OtapPayload;
 
 use super::codec::decode_pdata_to_observations;
 use crate::answers::{
-    exact_answers, sketch_answers, CardinalityClassifier, CardinalityRules, Event, ExactWindow,
+    exact_answers, ks_distance, sketch_answers, CardinalityClassifier, CardinalityRules,
+    DriftRules, DriftTracker, Event, ExactWindow,
 };
 use crate::envelope::SketchEnvelope;
+use crate::envelope::SketchType;
 use crate::observation::{KeyValue, Observation, ObservationValueKind};
+use crate::precompute::QuantileSketch;
+use crate::sketches::KLLWrapper;
+use std::collections::{HashMap, VecDeque};
 
 /// Registered URN of the sketch detector.
 pub const SKETCH_DETECTOR_URN: &str = "urn:asap:processor:sketch_detector";
@@ -90,6 +95,9 @@ pub struct DetectorConfig {
     /// Thresholds for classifying per-attribute cardinality.
     #[serde(default)]
     pub cardinality: CardinalityRules,
+    /// Thresholds for distribution drift of quantile series.
+    #[serde(default)]
+    pub drift: DriftRules,
 }
 
 /// `exact_answers` configuration.
@@ -117,6 +125,9 @@ pub struct ExactConfig {
     /// Thresholds for classifying per-attribute cardinality.
     #[serde(default)]
     pub cardinality: CardinalityRules,
+    /// Thresholds for distribution drift of quantile series.
+    #[serde(default)]
+    pub drift: DriftRules,
 }
 
 fn now_ms() -> u64 {
@@ -249,6 +260,54 @@ async fn observations(
 struct SketchDetector {
     config: DetectorConfig,
     cardinality: CardinalityClassifier,
+    drift: DriftTracker,
+    /// Recent KLL windows per series, the drift baseline.
+    baselines: HashMap<String, VecDeque<KLLWrapper>>,
+}
+
+fn series_key(env: &SketchEnvelope) -> String {
+    format!("{}|{:?}", env.metric_name, env.labels)
+}
+
+/// Percentile grid of both distributions: where their CDFs are compared.
+fn grid(a: &dyn QuantileSketch, b: &dyn QuantileSketch) -> Vec<f64> {
+    (1..100)
+        .flat_map(|i| [a.quantile(i as f64 / 100.0), b.quantile(i as f64 / 100.0)])
+        .collect()
+}
+
+impl SketchDetector {
+    /// KS distance of this KLL window from the merge of the previous ones.
+    fn drift(&mut self, env: &SketchEnvelope) -> (Vec<SketchEnvelope>, Option<Event>) {
+        let Ok(current) = KLLWrapper::from_snapshot(&env.payload) else {
+            return (Vec::new(), None);
+        };
+        let window = self.baselines.entry(series_key(env)).or_default();
+        let mut result = (Vec::new(), None);
+        if window.len() >= 3 && env.count >= self.drift.rules.min_points {
+            let mut baseline = KLLWrapper::new(current.inner().k() as i32, None);
+            for w in window.iter() {
+                let _ = crate::precompute::Sketch::merge(&mut baseline, w);
+            }
+            fn cdf(s: &KLLWrapper) -> impl Fn(f64) -> f64 + '_ {
+                let n = s.inner().count().max(1) as f64;
+                move |x| s.inner().rank(x) as f64 / n
+            }
+            let ks = ks_distance(cdf(&current), cdf(&baseline), &grid(&current, &baseline));
+            result = self.drift.observe(
+                env,
+                "sketch",
+                ks,
+                baseline.quantile(0.5),
+                current.quantile(0.5),
+            );
+        }
+        window.push_back(current);
+        if window.len() > self.drift.rules.baseline_windows {
+            window.pop_front();
+        }
+        result
+    }
 }
 
 /// The `bound="estimate"` value of `metric` among `answers`.
@@ -283,6 +342,11 @@ impl local::Processor<OtapPdata> for SketchDetector {
             };
             match sketch_answers(&env, &self.config.quantiles, self.config.k) {
                 Ok(answers) => {
+                    if env.sketch_type == SketchType::KLLSketch {
+                        let (gauges, event) = self.drift(&env);
+                        out.extend(gauges);
+                        events.extend(event);
+                    }
                     if let Some(d) = estimate_of(&answers, "asap.distinct") {
                         let (gauges, event) =
                             self.cardinality.observe(&env, "sketch", d, env.count);
@@ -307,6 +371,9 @@ struct ExactState {
 struct ExactAnswers {
     config: ExactConfig,
     cardinality: CardinalityClassifier,
+    drift: DriftTracker,
+    /// Recent windows' sorted values per quantile series.
+    baselines: HashMap<String, VecDeque<Vec<f64>>>,
     /// Open windows by window end.
     windows: BTreeMap<u64, ExactState>,
     emitted_until: u64,
@@ -403,12 +470,54 @@ impl ExactAnswers {
                     &self.config.quantiles,
                     self.config.k,
                 ));
+                if !w.values.is_empty() {
+                    // exact_answers left the values sorted.
+                    let (gauges, event) =
+                        self.exact_drift(&template, std::mem::take(&mut w.values));
+                    out.extend(gauges);
+                    events.extend(event);
+                }
             }
         }
         (out, events)
     }
 }
 
+impl ExactAnswers {
+    /// Exact KS distance of a window's sorted values from the previous windows.
+    fn exact_drift(
+        &mut self,
+        template: &SketchEnvelope,
+        sorted: Vec<f64>,
+    ) -> (Vec<SketchEnvelope>, Option<Event>) {
+        let window = self.baselines.entry(series_key(template)).or_default();
+        let mut result = (Vec::new(), None);
+        if window.len() >= 3 && sorted.len() as u64 >= self.drift.rules.min_points {
+            let mut baseline: Vec<f64> = window.iter().flatten().copied().collect();
+            baseline.sort_unstable_by(f64::total_cmp);
+            let cdf = |v: &[f64], x: f64| v.partition_point(|&y| y <= x) as f64 / v.len() as f64;
+            let at = |v: &[f64], q: f64| v[((v.len() - 1) as f64 * q).round() as usize];
+            let xs: Vec<f64> = (1..100)
+                .flat_map(|i| {
+                    [
+                        at(&sorted, i as f64 / 100.0),
+                        at(&baseline, i as f64 / 100.0),
+                    ]
+                })
+                .collect();
+            let ks = ks_distance(|x| cdf(&sorted, x), |x| cdf(&baseline, x), &xs);
+            result =
+                self.drift
+                    .observe(template, "exact", ks, at(&baseline, 0.5), at(&sorted, 0.5));
+        }
+        window.push_back(sorted);
+        if window.len() > self.drift.rules.baseline_windows {
+            window.pop_front();
+        }
+        result
+    }
+}
+
 #[async_trait(?Send)]
 impl local::Processor<OtapPdata> for ExactAnswers {
     async fn process(
@@ -464,6 +573,8 @@ static SKETCH_DETECTOR: ProcessorFactory<OtapPdata> = ProcessorFactory {
         Ok(ProcessorWrapper::local(
             SketchDetector {
                 cardinality: CardinalityClassifier::new(config.cardinality.clone()),
+                drift: DriftTracker::new(config.drift.clone()),
+                baselines: HashMap::new(),
                 config,
             },
             node,
@@ -489,6 +600,8 @@ static EXACT_ANSWERS: ProcessorFactory<OtapPdata> = ProcessorFactory {
         Ok(ProcessorWrapper::local(
             ExactAnswers {
                 cardinality: CardinalityClassifier::new(config.cardinality.clone()),
+                drift: DriftTracker::new(config.drift.clone()),
+                baselines: HashMap::new(),
                 config,
                 windows: BTreeMap::new(),
                 emitted_until: 0,
@@ -529,6 +642,8 @@ mod tests {
         .unwrap();
         let mut exact = ExactAnswers {
             cardinality: CardinalityClassifier::new(config.cardinality.clone()),
+            drift: DriftTracker::new(config.drift.clone()),
+            baselines: HashMap::new(),
             config,
             windows: BTreeMap::new(),
             emitted_until: 0,