diff --git a/Cargo.lock b/Cargo.lock index 1523a0f4c..a4600e45b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -364,7 +364,7 @@ dependencies = [ [[package]] name = "asap-aware-mapping" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=76fbbf16cc44b19f56a780bfdb47a95327e84711#76fbbf16cc44b19f56a780bfdb47a95327e84711" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=c98281a59df59740f006ca0b5b8da1ea9bb7fc41#c98281a59df59740f006ca0b5b8da1ea9bb7fc41" dependencies = [ "asap-types", "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib)", @@ -376,7 +376,7 @@ dependencies = [ [[package]] name = "asap-frontend-promql" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=76fbbf16cc44b19f56a780bfdb47a95327e84711#76fbbf16cc44b19f56a780bfdb47a95327e84711" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=c98281a59df59740f006ca0b5b8da1ea9bb7fc41#c98281a59df59740f006ca0b5b8da1ea9bb7fc41" dependencies = [ "asap-types", "promql-parser 0.10.0 (git+https://github.com/ProjectASAP/promql-parser?rev=9fede7eecca923c9882fe256484d00d37f8706cb)", @@ -385,7 +385,7 @@ dependencies = [ [[package]] name = "asap-frontend-sql" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=76fbbf16cc44b19f56a780bfdb47a95327e84711#76fbbf16cc44b19f56a780bfdb47a95327e84711" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=c98281a59df59740f006ca0b5b8da1ea9bb7fc41#c98281a59df59740f006ca0b5b8da1ea9bb7fc41" dependencies = [ "asap-sql-function-catalog", "asap-types", @@ -396,36 +396,29 @@ dependencies = [ [[package]] name = "asap-physical-operators" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=76fbbf16cc44b19f56a780bfdb47a95327e84711#76fbbf16cc44b19f56a780bfdb47a95327e84711" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=c98281a59df59740f006ca0b5b8da1ea9bb7fc41#c98281a59df59740f006ca0b5b8da1ea9bb7fc41" dependencies = [ "asap-types", - "asap_sketch_codec", "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", - "base64 0.21.7", - "bincode", "futures", - "prost", - "rmp-serde", "serde", "serde_json", "thiserror 2.0.20", "tracing", - "xxhash-rust", ] [[package]] name = "asap-sql-function-catalog" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=76fbbf16cc44b19f56a780bfdb47a95327e84711#76fbbf16cc44b19f56a780bfdb47a95327e84711" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=c98281a59df59740f006ca0b5b8da1ea9bb7fc41#c98281a59df59740f006ca0b5b8da1ea9bb7fc41" [[package]] name = "asap-types" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=76fbbf16cc44b19f56a780bfdb47a95327e84711#76fbbf16cc44b19f56a780bfdb47a95327e84711" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=c98281a59df59740f006ca0b5b8da1ea9bb7fc41#c98281a59df59740f006ca0b5b8da1ea9bb7fc41" dependencies = [ "serde", "serde_json", - "sha2", "thiserror 2.0.20", ] @@ -443,7 +436,6 @@ dependencies = [ [[package]] name = "asap_sketch_codec" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=76fbbf16cc44b19f56a780bfdb47a95327e84711#76fbbf16cc44b19f56a780bfdb47a95327e84711" dependencies = [ "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", "prost", @@ -484,6 +476,25 @@ dependencies = [ "xxhash-rust", ] +[[package]] +name = "asap_summary_state" +version = "0.1.0" +dependencies = [ + "asap-physical-operators", + "asap-types", + "asap_sketch_codec", + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", + "base64 0.21.7", + "bincode", + "hex", + "prost", + "rmp-serde", + "serde", + "serde_json", + "thiserror 2.0.20", + "tracing", +] + [[package]] name = "asap_types" version = "0.1.0" @@ -492,6 +503,7 @@ dependencies = [ "asap-aware-mapping", "asap-physical-operators", "asap-types", + "asap_summary_state", "base64 0.21.7", "clap", "promql-parser 0.10.0 (git+https://github.com/ProjectASAP/promql-parser)", @@ -1185,6 +1197,7 @@ dependencies = [ "asap_otel_proto", "asap_sketch_codec", "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", + "asap_summary_state", "asap_types", "async-trait", "axum", diff --git a/Cargo.toml b/Cargo.toml index 2be96caea..121c5539a 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,6 +2,8 @@ resolver = "2" members = [ "crates/asap_otel_proto", + "crates/asap_sketch_codec", + "crates/asap_summary_state", "crates/asap_types", "data_plane", "control_plane", @@ -14,10 +16,10 @@ version = "0.1.0" [workspace.dependencies] # Keep Planner frontends, selection, and IR on the same immutable revision. # Alias upstream asap-types because this workspace also defines asap_types. -planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "76fbbf16cc44b19f56a780bfdb47a95327e84711" } -asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "76fbbf16cc44b19f56a780bfdb47a95327e84711" } -asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "76fbbf16cc44b19f56a780bfdb47a95327e84711" } -asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "76fbbf16cc44b19f56a780bfdb47a95327e84711" } +planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "c98281a59df59740f006ca0b5b8da1ea9bb7fc41" } +asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "c98281a59df59740f006ca0b5b8da1ea9bb7fc41" } +asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "c98281a59df59740f006ca0b5b8da1ea9bb7fc41" } +asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "c98281a59df59740f006ca0b5b8da1ea9bb7fc41" } # Shared external deps (used by 2+ crates) serde = { version = "1.0", features = ["derive"] } @@ -37,8 +39,9 @@ arc-swap = "1.7" reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } # Internal crates -asap-physical-operators = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "76fbbf16cc44b19f56a780bfdb47a95327e84711" } -asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "76fbbf16cc44b19f56a780bfdb47a95327e84711" } +asap-physical-operators = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "c98281a59df59740f006ca0b5b8da1ea9bb7fc41" } +asap_sketch_codec = { path = "crates/asap_sketch_codec" } +asap_summary_state = { path = "crates/asap_summary_state" } asap_types = { path = "crates/asap_types" } asap_otel_proto = { path = "crates/asap_otel_proto" } indexmap = { version = "2.0", features = ["serde"] } diff --git a/control_plane/src/main.rs b/control_plane/src/main.rs index ef5ac2254..326044635 100644 --- a/control_plane/src/main.rs +++ b/control_plane/src/main.rs @@ -199,7 +199,7 @@ struct CompileAndPublishPhysicalPlanRequest { workload_cost_evidence: Option, queries: Vec, data_workload: planner_types::workload::DataWorkload, - dataset_identity: planner_types::post_asap::LogicalDatasetIdentity, + dataset_identity: asap_types::semantic_fragment::LogicalDatasetIdentity, #[serde(rename = "collector_ids")] target_collector_ids: Vec, capability_snapshot_id: String, diff --git a/control_plane/src/physical/compiler.rs b/control_plane/src/physical/compiler.rs index af7c34bba..8078fad2c 100644 --- a/control_plane/src/physical/compiler.rs +++ b/control_plane/src/physical/compiler.rs @@ -3,6 +3,7 @@ //! Planner owns semantic candidates and guarantees. This module owns the //! deployment decision: evidence freshness, target capabilities, windows, the //! Collector execution projection, SummaryCatalog, and executable plans. +use asap_types::physical_plan_codec::PhysicalPlanCodec; use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::rc::Rc; @@ -39,6 +40,7 @@ use crate::query_plan::{ use crate::types::AccuracyTarget; use planner_types::pre_asap::Source; +mod rate_placement; mod windows; pub(super) use windows::gcd; pub use windows::{prepare_window_implementations, WindowCostModel}; @@ -96,7 +98,7 @@ impl QueryCompilationInput { pub(crate) fn retain_physical_candidate(&mut self) -> Result<(), CompileError> { use asap_physical_operators::physical_planner::{promql_rows, PhysicalCandidate}; let candidate = - promql_rows::compile_fixed_window_rate_aggregation(&self.selected_plan_root) + rate_placement::compile_fixed_window_rate_aggregation(&self.selected_plan_root) .or_else(|_| { promql_rows::compile_current_series_readout(&self.selected_plan_root) .or_else(|_| { @@ -377,7 +379,7 @@ impl ScopedAccuracyEvidence { #[serde(deny_unknown_fields)] pub struct PhysicalDeploymentContext { /// Semantic dataset served by this deployment's input channel; never an endpoint. - pub dataset_identity: planner_types::post_asap::LogicalDatasetIdentity, + pub dataset_identity: asap_types::semantic_fragment::LogicalDatasetIdentity, pub target: PhysicalDeploymentTarget, #[serde(rename = "collector_ids")] pub target_collector_ids: Vec, @@ -1043,12 +1045,16 @@ impl BackendLocalPlanningInput { ); proposed .candidates - .extend(strategy.fixed_window_rate_candidates(&typed).candidates); - proposed.candidates.extend( - strategy - .query_time_rate_aggregation_candidates(&typed) - .candidates, - ); + .extend(rate_placement::fixed_window_rate_candidates( + &direct.candidates, + &typed, + )); + proposed + .candidates + .extend(rate_placement::query_time_rate_aggregation_candidates( + &direct.candidates, + &typed, + )); proposed.candidates.extend(direct.candidates); proposed.rejected.extend(direct.rejected); for candidate in proposed.candidates { @@ -1059,7 +1065,7 @@ impl BackendLocalPlanningInput { .map_err(|error| CompileError::Snapshot(error.to_string()))?; let compiled = asap_physical_operators::physical_planner::promql_rows::compile_current_series_readout(&root) .or_else(|_| asap_physical_operators::physical_planner::promql_rows::compile_rate_ranking(&root).map(|(_, program)| program)); - if let Ok(physical) = asap_physical_operators::physical_planner::promql_rows::compile_fixed_window_rate_aggregation(&root) { + if let Ok(physical) = rate_placement::compile_fixed_window_rate_aggregation(&root) { planner_selection_trace.push(serde_json::json!({ "stage":"planner.physical_candidate", "query_id":query.query_id, "logical_root_id":crate::planner_selection::explained_root_id(&root, &query.accuracy_target), @@ -4653,7 +4659,7 @@ fn collect_selected_materializations( physical_source.as_ref().unwrap_or(node), None, composable, - asap_physical_operators::physical_planner::promql_rows::compile_fixed_window_rate_aggregation(node).is_ok(), + rate_placement::compile_fixed_window_rate_aggregation(node).is_ok(), None, &mut selected, )?; @@ -5599,7 +5605,7 @@ pub(crate) mod tests { fn environment(now: u64) -> PhysicalDeploymentContext { PhysicalDeploymentContext { - dataset_identity: planner_types::post_asap::LogicalDatasetIdentity { + dataset_identity: asap_types::semantic_fragment::LogicalDatasetIdentity { namespace: "test".into(), dataset: "metrics".into(), }, diff --git a/control_plane/src/physical/compiler/rate_placement.rs b/control_plane/src/physical/compiler/rate_placement.rs new file mode 100644 index 000000000..413a99d04 --- /dev/null +++ b/control_plane/src/physical/compiler/rate_placement.rs @@ -0,0 +1,225 @@ +//! Rate-placement variants of Planner heap and grouped Sum candidates. +//! +//! Planner no longer lists these: it treats placement as a lifecycle choice. +//! Until the compiler selects placement through lifecycle timing, it derives +//! the same two variants that Planner used to offer, with timing written into +//! the candidate root, and compiles them as before. +use asap_aware_mapping::{Replacement, ReplacementSubDAG}; +use asap_physical_operators::physical_planner::{promql_rows, PhysicalCandidate}; +use asap_physical_operators::Error; +use planner_types::post_asap::{ + compile_post_asap_dag, ExactKind, ExecutionTiming, PostAsapOperatorPayload, SketchAlgorithm, + SummaryExpr, SummaryFamilyType, SummaryNode, ValueOperation, +}; +use planner_types::pre_asap::{QueryExpr, Reduction}; +use std::rc::Rc; + +/// Compile a candidate whose Rate finalization timing is part of its root. +pub(crate) fn compile_fixed_window_rate_aggregation( + selected: &Rc, +) -> Result { + let dag = compile_post_asap_dag(selected).map_err(|e| Error::Invalid(e.to_string()))?; + promql_rows::compile_fixed_window_rate_aggregation(&dag) +} + +/// Fixed-window maintenance finalizes each series' counter state and builds a +/// fresh heap or grouped Sum for that evaluation window. Deployment must +/// provide a complete, synchronized population and bind the matching window; +/// this never adds one window's rates to another. +pub(crate) fn fixed_window_rate_candidates( + direct: &[ReplacementSubDAG], + root: &Rc, +) -> Vec { + fn place(node: &Rc) -> Option> { + let mut next = node.as_ref().clone(); + match &mut next.expr { + SummaryExpr::ValueOperation { + child, + operation: ValueOperation::FinalizeExactAccumulator, + timing, + } if matches!(&child.expr, SummaryExpr::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + reduction: Reduction::PerEntity, child: source, .. + } if matches!(&source.expr, SummaryExpr::KeepPreAsap(source) if matches!(source.as_ref(), QueryExpr::TimeRange { .. }))) => + { + *timing = ExecutionTiming::IngestionTime; + } + SummaryExpr::ValueOperation { child, .. } | SummaryExpr::SummaryAgg { child, .. } => { + *child = place(child)? + } + SummaryExpr::SummaryEstimate { summary_input, .. } => { + *summary_input = place(summary_input)? + } + _ => return None, + } + Some(Rc::new(next)) + } + let mut candidates = direct.to_vec(); + candidates.retain_mut(|candidate| { + let Replacement::Summary(node) = &candidate.replacement else { + return false; + }; + let Ok(dag) = compile_post_asap_dag(node) else { + return false; + }; + if !dag.nodes.iter().any(|node| match &node.payload { + PostAsapOperatorPayload::SummaryAgg { + family: SummaryFamilyType::Sketch(kind, _), + .. + } => matches!( + kind.algorithm(), + SketchAlgorithm::CmsWithHeap | SketchAlgorithm::CountSketchWithHeap + ), + PostAsapOperatorPayload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), + .. + } => true, + _ => false, + }) { + return false; + } + let Some(placed) = place(node) else { + return false; + }; + if compile_post_asap_dag(&placed).is_err() { + return false; + } + let Ok(placed) = asap_aware_mapping::replacement::finalize_query_candidate(placed, root) + else { + return false; + }; + candidate.replacement = Replacement::Summary(placed); + candidate + .rationale + .push_str("; fixed-window precompute over complete per-series counter states"); + true + }); + candidates +} + +/// Retain grouped Sum after a per-series Rate readout as a query-time +/// candidate alongside its complete-window maintenance placement. +pub(crate) fn query_time_rate_aggregation_candidates( + direct: &[ReplacementSubDAG], + root: &Rc, +) -> Vec { + fn query_time(node: &Rc) -> Rc { + let mut next = node.as_ref().clone(); + match &mut next.expr { + SummaryExpr::ValueOperation { + child, + operation: ValueOperation::FinalizeExactAccumulator, + timing, + } if matches!( + &child.expr, + SummaryExpr::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + .. + } + ) => + { + *timing = ExecutionTiming::QueryTime; + } + SummaryExpr::ValueOperation { child, .. } | SummaryExpr::SummaryAgg { child, .. } => { + *child = query_time(child) + } + _ => {} + } + Rc::new(next) + } + let mut candidates = fixed_window_rate_candidates(direct, root); + candidates.retain_mut(|candidate| { + let Replacement::Summary(node) = &candidate.replacement else { + return false; + }; + if !matches!(&node.expr, SummaryExpr::ValueOperation { child, operation: ValueOperation::FinalizeExactAccumulator, .. } + if matches!(&child.expr, SummaryExpr::SummaryAgg { family: SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), .. })) + { + return false; + } + candidate.replacement = Replacement::Summary(query_time(node)); + candidate.rationale = "query-time grouped Sum over complete per-series Rate readouts".into(); + true + }); + candidates +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_aware_mapping::{ReplacementStrategy, SketchAlgorithmStrategy, TargetSubDAG}; + use planner_types::types::AccuracyTarget; + + fn candidates(query: &str) -> (Vec, Vec) { + let root = + crate::query_parser::parse_query_expr_canonical(query, AccuracyTarget::Exact).unwrap(); + let typed = Rc::new(promql_rows::with_series_identity(&root).unwrap()); + let strategy = + SketchAlgorithmStrategy::new(&asap_aware_mapping::cost_model::DefaultCostModel); + let direct = strategy.propose(&TargetSubDAG::new(&typed)).candidates; + ( + fixed_window_rate_candidates(&direct, &typed), + query_time_rate_aggregation_candidates(&direct, &typed), + ) + } + + fn rate_timing(node: &Rc) -> Option { + match &node.expr { + SummaryExpr::ValueOperation { + child, + operation: ValueOperation::FinalizeExactAccumulator, + timing, + } if matches!( + &child.expr, + SummaryExpr::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + .. + } + ) => + { + Some(*timing) + } + SummaryExpr::ValueOperation { child, .. } | SummaryExpr::SummaryAgg { child, .. } => { + rate_timing(child) + } + SummaryExpr::SummaryEstimate { summary_input, .. } => rate_timing(summary_input), + _ => None, + } + } + + // Grouped Sum over Rate yields a fixed-window precompute placement and a + // query-time placement of the same computation. + #[test] + fn grouped_rate_sum_has_both_placements() { + let (fixed, query_time) = candidates("sum by (job) (rate(requests_total[1m]))"); + let [fixed] = fixed.as_slice() else { + panic!("expected one fixed-window candidate, got {}", fixed.len()); + }; + let Replacement::Summary(fixed) = &fixed.replacement else { + panic!("expected a summary replacement"); + }; + assert_eq!(rate_timing(fixed), Some(ExecutionTiming::IngestionTime)); + let physical = compile_fixed_window_rate_aggregation(fixed).unwrap(); + assert!(physical.precompute.is_some()); + + let [query_time] = query_time.as_slice() else { + panic!( + "expected one query-time candidate, got {}", + query_time.len() + ); + }; + let Replacement::Summary(query_time) = &query_time.replacement else { + panic!("expected a summary replacement"); + }; + assert_eq!(rate_timing(query_time), Some(ExecutionTiming::QueryTime)); + assert!(compile_fixed_window_rate_aggregation(query_time).is_err()); + } + + // Candidates without a heap or grouped Sum have no Rate placement variant. + #[test] + fn per_series_rate_has_no_placement_variant() { + let (fixed, query_time) = candidates("rate(requests_total[1m])"); + assert!(fixed.is_empty()); + assert!(query_time.is_empty()); + } +} diff --git a/control_plane/src/physical/compiler/windows.rs b/control_plane/src/physical/compiler/windows.rs index fcf2987ae..02d81bee1 100644 --- a/control_plane/src/physical/compiler/windows.rs +++ b/control_plane/src/physical/compiler/windows.rs @@ -157,7 +157,9 @@ pub fn prepare_window_implementations( reason, })?; let cohorts = cohort_nodes(&states); - let native_cohort = asap_physical_operators::physical_planner::promql_rows::compile_fixed_window_rate_aggregation(&query.selected_plan_root).is_ok(); + let native_cohort = + super::rate_placement::compile_fixed_window_rate_aggregation(&query.selected_plan_root) + .is_ok(); let requirements = states .iter() .map(|state| { diff --git a/control_plane/src/physical/executable_binding.rs b/control_plane/src/physical/executable_binding.rs index 7d6ad4724..7639251ea 100644 --- a/control_plane/src/physical/executable_binding.rs +++ b/control_plane/src/physical/executable_binding.rs @@ -1,4 +1,5 @@ //! Compiler checks relating the shared executable contract to QueryPlan. +use asap_types::physical_plan_codec::PhysicalPlanCodec; pub use asap_types::executable_plan::*; diff --git a/control_plane/src/physical/maintained_population.rs b/control_plane/src/physical/maintained_population.rs index 7bfee745f..c35e1c717 100644 --- a/control_plane/src/physical/maintained_population.rs +++ b/control_plane/src/physical/maintained_population.rs @@ -2,6 +2,7 @@ #[cfg(test)] use super::compiler::QueryCompilationInput; use super::compiler::{CompileError, PhysicalCompilationRequest}; +use asap_types::physical_plan_codec::PhysicalPlanCodec; use asap_types::query_plan::{ current_series::{SeriesPopulation, SeriesReadout}, query_time::{Grouping, LabelMatch, LabelMatcher, QueryTimeOperator}, diff --git a/control_plane/src/physical/plan_dot.rs b/control_plane/src/physical/plan_dot.rs index 98bf25509..2a7c5e980 100644 --- a/control_plane/src/physical/plan_dot.rs +++ b/control_plane/src/physical/plan_dot.rs @@ -2,6 +2,7 @@ //! //! This is intentionally a developer-facing view: JSON remains the complete //! representation, while DOT keeps labels compact enough to follow execution. +use asap_types::physical_plan_codec::PhysicalPlanCodec; use super::compiler::CompiledPhysicalPlan; use crate::query_plan::QueryPlanNode; diff --git a/control_plane/src/physical/workload_cost.rs b/control_plane/src/physical/workload_cost.rs index c6ce7e028..88b0a607b 100644 --- a/control_plane/src/physical/workload_cost.rs +++ b/control_plane/src/physical/workload_cost.rs @@ -41,7 +41,7 @@ pub struct CostComponentDemand { #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(deny_unknown_fields)] pub struct WorkloadCostManifest { - pub dataset_identity: planner_types::post_asap::LogicalDatasetIdentity, + pub dataset_identity: asap_types::semantic_fragment::LogicalDatasetIdentity, pub plan_id: u64, pub plan_version: u64, pub planner_revision: String, diff --git a/control_plane/src/query_plan.rs b/control_plane/src/query_plan.rs index e6aadf931..3cf562d80 100644 --- a/control_plane/src/query_plan.rs +++ b/control_plane/src/query_plan.rs @@ -1,5 +1,6 @@ //! Control-plane lowering from Planner IR to the shared installed query DAG. //! Serving consumes asap_types::query_plan; compilation stays in this component. +use asap_types::physical_plan_codec::PhysicalPlanCodec; mod clickhouse_exact; pub mod physical_values; diff --git a/control_plane/src/query_plan/clickhouse_exact.rs b/control_plane/src/query_plan/clickhouse_exact.rs index 6d0538c0b..f63510c4b 100644 --- a/control_plane/src/query_plan/clickhouse_exact.rs +++ b/control_plane/src/query_plan/clickhouse_exact.rs @@ -1,5 +1,7 @@ //! Render a supported canonical relational cut without changing its row population. //! Unsupported operators remain admission errors, never guessed SQL semantics. +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; use planner_types::pre_asap::{ AggIntent, ArithmeticOpKind, CompareOpKind, QueryExpr, Reduction, ScalarValue, Schema, Source, }; diff --git a/control_plane/src/query_plan/physical_values.rs b/control_plane/src/query_plan/physical_values.rs index f3855896c..ec14b64cd 100644 --- a/control_plane/src/query_plan/physical_values.rs +++ b/control_plane/src/query_plan/physical_values.rs @@ -1,6 +1,7 @@ //! Retain Planner-compiled scalar/vector fragments before publishing an installed plan. use super::*; use asap_physical_operators::physical_planner::{promql_values as physical, CompiledPhysicalDag}; +use asap_types::physical_plan_codec::PhysicalPlanCodec; use planner_types::{ post_asap::BinaryOperator, pre_asap::{AggIntent, ArithmeticOpKind, BinaryOpKind, ColumnRef, CompareOpKind, GroupKeys}, diff --git a/control_plane/src/query_plan/query_time.rs b/control_plane/src/query_plan/query_time.rs index 705871aad..1d755d5d1 100644 --- a/control_plane/src/query_plan/query_time.rs +++ b/control_plane/src/query_plan/query_time.rs @@ -6,6 +6,8 @@ use super::{ FallbackPolicy, InstantExecution, QueryNodeId, QueryPlanEntry, QueryPlanError, QueryPlanNode, }; +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; use promql_parser::{ label::MatchOp, parser::{self, Expr, LabelModifier, Offset, VectorSelector}, diff --git a/control_plane/tests/native_rate_topk.rs b/control_plane/tests/native_rate_topk.rs index e8cc2de57..88ed8725e 100644 --- a/control_plane/tests/native_rate_topk.rs +++ b/control_plane/tests/native_rate_topk.rs @@ -1,4 +1,5 @@ //! Planner owns Rate ranking; Backend binds durable state and prices candidates. +use asap_types::physical_plan_codec::PhysicalPlanCodec; use control_plane::physical::{ compiler::{BackendLocalPlanningInput, DeploymentPlanCompiler}, workload_cost::enumerate_exact_and_materialized_candidates, diff --git a/crates/asap_sketch_codec/Cargo.toml b/crates/asap_sketch_codec/Cargo.toml new file mode 100644 index 000000000..a4c2d651d --- /dev/null +++ b/crates/asap_sketch_codec/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "asap_sketch_codec" +version.workspace = true +edition.workspace = true + +[dependencies] +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "5f03ccbd798ed5fec62bdd839bcb331123cab369" } +prost = "0.13" diff --git a/crates/asap_sketch_codec/src/lib.rs b/crates/asap_sketch_codec/src/lib.rs new file mode 100644 index 000000000..00efb284e --- /dev/null +++ b/crates/asap_sketch_codec/src/lib.rs @@ -0,0 +1,80 @@ +//! Runtime-independent decoding of the sketchlib protobuf envelope. + +use asap_sketchlib::proto::sketchlib::{ + sketch_envelope::SketchState, DdSketchState, KllState, SketchEnvelope, +}; +use asap_sketchlib::DdSketch; +use prost::Message; + +pub fn envelope_state(bytes: &[u8]) -> Result, String> { + SketchEnvelope::decode(bytes) + .map(|envelope| envelope.sketch_state) + .map_err(|error| format!("decode SketchEnvelope: {error}")) +} + +pub fn ddsketch_state(bytes: &[u8]) -> Result<(DdSketchState, f64), String> { + let envelope = + SketchEnvelope::decode(bytes).map_err(|error| format!("decode SketchEnvelope: {error}"))?; + match envelope.sketch_state { + Some(SketchState::Ddsketch(state)) => Ok((state, envelope.sample_p)), + _ => Err("SketchEnvelope contains no DDSketch state".into()), + } +} + +pub fn reconstruct_ddsketch(bytes: &[u8]) -> Result<(DdSketch, f64), String> { + let (state, sample_p) = ddsketch_state(bytes)?; + if !state.alpha.is_finite() || !(0.0..1.0).contains(&state.alpha) || state.alpha == 0.0 { + return Err("DDSketch alpha must be finite and between zero and one".into()); + } + Ok(( + DdSketch::from_raw(state.alpha, state.store_counts, state.store_offset), + sample_p, + )) +} + +pub fn kll_state(bytes: &[u8]) -> Result { + let envelope = + SketchEnvelope::decode(bytes).map_err(|error| format!("decode SketchEnvelope: {error}"))?; + match envelope.sketch_state { + Some(SketchState::Kll(state)) => Ok(state), + _ => Err("SketchEnvelope contains no KLL state".into()), + } +} + +pub fn encode_ddsketch(sketch: &DdSketch) -> Vec { + let envelope = SketchEnvelope { + format_version: 1, + producer: None, + hash_spec: None, + sample_p: 0.0, + sketch_state: Some(SketchState::Ddsketch(sketch.to_proto())), + }; + envelope.encode_to_vec() +} + +pub fn encode_kll(sketch: &asap_sketchlib::sketches::kll::KLL) -> Vec { + use asap_sketchlib::proto::sketchlib::CoinState; + let (state, bit_cache, remaining_bits) = sketch.wire_coin(); + SketchEnvelope { + format_version: 1, + producer: None, + hash_spec: None, + sample_p: 0.0, + sketch_state: Some(SketchState::Kll(KllState { + k: sketch.wire_k(), + m: sketch.wire_m(), + num_levels: sketch.wire_num_levels(), + levels: sketch.wire_levels(), + items: sketch.wire_items(), + coin: Some(CoinState { + state, + bit_cache, + remaining_bits, + }), + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + })), + } + .encode_to_vec() +} diff --git a/crates/asap_summary_state/Cargo.toml b/crates/asap_summary_state/Cargo.toml new file mode 100644 index 000000000..0e53c07f0 --- /dev/null +++ b/crates/asap_summary_state/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "asap_summary_state" +version.workspace = true +edition.workspace = true + +[dependencies] +asap-physical-operators.workspace = true +planner-types.workspace = true +asap_sketch_codec.workspace = true +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "5f03ccbd798ed5fec62bdd839bcb331123cab369" } +serde = { version = "1", features = ["derive", "rc"] } +serde_json.workspace = true +tracing.workspace = true +base64 = "0.21" +bincode = "1.3" +rmp-serde = "1.3" +prost = "0.13" +thiserror = "2" + +[dev-dependencies] +hex = "0.4" + +[features] +default = [] +extra_debugging = [] diff --git a/crates/asap_summary_state/src/aggregation_type.rs b/crates/asap_summary_state/src/aggregation_type.rs new file mode 100644 index 000000000..647f7604f --- /dev/null +++ b/crates/asap_summary_state/src/aggregation_type.rs @@ -0,0 +1,215 @@ +//! Shared aggregation vocabulary for configuration and accumulator dispatch. +//! The wire shape combines aggregation type, subtype, and parameters. +//! `AccumulatorSpec` provides a typed representation at conversion boundaries. + +use serde::{Deserialize, Serialize}; +use std::fmt; +use std::str::FromStr; + +/// Concrete aggregation/sketch type used in precompute configs and accumulator dispatch. +/// +/// `Display` outputs the canonical PascalCase name used in YAML/JSON configs. +/// `FromStr` accepts the canonical name plus legacy aliases (e.g. "KLL" → `DatasketchesKLL`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum AggregationType { + // ---------- single-population (non-keyed) ---------- + Sum, + Count, + Increase, + Rate, + Min, + Max, + DatasketchesKLL, + // ---------- multi-population (keyed) ---------- + HydraKLL, + CountMinSketch, + CountMinSketchWithHeap, + CountSketch, + CountSketchWithHeap, + // ---------- cardinality / set tracking ---------- + HLL, + UnivMon, + DDSketch, + // ---------- legacy config wrapper names ---------- + SingleSubpopulation, + MultipleSubpopulation, +} + +impl AggregationType { + /// Adapt a storage/processor tag to Planner's exact family. Keyed storage + /// changes the payload layout, not the semantic family. + pub fn planner_exact_family(self) -> Option { + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + let (kind, params) = match self { + Self::Sum => (ExactKind::Sum, ExactParams::Sum), + Self::Count => (ExactKind::Count, ExactParams::Count), + Self::Increase => (ExactKind::Increase, ExactParams::Increase), + Self::Rate => (ExactKind::Rate, ExactParams::Rate), + Self::Min => (ExactKind::Min, ExactParams::Min), + Self::Max => (ExactKind::Max, ExactParams::Max), + _ => return None, + }; + Some(SummaryFamilyType::ExactAggregate(kind, params)) + } + + pub fn as_str(self) -> &'static str { + match self { + AggregationType::Sum => "Sum", + AggregationType::Count => "Count", + AggregationType::Increase => "Increase", + AggregationType::Rate => "Rate", + AggregationType::Min => "Min", + AggregationType::Max => "Max", + AggregationType::DatasketchesKLL => "DatasketchesKLL", + AggregationType::HydraKLL => "HydraKLL", + AggregationType::CountMinSketch => "CountMinSketch", + AggregationType::CountMinSketchWithHeap => "CountMinSketchWithHeap", + AggregationType::CountSketch => "CountSketch", + AggregationType::CountSketchWithHeap => "CountSketchWithHeap", + AggregationType::HLL => "HLL", + AggregationType::UnivMon => "UnivMon", + AggregationType::DDSketch => "DDSketch", + AggregationType::SingleSubpopulation => "SingleSubpopulation", + AggregationType::MultipleSubpopulation => "MultipleSubpopulation", + } + } + + /// Returns `true` if this type produces keyed (multi-population) accumulators. + pub fn is_keyed(self) -> bool { + matches!( + self, + AggregationType::MultipleSubpopulation + | AggregationType::CountMinSketch + | AggregationType::CountMinSketchWithHeap + | AggregationType::CountSketch + | AggregationType::CountSketchWithHeap + | AggregationType::HydraKLL + ) + } +} + +impl fmt::Display for AggregationType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +impl FromStr for AggregationType { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + // Canonical names + "Sum" => Ok(AggregationType::Sum), + "Count" => Ok(AggregationType::Count), + "Increase" => Ok(AggregationType::Increase), + "Rate" => Ok(AggregationType::Rate), + "Min" => Ok(AggregationType::Min), + "Max" => Ok(AggregationType::Max), + "DatasketchesKLL" => Ok(AggregationType::DatasketchesKLL), + "HydraKLL" => Ok(AggregationType::HydraKLL), + "CountMinSketch" => Ok(AggregationType::CountMinSketch), + "CountMinSketchWithHeap" => Ok(AggregationType::CountMinSketchWithHeap), + "CountSketch" => Ok(AggregationType::CountSketch), + "CountSketchWithHeap" => Ok(AggregationType::CountSketchWithHeap), + "HLL" | "HyperLogLog" => Ok(AggregationType::HLL), + "UnivMon" => Ok(AggregationType::UnivMon), + "DDSketch" | "DdSketch" => Ok(AggregationType::DDSketch), + "SingleSubpopulation" => Ok(AggregationType::SingleSubpopulation), + "MultipleSubpopulation" => Ok(AggregationType::MultipleSubpopulation), + // Legacy accumulator-suffixed aliases + "SumAccumulator" | "SumAggregator" | "sum" => Ok(AggregationType::Sum), + "IncreaseAccumulator" | "IncreaseAggregator" | "increase" => { + Ok(AggregationType::Increase) + } + "MinAccumulator" | "MinAggregator" | "min" => Ok(AggregationType::Min), + "MaxAccumulator" | "MaxAggregator" | "max" => Ok(AggregationType::Max), + "DatasketchesKLLAccumulator" | "KLL" | "kll" | "datasketches_kll" => { + Ok(AggregationType::DatasketchesKLL) + } + "HydraKllSketchAccumulator" | "hydra_kll" => Ok(AggregationType::HydraKLL), + "CountMinSketchAccumulator" | "CMS" | "cms" | "count_min_sketch" => { + Ok(AggregationType::CountMinSketch) + } + "CountMinSketchWithHeapAccumulator" => Ok(AggregationType::CountMinSketchWithHeap), + "CountSketchAccumulator" | "CS" | "cs" | "count_sketch" => { + Ok(AggregationType::CountSketch) + } + "CountSketchWithHeapAccumulator" => Ok(AggregationType::CountSketchWithHeap), + // Retired names. `MinMax` used to be one accumulator whose + // direction rode alongside in `aggregationSubType`; the two + // directions are separate types now, so there is no safe + // direction to guess here -- resolving a min workload as a + // max one is silently wrong, not merely imprecise. + "MinMax" + | "MinMaxAccumulator" + | "MinMaxAggregator" + | "min_max" + | "MultipleMinMax" + | "MultipleMinMaxAccumulator" + | "multiple_min_max" => Err(format!( + "Retired aggregation type: '{s}' -- min and max are separate types now, \ + use 'Min'/'Max'" + )), + _ => Err(format!("Unknown aggregation type: '{s}'")), + } + } +} + +impl Serialize for AggregationType { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_str(self.as_str()) + } +} + +impl<'de> Deserialize<'de> for AggregationType { + fn deserialize>(deserializer: D) -> Result { + let s = String::deserialize(deserializer)?; + s.parse().map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + + /// Removed layout tags cannot be installed as semantic families. + #[test] + fn rejects_keyed_family_aliases() { + for name in [ + "MultipleSum", + "MultipleIncrease", + "MultipleMin", + "MultipleMax", + ] { + assert!(name.parse::().is_err(), "{name}"); + } + } + + #[test] + fn storage_layout_tags_do_not_create_planner_families() { + for (storage, expected) in [ + (AggregationType::Sum, ExactKind::Sum), + (AggregationType::Count, ExactKind::Count), + (AggregationType::Increase, ExactKind::Increase), + (AggregationType::Rate, ExactKind::Rate), + ] { + let family = storage.planner_exact_family().unwrap(); + assert!( + matches!(family, SummaryFamilyType::ExactAggregate(kind, _) if kind == expected) + ); + } + assert_eq!( + AggregationType::Rate.planner_exact_family(), + Some(SummaryFamilyType::ExactAggregate( + ExactKind::Rate, + ExactParams::Rate + )) + ); + assert_ne!( + AggregationType::Rate.planner_exact_family(), + AggregationType::Increase.planner_exact_family() + ); + } +} diff --git a/crates/asap_summary_state/src/codec.rs b/crates/asap_summary_state/src/codec.rs new file mode 100644 index 000000000..856ec1486 --- /dev/null +++ b/crates/asap_summary_state/src/codec.rs @@ -0,0 +1,84 @@ +//! Stored encodings of Planner key and measurement values. +use crate::{KeyByLabelValues, Measurement}; + +pub trait KeyCodec: Sized { + fn serialize_to_json(&self) -> serde_json::Value; + fn deserialize_from_json(data: &serde_json::Value) -> Result; + fn serialize_to_bytes(&self) -> Vec; + fn deserialize_from_bytes(buffer: &[u8]) -> Result>; +} + +impl KeyCodec for KeyByLabelValues { + fn serialize_to_json(&self) -> serde_json::Value { + serde_json::to_value(&self.labels).unwrap_or(serde_json::Value::Null) + } + fn deserialize_from_json(data: &serde_json::Value) -> Result { + Ok(Self::new_with_labels(serde_json::from_value(data.clone())?)) + } + fn serialize_to_bytes(&self) -> Vec { + bincode::serialize(&self.labels).unwrap_or_default() + } + fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + Ok(Self::new_with_labels(bincode::deserialize(buffer)?)) + } +} + +pub trait MeasurementCodec: Sized { + fn serialize_to_json(&self) -> serde_json::Value; + fn deserialize_from_json(data: &serde_json::Value) -> Result; + fn serialize_to_bytes(&self) -> Vec; + fn deserialize_from_bytes(buffer: &[u8]) -> Result>; +} + +impl MeasurementCodec for Measurement { + fn serialize_to_json(&self) -> serde_json::Value { + serde_json::json!({ "value": self.value }) + } + fn deserialize_from_json(data: &serde_json::Value) -> Result { + let value = data["value"].as_f64().ok_or_else(|| { + serde_json::Error::io(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "Missing or invalid 'value' field", + )) + })?; + Ok(Self::new(value)) + } + fn serialize_to_bytes(&self) -> Vec { + self.value.to_le_bytes().to_vec() + } + fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let bytes: [u8; 8] = buffer + .get(..8) + .and_then(|b| b.try_into().ok()) + .ok_or("Buffer too short for f64")?; + Ok(Self::new(f64::from_le_bytes(bytes))) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Keys and measurements round-trip through their stored JSON and byte forms. + #[test] + fn key_and_measurement_encodings_roundtrip() { + let key = KeyByLabelValues::new_with_labels(vec!["a".into(), "b".into()]); + assert_eq!( + KeyByLabelValues::deserialize_from_json(&key.serialize_to_json()).unwrap(), + key + ); + assert_eq!( + KeyByLabelValues::deserialize_from_bytes(&key.serialize_to_bytes()).unwrap(), + key + ); + let m = Measurement::new(42.5); + assert_eq!( + Measurement::deserialize_from_json(&m.serialize_to_json()).unwrap(), + m + ); + assert_eq!( + Measurement::deserialize_from_bytes(&m.serialize_to_bytes()).unwrap(), + m + ); + } +} diff --git a/crates/asap_summary_state/src/lib.rs b/crates/asap_summary_state/src/lib.rs new file mode 100644 index 000000000..db5a1e9ac --- /dev/null +++ b/crates/asap_summary_state/src/lib.rs @@ -0,0 +1,19 @@ +//! Backend-owned summary state: the kernels that ingest and the sketch store +//! keep, their stored byte formats, and conversion to Planner physical states. +//! +//! Planner's `asap-physical-operators` keeps only in-memory computation state. +//! Storage formats, delta reconstruction and the legacy per-statistic kernels +//! are deployment concerns and live here. + +pub mod summary_kernels; +pub use summary_kernels::{factory, traits}; +pub use traits::*; + +mod aggregation_type; +pub use aggregation_type::AggregationType; + +pub mod codec; +pub mod physical; +pub mod stored_state; + +pub use asap_physical_operators::{KeyByLabelValues, Measurement, Statistic}; diff --git a/crates/asap_summary_state/src/physical.rs b/crates/asap_summary_state/src/physical.rs new file mode 100644 index 000000000..80149e7dd --- /dev/null +++ b/crates/asap_summary_state/src/physical.rs @@ -0,0 +1,212 @@ +//! Conversion between stored summary state and Planner physical state. +//! +//! Planner physical operators accept only their own in-memory kernels. Stored +//! kernels carry storage-only fields (for example edge `sample_p`), so only the +//! families Planner can bind as typed inputs convert, and only when those +//! fields are neutral. +use crate::summary_kernels as stored; +use crate::AggregateCore as StoredState; +use asap_physical_operators::summary_kernels as physical; +use asap_physical_operators::AggregateCore as PhysicalState; +use std::sync::Arc; + +type Error = Box; + +/// Bind a stored state as a Planner physical input. +/// +/// A stored `SumAccumulator` becomes Planner's unkeyed exact Sum state, which +/// is the only exact representation Planner operators accept. +pub fn to_physical(state: &dyn StoredState) -> Result, Error> { + let any = state.as_any(); + if let Some(s) = any.downcast_ref::() { + unsampled(s.sample_p)?; + return Ok(Arc::new(physical::DDSketchAccumulator { + inner: s.inner.clone(), + })); + } + if let Some(s) = any.downcast_ref::() { + unsampled(s.sample_p)?; + return Ok(Arc::new(physical::HllSketchAccumulator { + inner: s.inner.clone(), + })); + } + if let Some(s) = any.downcast_ref::() { + return Ok(Arc::new(physical::DatasketchesKLLAccumulator { + inner: s.inner.clone(), + })); + } + if let Some(s) = any.downcast_ref::() { + return Ok(Arc::new(exact_to_physical(s)?)); + } + if let Some(s) = any.downcast_ref::() { + return Ok(Arc::new(s.0.clone())); + } + if let Some(s) = any.downcast_ref::() { + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + let mut exact = physical::exact::ExactAccumulator::new(family, false)?; + exact.update(None, s.sum, 0); + return Ok(Arc::new(exact)); + } + Err(format!("{} has no Planner physical state", state.type_name()).into()) +} + +/// Keep a Planner physical output in the stored kernel family. +pub fn from_physical(state: &dyn PhysicalState) -> Result, Error> { + let any = state.as_any(); + if let Some(s) = any.downcast_ref::() { + return Ok(Box::new(stored::DDSketchAccumulator { + inner: s.inner.clone(), + sample_p: 1.0, + })); + } + if let Some(s) = any.downcast_ref::() { + return Ok(Box::new(stored::HllSketchAccumulator { + inner: s.inner.clone(), + sample_p: 1.0, + })); + } + if let Some(s) = any.downcast_ref::() { + return Ok(Box::new(stored::DatasketchesKLLAccumulator { + inner: s.inner.clone(), + })); + } + if let Some(s) = any.downcast_ref::() { + return Ok(Box::new(stored::weighted_frequency::WeightedFrequency( + s.clone(), + ))); + } + if let Some(s) = any.downcast_ref::() { + return Ok(Box::new( + stored::exact::ExactAccumulator::deserialize_from_bytes(&rmp_serde::to_vec_named(s)?)?, + )); + } + Err("Planner physical state has no stored kernel".into()) +} + +/// Stored kernel identity of a Planner physical state. +pub fn aggregation_type(state: &dyn PhysicalState) -> Result { + use crate::AggregationType as T; + let any = state.as_any(); + Ok(if any.is::() { + T::DDSketch + } else if any.is::() { + T::HLL + } else if any.is::() { + T::DatasketchesKLL + } else { + from_physical(state)?.get_accumulator_type() + }) +} + +/// Planner's weighted frequency state is serde-transparent over the sketchlib +/// kernel, which owns its persisted byte form. +pub(crate) fn frequency_kernel( + state: &physical::weighted_frequency::WeightedFrequency, +) -> Result { + Ok(rmp_serde::from_slice(&rmp_serde::to_vec(state)?)?) +} + +/// Both exact states share one serde shape; the stored decoder checks that the +/// payload matches its declared family. +fn exact_to_physical( + state: &stored::exact::ExactAccumulator, +) -> Result { + use crate::SerializableToSink; + Ok(rmp_serde::from_slice(&state.serialize_to_bytes())?) +} + +fn unsampled(sample_p: f64) -> Result<(), Error> { + if sample_p == 1.0 { + Ok(()) + } else { + Err("edge-sampled sketches have no Planner physical state".into()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{KeyByLabelValues, SerializableToSink, Statistic}; + use planner_types::post_asap::{ExactKind, ExactParams, SketchQuery, SummaryFamilyType}; + + // Sketches keep their estimates in both directions. + #[test] + fn sketches_convert_both_ways() { + let mut dd = stored::DDSketchAccumulator::new(0.01); + let mut kll = stored::DatasketchesKLLAccumulator::new(200); + for v in 1..=100 { + dd.inner.update(f64::from(v)); + kll.update(f64::from(v)); + } + let q = SketchQuery::Quantile { q: 0.5 }; + for state in [&dd as &dyn StoredState, &kll] { + let physical = to_physical(state).unwrap(); + let median = physical.estimate(&q).unwrap(); + assert!((median - 50.0).abs() <= 2.0, "{median}"); + let back = from_physical(physical.as_ref()).unwrap(); + assert_eq!(back.serialize_to_bytes(), state.serialize_to_bytes()); + } + } + + // Edge-sampled sketches must not become unscaled Planner inputs. + #[test] + fn sampled_sketch_is_rejected() { + let mut dd = stored::DDSketchAccumulator::new(0.01); + dd.sample_p = 0.5; + assert!(to_physical(&dd).is_err()); + } + + // Exact state keeps family and value; a stored Sum binds as Planner exact Sum. + #[test] + fn exact_and_sum_bind_as_planner_exact_state() { + let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + let mut exact = stored::exact::ExactAccumulator::new(family.clone(), false).unwrap(); + exact.update(None, 4.5, 10); + let physical = to_physical(&exact).unwrap(); + let read = |state: &Arc| { + state + .as_any() + .downcast_ref::() + .unwrap() + .readout(Statistic::Sum, None, None::<&KeyByLabelValues>) + .unwrap() + }; + assert_eq!(read(&physical), Some(4.5)); + let back = from_physical(physical.as_ref()).unwrap(); + assert_eq!(back.serialize_to_bytes(), exact.serialize_to_bytes()); + + let sum = stored::SumAccumulator::with_sum(7.0); + assert_eq!(read(&to_physical(&sum).unwrap()), Some(7.0)); + } + + // Heap state keeps its bytes and kernel identity through the stored form. + #[test] + fn weighted_frequency_converts_both_ways() { + use asap_physical_operators::values::Value; + use physical::weighted_frequency::{FrequencyAlgorithm, WeightedFrequency}; + let mut state = WeightedFrequency::new(FrequencyAlgorithm::Cms, 64, 5, 8).unwrap(); + state.update(&[Value::Utf8("a".into())], 0.5).unwrap(); + let stored = from_physical(&state).unwrap(); + assert_eq!( + stored.get_accumulator_type(), + crate::AggregationType::CountMinSketchWithHeap + ); + assert_eq!( + aggregation_type(&state).unwrap(), + crate::AggregationType::CountMinSketchWithHeap + ); + let back = to_physical(stored.as_ref()).unwrap(); + let back = back.as_any().downcast_ref::().unwrap(); + assert_eq!( + frequency_kernel(back).unwrap().to_bytes(), + frequency_kernel(&state).unwrap().to_bytes() + ); + } + + // Stored-only kernels have no Planner physical representation. + #[test] + fn stored_only_kernels_do_not_convert() { + assert!(to_physical(&stored::MinAccumulator::new()).is_err()); + } +} diff --git a/crates/asap_summary_state/src/stored_state/decoders.rs b/crates/asap_summary_state/src/stored_state/decoders.rs new file mode 100644 index 000000000..37928f8f4 --- /dev/null +++ b/crates/asap_summary_state/src/stored_state/decoders.rs @@ -0,0 +1,366 @@ +//! Shared sketch state reconstruction and decoding. +use asap_sketchlib::CountMinSketch; +use asap_sketchlib::CountMinSketchDelta; +use asap_sketchlib::CountMinSketchWithHeap; +use asap_sketchlib::CountSketch; +use asap_sketchlib::CountSketchDelta; +use asap_sketchlib::CountSketchWithHeap; +use asap_sketchlib::CsHeapItem; +use asap_sketchlib::MessagePackCodec; + +use crate::summary_kernels::count_min_sketch_with_heap::CountMinSketchWithHeapAccumulator; + +/// Decode a `CountMinSketch` from the modified-OTLP wire bytes. +/// MSGPACK path round-trips `CountMinSketch::deserialize_msgpack`; +/// PROTO path decodes a `SketchEnvelope{count_min: CountMinState}` +/// (or bare `CountMinState`) and re-projects to a flat matrix. Mirrors +/// `precompute_operators::count_min_sketch::from_sketchlib_proto_bytes`. +pub fn decode_cms_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::CountMin(st)) => st, + Some(_) => return Err("SketchEnvelope contains non-CountMin sketch".to_string()), + None => { + CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))? + } + }, + Err(_) => { + CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))? + } + }; + let rows = state.rows as usize; + let cols = state.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountMinState has zero dims (rows={rows}, cols={cols})" + )); + } + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type) + .map_err(|_| format!("CountMinState unknown counter_type {}", state.counter_type))?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountMinState counts_int has {} entries, expected {}", + state.counts_int.len(), + expected_len + )); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountMinState counts_float has {} entries, expected {}", + state.counts_float.len(), + expected_len + )); + } + state.counts_float.clone() + } + other => { + return Err(format!( + "CountMinState counter_type {other:?} not yet supported in reducer" + )); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(CountMinSketch::from_legacy_matrix(matrix, rows, cols)) +} + +/// Decode a `CountMinSketch` from msgpack bytes (sketch-core wire +/// format). Mirrors +/// `CountMinSketchAccumulator::from_msgpack_bytes`. +pub fn decode_cms_from_msgpack(buffer: &[u8]) -> Result { + CountMinSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketch msgpack: {e}")) +} + +/// Decode a `CountSketch` from the modified-OTLP proto wire bytes. +/// Mirrors +/// `precompute_operators::count_sketch::from_sketchlib_proto_bytes`. +pub fn decode_cs_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountSketchState, CounterType, SketchEnvelope, + }; + use prost::Message; + + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::CountSketch(st)) => st, + Some(_) => return Err("SketchEnvelope contains non-CountSketch sketch".to_string()), + None => CountSketchState::decode(buffer) + .map_err(|e| format!("decode CountSketchState: {e}"))?, + }, + Err(_) => { + CountSketchState::decode(buffer).map_err(|e| format!("decode CountSketchState: {e}"))? + } + }; + let rows = state.rows as usize; + let cols = state.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountSketchState has zero dims (rows={rows}, cols={cols})" + )); + } + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { + format!( + "CountSketchState unknown counter_type {}", + state.counter_type + ) + })?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountSketchState counts_int has {} entries, expected {}", + state.counts_int.len(), + expected_len + )); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountSketchState counts_float has {} entries, expected {}", + state.counts_float.len(), + expected_len + )); + } + state.counts_float.clone() + } + other => { + return Err(format!( + "CountSketchState counter_type {other:?} not yet supported in reducer" + )); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(CountSketch::from_legacy_matrix(matrix, rows, cols)) +} + +/// Decode a `CountSketch` from msgpack bytes (sketch-core wire format). +pub fn decode_cs_from_msgpack(buffer: &[u8]) -> Result { + CountSketch::from_msgpack(buffer).map_err(|e| format!("deserialize CountSketch msgpack: {e}")) +} + +/// Decode a `CountMinSketchWithHeap` from msgpack bytes — the OTLP +/// `CountMinSketch` wire bytes when the gateway/precompute layer +/// marked the sid as CmsWithHeap (heap embedded in the +/// `CountMinSketchWithHeapSerialized` outer wrapper). Delegates to +/// `asap_sketchlib::CountMinSketchWithHeap::deserialize_msgpack`. +pub fn decode_cms_with_heap_from_msgpack(buffer: &[u8]) -> Result { + CountMinSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketchWithHeap msgpack: {e}")) +} + +/// Decode a `CountSketchWithHeap` (median-estimator, Count Sketch family) +/// from msgpack bytes. Distinct wire type from `CountMinSketchWithHeap` +/// (min-estimator, Count-Min family) even though both are heap-bearing +/// frequency sketches — see `asap_sketchlib::CountSketchWithHeap`. +/// Delegates to `asap_sketchlib::CountSketchWithHeap::from_msgpack`. +pub fn decode_cs_with_heap_from_msgpack(buffer: &[u8]) -> Result { + CountSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketchWithHeap msgpack: {e}")) +} + +// --------------------------------------------------------------------------- +// Delta decoders. Under the per-window-reset (PWR) contract +// (`asap-precompute-go/window.go`: a delta is that window's own state +// applied onto a freshly-reset per-series sketch), each stored *Delta +// frame reconstructs into the FULL window state when applied onto an +// EMPTY base of the frame's declared dimensions. The reducer's +// `FrequencyEstimate` / `FrequencyTopk` paths are per-window evaluations, +// so "empty + apply(this window's delta)" yields exactly the window's +// matrix/heap — no cross-window stitching needed (mirrors how the ingest +// accumulators reset_to_empty per window before applying). +// +// The proto path reuses the PUBLIC `asap_sketchlib::{CountSketch, +// CountMinSketch}::apply_delta`; the proto `*Delta` message is decoded via +// `asap_sketchlib::proto::sketchlib::{CountSketchDelta, CountMinDelta}`, +// exactly as `precompute_operators::{count_sketch, +// count_min_sketch}::apply_proto_delta_bytes` does. +// --------------------------------------------------------------------------- + +/// Decode a `CountMinSketch` PROTO_DELTA frame into a FULL sketch by +/// applying the sparse cell delta onto an empty base of the frame's +/// declared dimensions. Mirrors +/// `precompute_operators::count_min_sketch::apply_proto_delta_bytes`. +pub fn decode_cms_from_proto_delta(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountMinDelta packed-array length mismatch: cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + )); + } + let rows = pb.rows as usize; + let cols = pb.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountMinDelta has zero dims (rows={rows}, cols={cols})" + )); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + // hh_keys is parsed off the wire by the precompute accumulator but + // intentionally dropped (the vendored Go proto bindings don't yet + // populate it); match that to keep behavior identical. + let delta = CountMinSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l1: pb.l1, + l2: pb.l2, + hh_keys: Vec::new(), + }; + let mut cms = CountMinSketch::from_legacy_matrix(vec![vec![0.0; cols]; rows], rows, cols); + cms.apply_delta(&delta) + .map_err(|e| format!("apply CountMinDelta onto empty base: {e}"))?; + Ok(cms) +} + +/// Decode a `CountSketch` PROTO_DELTA frame into a FULL sketch by applying +/// the sparse cell delta onto an empty base of the frame's declared +/// dimensions. Mirrors +/// `precompute_operators::count_sketch::apply_proto_delta_bytes`. +pub fn decode_cs_from_proto_delta(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountSketchDelta packed-array length mismatch: cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + )); + } + let rows = pb.rows as usize; + let cols = pb.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountSketchDelta has zero dims (rows={rows}, cols={cols})" + )); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + let delta = CountSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l2: pb.l2, + hh_keys: Vec::new(), + }; + let mut cs = CountSketch::from_legacy_matrix(vec![vec![0.0; cols]; rows], rows, cols); + cs.apply_delta(&delta) + .map_err(|e| format!("apply CountSketchDelta onto empty base: {e}"))?; + Ok(cs) +} + +/// Decode a heap-bearing CountSketch MSGPACK_DELTA frame into a FULL +/// `CountMinSketchWithHeap` by applying the sparse matrix delta + full +/// heap onto an empty base of the frame's declared dimensions. This +/// REUSES the ingest-side delta-heap apply logic +/// (`CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes` → +/// `apply_msgpack_heap_delta_bytes`), which decodes the frame generically +/// with `rmp_serde` — no `asap_sketchlib` delta API is added. +pub fn decode_cms_with_heap_from_msgpack_delta( + buffer: &[u8], +) -> Result { + let acc = CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes(buffer) + .map_err(|e| format!("reconstruct CountMinSketchWithHeap from delta: {e}"))?; + Ok(acc.inner) +} + +/// Decode a heap-bearing CountSketch (median-estimator) MSGPACK_DELTA frame +/// into a FULL `asap_sketchlib::CountSketchWithHeap` by applying the sparse +/// matrix delta + full heap onto an empty base of the frame's declared +/// dimensions. Same DELTA-HEAP wire shape as the CmsWithHeap delta frame +/// (see `HeapDeltaWire`/`MatrixDeltaWire` in +/// `count_min_sketch_with_heap.rs`), decoded here directly +/// with `rmp_serde` since there is no CountSketchWithHeap ingest +/// accumulator to delegate to. No `asap_sketchlib` delta API needed — the +/// public `from_legacy_matrix` rebuilds both the matrix and heap. +pub fn decode_cs_with_heap_from_msgpack_delta( + buffer: &[u8], +) -> Result { + #[derive(serde::Deserialize)] + struct HeapDeltaWire { + is_delta: bool, + matrix_delta: MatrixDeltaWire, + topk_heap: Vec<(String, f64)>, + heap_size: u64, + } + #[derive(serde::Deserialize)] + struct MatrixDeltaWire { + rows: u32, + cols: u32, + cells: Vec<(u32, u32, i64)>, + } + + let wire: HeapDeltaWire = rmp_serde::from_slice(buffer) + .map_err(|e| format!("decode CountSketchWithHeap delta msgpack: {e}"))?; + if !wire.is_delta { + return Err("CountSketchWithHeap delta frame has is_delta=false".to_string()); + } + let rows = wire.matrix_delta.rows as usize; + let cols = wire.matrix_delta.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountSketchWithHeap delta frame has zero dims (rows={rows}, cols={cols})" + )); + } + let mut matrix = vec![vec![0.0; cols]; rows]; + for (r, c, dc) in &wire.matrix_delta.cells { + let (r, c) = (*r as usize, *c as usize); + if r >= rows || c >= cols { + continue; + } + matrix[r][c] += *dc as f64; + } + let heap: Vec = wire + .topk_heap + .into_iter() + .map(|(key, value)| CsHeapItem { key, value }) + .collect(); + Ok(CountSketchWithHeap::from_legacy_matrix( + matrix, + heap, + rows, + cols, + wire.heap_size as usize, + )) +} diff --git a/crates/asap_summary_state/src/stored_state/delta_apply.rs b/crates/asap_summary_state/src/stored_state/delta_apply.rs new file mode 100644 index 000000000..3f638503c --- /dev/null +++ b/crates/asap_summary_state/src/stored_state/delta_apply.rs @@ -0,0 +1,1276 @@ +//! Shared sketch state reconstruction and decoding. +use asap_sketchlib::CountMinSketch; +use asap_sketchlib::CountMinSketchWithHeap; +use asap_sketchlib::CountSketch; +use asap_sketchlib::CountSketchWithHeap; +use asap_sketchlib::DdSketch; +use asap_sketchlib::HllSketch; +use asap_sketchlib::HllVariant; +use asap_sketchlib::KllSketch; +use asap_sketchlib::MessagePackCodec; + +use super::decoders::{ + decode_cms_from_msgpack, decode_cms_from_proto, decode_cms_from_proto_delta, + decode_cms_with_heap_from_msgpack, decode_cms_with_heap_from_msgpack_delta, + decode_cs_from_msgpack, decode_cs_from_proto, decode_cs_from_proto_delta, + decode_cs_with_heap_from_msgpack, decode_cs_with_heap_from_msgpack_delta, +}; +use super::{SketchEncoding, SketchSampleState}; + +/// Which sketch family a candidate is, and the parameters needed to +/// *bootstrap an empty state* — required by the per-window-reset (PWR) +/// delta model where a window's FIRST frame is a delta-from-empty (no +/// carry-in Full). Most families' deltas embed their own params in the +/// wire fragment (decoded independently, then merged in — see +/// `SummaryState::apply_delta_bytes`); HLL register deltas and DD's +/// bucket-index deltas are applied onto a pre-sized structure instead, +/// so those two need the params known up front to allocate it. +#[derive(Debug, Clone, Copy)] +pub enum DeltaSketchKind { + UnivMon { + heap_size: u32, + sketch_rows: u32, + sketch_cols: u32, + layers: u8, + }, + DDSketch { + alpha: f64, + }, + Hll { + precision: u32, + }, + Kll { + k: u32, + }, + Cms { + rows: usize, + cols: usize, + }, + CountSketch { + rows: usize, + cols: usize, + }, + /// `CmsWithHeap` wraps `asap_sketchlib::CountMinSketchWithHeap` + /// (min-over-rows estimator) and `CountSketchWithHeap` wraps the + /// distinct `asap_sketchlib::CountSketchWithHeap` (median-of-signed-rows + /// estimator) -- different algorithms that happen to share a storage + /// shape. Kept as two variants (not one shared `Heap`) so + /// `merge_same_family` rejects merging one into the other the same + /// way it already rejects e.g. merging a `Cms` into a `Kll`; now the + /// type system enforces it too, since the two variants hold different + /// Rust types. + CmsWithHeap { + rows: usize, + cols: usize, + heap_size: usize, + }, + CountSketchWithHeap { + rows: usize, + cols: usize, + heap_size: usize, + }, +} + +impl DeltaSketchKind { + /// Construct an EMPTY state for this kind, used to seed a new window + /// when its first frame is a delta-from-empty (PWR). A delta applied + /// onto this empty base reconstructs exactly that window's state + /// (delta-from-empty ⊕ empty = window state). + fn bootstrap_empty(&self) -> SummaryState { + match self { + Self::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + } => SummaryState::UnivMon( + crate::summary_kernels::univmon::UnivMonAccumulator::new( + *heap_size as usize, + *sketch_rows as usize, + *sketch_cols as usize, + *layers as usize, + ) + .expect("validated UnivMon catalog dimensions"), + ), + DeltaSketchKind::DDSketch { alpha } => SummaryState::Dd(DdSketch::new(*alpha)), + DeltaSketchKind::Kll { k } => SummaryState::Kll(KllSketch::new(*k as u16)), + DeltaSketchKind::Hll { precision } => { + SummaryState::Hll(HllSketch::new(HllVariant::Regular, *precision)) + } + DeltaSketchKind::Cms { rows, cols } => { + SummaryState::Cms(CountMinSketch::new(*rows, *cols)) + } + DeltaSketchKind::CountSketch { rows, cols } => { + SummaryState::CountSketch(CountSketch::new(*rows, *cols)) + } + DeltaSketchKind::CmsWithHeap { + rows, + cols, + heap_size, + } => SummaryState::CmsWithHeap(CountMinSketchWithHeap::new(*rows, *cols, *heap_size)), + DeltaSketchKind::CountSketchWithHeap { + rows, + cols, + heap_size, + } => SummaryState::CountSketchWithHeap(CountSketchWithHeap::new( + *rows, *cols, *heap_size, + )), + } + } +} + +/// Try to decode a "full" sketch from the bytes (used by both +/// per-window and cumulative modes when the encoding is `*Full`). +fn decode_full( + kind: &DeltaSketchKind, + bytes: &[u8], + encoding: SketchEncoding, +) -> Result { + match (kind, encoding) { + ( + DeltaSketchKind::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + }, + SketchEncoding::MsgpackFull, + ) => { + let state = crate::summary_kernels::univmon::UnivMonAccumulator::from_bytes(bytes) + .map_err(|e| e.to_string())?; + if state.dimensions() + != ( + *heap_size as usize, + *sketch_rows as usize, + *sketch_cols as usize, + *layers as usize, + ) + { + return Err("UnivMon payload dimensions differ from installed catalog".into()); + } + Ok(SummaryState::UnivMon(state)) + } + (DeltaSketchKind::DDSketch { .. }, SketchEncoding::ProtoFull) => { + let sk = dd_from_proto(bytes)?; + Ok(SummaryState::Dd(sk)) + } + (DeltaSketchKind::DDSketch { .. }, SketchEncoding::MsgpackFull) => { + let sk = DdSketch::from_msgpack(bytes) + .map_err(|e| format!("deserialize DDSketch msgpack: {e}"))?; + Ok(SummaryState::Dd(sk)) + } + (DeltaSketchKind::Hll { .. }, SketchEncoding::ProtoFull) => { + let sk = hll_from_proto(bytes)?; + Ok(SummaryState::Hll(sk)) + } + (DeltaSketchKind::Hll { .. }, SketchEncoding::MsgpackFull) => { + let sk = HllSketch::from_msgpack(bytes) + .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?; + Ok(SummaryState::Hll(sk)) + } + (DeltaSketchKind::Kll { .. }, SketchEncoding::ProtoFull) => { + let sk = kll_from_proto(bytes)?; + Ok(SummaryState::Kll(sk)) + } + (DeltaSketchKind::Kll { .. }, SketchEncoding::MsgpackFull) => { + let sk = KllSketch::from_msgpack(bytes) + .map_err(|e| format!("deserialize KllSketch msgpack: {e}"))?; + Ok(SummaryState::Kll(sk)) + } + (DeltaSketchKind::Cms { .. }, SketchEncoding::ProtoFull) => { + Ok(SummaryState::Cms(decode_cms_from_proto(bytes)?)) + } + (DeltaSketchKind::Cms { .. }, SketchEncoding::MsgpackFull) => { + Ok(SummaryState::Cms(decode_cms_from_msgpack(bytes)?)) + } + (DeltaSketchKind::CountSketch { .. }, SketchEncoding::ProtoFull) => { + Ok(SummaryState::CountSketch(decode_cs_from_proto(bytes)?)) + } + (DeltaSketchKind::CountSketch { .. }, SketchEncoding::MsgpackFull) => { + Ok(SummaryState::CountSketch(decode_cs_from_msgpack(bytes)?)) + } + // The heap-bearing wire format is msgpack-only in this + // deployment; `decode_cms_with_heap_from_msgpack` is the same + // "Full" decoder the reducer's existing per-frame dispatch falls + // through to for any non-MsgpackDelta encoding. + ( + DeltaSketchKind::CmsWithHeap { .. }, + SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull, + ) => Ok(SummaryState::CmsWithHeap( + decode_cms_with_heap_from_msgpack(bytes)?, + )), + ( + DeltaSketchKind::CountSketchWithHeap { .. }, + SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull, + ) => Ok(SummaryState::CountSketchWithHeap( + decode_cs_with_heap_from_msgpack(bytes)?, + )), + (_, e) => Err(format!("decode_full called with non-Full encoding {e:?}")), + } +} + +/// The reconstructed state one candidate sid contributes — either +/// folded across a window (or several) via delta application, or merged +/// in from another sid's own reconstruction. +pub enum SummaryState { + UnivMon(crate::summary_kernels::univmon::UnivMonAccumulator), + Dd(DdSketch), + Hll(HllSketch), + Kll(KllSketch), + Cms(CountMinSketch), + CountSketch(CountSketch), + /// See `DeltaSketchKind::CmsWithHeap`/`CountSketchWithHeap` for why + /// these are two variants holding two different sketchlib types. + CmsWithHeap(CountMinSketchWithHeap), + CountSketchWithHeap(CountSketchWithHeap), +} + +impl SummaryState { + /// Apply a delta-encoded payload from a window sample. For DD / KLL, + /// the delta is interpreted as a "mergeable fragment" decoded + /// through the same full-state decoder and merged into the + /// rolling state. For HLL, the wire delta is a sparse register + /// update applied via the sketch's `apply_delta`. + /// + /// On encoding mismatch (e.g. trying to apply an HllDelta to a + /// DDSketch rolling state) returns Err. + pub fn apply_delta_bytes( + &mut self, + bytes: &[u8], + encoding: SketchEncoding, + ) -> Result<(), String> { + if !matches!( + encoding, + SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta + ) { + return Err(format!( + "apply_delta_bytes called with non-Delta encoding {encoding:?}" + )); + } + match self { + SummaryState::UnivMon(_) => Err("UnivMon requires full pane snapshots".into()), + SummaryState::Dd(sk) => { + match encoding { + // PROTO_DELTA: dispatch on the payload SHAPE, mirroring the + // supported DDSketch frame decoder, which tries the + // full-envelope decode first, then falls back to + // the bucket-delta proto. Two wire shapes can arrive on the + // ProtoDelta channel: + // + // 1. `SketchEnvelope{DdSketchState}` — a full-state + // fragment, mergeable via `DdSketch::merge`. (The edge + // sends this when `compute_delta_against` hits the + // empty-current / undecodable-prior fallback and ships + // a full snapshot tagged as a delta.) + // 2. `DDSketchDelta { buckets: [{index, d_count}] }` — a + // bucket-index delta proto, applied additively. This is + // the COMMON delta_transmission frame the edge emits + // under per-window-reset (`compute_delta(&empty)`). + // + // Before this fix the reducer decoded ONLY shape (1) via + // `decode_full`. A real shape-(2) frame failed with a wire- + // type mismatch on field 1 (delta field 1 = repeated + // submessage; state field 1 = `double alpha`) → the whole + // `quantile_over_time` returned `No result` for every + // delta_transmission DDSketch stream. We wrap the rolling + // `DdSketch` in a transient accumulator so the bucket-delta + // apply lands on `sk` in place. + SketchEncoding::ProtoDelta => { + // Shape (1): full envelope fragment → merge. Try this + // first (cheap decode attempt; a bucket-delta proto + // fails it on the field-1 wire-type mismatch). + if let Ok(SummaryState::Dd(other)) = decode_full( + &DeltaSketchKind::DDSketch { alpha: 0.0 }, + bytes, + SketchEncoding::ProtoFull, + ) { + sk.merge(&other) + .map_err(|e| format!("merge DDSketch delta envelope: {e}"))?; + return Ok(()); + } + // Shape (2): bucket-delta proto → additive apply via the + // SAME decoder the ingest delta path uses. + use crate::summary_kernels::dd_sketch::DDSketchAccumulator; + let mut acc = DDSketchAccumulator { + inner: std::mem::replace(sk, DdSketch::new(sk.alpha)), + sample_p: 1.0, + }; + let res = acc.apply_proto_delta_bytes(bytes); + *sk = acc.inner; + res.map_err(|e| format!("apply DDSketch proto bucket-delta: {e}"))?; + Ok(()) + } + // MSGPACK_DELTA: a serialized full-sketch fragment, mergeable + // via the full-state decoder. Kept for completeness — the + // edge wires PROTO_DELTA for DDSketch today. + SketchEncoding::MsgpackDelta => { + let other = match decode_full( + &DeltaSketchKind::DDSketch { alpha: 0.0 }, + bytes, + SketchEncoding::MsgpackFull, + ) { + Ok(SummaryState::Dd(s)) => s, + Ok(_) => { + return Err( + "decode_full(DDSketch) returned non-DDSketch state".to_string() + ) + } + Err(e) => return Err(e), + }; + sk.merge(&other) + .map_err(|e| format!("merge DDSketch delta: {e}"))?; + Ok(()) + } + _ => unreachable!(), + } + } + SummaryState::Hll(sk) => { + // HLL has a true sparse register delta in the proto + // wire format. Use the same path the precompute + // accumulator uses (`apply_proto_delta_bytes`-style). + if encoding == SketchEncoding::ProtoDelta { + apply_hll_proto_delta(sk, bytes) + } else { + // MsgpackDelta for HLL isn't a sparse encoding; + // it's a serialized HllSketch fragment, mergeable + // via `HllSketch::merge`. + let other = HllSketch::from_msgpack(bytes) + .map_err(|e| format!("deserialize HllSketch (delta-as-msgpack): {e}"))?; + sk.merge(&other) + .map_err(|e| format!("merge HLL delta: {e}"))?; + Ok(()) + } + } + SummaryState::Kll(sk) => { + let full_enc = match encoding { + SketchEncoding::ProtoDelta => SketchEncoding::ProtoFull, + SketchEncoding::MsgpackDelta => SketchEncoding::MsgpackFull, + _ => unreachable!(), + }; + let other = match decode_full(&DeltaSketchKind::Kll { k: 0 }, bytes, full_enc) { + Ok(SummaryState::Kll(s)) => s, + Ok(_) => return Err("decode_full(Kll) returned non-Kll state".to_string()), + Err(e) => return Err(e), + }; + sk.merge(&other) + .map_err(|e| format!("merge KLL delta: {e}"))?; + Ok(()) + } + // CMS/CountSketch/Heap have no true sparse in-place delta + // (unlike DD's bucket-index proto or HLL's register proto, + // above) — every delta frame already decodes into a + // complete, standalone state on its own (the PWR wire + // contract resets to empty at the source), so applying one + // is always "decode independently, then merge". + SummaryState::Cms(sk) => { + if encoding != SketchEncoding::ProtoDelta { + return Err( + "CountMin (heap-less) MSGPACK_DELTA is not a valid producer encoding \ + (msgpack-delta is the heap-bearing form)" + .to_string(), + ); + } + let other = decode_cms_from_proto_delta(bytes)?; + sk.merge(&other) + .map_err(|e| format!("merge CountMinSketch delta: {e}")) + } + SummaryState::CountSketch(sk) => { + if encoding != SketchEncoding::ProtoDelta { + return Err( + "CountSketch (heap-less) MSGPACK_DELTA is not a valid producer encoding \ + (msgpack-delta is the heap-bearing form)" + .to_string(), + ); + } + let other = decode_cs_from_proto_delta(bytes)?; + sk.merge(&other) + .map_err(|e| format!("merge CountSketch delta: {e}")) + } + SummaryState::CmsWithHeap(sk) => { + // Matches the existing per-frame reducer dispatch: only + // MsgpackDelta gets true delta treatment; ProtoDelta (not + // produced for this family in this deployment) falls + // through to the full-msgpack decoder, same as `decode_full`. + let other = if encoding == SketchEncoding::MsgpackDelta { + decode_cms_with_heap_from_msgpack_delta(bytes)? + } else { + decode_cms_with_heap_from_msgpack(bytes)? + }; + sk.merge(&other) + .map_err(|e| format!("merge CmsWithHeap delta: {e}")) + } + SummaryState::CountSketchWithHeap(sk) => { + let other = if encoding == SketchEncoding::MsgpackDelta { + decode_cs_with_heap_from_msgpack_delta(bytes)? + } else { + decode_cs_with_heap_from_msgpack(bytes)? + }; + sk.merge(&other) + .map_err(|e| format!("merge CountSketchWithHeap delta: {e}")) + } + } + } + + pub fn quantile(&self, q: f64) -> f64 { + match self { + SummaryState::Dd(sk) => sk.quantile(q).unwrap_or(0.0), + SummaryState::Kll(sk) => sk.quantile(q), + _ => 0.0, + } + } + + pub fn cardinality(&self) -> f64 { + match self { + SummaryState::Hll(sk) => sk.estimate(), + _ => 0.0, + } + } + + /// The bucket TOTAL — sum of row 0 of the underlying matrix. What a + /// bare `count_over_time`/`sum by (item) (rate(...))`-shaped query + /// (no specific item key) reads out. `0.0` for non-Frequency-family + /// states. + pub fn total(&self) -> f64 { + let matrix = match self { + SummaryState::Cms(c) => c.sketch(), + SummaryState::CountSketch(c) => c.sketch().clone(), + SummaryState::CmsWithHeap(h) => h.sketch_matrix(), + SummaryState::CountSketchWithHeap(h) => h.sketch_matrix(), + _ => return 0.0, + }; + matrix + .first() + .map(|row| row.iter().copied().sum::()) + .unwrap_or(0.0) + } + + /// Per-key point estimate — `count(metric{item="x"})`-shaped queries. + /// Unlike [`Self::topk_items`], no heap is needed: all four Frequency + /// variants (heap-bearing or not) already carry a keyed `estimate` + /// over their matrix. `None` for the quantile/cardinality states, + /// which have no item universe at all. + pub fn estimate(&self, key: &str) -> Option { + match self { + SummaryState::Cms(c) => Some(c.estimate(key)), + SummaryState::CountSketch(c) => Some(c.estimate(key)), + SummaryState::CmsWithHeap(h) => Some(h.estimate(key)), + SummaryState::CountSketchWithHeap(h) => Some(h.estimate(key)), + _ => None, + } + } + + /// Top-k `(key, value)` pairs from the heap, descending by value. + /// `None` for anything other than a heap-bearing state — the + /// heap-less Frequency states (`Cms`/`CountSketch`) carry no item + /// universe to enumerate, and the quantile/cardinality states have + /// no heap at all. + pub fn topk_items(&self) -> Option> { + match self { + SummaryState::CmsWithHeap(h) => Some( + h.topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(), + ), + SummaryState::CountSketchWithHeap(h) => Some( + h.topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(), + ), + _ => None, + } + } + + /// Merge `other` into `self` in place — both must be the same sketch + /// family. Used to combine several sids' reconstructed states + /// (`cumulative_summary_state`/`per_window_summary_states`) into one + /// cross-sid answer. `CmsWithHeap`/`CountSketchWithHeap` fall through + /// to the catch-all mismatch arm below like any other mixed pair — + /// and since the two variants now hold distinct sketchlib types + /// (`CountMinSketchWithHeap` vs `CountSketchWithHeap`), there is no + /// arm that could accidentally match them together — see their doc + /// on `DeltaSketchKind`. + pub fn merge_same_family(&mut self, other: &SummaryState) -> Result<(), String> { + match (self, other) { + (SummaryState::UnivMon(a), SummaryState::UnivMon(b)) => { + a.merge_in_place(b).map_err(|e| e.to_string()) + } + (SummaryState::Dd(a), SummaryState::Dd(b)) => { + a.merge(b).map_err(|e| format!("merge DDSketch: {e}")) + } + (SummaryState::Hll(a), SummaryState::Hll(b)) => { + a.merge(b).map_err(|e| format!("merge HLL: {e}")) + } + (SummaryState::Kll(a), SummaryState::Kll(b)) => { + a.merge(b).map_err(|e| format!("merge KLL: {e}")) + } + (SummaryState::Cms(a), SummaryState::Cms(b)) => { + a.merge(b).map_err(|e| format!("merge CountMinSketch: {e}")) + } + (SummaryState::CountSketch(a), SummaryState::CountSketch(b)) => { + a.merge(b).map_err(|e| format!("merge CountSketch: {e}")) + } + (SummaryState::CmsWithHeap(a), SummaryState::CmsWithHeap(b)) => { + a.merge(b).map_err(|e| format!("merge CmsWithHeap: {e}")) + } + (SummaryState::CountSketchWithHeap(a), SummaryState::CountSketchWithHeap(b)) => a + .merge(b) + .map_err(|e| format!("merge CountSketchWithHeap: {e}")), + (a, _) => Err(format!( + "SummaryState family mismatch in merge_same_family (self is {})", + a.family_name() + )), + } + } + + /// Diagnostic family name for error messages — not used for dispatch. + fn family_name(&self) -> &'static str { + match self { + SummaryState::UnivMon(_) => "UnivMon", + SummaryState::Dd(_) => "DDSketch", + SummaryState::Hll(_) => "Hll", + SummaryState::Kll(_) => "Kll", + SummaryState::Cms(_) => "Cms", + SummaryState::CountSketch(_) => "CountSketch", + SummaryState::CmsWithHeap(_) => "CmsWithHeap", + SummaryState::CountSketchWithHeap(_) => "CountSketchWithHeap", + } + } +} + +/// Fold every in-range window's frames for ONE series into a single +/// merged `SummaryState` (cumulative over `[t0, t1]`), returning `None` +/// if no Full frame ever landed (every sample was a leading delta). The +/// per-sid building block for a cross-sid answer: reconstruct each +/// candidate sid's state this way, then merge them (`merge_same_family`) +/// before reading out a quantile/cardinality over the combined data. +pub fn cumulative_summary_state( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, +) -> Result, String> { + let mut rolling: Option = None; + visit_window_summary_states(samples, kind, |_, state| { + if let Some(acc) = rolling.as_mut() { + acc.merge_same_family(&state)?; + } else { + rolling = Some(state); + } + Ok(()) + })?; + Ok(rolling) +} + +#[cfg(test)] +/// Walk a sorted-by-window-end slice of samples in time order and +/// produce ONE per-window scalar `(window_end_ms, scalar)`. +/// +/// ## Per-window-reset (PWR) delta model +/// +/// The edge emits frames grouped by window (all frames of one window +/// share the same `window_end` key; the key changes across windows). +/// The edge RESETS its snapshot base at each window boundary, so each +/// window's state is built *from empty*: +/// +/// * Within a window, frames accumulate to the window total. The first +/// frame may be a `Full` (window 1, or a periodic re-snapshot) or a +/// `Delta`-from-empty (windows 2+ under PWR); subsequent frames are +/// `Delta` INCREMENTS applied onto the window's running base. +/// * Across windows, the base MUST reset — a new `window_end` discards +/// the previous window's rolling state and starts from empty. Never +/// carry one window's state into the next (that would inflate via +/// cross-window accumulation). +/// +/// Concretely this fixes two bugs in the old "single rolling Option that +/// only ever resets on a Full" walk: +/// 1. A query range whose Full lives only in window 1 (or out of +/// range) left windows 2+ as deltas with `rolling=None`, all +/// skipped → empty result. +/// 2. A window 2+ delta applied onto window 1's leftover rolling state +/// → cross-window inflation. +/// +/// For a `Delta` that is the window's FIRST frame (the PWR delta-from- +/// empty case), we bootstrap an EMPTY rolling state of `kind` and apply +/// the delta onto it (delta-from-empty ⊕ empty = that window's state). +/// +/// The delta-OFF path (exactly one `Full` per window) still produces one +/// correct value per window: the window opens with a Full, has no +/// further frames, and emits that Full's scalar. +/// +/// `eval` reads a scalar from the rolling state (`quantile(q)` / +/// `cardinality()`). `skipped` counts frames that could not contribute +/// (a delta we genuinely couldn't bootstrap from — should be rare). +/// +/// Returns `Ok(per_window_samples, skipped)`. +pub fn per_window_evaluate( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, + eval: E, +) -> Result<(Vec<(i64, f64)>, usize), String> +where + E: Fn(&SummaryState) -> f64, +{ + let (states, skipped) = per_window_summary_states(samples, kind)?; + Ok(( + states.into_iter().map(|(w, rs)| (w, eval(&rs))).collect(), + skipped, + )) +} + +/// Walk a sorted-by-window-end slice of samples in time order and +/// reconstruct ONE sid's per-window `SummaryState` (same per-window-reset +/// walk as [`per_window_evaluate`], generalized to return the +/// reconstructed state itself instead of an already-evaluated scalar). +/// The per-sid building block for cross-sid per-window merging (unlike +/// [`cumulative_summary_state`], which folds a whole `[t0, t1]` range +/// into one answer, this keeps each window separate so a caller can +/// merge same-window states across several sids before evaluating -- +/// needed for a matrix/range-query answer, where each output point is +/// itself a cross-sid merge for that one window). +/// +/// Returns `Ok((per_window_states, skipped))`. +pub fn per_window_summary_states( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, +) -> Result<(Vec<(i64, SummaryState)>, usize), String> { + let mut out: Vec<(i64, SummaryState)> = Vec::new(); + let skipped = visit_window_summary_states(samples, kind, |end, state| { + out.push((end, state)); + Ok(()) + })?; + Ok((out, skipped)) +} + +// Both readout modes must reconstruct the same final pane population. The +// visitor lets cumulative merging stream panes without retaining every state. +fn visit_window_summary_states( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, + mut emit: impl FnMut(i64, SummaryState) -> Result<(), String>, +) -> Result { + let mut skipped = 0usize; + + // Rolling state for the CURRENT window only. Reset to None whenever + // `window_end` changes (a new window establishes its own base from + // empty). `cur_end` tracks which window `rolling` belongs to. + let mut rolling: Option = None; + let mut cur_end: Option = None; + + for (window_end, state) in samples { + // Window boundary: flush the previous window's final accumulated + // state, then reset the base so this window starts from empty. + if cur_end != Some(*window_end) { + if let (Some(prev_end), Some(rs)) = (cur_end, rolling.take()) { + emit(prev_end, rs)?; + } + cur_end = Some(*window_end); + } + + match state.encoding { + SketchEncoding::NativeBatchV1 => { + return Err("native physical outputs require the bound native batch decoder".into()) + } + SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull => { + // A Full (re)sets this window's base. + rolling = Some(decode_full(&kind, &state.bytes, state.encoding)?); + } + SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta => { + // Apply onto this window's running base. If this is the + // window's first frame (PWR delta-from-empty), bootstrap + // an empty base and apply onto it. + if rolling.is_none() { + rolling = Some(kind.bootstrap_empty()); + } + match rolling.as_mut() { + Some(rs) => rs.apply_delta_bytes(&state.bytes, state.encoding)?, + None => skipped += 1, + } + } + } + } + + // Flush the final window. + if let (Some(prev_end), Some(rs)) = (cur_end, rolling.take()) { + emit(prev_end, rs)?; + } + + Ok(skipped) +} + +// --------------------------------------------------------------------------- +// Proto-envelope decoders — P2-4: ONE decoder per family. +// +// These delegate to the precompute-side accumulators' +// `from_sketchlib_proto_bytes`, which are the single source of truth for +// the modified-OTLP proto wire format (envelope unwrapping, alpha/k/ +// precision validation, and — critically for HLL — SPARSE +// `registers_sparse` expansion). Folding the warm read path onto the +// same decoder the ingest path uses means the sparse-register fix (and +// any future format change) can never drift between the two copies again +// — the bug class P2-3 / P2-4 closed. We extract the accumulator's +// public `inner` sketch for the rolling-state merge. +// --------------------------------------------------------------------------- + +fn dd_from_proto(buffer: &[u8]) -> Result { + use crate::summary_kernels::dd_sketch::DDSketchAccumulator; + DDSketchAccumulator::from_sketchlib_proto_bytes(buffer) + .map(|acc| acc.inner) + .map_err(|e| e.to_string()) +} + +fn kll_from_proto(buffer: &[u8]) -> Result { + use crate::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(buffer) + .map(|acc| acc.inner) + .map_err(|e| e.to_string()) +} + +fn hll_from_proto(buffer: &[u8]) -> Result { + use crate::summary_kernels::hll_sketch::HllSketchAccumulator; + HllSketchAccumulator::from_sketchlib_proto_bytes(buffer) + .map(|acc| acc.inner) + .map_err(|e| e.to_string()) +} + +/// Apply a proto-encoded `HllDelta` frame onto the HLL register vector — the +/// delta is a varint-packed (index_delta, value) blob; decode + apply +/// (register-wise max) via the shared sketch library so the unpacking stays a +/// single source of truth. +fn apply_hll_proto_delta(sk: &mut HllSketch, buffer: &[u8]) -> Result<(), String> { + sk.apply_delta_bytes(buffer) + .map_err(|e| format!("apply HLLDelta: {e}"))?; + Ok(()) +} + +#[cfg(test)] +mod tests { + //! P2-3 / P2-4 regression tests for the consolidated single-decoder + //! path. These exercise the family proto decoders that now delegate + //! to the precompute accumulators (the single source of truth), so a + //! divergence between the warm read path and the ingest path — + //! notably the SPARSE-register HLL handling the deleted dead decoder + //! got wrong — fails the build. + use super::*; + use asap_sketchlib::HllVariant; + + fn encode_dd(sk: &DdSketch) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + let state = DdSketchState { + alpha: sk.alpha, + store_counts: sk.store_counts.clone(), + store_offset: sk.store_offset, + ..Default::default() + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), + ..Default::default() + } + .encode_to_vec() + } + + fn encode_kll(k: u16, items: &[f64]) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + let state = KllState { + k: k as u32, + items: items.to_vec(), + levels: vec![], + num_levels: 0, + ..Default::default() + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + } + .encode_to_vec() + } + + fn encode_hll_dense(sk: &HllSketch) -> Vec { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision: sk.precision, + registers: sk.registers.clone(), + hip_kxq0: sk.hip_kxq0, + hip_kxq1: sk.hip_kxq1, + hip_est: sk.hip_est, + registers_sparse: None, + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + } + .encode_to_vec() + } + + /// Build a SPARSE HLL proto frame: dense `registers` left empty, + /// `registers_sparse.packed` = varint (index_delta, value) pairs. + /// This is exactly the wire form a low-cardinality producer emits + /// (sketchlib-go below its dense/sparse crossover) — the frame the + /// DELETED `HllSketch_from_sketchlib_proto_bytes` hard-rejected with + /// "registers has 0 bytes". + fn encode_hll_sparse(precision: u32, nonzero: &[(u64, u8)]) -> Vec { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllSparseRegisters, HllVariant as ProtoVariant, HyperLogLogState, + SketchEnvelope, + }; + use prost::Message; + // Varint-pack (index_delta, value), ascending index order. + let mut packed: Vec = Vec::new(); + let mut prev: u64 = 0; + let put_uvarint = |buf: &mut Vec, mut v: u64| loop { + let b = (v & 0x7f) as u8; + v >>= 7; + if v != 0 { + buf.push(b | 0x80); + } else { + buf.push(b); + break; + } + }; + let mut sorted = nonzero.to_vec(); + sorted.sort_by_key(|(i, _)| *i); + for (idx, val) in &sorted { + put_uvarint(&mut packed, idx - prev); + put_uvarint(&mut packed, *val as u64); + prev = *idx; + } + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision, + registers: Vec::new(), // dense field empty → sparse path + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + // `num_registers` is informational — the decoder expands + // against `expected_len` from precision, not this field. + registers_sparse: Some(HllSparseRegisters { + num_registers: 1u32 << precision, + packed, + }), + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + } + .encode_to_vec() + } + + #[test] + fn hll_from_proto_accepts_sparse_frame() { + // The consolidated decoder must accept the sparse wire form (the + // deleted dead decoder rejected it). Build a sparse frame setting + // a handful of registers, decode it, and confirm those register + // slots came back set in the dense array. + let precision = 12u32; + let nonzero = [(3u64, 5u8), (100, 2), (4000, 7)]; + let bytes = encode_hll_sparse(precision, &nonzero); + let sk = hll_from_proto(&bytes).expect("sparse HLL frame must decode (P2-3 regression)"); + assert_eq!(sk.registers.len(), 1usize << precision); + for (idx, val) in nonzero { + assert_eq!( + sk.registers[idx as usize], val, + "sparse register {idx} expanded to wrong value" + ); + } + } + + #[test] + fn hll_from_proto_matches_accumulator_decoder() { + // P2-4: the warm read path and the ingest accumulator must decode + // the SAME bytes to the SAME sketch (one source of truth). + use crate::summary_kernels::hll_sketch::HllSketchAccumulator; + let mut sk = HllSketch::new(HllVariant::Regular, 12); + for i in 0..500u64 { + sk.update(format!("item-{i}").as_bytes()); + } + let bytes = encode_hll_dense(&sk); + let via_delta = hll_from_proto(&bytes).expect("delta_apply hll decode"); + let via_acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("accumulator hll decode") + .inner; + assert_eq!( + via_delta.registers, via_acc.registers, + "delta_apply and accumulator must produce identical HLL registers" + ); + assert!((via_delta.estimate() - via_acc.estimate()).abs() < 1e-9); + } + + #[test] + fn dd_from_proto_matches_accumulator_decoder() { + use crate::summary_kernels::dd_sketch::DDSketchAccumulator; + let mut sk = DdSketch::new(0.01); + for v in [1.0, 2.0, 5.0, 5.0, 9.0, 42.0] { + sk.update(v); + } + let bytes = encode_dd(&sk); + let via_delta = dd_from_proto(&bytes).expect("delta_apply dd decode"); + let via_acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("accumulator dd decode") + .inner; + // Same quantile answers from the same bytes through both paths. + assert_eq!(via_delta.quantile(0.5), via_acc.quantile(0.5)); + assert_eq!(via_delta.quantile(0.99), via_acc.quantile(0.99)); + } + + #[test] + fn kll_from_proto_matches_accumulator_decoder() { + use crate::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; + let items: Vec = (0..200).map(|i| i as f64).collect(); + let bytes = encode_kll(256, &items); + let via_delta = kll_from_proto(&bytes).expect("delta_apply kll decode"); + let via_acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("accumulator kll decode") + .inner; + assert_eq!(via_delta.quantile(0.5), via_acc.quantile(0.5)); + } + + // ----------------------------------------------------------------- + // Per-window-reset (PWR) delta-apply regression tests. + // + // The edge resets its snapshot base at every window boundary, so a + // window's first frame is either a Full (window 1 / re-snapshot) or + // a Delta-from-empty (windows 2+). The query-side walk must: + // * reset the rolling base when `window_end` changes, + // * bootstrap an empty base for a window's leading Delta, + // * emit ONE value per window (the window's final accumulated + // state), never per-frame and never cross-window-accumulated. + // ----------------------------------------------------------------- + + fn full(bytes: Vec) -> SketchSampleState { + SketchSampleState { + bytes, + encoding: SketchEncoding::ProtoFull, + } + } + fn delta(bytes: Vec) -> SketchSampleState { + SketchSampleState { + bytes, + encoding: SketchEncoding::ProtoDelta, + } + } + + fn dd_over(alpha: f64, vals: &[f64]) -> DdSketch { + let mut sk = DdSketch::new(alpha); + for &v in vals { + sk.update(v); + } + sk + } + + /// A full re-snapshot replaces its pane's earlier frames; cumulative + /// readout must merge the finalized panes without counting updates twice. + #[test] + fn cumulative_readout_counts_resnapshot_population_once() { + let first = full(encode_dd(&dd_over(0.01, &[1., 2.]))); + let updated = full(encode_dd(&dd_over(0.01, &[1., 2., 3.]))); + let next = delta(encode_dd(&dd_over(0.01, &[9.]))); + let samples = [(1000, &first), (1000, &updated), (2000, &next)]; + let state = cumulative_summary_state(&samples, DeltaSketchKind::DDSketch { alpha: 0.01 }) + .unwrap() + .unwrap(); + let SummaryState::Dd(state) = state else { + panic!("expected DDSketch state"); + }; + assert_eq!(state.store_counts.iter().sum::(), 4); + } + + /// PWR across 3 windows: window 1 is `[Full]`, windows 2 & 3 are + /// `[Delta-from-empty]` (NO Full carry-in). Each window must + /// reconstruct its OWN distribution's median — not empty (the old + /// "skip delta with no base" bug) and not cross-window-inflated. + #[test] + fn pwr_ddsketch_three_windows_delta_from_empty() { + let alpha = 0.01; + let w1 = dd_over(alpha, &[1.0, 2.0, 3.0, 4.0, 5.0]); + let w2 = dd_over(alpha, &[10.0, 20.0, 30.0, 40.0, 50.0]); + let w3 = dd_over(alpha, &[100.0, 200.0, 300.0, 400.0, 500.0]); + + // window 1 ships a Full; windows 2+ ship a delta-from-empty. + let s1 = full(encode_dd(&w1)); + let s2 = delta(encode_dd(&w2)); + let s3 = delta(encode_dd(&w3)); + let samples = vec![(1000_i64, &s1), (2000, &s2), (3000, &s3)]; + + let kind = DeltaSketchKind::DDSketch { alpha }; + let (out, skipped) = + per_window_evaluate(&samples, kind, |rs| rs.quantile(0.5)).expect("pwr eval"); + assert_eq!(skipped, 0, "PWR must not skip delta-from-empty frames"); + assert_eq!(out.len(), 3, "one value per window"); + + // Each window's median ≈ that window's own distribution median, + // independent of the others (no carry-in inflation). + let truth = [ + w1.quantile(0.5).unwrap(), + w2.quantile(0.5).unwrap(), + w3.quantile(0.5).unwrap(), + ]; + for (i, (w_end, est)) in out.iter().enumerate() { + assert_eq!(*w_end, (i as i64 + 1) * 1000); + let rel = (est - truth[i]).abs() / truth[i].max(1e-9); + assert!( + rel < 0.05, + "window {i}: est={est} truth={} rel={rel}", + truth[i] + ); + } + // Cross-window-inflation guard: window 2's median must NOT have + // absorbed window 1 (would pull it well below 30). + assert!( + out[1].1 > 20.0, + "window 2 median {} suggests cross-window accumulation", + out[1].1 + ); + } + + /// Sub-window producer: a SINGLE window carries multiple frames + /// `[Full, Delta, Delta]`, where each later delta is an increment + /// since the previous emit in that window. The walk must COLLAPSE + /// them to ONE value = the window's running total, not emit three. + #[test] + fn pwr_ddsketch_subwindow_frames_collapse_to_window_total() { + let alpha = 0.01; + // Three sub-window increments that together cover 1..=15. + let a = dd_over(alpha, &[1.0, 2.0, 3.0, 4.0, 5.0]); + let b = dd_over(alpha, &[6.0, 7.0, 8.0, 9.0, 10.0]); + let c = dd_over(alpha, &[11.0, 12.0, 13.0, 14.0, 15.0]); + let s_a = full(encode_dd(&a)); + let s_b = delta(encode_dd(&b)); + let s_c = delta(encode_dd(&c)); + // All three share the same window_end (one window, sub-window frames). + let samples = vec![(5000_i64, &s_a), (5000, &s_b), (5000, &s_c)]; + + let kind = DeltaSketchKind::DDSketch { alpha }; + let (out, skipped) = + per_window_evaluate(&samples, kind, |rs| rs.quantile(0.5)).expect("subwindow eval"); + assert_eq!(skipped, 0); + assert_eq!(out.len(), 1, "sub-window frames collapse to ONE value"); + assert_eq!(out[0].0, 5000); + + let truth = dd_over(alpha, &(1..=15).map(|v| v as f64).collect::>()) + .quantile(0.5) + .unwrap(); + let rel = (out[0].1 - truth).abs() / truth.max(1e-9); + assert!(rel < 0.05, "window total est={} truth={truth}", out[0].1); + } + + /// Same sub-window collapse, but the window's FIRST frame is a + /// Delta-from-empty (PWR window 2+ with sub-window frames): + /// `[Delta-from-empty, Delta, Delta]`. + #[test] + fn pwr_ddsketch_subwindow_first_frame_delta_from_empty() { + let alpha = 0.01; + let a = dd_over(alpha, &[1.0, 2.0, 3.0, 4.0, 5.0]); + let b = dd_over(alpha, &[6.0, 7.0, 8.0, 9.0, 10.0]); + let c = dd_over(alpha, &[11.0, 12.0, 13.0, 14.0, 15.0]); + let s_a = delta(encode_dd(&a)); // first frame is delta-from-empty + let s_b = delta(encode_dd(&b)); + let s_c = delta(encode_dd(&c)); + let samples = vec![(9000_i64, &s_a), (9000, &s_b), (9000, &s_c)]; + + let kind = DeltaSketchKind::DDSketch { alpha }; + let (out, skipped) = + per_window_evaluate(&samples, kind, |rs| rs.quantile(0.5)).expect("eval"); + assert_eq!(skipped, 0); + assert_eq!(out.len(), 1); + let truth = dd_over(alpha, &(1..=15).map(|v| v as f64).collect::>()) + .quantile(0.5) + .unwrap(); + let rel = (out[0].1 - truth).abs() / truth.max(1e-9); + assert!(rel < 0.05, "est={} truth={truth}", out[0].1); + } + + /// PWR for HLL across 3 windows, each a Delta-from-empty (sparse + /// register delta). Bootstrapping an EMPTY HLL of the right precision + /// is required (register deltas index into a pre-sized array). Each + /// window's cardinality must reflect its OWN item set. + #[test] + fn pwr_hll_three_windows_delta_from_empty() { + let precision = 12u32; + // Build per-window HLLs, then encode each as a register-delta + // against an EMPTY sketch (= that window's full register state, + // the PWR delta-from-empty wire form). + let empty = HllSketch::new(HllVariant::Regular, precision); + let mut frames = Vec::new(); + let truths = [200usize, 800, 1500]; + for (w, &n) in truths.iter().enumerate() { + let mut sk = HllSketch::new(HllVariant::Regular, precision); + let base = (w as u64) * 100_000; // disjoint item sets per window + for i in 0..n as u64 { + sk.update(format!("u-{}", base + i).as_bytes()); + } + let bytes = sk.compute_delta(&empty, 0); + frames.push((((w as u64) + 1) * 1000, delta(bytes))); + } + let samples: Vec<(i64, &SketchSampleState)> = + frames.iter().map(|(t, s)| (*t as i64, s)).collect(); + + let kind = DeltaSketchKind::Hll { precision }; + let (out, skipped) = + per_window_evaluate(&samples, kind, |rs| rs.cardinality()).expect("hll pwr eval"); + assert_eq!(skipped, 0, "HLL delta-from-empty must bootstrap, not skip"); + assert_eq!(out.len(), 3); + for (i, (_w_end, est)) in out.iter().enumerate() { + let n = truths[i] as f64; + let rel = (est - n).abs() / n; + assert!( + rel < 0.15, + "window {i}: HLL est={est} truth={n} rel={rel} (each window independent)" + ); + } + } + + /// `CmsWithHeap` (min-over-rows estimator, `CountMinSketchWithHeap`) + /// and `CountSketchWithHeap` (median-of-signed-rows estimator, the + /// distinct `CountSketchWithHeap` type) are different sketch + /// algorithms that merely happen to share a storage shape — merging + /// one into the other must be rejected as a family mismatch, the + /// same as merging a `Cms` into a `Kll` would be. Since the two + /// `SummaryState` variants now hold genuinely different Rust types, + /// this is also enforced at compile time — there is no arm in + /// `merge_same_family` that type-checks a mixed pair together. + #[test] + fn cms_with_heap_and_count_sketch_with_heap_are_not_the_same_family() { + use asap_sketchlib::{CountMinSketchWithHeap, CountSketchWithHeap, MessagePackCodec}; + + let mut cms_heap = CountMinSketchWithHeap::new(4, 256, 10); + cms_heap.update("a", 1.0); + let mut cs_heap = CountSketchWithHeap::new(4, 256, 10); + cs_heap.update("b", 1.0); + + let mut a = SummaryState::CmsWithHeap( + CountMinSketchWithHeap::from_msgpack(&cms_heap.to_msgpack().unwrap()).unwrap(), + ); + let b = SummaryState::CountSketchWithHeap( + CountSketchWithHeap::from_msgpack(&cs_heap.to_msgpack().unwrap()).unwrap(), + ); + + match a.merge_same_family(&b) { + Err(msg) => assert!( + msg.contains("family mismatch"), + "expected a family-mismatch error, got: {msg}" + ), + Ok(()) => panic!( + "CmsWithHeap must not merge with CountSketchWithHeap -- \ + different algorithms sharing only a storage shape" + ), + } + } + + fn encode_delta_heap( + rows: u32, + cols: u32, + cells: &[(u32, u32, i64)], + heap: &[(&str, f64)], + heap_size: u64, + ) -> Vec { + #[derive(serde::Serialize)] + struct W<'a>( + bool, + (u32, u32, &'a [(u32, u32, i64)]), + Vec<(String, f64)>, + u64, + ); + let heap_owned: Vec<(String, f64)> = + heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); + let w = W(true, (rows, cols, cells), heap_owned, heap_size); + rmp_serde::to_vec(&w).expect("encode delta-heap") + } + + /// `SummaryState::CountSketchWithHeap` must decode both FULL and + /// DELTA-HEAP msgpack frames through the genuine + /// `asap_sketchlib::CountSketchWithHeap` (median-of-signed-rows + /// estimator) rather than the CMS-family `CountMinSketchWithHeap` + /// (min-over-rows estimator) it used to alias — the bug this split + /// fixed. Built via real `update()` calls (not a hand-crafted matrix) + /// so the sign-hashed row semantics are genuinely exercised, then + /// checks both decode paths reproduce the same matrix and the same + /// `estimate()` as the in-memory sketch they were encoded from. + #[test] + fn count_sketch_with_heap_full_and_delta_decode_via_new_asap_sketchlib_type() { + use asap_sketchlib::{CountSketchWithHeap, MessagePackCodec}; + + let mut built = CountSketchWithHeap::new(4, 64, 10); + for _ in 0..50 { + built.update("k", 1.0); + } + let expected_matrix = built.sketch_matrix(); + let expected_estimate = built.estimate("k"); + + // FULL path. + let full_bytes = built.to_msgpack().expect("encode full CountSketchWithHeap"); + let full_state = decode_full( + &DeltaSketchKind::CountSketchWithHeap { + rows: 4, + cols: 64, + heap_size: 10, + }, + &full_bytes, + SketchEncoding::MsgpackFull, + ) + .expect("decode_full CountSketchWithHeap"); + match full_state { + SummaryState::CountSketchWithHeap(inner) => { + assert_eq!(inner.sketch_matrix(), expected_matrix); + assert_eq!(inner.estimate("k"), expected_estimate); + } + other => panic!( + "expected CountSketchWithHeap state, got {}", + other.family_name() + ), + } + + // DELTA-HEAP path: same cells + heap against an empty base (PWR + // contract), encoded the way the Go producer does. + let cells: Vec<(u32, u32, i64)> = expected_matrix + .iter() + .enumerate() + .flat_map(|(r, row)| { + row.iter().enumerate().filter_map(move |(c, v)| { + if *v != 0.0 { + Some((r as u32, c as u32, *v as i64)) + } else { + None + } + }) + }) + .collect(); + let heap_pairs: Vec<(String, f64)> = built + .topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(); + assert!(!heap_pairs.is_empty(), "expected \"k\" in the top-k heap"); + let heap_refs: Vec<(&str, f64)> = + heap_pairs.iter().map(|(k, v)| (k.as_str(), *v)).collect(); + let delta_bytes = encode_delta_heap(4, 64, &cells, &heap_refs, 10); + + let mut rolling = DeltaSketchKind::CountSketchWithHeap { + rows: 4, + cols: 64, + heap_size: 10, + } + .bootstrap_empty(); + rolling + .apply_delta_bytes(&delta_bytes, SketchEncoding::MsgpackDelta) + .expect("apply CountSketchWithHeap delta"); + match rolling { + SummaryState::CountSketchWithHeap(inner) => { + assert_eq!( + inner.sketch_matrix(), + expected_matrix, + "delta path must reconstruct the identical matrix" + ); + assert_eq!(inner.estimate("k"), expected_estimate); + } + other => panic!( + "expected CountSketchWithHeap state, got {}", + other.family_name() + ), + } + } +} diff --git a/crates/asap_summary_state/src/stored_state/mod.rs b/crates/asap_summary_state/src/stored_state/mod.rs new file mode 100644 index 000000000..89ab95786 --- /dev/null +++ b/crates/asap_summary_state/src/stored_state/mod.rs @@ -0,0 +1,22 @@ +//! Portable stored-summary payloads and reconstruction, independent of storage engines. +pub mod decoders; +pub mod delta_apply; +pub mod native; +pub mod readout; + +#[derive(Debug, Clone)] +pub struct SketchSampleState { + pub bytes: Vec, + /// Wire-encoding hint from the OTLP DataPoint's `encoding` field. + pub encoding: SketchEncoding, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SketchEncoding { + ProtoFull, + ProtoDelta, + MsgpackFull, + MsgpackDelta, + /// Versioned typed physical output; never a legacy sketch frame. + NativeBatchV1, +} diff --git a/crates/asap_summary_state/src/stored_state/native.rs b/crates/asap_summary_state/src/stored_state/native.rs new file mode 100644 index 000000000..d41133660 --- /dev/null +++ b/crates/asap_summary_state/src/stored_state/native.rs @@ -0,0 +1,409 @@ +//! Versioned physical output batches. Deployment identities and coverage remain +//! outside this payload and must be checked before decoding with the bound schema. +use crate::summary_kernels::{ + datasketches_kll::DatasketchesKLLAccumulator, dd_sketch::DDSketchAccumulator, + exact::ExactAccumulator, hll_sketch::HllSketchAccumulator, SumAccumulator, +}; +use asap_physical_operators::{ + summary_kernels as physical, + values::{Batch, Schema, Value}, + AggregateCore, Error, +}; +use planner_types::post_asap::{SummaryFamilyType, SummarySchema}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; + +#[derive(Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct StoredBatch { + version: u32, + schema: SummarySchema, + rows: Vec>, +} + +#[derive(Serialize, Deserialize)] +enum Cell { + Plain(Value), + Summary { + family: SummaryFamilyType, + codec: StateCodec, + bytes: Vec, + }, +} + +// Codec identity is distinct from algorithm identity: old integer CMS/CS bytes +// must never be interpreted as Float64 weighted state with typed item tuples. +#[derive(Serialize, Deserialize)] +enum StateCodec { + WeightedFrequencyV1, + ExactAccumulatorV1, + /// Read-only: decodes to Planner's exact Sum state. + SumAccumulatorV1, + KllMsgpackV1, + DdMsgpackV1, + HllMsgpackV1, +} +fn invalid(message: impl ToString) -> Error { + Error::Invalid(message.to_string()) +} +type Kernel = asap_sketchlib::WeightedFrequency; +impl StateCodec { + fn encode(state: &dyn AggregateCore) -> Result<(Self, Vec), Error> { + let any = state.as_any(); + if let Some(state) = any.downcast_ref::() { + return Ok(( + Self::WeightedFrequencyV1, + crate::physical::frequency_kernel(state) + .map_err(invalid)? + .to_bytes(), + )); + } + if let Some(state) = any.downcast_ref::() { + return Ok(( + Self::ExactAccumulatorV1, + rmp_serde::to_vec_named(state).map_err(invalid)?, + )); + } + let codec = if any.is::() { + Self::KllMsgpackV1 + } else if any.is::() { + Self::DdMsgpackV1 + } else if any.is::() { + Self::HllMsgpackV1 + } else { + return Err(invalid("physical summary has no persisted native codec")); + }; + let stored = crate::physical::from_physical(state).map_err(invalid)?; + Ok((codec, stored.serialize_to_bytes())) + } + fn decode(&self, bytes: &[u8]) -> Result, Error> { + let stored: Box = match self { + Self::WeightedFrequencyV1 => { + let kernel = Kernel::from_bytes(bytes).map_err(|e| invalid(format!("{e:?}")))?; + let state: physical::weighted_frequency::WeightedFrequency = + rmp_serde::from_slice(&rmp_serde::to_vec(&kernel).map_err(invalid)?) + .map_err(invalid)?; + return Ok(Arc::new(state)); + } + Self::ExactAccumulatorV1 => { + Box::new(ExactAccumulator::deserialize_from_bytes(bytes).map_err(invalid)?) + } + Self::SumAccumulatorV1 => { + Box::new(SumAccumulator::deserialize_from_bytes(bytes).map_err(invalid)?) + } + Self::KllMsgpackV1 => { + Box::new(DatasketchesKLLAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + } + Self::DdMsgpackV1 => { + Box::new(DDSketchAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + } + Self::HllMsgpackV1 => { + Box::new(HllSketchAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + } + }; + crate::physical::to_physical(stored.as_ref()).map_err(invalid) + } +} + +/// Encode a validated physical output, preserving Float64 and typed identities. +/// This format is independent of the logical and physical plan wire formats. +pub fn encode_batch(batch: &Batch) -> Result, Error> { + let rows = batch + .rows() + .iter() + .map(|row| { + row.iter() + .map(|value| { + Ok(match value { + Value::Summary { family, state } => { + let (codec, bytes) = StateCodec::encode(state.as_ref())?; + Cell::Summary { + family: family.clone(), + codec, + bytes, + } + } + value => Cell::Plain(value.clone()), + }) + }) + .collect::, Error>>() + }) + .collect::, Error>>()?; + rmp_serde::to_vec_named(&StoredBatch { + version: 1, + schema: batch.schema().as_ref().clone(), + rows, + }) + .map_err(invalid) +} + +/// Decode only against the installed output contract. The caller supplies its +/// per-read payload limit; checking state parameters is part of Batch validation. +pub fn decode_batch(bytes: &[u8], expected: Schema, max_bytes: usize) -> Result { + if bytes.len() > max_bytes { + return Err(invalid("native output payload exceeds read budget")); + } + let stored: StoredBatch = rmp_serde::from_slice(bytes).map_err(invalid)?; + if stored.version != 1 { + return Err(invalid("unsupported native output format")); + } + if stored.schema != *expected { + return Err(invalid( + "native output schema differs from installed contract", + )); + } + let rows = stored + .rows + .into_iter() + .map(|row| { + row.into_iter() + .map(|cell| { + Ok(match cell { + Cell::Plain(value) => value, + Cell::Summary { + family, + codec, + bytes, + } => Value::Summary { + family, + state: codec.decode(&bytes)?, + }, + }) + }) + .collect::, Error>>() + }) + .collect::, Error>>()?; + Batch::try_new(expected, rows) +} + +#[cfg(test)] +mod tests { + use super::*; + use physical::weighted_frequency::{FrequencyAlgorithm, WeightedFrequency}; + use planner_types::{ + post_asap::{SketchAlgorithm, SketchKind, SketchParams, SummaryField}, + pre_asap::DataType, + }; + + fn weighted(algorithm: SketchAlgorithm) -> Batch { + let (native, params) = match algorithm { + SketchAlgorithm::CmsWithHeap => ( + FrequencyAlgorithm::Cms, + SketchParams::CmsWithHeap { + width: 64, + depth: 5, + heap_size: 8, + }, + ), + _ => ( + FrequencyAlgorithm::CountSketch, + SketchParams::CountSketchWithHeap { + width: 64, + depth: 5, + heap_size: 8, + }, + ), + }; + let family = + SummaryFamilyType::Sketch(SketchKind::new(algorithm, params), Default::default()); + let schema = Arc::new(SummarySchema { + fields: vec![ + SummaryField { + name: "group".into(), + dtype: SummaryFamilyType::Plain(DataType::Utf8), + nullable: false, + }, + SummaryField { + name: "state".into(), + dtype: family.clone(), + nullable: false, + }, + ], + time_index: None, + }); + let mut state = WeightedFrequency::new(native, 64, 5, 8).unwrap(); + state + .update(&[Value::Int64(7), Value::Utf8("service-a".into())], 0.125) + .unwrap(); + state + .update( + &[Value::Utf8("7".into()), Value::Utf8("service-b".into())], + 0.25, + ) + .unwrap(); + Batch::try_new( + schema, + vec![vec![ + Value::Utf8("job-a".into()), + Value::Summary { + family, + state: Arc::new(state), + }, + ]], + ) + .unwrap() + } + + // Fractional rates and distinct typed item tuples survive both heap codecs. + #[test] + fn weighted_outputs_roundtrip_without_integer_conversion() { + for algorithm in [ + SketchAlgorithm::CmsWithHeap, + SketchAlgorithm::CountSketchWithHeap, + ] { + let batch = weighted(algorithm); + let bytes = encode_batch(&batch).unwrap(); + let restored = decode_batch(&bytes, batch.schema().clone(), bytes.len()).unwrap(); + let scores = |b: &Batch| { + let Value::Summary { state, .. } = &b.rows()[0][1] else { + panic!() + }; + state + .as_any() + .downcast_ref::() + .unwrap() + .rows(8) + .iter() + .map(|row| row.iter().map(|v| v.key().unwrap()).collect::>()) + .collect::>() + }; + assert_eq!(scores(&batch), scores(&restored)); + assert!(decode_batch(&bytes, batch.schema().clone(), bytes.len() - 1).is_err()); + } + } + + // A storage tag cannot send weighted physical output through an integer heap decoder. + #[test] + fn legacy_sketch_reader_rejects_native_batch_frames() { + let sample = super::super::SketchSampleState { + bytes: encode_batch(&weighted(SketchAlgorithm::CmsWithHeap)).unwrap(), + encoding: super::super::SketchEncoding::NativeBatchV1, + }; + let result = super::super::delta_apply::per_window_summary_states( + &[(60_000, &sample)], + super::super::delta_apply::DeltaSketchKind::CmsWithHeap { + rows: 5, + cols: 64, + heap_size: 8, + }, + ); + assert!(result.is_err()); + } + + // Every admitted native summary codec survives the same typed boundary. + #[test] + fn native_summary_families_and_nonfinite_plain_values_roundtrip() { + use planner_types::post_asap::{ExactKind, ExactParams}; + let exact = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + let sketch = |algorithm, params| { + SummaryFamilyType::Sketch(SketchKind::new(algorithm, params), Default::default()) + }; + let cases: Vec<(SummaryFamilyType, Arc)> = vec![ + ( + exact.clone(), + Arc::new(physical::exact::ExactAccumulator::new(exact.clone(), false).unwrap()), + ), + ( + sketch(SketchAlgorithm::Kll, SketchParams::Kll { k: 200 }), + Arc::new(physical::DatasketchesKLLAccumulator::new(200)), + ), + ( + sketch( + SketchAlgorithm::DDSketch, + SketchParams::DDSketch { alpha: 0.01 }, + ), + Arc::new(physical::DDSketchAccumulator::new(0.01)), + ), + ( + sketch(SketchAlgorithm::Hll, SketchParams::Hll { precision: 12 }), + Arc::new(physical::HllSketchAccumulator::new( + asap_sketchlib::HllVariant::Regular, + 12, + )), + ), + ]; + for (family, state) in cases { + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "state".into(), + dtype: family.clone(), + nullable: false, + }], + time_index: None, + }); + let batch = + Batch::try_new(schema.clone(), vec![vec![Value::Summary { family, state }]]) + .unwrap(); + let bytes = encode_batch(&batch).unwrap(); + let restored = decode_batch(&bytes, schema, bytes.len()).unwrap(); + assert_eq!(encode_batch(&restored).unwrap(), bytes); + } + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }], + time_index: None, + }); + let batch = Batch::try_new( + schema.clone(), + vec![ + vec![Value::Float64(f64::NAN)], + vec![Value::Float64(f64::INFINITY)], + ], + ) + .unwrap(); + let restored = decode_batch(&encode_batch(&batch).unwrap(), schema, usize::MAX).unwrap(); + assert!(matches!(restored.rows()[0][0], Value::Float64(v) if v.is_nan())); + assert!(matches!(restored.rows()[1][0], Value::Float64(v) if v == f64::INFINITY)); + } + + // Recovery validates the format, bound schema and actual sketch parameters. + #[test] + fn corrupt_or_relabelled_output_is_rejected() { + let batch = weighted(SketchAlgorithm::CmsWithHeap); + let bytes = encode_batch(&batch).unwrap(); + let wrong = weighted(SketchAlgorithm::CountSketchWithHeap); + assert!(decode_batch(&bytes, wrong.schema().clone(), usize::MAX).is_err()); + let mut stored: StoredBatch = rmp_serde::from_slice(&bytes).unwrap(); + stored.version = 2; + assert!(decode_batch( + &rmp_serde::to_vec_named(&stored).unwrap(), + batch.schema().clone(), + usize::MAX + ) + .is_err()); + stored.version = 1; + let Cell::Summary { bytes: payload, .. } = &mut stored.rows[0][1] else { + panic!() + }; + *payload = b"legacy integer heap".to_vec(); + assert!(decode_batch( + &rmp_serde::to_vec_named(&stored).unwrap(), + batch.schema().clone(), + usize::MAX + ) + .is_err()); + } + + // A stored Sum payload is read back as Planner's exact Sum with the same value. + #[test] + fn stored_sum_payload_decodes_as_planner_exact_sum() { + use crate::SerializableToSink; + use planner_types::post_asap::{ExactKind, ExactParams}; + let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + let state = StateCodec::SumAccumulatorV1 + .decode(&SumAccumulator::with_sum(5.5).serialize_to_bytes()) + .unwrap(); + let exact = state + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(exact.family(), &family); + assert_eq!( + exact.readout(crate::Statistic::Sum, None, None).unwrap(), + Some(5.5) + ); + } +} diff --git a/crates/asap_summary_state/src/stored_state/readout.rs b/crates/asap_summary_state/src/stored_state/readout.rs new file mode 100644 index 000000000..df13f33c6 --- /dev/null +++ b/crates/asap_summary_state/src/stored_state/readout.rs @@ -0,0 +1,303 @@ +//! Planner-declared readouts over reconstructed summary states. +use super::delta_apply::SummaryState; +use planner_types::{post_asap::SketchQuery, pre_asap::ColumnRef}; +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("{0}")] + Unsupported(&'static str), +} +pub fn sketch_query_value(rs: &SummaryState, query: &SketchQuery) -> Result { + if let SummaryState::UnivMon(state) = rs { + use crate::AggregateCore; + let statistic = match query { + SketchQuery::Cardinality => crate::Statistic::Cardinality, + SketchQuery::FrequencyL2 => crate::Statistic::FrequencyL2, + SketchQuery::FrequencyEntropy => crate::Statistic::FrequencyEntropy, + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + } => crate::Statistic::Count, + _ => return Err(Error::Unsupported("unsupported UnivMon readout")), + }; + return state + .query_statistic(statistic, &None, &Default::default()) + .map_err(|_| Error::Unsupported("UnivMon readout failed")); + } + match query { + SketchQuery::FrequencyL2 | SketchQuery::FrequencyEntropy => Err(Error::Unsupported( + "frequency moment readout requires UnivMon", + )), + SketchQuery::Quantile { q } => match rs { + // Typed PromQL/continuous-percentile readout uses interpolation; + // portable DDS `quantile` deliberately retains lower-rank parity. + SummaryState::Dd(sketch) => sketch.quantile_interpolated(*q).ok_or(Error::Unsupported( + "DDS interpolated quantile is unavailable", + )), + _ => Ok(rs.quantile(*q)), + }, + SketchQuery::Cardinality => Ok(rs.cardinality()), + // `key: ColumnRef::SampleValue, value: None` means "no specific + // item" -- the bare bucket total. `key: Named(_), value: Some(v)` + // is a per-item point lookup (e.g. `count(cms_metric{item="x"})`) + // -- `value` is where the filter's actual value lives (see + // `planner_types::post_asap::SketchQuery::PointCount`'s doc for why `readout` + // can't resolve it itself). Any other combination (e.g. a `Named` + // key with no value, or `SampleValue` with a value) is a shape + // this executor doesn't expect to see and reports rather than + // silently misreading. + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + } => Ok(rs.total()), + SketchQuery::PointCount { + key: ColumnRef::Named(_) | ColumnRef::Qualified { .. }, + value: Some(v), + } => rs.estimate(v).ok_or(Error::Unsupported( + "PointCount by key requires a Frequency-family sketch (Cms/CountSketch/..WithHeap)", + )), + SketchQuery::PointCount { .. } => Err(Error::Unsupported( + "unrecognized PointCount shape (key/value combination not expected)", + )), + // Both readout callers branch on `TopK` before ever calling this + // function (see `readout_cumulative`/`readout_per_window`), so + // this arm is unreachable in practice; kept for match + // exhaustiveness (`SketchQuery` has no `#[non_exhaustive]`) and to + // fail loudly rather than panic if that invariant is ever broken. + SketchQuery::TopK { .. } => Err(Error::Unsupported( + "TopK must be read out via topk_ranked, not sketch_query_value", + )), + } +} + +/// Rank a merged `SummaryState`'s top-k heap items descending by value and +/// cap at the requested `k`. The sort is load-bearing, not defensive +/// polish: `SummaryState::topk_items` reads back a bounded min-heap's +/// backing array as-is (`HHHeap::heap()`, asap_sketchlib) -- it does NOT +/// actually guarantee order despite its own doc wording. Errors for a +/// heap-less family (`Dd`/`Hll`/`Kll`/`Cms`/`CountSketch` -- no item +/// universe to rank), not for an empty heap (a heap-bearing family that +/// simply never received any updates yields `Ok(vec![])`, not an error). +pub fn topk_ranked(rs: &SummaryState, k: usize) -> Result, Error> { + let mut items = rs.topk_items().ok_or(Error::Unsupported( + "TopK requires a heap-bearing family (CmsWithHeap/CountSketchWithHeap) -- \ + this state's family carries no item universe to rank", + ))?; + items.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.0.cmp(&b.0)) // deterministic tie-break for equal counts + }); + items.truncate(k); + Ok(items) +} + +/// Merge already selected exact panes and finalize using the shared accumulator contract. +pub fn exact_readout( + states: impl IntoIterator>, + statistic: crate::Statistic, + key: &Option, + parameters: &std::collections::HashMap, +) -> Result { + let merged = merge_exact_states(states)?; + merged + .query_statistic(statistic, key, parameters) + .map_err(|e| e.to_string()) +} + +fn merge_exact_states( + states: impl IntoIterator>, +) -> Result, String> { + let mut states = states.into_iter(); + let first = states + .next() + .ok_or_else(|| "empty exact state input".to_string())?; + if let Some(first) = first + .as_any() + .downcast_ref::() + { + let mut merged = first.clone(); + for state in states { + let other = state + .as_any() + .downcast_ref::() + .ok_or_else(|| "merge requires Planner exact state".to_string())?; + merged + .merge_from(other) + .map_err(|error| error.to_string())?; + } + return Ok(Box::new(merged)); + } + let mut merged = first.clone_boxed_core(); + for state in states { + merged = merged + .merge_with(state.as_ref()) + .map_err(|e| e.to_string())?; + } + Ok(merged) +} + +/// PromQL counter readouts omit a series with fewer than two samples. Other +/// state/type/range failures remain errors rather than empty results. +pub fn insufficient_counter_samples( + state: &dyn crate::AggregateCore, + statistic: crate::Statistic, +) -> bool { + matches!( + statistic, + crate::Statistic::Rate | crate::Statistic::Increase + ) && (state + .as_any() + .downcast_ref::() + .is_some_and(|state| { + state.sample_count < 2 || state.last_seen_timestamp == state.starting_timestamp + }) + || state + .as_any() + .downcast_ref::() + .is_some_and(|state| state.insufficient_counter_samples(statistic, &None))) +} + +pub fn exact_readout_optional( + states: impl IntoIterator>, + statistic: crate::Statistic, + key: &Option, + parameters: &std::collections::HashMap, +) -> Result, String> { + let merged = merge_exact_states(states)?; + let counter = key.as_ref().and_then(|key| { + merged + .as_any() + .downcast_ref::() + .and_then(|state| state.increases.get(key)) + }); + let exact_insufficient = merged + .as_any() + .downcast_ref::() + .is_some_and(|state| state.insufficient_counter_samples(statistic, key)); + if exact_insufficient + || insufficient_counter_samples(merged.as_ref(), statistic) + || counter.is_some_and(|counter| insufficient_counter_samples(counter, statistic)) + { + return Ok(None); + } + merged + .query_statistic(statistic, key, parameters) + .map(Some) + .map_err(|error| error.to_string()) +} + +#[cfg(test)] +mod counter_tests { + use super::*; + use crate::{summary_kernels::IncreaseAccumulator, AggregateCore, Measurement, Statistic}; + use std::sync::Arc; + + #[test] + fn planner_counter_population_omits_insufficient_samples() { + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + for (kind, params, statistic) in [ + (ExactKind::Rate, ExactParams::Rate, Statistic::Rate), + ( + ExactKind::Increase, + ExactParams::Increase, + Statistic::Increase, + ), + ] { + for keyed in [false, true] { + let mut state = crate::summary_kernels::exact::ExactAccumulator::new( + SummaryFamilyType::ExactAggregate(kind.clone(), params.clone()), + keyed, + ) + .unwrap(); + let key = keyed + .then(|| crate::KeyByLabelValues::new_with_labels(vec!["checkout".into()])); + state.update(key.as_ref(), 10., 10_000); + assert_eq!( + exact_readout_optional( + [Arc::new(state) as Arc], + statistic, + &key, + &Default::default() + ) + .unwrap(), + None + ); + } + } + } + + #[test] + fn sparse_counter_is_absent_but_invalid_ranges_still_fail() { + let mut state = + IncreaseAccumulator::new(Measurement::new(10.), 10_000, Measurement::new(10.), 10_000); + let parameters = std::collections::HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + assert_eq!( + exact_readout_optional( + [Arc::new(state.clone()) as Arc], + Statistic::Rate, + &None, + ¶meters + ) + .unwrap(), + None + ); + let mut repeated = state.clone(); + repeated.update(Measurement::new(10.), 10_000); + assert_eq!( + exact_readout_optional( + [Arc::new(repeated) as Arc], + Statistic::Rate, + &None, + ¶meters + ) + .unwrap(), + None + ); + let mut keyed = crate::summary_kernels::KeyedCounterState::new(); + let label = crate::KeyByLabelValues::new_with_labels(vec!["checkout".into()]); + keyed.update(label.clone(), state.clone()); + assert_eq!( + exact_readout_optional( + [Arc::new(keyed.clone()) as Arc], + Statistic::Rate, + &Some(label), + ¶meters + ) + .unwrap(), + None + ); + assert!(exact_readout_optional( + [Arc::new(keyed) as Arc], + Statistic::Rate, + &Some(crate::KeyByLabelValues::new_with_labels(vec![ + "missing".into() + ])), + ¶meters + ) + .is_err()); + state.update(Measurement::new(20.), 20_000); + assert!(exact_readout_optional( + [Arc::new(state.clone()) as Arc], + Statistic::Rate, + &None, + ¶meters + ) + .unwrap() + .is_some()); + let invalid = std::collections::HashMap::from([ + ("range_start_ms".into(), "60000".into()), + ("range_end_ms".into(), "0".into()), + ]); + assert!(exact_readout_optional( + [Arc::new(state) as Arc], + Statistic::Rate, + &None, + &invalid + ) + .is_err()); + assert!(exact_readout_optional([], Statistic::Rate, &None, ¶meters).is_err()); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/count_min_sketch.rs b/crates/asap_summary_state/src/summary_kernels/count_min_sketch.rs new file mode 100644 index 000000000..cffe90d28 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/count_min_sketch.rs @@ -0,0 +1,1323 @@ +use crate::summary_kernels::dd_sketch::normalize_sample_p; +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountMinSketch, CountMinSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Count-Min Sketch accumulator — wraps asap_sketchlib::CountMinSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. +/// This file retains QE-specific trait impls, legacy deserializers, and JSON output. +#[derive(Debug, Clone)] +pub struct CountMinSketchAccumulator { + pub inner: CountMinSketch, + /// Edge sampling probability `p ∈ (0,1]` carried on the producer's + /// `SketchEnvelope.sample_p`. The edge admits each insert with + /// probability `p`, so every stored cell count is ~`p`× the true count. + /// CMS is L1/additive and linear, so the unbiased rescale of BOTH a + /// point-frequency estimate (`query_key`) and the aggregate + /// total-event statistics (`Count`/`Sum`/`Increase`/`Rate`) is `×1/p`. + /// `1.0` (and the proto3 default `0.0`, dual-read as `1.0`) means no + /// sampling, so the rescale is a no-op and the behaviour is identical + /// to before. Mirrors `DDSketchAccumulator::sample_p`; set from the + /// envelope at the `from_sketchlib_proto_bytes` decode site and + /// preserved across `reset_to_empty` and `merge_with`. + pub sample_p: f64, +} + +impl CountMinSketchAccumulator { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + inner: CountMinSketch::new(row_num, col_num), + sample_p: 1.0, + } + } + + // Marked as _update and kept private; only called internally. + fn _update(&mut self, key: &KeyByLabelValues, value: f64) { + self.inner.update(&key.to_semicolon_str(), value); + } + + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + // The edge sampled inserts with probability `sample_p`, so the + // stored point-frequency estimate is ~`p`× the true frequency. + // CMS is linear/additive, so `×1/p` is the unbiased rescale. + // `sample_p == 1.0` (unsampled / legacy) makes this a no-op. + self.inner.estimate(&key.to_semicolon_str()) / self.sample_p + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let row_num = data["row_num"] + .as_f64() + .ok_or("Missing or invalid 'row_num' field")? as usize; + let col_num = data["col_num"] + .as_f64() + .ok_or("Missing or invalid 'col_num' field")? as usize; + + let sketch_data = data["sketch"] + .as_array() + .ok_or("Missing or invalid 'sketch' field")?; + + let mut sketch = Vec::new(); + for row in sketch_data { + let row_array = row.as_array().ok_or("Invalid row in sketch data")?; + let mut sketch_row = Vec::new(); + for cell in row_array { + let value = cell.as_f64().ok_or("Invalid cell value in sketch data")?; + sketch_row.push(value); + } + sketch.push(sketch_row); + } + + Ok(Self { + inner: CountMinSketch::from_legacy_matrix(sketch, row_num, col_num), + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `CountMinSketchDataPoint.sketch` bytes when + /// `encoding = COUNT_MIN_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `CountMinSketch` wire struct (same format the legacy Arroyo path + /// uses — this method is the modified-OTLP entrypoint for PR I). + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountMinSketch::from_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, + // The msgpack CountMinSketch struct carries no envelope/sample_p; + // the msgpack path is parity/test-only and is never edge-sampled. + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `CountMinSketchDataPoint.sketch` bytes — i.e. the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::CountMinState` message used by + /// DataCollector's `countminsketchprocessor` when emitting via + /// `Metric.data = CountMinSketch{…}` with + /// `encoding = COUNT_MIN_SKETCH_ENCODING_PROTO`. + /// + /// The resulting accumulator is constructed via + /// `CountMinSketch::from_legacy_matrix` after reshaping the flat + /// `counts_int` / `counts_float` field into a `Vec>`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + + // DataCollector's countminsketchprocessor wraps the state in a + // `SketchEnvelope{count_min: CountMinState}` via + // `SerializePortableFO` + `proto.Marshal`. Try decoding as envelope + // first, fall back to bare `CountMinState` for callers (e.g. unit + // tests) that encode the state directly. Capture the envelope's + // `sample_p` alongside the state so the point-frequency + // (`query_key`) and aggregate statistics rescale by `1/p`. Bare + // `CountMinState` bytes (no envelope) carry no sampling info → + // `sample_p` 1.0 (no rescale). Mirrors `DDSketchAccumulator`. + let (state, sample_p) = match SketchEnvelope::decode(buffer) { + Ok(env) => { + let sp = env.sample_p; + match env.sketch_state { + Some(sketch_envelope::SketchState::CountMin(st)) => (st, sp), + Some(other) => { + return Err(format!( + "SketchEnvelope contains non-CountMin sketch: {:?}", + std::mem::discriminant(&other) + ) + .into()); + } + // Envelope decoded but was empty (e.g. the buffer is a + // bare CountMinState that happened to parse as a default + // envelope). Fall through to bare decode. + None => ( + CountMinState::decode(buffer) + .map_err(|e| format!("decode CountMinState: {e}"))?, + 1.0, + ), + } + } + Err(_) => ( + CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))?, + 1.0, + ), + }; + let rows = state.rows as usize; + let cols = state.cols as usize; + // Defensive dim validation BEFORE reconstructing the matrix: + // reject degenerate / narrow-hash-budget-violating / absurdly + // oversized dims so a malformed payload fails gracefully (the + // ingest caller skips the data point) instead of building a + // degenerate or huge matrix. + validate_sketch_dims("CountMinState", rows, cols)?; + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { + format!( + "CountMinState has unknown counter_type tag {}", + state.counter_type + ) + })?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountMinState counts_int has {} entries, expected rows*cols = {}", + state.counts_int.len(), + expected_len + ) + .into()); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountMinState counts_float has {} entries, expected rows*cols = {}", + state.counts_float.len(), + expected_len + ) + .into()); + } + state.counts_float.clone() + } + // INT128 stores (hi, lo) pairs and would have 2 * rows * cols + // entries in counts_int; defer to PR C if a producer ever uses it. + other => { + return Err(format!( + "CountMinState counter_type {other:?} not yet supported \ + (PR C will extend coverage)" + ) + .into()); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(Self { + inner: CountMinSketch::from_legacy_matrix(matrix, rows, cols), + sample_p: normalize_sample_p(sample_p), + }) + } + + /// Apply a proto-encoded `CountMinDelta` frame to this + /// accumulator's inner sketch — the decode path for + /// `COUNT_MIN_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; + + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountMinDelta packed-array length mismatch: \ + cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + ) + .into()); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + let delta = CountMinSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l1: pb.l1, + l2: pb.l2, + // The Go-side CountMinDelta proto now carries an hh_keys field + // (heavy-hitter candidates), mirrored on asap_sketchlib's + // CountMinSketchDelta. The vendored Rust proto bindings here don't + // decode it yet, and CountMin has no TopK to rebuild, so pass an + // empty set — same handling as CountSketch's hh_keys. + hh_keys: Vec::new(), + }; + self.inner + .apply_delta(&delta) + .map_err(|e| format!("apply CountMinDelta: {e}"))?; + Ok(()) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 8 { + return Err("Buffer too short for row_num and col_num".into()); + } + + // TODO: this logic will need to be checked for i32 -> f64 + // Github Issue #11 + + let row_num = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; + let col_num = u32::from_le_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]) as usize; + + let expected_size = 8 + (row_num * col_num * 4); + if buffer.len() < expected_size { + return Err("Buffer too short for sketch data".into()); + } + + let mut sketch = Vec::new(); + let mut offset = 8; + + for _ in 0..row_num { + let mut row = Vec::new(); + for _ in 0..col_num { + let value = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + row.push(value); + offset += 8; + } + sketch.push(row); + } + + Ok(Self { + inner: CountMinSketch::from_legacy_matrix(sketch, row_num, col_num), + sample_p: 1.0, + }) + } + + /// Merge multiple accumulators efficiently without cloning all of them. + pub fn merge_multiple( + accumulators: &[Box], + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut cms_accumulators = Vec::with_capacity(accumulators.len()); + for acc in accumulators { + if acc.get_accumulator_type() != AggregationType::CountMinSketch { + return Err(format!( + "Cannot merge CountMinSketchAccumulator with {:?}", + acc.get_accumulator_type() + ) + .into()); + } + let cms_acc = acc + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountMinSketchAccumulator")?; + cms_accumulators.push(cms_acc); + } + + // Check dimensions are consistent + let rows = cms_accumulators[0].inner.rows(); + let cols = cms_accumulators[0].inner.cols(); + for acc in &cms_accumulators { + if acc.inner.rows() != rows || acc.inner.cols() != cols { + return Err( + "Cannot merge CountMinSketch accumulators with different dimensions".into(), + ); + } + } + + let inner_refs: Vec<&CountMinSketch> = + cms_accumulators.iter().map(|acc| &acc.inner).collect(); + let merged_inner = CountMinSketch::merge_refs(&inner_refs)?; + // sample_p is a per-series config constant, so all operands carry the + // same value in practice. Mirror DDSketch's merge policy: prefer a + // sampled factor (< 1.0) over the no-sampling default so a merge with + // a freshly-reset (1.0) base keeps the series' sampling rate. + let sample_p = cms_accumulators + .iter() + .map(|acc| acc.sample_p) + .find(|&p| p < 1.0) + .unwrap_or(cms_accumulators[0].sample_p); + Ok(Self { + inner: merged_inner, + sample_p, + }) + } +} + +/// Defensive upper bound on the number of matrix cells (`rows * cols`) +/// we'll reconstruct from an inbound wire-declared CMS / CountSketch +/// dimension pair. A malformed / hostile payload could declare absurd +/// dims (e.g. `rows = cols = u32::MAX`) and trick the decoder into a +/// huge `Vec` allocation before the `counts_*.len() != rows*cols` +/// check ever runs. Realistic sketches are at most a few hundred rows +/// by tens-of-thousands of columns, so 8M cells (~64 MiB of f64) is a +/// generous ceiling that no legitimate producer reaches. +pub(crate) const MAX_SKETCH_CELLS: usize = 8 * 1024 * 1024; + +/// Validate an inbound, wire-declared `(rows, cols)` pair for a +/// matrix-backed frequency sketch (CMS / CountSketch) BEFORE any matrix +/// is reconstructed from it. Returns `Ok(())` for dimensions a +/// legitimate producer could have emitted, and an `Err` (never a panic) +/// for malformed / degenerate ones so the ingest path can skip the data +/// point and fall through to its existing decode-failure accounting. +/// +/// Rejections: +/// 1. `rows < 1` or `cols < 1` — a zero-dim matrix has no cells. +/// 2. Narrow-hash-budget violation. The cross-language wire hasher +/// (`sketchlib`'s `MatrixHashType::Packed64`) derives every row's +/// column index from disjoint bit-fields of a single 64-bit hash +/// word: row `r` reads `mask_bits = ceil(log2(cols))` bits at offset +/// `r * mask_bits`. Once `rows * mask_bits > 64` the per-row column +/// slices overflow / alias the 64-bit word and the matrix-cell +/// layout is no longer the one the producer hashed into — the sketch +/// is internally degenerate. This mirrors sketchlib's own +/// `MatrixFastHash::assert_compatible` budget (`rows * (mask_bits + 1) <= 64`); we check the column-index bits alone so realistic +/// configs (5x2048, 5x4096, 5x2000) — for which the sign bits share +/// the top of the word without affecting the cell layout — still +/// pass. +/// 3. Obviously-oversized dims: `rows * cols > MAX_SKETCH_CELLS`, +/// guarding against a huge allocation from a malformed payload. +/// +/// `what` names the wire struct for the error message (e.g. +/// `"CountMinState"`). +pub(crate) fn validate_sketch_dims(what: &str, rows: usize, cols: usize) -> Result<(), String> { + if rows < 1 || cols < 1 { + return Err(format!( + "{what} has degenerate dims (rows={rows}, cols={cols}); rejecting" + )); + } + // mask_bits = ceil(log2(cols)); cols >= 1 here. ilog2 is floor(log2). + let mask_bits = if cols.is_power_of_two() { + cols.ilog2() as usize + } else { + cols.ilog2() as usize + 1 + }; + if rows.saturating_mul(mask_bits) > 64 { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) exceed the 64-bit \ + packed-hash column budget (rows * ceil(log2(cols)) = {} > 64); \ + the sketch's matrix-cell layout is degenerate, rejecting", + rows.saturating_mul(mask_bits) + )); + } + if rows.saturating_mul(cols) > MAX_SKETCH_CELLS { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) declare {} cells, \ + exceeding the {MAX_SKETCH_CELLS}-cell ingest cap; rejecting to \ + avoid a huge allocation from a malformed payload", + rows.saturating_mul(cols) + )); + } + Ok(()) +} + +impl SerializableToSink for CountMinSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), + "sketch": self.inner.sketch() + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountMinSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountMinSketchAccumulator" + } + + /// Per-window base rotation: rebuild an empty counter matrix with + /// the same (rows, cols) so the next window's additive cell deltas + /// align to the identical hash geometry. `sample_p` is a per-series + /// config constant (not per-window data), so it is intentionally + /// preserved across the rotation — mirrors `DDSketchAccumulator`. + fn reset_to_empty(&mut self) { + self.inner = CountMinSketch::new(self.inner.rows(), self.inner.cols()); + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountMinSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_cms = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountMinSketchAccumulator")?; + + let merged_inner = CountMinSketch::merge_refs(&[&self.inner, &other_cms.inner])?; + // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a + // per-series config constant, so both operands carry the same value + // in practice. Prefer a sampled factor over the no-sampling default + // so a merge with a freshly-reset (1.0) base keeps the series' + // sampling rate. + let sample_p = if self.sample_p < 1.0 { + self.sample_p + } else { + other_cms.sample_p + }; + Ok(Box::new(Self { + inner: merged_inner, + sample_p, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountMinSketch + } + + fn approx_memory_bytes(&self) -> usize { + // Conservative constant for the CountMinSketch counter matrix. + // Real per-instance sizing would require exposing rows/cols on + // the inner sketch; 16 KiB is a reasonable v1 default. + 16 * 1024 + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + use crate::Statistic; + + // Key-provided path: route to MultipleSubpopulationAggregate::query + // (the canonical "what's the count of this key?" lookup). + if let Some(key_val) = key.as_ref() { + return self.query(statistic, key_val, Some(query_kwargs)); + } + if let Some(k) = query_kwargs.get("key") { + let key_val = crate::KeyByLabelValues::new_with_labels(vec![k.clone()]); + return self.query(statistic, &key_val, Some(query_kwargs)); + } + + // No-key path: return total event volume. The min-row-sum is the + // canonical CMS estimator for "how many inserts were observed" — + // each insert increments exactly one cell per row, so every row + // sums to the true insert count (modulo collisions, which CMS + // never *underestimates*; min is the tightest upper bound). + // + // When the edge sampled this series (sample_p < 1.0), each insert + // was admitted w.p. `p`, so the stored min-row-sum is ~`p`× the + // true event count. CMS is L1/additive and linear, so rescale by + // `1/sample_p` for an unbiased estimate. `sample_p == 1.0` + // (unsampled / legacy) makes this a no-op. This rescales BOTH the + // Count/Sum/Increase statistics and (via the same closure) the + // Rate per-second readout. + let total_events = || -> f64 { + let matrix = self.inner.sketch(); + if matrix.is_empty() || matrix[0].is_empty() { + return 0.0; + } + let row_totals = matrix.iter().map(|r| r.iter().sum::()); + let min_total = row_totals.fold(f64::INFINITY, f64::min); + if min_total.is_finite() { + min_total / self.sample_p + } else { + 0.0 + } + }; + match statistic { + Statistic::Count | Statistic::Sum => Ok(total_events()), + // PR #111 honest-gap closure (in-the-bag for ASAP tier). + // CMS records insert counts but not timestamps, so per-second + // `rate(metric[range])` requires the engine to push the + // range duration via `query_kwargs["range_ms"]`. When + // present, divide the min-row-sum by `range_ms / 1000`. When + // absent (the engine has not been wired to inject range_ms + // for this query, e.g. instant `rate` calls outside the + // PromQL range-vector pattern), fall back to the raw event + // count so the answer is at least non-empty — the caller's + // caveat is that the units are events/window rather than + // events/second. Increase carries the same caveat. + Statistic::Rate => { + let total = total_events(); + let range_ms_str = query_kwargs.get("range_ms").map(String::as_str); + let Some(s) = range_ms_str else { + return Ok(total); + }; + let range_ms: f64 = s + .parse() + .map_err(|e| format!("CountMinSketchAccumulator: bad range_ms='{s}': {e}"))?; + if range_ms <= 0.0 { + return Err("CountMinSketchAccumulator: range_ms must be positive".into()); + } + Ok(total * 1000.0 / range_ms) + } + Statistic::Increase => Ok(total_events()), + other => Err(format!( + "CountMinSketchAccumulator: statistic {:?} not supported \ + without a key (only Count / Sum / Rate / Increase aggregate \ + over the whole sketch)", + other, + ) + .into()), + } + } +} + +impl MultipleSubpopulationAggregate for CountMinSketchAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountMinSketchAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_min_sketch_creation() { + let cms = CountMinSketchAccumulator::new(4, 1000); + assert_eq!(cms.inner.rows(), 4); + assert_eq!(cms.inner.cols(), 1000); + let sketch = cms.inner.sketch(); + assert_eq!(sketch.len(), 4); + assert_eq!(sketch[0].len(), 1000); + + for row in &sketch { + for &value in row { + assert_eq!(value, 0.0); + } + } + } + + #[test] + fn test_count_min_sketch_update() { + let mut cms = CountMinSketchAccumulator::new(2, 10); + let key = KeyByLabelValues::new(); + cms._update(&key, 1.0); + let result = cms.query_key(&key); + assert!(result >= 1.0); + } + + #[test] + fn test_count_min_sketch_query() { + let cms = CountMinSketchAccumulator::new(2, 10); + let key = KeyByLabelValues::new(); + assert_eq!(cms.query_key(&key), 0.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); + } + + #[test] + fn test_count_min_sketch_merge() { + // Build controlled state via from_legacy_matrix (works for both Legacy and Sketchlib backends). + let cms1 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![5.0, 0.0, 0.0], vec![0.0, 0.0, 10.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let cms2 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![3.0, 7.0, 0.0], vec![0.0, 0.0, 0.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + + let merged = CountMinSketchAccumulator::merge_accumulators(vec![cms1, cms2]).unwrap(); + + let merged_sketch = merged.inner.sketch(); + assert_eq!(merged_sketch[0][0], 8.0); + assert_eq!(merged_sketch[0][1], 7.0); + assert_eq!(merged_sketch[1][2], 10.0); + } + + #[test] + fn test_count_min_sketch_merge_dimension_mismatch() { + let cms1 = CountMinSketchAccumulator::new(2, 3); + let cms2 = CountMinSketchAccumulator::new(3, 3); + let result = CountMinSketchAccumulator::merge_accumulators(vec![cms1, cms2]); + assert!(result.is_err()); + } + + #[test] + fn test_count_min_sketch_as_aggregate_core() { + let cms = CountMinSketchAccumulator::new(2, 3); + assert_eq!(cms.type_name(), "CountMinSketchAccumulator"); + } + + #[test] + fn test_trait_object() { + let cms = CountMinSketchAccumulator::new(2, 3); + let trait_obj: Box = Box::new(cms); + assert_eq!(trait_obj.type_name(), "CountMinSketchAccumulator"); + } + + #[test] + fn test_count_min_sketch_key_query() { + let mut cms = CountMinSketchAccumulator::new(4, 100); + let key = KeyByLabelValues::new(); + assert_eq!(cms.query_key(&key), 0.0); + cms._update(&key, 5.0); + let result = cms.query_key(&key); + assert!(result >= 5.0); + } + + #[test] + fn test_update_and_query_use_same_key_encoding() { + // Regression test: _update and query_key must hash the same key string. + // Previously _update went through serialize_to_json (which returns a JSON + // array, so as_object() is always None) and always stored under key "". + // query_key correctly used key.labels.join(";"), so they never matched. + let mut cms = CountMinSketchAccumulator::new(4, 1000); + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string(), "prod".to_string()]); + cms._update(&key, 5.0); + let result = cms.query_key(&key); + assert!( + result >= 5.0, + "_update and query_key used different key encodings: got {result}" + ); + + // Also verify a different key does not interfere. + let other_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + // other_key was never updated; its estimate should be lower than key's. + let other_result = cms.query_key(&other_key); + // In a sketch this large there should be no collision, so other_result == 0. + assert_eq!( + other_result, 0.0, + "unrelated key returned non-zero: {other_result}" + ); + } + + #[test] + fn test_multiple_subpopulation_aggregate() { + let mut cms = CountMinSketchAccumulator::new(3, 50); + let key = KeyByLabelValues::new(); + cms._update(&key, 10.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert!(result >= 10.0); + + let keys = multi_trait.get_keys(); + assert!(keys.is_none()); + } + + #[test] + fn test_count_min_sketch_merge_multiple() { + // Build controlled state via from_legacy_matrix (works for both Legacy and Sketchlib backends). + let cms1 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![5.0, 0.0, 0.0], vec![0.0, 0.0, 10.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let cms2 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![3.0, 7.0, 0.0], vec![0.0, 0.0, 0.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let cms3 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![2.0, 0.0, 0.0], vec![0.0, 0.0, 5.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + + let boxed_accs: Vec> = + vec![Box::new(cms1), Box::new(cms2), Box::new(cms3)]; + + let merged = CountMinSketchAccumulator::merge_multiple(&boxed_accs).unwrap(); + + let merged_sketch = merged.inner.sketch(); + assert_eq!(merged_sketch[0][0], 10.0); + assert_eq!(merged_sketch[0][1], 7.0); + assert_eq!(merged_sketch[1][2], 15.0); + } + + #[test] + fn test_count_min_sketch_merge_multiple_error_cases() { + let empty: Vec> = vec![]; + assert!(CountMinSketchAccumulator::merge_multiple(&empty).is_err()); + + let cms1 = CountMinSketchAccumulator::new(2, 3); + let cms2 = CountMinSketchAccumulator::new(3, 3); + let boxed_accs: Vec> = vec![Box::new(cms1), Box::new(cms2)]; + assert!(CountMinSketchAccumulator::merge_multiple(&boxed_accs).is_err()); + + use crate::summary_kernels::sum::SumAccumulator; + let cms = CountMinSketchAccumulator::new(2, 3); + let sum = SumAccumulator::new(); + let mixed_accs: Vec> = vec![Box::new(cms), Box::new(sum)]; + assert!(CountMinSketchAccumulator::merge_multiple(&mixed_accs).is_err()); + } + + #[test] + fn test_from_sketchlib_proto_bytes_int64() { + // Hand-build a CountMinState proto with INT64 counters and verify + // round-tripping through from_sketchlib_proto_bytes yields the same + // matrix that the modified-OTLP wire format would carry. + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + + let rows = 2u32; + let cols = 3u32; + // Row-major: row 0 = [1,2,3], row 1 = [4,5,6] + let counts_int: Vec = vec![1, 2, 3, 4, 5, 6]; + let state = CountMinState { + rows, + cols, + counter_type: CounterType::Int64 as i32, + counts_int: counts_int.clone(), + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let bytes = state.encode_to_vec(); + + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix.len(), rows as usize); + assert_eq!(matrix[0], vec![1.0, 2.0, 3.0]); + assert_eq!(matrix[1], vec![4.0, 5.0, 6.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's countminsketchprocessor emits: + // the state is wrapped in a `SketchEnvelope{count_min: ...}` + // via sketchlib-go's `SerializePortableFO` + `proto.Marshal`. + // Before the fix, the Rust decoder decoded the envelope bytes as + // a bare CountMinState, which produced "invalid wire type" + // errors on field `cols` and silently fell through to §5.2. + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + + let state = CountMinState { + rows: 2, + cols: 3, + counter_type: CounterType::Int64 as i32, + counts_int: vec![7, 8, 9, 10, 11, 12], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![7.0, 8.0, 9.0]); + assert_eq!(matrix[1], vec![10.0, 11.0, 12.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + // An envelope carrying a non-CountMin sketch should be rejected + // with a clear error rather than silently producing garbage. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + + let kll = KllState::default(); + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(kll)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_float64() { + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + + let state = CountMinState { + rows: 2, + cols: 2, + counter_type: CounterType::Float64 as i32, + counts_int: Vec::new(), + counts_float: vec![1.5, 2.5, 3.5, 4.5], + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let bytes = state.encode_to_vec(); + + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.5, 2.5]); + assert_eq!(matrix[1], vec![3.5, 4.5]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_dimension_mismatch() { + // counts_int has 5 entries but rows*cols = 6 → expect error + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + + let state = CountMinState { + rows: 2, + cols: 3, + counter_type: CounterType::Int64 as i32, + counts_int: vec![1, 2, 3, 4, 5], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let bytes = state.encode_to_vec(); + + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!( + result.unwrap_err().to_string().contains("counts_int"), + "error should mention counts_int dim mismatch" + ); + } + + #[test] + fn test_from_sketchlib_proto_bytes_zero_dims_rejected() { + use asap_sketchlib::proto::sketchlib::CountMinState; + use prost::Message; + + let state = CountMinState::default(); + let bytes = state.encode_to_vec(); + + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("degenerate dims")); + } + + #[test] + fn test_apply_proto_delta_bytes_round_trip() { + use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + use prost::Message; + + let mut acc = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let bytes = PbDelta { + rows: 2, + cols: 3, + cell_rows: vec![0, 1], + cell_cols: vec![0, 2], + d_counts: vec![10, 100], + l1: vec![], + l2: vec![], + ..Default::default() + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!( + acc.inner.sketch(), + vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 106.0]] + ); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = CountMinSketchAccumulator::new(2, 3); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ---------------------------------------------------------------- + // Statistic::Rate / Statistic::Increase — PR #111 honest-gap closure. + // CMS records insert counts but not timestamps. The Rate readout + // requires the engine to push `range_ms` via query_kwargs; without + // it the accumulator falls back to the raw event count (units of + // events/window) so the answer is at least non-empty. + // ---------------------------------------------------------------- + + #[test] + fn test_query_statistic_rate_with_range_ms() { + // Build a CMS whose min-row-sum is 100 events. With a 5-minute + // (300_000 ms) range, the per-second rate is 100 / 300 ≈ 0.333. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![100.0, 0.0], vec![100.0, 0.0]], + 2, + 2, + ), + sample_p: 1.0, + }; + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "300000".to_string()); + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect("Rate with range_ms is supported"); + assert!( + (v - (100.0 / 300.0)).abs() < 1e-9, + "expected 100/300 = {}, got {v}", + 100.0 / 300.0, + ); + } + + #[test] + fn test_query_statistic_rate_without_range_ms_falls_back_to_count() { + // Without `range_ms` in kwargs the accumulator returns the raw + // event volume (events/window units). Caller is responsible for + // surfacing that caveat to the user; this avoids `status=error` + // for instant rate-shape queries that bypass the matrix-selector + // code path. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix(vec![vec![42.0, 0.0], vec![42.0, 0.0]], 2, 2), + sample_p: 1.0, + }; + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Rate, &None, &HashMap::new()) + .expect("Rate without range_ms still answers (fallback)"); + assert_eq!(v, 42.0); + } + + #[test] + fn test_query_statistic_increase_returns_total_count() { + // Increase semantics on CMS: total events in the window — the + // same min-row-sum as Sum / Count. Differs from Rate only in + // that it never divides by range. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix(vec![vec![5.0, 7.0], vec![3.0, 9.0]], 2, 2), + sample_p: 1.0, + }; + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Increase, &None, &HashMap::new()) + .expect("Increase is supported"); + // min-row-sum: row0 = 12, row1 = 12, min = 12. + assert_eq!(v, 12.0); + } + + // ---------------------------------------------------------------- + // Defensive inbound-dimension validation (harden/sketch-dim-validation). + // Malformed / degenerate / narrow-hash-budget-violating CMS dims must + // be rejected gracefully (Err, never a panic); valid configs the + // backend actually uses (5x2048, 5x4096, 5x2000) must still decode. + // ---------------------------------------------------------------- + + /// Build a bare `CountMinState` proto carrying the given dims and a + /// row-major INT64 counts vector sized to `rows*cols` so that, IF the + /// dims pass validation, the reshape also succeeds. Used to prove a + /// malformed-dim payload is rejected at the dim gate, not later. + fn cms_state_bytes(rows: u32, cols: u32) -> Vec { + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + let n = (rows as usize).saturating_mul(cols as usize); + let state = CountMinState { + rows, + cols, + counter_type: CounterType::Int64 as i32, + counts_int: vec![0i64; n], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + state.encode_to_vec() + } + + #[test] + fn test_validate_sketch_dims_accepts_valid_configs() { + // The realistic configs the backend uses must pass unchanged. + for (r, c) in [(5usize, 2048usize), (5, 4096), (5, 2000), (4, 1000), (2, 3)] { + assert!( + validate_sketch_dims("CountMinState", r, c).is_ok(), + "valid config {r}x{c} was wrongly rejected" + ); + } + } + + #[test] + fn test_validate_sketch_dims_rejects_malformed() { + // Zero dims. + assert!(validate_sketch_dims("CountMinState", 0, 2048).is_err()); + assert!(validate_sketch_dims("CountMinState", 5, 0).is_err()); + // Narrow-hash-budget violation: 5 * ceil(log2(8192))=5*13=65 > 64. + let err = validate_sketch_dims("CountMinState", 5, 8192).unwrap_err(); + assert!(err.contains("budget"), "expected budget error, got: {err}"); + // Absurdly oversized: 1 x 16,777,216 = 16M cells > 8M cap. (1 row + // keeps the hash budget tiny — 1*24=24 — so the cap check, not the + // budget check, is what fires here.) + let err = validate_sketch_dims("CountMinState", 1, 16_777_216).unwrap_err(); + assert!(err.contains("cap"), "expected cell-cap error, got: {err}"); + // No panic on extreme dims (saturating_mul guards the products). + assert!(validate_sketch_dims("CountMinState", usize::MAX, usize::MAX).is_err()); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { + // A data point declaring narrow-hash-budget-violating dims must be + // skipped (Err returned, NOT a panic). The ingest caller turns + // this Err into a dropped data point + WARN log. + let bytes = cms_state_bytes(5, 8192); + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "budget-violating dims should be rejected"); + assert!(result.unwrap_err().to_string().contains("rejecting")); + + // A valid neighbour (5x4096) on the same path still decodes fine. + let ok_bytes = cms_state_bytes(5, 4096); + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) + .expect("valid 5x4096 CMS should still decode"); + assert_eq!(acc.inner.rows(), 5); + assert_eq!(acc.inner.cols(), 4096); + } + + #[test] + fn test_query_statistic_rate_rejects_invalid_range_ms() { + let cms = CountMinSketchAccumulator::new(2, 2); + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "0".to_string()); + let trait_obj: &dyn AggregateCore = &cms; + let err = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect_err("range_ms=0 should error"); + assert!(err.to_string().contains("positive")); + + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "not-a-number".to_string()); + let err = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect_err("non-numeric range_ms should error"); + assert!(err.to_string().contains("bad range_ms")); + } + + // ---------------------------------------------------------------- + // sample_p rescale. The edge admits each insert with probability `p`, + // so every stored cell is ~p× the true count. CMS is L1/additive and + // linear, so BOTH the point-frequency (query_key) and the aggregate + // total-event statistics (Count/Sum/Increase/Rate) rescale by 1/p. + // ---------------------------------------------------------------- + + #[test] + fn test_query_key_rescaled_by_sample_p() { + // Same stored cell counts, two sample_p values: the p=0.25 sketch + // must report 4× the point-frequency of the unsampled one. + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let mut unsampled = CountMinSketchAccumulator::new(4, 1000); + unsampled._update(&key, 10.0); + let mut sampled = CountMinSketchAccumulator::new(4, 1000); + sampled._update(&key, 10.0); + sampled.sample_p = 0.25; + + let raw = unsampled.query_key(&key); + let rescaled = sampled.query_key(&key); + assert!( + raw >= 10.0, + "raw estimate should be >= inserted 10, got {raw}" + ); + assert!( + (rescaled - raw * 4.0).abs() < 1e-9, + "expected point-frequency rescaled ≈ 4×raw ({}), got {rescaled}", + raw * 4.0 + ); + } + + #[test] + fn test_aggregate_statistics_rescaled_by_sample_p() { + use crate::Statistic; + // Build a CMS with a known min-row-sum of 12 events, sampled at + // p=0.25 → every aggregate statistic should report 12 / 0.25 = 48. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix(vec![vec![5.0, 7.0], vec![3.0, 9.0]], 2, 2), + sample_p: 0.25, + }; + let trait_obj: &dyn AggregateCore = &cms; + for stat in [Statistic::Count, Statistic::Sum, Statistic::Increase] { + let v = trait_obj + .query_statistic(stat, &None, &HashMap::new()) + .unwrap_or_else(|e| panic!("{stat:?} should be supported: {e}")); + // min-row-sum = 12, rescaled by 1/0.25 = 48. + assert!( + (v - 48.0).abs() < 1e-9, + "{stat:?}: expected rescaled 48, got {v}" + ); + } + // Rate also divides through the rescaled total: 48 events over a + // 6-second (6000 ms) range = 8 events/s. + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "6000".to_string()); + let r = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect("rate ok"); + assert!((r - 8.0).abs() < 1e-9, "expected rate 8.0, got {r}"); + } + + #[test] + fn test_sample_p_unset_behaves_as_one() { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + // An envelope with no sample_p (proto3 default 0.0) must normalize + // to 1.0 (no rescale) — byte-compatible with legacy frames. + let state = CountMinState { + rows: 2, + cols: 2, + counter_type: CounterType::Int64 as i32, + counts_int: vec![1, 2, 3, 4], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let env = SketchEnvelope { + // sample_p left at proto3 default 0.0. + sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 1.0, "unset sample_p must normalize to 1.0"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { + use crate::Statistic; + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + // min-row-sum = 12 raw; sample_p 0.25 → Count = 48. + let state = CountMinState { + rows: 2, + cols: 2, + counter_type: CounterType::Float64 as i32, + counts_int: Vec::new(), + counts_float: vec![5.0, 7.0, 3.0, 9.0], + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let env = SketchEnvelope { + sample_p: 0.25, + sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 0.25); + let trait_obj: &dyn AggregateCore = &acc; + let v = trait_obj + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + assert!((v - 48.0).abs() < 1e-9, "expected rescaled 48, got {v}"); + } + + #[test] + fn test_reset_to_empty_preserves_sample_p() { + let mut acc = CountMinSketchAccumulator::new(2, 3); + acc.sample_p = 0.25; + acc.reset_to_empty(); + assert_eq!(acc.sample_p, 0.25, "window rotation must keep sample_p"); + } + + #[test] + fn test_merge_prefers_sampled_factor() { + let mut a = CountMinSketchAccumulator::new(2, 3); + a.sample_p = 0.25; + let b = CountMinSketchAccumulator::new(2, 3); // sample_p 1.0 + let merged = a.merge_with(&b).expect("merge ok"); + let merged = merged + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.sample_p, 0.25); + + // merge_multiple mirrors the same policy. + let mut c = CountMinSketchAccumulator::new(2, 3); + c.sample_p = 0.25; + let d = CountMinSketchAccumulator::new(2, 3); + let boxed: Vec> = vec![Box::new(d), Box::new(c)]; + let merged = CountMinSketchAccumulator::merge_multiple(&boxed).expect("merge ok"); + assert_eq!(merged.sample_p, 0.25); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/count_min_sketch_with_heap.rs b/crates/asap_summary_state/src/summary_kernels/count_min_sketch_with_heap.rs new file mode 100644 index 000000000..f5d1369d1 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/count_min_sketch_with_heap.rs @@ -0,0 +1,832 @@ +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CmsHeapItem, CountMinSketchWithHeap, MessagePackCodec}; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Local serde view of the DELTA-HEAP wire frame produced by sketchlib-go's +/// `CountSketch.SerializeMsgpackWithHeapDelta` (encoding `MSGPACK_DELTA`). +/// Decoded with `rmp_serde` directly in the backend so NO delta API needs to +/// be added to the public `asap_sketchlib`. +/// +/// rmp_serde compact layout — a 4-element positional array: +/// +/// [ +/// is_delta: bool (always true), +/// matrix_delta: ( rows:u32, cols:u32, cells: Vec<(u32,u32,i64)> ), +/// topk_heap: Vec<(String, f64)>, // FULL heap, [key, value] pairs +/// heap_size: u64, +/// ] +/// +/// Tuple structs deserialize from msgpack fixed arrays positionally, so this +/// matches the Go encoder's byte layout exactly (no field names on the wire). +#[derive(Debug, Deserialize)] +struct HeapDeltaWire { + is_delta: bool, + matrix_delta: MatrixDeltaWire, + topk_heap: Vec<(String, f64)>, + #[allow(dead_code)] + heap_size: u64, +} + +#[derive(Debug, Deserialize)] +struct MatrixDeltaWire { + rows: u32, + cols: u32, + cells: Vec<(u32, u32, i64)>, +} + +/// Validated/flattened view of a decoded DELTA-HEAP frame. +struct HeapDeltaFrame { + rows: u32, + cols: u32, + heap_size: u64, + cells: Vec<(u32, u32, i64)>, + heap: Vec<(String, f64)>, +} + +impl HeapDeltaFrame { + fn from_msgpack(buffer: &[u8]) -> Result> { + let wire: HeapDeltaWire = rmp_serde::from_slice(buffer) + .map_err(|e| format!("decode CountSketchWithHeap delta msgpack: {e}"))?; + if !wire.is_delta { + return Err("CountSketchWithHeap delta frame has is_delta=false".into()); + } + Ok(Self { + rows: wire.matrix_delta.rows, + cols: wire.matrix_delta.cols, + heap_size: wire.heap_size, + cells: wire.matrix_delta.cells, + heap: wire.topk_heap, + }) + } +} + +/// Count-Min Sketch with Heap accumulator — wraps `asap_sketchlib::CountMinSketchWithHeap`. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::message_pack_format::portable::countminsketch_topk`. +/// This file retains QE-specific trait impls, legacy deserializers, and JSON output. +#[derive(Debug, Clone)] +pub struct CountMinSketchWithHeapAccumulator { + pub inner: CountMinSketchWithHeap, +} + +// Re-export HeapItem so existing code using CountMinSketchWithHeapAccumulator::HeapItem still works. +pub use asap_sketchlib::CmsHeapItem as HeapItemReexport; + +impl CountMinSketchWithHeapAccumulator { + pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { + Self { + inner: CountMinSketchWithHeap::new(row_num, col_num, heap_size), + } + } + + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + let key_string = key.labels.join(";"); + self.inner.estimate(&key_string) + } + + /// Decode a heap-bearing CountSketch FULL msgpack frame + /// (`{sketch:[matrix,rows,cols], topk_heap, heap_size}`) into a heap + /// accumulator. This is the window-1 / full-frame base for the + /// DELTA-HEAP delta path: the backend caches THIS accumulator as the + /// per-series base so a later `MSGPACK_DELTA` frame applies its sparse + /// matrix delta onto a heap accumulator (not a plain CountSketch). + /// + /// Delegates to the PUBLIC `asap_sketchlib::CountMinSketchWithHeap:: + /// from_msgpack` (both heap-bearing frequency variants share the wire + /// shape; the CountSketch-with-heap promotion is decided by the ingest + /// router, not the bytes). + pub fn from_msgpack_with_heap_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountMinSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketchWithHeap msgpack: {e}"))?, + }) + } + + /// Apply a DELTA-HEAP msgpack frame (encoding `MSGPACK_DELTA`) onto this + /// accumulator IN PLACE, WITHOUT any change to the public + /// `asap_sketchlib`: the frame is decoded generically with `rmp_serde` + /// into local serde structs, the sparse signed cell deltas are added to + /// the stored matrix (read back via the public `sketch_matrix()`), and + /// the top-k heap is REPLACED with the frame's full heap. The rebuilt + /// inner is produced via the public `from_legacy_matrix`, which rounds + /// cells to the i64 storage and re-seeds the heap. + /// + /// Under the per-window-reset model (`docs/delta-baseline-contract.md` + /// §3) the ingest caller resets this accumulator to empty at a window + /// boundary before applying, so the delta — which is the window's own + /// matrix against an empty base — reconstructs the window's state. + pub fn apply_msgpack_heap_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + + let rows = self.inner.rows(); + let cols = self.inner.cols(); + let heap_size = self.inner.heap_size; + + // Read the current (post-reset, possibly empty) matrix and apply the + // sparse signed deltas additively. Cells outside the stored + // dimensions are skipped defensively (mirrors the plain-CountSketch + // delta apply). + let mut matrix = self.inner.sketch_matrix(); + for (r, c, dc) in &frame.cells { + let (r, c) = (*r as usize, *c as usize); + if r >= rows || c >= cols { + continue; + } + matrix[r][c] += *dc as f64; + } + + // Replace the heap with the frame's full heap. `from_legacy_matrix` + // re-seeds both the matrix and the heap from these inputs. + let heap: Vec = frame + .heap + .into_iter() + .map(|(key, value)| CmsHeapItem { key, value }) + .collect(); + + self.inner = + CountMinSketchWithHeap::from_legacy_matrix(matrix, heap, rows, cols, heap_size); + Ok(()) + } + + /// Reconstruct a heap accumulator STANDALONE from a single DELTA-HEAP + /// msgpack frame (encoding `MSGPACK_DELTA`), with NO cached per-series + /// base. Used by the read-side reducer's `FrequencyTopk` path, where — + /// unlike the ingest accumulator — there is no rolling base to apply + /// onto: under the per-window-reset contract + /// (`docs/delta-baseline-contract.md` §3) each window's delta encodes + /// that window's own state against an EMPTY base, so reconstruction is + /// "empty(dims) + apply(delta)". + /// + /// Reuses the exact ingest-side apply logic: read the (rows, cols, + /// heap_size) the frame declares, build an empty accumulator of those + /// dims (equivalent to `reset_to_empty` on a same-shape base), then + /// fold the frame in via `apply_msgpack_heap_delta_bytes`. No + /// `asap_sketchlib` change — the frame is decoded generically with + /// `rmp_serde`. + pub fn from_msgpack_heap_delta_bytes( + buffer: &[u8], + ) -> Result> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + if frame.rows == 0 || frame.cols == 0 { + return Err(format!( + "CountSketchWithHeap delta frame has zero dims (rows={}, cols={})", + frame.rows, frame.cols + ) + .into()); + } + let mut acc = Self::new( + frame.rows as usize, + frame.cols as usize, + frame.heap_size as usize, + ); + acc.apply_msgpack_heap_delta_bytes(buffer)?; + Ok(acc) + } + + /// This function seems will never be used anymore. Keep it for possible future use. + pub fn deserialize_from_json(data: &Value) -> Result> { + let row_num = data["row_num"] + .as_f64() + .ok_or("Missing or invalid 'row_num' field")? as usize; + let col_num = data["col_num"] + .as_f64() + .ok_or("Missing or invalid 'col_num' field")? as usize; + let heap_size = data["heap_size"] + .as_f64() + .ok_or("Missing or invalid 'heap_size' field")? as usize; + + let sketch_data = data["sketch"] + .as_array() + .ok_or("Missing or invalid 'sketch' field")?; + + let mut sketch = Vec::new(); + for row in sketch_data { + let row_array = row.as_array().ok_or("Invalid row in sketch data")?; + let mut sketch_row = Vec::new(); + for cell in row_array { + let value = cell.as_f64().ok_or("Invalid cell value in sketch data")?; + sketch_row.push(value); + } + sketch.push(sketch_row); + } + + let topk_heap_data = data["topk_heap"] + .as_array() + .ok_or("Missing or invalid 'topk_heap' field")?; + + let mut topk_heap = Vec::new(); + for item in topk_heap_data { + let key = item["key"] + .as_str() + .ok_or("Missing or invalid 'key' in heap item")? + .to_string(); + let value = item["value"] + .as_f64() + .ok_or("Missing or invalid 'value' in heap item")?; + topk_heap.push(CmsHeapItem { key, value }); + } + + Ok(Self { + inner: CountMinSketchWithHeap::from_legacy_matrix( + sketch, topk_heap, row_num, col_num, heap_size, + ), + }) + } + + pub fn deserialize_from_bytes(_buffer: &[u8]) -> Result> { + Err("deserialize_from_bytes for CountMinSketchWithHeapAccumulator not implemented".into()) + } + + /// VALUE-WEIGHTED heavy-hitter update (FIX: CountSketch/CMS topk + /// recall-0). The default ingest path inserts `+1` per occurrence keyed + /// by the raw `item`, so the heap ranks groups by OCCURRENCE COUNT — the + /// wrong answer for `topk(k, sum by (label) (metric))`, which asks for + /// the top groups by SUM OF VALUE. This update adds the sample `value` + /// (not `+1`) into both the CMS matrix and the top-k heap, keyed by the + /// GROUP LABEL (e.g. the `host` / `zone` value), so the heap's ranking is + /// by summed value. Repeated calls for the same `group_label` accumulate, + /// so after folding a window the heap holds Σvalue per group. + /// + /// Delegates to the library's value-weighted `CountMinSketchWithHeap:: + /// update(key, value)` (`sketchlib_cms_heap_update` → `insert_many(key, + /// round(value))`), which is the "separate update path" the evaluation + /// plan (Fig 3c) called for. + pub fn insert_value(&mut self, group_label: &str, value: f64) { + self.inner.update(group_label, value); + } + + /// Read the top-`k` GROUPS ranked by summed VALUE (descending), keyed by + /// the group label. Pairs with [`Self::insert_value`]: the heap built by + /// value-weighted updates ranks by Σvalue, so this returns the + /// value-weighted top-k (not the occurrence-count top-k the raw `item` + /// heap would give). Sorted descending by value; ties broken by key for + /// determinism; truncated to `k`. + pub fn topk_by_value(&self, k: usize) -> Vec<(String, f64)> { + let mut items: Vec<(String, f64)> = self + .inner + .topk_heap_items() + .into_iter() + .map(|it| (it.key, it.value)) + .collect(); + items.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.0.cmp(&b.0)) + }); + items.truncate(k); + items + } + + /// Get all keys from the top-k heap. + pub fn get_topk_keys(&self) -> Vec { + self.inner + .topk_heap_items() + .iter() + .map(|item| { + let labels: Vec = item.key.split(';').map(|s| s.to_string()).collect(); + KeyByLabelValues { labels } + }) + .collect() + } +} + +impl SerializableToSink for CountMinSketchWithHeapAccumulator { + fn serialize_to_json(&self) -> Value { + let heap_items: Vec = self + .inner + .topk_heap_items() + .iter() + .map(|item| { + serde_json::json!({ + "key": item.key, + "value": item.value + }) + }) + .collect(); + + serde_json::json!({ + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), + "heap_size": self.inner.heap_size, + "sketch": self.inner.sketch_matrix(), + "topk_heap": heap_items + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountMinSketchWithHeapAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountMinSketchWithHeapAccumulator" + } + + /// Per-window base rotation (`docs/delta-baseline-contract.md` §3): + /// rebuild an empty heap accumulator with the same (rows, cols, + /// heap_size) so the next window's DELTA-HEAP frame applies onto a clean, + /// same-shape base. Without this override the trait default is a no-op, + /// which would let the additive matrix delta accumulate across windows + /// (over-counting). Mirrors `CountSketchAccumulator::reset_to_empty`. + fn reset_to_empty(&mut self) { + self.inner = + CountMinSketchWithHeap::new(self.inner.rows(), self.inner.cols(), self.inner.heap_size); + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountMinSketchWithHeapAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_cms = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountMinSketchWithHeapAccumulator")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_cms.clone()])?; + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountMinSketchWithHeap + } + + fn get_keys(&self) -> Option> { + Some(self.get_topk_keys()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for CountMinSketchWithHeapAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for CountMinSketchWithHeapAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountMinSketchWithHeapAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_min_sketch_with_heap_creation() { + let cms = CountMinSketchWithHeapAccumulator::new(4, 1000, 20); + assert_eq!(cms.inner.rows(), 4); + assert_eq!(cms.inner.cols(), 1000); + assert_eq!(cms.inner.heap_size, 20); + assert_eq!(cms.inner.topk_heap_items().len(), 0); + } + + #[test] + fn test_count_min_sketch_with_heap_query() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 10, 5); + let key = KeyByLabelValues::new(); + assert_eq!(cms.query_key(&key), 0.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); + } + + #[test] + fn test_count_min_sketch_with_heap_merge() { + // Build controlled state via from_legacy_matrix (works regardless of backend config). + let sketch1 = vec![ + vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap1 = vec![ + CmsHeapItem { + key: "key1".to_string(), + value: 100.0, + }, + CmsHeapItem { + key: "key2".to_string(), + value: 50.0, + }, + ]; + let sketch2 = vec![ + vec![5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap2 = vec![ + CmsHeapItem { + key: "key3".to_string(), + value: 75.0, + }, + CmsHeapItem { + key: "key1".to_string(), + value: 80.0, + }, + ]; + + let cms1 = CountMinSketchWithHeapAccumulator { + inner: CountMinSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), + }; + let cms2 = CountMinSketchWithHeapAccumulator { + inner: CountMinSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), + }; + + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.sketch_matrix()[0][0], 15.0); + assert_eq!(merged.inner.sketch_matrix()[1][1], 35.0); + assert_eq!(merged.inner.heap_size, 3); + assert!(merged.inner.topk_heap_items().len() <= 3); + } + + #[test] + fn test_count_min_sketch_with_heap_merge_single() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms.clone()]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.rows(), cms.inner.rows()); + assert_eq!(merged.inner.cols(), cms.inner.cols()); + assert_eq!(merged.inner.heap_size, cms.inner.heap_size); + } + + #[test] + fn test_count_min_sketch_with_heap_merge_dimension_mismatch() { + let cms1 = CountMinSketchWithHeapAccumulator::new(2, 10, 5); + let cms2 = CountMinSketchWithHeapAccumulator::new(3, 10, 5); + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("dimension")); + } + + #[test] + fn test_count_min_sketch_with_heap_as_aggregate_core() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + assert_eq!(cms.type_name(), "CountMinSketchWithHeapAccumulator"); + } + + #[test] + fn test_get_topk_keys() { + let mut cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + cms.inner.update("label1;label2", 100.0); + cms.inner.update("label3;label4", 50.0); + + let keys = cms.get_topk_keys(); + assert_eq!(keys.len(), 2); + // Top-k order can differ between Legacy and Sketchlib backends (heap ordering / estimates). + let label_sets: std::collections::HashSet<_> = + keys.iter().map(|k| k.labels.clone()).collect(); + assert!(label_sets.contains(&vec!["label1".to_string(), "label2".to_string()])); + assert!(label_sets.contains(&vec!["label3".to_string(), "label4".to_string()])); + } + + #[test] + fn test_multiple_subpopulation_aggregate() { + let cms = CountMinSketchWithHeapAccumulator::new(3, 50, 10); + let key = KeyByLabelValues::new(); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, 0.0); + + let keys = multi_trait.get_keys(); + assert!(keys.is_some()); + assert_eq!(keys.unwrap().len(), 0); + } + + // ---------------------------------------------------------------- + // DELTA-HEAP wire form (encoding MSGPACK_DELTA): apply a sparse matrix + // delta + replace the heap, decoded generically (rmp_serde) WITHOUT any + // asap_sketchlib delta API. The first test feeds a frame produced by the + // Go encoder (sketchlib-go `MarshalCountSketchWithHeapDelta`) to prove + // cross-language byte parity — mirrors how the full-heap parity is + // proven. The second proves PWR full -> delta -> delta reconstruction. + // ---------------------------------------------------------------- + + /// Cross-language byte-parity: this hex is the exact output of + /// sketchlib-go's `asapmsgpack.MarshalCountSketchWithHeapDelta(5, 1024, + /// cells=[(0,1,50),(1,3,-4),(4,1023,1_000_000)], + /// heap=[("/checkout",50),("/cart",20)], heap_size=20)` (captured via a + /// throw-away Go print test, identical methodology to the full-heap + /// golden in `sketchlib-go/.../count_sketch_with_heap_test.go`). If the + /// Go encoder or the rmp_serde layout ever shifts, this decode fails + /// loudly. + const GO_DELTA_HEAP_GOLDEN_HEX: &str = "94c39305cd04009393000132930103fc9304cd03ffce000f42409292a92f636865636b6f7574cb404900000000000092a52f63617274cb403400000000000014"; + + #[test] + fn test_apply_go_produced_delta_heap_frame_matrix_and_heap() { + let bytes = hex::decode(GO_DELTA_HEAP_GOLDEN_HEX).expect("hex"); + + // Base = empty heap accumulator with the frame's dims (what the + // ingest caller holds after the per-window base rotation). + let mut acc = CountMinSketchWithHeapAccumulator::new(5, 1024, 20); + acc.apply_msgpack_heap_delta_bytes(&bytes) + .expect("apply Go delta-heap frame"); + + // Matrix: the three sparse cells landed onto the empty base. + let m = acc.inner.sketch_matrix(); + assert_eq!(m.len(), 5); + assert_eq!(m[0].len(), 1024); + assert_eq!(m[0][1], 50.0, "cell (0,1)"); + assert_eq!(m[1][3], -4.0, "cell (1,3)"); + assert_eq!(m[4][1023], 1_000_000.0, "cell (4,1023)"); + // Everything else stays zero. + assert_eq!(m[2][2], 0.0); + assert_eq!(m[0][0], 0.0); + + // Heap: the frame's full heap, with /checkout ranked above /cart. + let mut items = acc.inner.topk_heap_items(); + items.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap()); + assert_eq!(items.len(), 2); + assert_eq!(items[0].key, "/checkout"); + assert_eq!(items[0].value, 50.0); + assert_eq!(items[1].key, "/cart"); + assert_eq!(items[1].value, 20.0); + } + + #[test] + fn test_pwr_full_then_delta_then_delta_reconstructs_per_window() { + use asap_sketchlib::MessagePackCodec; + + // Window 1 (full frame): build a heap-bearing CountSketch with mass + // and serialize the FULL `{sketch,topk_heap,heap_size}` frame, then + // decode it into a heap accumulator (the cached per-series base). + let w1 = CountMinSketchWithHeap::from_legacy_matrix( + vec![vec![300.0; 4]; 5], + vec![CmsHeapItem { + key: "k".into(), + value: 300.0, + }], + 5, + 4, + 20, + ); + let w1_bytes = w1.to_msgpack().expect("w1 full msgpack"); + let mut base = CountMinSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(&w1_bytes) + .expect("decode w1 full frame as heap accumulator"); + assert_eq!(base.inner.sketch_matrix()[0][0], 300.0); + + // Window 2 delta: this window's own state is matrix cells of value 50 + // against an EMPTY base + heap {k:50}. The DELTA-HEAP frame is encoded + // the same way the Go producer does (4-array, is_delta, sparse cells). + let w2_frame = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + // PWR: rotate base to empty at the window boundary, then apply. + base.reset_to_empty(); + assert_eq!( + base.inner.sketch_matrix()[0][0], + 0.0, + "reset_to_empty cleared matrix" + ); + base.apply_msgpack_heap_delta_bytes(&w2_frame) + .expect("apply w2 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 50.0, "window-2 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 50.0); + // No cross-window leakage from window 1's 300s. + assert_eq!(base.inner.sketch_matrix()[2][2], 0.0); + let h2: Vec<_> = base.inner.topk_heap_items(); + assert_eq!(h2.len(), 1); + assert_eq!(h2[0].key, "k"); + assert_eq!(h2[0].value, 50.0); + + // Window 3 delta: 80s against empty + heap {k:80}. + let w3_frame = encode_delta_heap(5, 4, &[(0, 0, 80)], &[("k", 80.0)], 20); + base.reset_to_empty(); + base.apply_msgpack_heap_delta_bytes(&w3_frame) + .expect("apply w3 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 80.0, "window-3 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 0.0, "no window-2 leakage"); + let h3 = base.inner.topk_heap_items(); + assert_eq!(h3.len(), 1); + assert_eq!(h3[0].value, 80.0); + } + + #[test] + fn test_rmp_serde_layout_is_byte_identical_to_go_encoder() { + // The rmp_serde positional encoding of the delta-heap frame must be + // BYTE-IDENTICAL to sketchlib-go's hand-rolled + // `MarshalCountSketchWithHeapDelta`. This hex is the Go encoder's + // output for (5, 4, cells=[(0,0,50),(1,1,50)], heap=[("k",50)], + // heap_size=20) — the same inputs `encode_delta_heap` uses below. + // Equality here proves both encode AND decode are cross-language + // byte-compatible (the decode path is exercised by the Go-golden + // test above). + const GO_PARITY_HEX: &str = "94c39305049293000032930101329192a16bcb404900000000000014"; + let rust_bytes = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + assert_eq!(hex::encode(&rust_bytes), GO_PARITY_HEX); + } + + #[test] + fn test_apply_delta_rejects_full_frame_and_garbage() { + use asap_sketchlib::MessagePackCodec; + let mut acc = CountMinSketchWithHeapAccumulator::new(2, 4, 5); + // A FULL frame (3-array, no is_delta marker) must NOT decode as a + // delta — the routing relies on the two shapes being distinct. + let full = CountMinSketchWithHeap::from_legacy_matrix( + vec![vec![1.0; 4]; 2], + vec![CmsHeapItem { + key: "a".into(), + value: 1.0, + }], + 2, + 4, + 5, + ) + .to_msgpack() + .unwrap(); + assert!(acc.apply_msgpack_heap_delta_bytes(&full).is_err()); + assert!(acc.apply_msgpack_heap_delta_bytes(b"not msgpack").is_err()); + } + + /// Encode a DELTA-HEAP frame the same way sketchlib-go's + /// `MarshalCountSketchWithHeapDelta` does (rmp_serde positional layout), + /// so the test exercises the real decode path. Tuple structs serialize + /// as msgpack fixed arrays — byte-identical to the Go hand-rolled writer. + fn encode_delta_heap( + rows: u32, + cols: u32, + cells: &[(u32, u32, i64)], + heap: &[(&str, f64)], + heap_size: u64, + ) -> Vec { + #[derive(serde::Serialize)] + struct W<'a>( + bool, + (u32, u32, &'a [(u32, u32, i64)]), + Vec<(String, f64)>, + u64, + ); + let heap_owned: Vec<(String, f64)> = + heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); + let w = W(true, (rows, cols, cells), heap_owned, heap_size); + rmp_serde::to_vec(&w).expect("encode delta-heap") + } + + // ---------------------------------------------------------------- + // FIX 1 — VALUE-WEIGHTED top-k (recall 0 → correct). + // + // `topk(k, sum by (host) (cpu_load))` asks for the top-k hosts by + // SUM OF VALUE. The heavy-hitter heap built by the default `+1`-per- + // occurrence update ranks by COUNT keyed by `item`, so its recall + // against the value-weighted ground truth is 0 when the busiest host + // (most samples) is NOT the heaviest host (largest Σvalue). + // `insert_value(group_label, value)` adds the sample VALUE keyed by the + // GROUP LABEL, so `topk_by_value` ranks by Σvalue — correct recall. + // ---------------------------------------------------------------- + + /// Crafted adversarial dataset: the host with the MOST samples + /// (`h_chatty`, 100 tiny samples) is NOT the host with the largest + /// value-sum (`h_heavy`, a handful of huge samples). A COUNT-ranked + /// heap would surface `h_chatty`; the value-weighted top-k must surface + /// the true heavy hitters by Σvalue, giving recall 1.0 against the + /// ground-truth top-k-by-value-sum. + #[test] + fn value_weighted_topk_has_full_recall_vs_count_topk() { + // (host, per-sample value, sample count) → true Σvalue: + // h_heavy : 1000 × 3 = 3000 (few samples, huge value) + // h_mid : 200 × 5 = 1000 + // h_small : 50 × 6 = 300 + // h_chatty: 1 × 100 = 100 (MOST samples, tiny value) + let data: &[(&str, f64, usize)] = &[ + ("h_heavy", 1000.0, 3), + ("h_mid", 200.0, 5), + ("h_small", 50.0, 6), + ("h_chatty", 1.0, 100), + ]; + + // Wide CMS + heap large enough to hold every group exactly (4 groups) + // so the estimate equals the true Σvalue with no hash collisions. + let mut acc = CountMinSketchWithHeapAccumulator::new(5, 4096, 16); + let mut truth: std::collections::HashMap<&str, f64> = std::collections::HashMap::new(); + for (host, value, count) in data { + for _ in 0..*count { + acc.insert_value(host, *value); + } + *truth.entry(*host).or_insert(0.0) += value * (*count as f64); + } + + // Ground-truth top-2 by value-sum: h_heavy (3000), h_mid (1000). + let mut truth_ranked: Vec<(&str, f64)> = truth.into_iter().collect(); + truth_ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + let truth_top2: std::collections::HashSet<&str> = + truth_ranked.iter().take(2).map(|(k, _)| *k).collect(); + assert!( + truth_top2.contains("h_heavy") && truth_top2.contains("h_mid"), + "ground-truth top-2 by value-sum should be h_heavy + h_mid" + ); + + // Value-weighted top-2 from the heap. + let got = acc.topk_by_value(2); + assert_eq!(got.len(), 2, "k=2 → two groups: {got:?}"); + let got_keys: std::collections::HashSet<&str> = + got.iter().map(|(k, _)| k.as_str()).collect(); + + // RECALL = |got ∩ truth| / |truth| must be 1.0. + let hits = got_keys.intersection(&truth_top2).count(); + let recall = hits as f64 / truth_top2.len() as f64; + assert_eq!( + recall, 1.0, + "value-weighted top-k recall must be 1.0 (count-ranked heap would \ + surface h_chatty and miss h_heavy → recall < 1): got={got:?}" + ); + + // The busiest-by-count host (h_chatty) must NOT be in the top-2, + // proving we rank by value-sum, not occurrence count. + assert!( + !got_keys.contains("h_chatty"), + "h_chatty (most samples, smallest value-sum) must be excluded: {got:?}" + ); + + // Estimates are exact here (no collisions, heap holds all groups): + // top-1 must be h_heavy with Σvalue 3000. + assert_eq!(got[0].0, "h_heavy"); + assert!( + (got[0].1 - 3000.0).abs() < 1e-6, + "h_heavy value-sum estimate ≈ 3000, got {}", + got[0].1 + ); + assert_eq!(got[1].0, "h_mid"); + assert!( + (got[1].1 - 1000.0).abs() < 1e-6, + "h_mid value-sum estimate ≈ 1000, got {}", + got[1].1 + ); + } + + /// A single value-weighted insert must put the full value (not +1) into + /// the heap, and repeated inserts for the same group must accumulate. + #[test] + fn insert_value_accumulates_summed_value_in_heap() { + let mut acc = CountMinSketchWithHeapAccumulator::new(4, 1024, 8); + acc.insert_value("g", 10.0); + acc.insert_value("g", 25.0); + let top = acc.topk_by_value(1); + assert_eq!(top.len(), 1); + assert_eq!(top[0].0, "g"); + assert!( + (top[0].1 - 35.0).abs() < 1e-6, + "summed value should be 35 (10+25), got {}", + top[0].1 + ); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/count_sketch.rs b/crates/asap_summary_state/src/summary_kernels/count_sketch.rs new file mode 100644 index 000000000..78c1f7d72 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/count_sketch.rs @@ -0,0 +1,678 @@ +//! CountSketch accumulator backed by `asap_sketchlib::CountSketch`. +//! +//! Supports worker merge, persistence serialization, and modified-OTLP proto +//! decoding. Per-key queries delegate to sketchlib's median-of-signed-rows +//! estimator so query and ingest use the same hash specification. Top-k +//! requires the separate heap-bearing accumulator. + +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountSketch, CountSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Count Sketch accumulator — inner matrix of signed counts. +#[derive(Debug, Clone)] +pub struct CountSketchAccumulator { + pub inner: CountSketch, +} + +impl CountSketchAccumulator { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + inner: CountSketch::new(row_num, col_num), + } + } + + /// Median-of-signed-rows point estimate for `key`, via the real + /// `asap_sketchlib::CountSketch::estimate` — the canonical, hash-spec- + /// compatible estimator (see `AggregateCore::query_statistic`'s doc for + /// why this replaced a hand-rolled, non-compatible hash). + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + self.inner.estimate(&key.to_semicolon_str()) + } + + /// Decode from the modified OTLP wire format's + /// `CountSketchDataPoint.sketch` bytes when + /// `encoding = COUNT_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `CountSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketch msgpack: {e}"))?, + }) + } + + /// Decode from the modified OTLP wire format's + /// `CountSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::CountSketchState` message + /// that DataCollector's `countsketchprocessor` emits when + /// `encoding = COUNT_SKETCH_ENCODING_PROTO`. + /// + /// Mirrors `CountMinSketchAccumulator::from_sketchlib_proto_bytes` + /// but on the signed-counter `CountSketchState`. The resulting + /// accumulator is constructed via + /// `CountSketch::from_legacy_matrix` after reshaping the flat + /// `counts_int` / `counts_float` field into a `Vec>`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountSketchState, CounterType, SketchEnvelope, + }; + use prost::Message; + + // DataCollector's countsketchprocessor wraps the state in a + // `SketchEnvelope{count_sketch: CountSketchState}` via + // sketchlib-go's `SerializePortableFO` + `proto.Marshal`. Try + // decoding as envelope first, fall back to bare + // `CountSketchState` for callers (e.g. unit tests) that + // encode the state directly. Mirrors the PR #14 fix on + // `CountMinSketchAccumulator::from_sketchlib_proto_bytes`. + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::CountSketch(st)) => st, + Some(other) => { + return Err(format!( + "SketchEnvelope contains non-CountSketch sketch: {:?}", + std::mem::discriminant(&other) + ) + .into()); + } + None => CountSketchState::decode(buffer) + .map_err(|e| format!("decode CountSketchState: {e}"))?, + }, + Err(_) => CountSketchState::decode(buffer) + .map_err(|e| format!("decode CountSketchState: {e}"))?, + }; + let rows = state.rows as usize; + let cols = state.cols as usize; + // Defensive dim validation BEFORE reconstructing the matrix: + // reject degenerate / narrow-hash-budget-violating / absurdly + // oversized dims so a malformed payload fails gracefully (the + // ingest caller skips the data point) instead of building a + // degenerate or huge matrix. Shares the CMS validator since the + // CountSketch matrix uses the same packed-hash column layout. + crate::summary_kernels::count_min_sketch::validate_sketch_dims( + "CountSketchState", + rows, + cols, + )?; + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { + format!( + "CountSketchState has unknown counter_type tag {}", + state.counter_type + ) + })?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountSketchState counts_int has {} entries, expected rows*cols = {}", + state.counts_int.len(), + expected_len + ) + .into()); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountSketchState counts_float has {} entries, expected rows*cols = {}", + state.counts_float.len(), + expected_len + ) + .into()); + } + state.counts_float.clone() + } + other => { + return Err(format!( + "CountSketchState counter_type {other:?} not yet supported \ + (INT128 stores interleaved hi/lo pairs; will be added when needed)" + ) + .into()); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(Self { + inner: CountSketch::from_legacy_matrix(matrix, rows, cols), + }) + } + + /// Apply a proto-encoded `CountSketchDelta` frame to this + /// accumulator's inner sketch — the decode path for + /// `COUNT_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + /// + /// Cells apply additively: `matrix[cell_rows[i]][cell_cols[i]] + /// += d_counts[i]`. Per-row L2 is parsed off the wire but + /// ignored at application time — it's a downstream error- + /// accounting signal, not a merge input. + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; + + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountSketchDelta packed-array length mismatch: \ + cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + ) + .into()); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + // This is the heap-less matrix kernel; ranked membership is handled + // by the explicit heap-bearing operator, not inferred from delta keys. + let delta = CountSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l2: pb.l2, + hh_keys: Vec::new(), + }; + self.inner + .apply_delta(&delta) + .map_err(|e| format!("apply CountSketchDelta: {e}"))?; + Ok(()) + } +} + +impl SerializableToSink for CountSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "row_num": self.inner.rows, + "col_num": self.inner.cols, + "sketch": self.inner.sketch(), + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountSketchAccumulator" + } + + /// Per-window base rotation: rebuild an empty signed-counter matrix + /// with the same (rows, cols) so the next window's additive cell + /// deltas align to the identical hash geometry. + fn reset_to_empty(&mut self) { + self.inner = CountSketch::new(self.inner.rows, self.inner.cols); + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_cs = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountSketchAccumulator")?; + + let merged_inner = CountSketch::merge_refs(&[&self.inner, &other_cs.inner])?; + Ok(Box::new(Self { + inner: merged_inner, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountSketch + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &HashMap, + ) -> Result> { + use crate::Statistic; + // Key-provided path: route to MultipleSubpopulationAggregate::query + // (the canonical "what's the count of this key?" lookup), same + // pattern as CountMinSketchAccumulator. Fixed from a hand-rolled + // `DefaultHasher`-based estimator that did NOT use the sketchlib + // hash spec (its own doc admitted this — "not the sketchlib hash + // spec... the canonical compatibility path requires plumbing the + // sketchlib seeds through") — `asap_sketchlib::CountSketch::estimate` + // already hashes against the correct portable spec, so this is a + // genuine correctness fix, not just a refactor. + if let Some(key_val) = key.as_ref() { + return self.query(statistic, key_val, Some(query_kwargs)); + } + if let Some(k) = query_kwargs.get("key") { + let key_val = KeyByLabelValues::new_with_labels(vec![k.clone()]); + return self.query(statistic, &key_val, Some(query_kwargs)); + } + // No-key path: unchanged from before this fix -- CountSketch's + // signed rows have no CMS-style "min-row-sum = true total" + // property, so these are documented approximations, not a + // heavy-hitter answer. Not touched by this fix (only the + // key-provided path above had the hash-compatibility bug). + match statistic { + Statistic::Topk | Statistic::Count => { + let matrix = self.inner.sketch(); + let total: f64 = matrix.iter().flatten().map(|v| v.abs()).sum(); + let rows = matrix.len() as f64; + Ok(if rows > 0.0 { total / rows } else { 0.0 }) + } + Statistic::Sum => { + let matrix = self.inner.sketch(); + let total: f64 = matrix.iter().flatten().sum(); + let rows = matrix.len() as f64; + Ok(if rows > 0.0 { total / rows } else { 0.0 }) + } + other => Err(format!( + "CountSketchAccumulator: statistic {:?} not supported (only Topk / Count / Sum, with optional `key` in query_kwargs)", + other, + ) + .into()), + } + } +} + +impl MultipleSubpopulationAggregate for CountSketchAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountSketchAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_query_key_uses_real_sketchlib_estimator() { + // `query_key` must match sketchlib's estimator and hash specification. + let mut cs = CountSketchAccumulator::new(4, 1000); + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + cs.inner.update(&key.to_semicolon_str(), 10.0); + assert_eq!( + cs.query_key(&key), + cs.inner.estimate(&key.to_semicolon_str()) + ); + } + + #[test] + fn test_multiple_subpopulation_aggregate_query() { + let mut cs = CountSketchAccumulator::new(4, 1000); + let key = KeyByLabelValues::new_with_labels(vec!["checkout".to_string()]); + cs.inner.update(&key.to_semicolon_str(), 25.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, cs.query_key(&key)); + + // query_statistic (the AggregateCore entry point) must route a + // provided key through the same path. + let core: &dyn AggregateCore = &cs; + let via_core = core + .query_statistic(Statistic::Sum, &Some(key.clone()), &HashMap::new()) + .unwrap(); + assert_eq!(via_core, cs.query_key(&key)); + } + + #[test] + fn test_mergeable_accumulator_merge_accumulators() { + let cs1 = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![1.0, -2.0], vec![3.0, -4.0]], 2, 2), + }; + let cs2 = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![-1.0, 2.0], vec![-3.0, 4.0]], 2, 2), + }; + let merged = CountSketchAccumulator::merge_accumulators(vec![cs1, cs2]).unwrap(); + assert_eq!(merged.inner.sketch(), &vec![vec![0.0, 0.0], vec![0.0, 0.0]]); + } + + #[test] + fn test_mergeable_accumulator_rejects_empty() { + let result = CountSketchAccumulator::merge_accumulators(vec![]); + assert!(result.is_err()); + } + + fn encode_state( + rows: u32, + cols: u32, + counter_type: i32, + counts_int: Vec, + counts_float: Vec, + ) -> Vec { + use asap_sketchlib::proto::sketchlib::CountSketchState; + use prost::Message; + let state = CountSketchState { + rows, + cols, + counter_type, + counts_int, + counts_float, + l2: Vec::new(), + topk: None, + }; + state.encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_int64() { + use asap_sketchlib::proto::sketchlib::CounterType; + // Signed 2x3 matrix: row 0 = [1,-2,3], row 1 = [-4,5,-6] + let bytes = encode_state( + 2, + 3, + CounterType::Int64 as i32, + vec![1, -2, 3, -4, 5, -6], + Vec::new(), + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.0, -2.0, 3.0]); + assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's countsketchprocessor emits: + // the state wrapped in a `SketchEnvelope{count_sketch: ...}` + // via sketchlib-go's `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountSketchState, CounterType, SketchEnvelope, + }; + use prost::Message; + + let state = CountSketchState { + rows: 2, + cols: 3, + counter_type: CounterType::Int64 as i32, + counts_int: vec![1, -2, 3, -4, 5, -6], + counts_float: Vec::new(), + ..Default::default() + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::CountSketch(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.0, -2.0, 3.0]); + assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + // An envelope carrying a non-CountSketch sketch should be + // rejected with a clear error rather than silently producing + // garbage. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_float64() { + use asap_sketchlib::proto::sketchlib::CounterType; + let bytes = encode_state( + 2, + 2, + CounterType::Float64 as i32, + Vec::new(), + vec![1.5, -2.5, 3.5, -4.5], + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.5, -2.5]); + assert_eq!(matrix[1], vec![3.5, -4.5]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_dimension_mismatch() { + use asap_sketchlib::proto::sketchlib::CounterType; + // 2x3 declared but only 5 int entries + let bytes = encode_state( + 2, + 3, + CounterType::Int64 as i32, + vec![1, 2, 3, 4, 5], + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!( + result.unwrap_err().to_string().contains("counts_int"), + "error should mention counts_int dim mismatch" + ); + } + + #[test] + fn test_from_sketchlib_proto_bytes_zero_dims_rejected() { + use asap_sketchlib::proto::sketchlib::CountSketchState; + use prost::Message; + let state = CountSketchState::default(); + let bytes = state.encode_to_vec(); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("degenerate dims")); + } + + #[test] + fn test_aggregate_core_merge_matches_matrix_add() { + let a = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![1.0, -2.0], vec![3.0, -4.0]], 2, 2), + }; + let b = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![-1.0, 2.0], vec![-3.0, 4.0]], 2, 2), + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + let m = merged.inner.sketch(); + assert_eq!(m[0], vec![0.0, 0.0]); + assert_eq!(m[1], vec![0.0, 0.0]); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::summary_kernels::count_min_sketch::CountMinSketchAccumulator; + let cs = CountSketchAccumulator::new(2, 3); + let cms = CountMinSketchAccumulator::new(2, 3); + let result = cs.merge_with(&cms); + assert!(result.is_err()); + } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = CountSketch::from_legacy_matrix( + vec![vec![1.0, -2.0, 3.0], vec![-4.0, 5.0, -6.0]], + 2, + 3, + ); + let bytes = original.to_msgpack().unwrap(); + let acc = CountSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.rows, 2); + assert_eq!(acc.inner.cols, 3); + assert_eq!(acc.inner.sketch(), original.sketch()); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = CountSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } + + #[test] + fn test_apply_proto_delta_bytes_round_trip() { + use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + use prost::Message; + + let mut acc = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix( + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], + 2, + 3, + ), + }; + let bytes = PbDelta { + rows: 2, + cols: 3, + cell_rows: vec![0, 1], + cell_cols: vec![0, 2], + d_counts: vec![10, -6], + l2: vec![], + ..Default::default() + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!( + acc.inner.sketch(), + &vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 0.0]] + ); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = CountSketchAccumulator::new(2, 3); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ---------------------------------------------------------------- + // Defensive inbound-dimension validation (harden/sketch-dim-validation). + // Malformed / narrow-hash-budget-violating CountSketch dims must be + // rejected gracefully (Err, never a panic); valid configs the backend + // actually uses (5x2048, 5x4096, 5x2000) must still decode. + // ---------------------------------------------------------------- + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { + use asap_sketchlib::proto::sketchlib::CounterType; + // 5 * ceil(log2(8192))=5*13=65 > 64 — narrow-hash-budget violation. + // counts sized to rows*cols so rejection is on dims, not length. + let n = 5usize * 8192usize; + let bytes = encode_state( + 5, + 8192, + CounterType::Int64 as i32, + vec![0i64; n], + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "budget-violating dims should be rejected"); + assert!(result.unwrap_err().to_string().contains("rejecting")); + + // A valid neighbour (5x4096) on the same path still decodes fine. + let n_ok = 5usize * 4096usize; + let ok_bytes = encode_state( + 5, + 4096, + CounterType::Int64 as i32, + vec![0i64; n_ok], + Vec::new(), + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) + .expect("valid 5x4096 CountSketch should still decode"); + assert_eq!(acc.inner.rows, 5); + assert_eq!(acc.inner.cols, 4096); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_oversized_dims() { + use asap_sketchlib::proto::sketchlib::CounterType; + // Declare 1 x 16,777,216 = 16M cells (> 8M cap) but send an empty + // counts vector: validation must reject on the dim cap BEFORE the + // decoder tries to allocate/reshape a 16M-entry matrix. (1 row keeps + // the hash budget tiny so the cap check, not the budget check, fires.) + let bytes = encode_state( + 1, + 16_777_216, + CounterType::Int64 as i32, + Vec::new(), + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "oversized dims should be rejected"); + let msg = result.unwrap_err().to_string(); + assert!(msg.contains("cap"), "expected cell-cap error, got: {msg}"); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/count_sketch_with_heap.rs b/crates/asap_summary_state/src/summary_kernels/count_sketch_with_heap.rs new file mode 100644 index 000000000..1f614bcb5 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/count_sketch_with_heap.rs @@ -0,0 +1,575 @@ +//! Count Sketch with Heap accumulator — wraps +//! `asap_sketchlib::CountSketchWithHeap`. +//! +//! Port of `count_min_sketch_with_heap.rs` for the distinct +//! `CountSketchWithHeap` (median-of-signed-rows estimator) rather than +//! `CountMinSketchWithHeap` (min-over-rows estimator). The two are +//! different sketch algorithms that happen to share a storage shape and +//! wire layout -- see `asap_sketchlib::CountSketchWithHeap`'s own doc and +//! this session's `delta_apply.rs`/`decoders.rs` fix on the read side. +//! Before this file existed, `accumulator_factory.rs`'s raw-metric +//! ingest dispatch built a `CountMinSketchWithHeapAccumulator` (CMS math) +//! for `SketchAlgorithm::CountSketchWithHeap` sids -- the same conflation bug +//! already fixed on the read side, now closed on the write side too. + +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountSketchWithHeap, CsHeapItem, MessagePackCodec}; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Local serde view of the DELTA-HEAP wire frame (encoding `MSGPACK_DELTA`). +/// Identical shape to `count_min_sketch_with_heap.rs`'s +/// `HeapDeltaWire`/`MatrixDeltaWire` -- the wire frame is generic (sparse +/// cell deltas + a full heap), not CMS-specific. See that file's doc for +/// the exact rmp_serde positional layout. +#[derive(Debug, Deserialize)] +struct HeapDeltaWire { + is_delta: bool, + matrix_delta: MatrixDeltaWire, + topk_heap: Vec<(String, f64)>, + #[allow(dead_code)] + heap_size: u64, +} + +#[derive(Debug, Deserialize)] +struct MatrixDeltaWire { + rows: u32, + cols: u32, + cells: Vec<(u32, u32, i64)>, +} + +/// Validated/flattened view of a decoded DELTA-HEAP frame. +struct HeapDeltaFrame { + rows: u32, + cols: u32, + heap_size: u64, + cells: Vec<(u32, u32, i64)>, + heap: Vec<(String, f64)>, +} + +impl HeapDeltaFrame { + fn from_msgpack(buffer: &[u8]) -> Result> { + let wire: HeapDeltaWire = rmp_serde::from_slice(buffer) + .map_err(|e| format!("decode CountSketchWithHeap delta msgpack: {e}"))?; + if !wire.is_delta { + return Err("CountSketchWithHeap delta frame has is_delta=false".into()); + } + Ok(Self { + rows: wire.matrix_delta.rows, + cols: wire.matrix_delta.cols, + heap_size: wire.heap_size, + cells: wire.matrix_delta.cells, + heap: wire.topk_heap, + }) + } +} + +/// Count Sketch with Heap accumulator — wraps `asap_sketchlib::CountSketchWithHeap`. +/// Core struct, update/merge/serde logic live in +/// `asap_sketchlib::message_pack_format::portable::countsketch_topk`. This +/// file retains QE-specific trait impls, legacy deserializers, and JSON +/// output -- same split as `CountMinSketchWithHeapAccumulator`. +#[derive(Debug, Clone)] +pub struct CountSketchWithHeapAccumulator { + pub inner: CountSketchWithHeap, +} + +impl CountSketchWithHeapAccumulator { + pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { + Self { + inner: CountSketchWithHeap::new(row_num, col_num, heap_size), + } + } + + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + let key_string = key.labels.join(";"); + self.inner.estimate(&key_string) + } + + /// Decode a heap-bearing CountSketch FULL msgpack frame into a heap + /// accumulator -- the window-1 / full-frame base for the DELTA-HEAP + /// delta path. Mirrors `CountMinSketchWithHeapAccumulator::from_msgpack_with_heap_bytes`. + pub fn from_msgpack_with_heap_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketchWithHeap msgpack: {e}"))?, + }) + } + + /// Apply a DELTA-HEAP msgpack frame (encoding `MSGPACK_DELTA`) onto this + /// accumulator IN PLACE. Mirrors + /// `CountMinSketchWithHeapAccumulator::apply_msgpack_heap_delta_bytes` + /// exactly -- the frame decode/apply logic is generic, not tied to + /// which estimator the rebuilt sketch uses. + pub fn apply_msgpack_heap_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + + let rows = self.inner.rows(); + let cols = self.inner.cols(); + let heap_size = self.inner.heap_size; + + let mut matrix = self.inner.sketch_matrix(); + for (r, c, dc) in &frame.cells { + let (r, c) = (*r as usize, *c as usize); + if r >= rows || c >= cols { + continue; + } + matrix[r][c] += *dc as f64; + } + + let heap: Vec = frame + .heap + .into_iter() + .map(|(key, value)| CsHeapItem { key, value }) + .collect(); + + self.inner = CountSketchWithHeap::from_legacy_matrix(matrix, heap, rows, cols, heap_size); + Ok(()) + } + + /// Reconstruct a heap accumulator STANDALONE from a single DELTA-HEAP + /// msgpack frame, with no cached per-series base. Mirrors + /// `CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes`. + pub fn from_msgpack_heap_delta_bytes( + buffer: &[u8], + ) -> Result> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + if frame.rows == 0 || frame.cols == 0 { + return Err(format!( + "CountSketchWithHeap delta frame has zero dims (rows={}, cols={})", + frame.rows, frame.cols + ) + .into()); + } + let mut acc = Self::new( + frame.rows as usize, + frame.cols as usize, + frame.heap_size as usize, + ); + acc.apply_msgpack_heap_delta_bytes(buffer)?; + Ok(acc) + } + + /// Value-weighted heavy-hitter update -- see + /// `CountMinSketchWithHeapAccumulator::insert_value`'s doc for why + /// this (not a `+1`-per-occurrence update) is the correct semantics + /// for `topk(k, sum by (label) (metric))`-shaped queries. + pub fn insert_value(&mut self, group_label: &str, value: f64) { + self.inner.update(group_label, value); + } + + /// Read the top-`k` groups ranked by summed value (descending, tie-broken + /// by key for determinism). Mirrors `CountMinSketchWithHeapAccumulator::topk_by_value`. + pub fn topk_by_value(&self, k: usize) -> Vec<(String, f64)> { + let mut items: Vec<(String, f64)> = self + .inner + .topk_heap_items() + .into_iter() + .map(|it| (it.key, it.value)) + .collect(); + items.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.0.cmp(&b.0)) + }); + items.truncate(k); + items + } + + /// Get all keys from the top-k heap. + pub fn get_topk_keys(&self) -> Vec { + self.inner + .topk_heap_items() + .iter() + .map(|item| { + let labels: Vec = item.key.split(';').map(|s| s.to_string()).collect(); + KeyByLabelValues { labels } + }) + .collect() + } +} + +impl SerializableToSink for CountSketchWithHeapAccumulator { + fn serialize_to_json(&self) -> Value { + let heap_items: Vec = self + .inner + .topk_heap_items() + .iter() + .map(|item| { + serde_json::json!({ + "key": item.key, + "value": item.value + }) + }) + .collect(); + + serde_json::json!({ + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), + "heap_size": self.inner.heap_size, + "sketch": self.inner.sketch_matrix(), + "topk_heap": heap_items + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountSketchWithHeapAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountSketchWithHeapAccumulator" + } + + /// Per-window base rotation -- mirrors + /// `CountMinSketchWithHeapAccumulator::reset_to_empty`. + fn reset_to_empty(&mut self) { + self.inner = + CountSketchWithHeap::new(self.inner.rows(), self.inner.cols(), self.inner.heap_size); + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountSketchWithHeapAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_cs = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountSketchWithHeapAccumulator")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_cs.clone()])?; + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountSketchWithHeap + } + + fn get_keys(&self) -> Option> { + Some(self.get_topk_keys()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for CountSketchWithHeapAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for CountSketchWithHeapAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountSketchWithHeapAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_sketch_with_heap_creation() { + let cs = CountSketchWithHeapAccumulator::new(4, 1000, 20); + assert_eq!(cs.inner.rows(), 4); + assert_eq!(cs.inner.cols(), 1000); + assert_eq!(cs.inner.heap_size, 20); + assert_eq!(cs.inner.topk_heap_items().len(), 0); + } + + #[test] + fn test_count_sketch_with_heap_query() { + let cs = CountSketchWithHeapAccumulator::new(2, 10, 5); + let key = KeyByLabelValues::new(); + assert_eq!(cs.query_key(&key), 0.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); + } + + #[test] + fn test_count_sketch_with_heap_merge() { + let sketch1 = vec![ + vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap1 = vec![ + CsHeapItem { + key: "key1".to_string(), + value: 100.0, + }, + CsHeapItem { + key: "key2".to_string(), + value: 50.0, + }, + ]; + let sketch2 = vec![ + vec![5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap2 = vec![ + CsHeapItem { + key: "key3".to_string(), + value: 75.0, + }, + CsHeapItem { + key: "key1".to_string(), + value: 80.0, + }, + ]; + + let cs1 = CountSketchWithHeapAccumulator { + inner: CountSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), + }; + let cs2 = CountSketchWithHeapAccumulator { + inner: CountSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), + }; + + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs1, cs2]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.sketch_matrix()[0][0], 15.0); + assert_eq!(merged.inner.sketch_matrix()[1][1], 35.0); + assert_eq!(merged.inner.heap_size, 3); + assert!(merged.inner.topk_heap_items().len() <= 3); + } + + #[test] + fn test_count_sketch_with_heap_merge_single() { + let cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs.clone()]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.rows(), cs.inner.rows()); + assert_eq!(merged.inner.cols(), cs.inner.cols()); + assert_eq!(merged.inner.heap_size, cs.inner.heap_size); + } + + #[test] + fn test_count_sketch_with_heap_merge_dimension_mismatch() { + let cs1 = CountSketchWithHeapAccumulator::new(2, 10, 5); + let cs2 = CountSketchWithHeapAccumulator::new(3, 10, 5); + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs1, cs2]); + assert!(result.is_err()); + } + + #[test] + fn test_count_sketch_with_heap_as_aggregate_core() { + let cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + assert_eq!(cs.type_name(), "CountSketchWithHeapAccumulator"); + } + + #[test] + fn test_get_topk_keys() { + let mut cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + cs.inner.update("label1;label2", 100.0); + cs.inner.update("label3;label4", 50.0); + + let keys = cs.get_topk_keys(); + assert_eq!(keys.len(), 2); + let label_sets: std::collections::HashSet<_> = + keys.iter().map(|k| k.labels.clone()).collect(); + assert!(label_sets.contains(&vec!["label1".to_string(), "label2".to_string()])); + assert!(label_sets.contains(&vec!["label3".to_string(), "label4".to_string()])); + } + + #[test] + fn test_multiple_subpopulation_aggregate() { + let cs = CountSketchWithHeapAccumulator::new(3, 50, 10); + let key = KeyByLabelValues::new(); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, 0.0); + + let keys = multi_trait.get_keys(); + assert!(keys.is_some()); + assert_eq!(keys.unwrap().len(), 0); + } + + #[test] + fn test_pwr_full_then_delta_then_delta_reconstructs_per_window() { + use asap_sketchlib::MessagePackCodec; + + let w1 = CountSketchWithHeap::from_legacy_matrix( + vec![vec![300.0; 4]; 5], + vec![CsHeapItem { + key: "k".into(), + value: 300.0, + }], + 5, + 4, + 20, + ); + let w1_bytes = w1.to_msgpack().expect("w1 full msgpack"); + let mut base = CountSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(&w1_bytes) + .expect("decode w1 full frame as heap accumulator"); + assert_eq!(base.inner.sketch_matrix()[0][0], 300.0); + + let w2_frame = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + base.reset_to_empty(); + assert_eq!( + base.inner.sketch_matrix()[0][0], + 0.0, + "reset_to_empty cleared matrix" + ); + base.apply_msgpack_heap_delta_bytes(&w2_frame) + .expect("apply w2 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 50.0, "window-2 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 50.0); + assert_eq!(base.inner.sketch_matrix()[2][2], 0.0); + let h2: Vec<_> = base.inner.topk_heap_items(); + assert_eq!(h2.len(), 1); + assert_eq!(h2[0].key, "k"); + assert_eq!(h2[0].value, 50.0); + + let w3_frame = encode_delta_heap(5, 4, &[(0, 0, 80)], &[("k", 80.0)], 20); + base.reset_to_empty(); + base.apply_msgpack_heap_delta_bytes(&w3_frame) + .expect("apply w3 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 80.0, "window-3 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 0.0, "no window-2 leakage"); + let h3 = base.inner.topk_heap_items(); + assert_eq!(h3.len(), 1); + assert_eq!(h3[0].value, 80.0); + } + + #[test] + fn test_apply_delta_rejects_full_frame_and_garbage() { + use asap_sketchlib::MessagePackCodec; + let mut acc = CountSketchWithHeapAccumulator::new(2, 4, 5); + let full = CountSketchWithHeap::from_legacy_matrix( + vec![vec![1.0; 4]; 2], + vec![CsHeapItem { + key: "a".into(), + value: 1.0, + }], + 2, + 4, + 5, + ) + .to_msgpack() + .unwrap(); + assert!(acc.apply_msgpack_heap_delta_bytes(&full).is_err()); + assert!(acc.apply_msgpack_heap_delta_bytes(b"not msgpack").is_err()); + } + + fn encode_delta_heap( + rows: u32, + cols: u32, + cells: &[(u32, u32, i64)], + heap: &[(&str, f64)], + heap_size: u64, + ) -> Vec { + #[derive(serde::Serialize)] + struct W<'a>( + bool, + (u32, u32, &'a [(u32, u32, i64)]), + Vec<(String, f64)>, + u64, + ); + let heap_owned: Vec<(String, f64)> = + heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); + let w = W(true, (rows, cols, cells), heap_owned, heap_size); + rmp_serde::to_vec(&w).expect("encode delta-heap") + } + + #[test] + fn insert_value_accumulates_summed_value_in_heap() { + let mut acc = CountSketchWithHeapAccumulator::new(4, 1024, 8); + acc.insert_value("g", 10.0); + acc.insert_value("g", 25.0); + let top = acc.topk_by_value(1); + assert_eq!(top.len(), 1); + assert_eq!(top[0].0, "g"); + assert!( + (top[0].1 - 35.0).abs() < 1e-6, + "summed value should be 35 (10+25), got {}", + top[0].1 + ); + } + + /// The core proof this file exists at all: `CountSketchWithHeapAccumulator` + /// wraps the real, distinct `asap_sketchlib::CountSketchWithHeap` -- + /// not the CMS-family `CountMinSketchWithHeap` a collapsed dispatch + /// used to substitute (the exact bug this file fixes on the ingest + /// side, mirroring the already-fixed read side). Two different Rust + /// types means `merge_with` rejects mixing them at the type-check + /// level, same as any other mismatched-family merge attempt -- + /// verified directly rather than via a numeric estimate comparison + /// (asap_sketchlib's own test suite already proves the median vs + /// min-over-rows divergence at the sketch-math level). + #[test] + fn test_rejects_merge_with_cms_family_accumulator() { + use crate::summary_kernels::count_min_sketch_with_heap::CountMinSketchWithHeapAccumulator; + + let cs = CountSketchWithHeapAccumulator::new(4, 64, 10); + let cms = CountMinSketchWithHeapAccumulator::new(4, 64, 10); + let result = cs.merge_with(&cms); + assert!( + result.is_err(), + "CountSketchWithHeapAccumulator must not merge with CountMinSketchWithHeapAccumulator \ + -- different algorithms sharing only a storage shape" + ); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/datasketches_kll.rs b/crates/asap_summary_state/src/summary_kernels/datasketches_kll.rs new file mode 100644 index 000000000..1ab4df791 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/datasketches_kll.rs @@ -0,0 +1,727 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use asap_sketchlib::{KllSketch, MessagePackCodec}; +use base64::{engine::general_purpose, Engine as _}; +use serde_json::Value; +use std::collections::HashMap; +#[cfg(feature = "extra_debugging")] +use std::time::Instant; +use tracing::debug; + +use crate::Statistic; + +/// KLL sketch accumulator — wraps asap_sketchlib::KllSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. +/// This file retains QE-specific trait impls and JSON output. +pub struct DatasketchesKLLAccumulator { + pub inner: KllSketch, +} + +impl DatasketchesKLLAccumulator { + pub fn new(k: u16) -> Self { + Self { + inner: KllSketch::new(k), + } + } + + pub fn update(&mut self, value: f64) { + self.inner.update(value); + } + + pub fn get_quantile(&self, quantile: f64) -> f64 { + self.inner.quantile(quantile) + } + + /// Decode from the modified OTLP wire format's + /// `KLLSketchDataPoint.sketch` bytes when + /// `encoding = KLL_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `KllSketch` struct — PR I parity entrypoint. Unlike the + /// `_ENCODING_PROTO` path (which does lossy statistical + /// reconstruction via `update()` replay), the msgpack path is a + /// bit-identical round-trip because sketch-core's `KllSketch` + /// serializes its full internal state to msgpack. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: KllSketch::from_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, + }) + } + + /// Decode from the modified OTLP wire format's + /// `KLLSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::KllState` message that + /// DataCollector's `kllprocessor` emits when + /// `encoding = KLL_SKETCH_ENCODING_PROTO`. + /// + /// The neutral codec decodes the sketchlib envelope. + /// The level-aware constructor below preserves the supplied retained + /// sample layout without replaying updates. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + let state = asap_sketch_codec::kll_state(buffer)?; + if state.k < 8 { + return Err(format!("KllState.k must be >= 8 (got {})", state.k).into()); + } + if state.k > u16::MAX as u32 { + return Err(format!( + "KllState.k does not fit in u16 (got {}, max {})", + state.k, + u16::MAX + ) + .into()); + } + // Validate the levels[] boundary array if it is populated. The + // proto contract says `levels[0] == 0` and + // `levels[num_levels] == items.len()`. If the producer left + // levels empty (common when num_levels is zero), skip. + if !state.levels.is_empty() { + if state.levels.len() as u32 != state.num_levels + 1 { + return Err(format!( + "KllState levels length = {}, expected num_levels+1 = {}", + state.levels.len(), + state.num_levels + 1 + ) + .into()); + } + if state.levels[0] != 0 { + return Err(format!("KllState.levels[0] = {}, expected 0", state.levels[0]).into()); + } + if *state.levels.last().unwrap() as usize != state.items.len() { + return Err(format!( + "KllState.levels[{}] = {}, expected items.len() = {}", + state.num_levels, + state.levels.last().unwrap(), + state.items.len() + ) + .into()); + } + } + let k = state.k as u16; + // Direct, bit-exact reconstruction from the portable state (no per-item + // `update()` replay) whenever the producer supplied the `levels[]` + // boundary array — which it does for any non-empty sketch. Falls back to + // the statistical replay only when `levels` is absent (empty sketch). + if !state.levels.is_empty() { + // KllState is highest-level first; the in-memory constructor + // expects L0 first. Replaying or copying the wire order changes + // retained-item weights after the first compaction. + let mut items = Vec::with_capacity(state.items.len()); + let mut levels = vec![0]; + if state + .levels + .windows(2) + .any(|bounds| bounds[0] > bounds[1] || bounds[1] as usize > state.items.len()) + { + return Err("KllState levels must be monotonic and within items".into()); + } + for bounds in state.levels.windows(2).rev() { + items.extend_from_slice(&state.items[bounds[0] as usize..bounds[1] as usize]); + levels.push(items.len()); + } + return Ok(Self { + inner: KllSketch::from_portable_state( + k, + &items, + &levels, + state.num_levels as usize, + ) + .map_err(|e| -> Box { e.into() })?, + }); + } + let mut acc = Self::new(k); + for item in &state.items { + acc.update(*item); + } + Ok(acc) + } + + /// Merge multiple accumulators efficiently without cloning all of them. + pub fn merge_multiple( + accumulators: &[Box], + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut kll_accumulators = Vec::with_capacity(accumulators.len()); + for acc in accumulators { + if acc.get_accumulator_type() != AggregationType::DatasketchesKLL { + return Err(format!( + "Cannot merge DatasketchesKLLAccumulator with {:?}", + acc.get_accumulator_type() + ) + .into()); + } + let kll_acc = acc + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to DatasketchesKLLAccumulator")?; + kll_accumulators.push(kll_acc); + } + + let inner_refs: Vec<&KllSketch> = kll_accumulators.iter().map(|acc| &acc.inner).collect(); + let merged_inner = KllSketch::merge_refs(&inner_refs)?; + Ok(Self { + inner: merged_inner, + }) + } +} + +// Manual trait implementations since the C++ library doesn't provide them +impl Clone for DatasketchesKLLAccumulator { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl std::fmt::Debug for DatasketchesKLLAccumulator { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DatasketchesKLLAccumulator") + .field("k", &self.inner.k) + .field("sketch_n", &self.inner.count()) + .finish() + } +} + +// TODO: verify this +// Thread safety: The C++ library is not thread-safe by default, but since we're using it +// in a single-threaded context per accumulator instance and only sharing read-only operations, +// this should be safe. +unsafe impl Send for DatasketchesKLLAccumulator {} +unsafe impl Sync for DatasketchesKLLAccumulator {} + +impl SerializableToSink for DatasketchesKLLAccumulator { + fn serialize_to_json(&self) -> Value { + // Mirror Python implementation: {"sketch": base64_encoded_string} + let sketch_bytes = self.inner.sketch_bytes(); + let sketch_b64 = general_purpose::STANDARD.encode(&sketch_bytes); + serde_json::json!({ "sketch": sketch_b64 }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for DatasketchesKLLAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "DatasketchesKLLAccumulator" + } + + 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> { + #[cfg(feature = "extra_debugging")] + let merge_with_start = Instant::now(); + #[cfg(feature = "extra_debugging")] + debug!( + "[PERF] DatasketchesKLLAccumulator::merge_with() started - self.k={}, self.n={}", + self.inner.k, + self.inner.count() + ); + + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge DatasketchesKLLAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_kll = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to DatasketchesKLLAccumulator")?; + + let merged_inner = KllSketch::merge_refs(&[&self.inner, &other_kll.inner])?; + let merged = Self { + inner: merged_inner, + }; + + #[cfg(feature = "extra_debugging")] + debug!( + "[PERF] DatasketchesKLLAccumulator::merge_with() TOTAL TIME: {:?}", + merge_with_start.elapsed() + ); + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::DatasketchesKLL + } + + fn approx_memory_bytes(&self) -> usize { + // KLL with default k=200 holds ~2*k items (~3 KiB). Round up + // for overhead. + 4 * 1024 + } + + fn aux_stats(&self) -> AuxStats { + // KLL natively tracks `count` (n, samples observed). min/max + // are available from the underlying sketch but only via a + // O(k) quantile extraction at quantile=0/1, which is not + // a cheap trait-method call. sum is not retained by KLL. + // + // Surface only count here; follow-up PR may add min/max via a + // dedicated accessor on sketch-core. `sum_over_time` queries + // on KLL fall back to query_statistic as they do today. + AuxStats { + count: Some(self.inner.count()), + ..AuxStats::empty() + } + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query(statistic, Some(query_kwargs)) + } +} + +impl SingleSubpopulationAggregate for DatasketchesKLLAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Quantile => { + debug!( + "Querying DatasketchesKLLAccumulator for quantile with kwargs: {:?}", + query_kwargs + ); + let quantile = query_kwargs + .and_then(|kwargs| kwargs.get("quantile")) + .ok_or("Missing quantile parameter for quantile query")? + .parse::() + .map_err(|_| "Invalid quantile parameter format")?; + + if !(0.0..=1.0).contains(&quantile) { + return Err("Quantile must be between 0.0 and 1.0".into()); + } + + Ok(self.get_quantile(quantile)) + } + _ => Err( + format!("Unsupported statistic in DatasketchesKLLAccumulator: {statistic:?}") + .into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for DatasketchesKLLAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use prost::Message; + + fn encode_state(state: asap_sketchlib::proto::sketchlib::KllState) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + } + .encode_to_vec() + } + + #[test] + fn test_datasketches_kll_creation() { + let kll = DatasketchesKLLAccumulator::new(200); + assert!(kll.inner.count() == 0); + assert_eq!(kll.inner.k, 200); + } + + #[test] + fn test_datasketches_kll_update() { + let mut kll = DatasketchesKLLAccumulator::new(200); + kll.update(10.0); + kll.update(20.0); + kll.update(15.0); + assert_eq!(kll.inner.count(), 3); + } + + #[test] + fn test_datasketches_kll_quantile() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + assert_eq!(kll.get_quantile(0.0), 1.0); + assert_eq!(kll.get_quantile(1.0), 10.0); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + let q50 = kll.get_quantile(0.5); + assert!((q50 - 6.0).abs() <= 1.0, "expected median ~6, got {q50}"); + } + + #[test] + fn test_datasketches_kll_query() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + + let mut query_kwargs = HashMap::new(); + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + assert!( + (result - 6.0).abs() <= 1.0, + "expected median ~6, got {result}" + ); + + assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); + } + + #[test] + fn test_datasketches_kll_merge() { + let mut kll1 = DatasketchesKLLAccumulator::new(200); + let mut kll2 = DatasketchesKLLAccumulator::new(200); + + for i in 1..=5 { + kll1.update(i as f64); + } + for i in 6..=10 { + kll2.update(i as f64); + } + + let merged = DatasketchesKLLAccumulator::merge_accumulators(vec![kll1, kll2]).unwrap(); + assert_eq!(merged.inner.count(), 10); + assert_eq!(merged.get_quantile(0.0), 1.0); + assert_eq!(merged.get_quantile(1.0), 10.0); + } + + #[test] + fn test_datasketches_kll_get_keys() { + let kll = DatasketchesKLLAccumulator::new(200); + assert_eq!(kll.type_name(), "DatasketchesKLLAccumulator"); + } + + #[test] + fn test_trait_object() { + let mut kll = DatasketchesKLLAccumulator::new(200); + kll.update(5.0); + let trait_obj: Box = Box::new(kll); + assert_eq!(trait_obj.type_name(), "DatasketchesKLLAccumulator"); + } + + #[test] + fn test_datasketches_kll_query_with_kwargs() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + + let mut query_kwargs = HashMap::new(); + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + assert!( + (result - 6.0).abs() <= 1.0, + "expected median ~6, got {result}" + ); + + query_kwargs.insert("quantile".to_string(), "0.9".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.9 quantile of 1..10 may be 9 or 10. + assert!( + (9.0..=10.0).contains(&result), + "expected 0.9 quantile in [9,10], got {result}" + ); + + query_kwargs.insert("quantile".to_string(), "0.0".to_string()); + assert_eq!( + kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), + 1.0 + ); + + query_kwargs.insert("quantile".to_string(), "1.0".to_string()); + assert_eq!( + kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), + 10.0 + ); + + assert!(kll.query(Statistic::Quantile, None).is_err()); + + query_kwargs.insert("quantile".to_string(), "invalid".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "1.5".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "-0.1".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); + } + + #[test] + fn test_datasketches_kll_merge_multiple() { + let mut kll1 = DatasketchesKLLAccumulator::new(200); + let mut kll2 = DatasketchesKLLAccumulator::new(200); + let mut kll3 = DatasketchesKLLAccumulator::new(200); + + for i in 1..=5 { + kll1.update(i as f64); + } + for i in 6..=10 { + kll2.update(i as f64); + } + for i in 11..=15 { + kll3.update(i as f64); + } + + let boxed_accs: Vec> = + vec![Box::new(kll1), Box::new(kll2), Box::new(kll3)]; + + let merged = DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).unwrap(); + assert_eq!(merged.inner.count(), 15); + assert_eq!(merged.get_quantile(0.0), 1.0); + assert_eq!(merged.get_quantile(1.0), 15.0); + assert_eq!(merged.get_quantile(0.5), 8.0); + } + + #[test] + fn test_datasketches_kll_merge_multiple_error_cases() { + let empty: Vec> = vec![]; + assert!(DatasketchesKLLAccumulator::merge_multiple(&empty).is_err()); + + let kll1 = DatasketchesKLLAccumulator::new(200); + let kll2 = DatasketchesKLLAccumulator::new(100); + let boxed_accs: Vec> = vec![Box::new(kll1), Box::new(kll2)]; + assert!(DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).is_err()); + + use crate::summary_kernels::sum::SumAccumulator; + let kll = DatasketchesKLLAccumulator::new(200); + let sum = SumAccumulator::new(); + let mixed_accs: Vec> = vec![Box::new(kll), Box::new(sum)]; + assert!(DatasketchesKLLAccumulator::merge_multiple(&mixed_accs).is_err()); + } + + #[test] + fn test_from_sketchlib_proto_bytes_reconstructs_quantiles() { + // Build a KllState with 64 items in level order; the decoder + // replays every item through `update()` so the reconstructed + // sketch is statistically equivalent — quantile estimates + // match the ground truth (sorted items) within KLL's own + // rank-error bound for k=200. + use asap_sketchlib::proto::sketchlib::KllState; + + let items: Vec = (0..64).map(|i| i as f64).collect(); + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 64], + items: items.clone(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + + let acc = + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.count(), 64); + // For 64 values 0..63, the true median is 31.5 and quantile + // error is ~1% × range = 0.63. KLL's own point query can + // legally be off by up to ε × N ~= 0.01 × 64 = 0.64. Allow a + // generous tolerance since the important invariant is "the + // decoded sketch is queryable and returns a sensible value". + let median = acc.get_quantile(0.5); + assert!( + (median - 31.5).abs() <= 10.0, + "reconstructed median {median} is outside tolerance of true median 31.5" + ); + let q01 = acc.get_quantile(0.01); + let q99 = acc.get_quantile(0.99); + assert!( + q01 <= q99, + "quantile monotonicity violated: q01={q01}, q99={q99}" + ); + } + + // Compacted portable state is highest-level first, unlike the runtime buffer. + #[test] + fn compacted_wire_state_preserves_count_and_quantiles() { + use asap_sketchlib::{proto::sketchlib::KllState, sketches::KLL}; + let mut source = KLL::::init_kll_with_seed(32, 123); + for i in 0..1000 { + source.update(&(((i * 7919 + 17) % 1009) as f64 / 1009.0)); + } + assert!(source.wire_num_levels() > 1); + let state = KllState { + k: 32, + m: source.wire_m(), + num_levels: source.wire_num_levels(), + levels: source.wire_levels(), + items: source.wire_items(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: vec![], + }; + let decoded = + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&encode_state(state)).unwrap(); + assert_eq!(decoded.inner.count(), source.count() as u64); + for q in [0.0, 0.1, 0.5, 0.9, 1.0] { + assert_eq!(decoded.inner.quantile(q), source.quantile(q), "q={q}"); + } + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's kllprocessor emits: the state + // wrapped in a `SketchEnvelope{kll: ...}` via sketchlib-go's + // `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + + let items: Vec = (0..64).map(|i| i as f64).collect(); + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 64], + items, + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + assert_eq!(acc.inner.count(), 64); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, CountMinState, SketchEnvelope}; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::CountMin( + CountMinState::default(), + )), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_small_k() { + use asap_sketchlib::proto::sketchlib::KllState; + let state = KllState { + k: 4, // < minimum of 8 + m: 2, + num_levels: 0, + levels: Vec::new(), + items: Vec::new(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("k must be >= 8")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_inconsistent_levels() { + use asap_sketchlib::proto::sketchlib::KllState; + // num_levels=1 but levels array has 3 entries instead of 2 + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 5, 10], + items: vec![1.0, 2.0, 3.0, 4.0, 5.0], + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("levels length")); + } + + #[test] + fn aux_stats_exposes_count_via_kll_n() { + let mut acc = DatasketchesKLLAccumulator::new(200); + for i in 0..50 { + acc.update(i as f64); + } + let aux = acc.aux_stats(); + assert_eq!(aux.count, Some(50)); + // KLL doesn't natively expose min/max cheaply and doesn't + // track sum at all — those fields must be None so callers + // fall through to query_statistic. + assert_eq!(aux.sum, None); + assert_eq!(aux.min, None); + assert_eq!(aux.max, None); + } + + #[test] + fn aux_stats_empty_kll_has_zero_count() { + let acc = DatasketchesKLLAccumulator::new(200); + assert_eq!(acc.aux_stats().count, Some(0)); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/dd_sketch.rs b/crates/asap_summary_state/src/summary_kernels/dd_sketch.rs new file mode 100644 index 000000000..1c25d940c --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/dd_sketch.rs @@ -0,0 +1,726 @@ +//! DDSketch accumulator — wraps `asap_sketchlib::DdSketch`. +//! +//! Concrete accumulator reached from the modified-OTLP +//! `Metric.data = DDSketch{…}` hot path (PR C-CountSketch follow-up). +//! Merge via bucket-index alignment on the inner sketch, serialize as +//! MessagePack for the sink, and decode from the sketchlib +//! `DDSketchState` proto. +//! +//! Query semantics follow the STRICT policy after the DataPoint-level +//! METRIC scalars were dropped from the wire format +//! (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57): the sketch serves +//! Quantile (log-bucket estimation) and Count (sum of bucket counts). +//! Sum/Min/Max are no longer derivable from the wire bytes and are +//! served by controller-provisioned exact aggregations — `query_statistic` +//! returns the unavailable-statistic error for them. + +use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::{DdSketch, DdSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +/// DDSketch accumulator — inner log-bucketed sketch. +#[derive(Debug, Clone)] +pub struct DDSketchAccumulator { + pub inner: DdSketch, + /// Edge sampling probability `p ∈ (0,1]` carried on the producer's + /// `SketchEnvelope.sample_p`. The edge admits each value with probability + /// `p` (NitroSketch geometric skip), so `inner.total_count()` is ~`p`× the + /// true count and a `Count` query must rescale by `1/p`. Quantiles are + /// rank-preserving and need NO rescale. `1.0` (and the proto3 default `0.0`, + /// dual-read as `1.0`) means no sampling, so the rescale is a no-op and the + /// behaviour is identical to before. The factor is a per-series config + /// constant: it is set from the first (always-full, otel.rs ingest + /// contract) frame and preserved across delta applies, window-boundary + /// `reset_to_empty`, and `merge_with`. + pub sample_p: f64, +} + +/// Normalize a wire `sample_p` to a usable rescale denominator. `0.0` (proto3 +/// default), `>= 1.0`, and non-finite all collapse to `1.0` (no sampling), so a +/// `Count` rescale by `1/p` is a no-op on unsampled / legacy frames. +pub(crate) fn normalize_sample_p(p: f64) -> f64 { + if p.is_finite() && p > 0.0 && p < 1.0 { + p + } else { + 1.0 + } +} + +impl DDSketchAccumulator { + pub fn new(alpha: f64) -> Self { + Self { + inner: DdSketch::new(alpha), + sample_p: 1.0, + } + } + + /// Read the normalized edge sampling probability from a full-frame + /// `SketchEnvelope`'s `sample_p`. Returns `1.0` (no sampling) for bare + /// `DdSketchState` bytes or any decode failure — the primary production + /// decode path (`reconstruct_via_runtime`) discards the envelope's + /// `sample_p`, so the ingest call site re-reads it from the same bytes. + pub fn sample_p_from_envelope_bytes(buffer: &[u8]) -> f64 { + use asap_sketchlib::proto::sketchlib::SketchEnvelope; + use prost::Message; + SketchEnvelope::decode(buffer) + .map(|env| normalize_sample_p(env.sample_p)) + .unwrap_or(1.0) + } + + /// Decode from the modified OTLP wire format's + /// `DDSketchDataPoint.sketch` bytes when + /// `encoding = DDSKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `DdSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: DdSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize DdSketch msgpack: {e}"))?, + // The msgpack DdSketch struct carries no envelope/sample_p; the + // msgpack path is parity/test-only and is never edge-sampled. + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `DDSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::DDSketchState` message that + /// DataCollector's `ddsketchprocessor` emits when + /// `encoding = DD_SKETCH_ENCODING_PROTO`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + let (state, sample_p) = asap_sketch_codec::ddsketch_state(buffer)?; + if !(state.alpha > 0.0 && state.alpha < 1.0) { + return Err(format!( + "DDSketchState alpha {} out of range (expected 0 < alpha < 1)", + state.alpha + ) + .into()); + } + // Preserve positive, negative and zero stores from the sketchlib wire state. + let inner = DdSketch::from_proto(state); + Ok(Self { + inner, + sample_p: normalize_sample_p(sample_p), + }) + } + + /// Apply a proto-encoded `DDSketchDelta` frame to this + /// accumulator's inner sketch — the decode path for + /// `DD_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + /// + /// Called against an accumulator that already carries the base + /// sketch state; the caller is the per-series snapshot cache in + /// the ingest path. Bytes are the + /// `asap_sketchlib::proto::sketchlib::DdSketchDelta` message. + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + use asap_sketchlib::proto::sketchlib::DdSketchDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode DDSketchDelta: {e}"))?; + + // The delta no longer carries d_count/d_sum/min/max + // (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). Apply the + // bucket deltas only; `DdSketch` recomputes its total count from + // the merged bucket counts (`total_count()`). + let buckets = pb + .buckets + .into_iter() + .map(|b| (b.index, b.d_count)) + .collect(); + let delta = DdSketchDelta { + buckets, + negative_buckets: pb + .negative_buckets + .into_iter() + .map(|b| (b.index, b.d_count)) + .collect(), + zero_count: pb.zero_count, + ..Default::default() + }; + self.inner + .apply_delta(&delta) + .map_err(|error| format!("apply DDSketchDelta: {error}"))?; + Ok(()) + } +} + +impl SerializableToSink for DDSketchAccumulator { + fn serialize_to_json(&self) -> Value { + // The DataPoint-level scalars (sum/min/max) are no longer carried + // by `DdSketch` (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). + // `count` is the bucket-derived total via `total_count()`. + serde_json::json!({ + "alpha": self.inner.alpha, + "store_offset": self.inner.store_offset, + "bucket_count": self.inner.store_counts.len(), + // Raw bucket-derived count (admitted samples). `sample_p` is the + // scale factor a consumer applies (count / sample_p) to estimate + // the true count; `query_statistic(Count)` already does this. + "count": self.inner.total_count(), + "sample_p": self.sample_p, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for DDSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "DDSketchAccumulator" + } + + /// Per-window base rotation: drop all bucket counts but keep the + /// relative-accuracy parameter so the next window's bucket deltas + /// index into the same log-bucket layout. `sample_p` is a per-series + /// config constant (not per-window data), so it is intentionally + /// preserved across the rotation — the next window's deltas are sampled + /// at the same rate and must rescale identically. + fn reset_to_empty(&mut self) { + self.inner = DdSketch::new(self.inner.alpha); + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge DDSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_dd = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to DDSketchAccumulator")?; + let merged_inner = DdSketch::merge_refs(&[&self.inner, &other_dd.inner])?; + // sample_p is a per-series config constant, so both operands carry the + // same value in practice. Prefer a sampled factor over the no-sampling + // default so a merge with a freshly-reset (1.0) base keeps the series' + // sampling rate. + let sample_p = if self.sample_p < 1.0 { + self.sample_p + } else { + other_dd.sample_p + }; + Ok(Box::new(Self { + inner: merged_inner, + sample_p, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::DDSketch + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + query_kwargs: &HashMap, + ) -> Result> { + use crate::Statistic; + + match statistic { + Statistic::Quantile => { + // PromQL `histogram_quantile(q, …)` and + // `quantile_over_time(q, …)` both land here with + // `q` in `query_kwargs["quantile"]`. Default to + // 0.99 when the caller didn't provide one + // (defensive — pattern-matched queries in + // `inference_config.yaml` always populate it). + let q: f64 = query_kwargs + .get("quantile") + .and_then(|s| s.parse().ok()) + .unwrap_or(0.99); + if !(0.0..=1.0).contains(&q) { + return Err(format!("DDSketchAccumulator: quantile {q} out of [0,1]").into()); + } + self.inner.quantile(q).ok_or_else(|| { + "DDSketchAccumulator: quantile() returned None (sketch empty?)".into() + }) + } + // Count is derived by summing the bucket store counts — the only + // DataPoint-level scalar that survives the wire-format trim + // (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). When the edge + // sampled this series (sample_p < 1.0), the stored count is ~p× the + // true count, so rescale by 1/sample_p to recover an unbiased + // estimate. sample_p == 1.0 (unsampled / legacy) makes this a no-op. + Statistic::Count => Ok(self.inner.total_count() as f64 / self.sample_p), + // STRICT policy: the Sum/Min/Max scalars were removed from + // the DDSketch wire format. They are now served by the + // controller-provisioned exact aggregations (an exact `Sum` + // and an exact `MinMax`), NOT estimated from the buckets. + // Surface the unavailable-statistic error so the query path + // routes to those aggregations instead of returning a wrong + // (0 / panicked) value. + Statistic::Sum => Err( + "DDSketchAccumulator: Sum not available from DDSketch wire format \ + (ProjectASAP/sketchlib-go#243); use an exact Sum aggregation" + .into(), + ), + Statistic::Min | Statistic::Max => Err(format!( + "DDSketchAccumulator: {statistic:?} not available from DDSketch wire format \ + (ProjectASAP/sketchlib-go#243); use an exact MinMax aggregation", + ) + .into()), + other => Err(format!( + "DDSketchAccumulator: statistic {other:?} not supported (only Quantile / Count)", + ) + .into()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // The DataPoint-level METRIC scalars (count/sum/min/max) were dropped + // from `DdSketchState` (ProjectASAP/sketchlib-go#243 / + // asap_sketchlib#57); the proto now carries only + // `alpha`/`store_counts`/`store_offset`. + fn encode_state(alpha: f64, store_counts: Vec, store_offset: i32) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + let state = DdSketchState { + alpha, + store_counts, + store_offset, + ..Default::default() + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), + ..Default::default() + } + .encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_round_trip() { + let bytes = encode_state(0.01, vec![1, 2, 3, 4], -2); + let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.alpha, 0.01); + assert_eq!(acc.inner.store_counts, vec![1, 2, 3, 4]); + assert_eq!(acc.inner.store_offset, -2); + // `count` is recovered by summing the bucket store counts. + assert_eq!(acc.inner.total_count(), 10); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_invalid_alpha() { + let bytes = encode_state(0.0, vec![1], 0); + let result = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("alpha")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's ddsketchprocessor emits: the + // state wrapped in a `SketchEnvelope{ddsketch: ...}` via + // sketchlib-go's `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + + let state = DdSketchState { + alpha: 0.01, + store_counts: vec![1, 2, 3, 4], + store_offset: -2, + ..Default::default() + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + assert_eq!(acc.inner.alpha, 0.01); + assert_eq!(acc.inner.total_count(), 10); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_aggregate_core_merge_aligns_buckets() { + let a = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 1, 1], -1), + sample_p: 1.0, + }; + let b = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![10, 10, 10], 0), + sample_p: 1.0, + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.inner.store_counts, vec![1, 11, 11, 10]); + assert_eq!(merged.inner.store_offset, -1); + assert_eq!(merged.inner.total_count(), 33); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::summary_kernels::count_sketch::CountSketchAccumulator; + let dd = DDSketchAccumulator::new(0.01); + let cs = CountSketchAccumulator::new(2, 3); + assert!(dd.merge_with(&cs).is_err()); + } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = DdSketch::from_raw(0.01, vec![5, 10, 15, 20], -2); + let bytes = original.to_msgpack().unwrap(); + let acc = DDSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.alpha, 0.01); + assert_eq!(acc.inner.store_counts, vec![5, 10, 15, 20]); + assert_eq!(acc.inner.store_offset, -2); + // `count` is recovered by summing the bucket store counts. + assert_eq!(acc.inner.total_count(), 50); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = DDSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } + + #[test] + fn test_apply_proto_delta_bytes_round_trip() { + use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use prost::Message; + + let mut acc = DDSketchAccumulator::new(0.01); + acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); + + // The wire delta now carries only bucket deltas (tags 2-7 + // reserved); `DdSketchBucketDelta` has just `index` + `d_count`. + let bytes = PbDelta { + buckets: vec![ + DdSketchBucketDelta { + index: 0, + d_count: 10, + }, + DdSketchBucketDelta { + index: 2, + d_count: 20, + }, + ], + ..Default::default() + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!(acc.inner.store_counts, vec![11, 2, 23]); + // `count` recomputed from the merged buckets: 11 + 2 + 23 = 36. + assert_eq!(acc.inner.total_count(), 36); + } + + /// A valid protobuf with an inadmissible span must not acknowledge a dropped update. + #[test] + fn test_apply_proto_delta_rejects_span_without_mutating_state() { + use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use prost::Message; + let mut acc = DDSketchAccumulator::new(0.01); + acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); + let bytes = PbDelta { + buckets: vec![DdSketchBucketDelta { + index: i32::MAX, + d_count: 1, + }], + ..Default::default() + } + .encode_to_vec(); + assert!(acc.apply_proto_delta_bytes(&bytes).is_err()); + assert_eq!(acc.inner.store_counts, vec![1, 2, 3]); + assert_eq!(acc.inner.store_offset, 0); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = DDSketchAccumulator::new(0.01); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ----- query_statistic STRICT policy ----- + // + // After the DataPoint-level METRIC scalars were dropped from the + // DDSketch wire format (ProjectASAP/sketchlib-go#243 / + // asap_sketchlib#57), DDSketch serves only quantiles and Count. + // Sum/Min/Max move to controller-provisioned exact aggregations and + // MUST surface the unavailable-statistic error (never a panic / 0). + + fn sample_accumulator() -> DDSketchAccumulator { + // Build the in-memory sketch from bucket counts only — no scalars. + DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), + sample_p: 1.0, + } + } + + #[test] + fn test_query_statistic_quantile_is_sketch_derived() { + use crate::Statistic; + let acc = sample_accumulator(); + let mut kwargs = HashMap::new(); + kwargs.insert("quantile".to_string(), "0.5".to_string()); + let v = acc + .query_statistic(Statistic::Quantile, &None, &kwargs) + .expect("quantile should be served from the sketch buckets"); + assert!( + v.is_finite() && v > 0.0, + "quantile estimate should be positive finite, got {v}" + ); + } + + #[test] + fn test_query_statistic_count_is_bucket_derived() { + use crate::Statistic; + let acc = sample_accumulator(); + let v = acc + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count should be derivable from the bucket store"); + // 1 + 2 + 3 + 4 = 10. + assert_eq!(v, 10.0); + } + + #[test] + fn test_query_statistic_sum_min_max_return_unavailable_error() { + use crate::Statistic; + let acc = sample_accumulator(); + for stat in [Statistic::Sum, Statistic::Min, Statistic::Max] { + let result = acc.query_statistic(stat, &None, &HashMap::new()); + assert!( + result.is_err(), + "{stat:?} must return the unavailable-statistic error (not a panic / 0)" + ); + let msg = result.unwrap_err().to_string(); + assert!( + msg.contains("not available"), + "{stat:?} error should explain the statistic is unavailable, got: {msg}" + ); + } + } + + // ----- sample_p count rescale ----- + // + // When the edge sampled a DDSketch (sample_p < 1.0), the stored count is + // ~p× the true count, so Count rescales by 1/p. Quantiles are + // rank-preserving and must NOT be rescaled. + + #[test] + fn test_count_is_rescaled_by_sample_p() { + use crate::Statistic; + let acc = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), + sample_p: 0.1, + }; + let c = acc + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + // Raw bucket sum 10, rescaled by 1/0.1 = 100. + assert!((c - 100.0).abs() < 1e-9, "expected rescaled 100, got {c}"); + } + + #[test] + fn test_quantile_ignores_sample_p() { + use crate::Statistic; + let mut kwargs = HashMap::new(); + kwargs.insert("quantile".to_string(), "0.5".to_string()); + let unsampled = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), + sample_p: 1.0, + }; + let sampled = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), + sample_p: 0.1, + }; + let qu = unsampled + .query_statistic(Statistic::Quantile, &None, &kwargs) + .expect("q ok"); + let qs = sampled + .query_statistic(Statistic::Quantile, &None, &kwargs) + .expect("q ok"); + assert_eq!(qu, qs, "quantile must be sample_p-invariant"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { + use crate::Statistic; + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + + let env = SketchEnvelope { + sample_p: 0.25, + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(DdSketchState { + alpha: 0.01, + store_counts: vec![2, 4, 6, 8], + store_offset: -2, + ..Default::default() + })), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 0.25); + // Raw 20, rescaled 20 / 0.25 = 80. + let c = acc + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + assert!((c - 80.0).abs() < 1e-9, "expected rescaled 80, got {c}"); + } + + #[test] + fn test_sample_p_normalization() { + // proto3 default (0.0), >=1.0, and non-finite all mean no sampling. + assert_eq!(normalize_sample_p(0.0), 1.0); + assert_eq!(normalize_sample_p(1.0), 1.0); + assert_eq!(normalize_sample_p(1.5), 1.0); + assert_eq!(normalize_sample_p(f64::NAN), 1.0); + assert_eq!(normalize_sample_p(-0.1), 1.0); + assert_eq!(normalize_sample_p(0.5), 0.5); + } + + #[test] + fn test_sample_p_from_envelope_bytes_defaults_to_one() { + use asap_sketchlib::proto::sketchlib::DdSketchState; + use prost::Message; + // Bare DdSketchState bytes (no envelope) → no sampling info → 1.0. + let bare = DdSketchState { + alpha: 0.01, + store_counts: vec![1, 2, 3], + store_offset: 0, + ..Default::default() + } + .encode_to_vec(); + assert_eq!( + DDSketchAccumulator::sample_p_from_envelope_bytes(&bare), + 1.0 + ); + } + + #[test] + fn test_reset_to_empty_preserves_sample_p() { + let mut acc = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3], 0), + sample_p: 0.2, + }; + acc.reset_to_empty(); + assert_eq!(acc.sample_p, 0.2, "window rotation must keep sample_p"); + assert_eq!(acc.inner.total_count(), 0, "buckets cleared"); + } + + #[test] + fn test_merge_prefers_sampled_factor() { + // A sampled base merged with a freshly-reset (1.0) operand keeps the + // series' sampling rate. + let a = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 1, 1], 0), + sample_p: 0.1, + }; + let b = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 1, 1], 0), + sample_p: 1.0, + }; + let merged = a.merge_with(&b).expect("merge ok"); + let merged = merged + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.sample_p, 0.1); + } +} + +#[cfg(test)] +mod dependency_upgrade_tests { + use super::*; + // The upgraded sketchlib state must retain negative and zero stores through both adapters. + #[test] + fn signed_state_survives_codec_and_accumulator_roundtrip() { + let mut inner = DdSketch::new(0.01); + for value in [-4.0, 0.0, 8.0] { + inner.update(value); + } + let bytes = asap_sketch_codec::encode_ddsketch(&inner); + let (wire, _) = asap_sketch_codec::ddsketch_state(&bytes).unwrap(); + assert_eq!(wire.zero_count, 1); + assert_eq!(wire.negative_store_counts.iter().sum::(), 1); + let restored = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).unwrap(); + assert_eq!(restored.inner.total_count(), 3); + assert_eq!(restored.inner.alpha, inner.wire_alpha()); + assert_eq!(restored.inner.store_counts, inner.store_counts); + assert_eq!(restored.inner.store_offset, inner.store_offset); + assert_eq!( + restored.inner.negative_store_counts, + inner.negative_store_counts + ); + assert_eq!( + restored.inner.negative_store_offset, + inner.negative_store_offset + ); + assert_eq!(restored.inner.zero_count, inner.zero_count); + } + // Negative and zero delta fields added by sketchlib must not be discarded by the adapter. + #[test] + fn signed_delta_survives_adapter() { + use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use prost::Message; + let mut accumulator = DDSketchAccumulator::new(0.01); + let bytes = PbDelta { + negative_buckets: vec![DdSketchBucketDelta { + index: 0, + d_count: 2, + }], + zero_count: 3, + ..Default::default() + } + .encode_to_vec(); + accumulator.apply_proto_delta_bytes(&bytes).unwrap(); + assert_eq!(accumulator.inner.total_count(), 5); + assert_eq!(accumulator.inner.zero_count, 3); + assert_eq!( + accumulator.inner.negative_store_counts.iter().sum::(), + 2 + ); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/exact.rs b/crates/asap_summary_state/src/summary_kernels/exact.rs new file mode 100644 index 000000000..8652bec47 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/exact.rs @@ -0,0 +1,350 @@ +//! Exact summary state identified by Planner family, independent of keyed layout. +use super::increase::IncreaseAccumulator; +use crate::Statistic; +use crate::{ + AggregateCore, AggregationType, AuxStats, KeyByLabelValues, Measurement, SerializableToSink, +}; +use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +type Error = Box; + +#[derive(Debug, Clone, Serialize, Deserialize)] +enum ScalarState { + Sum(f64), + Count(u64), + Min(Option), + Max(Option), + Counter(Option), +} + +/// Both the family and population layout survive persistence. Sharing counter +/// arithmetic never authorizes a Rate state to answer an Increase readout. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExactAccumulator { + family: SummaryFamilyType, + scalar: ScalarState, + keyed: Option>, +} + +impl ExactAccumulator { + /// Accumulate into run-local scratch state. Persistent input states remain + /// immutable; a failed merge discards this scratch state. + pub(crate) fn merge_from(&mut self, other: &Self) -> Result<(), Error> { + if self.family != other.family || self.is_keyed() != other.is_keyed() { + return Err("cannot merge different Planner families or layouts".into()); + } + if let (Some(target), Some(source)) = (&mut self.keyed, &other.keyed) { + for (key, state) in source { + let combined = match target.get(key) { + Some(old) => merge_scalar(old, state)?, + None => state.clone(), + }; + target.insert(key.clone(), combined); + } + } else { + self.scalar = merge_scalar(&self.scalar, &other.scalar)?; + } + Ok(()) + } + + pub fn new(family: SummaryFamilyType, keyed: bool) -> Result { + use ExactKind as K; + use ExactParams as P; + let scalar = match &family { + SummaryFamilyType::ExactAggregate(K::Sum, P::Sum) => ScalarState::Sum(0.0), + SummaryFamilyType::ExactAggregate(K::Count, P::Count) => ScalarState::Count(0), + SummaryFamilyType::ExactAggregate(K::Min, P::Min) => ScalarState::Min(None), + SummaryFamilyType::ExactAggregate(K::Max, P::Max) => ScalarState::Max(None), + SummaryFamilyType::ExactAggregate(K::Rate, P::Rate) + | SummaryFamilyType::ExactAggregate(K::Increase, P::Increase) => { + ScalarState::Counter(None) + } + _ => return Err(format!("unsupported exact Planner family: {family:?}")), + }; + Ok(Self { + family, + scalar, + keyed: keyed.then(HashMap::new), + }) + } + + pub fn family(&self) -> &SummaryFamilyType { + &self.family + } + pub(crate) fn insufficient_counter_samples( + &self, + statistic: Statistic, + key: &Option, + ) -> bool { + if statistic != self.statistic() { + return false; + } + let state = match (&self.keyed, key) { + (Some(states), Some(key)) => states.get(key), + (None, None) => Some(&self.scalar), + _ => None, + }; + match state { + Some(ScalarState::Counter(None)) => true, + Some(ScalarState::Counter(Some(counter))) => { + counter.sample_count < 2 + || counter.last_seen_timestamp == counter.starting_timestamp + } + _ => false, + } + } + pub fn is_keyed(&self) -> bool { + self.keyed.is_some() + } + + pub fn update(&mut self, key: Option<&KeyByLabelValues>, value: f64, timestamp: i64) { + let state = match (&mut self.keyed, key) { + (Some(states), Some(key)) => states + .entry(key.clone()) + .or_insert_with(|| self.scalar.clone()), + (None, None) => &mut self.scalar, + _ => panic!("exact update population layout differs from installed DAG"), + }; + match state { + ScalarState::Sum(sum) => *sum += value, + ScalarState::Count(count) => { + *count = count.checked_add(1).expect("exact count overflow") + } + ScalarState::Min(current) => { + *current = Some(current.map_or(value, |old| old.min(value))) + } + ScalarState::Max(current) => { + *current = Some(current.map_or(value, |old| old.max(value))) + } + ScalarState::Counter(current) => match current { + Some(counter) => counter.update(Measurement::new(value), timestamp), + None => { + *current = Some(IncreaseAccumulator::new( + Measurement::new(value), + timestamp, + Measurement::new(value), + timestamp, + )) + } + }, + } + } + + pub fn deserialize_from_bytes(bytes: &[u8]) -> Result { + let state: Self = rmp_serde::from_slice(bytes)?; + let expected = Self::new(state.family.clone(), state.is_keyed())?; + let same_variant = |value: &ScalarState| { + std::mem::discriminant(value) == std::mem::discriminant(&expected.scalar) + }; + if !same_variant(&state.scalar) + || state + .keyed + .as_ref() + .is_some_and(|states| states.values().any(|s| !same_variant(s))) + { + return Err("exact payload differs from declared Planner family".into()); + } + Ok(state) + } + + fn statistic(&self) -> Statistic { + match self.family { + SummaryFamilyType::ExactAggregate(ExactKind::Sum, _) => Statistic::Sum, + SummaryFamilyType::ExactAggregate(ExactKind::Count, _) => Statistic::Count, + SummaryFamilyType::ExactAggregate(ExactKind::Min, _) => Statistic::Min, + SummaryFamilyType::ExactAggregate(ExactKind::Max, _) => Statistic::Max, + SummaryFamilyType::ExactAggregate(ExactKind::Rate, _) => Statistic::Rate, + SummaryFamilyType::ExactAggregate(ExactKind::Increase, _) => Statistic::Increase, + _ => unreachable!("validated exact family"), + } + } +} + +fn merge_scalar(left: &ScalarState, right: &ScalarState) -> Result { + Ok(match (left, right) { + (ScalarState::Sum(a), ScalarState::Sum(b)) => ScalarState::Sum(a + b), + (ScalarState::Count(a), ScalarState::Count(b)) => { + ScalarState::Count(a.checked_add(*b).ok_or("exact count overflow")?) + } + (ScalarState::Min(a), ScalarState::Min(b)) => { + ScalarState::Min(a.iter().chain(b).copied().reduce(f64::min)) + } + (ScalarState::Max(a), ScalarState::Max(b)) => { + ScalarState::Max(a.iter().chain(b).copied().reduce(f64::max)) + } + (ScalarState::Counter(a), ScalarState::Counter(b)) => ScalarState::Counter(match (a, b) { + (Some(a), Some(b)) => Some(IncreaseAccumulator::merge_pair(a, b)), + (a, b) => a.clone().or_else(|| b.clone()), + }), + _ => return Err("exact scalar state families differ".into()), + }) +} + +impl SerializableToSink for ExactAccumulator { + fn serialize_to_json(&self) -> serde_json::Value { + serde_json::json!({"family": self.family, "scalar": self.scalar, "keyed": self.keyed.as_ref().map(|m|m.iter().collect::>())}) + } + fn serialize_to_bytes(&self) -> Vec { + rmp_serde::to_vec_named(self).expect("exact state encoding") + } +} + +impl AggregateCore for ExactAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + fn type_name(&self) -> &'static str { + "PlannerExactAccumulatorV1" + } + 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, Error> { + let other = other + .as_any() + .downcast_ref::() + .ok_or("merge requires Planner exact state")?; + let mut merged = self.clone(); + merged.merge_from(other)?; + Ok(Box::new(merged)) + } + fn get_accumulator_type(&self) -> AggregationType { + match self.statistic() { + Statistic::Sum => AggregationType::Sum, + Statistic::Count => AggregationType::Count, + Statistic::Min => AggregationType::Min, + Statistic::Max => AggregationType::Max, + Statistic::Rate => AggregationType::Rate, + Statistic::Increase => AggregationType::Increase, + _ => unreachable!(), + } + } + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::() + + self.keyed.as_ref().map_or(0, |m| { + m.keys() + .map(|k| { + std::mem::size_of::() + + k.labels.iter().map(String::len).sum::() + }) + .sum::() + }) + } + fn aux_stats(&self) -> AuxStats { + if self.is_keyed() { + return AuxStats::empty(); + } + match self.scalar { + ScalarState::Sum(value) => AuxStats { + sum: Some(value), + ..AuxStats::empty() + }, + ScalarState::Count(value) => AuxStats { + count: Some(value), + ..AuxStats::empty() + }, + ScalarState::Min(value) => AuxStats { + min: value, + ..AuxStats::empty() + }, + ScalarState::Max(value) => AuxStats { + max: value, + ..AuxStats::empty() + }, + ScalarState::Counter(_) => AuxStats::empty(), + } + } + fn get_keys(&self) -> Option> { + self.keyed.as_ref().map(|m| m.keys().cloned().collect()) + } + fn query_statistic( + &self, + statistic: Statistic, + key: &Option, + kwargs: &HashMap, + ) -> Result { + if statistic != self.statistic() { + return Err("readout differs from Planner exact family".into()); + } + let state = match (&self.keyed, key) { + (Some(states), Some(key)) => states.get(key).ok_or("unknown exact population")?, + (None, None) => &self.scalar, + _ => return Err("readout population differs from installed layout".into()), + }; + match state { + ScalarState::Sum(sum) => Ok(*sum), + ScalarState::Count(count) => Ok(*count as f64), + ScalarState::Min(value) | ScalarState::Max(value) => { + value.ok_or_else(|| "empty exact population".into()) + } + ScalarState::Counter(Some(counter)) => { + counter.query_statistic(statistic, &None, kwargs) + } + ScalarState::Counter(None) => Err("empty counter population".into()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Identity, population isolation, and readout survive the persisted format. + #[test] + fn exact_families_roundtrip_and_reject_cross_family_operations() { + let families = [ + (ExactKind::Sum, ExactParams::Sum, Statistic::Sum, 16.0), + (ExactKind::Count, ExactParams::Count, Statistic::Count, 3.0), + (ExactKind::Min, ExactParams::Min, Statistic::Min, 2.0), + (ExactKind::Max, ExactParams::Max, Statistic::Max, 8.0), + (ExactKind::Rate, ExactParams::Rate, Statistic::Rate, 3.0), + ( + ExactKind::Increase, + ExactParams::Increase, + Statistic::Increase, + 6.0, + ), + ]; + for keyed in [false, true] { + let key = keyed.then(|| KeyByLabelValues::new_with_labels(vec!["a".into()])); + let mut states = Vec::new(); + for (kind, params, stat, value) in &families { + let mut state = ExactAccumulator::new( + SummaryFamilyType::ExactAggregate(kind.clone(), params.clone()), + keyed, + ) + .unwrap(); + for (ts, v) in [(1000, 8.0), (2000, 2.0), (3000, 6.0)] { + state.update(key.as_ref(), v, ts); + } + let restored = + ExactAccumulator::deserialize_from_bytes(&state.serialize_to_bytes()).unwrap(); + assert_eq!(restored.family(), state.family()); + assert_eq!( + restored + .query_statistic(*stat, &key, &HashMap::new()) + .unwrap(), + *value + ); + for (_, _, wrong, _) in &families { + if wrong != stat { + assert!(restored + .query_statistic(*wrong, &key, &HashMap::new()) + .is_err()); + } + } + states.push(restored); + } + for (i, a) in states.iter().enumerate() { + for (j, b) in states.iter().enumerate() { + assert_eq!(a.merge_with(b).is_ok(), i == j); + } + } + } + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/factory.rs b/crates/asap_summary_state/src/summary_kernels/factory.rs new file mode 100644 index 000000000..c4e7ef829 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/factory.rs @@ -0,0 +1,1176 @@ +use crate::summary_kernels::{ + CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, + CountSketchWithHeapAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, + HydraKllSketchAccumulator, IncreaseAccumulator, KeyedCounterState, KeyedMaxState, + KeyedMinState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, +}; +use crate::{AggregateCore, KeyByLabelValues, Measurement}; +// Production dispatch consumes Planner SummaryAgg payloads directly. The +// config adapter below is compiled only for isolated historical kernel tests. +use crate::summary_kernels::hll_sketch::HllSketchAccumulator; +use crate::summary_kernels::univmon::UnivMonAccumulator; +use planner_types::post_asap::{ExactKind, SketchAlgorithm, SketchParams, SummaryFamilyType}; + +/// Generate the two boilerplate clone-based `AccumulatorUpdater` methods +/// for updaters whose inner `acc` field implements `Clone + AggregateCore`. +/// Not applicable to `IncreaseAccumulatorUpdater` (its `acc` is `Option<_>` +/// with non-trivial `None` handling). +macro_rules! impl_clone_accumulator_methods { + ($acc_field:ident) => { + fn take_accumulator(&mut self) -> Box { + let result = Box::new(self.$acc_field.clone()); + self.reset(); + result + } + + fn snapshot_accumulator(&self) -> Box { + Box::new(self.$acc_field.clone()) + } + + fn into_accumulator(self: Box) -> Box { + // Consume the updater and MOVE the accumulator out — no clone. + // Avoids the expensive `Clone` (a full msgpack serialize/deserialize + // round-trip for sketch accumulators) when a pane is evicted at + // window close. + let this = *self; + Box::new(this.$acc_field) + } + }; +} + +/// Shared update interface for query-time and precompute-time accumulation. +/// +/// This provides a uniform interface over all accumulator types so that the +/// worker loop doesn't need to know which concrete type it's dealing with. +pub trait AccumulatorUpdater: Send { + /// Validate an immutable precompute input before an updater can silently + /// discard a value outside its representable domain. + fn validate_single_input(&self, value: f64) -> Result<(), String> { + if value.is_finite() { + Ok(()) + } else { + Err("accumulator input must be finite".into()) + } + } + + /// Feed a single (value, timestamp_ms) pair — for SingleSubpopulation types. + fn update_single(&mut self, value: f64, timestamp_ms: i64); + + /// Feed a keyed (key, value, timestamp_ms) triple — for MultipleSubpopulation types. + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp_ms: i64); + + /// Extract the final accumulator as a boxed `AggregateCore`. + fn take_accumulator(&mut self) -> Box; + + /// Non-destructive read of the current accumulator state (clone without reset). + /// Used by pane-based sliding windows to read shared panes. + fn snapshot_accumulator(&self) -> Box; + + /// Consume the updater and return its accumulator BY MOVE, avoiding the + /// `Clone` that `take_accumulator`/`snapshot_accumulator` pay (for sketch + /// accumulators that clone is a full msgpack serialize/deserialize + /// round-trip). Used by `merge_panes_for_window` when a pane is evicted at + /// window close. Default falls back to a clone for updaters that can't + /// cheaply move their inner accumulator out. + fn into_accumulator(self: Box) -> Box { + self.snapshot_accumulator() + } + + /// Reset internal state for reuse (avoids re-allocation). + fn reset(&mut self); + + /// Whether this updater is keyed (MultipleSubpopulation). + fn is_keyed(&self) -> bool; + + /// Estimated memory usage in bytes. + fn memory_usage_bytes(&self) -> usize; +} + +// --------------------------------------------------------------------------- +// SumAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct SumAccumulatorUpdater { + acc: SumAccumulator, +} + +impl SumAccumulatorUpdater { + pub fn new() -> Self { + Self { + acc: SumAccumulator::new(), + } + } +} + +impl Default for SumAccumulatorUpdater { + fn default() -> Self { + Self::new() + } +} + +impl AccumulatorUpdater for SumAccumulatorUpdater { + fn update_single(&mut self, value: f64, _timestamp_ms: i64) { + self.acc.update(value); + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = SumAccumulator::new(); + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + } +} + +// --------------------------------------------------------------------------- +// MinAccumulatorUpdater / MaxAccumulatorUpdater +// --------------------------------------------------------------------------- + +macro_rules! extremum_updater { + ($updater:ident, $acc:ty) => { + #[derive(Default)] + pub struct $updater { + acc: $acc, + } + + impl $updater { + pub fn new() -> Self { + Self::default() + } + } + + impl AccumulatorUpdater for $updater { + fn update_single(&mut self, value: f64, _timestamp_ms: i64) { + self.acc.update(value); + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = <$acc>::new(); + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::<$acc>() + } + } + }; +} + +extremum_updater!(MinAccumulatorUpdater, MinAccumulator); +extremum_updater!(MaxAccumulatorUpdater, MaxAccumulator); + +// --------------------------------------------------------------------------- +// IncreaseAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct IncreaseAccumulatorUpdater { + acc: Option, +} + +impl IncreaseAccumulatorUpdater { + pub fn new() -> Self { + Self { acc: None } + } +} + +impl Default for IncreaseAccumulatorUpdater { + fn default() -> Self { + Self::new() + } +} + +impl AccumulatorUpdater for IncreaseAccumulatorUpdater { + fn update_single(&mut self, value: f64, timestamp_ms: i64) { + let measurement = Measurement::new(value); + match &mut self.acc { + Some(acc) => acc.update(measurement, timestamp_ms), + None => { + self.acc = Some(IncreaseAccumulator::new( + measurement.clone(), + timestamp_ms, + measurement, + timestamp_ms, + )); + } + } + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + // Hand-written: acc is Option<_> with non-trivial None handling. + fn take_accumulator(&mut self) -> Box { + let acc = self.acc.take().unwrap_or_else(|| { + IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(0.0), 0) + }); + let result = Box::new(acc); + self.reset(); + result + } + + fn snapshot_accumulator(&self) -> Box { + match &self.acc { + Some(acc) => Box::new(acc.clone()), + None => Box::new(IncreaseAccumulator::new( + Measurement::new(0.0), + 0, + Measurement::new(0.0), + 0, + )), + } + } + + fn reset(&mut self) { + self.acc = None; + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::>() + } +} + +// --------------------------------------------------------------------------- +// KllAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct KllAccumulatorUpdater { + acc: DatasketchesKLLAccumulator, + k: u16, +} + +impl KllAccumulatorUpdater { + pub fn new(k: u16) -> Self { + Self { + acc: DatasketchesKLLAccumulator::new(k), + k, + } + } +} + +impl AccumulatorUpdater for KllAccumulatorUpdater { + fn update_single(&mut self, value: f64, _timestamp_ms: i64) { + self.acc.update(value); + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = DatasketchesKLLAccumulator::new(self.k); + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + // KLL sketch size is hard to estimate precisely; use a rough estimate + std::mem::size_of::() + 4096 + } +} + +// --------------------------------------------------------------------------- +// DDSketchAccumulatorUpdater — pendant to KllAccumulatorUpdater +// --------------------------------------------------------------------------- +// +// Drives the agent-aggregated DDSketch path: the worker either +// (a) merges an inbound `DDSketchAccumulator` from the +// modified-OTLP `Data::Ddsketch` ingest (via the worker's +// `merge_with`), or (b) consumes raw values via `update_single` +// when an OTLP scalar datapoint matches an aggregation typed as +// DDSketch. (b) is the less common path but it lets the same +// aggregation slot serve both pre-aggregated agent sketches and +// raw OTLP gauges. +pub struct DDSketchAccumulatorUpdater { + acc: DDSketchAccumulator, + alpha: f64, +} + +impl DDSketchAccumulatorUpdater { + pub fn new(alpha: f64) -> Self { + Self { + acc: DDSketchAccumulator::new(alpha), + alpha, + } + } +} + +impl AccumulatorUpdater for DDSketchAccumulatorUpdater { + fn validate_single_input(&self, value: f64) -> Result<(), String> { + let (minimum, maximum) = + asap_sketchlib::sketches::ddsketch::ddsketch_indexable_bounds(self.alpha); + if value.is_finite() && value > 0.0 && value >= minimum && value <= maximum { + Ok(()) + } else { + Err("DDS maintenance input is outside its positive representable domain".into()) + } + } + + fn update_single(&mut self, value: f64, _timestamp_ms: i64) { + // sketch-core's DdSketch (the inner of DDSketchAccumulator) + // exposes `update(f64)` for single-value ingestion. The + // worker calls this when a raw OTLP datapoint matches an + // aggregation typed as DDSketch — the sketch-merge path + // uses `merge_with` directly. + self.acc.inner.update(value); + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = DDSketchAccumulator::new(self.alpha); + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + // Bucket store is variable; rough estimate matches KLL. + std::mem::size_of::() + 4096 + } +} + +// --------------------------------------------------------------------------- +// KeyedSumCountAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct KeyedSumCountAccumulatorUpdater { + acc: KeyedSumCountAccumulator, +} + +impl KeyedSumCountAccumulatorUpdater { + pub fn new() -> Self { + Self::for_family(ExactKind::Sum) + } + + pub fn for_family(family: ExactKind) -> Self { + Self { + acc: KeyedSumCountAccumulator::for_family(family), + } + } +} + +impl Default for KeyedSumCountAccumulatorUpdater { + fn default() -> Self { + Self::new() + } +} + +impl AccumulatorUpdater for KeyedSumCountAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.update(key.clone(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = KeyedSumCountAccumulator::for_family(self.acc.family.clone()); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.acc.sums.len() * (std::mem::size_of::() + 16) + } +} + +// --------------------------------------------------------------------------- +// KeyedMinStateUpdater / KeyedMaxStateUpdater +// --------------------------------------------------------------------------- + +macro_rules! multiple_extremum_updater { + ($updater:ident, $acc:ty) => { + #[derive(Default)] + pub struct $updater { + acc: $acc, + } + + impl $updater { + pub fn new() -> Self { + Self::default() + } + } + + impl AccumulatorUpdater for $updater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.update(key.clone(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = <$acc>::new(); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::<$acc>() + + self.acc.values.len() * (std::mem::size_of::() + 8) + } + } + }; +} + +multiple_extremum_updater!(KeyedMinStateUpdater, KeyedMinState); +multiple_extremum_updater!(KeyedMaxStateUpdater, KeyedMaxState); + +// --------------------------------------------------------------------------- +// KeyedCounterStateUpdater +// --------------------------------------------------------------------------- + +pub struct KeyedCounterStateUpdater { + acc: KeyedCounterState, +} + +impl KeyedCounterStateUpdater { + pub fn new() -> Self { + Self { + acc: KeyedCounterState::new(), + } + } +} + +impl Default for KeyedCounterStateUpdater { + fn default() -> Self { + Self::new() + } +} + +impl AccumulatorUpdater for KeyedCounterStateUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + let measurement = Measurement::new(value); + match self.acc.increases.entry(key.clone()) { + std::collections::hash_map::Entry::Occupied(mut e) => { + e.get_mut().update(measurement, timestamp_ms); + } + std::collections::hash_map::Entry::Vacant(e) => { + e.insert(IncreaseAccumulator::new( + measurement.clone(), + timestamp_ms, + measurement, + timestamp_ms, + )); + } + } + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = KeyedCounterState::new(); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.acc.increases.len() + * (std::mem::size_of::() + + std::mem::size_of::()) + } +} + +// --------------------------------------------------------------------------- +// CmsAccumulatorUpdater (CountMinSketch) +// --------------------------------------------------------------------------- + +/// Keyed weighted-frequency updater. +/// +/// A raw Prometheus sample represents the observed metric value, so a bare CMS +/// adds `value` for its key. Counting each received sample as one is a distinct +/// event-count operation and requires an explicit typed plan contract; it must +/// not be inferred from the sketch algorithm alone. +pub struct CmsAccumulatorUpdater { + acc: CountMinSketchAccumulator, + row_num: usize, + col_num: usize, +} + +impl CmsAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + acc: CountMinSketchAccumulator::new(row_num, col_num), + row_num, + col_num, + } + } +} + +impl AccumulatorUpdater for CmsAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.inner.update(&key.to_semicolon_str(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = CountMinSketchAccumulator::new(self.row_num, self.col_num); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + } +} + +// --------------------------------------------------------------------------- +// CmsHeapAccumulatorUpdater — value-weighted / count-weighted top-k +// --------------------------------------------------------------------------- + +/// What quantity the top-k heap ranks keys by. +/// +/// These are DIFFERENT query semantics and must be chosen explicitly: +/// +/// * [`TopkWeight::Value`] — accumulate **Σ of the datapoint value** per key. +/// This answers "top-k by total " (e.g. "top-k hosts by +/// total CPU"). The heap value is the summed metric value, so the read-side +/// reducer's "sort heap descending by value" yields the correct ranking. +/// +/// * [`TopkWeight::Count`] — accumulate **+1 per event** per key (occurrence +/// frequency), the textbook heavy-hitter / frequency-top-k semantics +/// ("which keys appear most often"). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TopkWeight { + /// Σ datapoint value per key (value-weighted top-k). + Value, + /// +1 per event per key (count-weighted / frequency top-k). + Count, +} + +/// Keyed top-k updater backed by a real `CountMinSketchWithHeap` (a CMS +/// matrix PLUS a size-`heap_size` top-k heap). Unlike the heap-LESS +/// `CmsAccumulatorUpdater`, this enumerates top-k keys at read time +/// (`get_topk_keys` / `topk_heap_items`), which is what `topk(...)` queries +/// need. +/// +/// The key is the configured group-by (`aggregated_labels`) value vector — +/// e.g. `host` — formed by `extract_aggregated_key_from_series` in the worker, +/// NOT the hardcoded metric label `item`. The accumulated quantity is selected +/// by [`TopkWeight`]: +/// * `Value` → `inner.update(key, value)` adds the datapoint value (Σ value). +/// * `Count` → `inner.update(key, 1.0)` adds one per event (Σ count). +/// +/// Both `CountMinSketchWithHeap` and `CountSketchWithHeap` raw-input policies +/// route here; the heap is the shared distinguishing payload. +pub struct CmsHeapAccumulatorUpdater { + acc: CountMinSketchWithHeapAccumulator, + row_num: usize, + col_num: usize, + heap_size: usize, + weight: TopkWeight, +} + +impl CmsHeapAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize, heap_size: usize, weight: TopkWeight) -> Self { + Self { + acc: CountMinSketchWithHeapAccumulator::new(row_num, col_num, heap_size), + row_num, + col_num, + heap_size, + weight, + } + } +} + +impl AccumulatorUpdater for CmsHeapAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + // Heap key = the group-by label-value vector (e.g. `host`), joined the + // same way the read-side `get_topk_keys` splits it back apart (`;`). + let weighted = match self.weight { + // Σ value: feed the datapoint value. sketchlib's CMS-heap + // `update(key, w)` adds `w.round()` occurrences of `key`, so the + // heap value accumulates the (rounded) summed metric value. + TopkWeight::Value => value, + // Σ count: one occurrence per event, regardless of value. + TopkWeight::Count => 1.0, + }; + self.acc.inner.update(&key.to_semicolon_str(), weighted); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = + CountMinSketchWithHeapAccumulator::new(self.row_num, self.col_num, self.heap_size); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + + self.heap_size * (std::mem::size_of::() + 32) + } +} + +// --------------------------------------------------------------------------- +// CountSketchAccumulatorUpdater (real median-of-signed-rows CountSketch) +// --------------------------------------------------------------------------- + +/// Keyed point-frequency updater backed by a real `asap_sketchlib::CountSketch` +/// (signed rows, median-of-rows estimator) — distinct math from +/// `CmsAccumulatorUpdater`'s CMS (min-of-rows). Closes, on the raw-metric +/// ingest path, the conflation bug where `SketchAlgorithm::CountSketch` silently +/// shared `CmsAccumulatorUpdater` with bare CMS. +/// +/// As with bare CMS, each raw Prometheus sample contributes its `value`. +/// Unit event counting must be selected explicitly by a future typed plan +/// contract rather than being implied by `SketchAlgorithm::CountSketch`. +pub struct CountSketchAccumulatorUpdater { + acc: CountSketchAccumulator, + row_num: usize, + col_num: usize, +} + +impl CountSketchAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + acc: CountSketchAccumulator::new(row_num, col_num), + row_num, + col_num, + } + } +} + +impl AccumulatorUpdater for CountSketchAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.inner.update(&key.to_semicolon_str(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = CountSketchAccumulator::new(self.row_num, self.col_num); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + } +} + +// --------------------------------------------------------------------------- +// CountSketchWithHeapAccumulatorUpdater (real CountSketch + top-k heap) +// --------------------------------------------------------------------------- + +/// Keyed top-k updater backed by a real `CountSketchWithHeap` (signed-row +/// CountSketch matrix PLUS a size-`heap_size` top-k heap). Distinct math from +/// `CmsHeapAccumulatorUpdater`'s CMS-with-heap (min-of-rows); shares the same +/// [`TopkWeight`] semantics and heap payload shape. +pub struct CountSketchWithHeapAccumulatorUpdater { + acc: CountSketchWithHeapAccumulator, + row_num: usize, + col_num: usize, + heap_size: usize, + weight: TopkWeight, +} + +impl CountSketchWithHeapAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize, heap_size: usize, weight: TopkWeight) -> Self { + Self { + acc: CountSketchWithHeapAccumulator::new(row_num, col_num, heap_size), + row_num, + col_num, + heap_size, + weight, + } + } +} + +impl AccumulatorUpdater for CountSketchWithHeapAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + let weighted = match self.weight { + TopkWeight::Value => value, + TopkWeight::Count => 1.0, + }; + self.acc.inner.update(&key.to_semicolon_str(), weighted); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = CountSketchWithHeapAccumulator::new(self.row_num, self.col_num, self.heap_size); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + + self.heap_size * (std::mem::size_of::() + 32) + } +} + +// --------------------------------------------------------------------------- +// HydraKllAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct HydraKllAccumulatorUpdater { + acc: HydraKllSketchAccumulator, + row_num: usize, + col_num: usize, + k: u16, +} + +impl HydraKllAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize, k: u16) -> Self { + Self { + acc: HydraKllSketchAccumulator::new(row_num, col_num, k), + row_num, + col_num, + k, + } + } +} + +impl AccumulatorUpdater for HydraKllAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.update(key, value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = HydraKllSketchAccumulator::new(self.row_num, self.col_num, self.k); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + // Rough estimate: each cell is a KLL sketch + std::mem::size_of::() + self.row_num * self.col_num * 4096 + } +} + +// --------------------------------------------------------------------------- +// Config helpers +// --------------------------------------------------------------------------- + +fn cms_dims(params: &SketchParams) -> (usize, usize) { + match params { + SketchParams::Cms { width, depth } | SketchParams::CountSketch { width, depth } => { + (*depth as usize, *width as usize) + } + other => unreachable!( + "accumulator_spec() paired SketchAlgorithm::Cms/CountSketch with unexpected params: {other:?}" + ), + } +} + +/// Read `(rows = depth, columns = width, heap_size)` out of `SketchParams::CmsWithHeap` +/// or `::CountSketchWithHeap`. +fn cms_heap_dims(params: &SketchParams) -> (usize, usize, usize) { + match params { + SketchParams::CmsWithHeap { + width, + depth, + heap_size, + } + | SketchParams::CountSketchWithHeap { + width, + depth, + heap_size, + } => (*depth as usize, *width as usize, *heap_size as usize), + other => unreachable!( + "accumulator_spec() paired a WithHeap SketchAlgorithm with unexpected params: {other:?}" + ), + } +} + +/// Construct the kernel declared by a Planner SummaryAgg. No backend config +/// tags participate in this dispatch and unsupported payloads are errors. +pub fn create_planner_accumulator( + family: &SummaryFamilyType, + input: &planner_types::post_asap::SummaryUpdate, + grouping: &planner_types::post_asap::GroupingStrategy, +) -> Result, String> { + if input.item.is_some() + && matches!( + input.weight_domain, + planner_types::post_asap::WeightDomain::NonNegative { + proof: + planner_types::post_asap::NonNegativeWeightProof::ResetAwareCounterDerivative + } + ) + { + return Err("window-weighted summaries require typed DAG binding; integer heap updaters cannot consume rates".into()); + } + + asap_physical_operators::capability::validate_summary_kernel(family, input, grouping)?; + use planner_types::post_asap::GroupingStrategy; + if grouping != &GroupingStrategy::PerSubpopulationInstance { + return Err("shared summary grouping requires a supported Planner Hydra kernel".into()); + } + if matches!(family, SummaryFamilyType::ExactAggregate(..)) { + return Ok(Box::new(PlannerExactUpdater { + acc: crate::summary_kernels::exact::ExactAccumulator::new( + family.clone(), + input.item.is_some(), + )?, + })); + } + let SummaryFamilyType::Sketch(kind, family_grouping) = family else { + return Err(format!("unsupported Planner summary family {family:?}")); + }; + if family_grouping != grouping { + return Err("Planner family and operator grouping disagree".into()); + } + let updater: Box = match (kind.algorithm(), kind.params()) { + (SketchAlgorithm::Kll, SketchParams::Kll { k }) => Box::new(KllAccumulatorUpdater::new( + u16::try_from(*k).map_err(|_| "KLL k exceeds runtime bound")?, + )), + (SketchAlgorithm::DDSketch, SketchParams::DDSketch { alpha }) => { + Box::new(DDSketchAccumulatorUpdater::new(*alpha)) + } + (SketchAlgorithm::Cms, params @ SketchParams::Cms { .. }) => { + let (r, c) = cms_dims(params); + Box::new(CmsAccumulatorUpdater::new(r, c)) + } + (SketchAlgorithm::CountSketch, params @ SketchParams::CountSketch { .. }) => { + let (r, c) = cms_dims(params); + Box::new(CountSketchAccumulatorUpdater::new(r, c)) + } + (SketchAlgorithm::CmsWithHeap, params @ SketchParams::CmsWithHeap { .. }) => { + let (r, c, h) = cms_heap_dims(params); + Box::new(CmsHeapAccumulatorUpdater::new(r, c, h, TopkWeight::Value)) + } + ( + SketchAlgorithm::CountSketchWithHeap, + params @ SketchParams::CountSketchWithHeap { .. }, + ) => { + let (r, c, h) = cms_heap_dims(params); + Box::new(CountSketchWithHeapAccumulatorUpdater::new( + r, + c, + h, + TopkWeight::Value, + )) + } + (SketchAlgorithm::Hll, SketchParams::Hll { precision }) => Box::new(HllUpdater { + acc: HllSketchAccumulator::new( + asap_sketchlib::HllVariant::Regular, + u32::from(*precision), + ), + }), + ( + SketchAlgorithm::UnivMon, + SketchParams::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + }, + ) => Box::new(UnivMonUpdater { + acc: UnivMonAccumulator::new( + *heap_size as usize, + *sketch_rows as usize, + *sketch_cols as usize, + *layers as usize, + ) + .map_err(|e| e.to_string())?, + }), + _ => { + return Err(format!( + "unsupported Planner algorithm/parameters: {kind:?}" + )) + } + }; + if updater.is_keyed() != input.item.is_some() && !is_unit_sample_frequency(input) { + return Err("Planner item expression does not match the selected kernel layout".into()); + } + Ok(updater) +} + +/// A unit-weight sample-value frequency may use a keyed kernel without an item column. +fn is_unit_sample_frequency(update: &planner_types::post_asap::SummaryUpdate) -> bool { + use planner_types::post_asap::{NonNegativeWeightProof, SummaryInputExpr, WeightDomain}; + matches!( + update.item, + Some(SummaryInputExpr::Column( + planner_types::pre_asap::ColumnRef::SampleValue + )) + ) && matches!(update.weight, SummaryInputExpr::Constant(1.0)) + && matches!( + update.weight_domain, + WeightDomain::NonNegative { + proof: NonNegativeWeightProof::UnitCount + } + ) +} + +struct PlannerExactUpdater { + acc: crate::summary_kernels::exact::ExactAccumulator, +} +impl AccumulatorUpdater for PlannerExactUpdater { + fn update_single(&mut self, value: f64, timestamp: i64) { + self.acc.update(None, value, timestamp); + } + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp: i64) { + self.acc.update(Some(key), value, timestamp); + } + impl_clone_accumulator_methods!(acc); + fn reset(&mut self) { + self.acc = crate::summary_kernels::exact::ExactAccumulator::new( + self.acc.family().clone(), + self.acc.is_keyed(), + ) + .expect("installed exact family"); + } + fn is_keyed(&self) -> bool { + self.acc.is_keyed() + } + fn memory_usage_bytes(&self) -> usize { + self.acc.approx_memory_bytes() + } +} + +struct UnivMonUpdater { + acc: UnivMonAccumulator, +} + +struct HllUpdater { + acc: HllSketchAccumulator, +} + +impl AccumulatorUpdater for HllUpdater { + fn is_keyed(&self) -> bool { + false + } + fn memory_usage_bytes(&self) -> usize { + self.acc.approx_memory_bytes() + } + fn update_single(&mut self, value: f64, _: i64) { + if !value.is_nan() { + let bits = if value == 0.0 { 0 } else { value.to_bits() }; + self.acc.inner.update(&bits.to_le_bytes()); + } + } + fn update_keyed(&mut self, _: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + impl_clone_accumulator_methods!(acc); + fn reset(&mut self) { + self.acc.reset_to_empty(); + } +} + +impl AccumulatorUpdater for UnivMonUpdater { + fn is_keyed(&self) -> bool { + false + } + fn memory_usage_bytes(&self) -> usize { + self.acc.approx_memory_bytes() + } + fn update_single(&mut self, value: f64, _: i64) { + self.acc + .insert_sample(value) + .expect("UnivMon sample counter overflow"); + } + fn update_keyed(&mut self, _: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + impl_clone_accumulator_methods!(acc); + fn reset(&mut self) { + self.acc.reset_to_empty(); + } +} + +#[cfg(test)] +mod planner_parameter_regression { + use super::*; + use planner_types::post_asap::{SketchKind, SummaryInputExpr, SummaryUpdate}; + + // Planner width is the bucket count; depth is the independent hash-row count. + #[test] + fn planner_sketch_dimensions_are_not_transposed() { + for (algorithm, params) in [ + ( + SketchAlgorithm::Cms, + SketchParams::Cms { + width: 128, + depth: 3, + }, + ), + ( + SketchAlgorithm::CountSketch, + SketchParams::CountSketch { + width: 128, + depth: 3, + }, + ), + ( + SketchAlgorithm::CmsWithHeap, + SketchParams::CmsWithHeap { + width: 128, + depth: 3, + heap_size: 8, + }, + ), + ( + SketchAlgorithm::CountSketchWithHeap, + SketchParams::CountSketchWithHeap { + width: 128, + depth: 3, + heap_size: 8, + }, + ), + ] { + let family = SummaryFamilyType::Sketch( + SketchKind::new(algorithm.clone(), params), + Default::default(), + ); + let update = SummaryUpdate { + item: Some(SummaryInputExpr::Column( + planner_types::pre_asap::ColumnRef::Named("host".into()), + )), + weight: SummaryInputExpr::Constant(1.0), + weight_domain: Default::default(), + }; + let state = create_planner_accumulator(&family, &update, &Default::default()) + .unwrap() + .snapshot_accumulator(); + let dims = match algorithm { + SketchAlgorithm::Cms => { + let s = state + .as_any() + .downcast_ref::() + .unwrap(); + (s.inner.rows(), s.inner.cols()) + } + SketchAlgorithm::CountSketch => { + let s = state + .as_any() + .downcast_ref::() + .unwrap(); + (s.inner.rows, s.inner.cols) + } + SketchAlgorithm::CmsWithHeap => { + let s = state + .as_any() + .downcast_ref::() + .unwrap(); + (s.inner.rows(), s.inner.cols()) + } + SketchAlgorithm::CountSketchWithHeap => { + let s = state + .as_any() + .downcast_ref::() + .unwrap(); + (s.inner.rows(), s.inner.cols()) + } + _ => unreachable!(), + }; + assert_eq!(dims, (3, 128), "{algorithm:?}"); + } + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/hll_sketch.rs b/crates/asap_summary_state/src/summary_kernels/hll_sketch.rs new file mode 100644 index 000000000..4c16156d6 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/hll_sketch.rs @@ -0,0 +1,788 @@ +//! HLL accumulator — wraps `asap_sketchlib::HllSketch`. +//! +//! Concrete accumulator reached from the modified-OTLP +//! `Metric.data = HLLSketch{…}` hot path (PR C-CountSketch follow-up). +//! Mirrors the CountSketch accumulator's shape: merge via register-wise +//! max on the inner sketch, serialize as MessagePack for the sink, and +//! decode from the sketchlib `HyperLogLogState` proto. +//! +//! Query semantics (cardinality estimation via the three HLL variants' +//! estimators) are intentionally deferred — the wire format carries the +//! registers + variant + HIP accumulators losslessly, so the merge + +//! store round-trip works end-to-end without that richer query surface. + +use crate::summary_kernels::dd_sketch::normalize_sample_p; +use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::{HllSketch, HllVariant, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +/// Decode one protobuf base-128 varint (LEB128) from the front of `buf`. +/// Returns `(value, bytes_consumed)`, or `None` if the buffer is truncated +/// or the varint overflows u64. +pub(crate) fn read_uvarint(buf: &[u8]) -> Option<(u64, usize)> { + let mut result: u64 = 0; + let mut shift: u32 = 0; + for (i, &b) in buf.iter().enumerate() { + if shift >= 64 { + return None; + } + result |= u64::from(b & 0x7f) << shift; + if b & 0x80 == 0 { + return Some((result, i + 1)); + } + shift += 7; + } + None +} + +/// Expand sketchlib-go's sparse HLL register encoding +/// (`HLLSparseRegisters.packed`) into the dense `num_registers`-byte array. +/// +/// Layout (sketchlib-go `proto/hll/hll.proto`): varint-packed +/// `(index_delta, value)` pairs in ascending index order; `prev_index` +/// starts at 0, so each register's absolute index is the running sum of the +/// deltas. Mirrors the Go encoder in `sketches/HLL/sparse.go` +/// (`encodeSparseRegisters`). The reconstructed array is byte-identical to +/// the dense `registers` field a high-cardinality producer would have sent. +pub(crate) fn expand_sparse_hll_registers( + packed: &[u8], + num_registers: usize, +) -> Result, Box> { + let mut regs = vec![0u8; num_registers]; + let mut prev: u64 = 0; + let mut pos = 0usize; + while pos < packed.len() { + let (delta, n1) = read_uvarint(&packed[pos..]) + .ok_or("HLLSparseRegisters.packed: truncated index_delta varint")?; + pos += n1; + let (value, n2) = read_uvarint(&packed[pos..]) + .ok_or("HLLSparseRegisters.packed: truncated value varint")?; + pos += n2; + let idx = prev + delta; + let i = usize::try_from(idx) + .map_err(|_| format!("HLLSparseRegisters: index {idx} overflows usize"))?; + if i >= num_registers { + return Err(format!( + "HLLSparseRegisters: register index {i} >= num_registers {num_registers}" + ) + .into()); + } + regs[i] = u8::try_from(value) + .map_err(|_| format!("HLLSparseRegisters: register value {value} > 255"))?; + prev = idx; + } + Ok(regs) +} + +/// HLL accumulator — inner register array + variant metadata. +#[derive(Debug, Clone)] +pub struct HllSketchAccumulator { + pub inner: HllSketch, + /// Edge sampling probability `p ∈ (0,1]` carried on the producer's + /// `SketchEnvelope.sample_p`. HLL uses HASH-THRESHOLD sampling — each + /// DISTINCT key is admitted into the sketch with probability `p`, so the + /// register-derived distinct-count estimate is ~`p`× the true + /// cardinality and a `Cardinality`/`Count` query must rescale by `1/p`. + /// `1.0` (and the proto3 default `0.0`, dual-read as `1.0`) means no + /// sampling, so the rescale is a no-op and the behaviour is identical to + /// before. Mirrors `DDSketchAccumulator::sample_p`; set from the envelope + /// at the `from_sketchlib_proto_bytes` decode site and preserved across + /// `reset_to_empty` and `merge_with`. + /// + /// NOTE: HLL edge sampling is currently force-disabled in the edge + /// (`warm_sketch.go` HLL case always emits `sample_p = 1.0`), so in + /// practice `p = 1.0` today and this is a latent-correctness fix that + /// activates if HLL sampling is ever enabled. + pub sample_p: f64, +} + +impl HllSketchAccumulator { + pub fn new(variant: HllVariant, precision: u32) -> Self { + Self { + inner: HllSketch::new(variant, precision), + sample_p: 1.0, + } + } + + /// Decode from the modified OTLP wire format's + /// `HLLSketchDataPoint.sketch` bytes when + /// `encoding = HLL_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `HllSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: HllSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?, + // The msgpack HllSketch struct carries no envelope/sample_p; the + // msgpack path is parity/test-only and is never edge-sampled. + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `HLLSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::HyperLogLogState` message + /// that DataCollector's `hllprocessor` emits when + /// `encoding = HLL_SKETCH_ENCODING_PROTO`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + + // DataCollector's hllprocessor wraps the state in a + // `SketchEnvelope{hll: HyperLogLogState}` via sketchlib-go's + // `SerializePortableFO` + `proto.Marshal`. Try envelope first, + // fall back to bare `HyperLogLogState` for callers (e.g. unit + // tests) that encode the state directly. Mirrors the PR #14 + // fix on `CountMinSketchAccumulator::from_sketchlib_proto_bytes`. + // Capture the envelope's `sample_p` alongside the state so a + // Cardinality query can rescale the distinct-count estimate by + // `1/p`. Bare `HyperLogLogState` bytes (no envelope) carry no + // sampling info → `sample_p` 1.0 (no rescale). Mirrors + // `DDSketchAccumulator`. + let (state, sample_p) = match SketchEnvelope::decode(buffer) { + Ok(env) => { + let sp = env.sample_p; + match env.sketch_state { + Some(sketch_envelope::SketchState::Hll(st)) => (st, sp), + Some(other) => { + return Err(format!( + "SketchEnvelope contains non-HLL sketch: {:?}", + std::mem::discriminant(&other) + ) + .into()); + } + None => ( + HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?, + 1.0, + ), + } + } + Err(_) => ( + HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?, + 1.0, + ), + }; + if state.precision == 0 || state.precision > 20 { + return Err(format!( + "HyperLogLogState precision {} out of range (expected 1..=20)", + state.precision + ) + .into()); + } + let expected_len = 1usize << state.precision; + // Register resolution. sketchlib-go emits the SPARSE + // `registers_sparse` (proto tag 7) form below its dense/sparse + // crossover (~6000 non-zero registers — see + // sketchlib-go/sketches/HLL/sparse.go); low-cardinality producers + // (the common case) therefore leave the dense `registers` (tag 3) + // field empty. The proto contract (hll.proto) is: read whichever of + // `registers` / `registers_sparse` is present; if both are empty the + // sketch is all-zero. Reconstruct the dense 2^precision array in all + // three cases so the inner `HllSketch` always gets a full register + // vector. + let dense_registers: Vec = if state.registers.len() == expected_len { + state.registers.clone() + } else if !state.registers.is_empty() { + // A non-empty dense field of the wrong length is a malformed frame. + return Err(format!( + "HyperLogLogState registers has {} bytes, expected 2^precision = {}", + state.registers.len(), + expected_len + ) + .into()); + } else if let Some(sparse) = state.registers_sparse.as_ref() { + expand_sparse_hll_registers(&sparse.packed, expected_len)? + } else { + // Neither representation populated → all-zero register array. + vec![0u8; expected_len] + }; + let proto_variant = ProtoVariant::try_from(state.variant) + .map_err(|_| format!("HyperLogLogState has unknown variant tag {}", state.variant))?; + let variant = match proto_variant { + ProtoVariant::Unspecified => HllVariant::Unspecified, + ProtoVariant::Regular => HllVariant::Regular, + ProtoVariant::ErtlMle => HllVariant::Datafusion, + ProtoVariant::Hip => HllVariant::Hip, + }; + let inner = HllSketch::from_raw( + variant, + state.precision, + dense_registers, + state.hip_kxq0, + state.hip_kxq1, + state.hip_est, + ); + Ok(Self { + inner, + sample_p: normalize_sample_p(sample_p), + }) + } + + /// Apply a proto-encoded `HLLDelta` frame to this accumulator's + /// inner sketch — the decode path for + /// `HLL_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + /// + /// Called against an accumulator that already carries the base + /// sketch state; the caller is the per-series snapshot cache in + /// the ingest path. Bytes are the + /// `asap_sketchlib::proto::sketchlib::HllDelta` message. + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + // The HLLDelta wire format is a varint-packed (index_delta, value) blob; + // decode + apply (register-wise max) via the shared sketch library so + // the unpacking stays a single source of truth. + self.inner + .apply_delta_bytes(buffer) + .map_err(|e| format!("apply HLLDelta: {e}"))?; + Ok(()) + } +} + +impl SerializableToSink for HllSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "variant": format!("{:?}", self.inner.variant), + "precision": self.inner.precision, + "register_bytes": self.inner.registers.len(), + "hip_kxq0": self.inner.hip_kxq0, + "hip_kxq1": self.inner.hip_kxq1, + "hip_est": self.inner.hip_est, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for HllSketchAccumulator { + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::().saturating_add(self.inner.registers.capacity()) + } + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "HllSketchAccumulator" + } + + /// Per-window base rotation: zero the registers but keep the variant + /// and precision. Critical for HLL — its register-wise `max` merge + /// has no inverse, so a never-reset base accumulates the all-time-max + /// across windows (`docs/delta-baseline-contract.md` §1.5); rotating + /// to an empty register array makes per-window cardinality correct. + /// `sample_p` is a per-series config constant (not per-window data), so + /// it is intentionally preserved across the rotation — mirrors + /// `DDSketchAccumulator`. + fn reset_to_empty(&mut self) { + self.inner = HllSketch::new(self.inner.variant, self.inner.precision); + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge HllSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_hll = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to HllSketchAccumulator")?; + let merged_inner = HllSketch::merge_refs(&[&self.inner, &other_hll.inner])?; + // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a + // per-series config constant, so both operands carry the same value + // in practice. Prefer a sampled factor over the no-sampling default + // so a merge with a freshly-reset (1.0) base keeps the series' + // sampling rate. + let sample_p = if self.sample_p < 1.0 { + self.sample_p + } else { + other_hll.sample_p + }; + Ok(Box::new(Self { + inner: merged_inner, + sample_p, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::HLL + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + _query_kwargs: &HashMap, + ) -> Result> { + use crate::Statistic; + match statistic { + // HLL's natural answer is unique-cardinality. PromQL's + // `count_over_time(...)` and `count(...)` both surface + // as `Statistic::Count` after pattern matching but + // semantically they mean "how many distinct values + // were observed in this window" when the underlying + // aggregator is HLL — that's the cardinality estimate, + // not a sample-count. Accept both. + Statistic::Cardinality | Statistic::Count => { + // HLL uses hash-threshold sampling — each distinct key is + // admitted with probability `sample_p`, so the register- + // derived distinct-count estimate is ~`p`× the true + // cardinality. Rescale by `1/sample_p` for an unbiased + // estimate. `sample_p == 1.0` (unsampled / legacy / edge + // HLL sampling currently force-disabled) makes this a no-op. + Ok(hll_cardinality_estimate(&self.inner.registers) / self.sample_p) + } + other => Err(format!( + "HllSketchAccumulator: statistic {:?} not supported (only Cardinality / Count)", + other, + ) + .into()), + } + } +} + +/// Standard HyperLogLog cardinality estimate with the canonical +/// `α_m × m² / Σ 2^(-register[i])` formula plus the small-range +/// (linear-counting) and large-range (32-bit space) corrections +/// from the original Flajolet et al. paper. +/// +/// Inlined here rather than added as a method on `asap_sketchlib::HllSketch` +/// because the existing `asap_sketchlib::asap` types only expose merge / +/// serialize today; adding a query method there would force a +/// cross-crate change. +fn hll_cardinality_estimate(registers: &[u8]) -> f64 { + let m = registers.len() as f64; + if m == 0.0 { + return 0.0; + } + let alpha = match registers.len() { + 16 => 0.673, + 32 => 0.697, + 64 => 0.709, + _ => 0.7213 / (1.0 + 1.079 / m), + }; + + let mut sum = 0.0f64; + let mut zero_registers = 0usize; + for &r in registers { + sum += 2f64.powi(-(r as i32)); + if r == 0 { + zero_registers += 1; + } + } + let raw = alpha * m * m / sum; + + // Small-range (linear-counting) correction. + if raw <= 2.5 * m && zero_registers > 0 { + return m * (m / zero_registers as f64).ln(); + } + + // Large-range correction (only meaningful with 32-bit register + // spaces; sketch-core uses up to 64-bit hashes so this branch + // rarely fires in practice — kept for completeness). + let two_pow_32 = 4_294_967_296f64; + if raw > two_pow_32 / 30.0 { + return -two_pow_32 * (1.0 - raw / two_pow_32).ln(); + } + raw +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode_state( + variant: i32, + precision: u32, + registers: Vec, + hip_kxq0: f64, + hip_kxq1: f64, + hip_est: f64, + ) -> Vec { + use asap_sketchlib::proto::sketchlib::HyperLogLogState; + use prost::Message; + let state = HyperLogLogState { + variant, + precision, + registers, + hip_kxq0, + hip_kxq1, + hip_est, + registers_sparse: None, + }; + state.encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_regular() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + let bytes = encode_state( + ProtoVariant::Regular as i32, + 2, + vec![1, 2, 3, 4], + 0.0, + 0.0, + 0.0, + ); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Regular); + assert_eq!(acc.inner.precision, 2); + assert_eq!(acc.inner.registers, vec![1, 2, 3, 4]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_hip_preserves_accumulators() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + let bytes = encode_state( + ProtoVariant::Hip as i32, + 2, + vec![0, 0, 0, 0], + 1.5, + 2.5, + 42.0, + ); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Hip); + assert_eq!(acc.inner.hip_kxq0, 1.5); + assert_eq!(acc.inner.hip_kxq1, 2.5); + assert_eq!(acc.inner.hip_est, 42.0); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's hllprocessor emits: the state + // wrapped in a `SketchEnvelope{hll: ...}` via sketchlib-go's + // `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision: 2, + registers: vec![1, 2, 3, 4], + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + registers_sparse: None, + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + assert_eq!(acc.inner.variant, HllVariant::Regular); + assert_eq!(acc.inner.registers, vec![1, 2, 3, 4]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_register_length_mismatch() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + // precision=2 → expected 4 registers; supply only 3 + let bytes = encode_state( + ProtoVariant::Regular as i32, + 2, + vec![1, 2, 3], + 0.0, + 0.0, + 0.0, + ); + let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("registers")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_zero_precision_rejected() { + use asap_sketchlib::proto::sketchlib::HyperLogLogState; + use prost::Message; + let state = HyperLogLogState::default(); + let bytes = state.encode_to_vec(); + let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + } + + #[test] + fn test_aggregate_core_merge_matches_register_max() { + let a = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 5, 3, 7], 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let b = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![4, 2, 6, 0], 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.inner.registers, vec![4, 5, 6, 7]); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::summary_kernels::count_sketch::CountSketchAccumulator; + let hll = HllSketchAccumulator::new(HllVariant::Regular, 2); + let cs = CountSketchAccumulator::new(2, 3); + assert!(hll.merge_with(&cs).is_err()); + } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = HllSketch::from_raw( + HllVariant::Hip, + 3, + vec![0, 1, 2, 3, 4, 5, 6, 7], + 1.5, + 2.5, + 42.0, + ); + let bytes = original.to_msgpack().unwrap(); + let acc = HllSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Hip); + assert_eq!(acc.inner.precision, 3); + assert_eq!(acc.inner.registers, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(acc.inner.hip_kxq0, 1.5); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = HllSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } + + #[test] + fn test_apply_proto_delta_bytes_round_trip() { + use asap_sketchlib::proto::sketchlib::HllDelta as PbDelta; + use prost::Message; + + let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); + acc.inner.registers = vec![1, 5, 3, 7]; + + // Packed (index_delta, value) blob for updates {0:4, 2:6}: + // varint(0),varint(4),varint(2),varint(6). + let delta_bytes = PbDelta { + packed_updates: vec![0, 4, 2, 6], + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&delta_bytes).expect("apply ok"); + // Max semantics: reg[0]=max(1,4)=4, reg[2]=max(3,6)=6; others unchanged. + assert_eq!(acc.inner.registers, vec![4, 5, 6, 7]); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ----- sample_p cardinality rescale ----- + // + // HLL uses hash-threshold sampling: each distinct key is admitted into + // the sketch with probability `p`, so the register-derived cardinality + // estimate is ~p× the true distinct count and must be rescaled by 1/p. + + #[test] + fn test_cardinality_is_rescaled_by_sample_p() { + use crate::Statistic; + // Build two accumulators with identical registers but different + // sample_p. The sampled one (p=0.25) must report ~4× the unsampled + // estimate. Use precision 8 (256 registers) with a spread of + // register values so the estimate is a non-trivial positive number. + let mut registers = vec![0u8; 256]; + for (i, r) in registers.iter_mut().enumerate() { + *r = ((i % 7) + 1) as u8; + } + let unsampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 8, registers.clone(), 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let sampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 8, registers, 0.0, 0.0, 0.0), + sample_p: 0.25, + }; + let raw = unsampled + .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) + .expect("cardinality ok"); + let rescaled = sampled + .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) + .expect("cardinality ok"); + assert!(raw > 0.0, "raw estimate should be positive, got {raw}"); + // Exact algebraic relationship: rescaled == raw / 0.25 == raw * 4. + assert!( + (rescaled - raw * 4.0).abs() < 1e-9, + "expected rescaled ≈ 4×raw ({}), got {rescaled}", + raw * 4.0 + ); + } + + #[test] + fn test_count_statistic_also_rescaled_by_sample_p() { + use crate::Statistic; + // Count maps to the same cardinality estimate for HLL, so it must + // rescale identically. + let registers = vec![3u8; 16]; + let unsampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 4, registers.clone(), 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let sampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 4, registers, 0.0, 0.0, 0.0), + sample_p: 0.25, + }; + let raw = unsampled + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + let rescaled = sampled + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + assert!((rescaled - raw * 4.0).abs() < 1e-9); + } + + #[test] + fn test_sample_p_unset_behaves_as_one() { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + // An envelope with no sample_p set (proto3 default 0.0) must + // normalize to 1.0 (no rescale) — byte-compatible with legacy frames. + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision: 4, + registers: vec![2u8; 16], + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + registers_sparse: None, + }; + let env = SketchEnvelope { + // sample_p left at proto3 default 0.0. + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 1.0, "unset sample_p must normalize to 1.0"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { + use crate::Statistic; + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + + let registers = vec![3u8; 16]; + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision: 4, + registers: registers.clone(), + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + registers_sparse: None, + }; + let env = SketchEnvelope { + sample_p: 0.25, + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 0.25); + + // Compare against the unsampled estimate over the same registers. + let unsampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 4, registers, 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let raw = unsampled + .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) + .expect("cardinality ok"); + let rescaled = acc + .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) + .expect("cardinality ok"); + assert!( + (rescaled - raw * 4.0).abs() < 1e-9, + "expected 4×raw rescale" + ); + } + + #[test] + fn test_reset_to_empty_preserves_sample_p() { + let mut acc = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 4, vec![3u8; 16], 0.0, 0.0, 0.0), + sample_p: 0.25, + }; + acc.reset_to_empty(); + assert_eq!(acc.sample_p, 0.25, "window rotation must keep sample_p"); + assert_eq!(acc.inner.registers, vec![0u8; 16], "registers cleared"); + } + + #[test] + fn test_merge_prefers_sampled_factor() { + let a = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 1, 1, 1], 0.0, 0.0, 0.0), + sample_p: 0.25, + }; + let b = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 1, 1, 1], 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let merged = a.merge_with(&b).expect("merge ok"); + let merged = merged + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.sample_p, 0.25); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/hydra_kll.rs b/crates/asap_summary_state/src/summary_kernels/hydra_kll.rs new file mode 100644 index 000000000..167bcae38 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/hydra_kll.rs @@ -0,0 +1,165 @@ +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{HydraKllSketch, MessagePackCodec}; +use base64::{engine::general_purpose, Engine as _}; +use std::collections::HashMap; + +use crate::Statistic; + +/// HydraKLL sketch accumulator — wraps asap_sketchlib::HydraKllSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. +/// This file retains QE-specific trait impls and JSON output. +#[derive(Debug, Clone)] +pub struct HydraKllSketchAccumulator { + pub inner: HydraKllSketch, +} + +impl HydraKllSketchAccumulator { + pub fn new(row_num: usize, col_num: usize, k: u16) -> Self { + Self { + inner: HydraKllSketch::new(row_num, col_num, k), + } + } + + pub fn update(&mut self, key: &KeyByLabelValues, value: f64) { + self.inner.update(&key.to_semicolon_str(), value); + } + + pub fn deserialize_from_bytes(_buffer: &[u8]) -> Result> { + Err("deserialize_from_bytes for HydraKllSketchAccumulator not implemented".into()) + } + + pub fn query_key(&self, key: &KeyByLabelValues, quantile: f64) -> f64 { + self.inner.quantile(&key.to_semicolon_str(), quantile) + } +} + +impl SerializableToSink for HydraKllSketchAccumulator { + fn serialize_to_json(&self) -> serde_json::Value { + // Mirror Python implementation: {"sketch": base64_encoded_string} + let sketch_bytes = self.inner.to_msgpack().unwrap_or_default(); + let sketch_b64 = general_purpose::STANDARD.encode(&sketch_bytes); + serde_json::json!({ "sketch": sketch_b64 }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl MergeableAccumulator for HydraKllSketchAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +impl AggregateCore for HydraKllSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "HydraKllSketchAccumulator" + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge HydraKllSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let hk = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to HydraKllSketchAccumulator")?; + + let merged = Self::merge_accumulators(vec![self.clone(), hk.clone()])?; + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::HydraKLL + } + + fn approx_memory_bytes(&self) -> usize { + // HydraKLL is a row*col grid of KLL sketches; typical instances + // are on the order of tens of KiB. 32 KiB is a conservative + // per-instance default. + 32 * 1024 + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for HydraKllSketchAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for HydraKllSketchAccumulator { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Quantile => { + let quantile = query_kwargs + .and_then(|kwargs| kwargs.get("quantile")) + .ok_or("Missing quantile parameter for quantile query")? + .parse::() + .map_err(|_| "Invalid quantile parameter format")?; + + if !(0.0..=1.0).contains(&quantile) { + return Err("Quantile must be between 0.0 and 1.0".into()); + } + + Ok(self.query_key(key, quantile)) + } + _ => Err( + format!("Unsupported statistic in HydraKllSketchAccumulator: {statistic:?}").into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/increase.rs b/crates/asap_summary_state/src/summary_kernels/increase.rs new file mode 100644 index 000000000..c6b988325 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/increase.rs @@ -0,0 +1,737 @@ +use crate::codec::MeasurementCodec; +use crate::{ + AggregateCore, AggregationType, Measurement, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +const RESET_AWARE_WIRE_MAGIC: &[u8; 8] = b"ASAPINC2"; +const RESET_AWARE_WIRE_EXTENSION_LEN: usize = 8 + 8 + 8; + +/// Accumulator for tracking increases in counter metrics +/// Stores the starting and last seen measurements with timestamps +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct IncreaseAccumulator { + pub starting_measurement: Measurement, + pub starting_timestamp: i64, + pub last_seen_measurement: Measurement, + pub last_seen_timestamp: i64, + /// Sum of monotonic deltas, adding the post-reset value whenever the + /// counter decreases. This is the reset correction Prometheus applies. + #[serde(default)] + pub total_increase: f64, + #[serde(default)] + pub sample_count: u64, +} + +impl IncreaseAccumulator { + /// Merge two counter intervals without a temporary collection. Ties retain + /// the left input, matching the stable ordering of multi-pane merges. + pub(crate) fn merge_pair(left: &Self, right: &Self) -> Self { + let (first, second) = if left.starting_timestamp <= right.starting_timestamp { + (left, right) + } else { + (right, left) + }; + let mut merged = first.clone(); + if second.starting_timestamp > merged.last_seen_timestamp { + merged.total_increase += + if second.starting_measurement.value >= merged.last_seen_measurement.value { + second.starting_measurement.value - merged.last_seen_measurement.value + } else { + second.starting_measurement.value + }; + } + merged.total_increase += second.total_increase; + merged.sample_count = merged.sample_count.saturating_add(second.sample_count); + if second.last_seen_timestamp > merged.last_seen_timestamp { + merged.last_seen_measurement = second.last_seen_measurement.clone(); + merged.last_seen_timestamp = second.last_seen_timestamp; + } + + merged + } + + /// Return the number of bytes occupied by one accumulator at the start of + /// `buffer`. Old persisted values end after `last_seen_timestamp`; reset- + /// aware values carry a magic-prefixed extension. The magic makes this + /// safe when the buffer also contains the next keyed entry. + pub(crate) fn serialized_len_from_prefix( + buffer: &[u8], + ) -> Result> { + if buffer.len() < 4 { + return Err("Buffer too short for starting measurement length".into()); + } + let starting_len = u32::from_le_bytes(buffer[0..4].try_into()?) as usize; + let last_len_offset = 4usize + .checked_add(starting_len) + .and_then(|offset| offset.checked_add(8)) + .ok_or("IncreaseAccumulator length overflow")?; + if buffer.len() < last_len_offset + 4 { + return Err("Buffer too short for last seen measurement length".into()); + } + let last_len = + u32::from_le_bytes(buffer[last_len_offset..last_len_offset + 4].try_into()?) as usize; + let legacy_len = last_len_offset + .checked_add(4) + .and_then(|offset| offset.checked_add(last_len)) + .and_then(|offset| offset.checked_add(8)) + .ok_or("IncreaseAccumulator length overflow")?; + if buffer.len() < legacy_len { + return Err("Buffer too short for last seen timestamp".into()); + } + let has_extension = buffer.len() >= legacy_len + RESET_AWARE_WIRE_EXTENSION_LEN + && &buffer[legacy_len..legacy_len + RESET_AWARE_WIRE_MAGIC.len()] + == RESET_AWARE_WIRE_MAGIC; + Ok(legacy_len + + if has_extension { + RESET_AWARE_WIRE_EXTENSION_LEN + } else { + 0 + }) + } + + pub fn new( + starting_measurement: Measurement, + starting_timestamp: i64, + last_seen_measurement: Measurement, + last_seen_timestamp: i64, + ) -> Self { + let total_increase = if last_seen_timestamp <= starting_timestamp { + 0.0 + } else if last_seen_measurement.value >= starting_measurement.value { + last_seen_measurement.value - starting_measurement.value + } else { + last_seen_measurement.value + }; + let sample_count = if last_seen_timestamp > starting_timestamp { + 2 + } else { + 1 + }; + Self { + starting_measurement, + starting_timestamp, + last_seen_measurement, + last_seen_timestamp, + total_increase, + sample_count, + } + } + + pub fn update(&mut self, measurement: Measurement, timestamp: i64) { + if timestamp < self.last_seen_timestamp { + return; + } + if timestamp == self.last_seen_timestamp { + return; + } + if measurement.value >= self.last_seen_measurement.value { + self.total_increase += measurement.value - self.last_seen_measurement.value; + } else { + self.total_increase += measurement.value; + } + self.last_seen_measurement = measurement; + self.last_seen_timestamp = timestamp; + self.sample_count = self.sample_count.saturating_add(1); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let starting_measurement = + Measurement::deserialize_from_json(&data["starting_measurement"])?; + let starting_timestamp = data["starting_timestamp"] + .as_i64() + .ok_or("Missing or invalid 'starting_timestamp' field")?; + let last_seen_measurement = + Measurement::deserialize_from_json(&data["last_seen_measurement"])?; + let last_seen_timestamp = data["last_seen_timestamp"] + .as_i64() + .ok_or("Missing or invalid 'last_seen_timestamp' field")?; + + let mut accumulator = Self::new( + starting_measurement, + starting_timestamp, + last_seen_measurement, + last_seen_timestamp, + ); + accumulator.total_increase = data["total_increase"] + .as_f64() + .unwrap_or(accumulator.total_increase); + accumulator.sample_count = data["sample_count"] + .as_u64() + .unwrap_or(accumulator.sample_count); + Ok(accumulator) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut offset = 0; + + // Read starting measurement length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for starting measurement length".into()); + } + let starting_measurement_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + starting_measurement_length { + return Err("Buffer too short for starting measurement".into()); + } + let starting_measurement = Measurement::deserialize_from_bytes( + &buffer[offset..offset + starting_measurement_length], + )?; + offset += starting_measurement_length; + + // Read starting timestamp + if buffer.len() < offset + 8 { + return Err("Buffer too short for starting timestamp".into()); + } + let starting_timestamp = i64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + + // Read last seen measurement length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for last seen measurement length".into()); + } + let last_seen_measurement_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + last_seen_measurement_length { + return Err("Buffer too short for last seen measurement".into()); + } + let last_seen_measurement = Measurement::deserialize_from_bytes( + &buffer[offset..offset + last_seen_measurement_length], + )?; + offset += last_seen_measurement_length; + + // Read last seen timestamp + if buffer.len() < offset + 8 { + return Err("Buffer too short for last seen timestamp".into()); + } + let last_seen_timestamp = i64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + + let mut accumulator = Self::new( + starting_measurement, + starting_timestamp, + last_seen_measurement, + last_seen_timestamp, + ); + offset += 8; + if buffer.len() >= offset + RESET_AWARE_WIRE_EXTENSION_LEN + && &buffer[offset..offset + RESET_AWARE_WIRE_MAGIC.len()] == RESET_AWARE_WIRE_MAGIC + { + offset += RESET_AWARE_WIRE_MAGIC.len(); + accumulator.total_increase = f64::from_le_bytes( + buffer[offset..offset + 8] + .try_into() + .expect("checked total-increase bytes"), + ); + offset += 8; + accumulator.sample_count = u64::from_le_bytes( + buffer[offset..offset + 8] + .try_into() + .expect("checked sample-count bytes"), + ); + } + Ok(accumulator) + } +} + +impl SerializableToSink for IncreaseAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "starting_measurement": self.starting_measurement.serialize_to_json(), + "starting_timestamp": self.starting_timestamp, + "last_seen_measurement": self.last_seen_measurement.serialize_to_json(), + "last_seen_timestamp": self.last_seen_timestamp, + "total_increase": self.total_increase, + "sample_count": self.sample_count, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + let starting_measurement_bytes = self.starting_measurement.serialize_to_bytes(); + let last_seen_measurement_bytes = self.last_seen_measurement.serialize_to_bytes(); + + let mut buffer = Vec::new(); + + // Starting measurement length and data + buffer.extend_from_slice(&(starting_measurement_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&starting_measurement_bytes); + + // Starting timestamp + buffer.extend_from_slice(&self.starting_timestamp.to_le_bytes()); + + // Last seen measurement length and data + buffer.extend_from_slice(&(last_seen_measurement_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&last_seen_measurement_bytes); + + // Last seen timestamp + buffer.extend_from_slice(&self.last_seen_timestamp.to_le_bytes()); + buffer.extend_from_slice(RESET_AWARE_WIRE_MAGIC); + buffer.extend_from_slice(&self.total_increase.to_le_bytes()); + buffer.extend_from_slice(&self.sample_count.to_le_bytes()); + + buffer + } +} + +impl MergeableAccumulator for IncreaseAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut accumulators = accumulators; + accumulators.sort_by_key(|accumulator| accumulator.starting_timestamp); + let mut result = accumulators[0].clone(); + + for acc in &accumulators[1..] { + result = Self::merge_pair(&result, acc); + } + + Ok(result) + } +} + +impl AggregateCore for IncreaseAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "IncreaseAccumulator" + } + + 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> { + // Check if other is also an IncreaseAccumulator + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge IncreaseAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + // Downcast to IncreaseAccumulator + let other_increase = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to IncreaseAccumulator")?; + + let merged = Self::merge_pair(self, other_increase); + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Increase + } + + fn approx_memory_bytes(&self) -> usize { + // Two Measurements + two i64s. Measurements are a few f64 fields. + std::mem::size_of::() + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query( + statistic, + (!query_kwargs.is_empty()).then_some(query_kwargs), + ) + } +} + +impl SingleSubpopulationAggregate for IncreaseAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Increase => Ok(self.extrapolated_value(query_kwargs, false)?), + Statistic::Rate => Ok(self.extrapolated_value(query_kwargs, true)?), + // For instant `sum [by (...)] (counter_metric)` Prometheus + // sums the latest cumulative value of each matching series. + // The IncreaseAccumulator already tracks that latest value + // in `last_seen_measurement`, so per-series Sum is just + // that scalar; the engine's outer aggregation groups by the + // `by` labels and adds the per-series totals across keys. + // + // See PR #108 audit conclusion (commit 4359e10) and issue + // ProjectASAP/ASAPCollector#46: pre-fix the ASAP tier ingested + // counters as IncreaseAccumulator and bare `sum by (...) ()` + // capability-missed because this trait did not answer Sum. + Statistic::Sum => Ok(self.last_seen_measurement.value), + _ => Err(format!("Unsupported statistic in IncreaseAccumulator: {statistic:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl IncreaseAccumulator { + fn extrapolated_value( + &self, + query_kwargs: Option<&HashMap>, + is_rate: bool, + ) -> Result> { + if self.sample_count < 2 || self.last_seen_timestamp <= self.starting_timestamp { + return Err("at least two ordered counter samples are required".into()); + } + let sampled_interval = (self.last_seen_timestamp - self.starting_timestamp) as f64 / 1000.0; + let Some(kwargs) = query_kwargs else { + return Ok(if is_rate { + self.total_increase / sampled_interval + } else { + self.total_increase + }); + }; + let range_start = kwargs + .get("range_start_ms") + .ok_or("missing range_start_ms")? + .parse::()?; + let range_end = kwargs + .get("range_end_ms") + .ok_or("missing range_end_ms")? + .parse::()?; + if range_end <= range_start { + return Err("invalid counter evaluation range".into()); + } + + let mut duration_to_start = + (self.starting_timestamp.saturating_sub(range_start)) as f64 / 1000.0; + let duration_to_end = (range_end.saturating_sub(self.last_seen_timestamp)) as f64 / 1000.0; + let average_sample_interval = sampled_interval / (self.sample_count - 1) as f64; + let extrapolation_threshold = average_sample_interval * 1.1; + + if self.total_increase > 0.0 && self.starting_measurement.value >= 0.0 { + let duration_to_zero = + sampled_interval * (self.starting_measurement.value / self.total_increase); + duration_to_start = duration_to_start.min(duration_to_zero); + } + let mut extrapolate_to = sampled_interval; + extrapolate_to += if duration_to_start < extrapolation_threshold { + duration_to_start.max(0.0) + } else { + average_sample_interval / 2.0 + }; + extrapolate_to += if duration_to_end < extrapolation_threshold { + duration_to_end.max(0.0) + } else { + average_sample_interval / 2.0 + }; + let mut factor = extrapolate_to / sampled_interval; + if is_rate { + factor /= (range_end - range_start) as f64 / 1000.0; + } + Ok(self.total_increase * factor) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_increase_accumulator_creation() { + let starting_measurement = Measurement::new(10.0); + let last_seen_measurement = Measurement::new(25.0); + let acc = IncreaseAccumulator::new( + starting_measurement.clone(), + 1000, + last_seen_measurement.clone(), + 2000, + ); + + assert_eq!(acc.starting_measurement.value, 10.0); + assert_eq!(acc.starting_timestamp, 1000); + assert_eq!(acc.last_seen_measurement.value, 25.0); + assert_eq!(acc.last_seen_timestamp, 2000); + } + + #[test] + fn test_increase_accumulator_update() { + let starting_measurement = Measurement::new(10.0); + let mut acc = IncreaseAccumulator::new( + starting_measurement.clone(), + 1000, + starting_measurement.clone(), + 1000, + ); + + let new_measurement = Measurement::new(25.0); + acc.update(new_measurement.clone(), 2000); + + assert_eq!(acc.last_seen_measurement.value, 25.0); + assert_eq!(acc.last_seen_timestamp, 2000); + assert_eq!(acc.starting_measurement.value, 10.0); // Should remain unchanged + } + + #[test] + fn test_increase_accumulator_query() { + let starting_measurement = Measurement::new(10.0); + let last_seen_measurement = Measurement::new(25.0); + let acc = IncreaseAccumulator::new( + starting_measurement, + 1000, + last_seen_measurement, + 3000, // 2 second difference + ); + + // Test increase calculation + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, None).unwrap(), + 15.0 + ); + + // Test rate calculation (per second) + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, None).unwrap(), + 7.5 + ); // 15.0 / 2.0 + + // Statistic::Sum returns the latest cumulative counter value, + // matching Prometheus semantics for instant `sum()`. + // (Issue ProjectASAP/ASAPCollector#46, PR #108 diagnosis.) + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), + 25.0 + ); + + // Unsupported statistics still error. + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); + } + + #[test] + fn prometheus_counter_reset_and_boundary_extrapolation() { + let mut acc = IncreaseAccumulator::new( + Measurement::new(10.0), + 10_000, + Measurement::new(10.0), + 10_000, + ); + acc.update(Measurement::new(20.0), 20_000); + acc.update(Measurement::new(3.0), 30_000); + acc.update(Measurement::new(13.0), 50_000); + assert_eq!(acc.total_increase, 23.0); + assert_eq!(acc.sample_count, 4); + + let kwargs = HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + let increase = + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, Some(&kwargs)) + .unwrap(); + let rate = crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, Some(&kwargs)) + .unwrap(); + assert!((increase - 34.5).abs() < 1e-12); + assert!((rate - 0.575).abs() < 1e-12); + } + + #[test] + fn pane_merge_preserves_resets_and_prometheus_extrapolation() { + let mut left = IncreaseAccumulator::new( + Measurement::new(10.0), + 10_000, + Measurement::new(10.0), + 10_000, + ); + left.update(Measurement::new(20.0), 20_000); + let mut right = + IncreaseAccumulator::new(Measurement::new(3.0), 30_000, Measurement::new(3.0), 30_000); + right.update(Measurement::new(13.0), 50_000); + let merged = IncreaseAccumulator::merge_accumulators(vec![right, left]).unwrap(); + assert_eq!(merged.total_increase, 23.0); + assert_eq!(merged.sample_count, 4); + let kwargs = HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&merged, Statistic::Increase, Some(&kwargs)) + .unwrap(), + 34.5 + ); + } + + #[test] + fn counter_sds_state_is_constant_size_per_pane() { + let mut acc = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(0.0), 0); + let initial = acc.serialize_to_bytes().len(); + for second in 1..=86_400 { + acc.update(Measurement::new(second as f64), second * 1_000); + } + assert_eq!(acc.serialize_to_bytes().len(), initial); + assert_eq!(acc.sample_count, 86_401); + assert_eq!( + acc.approx_memory_bytes(), + std::mem::size_of::() + ); + } + + #[test] + fn test_increase_accumulator_sum_is_latest_cumulative_value() { + // Instant `sum ()` semantics: the per-series summand is + // the latest cumulative counter value. Two series with latest + // values 100 and 50 (started at 10 and 5 respectively) should + // each report Sum = 100 and Sum = 50 — the engine's `sum by` + // outer aggregation does the cross-series total. + let acc_a = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000); + let acc_b = + IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc_a, Statistic::Sum, None).unwrap(), + 100.0 + ); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc_b, Statistic::Sum, None).unwrap(), + 50.0 + ); + } + + #[test] + fn test_increase_accumulator_merge() { + let acc1 = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(20.0), 2000); + let acc2 = IncreaseAccumulator::new( + Measurement::new(5.0), + 500, // Earlier start + Measurement::new(15.0), + 1500, + ); + let acc3 = IncreaseAccumulator::new( + Measurement::new(20.0), + 2000, + Measurement::new(30.0), + 3000, // Later end + ); + + let merged = + >::merge_accumulators( + vec![acc1, acc2, acc3], + ) + .unwrap(); + + // Should use earliest start and latest end + assert_eq!(merged.starting_measurement.value, 5.0); + assert_eq!(merged.starting_timestamp, 500); + assert_eq!(merged.last_seen_measurement.value, 30.0); + assert_eq!(merged.last_seen_timestamp, 3000); + } + + #[test] + fn test_increase_accumulator_serialization() { + let acc = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(25.0), 2000); + + // Test JSON serialization + let json = acc.serialize_to_json(); + let deserialized = IncreaseAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!( + acc.starting_measurement.value, + deserialized.starting_measurement.value + ); + assert_eq!(acc.starting_timestamp, deserialized.starting_timestamp); + assert_eq!( + acc.last_seen_measurement.value, + deserialized.last_seen_measurement.value + ); + assert_eq!(acc.last_seen_timestamp, deserialized.last_seen_timestamp); + + // Test byte serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = IncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + acc.starting_measurement.value, + deserialized_bytes.starting_measurement.value + ); + assert_eq!( + acc.starting_timestamp, + deserialized_bytes.starting_timestamp + ); + assert_eq!( + acc.last_seen_measurement.value, + deserialized_bytes.last_seen_measurement.value + ); + assert_eq!( + acc.last_seen_timestamp, + deserialized_bytes.last_seen_timestamp + ); + assert_eq!(acc.total_increase, deserialized_bytes.total_increase); + assert_eq!(acc.sample_count, deserialized_bytes.sample_count); + + let legacy = &bytes[..bytes.len() - RESET_AWARE_WIRE_EXTENSION_LEN]; + let legacy_value = IncreaseAccumulator::deserialize_from_bytes(legacy).unwrap(); + assert_eq!(legacy_value.total_increase, 15.0); + assert_eq!(legacy_value.sample_count, 2); + } + + #[test] + fn test_trait_object() { + let acc: Box = Box::new(IncreaseAccumulator::new( + Measurement::new(10.0), + 1000, + Measurement::new(25.0), + 2000, + )); + + assert_eq!(acc.type_name(), "IncreaseAccumulator"); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/keyed_counter_state.rs b/crates/asap_summary_state/src/summary_kernels/keyed_counter_state.rs new file mode 100644 index 000000000..d65860ad3 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/keyed_counter_state.rs @@ -0,0 +1,530 @@ +use crate::codec::KeyCodec; +use crate::summary_kernels::IncreaseAccumulator; +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Accumulator that maintains separate increase accumulators for multiple keys +/// Allows tracking rate/increase for different label combinations +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KeyedCounterState { + pub increases: HashMap, +} + +impl KeyedCounterState { + pub fn new() -> Self { + Self { + increases: HashMap::new(), + } + } + + pub fn update(&mut self, key: KeyByLabelValues, accumulator: IncreaseAccumulator) { + self.increases.insert(key, accumulator); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let mut accumulator = Self::new(); + + if let Some(entries) = data["entries"].as_array() { + for entry in entries { + let key = KeyByLabelValues::deserialize_from_json(&entry["key"])?; + let increase_data = + IncreaseAccumulator::deserialize_from_json(&entry["increase_data"])?; + accumulator.increases.insert(key, increase_data); + } + } + + Ok(accumulator) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut accumulator = Self::new(); + let mut offset = 0; + + // Read number of entries + if buffer.len() < 4 { + return Err("Buffer too short for entry count".into()); + } + let num_entries = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; + offset += 4; + + for _ in 0..num_entries { + // Read key length and key + if offset + 4 > buffer.len() { + return Err("Buffer too short for key length".into()); + } + let key_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if offset + key_length > buffer.len() { + return Err("Buffer too short for key data".into()); + } + let key = + KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; + offset += key_length; + + // Read IncreaseAccumulator data + if offset >= buffer.len() { + return Err("Buffer too short for increase accumulator data".into()); + } + let consumed_bytes = + IncreaseAccumulator::serialized_len_from_prefix(&buffer[offset..])?; + let increase_data = IncreaseAccumulator::deserialize_from_bytes( + &buffer[offset..offset + consumed_bytes], + )?; + offset += consumed_bytes; + + accumulator.increases.insert(key, increase_data); + } + + Ok(accumulator) + } +} + +impl Default for KeyedCounterState { + fn default() -> Self { + Self::new() + } +} + +impl SerializableToSink for KeyedCounterState { + fn serialize_to_json(&self) -> Value { + let entries: Vec = self + .increases + .iter() + .map(|(key, data)| { + serde_json::json!({ + "key": key.serialize_to_json(), + "increase_data": data.serialize_to_json() + }) + }) + .collect(); + + serde_json::json!({ + "entries": entries + }) + } + + fn serialize_to_bytes(&self) -> Vec { + let mut buffer = Vec::new(); + + // Write number of entries + buffer.extend_from_slice(&(self.increases.len() as u32).to_le_bytes()); + + // Write each key-value pair + for (key, data) in &self.increases { + let key_bytes = key.serialize_to_bytes(); + buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&key_bytes); + + let data_bytes = data.serialize_to_bytes(); + buffer.extend_from_slice(&data_bytes); + } + + buffer + } +} + +impl AggregateCore for KeyedCounterState { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "KeyedCounterState" + } + + 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> { + // Check if other is also a KeyedCounterState + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge KeyedCounterState with {}", + other.get_accumulator_type() + ) + .into()); + } + + // Downcast to KeyedCounterState + let other_multiple_increase = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to KeyedCounterState")?; + + // Clone self once, then merge each matching counter with the same + // reset-aware, boundary-aware implementation used by the unkeyed path. + let mut merged = self.clone(); + for (key, data) in &other_multiple_increase.increases { + if let Some(existing_data) = merged.increases.get_mut(key) { + *existing_data = IncreaseAccumulator::merge_accumulators(vec![ + existing_data.clone(), + data.clone(), + ])?; + } else { + merged.increases.insert(key.clone(), data.clone()); + } + } + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Increase + } + + fn approx_memory_bytes(&self) -> usize { + // HashMap. IncreaseAccumulator is ~64 B, + // per-entry key/overhead is ~96 B. + const BYTES_PER_ENTRY: usize = 160; + std::mem::size_of::() + self.increases.len() * BYTES_PER_ENTRY + } + + fn get_keys(&self) -> Option> { + Some(self.increases.keys().cloned().collect()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedCounterState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedCounterState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + query_kwargs: Option<&HashMap>, + ) -> Result> { + let data = self + .increases + .get(key) + .ok_or_else(|| format!("Key {key} not found in KeyedCounterState"))?; + + data.query(statistic, query_kwargs) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedCounterState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedCounterState::new(); + + for accumulator in accumulators { + for (key, data) in accumulator.increases { + if let Some(existing_data) = result.increases.get_mut(&key) { + *existing_data = + IncreaseAccumulator::merge_accumulators(vec![existing_data.clone(), data])?; + } else { + result.increases.insert(key, data); + } + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Measurement; + + fn create_test_increase_accumulator(start_val: f64, end_val: f64) -> IncreaseAccumulator { + IncreaseAccumulator::new( + Measurement::new(start_val), + 1000, + Measurement::new(end_val), + 2000, + ) + } + + fn create_test_increase_accumulator_with_time( + start_val: f64, + start_time: i64, + end_val: f64, + end_time: i64, + ) -> IncreaseAccumulator { + IncreaseAccumulator::new( + Measurement::new(start_val), + start_time, + Measurement::new(end_val), + end_time, + ) + } + + #[test] + fn test_keyed_counter_state_creation() { + let acc = KeyedCounterState::new(); + assert!(acc.increases.is_empty()); + } + + #[test] + fn test_keyed_counter_state_update() { + let mut acc = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + let increase1 = create_test_increase_accumulator(10.0, 25.0); + let increase2 = create_test_increase_accumulator(5.0, 15.0); + + acc.update(key1.clone(), increase1); + acc.update(key2.clone(), increase2); + + assert_eq!(acc.increases.len(), 2); + assert!(acc.increases.contains_key(&key1)); + assert!(acc.increases.contains_key(&key2)); + } + + #[test] + fn test_keyed_counter_state_query() { + let mut acc = KeyedCounterState::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let increase_acc = create_test_increase_accumulator(10.0, 25.0); + acc.update(key.clone(), increase_acc); + + // Test increase query + assert_eq!(acc.query(Statistic::Increase, &key, None).unwrap(), 15.0); + + // Test rate query (15.0 increase over 1 second = 15.0 per second) + assert_eq!(acc.query(Statistic::Rate, &key, None).unwrap(), 15.0); + + // Sum returns the latest cumulative counter value for the + // queried key (per-series Prometheus `sum()` semantics; + // see issue ProjectASAP/ASAPCollector#46 and PR #108 diagnosis). + // The series here was created with last_seen=25.0. + assert_eq!(acc.query(Statistic::Sum, &key, None).unwrap(), 25.0); + + // Unsupported statistic still errors. + assert!(acc.query(Statistic::Min, &key, None).is_err()); + + let unknown_key = KeyByLabelValues::new(); + assert!(acc.query(Statistic::Increase, &unknown_key, None).is_err()); + } + + #[test] + fn test_keyed_counter_state_sum_per_key() { + // `sum by (zone) (counter)` reaches KeyedCounterState + // only when the ASAP-tier ingest groups multiple series under + // a single accumulator (the `Multiple*` variant). In that case + // each per-key Sum should be the series' latest cumulative + // value; the engine's outer `by` aggregation does the cross-key + // grouping. (Issue ProjectASAP/ASAPCollector#46.) + let mut acc = KeyedCounterState::new(); + let east = KeyByLabelValues::new_with_labels(vec!["us-east-1".to_string()]); + let west = KeyByLabelValues::new_with_labels(vec!["us-west-2".to_string()]); + + acc.update( + east.clone(), + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000), + ); + acc.update( + west.clone(), + IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000), + ); + + assert_eq!(acc.query(Statistic::Sum, &east, None).unwrap(), 100.0); + assert_eq!(acc.query(Statistic::Sum, &west, None).unwrap(), 50.0); + } + + #[test] + fn test_keyed_counter_state_merge() { + let mut acc1 = KeyedCounterState::new(); + let mut acc2 = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + // Add different keys to each accumulator + acc1.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); + acc2.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); + + // Also add overlapping key with different time ranges (later timestamps) + acc2.update( + key1.clone(), + create_test_increase_accumulator_with_time(15.0, 2000, 30.0, 3000), + ); // Later time range + + let merged = KeyedCounterState::merge_accumulators(vec![acc1, acc2]).unwrap(); + + assert_eq!(merged.increases.len(), 2); + assert!(merged.increases.contains_key(&key1)); + assert!(merged.increases.contains_key(&key2)); + + // The merged key1 should have the full range (earliest start to latest end) + let merged_key1 = merged.increases.get(&key1).unwrap(); + assert_eq!(merged_key1.starting_measurement.value, 10.0); // Earlier start + assert_eq!(merged_key1.last_seen_measurement.value, 30.0); // Later end + } + + #[test] + fn test_keyed_counter_state_serialization() { + let mut acc = KeyedCounterState::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let second_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + let mut reset_aware = create_test_increase_accumulator(10.0, 25.0); + reset_aware.update(Measurement::new(3.0), 3000); + acc.update(key.clone(), reset_aware); + acc.update( + second_key.clone(), + create_test_increase_accumulator(4.0, 9.0), + ); + + // Test JSON serialization + let json_value = acc.serialize_to_json(); + let deserialized = KeyedCounterState::deserialize_from_json(&json_value).unwrap(); + + assert_eq!(deserialized.increases.len(), 2); + let deserialized_acc = deserialized.increases.get(&key).unwrap(); + assert_eq!(deserialized_acc.starting_measurement.value, 10.0); + assert_eq!(deserialized_acc.last_seen_measurement.value, 3.0); + assert_eq!(deserialized_acc.total_increase, 18.0); + + // Test binary serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = KeyedCounterState::deserialize_from_bytes(&bytes).unwrap(); + + assert_eq!(deserialized_bytes.increases.len(), 2); + let deserialized_acc_bytes = deserialized_bytes.increases.get(&key).unwrap(); + assert_eq!(deserialized_acc_bytes.starting_measurement.value, 10.0); + assert_eq!(deserialized_acc_bytes.last_seen_measurement.value, 3.0); + assert_eq!(deserialized_acc_bytes.total_increase, 18.0); + assert_eq!( + deserialized_bytes + .increases + .get(&second_key) + .unwrap() + .last_seen_measurement + .value, + 9.0 + ); + } + + #[test] + fn test_keyed_counter_state_get_keys() { + let mut acc = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); + acc.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); + + let keys = acc.get_keys().unwrap(); + assert_eq!(keys.len(), 2); + assert!(keys.contains(&key1)); + assert!(keys.contains(&key2)); + } + + #[test] + fn test_trait_object() { + let mut acc = KeyedCounterState::new(); + let key = KeyByLabelValues::new(); + acc.update(key.clone(), create_test_increase_accumulator(10.0, 25.0)); + + let trait_obj: Box = Box::new(acc); + assert_eq!( + trait_obj.query(Statistic::Increase, &key, None).unwrap(), + 15.0 + ); + + let keys = trait_obj.get_keys().unwrap(); + assert_eq!(keys.len(), 1); + } + + // #[test] + // fn test_keyed_counter_state_arroyo_deserialization() { + // // Create test data in Arroyo MessagePack format + // // Format: {key: [starting_value, starting_timestamp, last_seen_value, last_seen_timestamp]} + // let mut test_data = std::collections::HashMap::new(); + // test_data.insert("web;service".to_string(), vec![10.0, 1000.0, 25.0, 2000.0]); + // test_data.insert("api;service".to_string(), vec![5.0, 1500.0, 15.0, 2500.0]); + + // // Serialize to MessagePack + // let arroyo_buffer = rmp_serde::to_vec(&test_data).unwrap(); + + // // Test Arroyo deserialization + // let deserialized_acc = + // KeyedCounterState::deserialize_from_bytes_arroyo(&arroyo_buffer).unwrap(); + + // // Verify the deserialized accumulator has the correct data + // assert_eq!(deserialized_acc.increases.len(), 2); + + // // Check first key (web;service) + // let keys: Vec<_> = deserialized_acc.increases.keys().collect(); + // let key1 = keys + // .iter() + // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "web")) + // .unwrap(); + + // let increase1 = deserialized_acc.increases.get(key1).unwrap(); + // assert_eq!(increase1.starting_measurement.value, 10.0); + // assert_eq!(increase1.starting_timestamp, 1000); + // assert_eq!(increase1.last_seen_measurement.value, 25.0); + // assert_eq!(increase1.last_seen_timestamp, 2000); + + // // Check second key (api;service) + // let key2 = keys + // .iter() + // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "api")) + // .unwrap(); + + // let increase2 = deserialized_acc.increases.get(key2).unwrap(); + // assert_eq!(increase2.starting_measurement.value, 5.0); + // assert_eq!(increase2.starting_timestamp, 1500); + // assert_eq!(increase2.last_seen_measurement.value, 15.0); + // assert_eq!(increase2.last_seen_timestamp, 2500); + + // // Test querying + // assert_eq!( + // deserialized_acc.query(Statistic::Increase, key1).unwrap(), + // 15.0 + // ); // 25.0 - 10.0 + // assert_eq!( + // deserialized_acc.query(Statistic::Increase, key2).unwrap(), + // 10.0 + // ); // 15.0 - 5.0 + // } +} diff --git a/crates/asap_summary_state/src/summary_kernels/keyed_max_state.rs b/crates/asap_summary_state/src/summary_kernels/keyed_max_state.rs new file mode 100644 index 000000000..38b9094b4 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/keyed_max_state.rs @@ -0,0 +1,336 @@ +use crate::codec::KeyCodec; +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Exact per-key maximum over many populations, mergeable by comparison. +/// +/// The minimum direction is +/// [`KeyedMinState`](super::keyed_min_state::KeyedMinState), +/// a separate type: these used to be one `MultipleMinMaxAccumulator` whose +/// direction lived in a `sub_type` string that every layer above had to carry +/// alongside the family. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct KeyedMaxState { + pub values: HashMap, +} + +impl KeyedMaxState { + pub fn new() -> Self { + Self::default() + } + + pub fn new_with_values(values: HashMap) -> Self { + Self { values } + } + + pub fn update(&mut self, key: KeyByLabelValues, value: f64) { + let current = self.values.entry(key).or_insert(f64::NEG_INFINITY); + if value > *current { + *current = value; + } + } + + pub fn add_value(&mut self, key: KeyByLabelValues, value: f64) { + self.values.insert(key, value); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let values_data = data["values"] + .as_object() + .ok_or("Missing or invalid 'values' field")?; + + let mut values = HashMap::new(); + for (key_str, value) in values_data { + let key_json: Value = serde_json::from_str(key_str)?; + let key = KeyByLabelValues::deserialize_from_json(&key_json)?; + let val = value.as_f64().ok_or("Invalid value")?; + values.insert(key, val); + } + + Ok(Self { values }) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut offset = 0; + + // Read number of entries + if buffer.len() < 4 { + return Err("Buffer too short for entry count".into()); + } + let num_entries = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + let mut values = HashMap::new(); + + for _ in 0..num_entries { + // Read key length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for key length".into()); + } + let key_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + key_length { + return Err("Buffer too short for key data".into()); + } + let key = + KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; + offset += key_length; + + // Read value + if buffer.len() < offset + 8 { + return Err("Buffer too short for value".into()); + } + let value = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + + values.insert(key, value); + } + + Ok(Self { values }) + } +} + +impl SerializableToSink for KeyedMaxState { + fn serialize_to_json(&self) -> Value { + let mut values_obj = serde_json::Map::new(); + for (key, value) in &self.values { + let key_json = key.serialize_to_json(); + let key_str = serde_json::to_string(&key_json).unwrap(); + values_obj.insert( + key_str, + Value::Number(serde_json::Number::from_f64(*value).unwrap()), + ); + } + + serde_json::json!({ "values": values_obj }) + } + + fn serialize_to_bytes(&self) -> Vec { + let mut buffer = Vec::new(); + + // Write number of entries + buffer.extend_from_slice(&(self.values.len() as u32).to_le_bytes()); + + // Write each key-value pair + for (key, value) in &self.values { + let key_bytes = key.serialize_to_bytes(); + + // Write key length and data + buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&key_bytes); + + // Write value + buffer.extend_from_slice(&value.to_le_bytes()); + } + + buffer + } +} + +impl AggregateCore for KeyedMaxState { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "KeyedMaxState" + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge KeyedMaxState with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_multiple = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to KeyedMaxState")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_multiple.clone()])?; + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Max + } + + fn approx_memory_bytes(&self) -> usize { + const BYTES_PER_ENTRY: usize = 96; + std::mem::size_of::() + self.values.len() * BYTES_PER_ENTRY + } + + fn get_keys(&self) -> Option> { + Some(self.values.keys().cloned().collect()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedMaxState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedMaxState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Max => self + .values + .get(key) + .copied() + .ok_or_else(|| format!("Key {key} not found in KeyedMaxState").into()), + other => Err(format!("Unsupported statistic in KeyedMaxState: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedMaxState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedMaxState::new(); + + for acc in accumulators { + for (key, value) in acc.values { + result.update(key, value); + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(value: &str) -> KeyByLabelValues { + KeyByLabelValues::new_with_labels(vec![value.to_string()]) + } + + #[test] + fn keeps_the_largest_per_key() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 10.0); + acc.update(key("a"), 5.0); + acc.update(key("a"), 15.0); + acc.update(key("b"), 7.0); + + assert_eq!(acc.query(Statistic::Max, &key("a"), None).unwrap(), 15.0); + assert_eq!(acc.query(Statistic::Max, &key("b"), None).unwrap(), 7.0); + } + + #[test] + fn refuses_the_opposite_statistic_and_unknown_keys() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 1.0); + assert!(acc.query(Statistic::Min, &key("a"), None).is_err()); + assert!(acc.query(Statistic::Max, &key("missing"), None).is_err()); + } + + #[test] + fn merges_per_key() { + let mut left = KeyedMaxState::new(); + left.update(key("a"), 10.0); + let mut right = KeyedMaxState::new(); + right.update(key("a"), 5.0); + right.update(key("b"), 3.0); + + let merged = + >::merge_accumulators(vec![ + left, right, + ]) + .unwrap(); + + assert_eq!(merged.query(Statistic::Max, &key("a"), None).unwrap(), 10.0); + assert_eq!(merged.query(Statistic::Max, &key("b"), None).unwrap(), 3.0); + } + + #[test] + fn refuses_to_merge_with_the_opposite_direction() { + use super::super::keyed_min_state::KeyedMinState; + let mine = KeyedMaxState::new(); + let theirs = KeyedMinState::new(); + assert!(mine.merge_with(&theirs).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 4.0); + + let json = acc.serialize_to_json(); + let from_json = KeyedMaxState::deserialize_from_json(&json).unwrap(); + assert_eq!( + from_json.query(Statistic::Max, &key("a"), None).unwrap(), + 4.0 + ); + + let bytes = acc.serialize_to_bytes(); + let from_bytes = KeyedMaxState::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + from_bytes.query(Statistic::Max, &key("a"), None).unwrap(), + 4.0 + ); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/keyed_min_state.rs b/crates/asap_summary_state/src/summary_kernels/keyed_min_state.rs new file mode 100644 index 000000000..b8fc2ddfb --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/keyed_min_state.rs @@ -0,0 +1,336 @@ +use crate::codec::KeyCodec; +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Exact per-key minimum over many populations, mergeable by comparison. +/// +/// The maximum direction is +/// [`KeyedMaxState`](super::keyed_max_state::KeyedMaxState), +/// a separate type: these used to be one `MultipleMinMaxAccumulator` whose +/// direction lived in a `sub_type` string that every layer above had to carry +/// alongside the family. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct KeyedMinState { + pub values: HashMap, +} + +impl KeyedMinState { + pub fn new() -> Self { + Self::default() + } + + pub fn new_with_values(values: HashMap) -> Self { + Self { values } + } + + pub fn update(&mut self, key: KeyByLabelValues, value: f64) { + let current = self.values.entry(key).or_insert(f64::INFINITY); + if value < *current { + *current = value; + } + } + + pub fn add_value(&mut self, key: KeyByLabelValues, value: f64) { + self.values.insert(key, value); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let values_data = data["values"] + .as_object() + .ok_or("Missing or invalid 'values' field")?; + + let mut values = HashMap::new(); + for (key_str, value) in values_data { + let key_json: Value = serde_json::from_str(key_str)?; + let key = KeyByLabelValues::deserialize_from_json(&key_json)?; + let val = value.as_f64().ok_or("Invalid value")?; + values.insert(key, val); + } + + Ok(Self { values }) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut offset = 0; + + // Read number of entries + if buffer.len() < 4 { + return Err("Buffer too short for entry count".into()); + } + let num_entries = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + let mut values = HashMap::new(); + + for _ in 0..num_entries { + // Read key length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for key length".into()); + } + let key_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + key_length { + return Err("Buffer too short for key data".into()); + } + let key = + KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; + offset += key_length; + + // Read value + if buffer.len() < offset + 8 { + return Err("Buffer too short for value".into()); + } + let value = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + + values.insert(key, value); + } + + Ok(Self { values }) + } +} + +impl SerializableToSink for KeyedMinState { + fn serialize_to_json(&self) -> Value { + let mut values_obj = serde_json::Map::new(); + for (key, value) in &self.values { + let key_json = key.serialize_to_json(); + let key_str = serde_json::to_string(&key_json).unwrap(); + values_obj.insert( + key_str, + Value::Number(serde_json::Number::from_f64(*value).unwrap()), + ); + } + + serde_json::json!({ "values": values_obj }) + } + + fn serialize_to_bytes(&self) -> Vec { + let mut buffer = Vec::new(); + + // Write number of entries + buffer.extend_from_slice(&(self.values.len() as u32).to_le_bytes()); + + // Write each key-value pair + for (key, value) in &self.values { + let key_bytes = key.serialize_to_bytes(); + + // Write key length and data + buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&key_bytes); + + // Write value + buffer.extend_from_slice(&value.to_le_bytes()); + } + + buffer + } +} + +impl AggregateCore for KeyedMinState { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "KeyedMinState" + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge KeyedMinState with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_multiple = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to KeyedMinState")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_multiple.clone()])?; + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Min + } + + fn approx_memory_bytes(&self) -> usize { + const BYTES_PER_ENTRY: usize = 96; + std::mem::size_of::() + self.values.len() * BYTES_PER_ENTRY + } + + fn get_keys(&self) -> Option> { + Some(self.values.keys().cloned().collect()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedMinState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedMinState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Min => self + .values + .get(key) + .copied() + .ok_or_else(|| format!("Key {key} not found in KeyedMinState").into()), + other => Err(format!("Unsupported statistic in KeyedMinState: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedMinState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedMinState::new(); + + for acc in accumulators { + for (key, value) in acc.values { + result.update(key, value); + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(value: &str) -> KeyByLabelValues { + KeyByLabelValues::new_with_labels(vec![value.to_string()]) + } + + #[test] + fn keeps_the_smallest_per_key() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 10.0); + acc.update(key("a"), 5.0); + acc.update(key("a"), 15.0); + acc.update(key("b"), 7.0); + + assert_eq!(acc.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); + assert_eq!(acc.query(Statistic::Min, &key("b"), None).unwrap(), 7.0); + } + + #[test] + fn refuses_the_opposite_statistic_and_unknown_keys() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 1.0); + assert!(acc.query(Statistic::Max, &key("a"), None).is_err()); + assert!(acc.query(Statistic::Min, &key("missing"), None).is_err()); + } + + #[test] + fn merges_per_key() { + let mut left = KeyedMinState::new(); + left.update(key("a"), 10.0); + let mut right = KeyedMinState::new(); + right.update(key("a"), 5.0); + right.update(key("b"), 3.0); + + let merged = + >::merge_accumulators(vec![ + left, right, + ]) + .unwrap(); + + assert_eq!(merged.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); + assert_eq!(merged.query(Statistic::Min, &key("b"), None).unwrap(), 3.0); + } + + #[test] + fn refuses_to_merge_with_the_opposite_direction() { + use super::super::keyed_max_state::KeyedMaxState; + let mine = KeyedMinState::new(); + let theirs = KeyedMaxState::new(); + assert!(mine.merge_with(&theirs).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 4.0); + + let json = acc.serialize_to_json(); + let from_json = KeyedMinState::deserialize_from_json(&json).unwrap(); + assert_eq!( + from_json.query(Statistic::Min, &key("a"), None).unwrap(), + 4.0 + ); + + let bytes = acc.serialize_to_bytes(); + let from_bytes = KeyedMinState::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + from_bytes.query(Statistic::Min, &key("a"), None).unwrap(), + 4.0 + ); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/keyed_sum_count.rs b/crates/asap_summary_state/src/summary_kernels/keyed_sum_count.rs new file mode 100644 index 000000000..7ff044d7a --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/keyed_sum_count.rs @@ -0,0 +1,559 @@ +use crate::codec::KeyCodec; +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; +use planner_types::post_asap::ExactKind; + +fn sum_family() -> ExactKind { + ExactKind::Sum +} + +/// Accumulator that maintains separate sum values for multiple keys +/// Allows querying sums for specific label combinations +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KeyedSumCountAccumulator { + #[serde(default = "sum_family")] + pub family: ExactKind, + pub sums: HashMap, + #[serde(default)] + pub counts: HashMap, +} + +impl KeyedSumCountAccumulator { + pub fn new() -> Self { + Self::for_family(ExactKind::Sum) + } + + pub fn for_family(family: ExactKind) -> Self { + assert!(matches!(family, ExactKind::Sum | ExactKind::Count)); + Self { + family, + sums: HashMap::new(), + counts: HashMap::new(), + } + } + + pub fn update(&mut self, key: KeyByLabelValues, value: f64) { + let is_new = !self.sums.contains_key(&key); + *self.sums.entry(key.clone()).or_insert(0.0) += value; + if let Some(count) = self.counts.get(&key).copied() { + if let Some(next) = count.checked_add(1).filter(|next| *next != u64::MAX) { + self.counts.insert(key, next); + } else { + self.counts.remove(&key); + } + } else if is_new { + self.counts.insert(key, 1); + } + } + + pub fn add_sum(&mut self, key: KeyByLabelValues, sum: f64) { + self.counts.remove(&key); + self.sums.insert(key, sum); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let sums_data = data["sums"] + .as_object() + .ok_or("Missing or invalid 'sums' field")?; + + let mut sums = HashMap::new(); + for (key_str, value) in sums_data { + let key_json: Value = serde_json::from_str(key_str)?; + let key = KeyByLabelValues::deserialize_from_json(&key_json)?; + let sum = value.as_f64().ok_or("Invalid sum value")?; + sums.insert(key, sum); + } + + let mut counts = HashMap::new(); + if let Some(counts_data) = data.get("counts").and_then(Value::as_object) { + for (key_str, value) in counts_data { + let key_json: Value = serde_json::from_str(key_str)?; + let key = KeyByLabelValues::deserialize_from_json(&key_json)?; + let count = value.as_u64().ok_or("Invalid count value")?; + if !sums.contains_key(&key) { + return Err("Count key missing from sums".into()); + } + counts.insert(key, count); + } + } + let family = match data.get("family").and_then(Value::as_str) { + None | Some("Sum") => ExactKind::Sum, + Some("Count") => ExactKind::Count, + _ => return Err("Invalid keyed additive family".into()), + }; + Ok(Self { + family, + sums, + counts, + }) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut offset = 0; + + // Read number of entries + if buffer.len() < 4 { + return Err("Buffer too short for entry count".into()); + } + let num_entries = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + let mut sums = HashMap::new(); + let mut keys = Vec::new(); + + for _ in 0..num_entries { + // Read key length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for key length".into()); + } + let key_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + key_length { + return Err("Buffer too short for key data".into()); + } + let key = + KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; + offset += key_length; + + // Read sum value + if buffer.len() < offset + 8 { + return Err("Buffer too short for sum value".into()); + } + let sum = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + + keys.push(key.clone()); + sums.insert(key, sum); + } + let remaining = buffer.len() - offset; + let count_bytes = num_entries + .checked_mul(8) + .ok_or("Count section too large")?; + if remaining != 0 && remaining != count_bytes && remaining != count_bytes + 1 { + return Err("Invalid count section length".into()); + } + let mut counts = HashMap::new(); + if count_bytes != 0 && remaining >= count_bytes { + for key in keys { + let count = u64::from_le_bytes(buffer[offset..offset + 8].try_into()?); + offset += 8; + if count != u64::MAX { + counts.insert(key, count); + } + } + } + let family = if remaining == count_bytes + 1 { + match buffer[offset] { + 0 => ExactKind::Sum, + 1 => ExactKind::Count, + _ => return Err("Invalid keyed additive family tag".into()), + } + } else { + ExactKind::Sum + }; + Ok(Self { + family, + sums, + counts, + }) + } +} + +impl Default for KeyedSumCountAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl SerializableToSink for KeyedSumCountAccumulator { + fn serialize_to_json(&self) -> Value { + let mut sums_obj = serde_json::Map::new(); + for (key, sum) in &self.sums { + let key_json = key.serialize_to_json(); + let key_str = serde_json::to_string(&key_json).unwrap(); + sums_obj.insert( + key_str, + Value::Number(serde_json::Number::from_f64(*sum).unwrap()), + ); + } + + let mut counts_obj = serde_json::Map::new(); + for (key, count) in &self.counts { + let key_str = serde_json::to_string(&key.serialize_to_json()).unwrap(); + counts_obj.insert(key_str, Value::from(*count)); + } + + serde_json::json!({ + "family": if self.family == ExactKind::Count { "Count" } else { "Sum" }, + "sums": sums_obj, + "counts": counts_obj + }) + } + + fn serialize_to_bytes(&self) -> Vec { + let mut buffer = Vec::new(); + + // Write number of entries + buffer.extend_from_slice(&(self.sums.len() as u32).to_le_bytes()); + + // Write each key-value pair + let mut ordered_keys = Vec::with_capacity(self.sums.len()); + for (key, sum) in &self.sums { + ordered_keys.push(key); + let key_bytes = key.serialize_to_bytes(); + + // Write key length and data + buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&key_bytes); + + // Write sum value + buffer.extend_from_slice(&sum.to_le_bytes()); + } + + for key in ordered_keys { + buffer.extend_from_slice( + &self + .counts + .get(key) + .copied() + .unwrap_or(u64::MAX) + .to_le_bytes(), + ); + } + + buffer.push(if self.family == ExactKind::Count { + 1 + } else { + 0 + }); + + buffer + } +} + +impl AggregateCore for KeyedSumCountAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "KeyedSumCountAccumulator" + } + + 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> { + // Check if other is also a KeyedSumCountAccumulator + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge KeyedSumCountAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + // Downcast to KeyedSumCountAccumulator + let other_multiple_sum = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to KeyedSumCountAccumulator")?; + + // Use the existing merge_accumulators method + let merged = Self::merge_accumulators(vec![self.clone(), other_multiple_sum.clone()])?; + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + if self.family == ExactKind::Count { + AggregationType::Count + } else { + AggregationType::Sum + } + } + + fn approx_memory_bytes(&self) -> usize { + // HashMap. Label strings dominate; use a + // conservative per-entry estimate plus HashMap overhead. + const BYTES_PER_ENTRY: usize = 112; + std::mem::size_of::() + self.sums.len() * BYTES_PER_ENTRY + } + + fn get_keys(&self) -> Option> { + Some(self.sums.keys().cloned().collect()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for KeyedSumCountAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedSumCountAccumulator { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match (&self.family, statistic) { + (ExactKind::Sum, Statistic::Sum) => self.sums.get(key).copied().ok_or_else(|| { + "Key not found in KeyedSumCountAccumulator" + .to_string() + .into() + }), + (ExactKind::Count, Statistic::Count) => self + .counts + .get(key) + .map(|count| *count as f64) + .ok_or_else(|| { + "Sample count unavailable in KeyedSumCountAccumulator" + .to_string() + .into() + }), + _ => Err( + format!("Unsupported statistic in KeyedSumCountAccumulator: {statistic:?}").into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedSumCountAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let family = accumulators[0].family.clone(); + if accumulators.iter().any(|acc| acc.family != family) { + return Err("Cannot merge different keyed additive families".into()); + } + let mut result = KeyedSumCountAccumulator::for_family(family); + + for acc in accumulators { + for key in acc.sums.keys() { + match ( + result.counts.get(key).copied(), + acc.counts.get(key).copied(), + ) { + (None, Some(count)) if !result.sums.contains_key(key) => { + result.counts.insert(key.clone(), count); + } + (Some(existing), Some(count)) => { + if let Some(total) = existing.checked_add(count) { + result.counts.insert(key.clone(), total); + } else { + result.counts.remove(key); + } + } + _ => { + result.counts.remove(key); + } + } + } + for (key, sum) in acc.sums { + *result.sums.entry(key).or_insert(0.0) += sum; + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use std::vec; + + use super::*; + + #[test] + fn test_keyed_sum_count_accumulator_creation() { + let acc = KeyedSumCountAccumulator::new(); + assert!(acc.sums.is_empty()); + } + + #[test] + fn test_keyed_sum_count_accumulator_update() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.update(key1.clone(), 10.0); + acc.update(key2.clone(), 20.0); + acc.update(key1.clone(), 5.0); // Should add to existing + + assert_eq!(acc.sums.get(&key1), Some(&15.0)); + assert_eq!(acc.sums.get(&key2), Some(&20.0)); + } + + #[test] + fn grouped_count_reads_sample_count_and_survives_merge_and_round_trip() { + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let mut first = KeyedSumCountAccumulator::for_family(ExactKind::Count); + first.update(key.clone(), 10.0); + first.update(key.clone(), 20.0); + let mut second = KeyedSumCountAccumulator::for_family(ExactKind::Count); + second.update(key.clone(), 7.0); + let merged = KeyedSumCountAccumulator::merge_accumulators(vec![first, second]).unwrap(); + for acc in [ + merged.clone(), + KeyedSumCountAccumulator::deserialize_from_json(&merged.serialize_to_json()).unwrap(), + KeyedSumCountAccumulator::deserialize_from_bytes(&merged.serialize_to_bytes()).unwrap(), + ] { + assert_eq!(acc.family, ExactKind::Count); + assert!(acc.query(Statistic::Sum, &key, None).is_err()); + assert_eq!(acc.query(Statistic::Count, &key, None).unwrap(), 3.0); + } + } + + #[test] + fn keyed_additive_merge_rejects_different_planner_families() { + assert!(KeyedSumCountAccumulator::merge_accumulators(vec![ + KeyedSumCountAccumulator::for_family(ExactKind::Sum), + KeyedSumCountAccumulator::for_family(ExactKind::Count), + ]) + .is_err()); + } + + #[test] + fn test_keyed_sum_count_accumulator_query() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); + + acc.add_sum(key.clone(), 42.0); + + // Test total queries (querying with the specific key) + assert_eq!( + crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Sum, &key, None).unwrap(), + 42.0 + ); + + // Test error cases + assert!( + crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Min, &key, None).is_err() + ); + } + + #[test] + fn test_keyed_sum_count_accumulator_get_keys() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.add_sum(key1.clone(), 10.0); + acc.add_sum(key2.clone(), 20.0); + + let keys = crate::AggregateCore::get_keys(&acc).unwrap(); + assert_eq!(keys.len(), 2); + assert!(keys.contains(&key1)); + assert!(keys.contains(&key2)); + } + + #[test] + fn test_keyed_sum_count_accumulator_merge() { + let mut acc1 = KeyedSumCountAccumulator::new(); + let mut acc2 = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc1.add_sum(key1.clone(), 10.0); + acc1.add_sum(key2.clone(), 20.0); + + acc2.add_sum(key1.clone(), 5.0); // Same key, different accumulator + + let merged = >::merge_accumulators(vec![acc1, acc2]).unwrap(); + + assert_eq!(merged.sums.get(&key1), Some(&15.0)); // Should be merged + assert_eq!(merged.sums.get(&key2), Some(&20.0)); // Should be preserved + } + + #[test] + fn test_keyed_sum_count_accumulator_serialization() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); + + acc.add_sum(key.clone(), 42.5); + + // Test JSON serialization + let json = acc.serialize_to_json(); + let deserialized = KeyedSumCountAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!(deserialized.sums.get(&key), Some(&42.5)); + + // Test byte serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = KeyedSumCountAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(deserialized_bytes.sums.get(&key), Some(&42.5)); + } + + #[test] + fn test_trait_object() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + acc.add_sum(key.clone(), 42.0); + + let trait_obj: Box = Box::new(acc); + + // Test type name through trait object + assert_eq!(trait_obj.type_name(), "KeyedSumCountAccumulator"); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/max.rs b/crates/asap_summary_state/src/summary_kernels/max.rs new file mode 100644 index 000000000..0c9547081 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/max.rs @@ -0,0 +1,248 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Exact maximum over one population, mergeable by comparison. +/// +/// See [`MinAccumulator`](super::min::MinAccumulator) for why the +/// two directions are separate types rather than one accumulator carrying a +/// `sub_type` string. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MaxAccumulator { + pub value: f64, +} + +impl Default for MaxAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl MaxAccumulator { + pub fn new() -> Self { + Self { + value: f64::NEG_INFINITY, + } + } + + pub fn with_value(value: f64) -> Self { + Self { value } + } + + pub fn update(&mut self, value: f64) { + if value > self.value { + self.value = value; + } + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let value = data["value"] + .as_f64() + .ok_or("Missing or invalid 'value' field")?; + Ok(Self::with_value(value)) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 8 { + return Err("Buffer too short".into()); + } + let value = f64::from_le_bytes([ + buffer[0], buffer[1], buffer[2], buffer[3], buffer[4], buffer[5], buffer[6], buffer[7], + ]); + Ok(Self::with_value(value)) + } +} + +impl SerializableToSink for MaxAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ "value": self.value }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.value.to_le_bytes().to_vec() + } +} + +impl MergeableAccumulator for MaxAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut result = MaxAccumulator::new(); + for acc in accumulators { + result.update(acc.value); + } + Ok(result) + } +} + +impl AggregateCore for MaxAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "MaxAccumulator" + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge MaxAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_max = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to MaxAccumulator")?; + let mut merged = self.clone(); + merged.update(other_max.value); + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Max + } + + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::() + } + + fn aux_stats(&self) -> AuxStats { + // The sentinel `f64::NEG_INFINITY` from `new()` is surfaced as-is; the + // query engine already treats it as "no data yet", the same way it + // does for `query_statistic`. + AuxStats { + max: Some(self.value), + ..AuxStats::empty() + } + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + _query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query(statistic, None) + } +} + +impl SingleSubpopulationAggregate for MaxAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + if query_kwargs.is_some() { + return Err("MaxAccumulator does not support query parameters".into()); + } + match statistic { + Statistic::Max => Ok(self.value), + other => Err(format!("Unsupported statistic in MaxAccumulator: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keeps_the_largest_update() { + let mut acc = MaxAccumulator::new(); + acc.update(10.0); + acc.update(5.0); + acc.update(15.0); + + assert_eq!(acc.value, 15.0); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Max, None).unwrap(), + 15.0 + ); + } + + #[test] + fn refuses_to_answer_a_minimum_query() { + let acc = MaxAccumulator::with_value(15.0); + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); + } + + #[test] + fn merges_by_taking_the_largest() { + let merged = + >::merge_accumulators(vec![ + MaxAccumulator::with_value(10.0), + MaxAccumulator::with_value(5.0), + MaxAccumulator::with_value(15.0), + ]) + .unwrap(); + assert_eq!(merged.value, 15.0); + } + + #[test] + fn refuses_to_merge_with_a_minimum() { + use super::super::min::MinAccumulator; + let max = MaxAccumulator::with_value(15.0); + let min = MinAccumulator::with_value(5.0); + assert!(max.merge_with(&min).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let acc = MaxAccumulator::with_value(42.5); + + let json = acc.serialize_to_json(); + assert_eq!( + MaxAccumulator::deserialize_from_json(&json).unwrap().value, + 42.5 + ); + + let bytes = acc.serialize_to_bytes(); + assert_eq!( + MaxAccumulator::deserialize_from_bytes(&bytes) + .unwrap() + .value, + 42.5 + ); + } + + #[test] + fn aux_stats_expose_max_only() { + let aux = MaxAccumulator::with_value(99.0).aux_stats(); + assert_eq!(aux.max, Some(99.0)); + assert_eq!(aux.min, None); + assert_eq!(aux.try_answer(Statistic::Max), Some(99.0)); + assert_eq!(aux.try_answer(Statistic::Min), None); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/min.rs b/crates/asap_summary_state/src/summary_kernels/min.rs new file mode 100644 index 000000000..ff2ad548a --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/min.rs @@ -0,0 +1,253 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Exact minimum over one population, mergeable by comparison. +/// +/// The sibling [`MaxAccumulator`](super::max::MaxAccumulator) is a +/// separate type on purpose: these two used to be one `MinMaxAccumulator` +/// whose direction lived in a `sub_type: String`, which meant every layer +/// above -- the wire `aggregationSubType`, the accumulator factory, the +/// summary catalog -- had to carry the direction alongside the family and +/// could silently answer a `min_over_time` read from maximum state. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MinAccumulator { + pub value: f64, +} + +impl Default for MinAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl MinAccumulator { + pub fn new() -> Self { + Self { + value: f64::INFINITY, + } + } + + pub fn with_value(value: f64) -> Self { + Self { value } + } + + pub fn update(&mut self, value: f64) { + if value < self.value { + self.value = value; + } + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let value = data["value"] + .as_f64() + .ok_or("Missing or invalid 'value' field")?; + Ok(Self::with_value(value)) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 8 { + return Err("Buffer too short".into()); + } + let value = f64::from_le_bytes([ + buffer[0], buffer[1], buffer[2], buffer[3], buffer[4], buffer[5], buffer[6], buffer[7], + ]); + Ok(Self::with_value(value)) + } +} + +impl SerializableToSink for MinAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ "value": self.value }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.value.to_le_bytes().to_vec() + } +} + +impl MergeableAccumulator for MinAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut result = MinAccumulator::new(); + for acc in accumulators { + result.update(acc.value); + } + Ok(result) + } +} + +impl AggregateCore for MinAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "MinAccumulator" + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge MinAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_min = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to MinAccumulator")?; + let mut merged = self.clone(); + merged.update(other_min.value); + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Min + } + + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::() + } + + fn aux_stats(&self) -> AuxStats { + // The sentinel `f64::INFINITY` from `new()` is surfaced as-is; the + // query engine already treats it as "no data yet", the same way it + // does for `query_statistic`. + AuxStats { + min: Some(self.value), + ..AuxStats::empty() + } + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + _query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query(statistic, None) + } +} + +impl SingleSubpopulationAggregate for MinAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + if query_kwargs.is_some() { + return Err("MinAccumulator does not support query parameters".into()); + } + match statistic { + Statistic::Min => Ok(self.value), + other => Err(format!("Unsupported statistic in MinAccumulator: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keeps_the_smallest_update() { + let mut acc = MinAccumulator::new(); + acc.update(10.0); + acc.update(5.0); + acc.update(15.0); + + assert_eq!(acc.value, 5.0); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).unwrap(), + 5.0 + ); + } + + #[test] + fn refuses_to_answer_a_maximum_query() { + let acc = MinAccumulator::with_value(5.0); + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Max, None).is_err()); + } + + #[test] + fn merges_by_taking_the_smallest() { + let merged = + >::merge_accumulators(vec![ + MinAccumulator::with_value(10.0), + MinAccumulator::with_value(5.0), + MinAccumulator::with_value(15.0), + ]) + .unwrap(); + assert_eq!(merged.value, 5.0); + } + + #[test] + fn refuses_to_merge_with_a_maximum() { + use super::super::max::MaxAccumulator; + let min = MinAccumulator::with_value(5.0); + let max = MaxAccumulator::with_value(15.0); + assert!(min.merge_with(&max).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let acc = MinAccumulator::with_value(42.5); + + let json = acc.serialize_to_json(); + assert_eq!( + MinAccumulator::deserialize_from_json(&json).unwrap().value, + 42.5 + ); + + let bytes = acc.serialize_to_bytes(); + assert_eq!( + MinAccumulator::deserialize_from_bytes(&bytes) + .unwrap() + .value, + 42.5 + ); + } + + #[test] + fn aux_stats_expose_min_only() { + let aux = MinAccumulator::with_value(3.5).aux_stats(); + assert_eq!(aux.min, Some(3.5)); + assert_eq!(aux.max, None); + assert_eq!(aux.count, None); + assert_eq!(aux.sum, None); + assert_eq!(aux.try_answer(Statistic::Min), Some(3.5)); + assert_eq!(aux.try_answer(Statistic::Max), None); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/mod.rs b/crates/asap_summary_state/src/summary_kernels/mod.rs new file mode 100644 index 000000000..8ebc5d55d --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/mod.rs @@ -0,0 +1,42 @@ +//! ASAP state adapters and exact accumulators; sketch algorithms live in `asap_sketchlib`. +pub mod count_min_sketch; +pub mod count_min_sketch_with_heap; +pub mod count_sketch; +pub mod count_sketch_with_heap; +pub mod datasketches_kll; +pub mod dd_sketch; +pub mod exact; +pub mod hll_sketch; +pub mod hydra_kll; +pub mod increase; +pub mod keyed_counter_state; +pub mod keyed_max_state; +pub mod keyed_min_state; +pub mod keyed_sum_count; +pub mod max; +pub mod min; +pub mod sketch_envelope; +pub mod sum; +pub mod univmon; + +pub use count_min_sketch::*; +pub use count_min_sketch_with_heap::*; +pub use count_sketch::*; +pub use count_sketch_with_heap::*; +pub use datasketches_kll::*; +pub use dd_sketch::*; +pub use hll_sketch::*; +pub use hydra_kll::*; +pub use increase::*; +pub use keyed_counter_state::*; +pub use keyed_max_state::*; +pub use keyed_min_state::*; +pub use keyed_sum_count::*; +pub use max::*; +pub use min::*; +pub use sketch_envelope::*; +pub use sum::*; + +pub mod factory; +pub mod traits; +pub mod weighted_frequency; diff --git a/crates/asap_summary_state/src/summary_kernels/sketch_envelope.rs b/crates/asap_summary_state/src/summary_kernels/sketch_envelope.rs new file mode 100644 index 000000000..475d7f51d --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/sketch_envelope.rs @@ -0,0 +1,154 @@ +//! SketchEnvelopeAccumulator — wraps a raw SketchEnvelope protobuf payload +//! received via OTLP ingest so it can be stored through the `Store` trait. +//! +//! The accumulator preserves the opaque proto bytes and decodes them lazily +//! (via `SketchEnvelope::decode`) only when merge or query operations need +//! the inner sketch type. + +use crate::{AggregateCore, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; +use prost::Message; +use serde_json::Value; +use std::collections::HashMap; + +use crate::AggregationType; +use crate::Statistic; + +/// Accumulator that stores a serialized `SketchEnvelope` protobuf. +/// +/// This is the simplest viable path for OTLP sketch ingest: the OTel Collector +/// has already computed the sketch, so the backend just stores the bytes and +/// serves them back at query time. +#[derive(Debug, Clone)] +pub struct SketchEnvelopeAccumulator { + /// Raw protobuf-encoded `SketchEnvelope`. + pub payload: Vec, + /// Sketch type string cached from decoding (e.g. "CountMin", "KLL"). + pub sketch_type: String, +} + +impl SketchEnvelopeAccumulator { + /// Create from raw protobuf bytes. Decodes the envelope once to cache + /// the sketch type; the full payload is kept for later use. + pub fn from_proto_bytes( + payload: Vec, + ) -> Result> { + let sketch_type = match SketchEnvelope::decode(payload.as_slice()) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::CountMin(_)) => "CountMin".to_string(), + Some(sketch_envelope::SketchState::CountSketch(_)) => "CountSketch".to_string(), + Some(sketch_envelope::SketchState::Kll(_)) => "KLL".to_string(), + Some(sketch_envelope::SketchState::Hll(_)) => "HLL".to_string(), + Some(sketch_envelope::SketchState::Ddsketch(_)) => "DDSketch".to_string(), + Some(sketch_envelope::SketchState::Univmon(_)) => "UnivMon".to_string(), + Some(sketch_envelope::SketchState::Hydra(_)) => "Hydra".to_string(), + Some(sketch_envelope::SketchState::Coco(_)) => "CocoSketch".to_string(), + Some(sketch_envelope::SketchState::Elastic(_)) => "Elastic".to_string(), + None => "Unknown".to_string(), + }, + Err(e) => { + return Err(format!("Failed to decode SketchEnvelope: {}", e).into()); + } + }; + + Ok(Self { + payload, + sketch_type, + }) + } +} + +// --------------------------------------------------------------------------- +// Trait implementations +// --------------------------------------------------------------------------- + +impl SerializableToSink for SketchEnvelopeAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "type": "SketchEnvelopeAccumulator", + "sketch_type": self.sketch_type, + "payload_bytes": self.payload.len(), + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.payload.clone() + } +} + +impl AggregateCore for SketchEnvelopeAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "SketchEnvelopeAccumulator" + } + + 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> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge SketchEnvelopeAccumulator with {:?}", + other.get_accumulator_type() + ) + .into()); + } + + // For now, merging opaque envelopes is not supported — each window is + // a self-contained sketch produced by the OTel Collector. Return self + // as-is so the store can still call merge_with without panicking. + Ok(Box::new(self.clone())) + } + + fn get_accumulator_type(&self) -> AggregationType { + // Opaque wrapper — report as the generic multi-subpopulation bucket. + // Direct dispatch is not supported; native sketch query path must + // decode the envelope and delegate to the correct accumulator. + AggregationType::MultipleSubpopulation + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + _statistic: Statistic, + _key: &Option, + _query_kwargs: &HashMap, + ) -> Result> { + Err( + "SketchEnvelopeAccumulator: query_statistic not supported; decode envelope first" + .into(), + ) + } +} + +impl crate::MultipleSubpopulationAggregate for SketchEnvelopeAccumulator { + fn query( + &self, + _statistic: Statistic, + _key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Err( + "SketchEnvelopeAccumulator: direct query not supported; use native sketch query path" + .into(), + ) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/sum.rs b/crates/asap_summary_state/src/summary_kernels/sum.rs new file mode 100644 index 000000000..4e74a45ae --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/sum.rs @@ -0,0 +1,413 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SumAccumulator { + pub sum: f64, + /// None for scalar-only payloads; a sum does not establish a sample count. + #[serde(default)] + pub observation_count: Option, +} + +impl SumAccumulator { + pub fn new() -> Self { + Self { + sum: 0.0, + observation_count: Some(0), + } + } + + pub fn with_sum(sum: f64) -> Self { + Self { + sum, + observation_count: None, + } + } + + pub fn update(&mut self, value: f64) { + self.sum += value; + self.observation_count = self + .observation_count + .and_then(|count| count.checked_add(1)); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let sum = data["sum"] + .as_f64() + .ok_or("Missing or invalid 'sum' field")?; + Ok(Self { + sum, + observation_count: data.get("observation_count").and_then(Value::as_u64), + }) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + match buffer.len() { + // Legacy Python scalar sums carry no sample-count evidence. + 4 => Ok(Self::with_sum(f32::from_le_bytes(buffer.try_into()?) as f64)), + // Counted sums use the same fixed layout as the Collector Sum payload. + 16 => Self::from_sum_bytes(buffer), + len => { + Err(format!("Invalid persisted Sum payload length: {len} (want 4 or 16)").into()) + } + } + } + + /// Decode the fixed Sum payload produced by the first-class Sum + /// AggregationType path (asap-precompute-go's SumWrapper): float64 sum + /// (little-endian) followed by uint64 count (little-endian), 16 bytes. + /// + /// Sum is an aggregation, NOT a sketch, so this deliberately does NOT + /// depend on the sketchlib sketch-envelope proto — the payload is a small + /// self-contained fixed layout. It decodes into the SAME + /// `AggregationType::Sum` accumulator as a plain-OTLP Sum, so the SumAgg + /// envelope and a plain Sum land on one identity (`exact_agg:Sum`) with no + /// new SketchAlgorithm. The supplied observation count is retained for + /// exact sample-count readouts; scalar-only legacy payloads leave it unknown. + pub fn from_sum_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 16 { + return Err(format!("Sum payload too short: {} bytes (want 16)", buffer.len()).into()); + } + let sum = f64::from_le_bytes(buffer[0..8].try_into().unwrap()); + let count = u64::from_le_bytes(buffer[8..16].try_into().unwrap()); + Ok(Self { + sum, + observation_count: Some(count), + }) + } +} + +impl Default for SumAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl SerializableToSink for SumAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "sum": self.sum, + "observation_count": self.observation_count + }) + } + + fn serialize_to_bytes(&self) -> Vec { + match self.observation_count { + Some(count) => { + let mut bytes = Vec::with_capacity(16); + bytes.extend_from_slice(&self.sum.to_le_bytes()); + bytes.extend_from_slice(&count.to_le_bytes()); + bytes + } + None => (self.sum as f32).to_le_bytes().to_vec(), + } + } +} + +impl AggregateCore for SumAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "SumAccumulator" + } + + 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> { + // Check if other is also a SumAccumulator + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge SumAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + // Downcast to SumAccumulator + let other_sum = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to SumAccumulator")?; + + // Use the existing merge_accumulators method + let merged = Self::merge_accumulators(vec![self.clone(), other_sum.clone()])?; + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Sum + } + + fn approx_memory_bytes(&self) -> usize { + // Single f64 + struct overhead. + std::mem::size_of::() + } + + fn aux_stats(&self) -> AuxStats { + AuxStats { + sum: Some(self.sum), + count: self.observation_count, + ..AuxStats::empty() + } + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + _query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query(statistic, None) + } +} + +impl SingleSubpopulationAggregate for SumAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + // SumAccumulator doesn't use query_kwargs, assert it's None + if query_kwargs.is_some() { + return Err("SumAccumulator does not support query parameters".into()); + } + + match statistic { + Statistic::Sum => Ok(self.sum), + Statistic::Count => self + .observation_count + .map(|count| count as f64) + .ok_or_else(|| "sample count is unavailable for this Sum payload".into()), + _ => Err(format!("Unsupported statistic in SumAccumulator: {statistic:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for SumAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + let total_sum = accumulators.iter().map(|acc| acc.sum).sum(); + let observation_count = accumulators + .iter() + .try_fold(0u64, |total, acc| total.checked_add(acc.observation_count?)); + Ok(SumAccumulator { + sum: total_sum, + observation_count, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Sample counts must survive updates and merges independently of the sum. + #[test] + fn observation_count_survives_merge() { + let mut first = SumAccumulator::new(); + first.update(10.0); + first.update(20.0); + let mut second = SumAccumulator::new(); + second.update(100.0); + let merged = SumAccumulator::merge_accumulators(vec![first, second]).unwrap(); + assert_eq!(merged.sum, 130.0); + assert_eq!(merged.aux_stats().count, Some(3)); + } + + // A legacy scalar sum has no evidence of how many observations produced it. + #[test] + fn legacy_sum_does_not_invent_observation_count() { + let mut raw = SumAccumulator::new(); + raw.update(10.0); + let merged = + SumAccumulator::merge_accumulators(vec![raw, SumAccumulator::with_sum(20.0)]).unwrap(); + assert_eq!(merged.aux_stats().count, None); + } + + // Persistence retains known counts, including zero and the full u64 range. + #[test] + fn counted_sum_binary_round_trip() { + for count in [0, 3, u64::MAX] { + let acc = SumAccumulator { + sum: 1.0000000000001, + observation_count: Some(count), + }; + let bytes = acc.serialize_to_bytes(); + assert_eq!(bytes.len(), 16); + let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(restored.sum, acc.sum); + assert_eq!(restored.observation_count, Some(count)); + } + } + + // Existing scalar-only files remain readable without inventing counts. + #[test] + fn legacy_binary_sum_has_unknown_count() { + let bytes = 42.5f32.to_le_bytes(); + let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(restored.sum, 42.5); + assert_eq!(restored.observation_count, None); + assert_eq!(restored.serialize_to_bytes(), bytes); + } + + // Truncated counted payloads must not silently decode as scalar sums. + #[test] + fn persisted_sum_rejects_invalid_lengths() { + for len in [0, 3, 5, 8, 15, 17] { + assert!(SumAccumulator::deserialize_from_bytes(&vec![0; len]).is_err()); + } + } + + #[test] + fn test_sum_accumulator_creation() { + let acc = SumAccumulator::new(); + assert_eq!(acc.sum, 0.0); + + let acc2 = SumAccumulator::with_sum(42.5); + assert_eq!(acc2.sum, 42.5); + } + + #[test] + fn test_sum_accumulator_update() { + let mut acc = SumAccumulator::new(); + acc.update(10.0); + acc.update(20.0); + assert_eq!(acc.sum, 30.0); + } + + #[test] + fn test_sum_accumulator_query() { + let acc = SumAccumulator::with_sum(42.0); + + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), + 42.0 + ); + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Count, None).is_err()); + + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); + // SumAccumulator is a single subpopulation accumulator, doesn't need key-based queries + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), + 42.0 + ); + } + + #[test] + fn count_readout_uses_observation_count_not_sum() { + let mut acc = SumAccumulator::new(); + acc.update(10.0); + acc.update(20.0); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Count, None).unwrap(), + 2.0 + ); + } + + #[test] + fn test_sum_accumulator_merge() { + let acc1 = SumAccumulator::with_sum(10.0); + let acc2 = SumAccumulator::with_sum(20.0); + let acc3 = SumAccumulator::with_sum(30.0); + + let merged = + >::merge_accumulators(vec![ + acc1, acc2, acc3, + ]) + .unwrap(); + assert_eq!(merged.sum, 60.0); + } + + #[test] + fn test_sum_accumulator_serialization() { + let acc = SumAccumulator::with_sum(42.5); + + // Test JSON serialization + let json = acc.serialize_to_json(); + let deserialized = SumAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!(acc.sum, deserialized.sum); + + // Test byte serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(acc.sum, deserialized_bytes.sum); + } + + #[test] + fn test_trait_object() { + let acc: Box = Box::new(SumAccumulator::with_sum(42.0)); + + assert_eq!(acc.type_name(), "SumAccumulator"); + } + + #[test] + fn from_sum_bytes_decodes_go_sum_payload() { + // GOLDEN: the 16-byte payload asap-precompute-go's + // SumWrapper{10,20,30,40}.Snapshot() emits — float64 sum (LE) followed + // by uint64 count (LE), sum=100, count=4. Proves the Rust backend + // decodes the first-class Sum payload the Go agent produces + // (cross-language wire parity, no sketchlib proto dependency). + let go_bytes: &[u8] = &[ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x59, 0x40, // 100.0 f64 LE + 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // 4 u64 LE + ]; + let acc = SumAccumulator::from_sum_bytes(go_bytes).expect("decode Go Sum payload"); + assert_eq!(acc.sum, 100.0, "decoded Go SumWrapper payload sum"); + } + + #[test] + fn from_sum_bytes_rejects_short_payload() { + // A short buffer is rejected (the ingest path then skips the point). + assert!(SumAccumulator::from_sum_bytes(&[]).is_err()); + assert!(SumAccumulator::from_sum_bytes(&[0u8; 8]).is_err()); + } + + #[test] + fn aux_stats_exposes_sum_only() { + let acc = SumAccumulator::with_sum(123.5); + let aux = acc.aux_stats(); + assert_eq!(aux.sum, Some(123.5)); + assert_eq!(aux.count, None); + assert_eq!(aux.min, None); + assert_eq!(aux.max, None); + } + + #[test] + fn aux_stats_try_answer_on_sum_statistic() { + use crate::Statistic; + let acc = SumAccumulator::with_sum(42.0); + // Sum statistic is covered by aux without deserialising. + assert_eq!(acc.aux_stats().try_answer(Statistic::Sum), Some(42.0)); + // Count is not tracked by SumAccumulator. + assert_eq!(acc.aux_stats().try_answer(Statistic::Count), None); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/traits.rs b/crates/asap_summary_state/src/summary_kernels/traits.rs new file mode 100644 index 000000000..ae05ca06b --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/traits.rs @@ -0,0 +1,357 @@ +use crate::KeyByLabelValues; +use std::collections::HashMap; + +use crate::AggregationType; +use crate::Statistic; + +use serde_json::Value; + +/// Trait for objects that can be serialized to different formats +pub trait SerializableToSink { + fn serialize_to_json(&self) -> Value; + fn serialize_to_bytes(&self) -> Vec; +} + +/// Core trait for all aggregates containing shared functionality +/// This trait provides common operations like serialization, cloning, and type identification +pub trait AggregateCore: SerializableToSink + Send + Sync { + /// Clone this accumulator into a boxed trait object + fn clone_boxed_core(&self) -> Box; + + /// Get the type name of this accumulator + fn type_name(&self) -> &'static str; + + /// Downcast to Any for type checking + fn as_any(&self) -> &dyn std::any::Any; + + /// Mutable downcast to Any. Used by ingest paths that need to + /// mutate a boxed accumulator in place — e.g. the PROTO_DELTA + /// delta-merge applier in `drivers::ingest::otel::apply_modified_otlp_delta_bytes`. + fn as_any_mut(&mut self) -> &mut dyn std::any::Any; + + /// Merge this accumulator with another accumulator of the same type + /// Returns a new merged accumulator, leaving the original unchanged + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box>; + + /// Get the accumulator type identifier for merge compatibility checking + fn get_accumulator_type(&self) -> AggregationType; + + /// Get all keys stored in this accumulator + fn get_keys(&self) -> Option>; + + /// Dispatch a statistic query without downcasting. + /// + /// Replaces the 12-arm `match get_accumulator_type()` in the engine. + /// Single-subpopulation types ignore `key`; multiple-subpopulation types + /// require it and return `Err` when it is `None`. + /// Special cases (DeltaSetAggregator, SetAggregator) fall back to a + /// cardinality value when `key` is `None`. + fn query_statistic( + &self, + statistic: Statistic, + key: &Option, + query_kwargs: &HashMap, + ) -> Result>; + + /// Approximate in-memory byte footprint of this accumulator. + /// + /// Used by the `SketchStore` persistence layer to drive its + /// memory-pressure trigger. Not required to be exact — the flusher + /// only needs rough proportionality. The default is a conservative + /// 4 KiB constant; concrete types should override it with a + /// type-aware estimate (e.g. KLL: `k * 8` plus overhead). + /// + /// Implementors must not call `serialize_to_bytes` here — this is + /// on the insert hot path. + fn approx_memory_bytes(&self) -> usize { + 4096 + } + + /// Typed auxiliary statistics — `count`, `sum`, `min`, `max` — + /// exposed as first-class scalars alongside the sketch payload. + /// + /// The overwhelming majority of production queries + /// (`count_over_time`, `sum_over_time`, `min_over_time`, + /// `max_over_time`, and the additive aggregations built on + /// them) only need these scalars. Returning them directly here + /// lets callers avoid deserialising the full sketch bytes. + /// + /// Returning fields as `None` means the accumulator doesn't + /// track that statistic exactly (e.g. a pure HLL doesn't carry + /// sum/min/max). Callers then fall back to the sketch's + /// `query_statistic` method. + /// + /// This is the phase-1 piece of the sketch DB design + /// (docs/design_docs/summary-storage.md). + fn aux_stats(&self) -> AuxStats { + AuxStats::empty() + } + + /// Reset the sketch state to empty **in place**, preserving its + /// shape / configuration (dimensions, relative accuracy, register + /// width, …) so a subsequent delta-apply lands on a clean, + /// same-shape base. + /// + /// Used by the OTLP ingest path's per-window base rotation: when a + /// delta frame opens a new tumbling window for a series, the cached + /// base is reset here before the new window's delta is applied, so + /// the reconstructed state reflects that window only rather than an + /// all-time accumulation across windows (see + /// `docs/delta-baseline-contract.md` §3). + /// + /// The default is a no-op: only the delta-capable, additive families + /// (DDSketch, CMS, CountSketch, HLL) ever reach the rotation path and + /// override this. KLL never deltas, and the non-sketch accumulators + /// are never cached as a delta base. + fn reset_to_empty(&mut self) {} +} + +/// Four typed auxiliary scalars tracked alongside every sketch entry: +/// `count`, `sum`, `min`, `max`. Exposed so the query engine can +/// serve Count / Sum / Min / Max statistics without touching sketch +/// bytes. +/// +/// Each field is `Option<…>` because not every accumulator tracks +/// every stat (e.g. HLL has cardinality but no meaningful +/// sum / min / max; DeltaSetAggregator tracks set transitions, not +/// numeric aggregates). +#[derive(Debug, Default, Clone, Copy, PartialEq)] +pub struct AuxStats { + pub count: Option, + pub sum: Option, + pub min: Option, + pub max: Option, +} + +impl AuxStats { + pub const fn empty() -> Self { + Self { + count: None, + sum: None, + min: None, + max: None, + } + } + + /// Attempt to fulfil a `Statistic` purely from the typed aux + /// columns, without needing to deserialise the sketch. Returns + /// `None` if the requested statistic isn't covered by aux + /// (e.g. Quantile, Cardinality, TopK) or if the corresponding + /// aux field is `None`. + pub fn try_answer(&self, statistic: Statistic) -> Option { + match statistic { + Statistic::Count => self.count.map(|c| c as f64), + Statistic::Sum => self.sum, + Statistic::Min => self.min, + Statistic::Max => self.max, + // Increase / Rate need two samples; aux columns carry + // window totals, so one entry's aux is insufficient. + // Cardinality / Quantile / Topk are sketch-native and + // must go through query_statistic. + _ => None, + } + } + + /// Merge two aux stats the way the corresponding sketch merge + /// would. Count / sum add, min / max take the extremum. When + /// either side is `None` the result is the other side (so a + /// window that only has partial aux still contributes). + pub fn merge(self, other: Self) -> Self { + fn add_opt_u(a: Option, b: Option) -> Option { + match (a, b) { + (Some(x), Some(y)) => Some(x.saturating_add(y)), + (x, None) => x, + (None, y) => y, + } + } + fn add_opt_f(a: Option, b: Option) -> Option { + match (a, b) { + (Some(x), Some(y)) => Some(x + y), + (x, None) => x, + (None, y) => y, + } + } + fn min_opt(a: Option, b: Option) -> Option { + match (a, b) { + (Some(x), Some(y)) => Some(x.min(y)), + (x, None) => x, + (None, y) => y, + } + } + fn max_opt(a: Option, b: Option) -> Option { + match (a, b) { + (Some(x), Some(y)) => Some(x.max(y)), + (x, None) => x, + (None, y) => y, + } + } + Self { + count: add_opt_u(self.count, other.count), + sum: add_opt_f(self.sum, other.sum), + min: min_opt(self.min, other.min), + max: max_opt(self.max, other.max), + } + } +} + +/// Trait for accumulators that support a single subpopulation +/// These accumulators store a single aggregate value (e.g., Sum, Increase) +pub trait SingleSubpopulationAggregate: AggregateCore { + /// Query the accumulator for a specific statistic + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result>; + + /// Clone this accumulator into a boxed trait object + fn clone_boxed(&self) -> Box; +} + +/// Trait for accumulators that support multiple subpopulations identified by keys +/// These accumulators store separate values for different label combinations +pub trait MultipleSubpopulationAggregate: AggregateCore { + /// Query the accumulator for a specific statistic and key + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + query_kwargs: Option<&HashMap>, + ) -> Result>; + + /// Clone this accumulator into a boxed trait object + fn clone_boxed(&self) -> Box; +} + +/// Trait for merging multiple accumulators of the same type +pub trait MergeableAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> + where + T: Sized; +} + +// Implement Clone for the new trait objects +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_boxed_core() + } +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_boxed() + } +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_boxed() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn aux_stats_empty_answers_nothing() { + let e = AuxStats::empty(); + assert_eq!(e.try_answer(Statistic::Count), None); + assert_eq!(e.try_answer(Statistic::Sum), None); + assert_eq!(e.try_answer(Statistic::Min), None); + assert_eq!(e.try_answer(Statistic::Max), None); + } + + #[test] + fn aux_stats_try_answer_covers_typed_stats() { + let a = AuxStats { + count: Some(7), + sum: Some(42.0), + min: Some(1.5), + max: Some(9.25), + }; + assert_eq!(a.try_answer(Statistic::Count), Some(7.0)); + assert_eq!(a.try_answer(Statistic::Sum), Some(42.0)); + assert_eq!(a.try_answer(Statistic::Min), Some(1.5)); + assert_eq!(a.try_answer(Statistic::Max), Some(9.25)); + } + + #[test] + fn aux_stats_try_answer_skips_sketch_native_stats() { + let a = AuxStats { + count: Some(100), + sum: Some(500.0), + min: Some(1.0), + max: Some(10.0), + }; + assert_eq!(a.try_answer(Statistic::Quantile), None); + assert_eq!(a.try_answer(Statistic::Cardinality), None); + assert_eq!(a.try_answer(Statistic::Topk), None); + assert_eq!(a.try_answer(Statistic::Increase), None); + assert_eq!(a.try_answer(Statistic::Rate), None); + } + + #[test] + fn aux_stats_merge_adds_count_and_sum_takes_extrema() { + let a = AuxStats { + count: Some(10), + sum: Some(50.0), + min: Some(1.0), + max: Some(9.0), + }; + let b = AuxStats { + count: Some(5), + sum: Some(20.0), + min: Some(0.5), + max: Some(12.0), + }; + let merged = a.merge(b); + assert_eq!(merged.count, Some(15)); + assert_eq!(merged.sum, Some(70.0)); + assert_eq!(merged.min, Some(0.5)); + assert_eq!(merged.max, Some(12.0)); + } + + #[test] + fn aux_stats_merge_handles_partial_sides() { + // HLL-like (count only) merged with Sum-only side. + let hll_like = AuxStats { + count: Some(100), + ..AuxStats::empty() + }; + let sum_like = AuxStats { + sum: Some(500.0), + ..AuxStats::empty() + }; + let merged = hll_like.merge(sum_like); + assert_eq!(merged.count, Some(100)); + assert_eq!(merged.sum, Some(500.0)); + assert_eq!(merged.min, None); + assert_eq!(merged.max, None); + } + + #[test] + fn aux_stats_merge_is_empty_plus_empty() { + let merged = AuxStats::empty().merge(AuxStats::empty()); + assert_eq!(merged, AuxStats::empty()); + } + + #[test] + fn aux_stats_count_saturates_on_overflow() { + let a = AuxStats { + count: Some(u64::MAX - 1), + ..AuxStats::empty() + }; + let b = AuxStats { + count: Some(100), + ..AuxStats::empty() + }; + let merged = a.merge(b); + assert_eq!(merged.count, Some(u64::MAX)); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/univmon.rs b/crates/asap_summary_state/src/summary_kernels/univmon.rs new file mode 100644 index 000000000..4a52bf1da --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/univmon.rs @@ -0,0 +1,234 @@ +//! One frequency state shared by count, distinct, L2 and entropy readouts. + +use crate::{AggregateCore, AuxStats, KeyByLabelValues, SerializableToSink}; +use crate::{AggregationType, Statistic}; +use asap_sketchlib::{DataInput, UnivMon}; +use serde_json::Value; +use std::collections::HashMap; + +type Error = Box; + +#[derive(Debug, Clone)] +pub struct UnivMonAccumulator { + inner: UnivMon, +} + +impl UnivMonAccumulator { + pub fn new(heap_size: usize, rows: usize, cols: usize, layers: usize) -> Result { + if heap_size == 0 || cols == 0 || !(1..=20).contains(&rows) || !(1..=64).contains(&layers) { + return Err("invalid UnivMon dimensions".into()); + } + rows.checked_mul(cols) + .and_then(|n| n.checked_mul(layers)) + .ok_or("UnivMon dimensions overflow")?; + Ok(Self { + inner: UnivMon::init_univmon(heap_size, rows, cols, layers), + }) + } + + /// Each non-NaN sample is one occurrence. Signed zero has one identity. + pub fn insert_sample(&mut self, value: f64) -> Result<(), Error> { + if value.is_nan() { + return Ok(()); + } + self.inner + .bucket_size + .checked_add(1) + .ok_or("UnivMon count overflow")?; + let bits = if value == 0.0 { 0 } else { value.to_bits() }; + self.inner.insert(&DataInput::U64(bits), 1); + Ok(()) + } + + pub fn from_bytes(bytes: &[u8]) -> Result { + let inner = UnivMon::deserialize_from_bytes(bytes) + .map_err(|e| format!("invalid UnivMon state: {e}"))?; + if !inner.accepts_standard_updates() { + return Err( + "terminal-mode UnivMon state cannot enter the standard-update accumulator".into(), + ); + } + Ok(Self { inner }) + } + + fn compatible(&self, other: &Self) -> bool { + ( + self.inner.heap_size, + self.inner.sketch_row, + self.inner.sketch_col, + self.inner.layer_size, + ) == ( + other.inner.heap_size, + other.inner.sketch_row, + other.inner.sketch_col, + other.inner.layer_size, + ) + } + + pub fn dimensions(&self) -> (usize, usize, usize, usize) { + ( + self.inner.heap_size, + self.inner.sketch_row, + self.inner.sketch_col, + self.inner.layer_size, + ) + } + + pub fn merge_in_place(&mut self, other: &Self) -> Result<(), Error> { + if !self.compatible(other) { + return Err("incompatible UnivMon dimensions".into()); + } + self.inner + .bucket_size + .checked_add(other.inner.bucket_size) + .ok_or("UnivMon count overflow")?; + self.inner.merge(&other.inner); + Ok(()) + } +} + +impl SerializableToSink for UnivMonAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({"count": self.inner.bucket_size}) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner + .serialize_to_bytes() + .expect("validated unit-frequency UnivMon state") + } +} + +impl AggregateCore for UnivMonAccumulator { + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::().saturating_add( + self.inner.layer_size.saturating_mul( + self.inner + .sketch_row + .saturating_mul(self.inner.sketch_col) + .saturating_mul(16) + .saturating_add(self.inner.heap_size.saturating_mul(256)), + ), + ) + } + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + fn type_name(&self) -> &'static str { + "UnivMonAccumulator" + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::UnivMon + } + fn get_keys(&self) -> Option> { + None + } + fn reset_to_empty(&mut self) { + self.inner.free(); + } + + fn merge_with(&self, other: &dyn AggregateCore) -> Result, Error> { + let other = other + .as_any() + .downcast_ref::() + .ok_or("expected UnivMon state")?; + let mut merged = self.clone(); + merged.merge_in_place(other)?; + Ok(Box::new(merged)) + } + + fn query_statistic( + &self, + statistic: Statistic, + key: &Option, + _: &HashMap, + ) -> Result { + if key.is_some() { + return Err("UnivMon population is selected by the catalog binding".into()); + } + match statistic { + Statistic::Count => Ok(self.inner.calc_l1()), + Statistic::Cardinality => Ok(self.inner.calc_card()), + Statistic::FrequencyL2 => Ok(self.inner.calc_l2()), + Statistic::FrequencyEntropy => Ok(self.inner.calc_entropy()), + _ => Err("unsupported UnivMon readout".into()), + } + } + + fn aux_stats(&self) -> AuxStats { + AuxStats { + count: Some(self.inner.bucket_size as u64), + ..AuxStats::empty() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn read(state: &dyn AggregateCore, stat: Statistic) -> f64 { + state.query_statistic(stat, &None, &HashMap::new()).unwrap() + } + + /// Duplicate samples affect frequency but not cardinality, including signed zero. + #[test] + fn shared_readouts_survive_serialization() { + let mut state = UnivMonAccumulator::new(32, 5, 1024, 4).unwrap(); + for value in [0.0, -0.0, 2.0, 2.0, f64::NAN] { + state.insert_sample(value).unwrap(); + } + let restored = UnivMonAccumulator::from_bytes(&state.serialize_to_bytes()).unwrap(); + for stat in [ + Statistic::Count, + Statistic::Cardinality, + Statistic::FrequencyL2, + Statistic::FrequencyEntropy, + ] { + assert_eq!(read(&state, stat), read(&restored, stat)); + } + assert_eq!(read(&restored, Statistic::Count), 4.0); + assert!((read(&restored, Statistic::Cardinality) - 2.0).abs() < 0.01); + assert!((read(&restored, Statistic::FrequencyL2) - 8.0f64.sqrt()).abs() < 0.01); + assert!((read(&restored, Statistic::FrequencyEntropy) - 1.0).abs() < 0.01); + } + + /// Terminal-mode serialization is valid sketchlib state but not this accumulator's update domain. + #[test] + fn terminal_state_is_rejected_before_ingestion_or_merge() { + let mut state = UnivMon::init_univmon(4, 3, 16, 2); + state.fast_insert(&DataInput::U64(1), 1); + let bytes = state.serialize_to_bytes().unwrap(); + assert!(UnivMonAccumulator::from_bytes(&bytes).is_err()); + state.free(); + assert!(UnivMonAccumulator::from_bytes(&state.serialize_to_bytes().unwrap()).is_ok()); + } + + /// Pane merge preserves overlapping keys and reset removes the previous window. + #[test] + fn merge_and_reset_preserve_frequency_semantics() { + let mut left = UnivMonAccumulator::new(32, 5, 1024, 4).unwrap(); + let mut right = left.clone(); + for value in [1.0, 2.0] { + left.insert_sample(value).unwrap(); + } + for value in [2.0, 3.0] { + right.insert_sample(value).unwrap(); + } + let merged = left.merge_with(&right).unwrap(); + assert_eq!(read(merged.as_ref(), Statistic::Count), 4.0); + assert!((read(merged.as_ref(), Statistic::Cardinality) - 3.0).abs() < 0.01); + left.reset_to_empty(); + assert_eq!(read(&left, Statistic::Count), 0.0); + assert_eq!(read(&left, Statistic::FrequencyEntropy), 0.0); + assert!(left + .merge_with(&UnivMonAccumulator::new(16, 5, 1024, 4).unwrap()) + .is_err()); + } +} diff --git a/crates/asap_summary_state/src/summary_kernels/weighted_frequency.rs b/crates/asap_summary_state/src/summary_kernels/weighted_frequency.rs new file mode 100644 index 000000000..34e142d99 --- /dev/null +++ b/crates/asap_summary_state/src/summary_kernels/weighted_frequency.rs @@ -0,0 +1,75 @@ +//! Stored form of Planner's weighted frequency state. It is read only through +//! typed heap rows, never through a scalar statistic. +use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink, Statistic}; +pub use asap_physical_operators::summary_kernels::weighted_frequency::{ + FrequencyAlgorithm, WeightedFrequency as PhysicalWeightedFrequency, +}; +use asap_physical_operators::AggregateCore as PhysicalState; +use std::collections::HashMap; + +#[derive(Debug, Clone)] +pub struct WeightedFrequency(pub PhysicalWeightedFrequency); + +impl SerializableToSink for WeightedFrequency { + fn serialize_to_json(&self) -> serde_json::Value { + serde_json::to_value(&self.0).expect("finite validated frequency state") + } + fn serialize_to_bytes(&self) -> Vec { + crate::physical::frequency_kernel(&self.0) + .expect("weighted frequency kernel encoding") + .to_bytes() + } +} + +impl AggregateCore for WeightedFrequency { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + fn type_name(&self) -> &'static str { + "WeightedFrequency" + } + 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> { + let other = other + .as_any() + .downcast_ref::() + .ok_or("weighted frequency state type mismatch")?; + let merged = self.0.merge_with(&other.0)?; + let merged = merged + .as_any() + .downcast_ref::() + .ok_or("weighted frequency merge changed state type")?; + Ok(Box::new(Self(merged.clone()))) + } + fn get_accumulator_type(&self) -> AggregationType { + match crate::physical::frequency_kernel(&self.0) + .expect("weighted frequency kernel encoding") + .algorithm() + { + FrequencyAlgorithm::Cms => AggregationType::CountMinSketchWithHeap, + FrequencyAlgorithm::CountSketch => AggregationType::CountSketchWithHeap, + } + } + fn get_keys(&self) -> Option> { + None + } + fn query_statistic( + &self, + _: Statistic, + _: &Option, + _: &HashMap, + ) -> Result> { + Err("weighted frequency uses typed row readout".into()) + } + fn approx_memory_bytes(&self) -> usize { + self.0.approx_memory_bytes() + } +} diff --git a/crates/asap_types/Cargo.toml b/crates/asap_types/Cargo.toml index d5c83ef67..4785a774b 100644 --- a/crates/asap_types/Cargo.toml +++ b/crates/asap_types/Cargo.toml @@ -23,3 +23,4 @@ planner-types.workspace = true asap-aware-mapping.workspace = true asap-physical-operators.workspace = true +asap_summary_state.workspace = true diff --git a/crates/asap_types/src/aggregation_type.rs b/crates/asap_types/src/aggregation_type.rs index 7ff2ce7ea..ebcc60d21 100644 --- a/crates/asap_types/src/aggregation_type.rs +++ b/crates/asap_types/src/aggregation_type.rs @@ -1,2 +1,2 @@ -//! Kernel identity is owned by the shared physical operator library. -pub use asap_physical_operators::AggregationType; +//! Kernel identity of backend-stored summary state. +pub use asap_summary_state::AggregationType; diff --git a/crates/asap_types/src/executable_plan.rs b/crates/asap_types/src/executable_plan.rs index 60220341c..d70baac8b 100644 --- a/crates/asap_types/src/executable_plan.rs +++ b/crates/asap_types/src/executable_plan.rs @@ -5,6 +5,7 @@ //! Planner payloads contain process-local `Rc` values, so they are decoded only //! when executing or validating a DAG. Compilation and QueryPlan cross-checks //! remain control-plane responsibilities. +use crate::physical_plan_codec::PhysicalPlanCodec; use std::collections::{BTreeMap, BTreeSet}; diff --git a/crates/asap_types/src/lib.rs b/crates/asap_types/src/lib.rs index 1d4af2fef..abb83ec16 100644 --- a/crates/asap_types/src/lib.rs +++ b/crates/asap_types/src/lib.rs @@ -9,6 +9,7 @@ pub mod executable_plan; pub mod grouping_projection; pub mod key_by_label_names; pub mod monitor_spec; +pub mod physical_plan_codec; pub mod plan_publication; pub mod policy_fingerprint; pub mod policy_registry; diff --git a/crates/asap_types/src/physical_plan_codec.rs b/crates/asap_types/src/physical_plan_codec.rs new file mode 100644 index 000000000..c0021b996 --- /dev/null +++ b/crates/asap_types/src/physical_plan_codec.rs @@ -0,0 +1,167 @@ +//! Persisted byte format of Planner physical plans. +//! +//! Planner exposes physical plans through serde only; the deployment owns the +//! versioned envelope. These formats are unchanged from when Planner encoded +//! them, so published plans and snapshots remain readable. +use asap_physical_operators::physical_planner::{ + CompiledPhysicalDag, InputContract, PhysicalCandidate, +}; +use asap_physical_operators::{plan::NodeId, Error}; +use serde_json::Value; +use std::collections::BTreeMap; + +const DAG_VERSION: u64 = 2; +const CANDIDATE_VERSION: u64 = 1; + +pub trait PhysicalPlanCodec: Sized { + fn encode(&self) -> Result, Error>; + fn decode(bytes: &[u8]) -> Result; +} + +fn invalid(error: impl ToString) -> Error { + Error::Invalid(error.to_string()) +} + +fn unversioned(version: u64, value: Value, format: &str) -> Result { + let Value::Object(mut fields) = value else { + return Err(invalid(format!("unsupported {format} format"))); + }; + if fields.remove("version").and_then(|v| v.as_u64()) != Some(version) { + return Err(invalid(format!("unsupported {format} format"))); + } + Ok(Value::Object(fields)) +} + +#[derive(serde::Serialize)] +struct StoredDag<'a> { + version: u64, + #[serde(flatten)] + dag: &'a CompiledPhysicalDag, +} + +#[derive(serde::Serialize)] +struct StoredCandidate<'a> { + version: u64, + precompute: Option, + query: Value, + materialized_outputs: &'a BTreeMap, +} + +fn dag_from_value(value: Value) -> Result { + serde_json::from_value(unversioned(DAG_VERSION, value, "physical plan")?).map_err(invalid) +} + +impl PhysicalPlanCodec for CompiledPhysicalDag { + /// Persist selected operators and input slots, never live state. + fn encode(&self) -> Result, Error> { + self.validate()?; + let bytes = serde_json::to_vec(&StoredDag { + version: DAG_VERSION, + dag: self, + }) + .map_err(invalid)?; + // JSON cannot preserve non-finite literal values. Fail at publication, + // rather than persisting a document that cannot be recovered. + Self::decode(&bytes)?; + Ok(bytes) + } + fn decode(bytes: &[u8]) -> Result { + dag_from_value(serde_json::from_slice(bytes).map_err(invalid)?) + } +} + +impl PhysicalPlanCodec for PhysicalCandidate { + fn encode(&self) -> Result, Error> { + self.validate()?; + let graph = |dag: &CompiledPhysicalDag| -> Result { + serde_json::from_slice(&dag.encode()?).map_err(invalid) + }; + serde_json::to_vec(&StoredCandidate { + version: CANDIDATE_VERSION, + precompute: self.precompute.as_ref().map(graph).transpose()?, + query: graph(&self.query)?, + materialized_outputs: &self.materialized_outputs, + }) + .map_err(invalid) + } + fn decode(bytes: &[u8]) -> Result { + let Value::Object(mut fields) = unversioned( + CANDIDATE_VERSION, + serde_json::from_slice(bytes).map_err(invalid)?, + "physical candidate", + )? + else { + unreachable!("unversioned returns an object"); + }; + for key in ["precompute", "query"] { + if let Some(graph) = fields.remove(key) { + let graph = match graph { + Value::Null => Value::Null, + graph => serde_json::to_value(dag_from_value(graph)?).map_err(invalid)?, + }; + fields.insert(key.into(), graph); + } + } + serde_json::from_value(Value::Object(fields)).map_err(invalid) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + + fn dag() -> CompiledPhysicalDag { + let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + let schema = + asap_physical_operators::physical_planner::precompute::population_schema(family); + CompiledPhysicalDag::from_operators( + BTreeMap::from([(1, InputContract::bounded(schema))]), + BTreeMap::new(), + vec![1], + ) + .unwrap() + } + + // A plan keeps its versioned envelope and round-trips unchanged. + #[test] + fn physical_dag_roundtrips_in_versioned_envelope() { + let bytes = dag().encode().unwrap(); + let value: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(value["version"], 2); + assert_eq!( + CompiledPhysicalDag::decode(&bytes) + .unwrap() + .encode() + .unwrap(), + bytes + ); + let mut other = value.clone(); + other["version"] = 3.into(); + assert!(CompiledPhysicalDag::decode(&serde_json::to_vec(&other).unwrap()).is_err()); + let mut unknown = value; + unknown["extra"] = 1.into(); + assert!(CompiledPhysicalDag::decode(&serde_json::to_vec(&unknown).unwrap()).is_err()); + } + + // A candidate nests versioned plans and rejects other candidate versions. + #[test] + fn physical_candidate_roundtrips_with_nested_plans() { + let candidate = PhysicalCandidate { + precompute: None, + query: dag(), + materialized_outputs: BTreeMap::new(), + }; + let bytes = candidate.encode().unwrap(); + let value: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(value["version"], 1); + assert_eq!(value["query"]["version"], 2); + assert_eq!( + PhysicalCandidate::decode(&bytes).unwrap().encode().unwrap(), + bytes + ); + let mut other = value; + other["version"] = 2.into(); + assert!(PhysicalCandidate::decode(&serde_json::to_vec(&other).unwrap()).is_err()); + } +} diff --git a/crates/asap_types/src/precompute_plan.rs b/crates/asap_types/src/precompute_plan.rs index e9cb9f253..3da190514 100644 --- a/crates/asap_types/src/precompute_plan.rs +++ b/crates/asap_types/src/precompute_plan.rs @@ -139,7 +139,7 @@ pub enum TimestampUnit { #[serde(deny_unknown_fields)] pub struct IngestContract { #[serde(default, skip_serializing_if = "Option::is_none")] - pub dataset_identity: Option, + pub dataset_identity: Option, pub protocol: IngestProtocol, pub endpoint_path: String, pub timestamp_unit: TimestampUnit, diff --git a/crates/asap_types/src/query_plan.rs b/crates/asap_types/src/query_plan.rs index 35c4349eb..2ded15630 100644 --- a/crates/asap_types/src/query_plan.rs +++ b/crates/asap_types/src/query_plan.rs @@ -4,6 +4,7 @@ //! maintained-summary leaf to one materialization and lowers edges to stable //! node IDs. Serving executes this graph without reconstructing Planner IR or //! searching for compatible materializations. +use crate::physical_plan_codec::PhysicalPlanCodec; pub mod current_series; mod native; diff --git a/crates/asap_types/src/query_plan/native.rs b/crates/asap_types/src/query_plan/native.rs index f8856486b..4f9feb627 100644 --- a/crates/asap_types/src/query_plan/native.rs +++ b/crates/asap_types/src/query_plan/native.rs @@ -1,5 +1,6 @@ //! Retained physical programs for SQL relations and PromQL vectors. use super::*; +use crate::physical_plan_codec::PhysicalPlanCodec; use asap_physical_operators::physical_planner::CompiledPhysicalDag; use planner_types::post_asap::SummarySchema; diff --git a/crates/asap_types/src/semantic_fragment.rs b/crates/asap_types/src/semantic_fragment.rs index 390354bd9..1e1f009ba 100644 --- a/crates/asap_types/src/semantic_fragment.rs +++ b/crates/asap_types/src/semantic_fragment.rs @@ -1,2 +1,657 @@ -//! Planner owns the versioned semantic description and its normalization. -pub use planner_types::post_asap::SummarySemanticFragment as SemanticFragment; +//! Persistable semantic dependency closure of a stored output, using Planner's +//! typed operation vocabulary. The backend owns this stored-definition identity. +use planner_types::post_asap::{ + EdgeRole, PostAsapDag, PostAsapNodeId, PostAsapOperatorPayload, SummarySchema, +}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::collections::{BTreeMap, BTreeSet}; + +pub use SummarySemanticFragment as SemanticFragment; + +/// Stable identity of a logical input dataset, independent of its endpoint. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LogicalDatasetIdentity { + pub namespace: String, + pub dataset: String, +} + +impl LogicalDatasetIdentity { + pub fn validate(&self) -> Result<(), String> { + if self.namespace.trim().is_empty() || self.dataset.trim().is_empty() { + return Err("dataset namespace and identity must be nonempty".into()); + } + Ok(()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SummarySemanticFragment { + pub format_version: u32, + /// Version 1 fragments are unbound structural descriptions. Persisted, + /// dataset-bound descriptions use version 2. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub dataset_identity: Option, + pub output: String, + pub nodes: BTreeMap, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SemanticOperation { + // Wire ownership must be Send + Sync. These values are checked against the + // Planner types on export and on recovery; arbitrary JSON is not accepted. + pub operation: serde_json::Value, + pub output_schema: serde_json::Value, + pub inputs: Vec, + /// The direct input range is supplied by the stored record, not by a query lookback. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub record_range: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SemanticInput { + pub role: EdgeRole, + pub node: String, +} + +pub(crate) fn canonical_bytes(value: &impl Serialize) -> Result, String> { + fn canonical(value: serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Object(values) => serde_json::Value::Object( + values + .into_iter() + .map(|(k, v)| (k, canonical(v))) + .collect::>() + .into_iter() + .collect(), + ), + serde_json::Value::Array(values) => { + serde_json::Value::Array(values.into_iter().map(canonical).collect()) + } + value => value, + } + } + serde_json::to_vec(&canonical( + serde_json::to_value(value).map_err(|e| e.to_string())?, + )) + .map_err(|e| e.to_string()) +} +fn hash(value: &impl Serialize) -> Result { + Ok(format!("{:x}", Sha256::digest(canonical_bytes(value)?))) +} +fn role(role: EdgeRole) -> u8 { + match role { + EdgeRole::Input => 0, + EdgeRole::Left => 1, + EdgeRole::Right => 2, + } +} + +impl SummarySemanticFragment { + pub fn from_stored_output(dag: &PostAsapDag, output: PostAsapNodeId) -> Result { + Self::export(dag, output, true) + } + + /// All source names in this DAG resolve within this logical dataset. + pub fn from_stored_output_in_dataset( + dag: &PostAsapDag, + output: PostAsapNodeId, + dataset: LogicalDatasetIdentity, + ) -> Result { + dataset.validate()?; + let mut fragment = Self::export(dag, output, true)?; + fragment.format_version = 2; + fragment.dataset_identity = Some(dataset); + fragment.validate()?; + Ok(fragment) + } + + pub fn from_dag(dag: &PostAsapDag, output: PostAsapNodeId) -> Result { + Self::export(dag, output, false) + } + + fn export( + dag: &PostAsapDag, + output: PostAsapNodeId, + parameterize_range: bool, + ) -> Result { + let mut included = BTreeSet::new(); + let mut pending = vec![output]; + while let Some(id) = pending.pop() { + if included.insert(id) { + pending.extend( + dag.edges + .iter() + .filter(|e| e.consumer == id) + .map(|e| e.producer), + ); + } + } + let dag = PostAsapDag { + nodes: dag + .nodes + .iter() + .filter(|n| included.contains(&n.id)) + .cloned() + .collect(), + edges: dag + .edges + .iter() + .filter(|e| included.contains(&e.consumer)) + .cloned() + .collect(), + root: output, + }; + dag.validate().map_err(|e| e.to_string())?; + if dag.nodes.len() > 4096 { + return Err("semantic fragment exceeds node budget".into()); + } + // Open PromQL entities carry all labels. Nullable label columns demanded + // only by a downstream consumer do not change a per-entity scalar state. + let mut dag = dag; + let sample_only = matches!( + &dag.nodes + .iter() + .find(|n| n.id == output) + .ok_or("missing output")? + .payload, + PostAsapOperatorPayload::SummaryAgg { + reduction: planner_types::pre_asap::Reduction::PerEntity, + grouping: planner_types::post_asap::GroupingStrategy::PerSubpopulationInstance, + input: planner_types::post_asap::SummaryUpdate { + item: None, + weight: planner_types::post_asap::SummaryInputExpr::Column( + planner_types::pre_asap::ColumnRef::SampleValue + ), + .. + }, + .. + } + ); + if parameterize_range && sample_only { + let direct: BTreeSet<_> = dag + .edges + .iter() + .filter(|e| e.consumer == output) + .map(|e| e.producer) + .collect(); + let mut normalized = false; + for node in &mut dag.nodes { + if !direct.contains(&node.id) { + continue; + } + if let PostAsapOperatorPayload::Fallback { expression } = &mut node.payload { + let source = match expression { + planner_types::pre_asap::QueryExpr::TimeRange { child, .. } => { + std::rc::Rc::make_mut(child) + } + other => other, + }; + if let planner_types::pre_asap::QueryExpr::Scan { + source: planner_types::pre_asap::Source::TimeSeries { .. }, + predicates, + schema, + } = source + { + if !schema.closed + && predicates.is_empty() + && schema.unique_keys.is_empty() + && schema + .columns + .iter() + .take_while(|c| { + !(c.nullable + && c.dtype == planner_types::pre_asap::DataType::Utf8) + }) + .count() + + schema + .columns + .iter() + .rev() + .take_while(|c| { + c.nullable + && c.dtype == planner_types::pre_asap::DataType::Utf8 + }) + .count() + == schema.columns.len() + { + schema.columns.retain(|c| { + !(c.nullable && c.dtype == planner_types::pre_asap::DataType::Utf8) + }); + node.output_schema.fields.retain(|c| { + !(c.nullable + && c.dtype + == planner_types::post_asap::SummaryFamilyType::Plain( + planner_types::pre_asap::DataType::Utf8, + )) + }); + normalized = true; + } + } + } + } + if normalized { + dag.nodes + .iter_mut() + .find(|n| n.id == output) + .unwrap() + .output_schema + .fields + .retain(|c| { + !(c.nullable + && c.dtype + == planner_types::post_asap::SummaryFamilyType::Plain( + planner_types::pre_asap::DataType::Utf8, + )) + }); + } + } + let nodes: BTreeMap<_, _> = dag.nodes.iter().map(|n| (n.id, n)).collect(); + let mut hashes: BTreeMap = BTreeMap::new(); + let mut result = Self { + format_version: 1, + dataset_identity: None, + output: String::new(), + nodes: BTreeMap::new(), + }; + let mut stack = vec![(output, false)]; + while let Some((id, finish)) = stack.pop() { + if hashes.contains_key(&id) { + continue; + } + let node = nodes.get(&id).ok_or("missing semantic output")?; + let edges: Vec<_> = dag.edges.iter().filter(|e| e.consumer == id).collect(); + if !finish { + stack.push((id, true)); + for edge in &edges { + stack.push((edge.producer, false)); + } + continue; + } + let mut inputs = edges + .iter() + .map(|e| SemanticInput { + role: e.role, + node: hashes[&e.producer].clone(), + }) + .collect::>(); + inputs.sort_by(|a, b| (role(a.role), &a.node).cmp(&(role(b.role), &b.node))); + let mut payload = node.payload.clone(); + let mut record_range = false; + if parameterize_range + && matches!( + nodes[&output].payload, + PostAsapOperatorPayload::SummaryAgg { .. } + ) + && dag + .edges + .iter() + .any(|e| e.consumer == output && e.producer == id) + { + if let PostAsapOperatorPayload::Fallback { + expression: planner_types::pre_asap::QueryExpr::TimeRange { child, .. }, + } = &payload + { + payload = PostAsapOperatorPayload::Fallback { + expression: child.as_ref().clone(), + }; + record_range = true; + } + } + if let PostAsapOperatorPayload::RelationalJoin { pruning, .. } = &mut payload { + *pruning = None; + } + let operation = SemanticOperation { + record_range, + operation: serde_json::to_value(&payload).map_err(|e| e.to_string())?, + output_schema: serde_json::to_value(&node.output_schema) + .map_err(|e| e.to_string())?, + inputs, + }; + let key = hash(&operation)?; + result.nodes.insert(key.clone(), operation); + hashes.insert(id, key); + } + result.output = hashes + .remove(&output) + .ok_or("missing semantic output hash")?; + result.validate()?; + Ok(result) + } + + pub fn validate(&self) -> Result<(), String> { + match (&self.dataset_identity, self.format_version) { + (None, 1) => (), + (Some(dataset), 2) => dataset.validate()?, + _ => return Err("semantic version and dataset binding disagree".into()), + } + if !matches!(self.format_version, 1 | 2) + || self.nodes.is_empty() + || self.nodes.len() > 4096 + || canonical_bytes(self)?.len() > 4 * 1024 * 1024 + { + return Err("unsupported semantic fragment version or size".into()); + } + for (key, node) in &self.nodes { + let payload: PostAsapOperatorPayload = + serde_json::from_value(node.operation.clone()).map_err(|e| e.to_string())?; + if node.record_range { + let root = self + .nodes + .get(&self.output) + .ok_or("missing semantic root")?; + let root_payload: PostAsapOperatorPayload = + serde_json::from_value(root.operation.clone()).map_err(|e| e.to_string())?; + if !matches!(payload, PostAsapOperatorPayload::Fallback { .. }) + || !matches!(root_payload, PostAsapOperatorPayload::SummaryAgg { .. }) + || !root.inputs.iter().any(|input| &input.node == key) + { + return Err("record range must belong to a direct summary input".into()); + } + } + let _: SummarySchema = + serde_json::from_value(node.output_schema.clone()).map_err(|e| e.to_string())?; + if hash(node)? != *key + || node + .inputs + .iter() + .any(|i| !self.nodes.contains_key(&i.node)) + { + return Err("semantic fragment hash or dependency mismatch".into()); + } + if node + .inputs + .windows(2) + .any(|p| (role(p[0].role), &p[0].node) > (role(p[1].role), &p[1].node)) + { + return Err("noncanonical semantic input order".into()); + } + } + let mut seen = BTreeSet::new(); + let mut stack = vec![self.output.as_str()]; + while let Some(id) = stack.pop() { + let node = self.nodes.get(id).ok_or("missing semantic fragment root")?; + if seen.insert(id) { + stack.extend(node.inputs.iter().map(|i| i.node.as_str())); + } + } + if seen.len() != self.nodes.len() { + return Err("unrelated semantic fragment nodes".into()); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use planner_types::post_asap::{compile_post_asap_dag, SummaryExpr, SummaryNode}; + use planner_types::pre_asap::{Column, DataType, QueryExpr, Schema, Source}; + use std::rc::Rc; + + fn fixture(metric: &str) -> PostAsapDag { + let scan = QueryExpr::Scan { + source: Source::TimeSeries { + metric: metric.into(), + }, + predicates: vec![], + schema: Schema::new(vec![Column::new("value", DataType::Float64, false)]), + }; + let schema = SummarySchema { + fields: vec![planner_types::post_asap::SummaryField { + name: "value".into(), + dtype: planner_types::post_asap::SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }], + time_index: None, + }; + compile_post_asap_dag(&Rc::new(SummaryNode { + expr: SummaryExpr::KeepPreAsap(Rc::new(scan)), + schema, + guarantee: None, + })) + .unwrap() + } + + // Equal source names in different datasets must not alias persisted meaning. + #[test] + fn dataset_identity_is_semantic_and_roundtrips() { + let dag = fixture("latency"); + let export = |namespace: &str| { + SummarySemanticFragment::from_stored_output_in_dataset( + &dag, + dag.root, + LogicalDatasetIdentity { + namespace: namespace.into(), + dataset: "requests".into(), + }, + ) + .unwrap() + }; + let a = export("tenant-a"); + assert_ne!( + canonical_bytes(&a).unwrap(), + canonical_bytes(&export("tenant-b")).unwrap() + ); + assert_eq!(a, export("tenant-a")); // No endpoint enters the semantic API. + let restored: SummarySemanticFragment = + serde_json::from_slice(&canonical_bytes(&a).unwrap()).unwrap(); + restored.validate().unwrap(); + assert_eq!(a, restored); + let mut bad = a.clone(); + bad.dataset_identity.as_mut().unwrap().namespace.clear(); + assert!(bad.validate().is_err()); + bad = a; + bad.dataset_identity = None; + assert!(bad.validate().is_err()); + } + + // Storage identity must ignore temporary identifiers and execution placement. + #[test] + fn identity_ignores_node_ids_and_phase() { + let dag = fixture("latency"); + let expected = SummarySemanticFragment::from_dag(&dag, dag.root).unwrap(); + let mut other = dag.clone(); + other.root = PostAsapNodeId(71); + other.nodes[0].id = other.root; + other.nodes[0].output_state.timing = + planner_types::post_asap::ExecutionTiming::IngestionTime; + assert_eq!( + canonical_bytes(&expected).unwrap(), + canonical_bytes(&SummarySemanticFragment::from_dag(&other, other.root).unwrap()) + .unwrap() + ); + } + + // Source identity and supported semantic format survive restart independently. + #[test] + fn semantics_roundtrip_and_reject_unknown_version() { + let a = fixture("latency"); + let b = fixture("bytes"); + let a = SummarySemanticFragment::from_dag(&a, a.root).unwrap(); + let b = SummarySemanticFragment::from_dag(&b, b.root).unwrap(); + assert_ne!(canonical_bytes(&a).unwrap(), canonical_bytes(&b).unwrap()); + let mut restored: SummarySemanticFragment = + serde_json::from_slice(&canonical_bytes(&a).unwrap()).unwrap(); + restored.validate().unwrap(); + restored.format_version += 1; + assert!(restored.validate().is_err()); + } + // A transformed value cannot share state identity with its source column. + #[test] + fn value_expression_is_semantic_and_nonfinite_constants_are_rejected() { + use planner_types::pre_asap::{ProjectItem, ScalarValue}; + let original = fixture("latency"); + let expected = SummarySemanticFragment::from_dag(&original, original.root).unwrap(); + let mut transformed = original.clone(); + let PostAsapOperatorPayload::Fallback { expression } = &mut transformed.nodes[0].payload + else { + unreachable!() + }; + *expression = QueryExpr::Project { + cols: vec![ProjectItem { + alias: Some("value".into()), + expr: QueryExpr::FunctionCall { + name: "ln".into(), + args: vec![QueryExpr::Column(0)], + }, + }], + qualifier: None, + child: Rc::new(expression.clone()), + }; + let logged = SummarySemanticFragment::from_dag(&transformed, transformed.root).unwrap(); + assert_ne!( + canonical_bytes(&expected).unwrap(), + canonical_bytes(&logged).unwrap() + ); + let PostAsapOperatorPayload::Fallback { + expression: QueryExpr::Project { cols, .. }, + } = &mut transformed.nodes[0].payload + else { + unreachable!() + }; + cols[0].expr = QueryExpr::Literal(ScalarValue::Float64(f64::NAN)); + assert!(SummarySemanticFragment::from_dag(&transformed, transformed.root).is_err()); + } + + // Changing a downstream consumer cannot change the persisted input definition. + #[test] + fn only_output_dependency_closure_is_exported() { + let mut dag = fixture("latency"); + let stored = dag.root; + let mut consumer = dag.nodes[0].clone(); + consumer.id = PostAsapNodeId(9); + consumer.payload = PostAsapOperatorPayload::Value { + operation: planner_types::post_asap::ValueOperation::Project { + cols: vec![], + qualifier: None, + }, + }; + dag.edges.push(planner_types::post_asap::PostAsapDagEdge { + producer: stored, + consumer: consumer.id, + role: EdgeRole::Input, + intermediate_schema: dag.nodes[0].output_schema.clone(), + data_state: dag.nodes[0].output_state, + grouping: planner_types::post_asap::GroupingEdgeCompatibility::NotApplicable, + window: planner_types::post_asap::WindowEdgeCompatibility::NotApplicable, + }); + dag.root = consumer.id; + dag.nodes.push(consumer); + let before = SummarySemanticFragment::from_dag(&dag, stored).unwrap(); + dag.nodes.reverse(); + let after = SummarySemanticFragment::from_dag(&dag, stored).unwrap(); + assert_eq!(before, after); + assert_eq!(before.nodes.len(), 1); + } + // Query lookback does not become the identity of each stored input pane. + #[test] + fn stored_input_range_is_parameterized_but_logical_range_is_preserved() { + use planner_types::post_asap::*; + use planner_types::pre_asap::{ColumnRef, Reduction}; + let make = |seconds| { + let mut dag = fixture("latency"); + let PostAsapOperatorPayload::Fallback { expression } = &mut dag.nodes[0].payload else { + unreachable!() + }; + *expression = QueryExpr::TimeRange { + range: std::time::Duration::from_secs(seconds), + child: Rc::new(expression.clone()), + }; + let mut output = dag.nodes[0].clone(); + output.id = PostAsapNodeId(1); + let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + output.payload = PostAsapOperatorPayload::SummaryAgg { + family: family.clone(), + input: SummaryUpdate { + item: None, + weight: SummaryInputExpr::Column(ColumnRef::SampleValue), + weight_domain: Default::default(), + }, + reduction: Reduction::PerEntity, + grouping: Default::default(), + }; + output.output_schema.fields[0].dtype = family; + output.output_state.primitive = DataPrimitive::SummaryState; + dag.edges.push(PostAsapDagEdge { + producer: dag.root, + consumer: output.id, + role: EdgeRole::Input, + intermediate_schema: dag.nodes[0].output_schema.clone(), + data_state: dag.nodes[0].output_state, + grouping: GroupingEdgeCompatibility::NotApplicable, + window: WindowEdgeCompatibility::NotApplicable, + }); + dag.root = output.id; + dag.nodes.push(output); + dag + }; + let one = make(60); + let five = make(300); + assert_ne!( + SummarySemanticFragment::from_dag(&one, one.root).unwrap(), + SummarySemanticFragment::from_dag(&five, five.root).unwrap() + ); + assert_eq!( + SummarySemanticFragment::from_stored_output(&one, one.root).unwrap(), + SummarySemanticFragment::from_stored_output(&five, five.root).unwrap() + ); + // Open entities retain their full label identity; consumer-demanded + // optional labels do not change the per-entity stored computation. + let mut open = one.clone(); + let PostAsapOperatorPayload::Fallback { + expression: QueryExpr::TimeRange { child, .. }, + } = &mut open.nodes[0].payload + else { + unreachable!() + }; + let QueryExpr::Scan { schema, .. } = Rc::make_mut(child) else { + unreachable!() + }; + schema.closed = false; + let expected = SummarySemanticFragment::from_stored_output(&open, open.root).unwrap(); + let PostAsapOperatorPayload::Fallback { + expression: QueryExpr::TimeRange { child, .. }, + } = &mut open.nodes[0].payload + else { + unreachable!() + }; + let QueryExpr::Scan { schema, .. } = Rc::make_mut(child) else { + unreachable!() + }; + schema + .columns + .push(Column::new("job", DataType::Utf8, true)); + let label = SummaryField { + name: "job".into(), + dtype: SummaryFamilyType::Plain(DataType::Utf8), + nullable: true, + }; + for node in &mut open.nodes { + node.output_schema.fields.push(label.clone()); + } + open.edges[0].intermediate_schema.fields.push(label); + assert_eq!( + expected, + SummarySemanticFragment::from_stored_output(&open, open.root).unwrap() + ); + let mut forged = SummarySemanticFragment::from_stored_output(&one, one.root).unwrap(); + forged.nodes.values_mut().next().unwrap().operation = + serde_json::json!({"kind": "unknown"}); + assert!(forged.validate().is_err()); + } + // Persisted semantic format changes require an explicit migration/version review. + #[test] + fn semantic_format_v1_has_stable_wire_identity() { + let dag = fixture("latency"); + let exported = SummarySemanticFragment::from_dag(&dag, dag.root).unwrap(); + assert_eq!( + exported.output, + "488a0550f37763997397403ae5ed3588dee5d2a09fa7840f4110b775095fe594" + ); + } +} diff --git a/crates/asap_types/src/traits.rs b/crates/asap_types/src/traits.rs index 51afd904f..4a91a97d5 100644 --- a/crates/asap_types/src/traits.rs +++ b/crates/asap_types/src/traits.rs @@ -1 +1 @@ -pub use asap_physical_operators::SerializableToSink; +pub use asap_summary_state::SerializableToSink; diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index 223275418..0f20e4625 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -8,6 +8,7 @@ edition.workspace = true asap_types.workspace = true asap-physical-operators.workspace = true asap_sketch_codec.workspace = true +asap_summary_state.workspace = true # Phase 9: the control plane is now an in-process library inside the # backend binary. Wiring up the in-process OpAMP server + capability-map # exposure is a follow-up after Phase 4 (centralized series_id @@ -103,4 +104,4 @@ default = [] # Enable lock profiling instrumentation lock_profiling = [] # Enable extra debugging output -extra_debugging = ["asap-physical-operators/extra_debugging"] +extra_debugging = ["asap_summary_state/extra_debugging"] diff --git a/data_plane/benches/sketch_db.rs b/data_plane/benches/sketch_db.rs index de4a448cf..0797aaf44 100644 --- a/data_plane/benches/sketch_db.rs +++ b/data_plane/benches/sketch_db.rs @@ -33,7 +33,7 @@ use asap_sketchlib::DdSketch; use asap_sketchlib::{HllSketch, HllVariant}; use prost::Message; -use asap_physical_operators::summary_kernels::SumAccumulator; +use asap_summary_state::summary_kernels::SumAccumulator; use data_plane::storage_engines::sketch_db::data::{ AccuracyBound, AggKind, AggregationType, Capability, SketchAlgorithm, SketchConfig, SketchEncoding, diff --git a/data_plane/examples/univmon_erp_artifact.rs b/data_plane/examples/univmon_erp_artifact.rs index c20687a1b..2ac5e07e1 100644 --- a/data_plane/examples/univmon_erp_artifact.rs +++ b/data_plane/examples/univmon_erp_artifact.rs @@ -1,7 +1,7 @@ //! Measure readout-specific ERP evidence from finite JSONL evaluation data. //! This offline tool retains samples; the production backend does not. -use asap_physical_operators::summary_kernels::hll_sketch::HllSketchAccumulator; -use asap_physical_operators::summary_kernels::univmon::UnivMonAccumulator; +use asap_summary_state::summary_kernels::hll_sketch::HllSketchAccumulator; +use asap_summary_state::summary_kernels::univmon::UnivMonAccumulator; use data_plane::storage_engines::types::{AggregateCore, SerializableToSink}; use serde_json::{json, Value}; use std::collections::{BTreeMap, HashMap}; diff --git a/data_plane/src/drivers/ingest/otel.rs b/data_plane/src/drivers/ingest/otel.rs index 81925b1a7..92e470c41 100644 --- a/data_plane/src/drivers/ingest/otel.rs +++ b/data_plane/src/drivers/ingest/otel.rs @@ -34,9 +34,9 @@ use asap_otel_proto::tonic::collector::metrics::v1::{ }; use asap_otel_proto::tonic::common::v1::any_value::Value as AnyValueVariant; use asap_otel_proto::tonic::metrics::v1::number_data_point::Value as NumberValue; -use asap_physical_operators::summary_kernels::sketch_envelope::SketchEnvelopeAccumulator; use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; use asap_sketchlib::MessagePackCodec; +use asap_summary_state::summary_kernels::sketch_envelope::SketchEnvelopeAccumulator; use axum::{body::Bytes, extract::State, routing::post, Json, Router}; use flate2::read::GzDecoder; use planner_types::post_asap::SketchAlgorithm; @@ -2110,7 +2110,7 @@ fn dp_carries_heap(dp: &ModifiedOtlpSketchDp) -> bool { .unwrap_or(false) } ENCODING_MSGPACK_DELTA => { - use asap_physical_operators::summary_kernels::CountMinSketchWithHeapAccumulator; + use asap_summary_state::summary_kernels::CountMinSketchWithHeapAccumulator; CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes(&dp.sketch) .map(|acc| !acc.inner.topk_heap_items().is_empty()) .unwrap_or(false) @@ -2539,7 +2539,7 @@ fn decode_modified_otlp_sketch_bytes( encoding: i32, bytes: &[u8], ) -> Result, Box> { - use asap_physical_operators::summary_kernels::{ + use asap_summary_state::summary_kernels::{ CountMinSketchAccumulator, CountSketchAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, HllSketchAccumulator, }; @@ -2610,7 +2610,7 @@ fn decode_modified_otlp_sketch_bytes( use asap_sketchlib::CountSketchWithHeap; if let Ok(heap) = CountSketchWithHeap::from_msgpack(bytes) { if !heap.topk_heap_items().is_empty() { - use asap_physical_operators::summary_kernels::CountSketchWithHeapAccumulator; + use asap_summary_state::summary_kernels::CountSketchWithHeapAccumulator; return Ok(Box::new( CountSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(bytes)?, )); @@ -2676,7 +2676,7 @@ fn empty_accumulator_for_delta_bootstrap( encoding: i32, ) -> Option> { use crate::storage_engines::sketch_db::index::SketchConfig; - use asap_physical_operators::summary_kernels::{ + use asap_summary_state::summary_kernels::{ CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, HllSketchAccumulator, }; @@ -2749,7 +2749,7 @@ pub(crate) fn apply_modified_otlp_delta_bytes( existing: &mut Box, bytes: &[u8], ) -> Result<(), Box> { - use asap_physical_operators::summary_kernels::{ + use asap_summary_state::summary_kernels::{ CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, DDSketchAccumulator, HllSketchAccumulator, }; @@ -2995,7 +2995,7 @@ fn otlp_to_metric_points_and_sketches(request: &ExportMetricsServiceRequest) -> // ExactAgg(Sum) path as a plain delta Sum — the backend sums // the per-window/per-shard partials for the same sid. for dp in &sa.data_points { - let value = match asap_physical_operators::summary_kernels::sum::SumAccumulator::from_sum_bytes(&dp.sketch) { + let value = match asap_summary_state::summary_kernels::sum::SumAccumulator::from_sum_bytes(&dp.sketch) { Ok(acc) => acc.sum, Err(e) => { debug!("asap_edge: SumAgg data point decode failed (skipping): {e}"); @@ -3436,9 +3436,9 @@ mod policy_fp_lookup_tests { mod dispatcher_tests { use super::*; use crate::storage_engines::types::AggregateCore; - use asap_physical_operators::summary_kernels::{DDSketchAccumulator, HllSketchAccumulator}; use asap_sketchlib::DdSketch; use asap_sketchlib::HllVariant; + use asap_summary_state::summary_kernels::{DDSketchAccumulator, HllSketchAccumulator}; #[test] fn apply_modified_otlp_delta_bytes_ddsketch_round_trip() { @@ -3789,8 +3789,8 @@ mod sid_resolution_tests { #[tokio::test] async fn delta_apply_rotates_per_series_base_at_window_boundary() { use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; - use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use asap_summary_state::summary_kernels::DDSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; @@ -4122,7 +4122,7 @@ mod sid_resolution_tests { #[tokio::test] async fn leading_cms_delta_bootstraps_onto_empty_base() { use asap_otel_proto::sketchlib::v1::CountMinDelta as PbDelta; - use asap_physical_operators::summary_kernels::CountMinSketchAccumulator; + use asap_summary_state::summary_kernels::CountMinSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; @@ -4208,7 +4208,7 @@ mod sid_resolution_tests { #[tokio::test] async fn leading_hll_delta_bootstraps_onto_empty_base() { use asap_otel_proto::sketchlib::v1::HllDelta as PbDelta; - use asap_physical_operators::summary_kernels::HllSketchAccumulator; + use asap_summary_state::summary_kernels::HllSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; diff --git a/data_plane/src/drivers/ingest/prometheus_remote_write.rs b/data_plane/src/drivers/ingest/prometheus_remote_write.rs index 3f05d9b2d..6acc03e5c 100644 --- a/data_plane/src/drivers/ingest/prometheus_remote_write.rs +++ b/data_plane/src/drivers/ingest/prometheus_remote_write.rs @@ -1858,7 +1858,7 @@ impl PrometheusRemoteWriteReceiver { if committed.contains_key(&policy_fp.0) { continue; } let program = RawDagProgram::from_plan(&plan.precompute_plan, config)?; let manager = WindowManager::with_layout(config.window_size, config.slide_interval, config.pane_origin_ms, &config.window_layout); - let mut windows: BTreeMap<(u64,u64), Box> = BTreeMap::new(); + let mut windows: BTreeMap<(u64,u64), Box> = BTreeMap::new(); samples.sort_by(|a,b| (&a.0,a.1).cmp(&(&b.0,b.1))); for (series,time,value) in samples { for start in manager.stored_bucket_starts(time) { diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index 158d1acb0..c7d3106df 100644 --- a/data_plane/src/lib.rs +++ b/data_plane/src/lib.rs @@ -42,7 +42,7 @@ pub use storage_engines::types::{ SerializableToSink, SingleSubpopulationAggregate, }; -pub use asap_physical_operators::summary_kernels::{ +pub use asap_summary_state::summary_kernels::{ IncreaseAccumulator, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, }; diff --git a/data_plane/src/precompute_engine/ingest_handler.rs b/data_plane/src/precompute_engine/ingest_handler.rs index 400797a7a..3d20001f4 100644 --- a/data_plane/src/precompute_engine/ingest_handler.rs +++ b/data_plane/src/precompute_engine/ingest_handler.rs @@ -367,8 +367,8 @@ mod tests { async fn delta_path_reconstitutes_cumulative_state() { use crate::drivers::ingest::otel::apply_modified_otlp_delta_bytes; use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; - use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_sketchlib::DdSketch; + use asap_summary_state::summary_kernels::DDSketchAccumulator; use planner_types::post_asap::SketchAlgorithm; use prost::Message; @@ -472,7 +472,7 @@ mod tests { /// survive; a stale entry from far in the past must be swept. #[tokio::test] async fn stale_snapshot_entry_is_evicted_by_sweep() { - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; let (state, drain) = setup_state(7, "evict_metric").await; diff --git a/data_plane/src/precompute_engine/maintenance_runtime.rs b/data_plane/src/precompute_engine/maintenance_runtime.rs index 9e6107e41..560c5f627 100644 --- a/data_plane/src/precompute_engine/maintenance_runtime.rs +++ b/data_plane/src/precompute_engine/maintenance_runtime.rs @@ -1,4 +1,7 @@ //! Bind immutable inputs to retained physical graphs and publish their stored outputs. +use asap_summary_state::codec::KeyCodec; +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; use super::output_sink::OutputSink; use crate::storage_engines::types::{ @@ -1560,7 +1563,7 @@ pub(crate) fn affected_materializations( mod tests { use super::*; - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; use planner_types::post_asap::{ EdgeRole, GroupingEdgeCompatibility, PostAsapDag, PostAsapDagEdge, SummarySchema, WindowEdgeCompatibility, @@ -1616,7 +1619,7 @@ mod tests { fn cohort_lineage_is_order_independent_and_binds_every_input() { use crate::storage_engines::sketch_db::index::FrozenExactWindows; let make = |sid, id, value| { - let mut state = asap_physical_operators::summary_kernels::SumAccumulator::new(); + let mut state = asap_summary_state::summary_kernels::SumAccumulator::new(); state.update(value); FrozenExactWindows { stored_output_reference: asap_types::sds::StoredOutputReference::for_output( @@ -3041,8 +3044,7 @@ pub(crate) fn execute_revision_outputs( } continue; } - let payload = - asap_physical_operators::stored_state::native::encode_batch(&batch)?; + let payload = asap_summary_state::stored_state::native::encode_batch(&batch)?; if payload.len() > limit { return Err(asap_physical_operators::Error::MemoryLimit.into()); } diff --git a/data_plane/src/precompute_engine/native_precompute.rs b/data_plane/src/precompute_engine/native_precompute.rs index a669e4042..9db27c275 100644 --- a/data_plane/src/precompute_engine/native_precompute.rs +++ b/data_plane/src/precompute_engine/native_precompute.rs @@ -1,4 +1,6 @@ //! Bind a complete durable counter cohort to a Planner-owned precompute graph. +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; use std::{collections::BTreeMap, sync::Arc}; use asap_physical_operators::{ @@ -71,7 +73,7 @@ pub(super) fn execute( ), Value::Summary { family: family.clone(), - state: Arc::clone(state), + state: asap_summary_state::physical::to_physical(state.as_ref())?, }, ]); } @@ -91,7 +93,8 @@ pub(super) fn execute( _, ) => Ok(Value::Summary { family: field.dtype.clone(), - state: Arc::clone(state), + state: asap_summary_state::physical::to_physical(state.as_ref()) + .map_err(|e| e.to_string())?, }), SummaryFamilyType::Plain(DataType::Timestamp) => Ok(Value::Timestamp( i64::try_from(window.1).map_err(|_| "native window overflow")?, @@ -213,7 +216,8 @@ pub(super) fn population_states( return Err("duplicate population label".into()); } } - if result.insert(group, Arc::clone(state)).is_some() { + let state = asap_summary_state::physical::from_physical(state.as_ref())?; + if result.insert(group, Arc::from(state)).is_some() { return Err("repeated precompute output population".into()); } } @@ -293,7 +297,7 @@ mod tests { snapshot_sha256: "0".repeat(64), }); let state = |value| { - let mut sum = asap_physical_operators::summary_kernels::SumAccumulator::new(); + let mut sum = asap_summary_state::summary_kernels::SumAccumulator::new(); sum.update(value); Arc::new(sum) as Arc }; @@ -331,7 +335,8 @@ mod tests { "shared merge must execute once, not once per output" ); assert_eq!( - state_at(3) + asap_summary_state::physical::from_physical(state_at(3).as_ref()) + .unwrap() .query_statistic(asap_types::Statistic::Sum, &None, &Default::default(),) .unwrap(), value + 3.0 diff --git a/data_plane/src/precompute_engine/output_sink.rs b/data_plane/src/precompute_engine/output_sink.rs index eab97b320..4b05edd1a 100644 --- a/data_plane/src/precompute_engine/output_sink.rs +++ b/data_plane/src/precompute_engine/output_sink.rs @@ -325,7 +325,7 @@ mod tests { use super::*; use crate::storage_engines::sketch_db::index::{AggKind, SeriesLookup}; use crate::storage_engines::types::{InstalledPrecomputePlan, KeyByLabelValues}; - use asap_physical_operators::summary_kernels::{DDSketchAccumulator, SumAccumulator}; + use asap_summary_state::summary_kernels::{DDSketchAccumulator, SumAccumulator}; use asap_types::aggregation_config::PrecomputeMaterialization; use asap_types::enums::WindowKind; use asap_types::AggregationType; diff --git a/data_plane/src/precompute_engine/raw_dag.rs b/data_plane/src/precompute_engine/raw_dag.rs index 993716d09..d549f1533 100644 --- a/data_plane/src/precompute_engine/raw_dag.rs +++ b/data_plane/src/precompute_engine/raw_dag.rs @@ -1,6 +1,6 @@ //! Bind raw ingestion to a selected Planner producer and its raw dependency edge. use crate::storage_engines::types::KeyByLabelValues; -use asap_physical_operators::factory::{create_planner_accumulator, AccumulatorUpdater}; +use asap_summary_state::factory::{create_planner_accumulator, AccumulatorUpdater}; use asap_types::{executable_plan::BackendNodeBinding, PrecomputeMaterialization}; use planner_types::post_asap::{ EdgeRole, GroupingStrategy, PostAsapNodeId, PostAsapOperatorPayload, SummaryFamilyType, diff --git a/data_plane/src/precompute_engine/revisions.rs b/data_plane/src/precompute_engine/revisions.rs index 8aca75d3d..21c9c55ae 100644 --- a/data_plane/src/precompute_engine/revisions.rs +++ b/data_plane/src/precompute_engine/revisions.rs @@ -431,10 +431,8 @@ fn validate_records(output: u64, records: &[RevisionRecord]) -> Result<(), Revis use crate::storage_engines::types::{ AggregateCore, InstalledPrecomputePlanHandle, RuntimePhysicalPlan, }; -use asap_physical_operators::{ - stored_state::native, - values::{Batch, Schema, Value}, -}; +use asap_physical_operators::values::{Batch, Schema, Value}; +use asap_summary_state::stored_state::native; use planner_types::post_asap::{SummaryFamilyType, SummaryField, SummarySchema}; fn state_schema(family: SummaryFamilyType) -> Schema { @@ -454,7 +452,10 @@ pub(crate) fn encode_state( ) -> Result, RevisionError> { Ok(native::encode_batch(&Batch::try_new( state_schema(family.clone()), - vec![vec![Value::Summary { family, state }]], + vec![vec![Value::Summary { + family, + state: asap_summary_state::physical::to_physical(state.as_ref())?, + }]], )?)?) } @@ -466,7 +467,9 @@ pub(crate) fn decode_state( let batch = native::decode_batch(&record.payload, state_schema(family), record.payload.len())?; match batch.rows() { [row] => match row.as_slice() { - [Value::Summary { state, .. }] => Ok(Arc::clone(state)), + [Value::Summary { state, .. }] => Ok(Arc::from( + asap_summary_state::physical::from_physical(state.as_ref())?, + )), _ => Err("revision record must contain exactly one typed summary".into()), }, _ => Err("revision record must contain exactly one row".into()), diff --git a/data_plane/src/precompute_engine/worker.rs b/data_plane/src/precompute_engine/worker.rs index f2da34982..2f6d3e085 100644 --- a/data_plane/src/precompute_engine/worker.rs +++ b/data_plane/src/precompute_engine/worker.rs @@ -9,8 +9,9 @@ use crate::storage_engines::types::{ }; #[cfg(test)] use crate::tests::accumulator_fixture::create_fixture_accumulator; -use asap_physical_operators::factory::AccumulatorUpdater; -use asap_physical_operators::summary_kernels::sum::SumAccumulator; +use asap_summary_state::codec::KeyCodec; +use asap_summary_state::factory::AccumulatorUpdater; +use asap_summary_state::summary_kernels::sum::SumAccumulator; use asap_types::aggregation_config::PrecomputeMaterialization; use asap_types::PolicyFingerprint; use asap_types::SampleUpdateRule; @@ -1971,10 +1972,10 @@ mod tests { use crate::precompute_engine::config::LateDataPolicy; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::InstalledPrecomputePlan; - use asap_physical_operators::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; - use asap_physical_operators::summary_kernels::keyed_sum_count::KeyedSumCountAccumulator; - use asap_physical_operators::summary_kernels::sum::SumAccumulator; use asap_sketchlib::KllSketch; + use asap_summary_state::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; + use asap_summary_state::summary_kernels::keyed_sum_count::KeyedSumCountAccumulator; + use asap_summary_state::summary_kernels::sum::SumAccumulator; use asap_types::enums::WindowKind; use asap_types::sds::StoredOutputId; use asap_types::AggregationType; @@ -3268,8 +3269,8 @@ mod tests { // OTLP ingest dispatch builds via `decode_modified_otlp_sketch_bytes`. // ----------------------------------------------------------------------- - use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_sketchlib::DdSketch; + use asap_summary_state::summary_kernels::DDSketchAccumulator; /// Build a fresh DDSketch holding `vals` so each test has a real, /// non-empty sketch to push through `process_accumulator_input`. @@ -4033,7 +4034,7 @@ mod tests { // A pooled Sum is correct only for an explicit cross-entity reduction. #[test] fn pooled_sum_does_not_preserve_per_entity_output_rows() { - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; let config = make_agg_config( 1, "gauge", @@ -4087,7 +4088,7 @@ mod tests { // The physical compiler rejects raw counter producers until series state is preserved. #[test] fn pooled_counter_samples_lose_independent_same_timestamp_reset() { - use asap_physical_operators::summary_kernels::IncreaseAccumulator; + use asap_summary_state::summary_kernels::IncreaseAccumulator; let config = make_agg_config( 1, "requests_total", @@ -4451,7 +4452,7 @@ mod dag_execution_tests { use super::*; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::InstalledPrecomputePlan; - use asap_physical_operators::summary_kernels::exact::ExactAccumulator; + use asap_summary_state::summary_kernels::exact::ExactAccumulator; use asap_types::query_plan::ExactReadout; fn plan(query: &str) -> control_plane::physical::compiler::CompiledPhysicalPlan { diff --git a/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs b/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs index 709372075..3e91bd5e5 100644 --- a/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs +++ b/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs @@ -1,6 +1,8 @@ //! Catalog-backed ClickHouse acceleration boundary. +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; -use asap_physical_operators::summary_kernels::SumAccumulator; +use asap_summary_state::summary_kernels::SumAccumulator; use async_trait::async_trait; use axum::{ body::Bytes, @@ -531,7 +533,7 @@ mod tests { } use crate::storage_engines::sketch_db::index::{AggKind, Capability, SummarySeriesMetadata}; - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; use asap_types::query_plan::{ ClickHousePlanningContext, ExactReadout, ExternalExactOutput, ExternalExactRequest, FallbackPolicy, FixedEvaluationRange, InstantExecution, MaterializationBinding, diff --git a/data_plane/src/query_engines/asap_clickhouse_query_engine/execution.rs b/data_plane/src/query_engines/asap_clickhouse_query_engine/execution.rs index 81afff3bf..1ae3909a4 100644 --- a/data_plane/src/query_engines/asap_clickhouse_query_engine/execution.rs +++ b/data_plane/src/query_engines/asap_clickhouse_query_engine/execution.rs @@ -9,6 +9,7 @@ use crate::{ }, storage_engines::sketch_db::index::SketchStore, }; +use asap_types::physical_plan_codec::PhysicalPlanCodec; use asap_types::query_plan::{QueryNodeId, QueryPlanEntry, QueryPlanNode}; use asap_types::summary_catalog::SummaryCatalog; use std::collections::{BTreeMap, BTreeSet}; diff --git a/data_plane/src/query_engines/asap_query_engine/engine.rs b/data_plane/src/query_engines/asap_query_engine/engine.rs index d263f71f1..66906aba4 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -1548,7 +1548,7 @@ mod sketch_query_tests { mod aux_pushdown_tests { use super::*; use crate::storage_engines::types::AggregationType; - use asap_physical_operators::summary_kernels::{ + use asap_summary_state::summary_kernels::{ max::MaxAccumulator, min::MinAccumulator, sum::SumAccumulator, }; use asap_types::Statistic; @@ -1767,7 +1767,7 @@ mod asap_tier_classify_tests { async fn execute_sum_by_zone_dispatches_to_exact_agg_reducer() { use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; - use asap_physical_operators::summary_kernels::sum::SumAccumulator; + use asap_summary_state::summary_kernels::sum::SumAccumulator; let idx = Arc::new(SketchStore::new()); // Mirror the acceptance-test setup: four ExactAgg(Sum) sids, one @@ -2421,7 +2421,7 @@ mod asap_tier_classify_tests { async fn execute_instant_sum_accumulates_all_windows_not_last() { use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; - use asap_physical_operators::summary_kernels::sum::SumAccumulator; + use asap_summary_state::summary_kernels::sum::SumAccumulator; let idx = Arc::new(SketchStore::new()); let now_ms = 600_000_u64; diff --git a/data_plane/src/query_engines/asap_query_engine/exact_subqueries.rs b/data_plane/src/query_engines/asap_query_engine/exact_subqueries.rs index 3f514158f..9ab038fc7 100644 --- a/data_plane/src/query_engines/asap_query_engine/exact_subqueries.rs +++ b/data_plane/src/query_engines/asap_query_engine/exact_subqueries.rs @@ -1,6 +1,8 @@ //! Fetch installed exact cuts from Prometheus before composing them with ASAP state. use super::logical_dag::{PreparedLeaf, PreparedLeaves, Value}; use crate::query_engines::EngineError; +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; use asap_types::query_plan::{ query_time::QueryTimeOperator, ExternalExactInput, ExternalExactRequest, QueryLanguage, QueryNodeId, QueryPlanEntry, QueryPlanNode, @@ -967,7 +969,7 @@ mod tests { index::{Capability, SummarySeriesMetadata}, }; use crate::storage_engines::types::Measurement; - use asap_physical_operators::summary_kernels::IncreaseAccumulator; + use asap_summary_state::summary_kernels::IncreaseAccumulator; use asap_types::query_plan::{ query_time::BinaryOperation, ExactReadout, MaterializationBinding, PhysicalGrouping, }; diff --git a/data_plane/src/query_engines/asap_query_engine/live_serve.rs b/data_plane/src/query_engines/asap_query_engine/live_serve.rs index fdf481272..a1c4c5aef 100644 --- a/data_plane/src/query_engines/asap_query_engine/live_serve.rs +++ b/data_plane/src/query_engines/asap_query_engine/live_serve.rs @@ -170,7 +170,7 @@ mod tests { 9, BTreeMap::new(), (start, end), - Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(value)), + Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(value)), ); } let entry = asap_types::query_plan::QueryPlanEntry { diff --git a/data_plane/src/query_engines/asap_query_engine/logical_dag.rs b/data_plane/src/query_engines/asap_query_engine/logical_dag.rs index c81ad52a4..1aa86c2aa 100644 --- a/data_plane/src/query_engines/asap_query_engine/logical_dag.rs +++ b/data_plane/src/query_engines/asap_query_engine/logical_dag.rs @@ -1,4 +1,6 @@ //! Executes the installed typed logical DAG. No serving-time PromQL parsing. +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; pub(super) mod native_values; use crate::query_engines::{ query_result::{InstantVectorElement, QueryResult}, @@ -1043,7 +1045,7 @@ mod topk_tests { // An overflowing sum cannot implement average, but zero/subnormal averages remain valid. #[test] fn finite_division_guards_temporal_average_without_rejecting_zero() { - let mut sum = asap_physical_operators::summary_kernels::sum::SumAccumulator::new(); + let mut sum = asap_summary_state::summary_kernels::sum::SumAccumulator::new(); sum.update(1e308); sum.update(1e308); assert!(binary( diff --git a/data_plane/src/query_engines/asap_query_engine/logical_dag/native_values.rs b/data_plane/src/query_engines/asap_query_engine/logical_dag/native_values.rs index 5c5ea3a83..c69ae80bd 100644 --- a/data_plane/src/query_engines/asap_query_engine/logical_dag/native_values.rs +++ b/data_plane/src/query_engines/asap_query_engine/logical_dag/native_values.rs @@ -5,6 +5,7 @@ use asap_physical_operators::dag::{ operators::Operator, values::{Batch, Schema, Value}, }; +use asap_types::physical_plan_codec::PhysicalPlanCodec; #[cfg(test)] use planner_types::post_asap::{SummaryField, SummarySchema}; use planner_types::{post_asap::SummaryFamilyType, pre_asap::DataType}; diff --git a/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs b/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs index 5bf9304a5..6c959d0cb 100644 --- a/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs +++ b/data_plane/src/query_engines/asap_query_engine/post_asap_readout.rs @@ -1,4 +1,6 @@ //! Execute installed query DAGs and project their values and coverage. +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; use std::collections::BTreeMap; @@ -1091,7 +1093,7 @@ mod tests { let groups = [1, 2] .into_iter() .map(|samples| { - let mut state = asap_physical_operators::summary_kernels::IncreaseAccumulator::new( + let mut state = asap_summary_state::summary_kernels::IncreaseAccumulator::new( Measurement::new(10.0), 10_000, Measurement::new(10.0), @@ -1165,7 +1167,7 @@ mod tests { policy_fp: policy, }); for (start, time, value) in [(20_000, 25_000, 10.), (40_000, 45_000, 20.)] { - let state = asap_physical_operators::summary_kernels::IncreaseAccumulator::new( + let state = asap_summary_state::summary_kernels::IncreaseAccumulator::new( Measurement::new(value), time, Measurement::new(value), @@ -1236,7 +1238,7 @@ mod tests { sid, BTreeMap::from([("instance".into(), instance.into())]), (0, 1_000), - Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(1.)), + Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(1.)), ); } let binding = asap_types::query_plan::MaterializationBinding { @@ -1808,7 +1810,7 @@ mod tests { 1, BTreeMap::new(), (1_000, 2_000), - Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(42.0)), + Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(42.0)), ); let config = test_plan::materialization("bytes_total", "Sum", serde_json::json!({}), &[], 1000); @@ -1906,9 +1908,7 @@ mod tests { BTreeMap::new(), bounds, Box::new( - asap_physical_operators::summary_kernels::SumAccumulator::with_sum( - sum, - ), + asap_summary_state::summary_kernels::SumAccumulator::with_sum(sum), ), ); } @@ -1976,7 +1976,7 @@ mod tests { BTreeMap::new(), (pane * 60_000, (pane + 1) * 60_000), Box::new( - asap_physical_operators::summary_kernels::SumAccumulator::with_sum( + asap_summary_state::summary_kernels::SumAccumulator::with_sum( (pane + 1) as f64, ), ), @@ -2051,7 +2051,7 @@ mod tests { BTreeMap::new(), (pane * 10_000, (pane + 1) * 10_000), Box::new( - asap_physical_operators::summary_kernels::SumAccumulator::with_sum( + asap_summary_state::summary_kernels::SumAccumulator::with_sum( (pane + 1) as f64, ), ), @@ -2119,7 +2119,7 @@ mod tests { 7, BTreeMap::new(), (pane * 10_000, (pane + 1) * 10_000), - Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(1.0)), + Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(1.0)), ); } assert!( @@ -2149,7 +2149,7 @@ mod tests { policy_fp: policy, }); use crate::storage_engines::types::Measurement; - let mut accumulator = asap_physical_operators::summary_kernels::IncreaseAccumulator::new( + let mut accumulator = asap_summary_state::summary_kernels::IncreaseAccumulator::new( Measurement::new(10.0), 10_000, Measurement::new(10.0), diff --git a/data_plane/src/query_engines/asap_query_engine/raw_source.rs b/data_plane/src/query_engines/asap_query_engine/raw_source.rs index 2456c3996..fa72e089f 100644 --- a/data_plane/src/query_engines/asap_query_engine/raw_source.rs +++ b/data_plane/src/query_engines/asap_query_engine/raw_source.rs @@ -12,6 +12,8 @@ use asap_physical_operators::{ values::{Batch, Schema}, Error, }; +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; use asap_types::query_plan::query_time::{LabelMatch, QueryTimeOperator}; use futures::{stream, StreamExt, TryStreamExt}; use planner_types::{ diff --git a/data_plane/src/query_engines/asap_query_engine/request_tests.rs b/data_plane/src/query_engines/asap_query_engine/request_tests.rs index 265e3b8cb..2ae8e83c6 100644 --- a/data_plane/src/query_engines/asap_query_engine/request_tests.rs +++ b/data_plane/src/query_engines/asap_query_engine/request_tests.rs @@ -202,7 +202,7 @@ async fn installed_local_range_accounts_for_accumulated_results() { 1, BTreeMap::new(), (pane * 1000, (pane + 1) * 1000), - Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(2.0)), + Box::new(asap_summary_state::summary_kernels::SumAccumulator::with_sum(2.0)), ); } engine diff --git a/data_plane/src/query_engines/asap_query_engine/summary_executor.rs b/data_plane/src/query_engines/asap_query_engine/summary_executor.rs index 0511e1aad..f23963d78 100644 --- a/data_plane/src/query_engines/asap_query_engine/summary_executor.rs +++ b/data_plane/src/query_engines/asap_query_engine/summary_executor.rs @@ -1,5 +1,5 @@ //! Deployment adapters for resolving, decoding, and reading installed materializations. -use asap_physical_operators::summary_kernels::{MaxAccumulator, MinAccumulator}; +use asap_summary_state::summary_kernels::{MaxAccumulator, MinAccumulator}; use std::collections::{BTreeMap, BTreeSet}; use std::rc::Rc; use std::sync::Arc; @@ -88,7 +88,7 @@ impl GroupState { SummaryFamilyType::ExactAggregate(ExactKind::Rate, _) => asap_types::Statistic::Rate, _ => return None, }; - asap_physical_operators::stored_state::readout::exact_readout( + asap_summary_state::stored_state::readout::exact_readout( entries.iter().flat_map(|windows| windows.values().cloned()), stat, key, @@ -133,7 +133,7 @@ impl GroupState { asap_types::query_plan::ExactReadout::Max => asap_types::Statistic::Max, }; - asap_physical_operators::stored_state::readout::exact_readout_optional( + asap_summary_state::stored_state::readout::exact_readout_optional( entries.iter().flat_map(|windows| windows.values().cloned()), stat, key, @@ -757,15 +757,15 @@ fn sketch_query_value( state: &SummaryState, query: &SketchQuery, ) -> Result { - asap_physical_operators::stored_state::readout::sketch_query_value(state, query).map_err( - |asap_physical_operators::stored_state::readout::Error::Unsupported(reason)| { + asap_summary_state::stored_state::readout::sketch_query_value(state, query).map_err( + |asap_summary_state::stored_state::readout::Error::Unsupported(reason)| { SummaryExecutorError::Unsupported(reason) }, ) } fn topk_ranked(state: &SummaryState, k: usize) -> Result, SummaryExecutorError> { - asap_physical_operators::stored_state::readout::topk_ranked(state, k).map_err( - |asap_physical_operators::stored_state::readout::Error::Unsupported(reason)| { + asap_summary_state::stored_state::readout::topk_ranked(state, k).map_err( + |asap_summary_state::stored_state::readout::Error::Unsupported(reason)| { SummaryExecutorError::Unsupported(reason) }, ) @@ -1007,7 +1007,7 @@ mod tests { #[test] fn keyed_count_state_follows_planner_family_and_query_readout() { - use asap_physical_operators::summary_kernels::KeyedSumCountAccumulator; + use asap_summary_state::summary_kernels::KeyedSumCountAccumulator; use asap_types::query_plan::ExactReadout; let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); @@ -1201,7 +1201,7 @@ mod tests { fn bound_univmon_merges_panes_for_four_readouts() { use crate::storage_engines::sketch_db::index::SketchEncoding; use crate::storage_engines::types::SerializableToSink; - use asap_physical_operators::summary_kernels::univmon::UnivMonAccumulator; + use asap_summary_state::summary_kernels::univmon::UnivMonAccumulator; use asap_types::query_plan::{MaterializationBinding, PhysicalGrouping}; let index = SketchStore::new(); let fp = asap_types::PolicyFingerprint(701); diff --git a/data_plane/src/storage_engines/sketch_db/backfill/processor.rs b/data_plane/src/storage_engines/sketch_db/backfill/processor.rs index f72eaa481..12f391aaa 100644 --- a/data_plane/src/storage_engines/sketch_db/backfill/processor.rs +++ b/data_plane/src/storage_engines/sketch_db/backfill/processor.rs @@ -980,7 +980,7 @@ mod tests { ), cfg.policy_fingerprint(), ); - let acc = asap_physical_operators::summary_kernels::sum::SumAccumulator::with_sum(1.0); + let acc = asap_summary_state::summary_kernels::sum::SumAccumulator::with_sum(1.0); let live_sid = store .ingest_precompute_for_agg_config( |_metric, attrs, _kind| { diff --git a/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs b/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs index ae89f9fd6..4944be0ef 100644 --- a/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs +++ b/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs @@ -11,7 +11,7 @@ use crate::storage_engines::types::AggregateCore; #[cfg(test)] use crate::tests::accumulator_fixture::create_fixture_accumulator; #[cfg(test)] -use asap_physical_operators::factory::AccumulatorUpdater; +use asap_summary_state::factory::AccumulatorUpdater; #[cfg(test)] use asap_types::aggregation_config::PrecomputeMaterialization; @@ -86,7 +86,7 @@ mod tests { // Replay must preserve each series and rank by the selected update mode. #[test] fn backfilled_topk_preserves_series_and_weight_mode() { - use asap_physical_operators::summary_kernels::{ + use asap_summary_state::summary_kernels::{ CountMinSketchWithHeapAccumulator, CountSketchWithHeapAccumulator, }; for kind in [ diff --git a/data_plane/src/storage_engines/sketch_db/data/mod.rs b/data_plane/src/storage_engines/sketch_db/data/mod.rs index 475fc6b14..5c8ed04e3 100644 --- a/data_plane/src/storage_engines/sketch_db/data/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/data/mod.rs @@ -472,7 +472,7 @@ impl AccuracyBound { /// Per-sample sketch state. Stored as the payload column inside the /// per-sid `SidStoreData` columnar storage. -pub use asap_physical_operators::stored_state::{SketchEncoding, SketchSampleState}; +pub use asap_summary_state::stored_state::{SketchEncoding, SketchSampleState}; /// One materialized series row returned by the query path. Resolved /// from the per-sid intern table at read time. diff --git a/data_plane/src/storage_engines/sketch_db/index/maintenance.rs b/data_plane/src/storage_engines/sketch_db/index/maintenance.rs index a7c2333a5..3991a5b68 100644 --- a/data_plane/src/storage_engines/sketch_db/index/maintenance.rs +++ b/data_plane/src/storage_engines/sketch_db/index/maintenance.rs @@ -795,7 +795,7 @@ impl SketchStore { mod tests { use super::*; use crate::storage_engines::types::PrecomputedOutput; - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; use asap_types::traits::SerializableToSink; #[test] @@ -1303,7 +1303,7 @@ mod tests { ); let assert_complete_output = |store: &SketchStore| { use crate::storage_engines::sketch_db::data::SketchEncoding; - use asap_physical_operators::summary_kernels::DDSketchAccumulator; + use asap_summary_state::summary_kernels::DDSketchAccumulator; let rows = store.query_range(target_sid, 0, 60_000); assert_eq!(rows.len(), 1); assert!(rows[0].series_label_values.is_empty()); diff --git a/data_plane/src/storage_engines/sketch_db/index/mod.rs b/data_plane/src/storage_engines/sketch_db/index/mod.rs index 3c0ed9c27..25270f1c6 100644 --- a/data_plane/src/storage_engines/sketch_db/index/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/index/mod.rs @@ -95,12 +95,18 @@ fn reconstruct_exact_agg( bytes: &[u8], ) -> Option> { use crate::storage_engines::types::AggregateCore; - use asap_physical_operators::summary_kernels::{ + use asap_summary_state::summary_kernels::{ IncreaseAccumulator, KeyedCounterState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, }; match type_name { - "PlannerExactAccumulatorV1" => asap_physical_operators::summary_kernels::exact::ExactAccumulator::deserialize_from_bytes(bytes).ok().map(|a|Box::new(a) as Box), + "PlannerExactAccumulatorV1" => { + asap_summary_state::summary_kernels::exact::ExactAccumulator::deserialize_from_bytes( + bytes, + ) + .ok() + .map(|a| Box::new(a) as Box) + } "SumAccumulator" => SumAccumulator::deserialize_from_bytes(bytes) .ok() .map(|a| Box::new(a) as Box), @@ -113,11 +119,9 @@ fn reconstruct_exact_agg( "MaxAccumulator" => MaxAccumulator::deserialize_from_bytes(bytes) .ok() .map(|a| Box::new(a) as Box), - "KeyedSumCountAccumulator" => { - KeyedSumCountAccumulator::deserialize_from_bytes(bytes) - .ok() - .map(|a| Box::new(a) as Box) - } + "KeyedSumCountAccumulator" => KeyedSumCountAccumulator::deserialize_from_bytes(bytes) + .ok() + .map(|a| Box::new(a) as Box), "KeyedCounterState" => KeyedCounterState::deserialize_from_bytes(bytes) .ok() .map(|a| Box::new(a) as Box), @@ -1712,12 +1716,12 @@ impl SketchStore { // string that had to agree with it. let rollup_value = payload .as_any() - .downcast_ref::() + .downcast_ref::() .map(|acc| (RollupReduction::Min, acc.value)) .or_else(|| { payload .as_any() - .downcast_ref::() + .downcast_ref::() .map(|acc| (RollupReduction::Max, acc.value)) }); let store = self @@ -4627,7 +4631,7 @@ mod tests { #[test] fn precompute_payload_round_trips_through_storage() { - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; let idx = SketchStore::new(); let cfg = SketchConfig::DDSketch { @@ -4668,7 +4672,7 @@ mod tests { #[test] fn query_precomputes_by_agg_returns_data_grouped_by_label_values() { - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; let idx = SketchStore::new(); let cfg = SketchConfig::DDSketch { @@ -4822,7 +4826,7 @@ mod tests { assert!(sketch.as_sketch().is_some()); assert!(sketch.as_exact_agg().is_none()); - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; let exact_agg = AggPayload::ExactAgg(Arc::new(SumAccumulator::with_sum(1.0))); assert!(exact_agg.as_sketch().is_none()); assert!(exact_agg.as_exact_agg().is_some()); @@ -5634,7 +5638,7 @@ mod tests { 850, BTreeMap::new(), (0, 30_000), - Box::new(asap_physical_operators::summary_kernels::SumAccumulator::new()) + Box::new(asap_summary_state::summary_kernels::SumAccumulator::new()) )); // A flusher that captured metadata before completion cannot reopen it. writer.upsert_all(&[stale_record]).unwrap(); @@ -6103,7 +6107,7 @@ mod tests { lv_zone("z0"), (s, s + 30_000), Box::new( - asap_physical_operators::summary_kernels::SumAccumulator::with_sum( + asap_summary_state::summary_kernels::SumAccumulator::with_sum( (i + 1) as f64, ), ), @@ -6322,9 +6326,7 @@ mod tests { lv_zone("z0"), (s, s + 30_000), Box::new( - asap_physical_operators::summary_kernels::SumAccumulator::with_sum( - (i + 1) as f64, - ), + asap_summary_state::summary_kernels::SumAccumulator::with_sum((i + 1) as f64), ), ); } @@ -6388,7 +6390,7 @@ mod tests { lv_zone("z0"), (s, s + 30_000), Box::new({ - let mut acc = asap_physical_operators::summary_kernels::SumAccumulator::new(); + let mut acc = asap_summary_state::summary_kernels::SumAccumulator::new(); acc.update((i + 1) as f64); acc.update(10.0); acc @@ -6582,7 +6584,7 @@ mod tests { #[test] fn planner_exact_families_survive_disk_eviction_and_restart() { use crate::storage_engines::types::{AggregateCore, AggregationType}; - use asap_physical_operators::summary_kernels::exact::ExactAccumulator; + use asap_summary_state::summary_kernels::exact::ExactAccumulator; let kinds = [ AggregationType::Sum, AggregationType::Count, diff --git a/data_plane/src/storage_engines/sketch_db/index/native.rs b/data_plane/src/storage_engines/sketch_db/index/native.rs index 397bb3e63..7488c7860 100644 --- a/data_plane/src/storage_engines/sketch_db/index/native.rs +++ b/data_plane/src/storage_engines/sketch_db/index/native.rs @@ -2,11 +2,13 @@ //! recovery path. Window snapshots are not additive hot-state updates. use super::*; use crate::drivers::ingest::series_resolver::SeriesIdResolver; -use asap_physical_operators::{ +use asap_physical_operators::values::{Batch, Schema, Value}; +use asap_summary_state::{ stored_state::native::{decode_batch, encode_batch}, - values::{Batch, Schema, Value}, AggregateCore, SerializableToSink, }; +#[cfg(test)] +use asap_types::physical_plan_codec::PhysicalPlanCodec; const NATIVE_OUTPUT_TYPE: &str = "NativePhysicalOutputV1"; const NATIVE_OUTPUT_TAG: u8 = persistence::part::encoding_tag::NATIVE_BATCH_V1; @@ -78,7 +80,8 @@ impl NativeSummaryOutput { let [state] = states.as_slice() else { return Err("native stored row requires one summary state".into()); }; - let row_kind = state.get_accumulator_type(); + let row_kind = asap_summary_state::physical::aggregation_type(state.as_ref()) + .map_err(|error| error.to_string())?; if kind.is_some_and(|kind| kind != row_kind) { return Err("native stored rows have different summary families".into()); } @@ -393,7 +396,7 @@ mod tests { use asap_physical_operators::{ dag::{operators::Operator, Limits, RunContext, Scope}, physical_planner::{CompiledPhysicalDag, InputContract, Source}, - summary_kernels::SumAccumulator, + summary_kernels::exact::ExactAccumulator, }; use futures::{executor::block_on, StreamExt}; use planner_types::{ @@ -457,8 +460,8 @@ mod tests { vec![vec![ Value::Utf8("api".into()), Value::Summary { - family, - state: Arc::new(SumAccumulator::new()), + family: family.clone(), + state: Arc::new(ExactAccumulator::new(family, false).unwrap()), }, ]], ) @@ -538,7 +541,7 @@ mod tests { let mut output = PrecomputedOutput::new(0, 60_000, None, source.policy_fingerprint()); output.population_labels = Some(group); output.catalog_generation = Some(generation.clone()); - let mut state = SumAccumulator::new(); + let mut state = asap_summary_state::summary_kernels::SumAccumulator::new(); state.update(value); store .publish_admitted_summary_update( @@ -671,8 +674,9 @@ mod tests { let readout = Operator::readout( batch.schema().clone(), 0, - asap_types::Statistic::Quantile, - HashMap::from([("quantile".into(), "1.0".into())]), + asap_physical_operators::operators::ReadoutQuery::Sketch( + planner_types::post_asap::SketchQuery::Quantile { q: 1.0 }, + ), ) .unwrap(); let values = run(batch, readout); diff --git a/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs b/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs index 4788aaca0..fd8e8b63e 100644 --- a/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs +++ b/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs @@ -194,7 +194,7 @@ pub fn warn_if_retention_inverted( mod tests { use super::*; use crate::storage_engines::types::{AggregationType, InstalledPrecomputePlan}; - use asap_physical_operators::summary_kernels::SumAccumulator; + use asap_summary_state::summary_kernels::SumAccumulator; use asap_types::aggregation_config::PrecomputeMaterialization; use asap_types::enums::WindowKind; use asap_types::KeyByLabelNames; diff --git a/data_plane/src/storage_engines/sketch_db/persistence/immutable_output.rs b/data_plane/src/storage_engines/sketch_db/persistence/immutable_output.rs index 76025600b..e4c5ee014 100644 --- a/data_plane/src/storage_engines/sketch_db/persistence/immutable_output.rs +++ b/data_plane/src/storage_engines/sketch_db/persistence/immutable_output.rs @@ -1,5 +1,6 @@ //! Crash-resumable publication of one immutable output window. The SID sidecar //! reserves the existing part ID before writing; no second payload store is used. +use asap_summary_state::codec::KeyCodec; use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; diff --git a/data_plane/src/storage_engines/sketch_db/persistence/part.rs b/data_plane/src/storage_engines/sketch_db/persistence/part.rs index 6eae5f043..526ef99fb 100644 --- a/data_plane/src/storage_engines/sketch_db/persistence/part.rs +++ b/data_plane/src/storage_engines/sketch_db/persistence/part.rs @@ -62,6 +62,7 @@ //! 8-byte alignment) and skips a bunch of the micro-optimizations in //! the design doc. v1 priorities are correctness, round-trip, and //! passing tests — not beating Prometheus TSDB. +use asap_summary_state::codec::KeyCodec; use std::fs::{self, File, OpenOptions}; use std::io::{BufWriter, Read, Write}; diff --git a/data_plane/src/storage_engines/sketch_db/query/decoders.rs b/data_plane/src/storage_engines/sketch_db/query/decoders.rs index d208cb57d..3e92b2814 100644 --- a/data_plane/src/storage_engines/sketch_db/query/decoders.rs +++ b/data_plane/src/storage_engines/sketch_db/query/decoders.rs @@ -1,2 +1,2 @@ //! Storage uses the Planner-owned state implementation. -pub use asap_physical_operators::stored_state::decoders::*; +pub use asap_summary_state::stored_state::decoders::*; diff --git a/data_plane/src/storage_engines/sketch_db/query/delta_apply.rs b/data_plane/src/storage_engines/sketch_db/query/delta_apply.rs index 09869df2e..f4be6f17e 100644 --- a/data_plane/src/storage_engines/sketch_db/query/delta_apply.rs +++ b/data_plane/src/storage_engines/sketch_db/query/delta_apply.rs @@ -1,10 +1,10 @@ //! Storage uses the Planner-owned state implementation. -pub use asap_physical_operators::stored_state::delta_apply::*; +pub use asap_summary_state::stored_state::delta_apply::*; #[cfg(test)] mod tests { use super::*; - use asap_physical_operators::stored_state::{SketchEncoding, SketchSampleState}; + use asap_summary_state::stored_state::{SketchEncoding, SketchSampleState}; #[test] fn native_batches_are_not_legacy_sketch_frames() { diff --git a/data_plane/src/storage_engines/types/mod.rs b/data_plane/src/storage_engines/types/mod.rs index f126e0db2..235b006ee 100644 --- a/data_plane/src/storage_engines/types/mod.rs +++ b/data_plane/src/storage_engines/types/mod.rs @@ -30,4 +30,4 @@ pub use crate::query_engines::routing::{ QueryOperatorShape, RoutingTarget, }; -pub use asap_physical_operators::{traits::*, KeyByLabelValues, Measurement}; +pub use asap_summary_state::{traits::*, KeyByLabelValues, Measurement}; diff --git a/data_plane/src/tests/accumulator_fixture.rs b/data_plane/src/tests/accumulator_fixture.rs index e8346cba4..9a2d92686 100644 --- a/data_plane/src/tests/accumulator_fixture.rs +++ b/data_plane/src/tests/accumulator_fixture.rs @@ -1,6 +1,6 @@ //! Config fixtures for backend integration tests; production binds Planner payloads. -use asap_physical_operators::factory::*; -use asap_physical_operators::{AggregateCore, AggregationType}; +use asap_summary_state::factory::*; +use asap_summary_state::{AggregateCore, AggregationType}; use asap_types::{accumulator_spec::cms_params, PrecomputeMaterialization}; use planner_types::post_asap::{ExactKind, SketchAlgorithm, SketchParams, SummaryFamilyType}; #[cfg(test)] @@ -232,7 +232,7 @@ pub fn create_fixture_accumulator( else { unreachable!("validated UnivMon family parameters") }; - asap_physical_operators::factory::create_planner_accumulator( + asap_summary_state::factory::create_planner_accumulator( &spec.family, &planner_types::post_asap::SummaryUpdate::column( planner_types::pre_asap::ColumnRef::SampleValue, @@ -248,7 +248,7 @@ pub fn create_fixture_accumulator( let SketchParams::Hll { precision } = kind.params() else { unreachable!("validated HLL family parameters") }; - asap_physical_operators::factory::create_planner_accumulator( + asap_summary_state::factory::create_planner_accumulator( &spec.family, &planner_types::post_asap::SummaryUpdate::column( planner_types::pre_asap::ColumnRef::SampleValue, diff --git a/data_plane/src/tests/trait_design_tests.rs b/data_plane/src/tests/trait_design_tests.rs index f8f5d6f04..8e2fafae6 100644 --- a/data_plane/src/tests/trait_design_tests.rs +++ b/data_plane/src/tests/trait_design_tests.rs @@ -2,7 +2,7 @@ use crate::storage_engines::types::{ KeyByLabelValues, MultipleSubpopulationAggregate, SingleSubpopulationAggregate, }; -use asap_physical_operators::summary_kernels::{KeyedSumCountAccumulator, SumAccumulator}; +use asap_summary_state::summary_kernels::{KeyedSumCountAccumulator, SumAccumulator}; use asap_types::Statistic; #[test] diff --git a/data_plane/tests/clickhouse_differential_e2e.rs b/data_plane/tests/clickhouse_differential_e2e.rs index 00ecd0bdf..2845169b1 100644 --- a/data_plane/tests/clickhouse_differential_e2e.rs +++ b/data_plane/tests/clickhouse_differential_e2e.rs @@ -1,6 +1,7 @@ //! Optional real-ClickHouse protocol and Grafana smoke coverage. //! //! Set `CLICKHOUSE_URL` (for example `http://127.0.0.1:8123`) to run it. +use asap_types::physical_plan_codec::PhysicalPlanCodec; #[path = "support/empty_physical_plan.rs"] mod empty_physical_plan; diff --git a/data_plane/tests/edge_sketch_codec.rs b/data_plane/tests/edge_sketch_codec.rs index 03f4e40f1..f4d758b98 100644 --- a/data_plane/tests/edge_sketch_codec.rs +++ b/data_plane/tests/edge_sketch_codec.rs @@ -35,7 +35,7 @@ fn ddsketch_bare_state_is_rejected_and_envelope_supports_query_readout() { let bare = prost::Message::encode_to_vec(&state); assert!(asap_sketch_codec::reconstruct_ddsketch(&bare).is_err()); let (decoded, _) = asap_sketch_codec::reconstruct_ddsketch(&envelope).unwrap(); - let accumulator = asap_physical_operators::summary_kernels::DDSketchAccumulator { + let accumulator = asap_summary_state::summary_kernels::DDSketchAccumulator { inner: decoded, sample_p: 1.0, }; @@ -62,7 +62,7 @@ fn kll_envelope_keeps_level_layout_for_backend_readout() { assert_eq!(state.k, 200); assert_eq!(state.items.len(), 50); let snapshot_bytes = bytes; - let accumulator = asap_physical_operators::summary_kernels::DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&snapshot_bytes).unwrap(); + let accumulator = asap_summary_state::summary_kernels::DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&snapshot_bytes).unwrap(); assert!(accumulator.get_quantile(0.5).is_finite()); } diff --git a/data_plane/tests/support/issue_701_702_process.rs b/data_plane/tests/support/issue_701_702_process.rs index 0c555bcaf..e617e575b 100644 --- a/data_plane/tests/support/issue_701_702_process.rs +++ b/data_plane/tests/support/issue_701_702_process.rs @@ -1,5 +1,6 @@ //! Issue workloads execute their selected Planner DAG on the production HTTP path. use super::*; +use asap_types::physical_plan_codec::PhysicalPlanCodec; use control_plane::physical::{ compiler::{ BackendLocalPlanningInput, DeploymentPlanCompiler, BACKEND_REVISION, PLANNER_REVISION, diff --git a/data_plane/tests/support/univmon_erp_process.rs b/data_plane/tests/support/univmon_erp_process.rs index 52a43ae79..73330a7b8 100644 --- a/data_plane/tests/support/univmon_erp_process.rs +++ b/data_plane/tests/support/univmon_erp_process.rs @@ -1,5 +1,5 @@ use super::*; -use asap_physical_operators::summary_kernels::univmon::UnivMonAccumulator; +use asap_summary_state::summary_kernels::univmon::UnivMonAccumulator; use control_plane::physical::erp::ErpShapeObserver; use data_plane::storage_engines::types::{AggregateCore, SerializableToSink}; diff --git a/docs/design_docs/physical-operators.md b/docs/design_docs/physical-operators.md index c403684da..a9f514763 100644 --- a/docs/design_docs/physical-operators.md +++ b/docs/design_docs/physical-operators.md @@ -6,10 +6,21 @@ The shared physical operator library lives in ASAPPlanner, alongside post-ASAP IR and physical lowering. Its canonical architecture and acceptance contract are in [the Planner design](https://github.com/ProjectASAP/ASAPPlanner/blob/feat/shared-physical-operators/docs/design_docs/physical-planning-and-deployment.md). -ASAPQuery-backend depends on `asap-physical-operators`, `asap_sketch_codec` and -Planner IR at the same immutable revision. It does not own a second copy of the -runtime or mathematical kernels. A new IR operation and its implementation can -be changed and tested together in Planner. +ASAPQuery-backend depends on `asap-physical-operators` and Planner IR at the +same immutable revision. It does not own a second copy of the runtime or of the +in-memory physical kernels. A new IR operation and its implementation can be +changed and tested together in Planner. + +Planner does not own storage formats. The backend owns them: + +- `crates/asap_summary_state`: the summary kernels that ingest and the sketch + store keep, their stored byte encodings, delta reconstruction, native batch + frames, and conversion to Planner physical states where stored state enters + or leaves a physical DAG. +- `crates/asap_sketch_codec`: the sketchlib `SketchEnvelope` codec. +- `asap_types::physical_plan_codec`: the versioned envelope of persisted + physical plans and candidates. +- `asap_types::semantic_fragment`: the stored-definition semantic identity. The query and precompute integration PRs both use the independent ASAP DAG runtime. Deployment code binds input sources, storage, ingestion windows, publication and protocol outputs. Execution