feat(asap-physical-operators): KMV distinct-count kernel (#697)

crates/asap-physical-operators/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2021"
 chrono = { version = "0.4.39", default-features = false, features = ["std"] }
 futures = "0.3"
 planner-types = { package = "asap-types", path = "../types" }
-asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "e6b03142628a3109b6482ebe8512bff7e75f5b35" }
+asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "e6b03142628a3109b6482ebe8512bff7e75f5b35", features = ["experimental"] }
 serde = { version = "1", features = ["derive", "rc"] }
 serde_json = "1"
 tracing = "0.1"

crates/asap-physical-operators/src/capability.rs

@@ -60,6 +60,7 @@ pub fn validate_summary_kernel(
                     false
                 }
                 (A::Hll, P::Hll { precision }) if (4..=18).contains(precision) => false,
+                (A::Kmv, P::Kmv { k }) if *k >= crate::summary_kernels::kmv::KMV_MIN_K => false,
                 (A::Cms, P::Cms { width, depth })
                 | (A::CountSketch, P::CountSketch { width, depth })
                     if valid_matrix(*width, *depth) =>
@@ -264,7 +265,7 @@ pub fn validate_native_family(family: &SummaryFamilyType) -> Result<(), Error> {
     match family {
         SummaryFamilyType::ExactAggregate(..) => {}
         SummaryFamilyType::Sketch(kind, _)
-            if matches!(kind.algorithm(), A::Kll | A::DDSketch | A::Hll) => {}
+            if matches!(kind.algorithm(), A::Kll | A::DDSketch | A::Hll | A::Kmv) => {}
         _ => {
             return Err(Error::Invalid(
                 "summary family has no native DAG state implementation".into(),
@@ -319,6 +320,8 @@ pub fn validate_sketch_evaluation(
                 }
                 true
             }
+            (A::Kmv, SketchStatistic::Cardinality) => true,
+            (A::Kmv, _) => bare_count,
             // Only count intents read a Count-Min bare count, and their
             // updates have unit weight; the evaluation is typed Int64 on that basis.
             (A::Cms, _) => bare_count,

crates/asap-physical-operators/src/summary_kernels/factory.rs

@@ -606,6 +606,9 @@ pub fn create_planner_accumulator(
                 u32::from(*precision),
             ),
         }),
+        (SketchAlgorithm::Kmv, SketchParams::Kmv { k }) => {
+            Box::new(crate::summary_kernels::kmv::KmvAccumulator::new(*k))
+        }
         (
             SketchAlgorithm::UnivMon,
             SketchParams::UnivMon {

crates/asap-physical-operators/src/summary_kernels/kmv.rs

@@ -0,0 +1,198 @@
+//! KMV (k minimum values) distinct-count summary over `asap_sketchlib::KMV`.
+//!
+//! Every item is hashed with the library's shared `CANONICAL_HASH_SEED`, and
+//! a merge keeps the `k` smallest distinct hashes of both states, so merging
+//! two states equals building one over the union of their inputs.
+use crate::summary_kernels::factory::AccumulatorUpdater;
+use crate::{AggregateCore, KernelError, KeyByLabelValues};
+use asap_sketchlib::{DataInput, KMV};
+use planner_types::post_asap::SketchStatistic;
+
+/// The smallest `k` the estimate `(k - 1) / U(k)` is defined for.
+pub const KMV_MIN_K: u32 = 2;
+
+#[derive(Debug, Clone)]
+pub struct KmvAccumulator {
+    pub inner: KMV,
+}
+
+impl KmvAccumulator {
+    pub fn new(k: u32) -> Self {
+        Self {
+            inner: KMV::new(k as usize),
+        }
+    }
+
+    /// Observe one sample value. Like HLL, a value is hashed by its bits
+    /// with `-0.0` folded into `0.0`; NaN is not an item.
+    pub fn insert_sample(&mut self, value: f64) {
+        if !value.is_nan() {
+            let bits = if value == 0.0 { 0 } else { value.to_bits() };
+            self.inner.insert(&DataInput::Bytes(&bits.to_le_bytes()));
+        }
+    }
+}
+
+impl AggregateCore for KmvAccumulator {
+    fn clone_boxed_core(&self) -> Box<dyn AggregateCore> {
+        Box::new(self.clone())
+    }
+
+    fn as_any(&self) -> &dyn std::any::Any {
+        self
+    }
+    fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
+        self
+    }
+
+    fn merge_with(&self, other: &dyn AggregateCore) -> Result<Box<dyn AggregateCore>, KernelError> {
+        let other = other
+            .as_any()
+            .downcast_ref::<Self>()
+            .ok_or("KMV merges only with KMV")?;
+        if self.inner.k != other.inner.k {
+            return Err(format!(
+                "KMV merges only with the same k ({} vs {})",
+                self.inner.k, other.inner.k
+            )
+            .into());
+        }
+        let mut merged = self.clone();
+        merged.inner.merge(&mut other.inner.clone());
+        Ok(Box::new(merged))
+    }
+
+    /// Distinct count. A bare `PointCount` over a KMV also reads the distinct count.
+    fn estimate(&self, query: &SketchStatistic) -> Result<f64, KernelError> {
+        match query {
+            // `KMV::estimate` takes `&mut self` without mutating.
+            SketchStatistic::Cardinality | SketchStatistic::PointCount { value: None, .. } => {
+                Ok(self.inner.clone().estimate())
+            }
+            other => Err(format!("KMV does not answer {other:?}").into()),
+        }
+    }
+
+    fn approx_memory_bytes(&self) -> usize {
+        std::mem::size_of::<Self>()
+            .saturating_add(self.inner.k_vals.len() * std::mem::size_of::<u64>())
+    }
+}
+
+/// A KMV state is its own updater: it observes unkeyed sample values.
+impl AccumulatorUpdater for KmvAccumulator {
+    fn is_keyed(&self) -> bool {
+        false
+    }
+    fn memory_usage_bytes(&self) -> usize {
+        self.approx_memory_bytes()
+    }
+    fn update_single(&mut self, value: f64, _: i64) {
+        self.insert_sample(value);
+    }
+    fn update_keyed(&mut self, _: &KeyByLabelValues, value: f64, _: i64) {
+        self.insert_sample(value);
+    }
+    fn take_accumulator(&mut self) -> Box<dyn AggregateCore> {
+        let result = Box::new(self.clone());
+        self.reset();
+        result
+    }
+    fn snapshot_accumulator(&self) -> Box<dyn AggregateCore> {
+        Box::new(self.clone())
+    }
+    fn into_accumulator(self: Box<Self>) -> Box<dyn AggregateCore> {
+        self
+    }
+    fn reset(&mut self) {
+        self.inner = KMV::new(self.inner.k);
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn build(k: u32, values: impl IntoIterator<Item = u32>) -> KmvAccumulator {
+        let mut kmv = KmvAccumulator::new(k);
+        for v in values {
+            kmv.insert_sample(f64::from(v));
+        }
+        kmv
+    }
+
+    fn cardinality(state: &dyn AggregateCore) -> f64 {
+        state.estimate(&SketchStatistic::Cardinality).unwrap()
+    }
+
+    // Below k distinct items the k smallest hashes are all the items: exact.
+    #[test]
+    fn counts_fewer_than_k_distinct_items_exactly() {
+        let kmv = build(1024, (0..500).chain(0..500));
+        assert_eq!(cardinality(&kmv), 500.0);
+        assert_eq!(
+            kmv.estimate(&SketchStatistic::PointCount {
+                key: planner_types::pre_asap::ColumnRef::SampleValue,
+                value: None,
+            })
+            .unwrap(),
+            500.0
+        );
+    }
+
+    // Above k the estimate is within KMV's relative standard error
+    // 1/sqrt(k - 2) times ten (the 99% Chebyshev bound Stage 1 sizes for).
+    #[test]
+    fn estimates_many_distinct_items_within_its_error() {
+        let k = 1024;
+        let bound = 10.0 / f64::from(k - 2).sqrt();
+        for exact in [10_000u32, 100_000] {
+            let estimate = cardinality(&build(k, (0..exact).chain(0..exact / 2)));
+            let error = (estimate - f64::from(exact)).abs() / f64::from(exact);
+            assert!(error < bound, "{exact}: {estimate} (error {error})");
+        }
+    }
+
+    // Merging two overlapping states keeps the same k smallest distinct
+    // hashes as building one over the union, so both estimate alike.
+    #[test]
+    fn merge_equals_the_union_build() {
+        let k = 256;
+        let (a, b) = (build(k, 0..6_000), build(k, 4_000..10_000));
+        let union = build(k, 0..10_000);
+        let merged = a.merge_with(&b).unwrap();
+        let merged = merged.as_any().downcast_ref::<KmvAccumulator>().unwrap();
+        let sorted = |s: &KmvAccumulator| {
+            let mut hashes: Vec<u64> = s.inner.k_vals.iter().copied().collect();
+            hashes.sort_unstable();
+            hashes
+        };
+        assert_eq!(sorted(merged), sorted(&union));
+        assert_eq!(cardinality(merged), cardinality(&union));
+        // The inputs are unchanged.
+        assert_eq!(sorted(&a), sorted(&build(k, 0..6_000)));
+    }
+
+    #[test]
+    fn rejects_merge_with_a_different_k() {
+        assert!(build(16, 0..10).merge_with(&build(32, 0..10)).is_err());
+    }
+
+    #[test]
+    fn rejects_quantile() {
+        assert!(KmvAccumulator::new(16)
+            .estimate(&SketchStatistic::Quantile { q: 0.5 })
+            .is_err());
+    }
+
+    #[test]
+    fn negative_zero_is_zero_and_nan_is_not_an_item() {
+        let mut kmv = KmvAccumulator::new(16);
+        for v in [0.0, -0.0, f64::NAN, 1.0] {
+            kmv.update_single(v, 0);
+        }
+        assert_eq!(cardinality(&kmv), 2.0);
+        assert_eq!(cardinality(kmv.take_accumulator().as_ref()), 2.0);
+        assert_eq!(cardinality(&kmv), 0.0);
+    }
+}

crates/asap-physical-operators/src/summary_kernels/mod.rs

@@ -10,6 +10,7 @@ pub mod hll_sketch;
 pub mod hydra_cms;
 pub mod hydra_kll;
 pub mod increase;
+pub mod kmv;
 pub mod univmon;
 
 pub use count_min_sketch::*;
@@ -22,6 +23,7 @@ pub use hll_sketch::*;
 pub use hydra_cms::*;
 pub use hydra_kll::*;
 pub use increase::*;
+pub use kmv::*;
 
 pub mod factory;
 pub mod traits;

crates/asap-physical-operators/src/values.rs

@@ -277,6 +277,7 @@ fn validate_state(family: &SummaryFamilyType, state: &dyn AggregateCore) -> Resu
     use crate::summary_kernels::{
         count_min_sketch::CountMinSketchAccumulator, datasketches_kll::DatasketchesKLLAccumulator,
         dd_sketch::DDSketchAccumulator, exact::ExactAccumulator, hll_sketch::HllSketchAccumulator,
+        kmv::KmvAccumulator,
     };
     use planner_types::post_asap::SketchParams;
     validate_family(family)?;
@@ -323,6 +324,10 @@ fn validate_state(family: &SummaryFamilyType, state: &dyn AggregateCore) -> Resu
                 .as_any()
                 .downcast_ref::<HllSketchAccumulator>()
                 .is_some_and(|s| s.inner.precision == u32::from(*precision)),
+            SketchParams::Kmv { k } => state
+                .as_any()
+                .downcast_ref::<KmvAccumulator>()
+                .is_some_and(|s| s.inner.k == *k as usize),
             SketchParams::Cms { width, depth } => state
                 .as_any()
                 .downcast_ref::<CountMinSketchAccumulator>()

crates/integration-tests/tests/planner_layering_example2.rs

@@ -12,6 +12,7 @@
 //! three statistics of one input over the same 1-min window, through the
 //! frontend's `distinct_over_time`, `entropy_over_time` and `l2_over_time`.
 
+mod physical_common;
 mod planner_layering_common;
 
 use std::collections::{BTreeMap, BTreeSet};
@@ -786,3 +787,98 @@ fn stage3_plan_selection_selected_plan_has_at_most_one_univmon() {
         .count();
     assert!(count <= 1, "{}: {count} UnivMons", selected.id);
 }
+
+/// Stage 3 over the candidates that answer Q1 with KMV and Q2 and Q3
+/// exactly selects one, and the executor runs its Q1: the distinct count of
+/// the window's sample values. Below KMV's `k` distinct values its estimate
+/// is exact. Only Q1 is compiled: the exact entropy and L2 of Q2 and Q3 do
+/// not compile yet ("per-entity aggregate requires an explicit entity
+/// binding").
+#[test]
+fn stage3_plan_selection_selected_kmv_plan_executes_end_to_end() {
+    use asap_aware_mapping::stage3_plan_selection::select_plan;
+    use asap_aware_mapping::PlanningModels;
+    use asap_physical_operators::physical_planner::{compile, promql_rows, InputContract};
+    use asap_physical_operators::runtime::Scope;
+    use asap_physical_operators::values::{Batch, Value};
+    use std::sync::Arc;
+
+    let workload = standin_workload();
+    let run = run_stages(&workload, lower_promql(&workload));
+    let kmv: Vec<_> = run
+        .logical
+        .iter()
+        .filter(|c| {
+            !shares_a_summary(c)
+                && (0..3)
+                    .map(|q| query_option(&c.dag, &c.query_roots, q))
+                    .eq(["Kmv", "exact", "exact"])
+        })
+        .flat_map(|c| run.physical_of(c).map(|p| p.candidate.clone()))
+        .collect();
+    assert!(!kmv.is_empty(), "a candidate answers Q1 with KMV");
+    let entries: Vec<_> = workload.query_workload.entries().collect();
+    let targets: Vec<_> = entries
+        .iter()
+        .map(|e| Some(e.requirements.accuracy.target()))
+        .collect();
+    let recurrence: Vec<_> = entries.iter().map(|e| e.recurrence.clone()).collect();
+    let selection = select_plan(
+        &kmv,
+        &targets,
+        &recurrence,
+        workload.data_workload.as_ref().unwrap(),
+        PlanningModels::builtin(),
+    )
+    .expect("Stage 3 selects a KMV plan");
+    let p = run.physical(&selection.selected);
+    assert_eq!(query_option(p.dag(), &p.query_roots, 0), "Kmv");
+
+    // 1,000 flow records over 50 s with 300 distinct source IPs.
+    let distinct = 300u32;
+    let samples: Vec<(i64, f64)> = (0..1_000u32)
+        .map(|i| (i64::from(i) * 50, f64::from(i * 7 % distinct)))
+        .collect();
+    let labels = BTreeMap::from([("__name__".to_string(), "flows_src_ip".to_string())]);
+    let dag = p.dag();
+    let mut inputs = BTreeMap::new();
+    let mut batches = BTreeMap::new();
+    for node in &dag.nodes {
+        if relational(&node.payload).as_deref() == Some("time_range") {
+            let schema = Arc::new(node.output_schema.clone());
+            let rows = samples
+                .iter()
+                .map(|(ts, v)| promql_rows::series_row(&schema, &labels, *ts, *v).unwrap())
+                .collect();
+            batches.insert(
+                node.id as u64,
+                Batch::try_new(schema.clone(), rows).unwrap(),
+            );
+            inputs.insert(node.id as u64, InputContract::bounded(schema));
+        }
+    }
+    let q1 = p.query_roots[0] as u64;
+    let plan = compile(dag, inputs, &[q1]).unwrap_or_else(|e| panic!("{}: {e}", p.id));
+    let batches = plan
+        .input_contracts()
+        .map(|(id, _)| (id, batches[&id].clone()))
+        .collect();
+    let result = physical_common::execute(
+        &plan,
+        batches,
+        Scope::Query {
+            evaluation_time_ms: 55_000,
+            revision: 1,
+        },
+    );
+    let values: Vec<f64> = result[0]
+        .iter()
+        .flat_map(|batch| batch.rows())
+        .flat_map(|row| row.iter())
+        .filter_map(|value| match value {
+            Value::Float64(v) => Some(*v),
+            _ => None,
+        })
+        .collect();
+    assert_eq!(values, [f64::from(distinct)], "{}", p.id);
+}