diff --git a/Cargo.lock b/Cargo.lock index acc50d509..5e0173958 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=bccc837f5caf21c888b2cce73c64a1be283b679a#bccc837f5caf21c888b2cce73c64a1be283b679a" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=4b0839ce1733aab8231eb90944c876063a4551f9#4b0839ce1733aab8231eb90944c876063a4551f9" 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=bccc837f5caf21c888b2cce73c64a1be283b679a#bccc837f5caf21c888b2cce73c64a1be283b679a" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=4b0839ce1733aab8231eb90944c876063a4551f9#4b0839ce1733aab8231eb90944c876063a4551f9" 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=bccc837f5caf21c888b2cce73c64a1be283b679a#bccc837f5caf21c888b2cce73c64a1be283b679a" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=4b0839ce1733aab8231eb90944c876063a4551f9#4b0839ce1733aab8231eb90944c876063a4551f9" dependencies = [ "asap-sql-function-catalog", "asap-types", @@ -393,15 +393,35 @@ dependencies = [ "serde_json", ] +[[package]] +name = "asap-physical-operators" +version = "0.1.0" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=4b0839ce1733aab8231eb90944c876063a4551f9#4b0839ce1733aab8231eb90944c876063a4551f9" +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=bccc837f5caf21c888b2cce73c64a1be283b679a#bccc837f5caf21c888b2cce73c64a1be283b679a" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=4b0839ce1733aab8231eb90944c876063a4551f9#4b0839ce1733aab8231eb90944c876063a4551f9" [[package]] name = "asap-types" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=bccc837f5caf21c888b2cce73c64a1be283b679a#bccc837f5caf21c888b2cce73c64a1be283b679a" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=4b0839ce1733aab8231eb90944c876063a4551f9#4b0839ce1733aab8231eb90944c876063a4551f9" dependencies = [ "serde", "serde_json", @@ -423,16 +443,18 @@ dependencies = [ [[package]] name = "asap_sketch_codec" version = "0.1.0" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=4b0839ce1733aab8231eb90944c876063a4551f9#4b0839ce1733aab8231eb90944c876063a4551f9" dependencies = [ - "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?branch=main)", + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", "prost", ] [[package]] name = "asap_sketchlib" version = "0.3.0" -source = "git+https://github.com/ProjectASAP/asap_sketchlib?branch=main#026cd18c7b8c23ae6c46d4d683151ba562b8cd3a" +source = "git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369#5f03ccbd798ed5fec62bdd839bcb331123cab369" dependencies = [ + "bincode", "bytes", "prost", "rand 0.9.5", @@ -468,6 +490,7 @@ version = "0.1.0" dependencies = [ "anyhow", "asap-aware-mapping", + "asap-physical-operators", "asap-types", "base64 0.21.7", "clap", @@ -946,6 +969,7 @@ dependencies = [ "asap-aware-mapping", "asap-frontend-promql", "asap-frontend-sql", + "asap-physical-operators", "asap-types", "asap_types", "axum", @@ -1156,10 +1180,11 @@ dependencies = [ "arrow", "asap-aware-mapping", "asap-frontend-promql", + "asap-physical-operators", "asap-types", "asap_otel_proto", "asap_sketch_codec", - "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?branch=main)", + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", "asap_types", "async-trait", "axum", diff --git a/Cargo.toml b/Cargo.toml index b5592550e..9040dbf29 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,7 +3,6 @@ resolver = "2" members = [ "crates/asap_otel_proto", "crates/asap_types", - "crates/asap_sketch_codec", "data_plane", "control_plane", ] @@ -15,10 +14,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 = "bccc837f5caf21c888b2cce73c64a1be283b679a" } -asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "bccc837f5caf21c888b2cce73c64a1be283b679a" } -asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "bccc837f5caf21c888b2cce73c64a1be283b679a" } -asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "bccc837f5caf21c888b2cce73c64a1be283b679a" } +planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "4b0839ce1733aab8231eb90944c876063a4551f9" } +asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "4b0839ce1733aab8231eb90944c876063a4551f9" } +asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "4b0839ce1733aab8231eb90944c876063a4551f9" } +asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "4b0839ce1733aab8231eb90944c876063a4551f9" } # Shared external deps (used by 2+ crates) serde = { version = "1.0", features = ["derive"] } @@ -38,6 +37,8 @@ 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 = "4b0839ce1733aab8231eb90944c876063a4551f9" } +asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "4b0839ce1733aab8231eb90944c876063a4551f9" } 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/Cargo.toml b/control_plane/Cargo.toml index ac83339b3..510d3aa9b 100644 --- a/control_plane/Cargo.toml +++ b/control_plane/Cargo.toml @@ -12,6 +12,7 @@ name = "control_plane" path = "src/main.rs" [dependencies] +asap-physical-operators.workspace = true tokio = { version = "1", features = ["full"] } axum = { version = "0.7", features = ["ws"] } futures-util = "0.3" diff --git a/control_plane/examples/calibration_candidates.rs b/control_plane/examples/calibration_candidates.rs index 5c95608bc..928dc1475 100644 --- a/control_plane/examples/calibration_candidates.rs +++ b/control_plane/examples/calibration_candidates.rs @@ -37,6 +37,7 @@ fn planner_forest(queries: &[control_plane::physical::compiler::QueryCompilation vec![lhs, rhs], json!({"operator_debug":format!("{operator:?}"),"timing_debug":format!("{timing:?}")}), ), + SummaryExpr::ValueOperation { child, operation, @@ -73,11 +74,11 @@ fn planner_forest(queries: &[control_plane::physical::compiler::QueryCompilation right, kind, pred, - pruning, + .. } => ( "RelationalJoin", vec![left, right], - json!({"kind_debug":format!("{kind:?}"),"predicate_debug":format!("{pred:?}"), "pruning_debug":format!("{pruning:?}")}), + json!({"kind_debug":format!("{kind:?}"),"predicate_debug":format!("{pred:?}")}), ), SummaryExpr::SummarySubtract { left, right } => { ("SummarySubtract", vec![left, right], json!({})) diff --git a/control_plane/src/emit/mod.rs b/control_plane/src/emit/mod.rs index bf51c16e5..a56536256 100644 --- a/control_plane/src/emit/mod.rs +++ b/control_plane/src/emit/mod.rs @@ -56,6 +56,7 @@ fn extract_from_node(node: &Rc) -> Option { SummaryExpr::SummaryEstimate { summary_input, .. } => extract_from_node(summary_input), SummaryExpr::SummaryMerge { children, .. } => children.iter().find_map(extract_from_node), SummaryExpr::ValueOperation { child, .. } => extract_from_node(child), + // Not surfaced by any `Bind*` path yet (gated on rules that // haven't landed — see `deployment_expr.rs`'s module docs). SummaryExpr::BinaryOp { .. } diff --git a/control_plane/src/physical/compiler.rs b/control_plane/src/physical/compiler.rs index cf9d6ea56..3bcbaf2c1 100644 --- a/control_plane/src/physical/compiler.rs +++ b/control_plane/src/physical/compiler.rs @@ -2094,6 +2094,7 @@ fn summary_agg_metric(node: &SummaryNode) -> Option { } } SummaryExpr::SummaryAgg { child, .. } => walk(child, metrics), + SummaryExpr::ValueOperation { child, .. } => walk(child, metrics), SummaryExpr::SummaryEstimate { summary_input, .. } => walk(summary_input, metrics), SummaryExpr::SummaryMerge { children, .. } => { @@ -2162,11 +2163,12 @@ fn observed_population_matches_root( }; if measures.is_empty() || !measures.iter().all(|intent| { - matches!( - intent, - AggIntent::FrequencyL2 { col: None, .. } - | AggIntent::FrequencyEntropy { col: None, .. } - ) || matches!(intent, AggIntent::Cardinality { cols, .. } if cols.is_empty()) + matches!(intent, AggIntent::Cardinality { cols, .. } if cols.is_empty()) + || matches!( + intent, + AggIntent::FrequencyL2 { col: None, .. } + | AggIntent::FrequencyEntropy { col: None, .. } + ) }) { return false; @@ -2369,6 +2371,7 @@ fn requires_exact_erp_fallback( walk(left, out); walk(right, out); } + SummaryExpr::KeepPreAsap(_) => {} } } @@ -3601,16 +3604,14 @@ fn collect_selected_materializations( } match &node.expr { SummaryExpr::RelationalJoin { - right: candidates, left: values, + right: candidates, kind: planner_types::pre_asap::JoinKind::Semi, pruning: Some(_), .. } => { walk(candidates, readout, composable, grouping.clone(), selected)?; - // In a hybrid TopK, the sketch is only a candidate-membership - // sidecar. Prometheus owns the authoritative value subtree; - // provisioning local exact state here duplicates that work. + // Explicit external authoritative values do not need duplicate local state. if !composable { walk(values, readout, composable, grouping.clone(), selected)?; } @@ -4479,6 +4480,7 @@ pub(crate) mod tests { } } + // The explicit rate-value frontier cannot be rebound as raw counter deltas. #[test] fn unsupported_rate_heap_uses_explicit_exact_route() { let query = "topk(2, sum by (job) (rate(m[1m])))"; @@ -4506,6 +4508,7 @@ pub(crate) mod tests { assert!(plan.precompute_plan.materializations.is_empty()); } + // Native exact values remain explicit; unsupported heaps publish no stored state. #[test] fn hybrid_rate_topk_preserves_the_original_exact_subquery() { let query = "topk(2, sum by (job) (rate(m[1m])))"; @@ -4530,7 +4533,7 @@ pub(crate) mod tests { assert!(matches!( &entry.nodes[&entry.root], QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { k: 2, .. }, + operator: ResidualQueryOperator::Limit { n: 2, .. }, .. } )); @@ -5753,6 +5756,11 @@ pub(crate) mod tests { let selected = match &selected_root.expr { SummaryExpr::SummaryEstimate { summary_input, .. } => summary_input.clone(), SummaryExpr::SummaryAgg { .. } => selected_root.clone(), + SummaryExpr::ValueOperation { + child, + operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, + .. + } => child.clone(), _ => panic!("expected maintained aggregate fixture"), }; request.queries[0].selected_plan_root = Rc::new(SummaryNode { @@ -5956,15 +5964,6 @@ pub(crate) mod tests { let (right, _) = right.into_physical_compilation_request().unwrap(); let left = request.queries[0].selected_plan_root.clone(); let right = right.queries[0].selected_plan_root.clone(); - let right = Rc::new(SummaryNode { - expr: SummaryExpr::ValueOperation { - timing: planner_types::post_asap::ExecutionTiming::QueryTime, - operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, - child: right.clone(), - }, - schema: right.schema.clone(), - guarantee: None, - }); request.queries[0].selected_plan_root = Rc::new(SummaryNode { expr: SummaryExpr::BinaryOp { timing: planner_types::post_asap::ExecutionTiming::QueryTime, diff --git a/control_plane/src/physical/plan_dot.rs b/control_plane/src/physical/plan_dot.rs index 19d602413..e0753c716 100644 --- a/control_plane/src/physical/plan_dot.rs +++ b/control_plane/src/physical/plan_dot.rs @@ -139,6 +139,16 @@ fn escape(value: &str) -> String { fn query_node_label(node: &QueryPlanNode) -> String { match node { + QueryPlanNode::PhysicalFragment { dag, .. } => { + asap_physical_operators::physical_planner::CompiledPhysicalDag::decode(dag) + .map(|plan| { + format!( + "Physical\n{}", + plan.operator_name(plan.roots()[0]).unwrap_or("Input") + ) + }) + .unwrap_or_else(|_| "Invalid physical DAG".into()) + } QueryPlanNode::RelationalJoin { .. } => "RelationalJoin".into(), QueryPlanNode::Relational { .. } => "Relational".into(), QueryPlanNode::Logical { operator, .. } => format!("Logical\n{}", residual_label(operator)), @@ -154,7 +164,6 @@ fn query_node_label(node: &QueryPlanNode) -> String { QueryPlanNode::SummaryEstimate { query, .. } => format!("SummaryEstimate\n{query:?}"), QueryPlanNode::ExactReadout { readout, .. } => format!("ExactReadout\n{readout:?}"), QueryPlanNode::SummaryMerge { .. } => "SummaryMerge".into(), - QueryPlanNode::MembershipFilter { .. } => "MembershipFilter".into(), QueryPlanNode::ExternalExact { .. } => "ExternalExact".into(), QueryPlanNode::ExactFallback { reason } => format!("ExactFallback\n{reason}"), } @@ -169,7 +178,7 @@ fn residual_label(operator: &ResidualQueryOperator) -> &'static str { ResidualQueryOperator::UnaryNegate => "UnaryNegate", ResidualQueryOperator::VectorToScalar => "VectorToScalar", ResidualQueryOperator::Aggregate { .. } => "Aggregate", - ResidualQueryOperator::TopKSelection { .. } => "TopKSelection", + ResidualQueryOperator::Limit { .. } => "Limit", ResidualQueryOperator::Binary { .. } => "Binary", ResidualQueryOperator::Temporal { .. } => "Temporal", ResidualQueryOperator::Sort { .. } => "Sort", diff --git a/control_plane/src/physical/post_asap/cost_model.rs b/control_plane/src/physical/post_asap/cost_model.rs index 5fa6c75d9..9610ed0a1 100644 --- a/control_plane/src/physical/post_asap/cost_model.rs +++ b/control_plane/src/physical/post_asap/cost_model.rs @@ -174,9 +174,20 @@ impl ControlPlaneCostModel { &self, candidate: &ReplacementSubDAG, ) -> Option { + if let Replacement::Rewrite(root) = &candidate.replacement { + // Price a semantic rewrite through its executable summary candidates, + // in the same retained-state units as a direct summary candidate. + // Unknown realizations remain uncosted rather than receiving zero. + use asap_aware_mapping::{ReplacementStrategy, SketchAlgorithmStrategy, TargetSubDAG}; + return SketchAlgorithmStrategy::new(self) + .replacements(&TargetSubDAG::new(root)) + .iter() + .filter(|candidate| matches!(candidate.replacement, Replacement::Summary(_))) + .filter_map(|candidate| self.candidate_cost_estimate(candidate)) + .min_by(|a, b| a.value.total_cmp(&b.value)); + } let Replacement::Summary(root) = &candidate.replacement else { - // Exact compositions have a separate measured rate model. Raw - // rewrites have no retained-state estimate in this model. + // Exact compositions have a separate measured rate model. return None; }; let dag = planner_types::post_asap::compile_executable_dag(root).ok()?; @@ -510,6 +521,23 @@ fn intent_accuracy(intent: &AggIntent) -> AccuracyTarget { } impl CostModel for ControlPlaneCostModel { + fn summary_support_evidence( + &self, + summary: &planner_types::post_asap::SummaryNode, + ) -> Option { + use planner_types::post_asap::{NonNegativeWeightProof, WeightDomain}; + let dag = + planner_types::post_asap::compile_executable_dag(&std::rc::Rc::new(summary.clone())) + .ok()?; + // Counter-weighted heaps now consume explicit rate values. The raw + // ingestion adapter cannot bind that frontier as counter deltas. + let requires_rate_values = dag.nodes.iter().any(|node| matches!(&node.payload, + ExecutableOperatorPayload::SummaryAgg { input, family: SummaryFamilyType::Sketch(kind, _), .. } + if matches!(kind.algorithm(), SketchAlgorithm::CmsWithHeap | SketchAlgorithm::CountSketchWithHeap) + && matches!(input.weight_domain, WeightDomain::NonNegative { proof: NonNegativeWeightProof::ResetAwareCounterDerivative }))); + requires_rate_values.then_some(false) + } + fn candidate_cost( &self, candidate: &ReplacementSubDAG, @@ -869,6 +897,13 @@ impl ForcedFamilyCostModel { } impl CostModel for ForcedFamilyCostModel { + fn summary_support_evidence( + &self, + summary: &planner_types::post_asap::SummaryNode, + ) -> Option { + self.inner.summary_support_evidence(summary) + } + fn candidate_cost( &self, candidate: &ReplacementSubDAG, diff --git a/control_plane/src/physical/post_asap/tests.rs b/control_plane/src/physical/post_asap/tests.rs index 344b60534..21bc52fff 100644 --- a/control_plane/src/physical/post_asap/tests.rs +++ b/control_plane/src/physical/post_asap/tests.rs @@ -18,6 +18,24 @@ use crate::types::AccuracyTarget; use planner_types::pre_asap::{AggIntent, QueryExpr, Reduction, Schema, Source}; use planner_types::pre_asap::{Column, DataType}; +// Query outputs are values; inspect the explicitly retained accumulator below them. +fn exact_query_state(node: &SummaryNode) -> &SummaryNode { + let SummaryExpr::ValueOperation { + child, + operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, + } = &node.expr + else { + panic!("query must finalize its exact state: {:?}", node.expr) + }; + assert!(node + .schema + .fields + .iter() + .all(|field| matches!(field.dtype, SummaryFamilyType::Plain(_)))); + child +} + fn sketch_family(kind: SketchAlgorithm, params: SketchParams) -> SummaryFamilyType { SummaryFamilyType::Sketch(SketchKind::new(kind, params), GroupingStrategy::default()) } @@ -116,6 +134,7 @@ fn node_is_archive(node: &Rc) -> bool { SummaryExpr::SummaryJoin { outer, inner, .. } => { node_is_archive(outer) || node_is_archive(inner) } + SummaryExpr::SummarySubtract { left, right } | SummaryExpr::RelationalJoin { left, right, .. } | SummaryExpr::BinaryOp { @@ -342,7 +361,7 @@ fn uncertified_hll_keeps_exact_execution() { #[test] fn sum_now_binds_to_exact_agg_after_pr_6_followup() { - // `AggIntent::Sum` binds to a bare `SummaryAgg` with `summary: + // `AggIntent::Sum` binds to a readout over `SummaryAgg` with `summary: // SummaryKind::Sum` and no `SummaryEstimate` wrapper (the partial // state *is* the value — see `asap_aware_mapping::replacement`'s module docs). The // old locally-defined `PhysicalExpr::ExactAgg { agg_type, .. }` @@ -360,16 +379,18 @@ fn sum_now_binds_to_exact_agg_after_pr_6_followup() { }; let bound = bind_query_expr(&expr, AccuracyTarget::Exact).expect("no error"); match bound { - PhysicalExpr::Committed(PostAsapPlan::Summary(node)) => match &node.expr { - SummaryExpr::SummaryAgg { family, .. } => { - assert_eq!( - family, - &SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum), - "Sum should bind to SummaryAgg(Sum)" - ); + PhysicalExpr::Committed(PostAsapPlan::Summary(node)) => { + match &exact_query_state(&node).expr { + SummaryExpr::SummaryAgg { family, .. } => { + assert_eq!( + family, + &SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum), + "Sum should bind to SummaryAgg(Sum)" + ); + } + other => panic!("expected Sum state below the query readout, got {other:?}"), } - other => panic!("expected bare SummaryAgg(Sum), got {other:?}"), - }, + } other => panic!("expected Committed(Summary(_)), got {other:?}"), } } @@ -440,7 +461,7 @@ fn phase_b_pattern_only_temporal_quantile_binds_to_sketch() { /// `ONLY_TEMPORAL` — `sum_over_time(m[5m])` (and the count/avg/min/max /// variants that legacy `single_query.rs` accepts). /// -/// Control plane path: `Aggregate{Sum}` over `Window` → binds to a bare +/// Control plane path: `Aggregate{Sum}` over `Window` → binds to a readout over /// `SummaryAgg{summary: SummaryKind::Sum}` (an exact mergeable /// accumulator — see `sum_now_binds_to_exact_agg_after_pr_6_followup`'s /// doc comment for the `ExactAgg` → `SummaryAgg` unification). @@ -455,15 +476,17 @@ fn phase_b_pattern_only_temporal_sum_binds_to_exact_agg() { }; let bound = bind_query_expr(&expr, AccuracyTarget::Epsilon(0.01)).unwrap(); match bound { - PhysicalExpr::Committed(PostAsapPlan::Summary(node)) => match &node.expr { - SummaryExpr::SummaryAgg { family, .. } => { - assert_eq!( - family, - &SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum) - ); + PhysicalExpr::Committed(PostAsapPlan::Summary(node)) => { + match &exact_query_state(&node).expr { + SummaryExpr::SummaryAgg { family, .. } => { + assert_eq!( + family, + &SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum) + ); + } + other => panic!("expected SummaryAgg(Sum), got {other:?}"), } - other => panic!("expected SummaryAgg(Sum), got {other:?}"), - }, + } other => panic!("expected Committed(Summary(_)), got {other:?}"), } } @@ -483,22 +506,24 @@ fn phase_b_pattern_only_spatial_aggregate_binds_to_grouped_sum() { }; let bound = bind_query_expr(&expr, AccuracyTarget::Epsilon(0.01)).unwrap(); match bound { - PhysicalExpr::Committed(PostAsapPlan::Summary(node)) => match &node.expr { - SummaryExpr::SummaryAgg { - family, reduction, .. - } => { - assert_eq!( - family, - &SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum) - ); - assert_eq!( - reduction.group_keys().map(|k| k.keys()), - Some(&[1][..]), - "Sum reduction must retain the group-by column" - ); + PhysicalExpr::Committed(PostAsapPlan::Summary(node)) => { + match &exact_query_state(&node).expr { + SummaryExpr::SummaryAgg { + family, reduction, .. + } => { + assert_eq!( + family, + &SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum) + ); + assert_eq!( + reduction.group_keys().map(|k| k.keys()), + Some(&[1][..]), + "Sum reduction must retain the group-by column" + ); + } + other => panic!("expected SummaryAgg(Sum, by=[1]), got {other:?}"), } - other => panic!("expected SummaryAgg(Sum, by=[1]), got {other:?}"), - }, + } other => panic!("expected Committed(Summary(_)), got {other:?}"), } } @@ -516,18 +541,20 @@ fn phase_b_pattern_temporal_and_spatial_combined_preserves_rate() { }; let bound = bind_query_expr(&expr, AccuracyTarget::Epsilon(0.01)).unwrap(); match bound { - PhysicalExpr::Committed(PostAsapPlan::Summary(node)) => match &node.expr { - SummaryExpr::SummaryAgg { - family, reduction, .. - } => { - assert_eq!( - family, - &SummaryFamilyType::ExactAggregate(ExactKind::Rate, ExactParams::Rate) - ); - assert_eq!(reduction.group_keys().map(|k| k.keys()), Some(&[1][..])); + PhysicalExpr::Committed(PostAsapPlan::Summary(node)) => { + match &exact_query_state(&node).expr { + SummaryExpr::SummaryAgg { + family, reduction, .. + } => { + assert_eq!( + family, + &SummaryFamilyType::ExactAggregate(ExactKind::Rate, ExactParams::Rate) + ); + assert_eq!(reduction.group_keys().map(|k| k.keys()), Some(&[1][..])); + } + other => panic!("expected SummaryAgg(Rate, by=[1]), got {other:?}"), } - other => panic!("expected SummaryAgg(Rate, by=[1]), got {other:?}"), - }, + } other => panic!("expected Committed(Summary(_)), got {other:?}"), } } diff --git a/control_plane/src/planner_selection.rs b/control_plane/src/planner_selection.rs index 37e17fd7f..81ead342b 100644 --- a/control_plane/src/planner_selection.rs +++ b/control_plane/src/planner_selection.rs @@ -475,7 +475,7 @@ fn select_workload_impl( .iter() .map(|(id, root)| { selection - .assemble_selected_dag(root) + .assemble_selected_query(root) .map_err(|error| SelectionError::Workload(error.to_string()))? .map(|node| (*id, node)) .ok_or_else(|| SelectionError::Workload(format!("missing query root {id}"))) @@ -509,7 +509,7 @@ pub fn select_query_with_models( ); space .global_selection(cost_model) - .assemble_selected_dag(&space.roots[0].1) + .assemble_selected_query(&space.roots[0].1) .map_err(|error| SelectionError::Workload(error.to_string()))? .ok_or(SelectionError::NoLegalCandidate) } @@ -744,7 +744,15 @@ mod workload_tests { } => child, _ => lhs, }; - assert!(Rc::ptr_eq(&roots[0].1, shared)); + let SummaryExpr::ValueOperation { + child: standalone, + operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, + .. + } = &roots[0].1.expr + else { + panic!("standalone query must finalize its shared state") + }; + assert!(Rc::ptr_eq(standalone, shared)); } // The registered set is exactly the four strategies this deployment @@ -901,6 +909,12 @@ mod workload_tests { SummaryExpr::SummaryEstimate { summary_input, .. } => { is_summary(summary_input) } + SummaryExpr::ValueOperation { + child, + operation: + planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, + .. + } => is_summary(child), _ => false, } } diff --git a/control_plane/src/query_plan.rs b/control_plane/src/query_plan.rs index 838966c05..a1e70c590 100644 --- a/control_plane/src/query_plan.rs +++ b/control_plane/src/query_plan.rs @@ -136,6 +136,82 @@ where Ok(entry) } +fn compile_native_fragment( + root: &Rc, + query_inputs: &[QueryNodeId], +) -> Result { + use asap_physical_operators::physical_planner::{compile, InputContract}; + use planner_types::post_asap::{compile_executable_dag, EdgeRole}; + let invalid = |e: String| QueryPlanError::Invalid(e); + let dag = compile_executable_dag(root).map_err(|e| invalid(e.to_string()))?; + let mut edges = dag + .edges + .iter() + .filter(|e| e.consumer == dag.root) + .collect::>(); + edges.sort_by_key(|e| match e.role { + EdgeRole::Left => 0, + EdgeRole::Input => 1, + EdgeRole::Right => 2, + }); + if edges.len() != query_inputs.len() { + return Err(invalid("physical frontier arity mismatch".into())); + } + let bindings = edges + .iter() + .zip(query_inputs) + .map(|(edge, &id)| (u64::from(edge.producer.0), id)) + .collect::>(); + let contracts = edges + .iter() + .map(|edge| { + ( + u64::from(edge.producer.0), + InputContract::bounded(std::sync::Arc::new(edge.intermediate_schema.clone())), + ) + }) + .collect(); + let physical = + compile(&dag, contracts, &[u64::from(dag.root.0)]).map_err(|e| invalid(e.to_string()))?; + let row_input = physical + .input_contracts() + .position(|(id, _)| bindings[&id] == query_inputs[0]) + .unwrap(); + let pruning = if let SummaryExpr::RelationalJoin { + left, + right, + pred, + pruning: Some(completeness), + .. + } = &root.expr + { + Some(asap_types::query_plan::PruningInputContract { + candidate_input: physical + .input_contracts() + .position(|(id, _)| bindings[&id] == query_inputs[1]) + .unwrap(), + keys: asap_physical_operators::physical_planner::equijoin_keys( + pred, + &left.schema, + &right.schema, + ) + .map_err(|e| invalid(e.to_string()))?, + completeness: completeness.clone(), + }) + } else { + None + }; + Ok(QueryPlanNode::PhysicalFragment { + pruning, + row_input, + inputs: physical + .input_contracts() + .map(|(id, _)| bindings[&id]) + .collect(), + dag: physical.encode().map_err(|e| invalid(e.to_string()))?, + }) +} + struct DagCompiler<'a, F> { next_id: u64, nodes: BTreeMap, @@ -172,15 +248,15 @@ where } for (local, mut physical) in nodes { match &mut physical { - QueryPlanNode::Logical { inputs, .. } + QueryPlanNode::PhysicalFragment { inputs, .. } + | QueryPlanNode::Logical { inputs, .. } | QueryPlanNode::SummaryMerge { inputs } | QueryPlanNode::ExternalExact { inputs, .. } => { for input in inputs { *input = remap[input]; } } - QueryPlanNode::MembershipFilter { inputs, .. } - | QueryPlanNode::Binary { inputs, .. } + QueryPlanNode::Binary { inputs, .. } | QueryPlanNode::RelationalJoin { inputs, .. } => { for input in inputs { *input = remap[input]; @@ -257,10 +333,11 @@ where right, kind, pred, - pruning: None, + pruning, } if self.preserve_relational => QueryPlanNode::RelationalJoin { inputs: [self.lower(left)?, self.lower(right)?], join_kind: kind.clone(), + pruning: pruning.clone(), pred: serde_json::to_value(pred).map_err(|error| { QueryPlanError::Invalid(format!( "cannot serialize relational join predicate: {error}" @@ -315,7 +392,6 @@ where AggIntent::Min { .. } => Some(residual::Aggregation::Min), AggIntent::Max { .. } => Some(residual::Aggregation::Max), AggIntent::Avg { .. } => Some(residual::Aggregation::Avg), - AggIntent::TopK { .. } => None, _ => { return Err(QueryPlanError::Invalid( "unsupported exact value aggregation".into(), @@ -347,16 +423,9 @@ where labels, without: keys.is_without(), }; - let operator = if let AggIntent::TopK { k, .. } = &measures[0] { - residual::ResidualQueryOperator::TopKSelection { - k: *k as u64, - grouping, - } - } else { - residual::ResidualQueryOperator::Aggregate { - operation: operation.expect("aggregate operation"), - grouping, - } + let operator = residual::ResidualQueryOperator::Aggregate { + operation: operation.expect("aggregate operation"), + grouping, }; QueryPlanNode::Logical { operator, @@ -364,91 +433,16 @@ where } } SummaryExpr::ValueOperation { - child: sort, + child, operation: - planner_types::post_asap::ValueOperation::Limit { - n, - offset: 0, - partition_by: limit_partition, - }, + planner_types::post_asap::ValueOperation::Limit { .. } + | planner_types::post_asap::ValueOperation::Sort { .. } + | planner_types::post_asap::ValueOperation::Filter { .. }, timing: planner_types::post_asap::ExecutionTiming::QueryTime, } => { - let SummaryExpr::ValueOperation { - child, - operation: planner_types::post_asap::ValueOperation::Sort { keys, partition_by }, - timing: planner_types::post_asap::ExecutionTiming::QueryTime, - } = &sort.expr - else { - return Err(QueryPlanError::Invalid( - "query-time Limit must consume a query-time Sort".into(), - )); - }; - if limit_partition != partition_by { - return Err(QueryPlanError::Invalid( - "TopK Limit and Sort grouping differ".into(), - )); - } - if keys.len() != 1 || keys[0].ascending { - return Err(QueryPlanError::Invalid( - "only descending value-ranked TopK is executable".into(), - )); - } - let planner_types::pre_asap::QueryExpr::Column(sort_column) = &keys[0].expr else { - return Err(QueryPlanError::Invalid( - "TopK sort key must reference the child value column".into(), - )); - }; - if !matches!( - child - .schema - .fields - .get(*sort_column) - .map(|field| &field.dtype), - Some(SummaryFamilyType::Plain( - planner_types::pre_asap::DataType::Float64 - )) | Some(SummaryFamilyType::ExactAggregate(..)) - ) { - return Err(QueryPlanError::Invalid( - "TopK sort key must produce a numeric value".into(), - )); - } - let labels = partition_by - .keys() - .iter() - .map(|&column| { - child - .schema - .fields - .get(column) - .map(|field| field.name.clone()) - .ok_or_else(|| { - QueryPlanError::Invalid("unresolved TopK partition column".into()) - }) - }) - .collect::, _>>()?; - QueryPlanNode::Logical { - operator: residual::ResidualQueryOperator::TopKSelection { - k: u64::try_from(*n).map_err(|_| { - QueryPlanError::Invalid("TopK limit exceeds u64".into()) - })?, - grouping: residual::Grouping { - labels, - without: partition_by.is_without(), - }, - }, - inputs: vec![self.lower(child)?], - } + let input = self.lower(child)?; + compile_native_fragment(node, &[input])? } - SummaryExpr::ValueOperation { - child, - operation: planner_types::post_asap::ValueOperation::Sort { keys, .. }, - timing: planner_types::post_asap::ExecutionTiming::QueryTime, - } if keys.len() == 1 => QueryPlanNode::Logical { - operator: residual::ResidualQueryOperator::Sort { - descending: !keys[0].ascending, - }, - inputs: vec![self.lower(child)?], - }, SummaryExpr::ValueOperation { .. } => QueryPlanNode::ExactFallback { reason: "unsupported post-ASAP value operation".into(), }, @@ -456,35 +450,14 @@ where right: candidates, left: values, kind: planner_types::pre_asap::JoinKind::Semi, - pruning: Some(completeness), - pred, - } if !self.preserve_relational => { - validate_membership_join(pred, &values.schema, &candidates.schema)?; + pruning, + .. + } => { let candidate_input = self.lower(candidates)?; - let value_input = if let Some(original) = &self.logical_source { + let value_input = if let Some(original) = + self.logical_source.as_ref().filter(|_| pruning.is_some()) + { let exact_expression = residual::selected_native_expression(original, values)?; - fn item_label(node: &SummaryNode) -> Option { - match &node.expr { - SummaryExpr::SummaryEstimate { summary_input, .. } => { - item_label(summary_input) - } - SummaryExpr::SummaryAgg { input, .. } => match &input.item { - Some(planner_types::post_asap::SummaryInputExpr::Column( - planner_types::pre_asap::ColumnRef::Named(label), - )) => Some(label.clone()), - Some(planner_types::post_asap::SummaryInputExpr::Column( - planner_types::pre_asap::ColumnRef::Qualified { name, .. }, - )) => Some(name.clone()), - _ => None, - }, - _ => None, - } - } - let item_label = item_label(candidates).ok_or_else(|| { - QueryPlanError::Invalid( - "MembershipFilter membership has no named item label".into(), - ) - })?; let value_id = QueryNodeId(self.next_id); self.next_id += 1; self.nodes.insert( @@ -497,25 +470,51 @@ where parameters: BTreeMap::new(), start_parameter: None, end_parameter: None, - input_contracts: vec![ExternalExactInput::CandidateMembership { - item_label, - }], + // The external source provides values; the Planner DAG owns matching. + input_contracts: vec![], }, - inputs: vec![candidate_input], + inputs: vec![], }, ); value_id } else { self.lower(values)? }; - QueryPlanNode::MembershipFilter { - inputs: [candidate_input, value_input], - completeness: completeness.clone(), - } + compile_native_fragment(node, &[value_input, candidate_input])? } SummaryExpr::RelationalJoin { .. } => QueryPlanNode::ExactFallback { - reason: "read-time relational join requires the relational compiler".into(), + reason: "unsupported join in vector adapter".into(), }, + SummaryExpr::BinaryOp { + lhs, + rhs, + operator, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, + } if matches!( + operator.kind, + planner_types::pre_asap::BinaryOpKind::Arithmetic(_) + ) && operator.vector_match.is_none() + && !operator.checked_relative_division + && !operator.checked_finite_division + && node.guarantee.as_ref().is_some_and(|g| !g.has_unknown()) + && [lhs, rhs].iter().all(|operand| { + matches!( + operand.expr, + SummaryExpr::SummaryEstimate { + query: planner_types::post_asap::SketchQuery::Quantile { .. }, + .. + } + ) + }) => + { + let planner_types::pre_asap::BinaryOpKind::Arithmetic(kind) = &operator.kind else { + unreachable!() + }; + QueryPlanNode::Binary { + inputs: [self.lower(lhs)?, self.lower(rhs)?], + operator: kind.clone(), + } + } SummaryExpr::BinaryOp { lhs, rhs, @@ -1337,6 +1336,82 @@ mod tests { assert_eq!(plan.lookup_clickhouse("shared").unwrap().query_id, "sql"); } + // A protocol-vector binding cannot silently invent fields or reconstruct changed rows. + #[test] + fn physical_binding_rejects_invented_samples_and_changed_rows() { + use asap_physical_operators::{ + operators::{Expression, Operator}, + physical_planner::{CompiledPhysicalDag, InputContract}, + }; + use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, + }; + for duplicate_sample in [false, true] { + let schema = std::sync::Arc::new(SummarySchema { + fields: (0..if duplicate_sample { 2 } else { 1 }) + .map(|i| SummaryField { + name: format!("v{i}"), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }) + .collect(), + time_index: None, + }); + let op = if duplicate_sample { + Operator::limit(schema.clone(), 1, 0, vec![]).unwrap() + } else { + Operator::project(schema.clone(), vec![("v0".into(), Expression::Column(0))]) + .unwrap() + }; + let compiled = CompiledPhysicalDag::from_operators( + [(0, InputContract::bounded(schema))].into(), + [(1, (vec![0], op))].into(), + vec![1], + ) + .unwrap(); + let entry = QueryPlanEntry { + language: QueryLanguage::PromQl, + query_id: "q".into(), + canonical_query: "topk(1, m)".into(), + fixed_evaluation: None, + root: QueryNodeId(1), + nodes: BTreeMap::from([ + ( + QueryNodeId(0), + QueryPlanNode::ExactFallback { + reason: "prepared source".into(), + }, + ), + ( + QueryNodeId(1), + QueryPlanNode::PhysicalFragment { + inputs: vec![QueryNodeId(0)], + dag: compiled.encode().unwrap(), + row_input: 0, + pruning: None, + }, + ), + ]), + instant: InstantExecution { + lookback_ms: 0, + full_history: false, + cumulative_readout: false, + }, + fallback: FallbackPolicy::Reject, + }; + let error = entry.validate(&BTreeSet::new()).unwrap_err().to_string(); + assert!( + error.contains(if duplicate_sample { + "one numeric sample" + } else { + "preserve its bound input rows" + }), + "{error}" + ); + } + } + #[test] fn graph_validation_rejects_cycles() { let mut nodes = BTreeMap::new(); @@ -1368,7 +1443,7 @@ mod tests { } #[test] - fn membership_filter_rejects_invalid_completeness_contract() { + fn semi_join_rejects_invalid_completeness_contract() { let leaf = QueryPlanNode::ExactFallback { reason: "prepared".into(), }; @@ -1381,11 +1456,33 @@ mod tests { nodes: BTreeMap::from([ (QueryNodeId(0), leaf.clone()), (QueryNodeId(1), leaf), - ( - QueryNodeId(2), - QueryPlanNode::MembershipFilter { - inputs: [QueryNodeId(0), QueryNodeId(1)], - completeness: CandidateCompleteness::Certified { + (QueryNodeId(2), { + let schema = planner_types::post_asap::SummarySchema { + fields: vec![planner_types::post_asap::SummaryField { + name: "pod".into(), + dtype: planner_types::post_asap::SummaryFamilyType::Plain( + planner_types::pre_asap::DataType::Utf8, + ), + nullable: false, + }], + time_index: None, + }; + QueryPlanNode::RelationalJoin { + inputs: [QueryNodeId(1), QueryNodeId(0)], + join_kind: planner_types::pre_asap::JoinKind::Semi, + pred: serde_json::to_value(planner_types::pre_asap::Predicate( + std::rc::Rc::new(planner_types::pre_asap::QueryExpr::Compare { + left: std::rc::Rc::new(planner_types::pre_asap::QueryExpr::Column( + 0, + )), + op: planner_types::pre_asap::CompareOpKind::Eq, + right: std::rc::Rc::new( + planner_types::pre_asap::QueryExpr::Column(1), + ), + }), + )) + .unwrap(), + pruning: Some(CandidateCompleteness::Certified { guarantee: planner_types::post_asap::ResultGuarantee { metric: planner_types::post_asap::ErrorMetric::Frequency, bound: planner_types::post_asap::BoundExpr::Unknown { @@ -1397,9 +1494,12 @@ mod tests { }, provenance: vec![], }, - }, - }, - ), + }), + left_schema: schema.clone(), + right_schema: schema.clone(), + output_schema: schema, + } + }), ]), instant: InstantExecution { lookback_ms: 300_000, @@ -1411,130 +1511,3 @@ mod tests { assert!(entry.validate(&BTreeSet::new()).is_err()); } } - -// The vector adapter matches complete label identities. Reject joins whose -// predicate would require a different projection instead of silently widening it. -fn validate_membership_join( - pred: &planner_types::pre_asap::Predicate, - left: &planner_types::post_asap::SummarySchema, - right: &planner_types::post_asap::SummarySchema, -) -> Result<(), QueryPlanError> { - use planner_types::post_asap::SummaryFamilyType; - use planner_types::pre_asap::{CompareOpKind, DataType, QueryExpr}; - use std::collections::BTreeSet; - fn collect( - expr: &QueryExpr, - width: usize, - keys: &mut Vec<(usize, usize)>, - ) -> Result<(), QueryPlanError> { - match expr { - QueryExpr::BoolAnd(parts) => { - for part in parts { - collect(part, width, keys)?; - } - } - QueryExpr::Compare { - left, - op: CompareOpKind::Eq, - right, - } => { - let (QueryExpr::Column(a), QueryExpr::Column(b)) = (left.as_ref(), right.as_ref()) - else { - return Err(QueryPlanError::Invalid( - "membership join requires column equality".into(), - )); - }; - let (a, b) = if a < b { (*a, *b) } else { (*b, *a) }; - if a >= width || b < width { - return Err(QueryPlanError::Invalid( - "membership join requires cross-input keys".into(), - )); - } - keys.push((a, b - width)); - } - _ => { - return Err(QueryPlanError::Invalid( - "unsupported membership join predicate".into(), - )) - } - } - Ok(()) - } - let labels = |schema: &planner_types::post_asap::SummarySchema| { - schema - .fields - .iter() - .filter(|f| { - f.name != "__name__" && matches!(f.dtype, SummaryFamilyType::Plain(DataType::Utf8)) - }) - .map(|f| f.name.clone()) - .collect::>() - }; - let mut keys = Vec::new(); - collect(&pred.0, left.fields.len(), &mut keys)?; - let mut matched = BTreeSet::new(); - for (a, b) in keys { - let Some((a, b)) = left.fields.get(a).zip(right.fields.get(b)) else { - return Err(QueryPlanError::Invalid( - "membership join key out of bounds".into(), - )); - }; - if a.name != b.name { - return Err(QueryPlanError::Invalid( - "membership join requires matching label names".into(), - )); - } - matched.insert(a.name.clone()); - } - if matched.is_empty() || matched != labels(left) || matched != labels(right) { - return Err(QueryPlanError::Invalid( - "membership join must match the complete label identity".into(), - )); - } - Ok(()) -} - -#[cfg(test)] -mod membership_binding_tests { - use super::*; - use planner_types::{ - post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, - pre_asap::{CompareOpKind, DataType, Predicate, QueryExpr}, - }; - - // This adapter can implement full-label equality only; reject narrower joins. - #[test] - fn membership_binding_rejects_other_join_semantics() { - let schema = |names: &[&str]| SummarySchema { - fields: names - .iter() - .map(|name| SummaryField { - name: (*name).into(), - dtype: SummaryFamilyType::Plain(DataType::Utf8), - nullable: false, - }) - .collect(), - time_index: None, - }; - let eq = |a, b| QueryExpr::Compare { - left: Rc::new(QueryExpr::Column(a)), - op: CompareOpKind::Eq, - right: Rc::new(QueryExpr::Column(b)), - }; - let one = schema(&["job"]); - assert!(validate_membership_join(&Predicate(Rc::new(eq(0, 1))), &one, &one).is_ok()); - assert!( - validate_membership_join(&Predicate(Rc::new(eq(0, 1))), &one, &schema(&["host"])) - .is_err() - ); - assert!(validate_membership_join(&Predicate(Rc::new(eq(0, 2))), &one, &one).is_err()); - assert!(validate_membership_join( - &Predicate(Rc::new(QueryExpr::BoolAnd(vec![]))), - &one, - &one - ) - .is_err()); - let two = schema(&["job", "host"]); - assert!(validate_membership_join(&Predicate(Rc::new(eq(0, 2))), &two, &two).is_err()); - } -} diff --git a/control_plane/src/query_plan/residual.rs b/control_plane/src/query_plan/residual.rs index 590bb9644..39926848c 100644 --- a/control_plane/src/query_plan/residual.rs +++ b/control_plane/src/query_plan/residual.rs @@ -170,12 +170,20 @@ impl Lower { )? } }; + let sorted = self.operation( + ResidualQueryOperator::Sort { + descending: true, + grouping: grouping.clone(), + }, + vec![input], + )?; return self.operation( - ResidualQueryOperator::TopKSelection { - k: u64::try_from(k).unwrap_or(0), + ResidualQueryOperator::Limit { + n: u64::try_from(k).unwrap_or(0), + offset: 0, grouping, }, - vec![input], + vec![sorted], ); } if a.param.is_some() { @@ -202,8 +210,20 @@ impl Lower { let operator = match c.func.name { "scalar" => ResidualQueryOperator::VectorToScalar, "histogram_quantile" => ResidualQueryOperator::HistogramQuantile, - "sort" => ResidualQueryOperator::Sort { descending: false }, - "sort_desc" => ResidualQueryOperator::Sort { descending: true }, + "sort" => ResidualQueryOperator::Sort { + descending: false, + grouping: Grouping { + labels: vec![], + without: false, + }, + }, + "sort_desc" => ResidualQueryOperator::Sort { + descending: true, + grouping: Grouping { + labels: vec![], + without: false, + }, + }, name => ResidualQueryOperator::Temporal { operation: match name { "rate" => TemporalOperation::Rate, @@ -840,6 +860,22 @@ mod planner_workload_tests { .unwrap_or_else(|error| panic!("{query}: {error}")) } + fn assert_local_limit(node: &QueryPlanNode) { + match node { + QueryPlanNode::PhysicalFragment { dag, .. } => { + let plan = + asap_physical_operators::physical_planner::CompiledPhysicalDag::decode(dag) + .unwrap(); + assert_eq!(plan.operator_name(plan.roots()[0]), Some("Limit")); + } + QueryPlanNode::Logical { + operator: ResidualQueryOperator::Limit { .. }, + .. + } => {} + _ => panic!("expected local Limit, got {node:?}"), + } + } + #[test] fn evaluation_topk_queries_retain_a_local_selection_root() { for query in [ @@ -851,17 +887,7 @@ mod planner_workload_tests { ] { let plan = compile_one(query); let entry = plan.query_plan.entries.values().next().unwrap(); - assert!( - matches!( - entry.nodes[&entry.root], - QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { .. }, - .. - } - ), - "{query}: {:?}", - entry.nodes[&entry.root] - ); + assert_local_limit(&entry.nodes[&entry.root]); assert!( !entry .nodes @@ -880,13 +906,7 @@ mod planner_workload_tests { ] { let plan = compile_one(query); let entry = plan.query_plan.entries.values().next().unwrap(); - assert!(matches!( - entry.nodes[&entry.root], - QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { .. }, - .. - } - )); + assert_local_limit(&entry.nodes[&entry.root]); assert!( !entry.materialization_bindings().is_empty(), "{query} must retain its SummaryStore child: {:?}", @@ -1003,6 +1023,14 @@ pub(crate) fn selected_range_max_materialization( node: &planner_types::post_asap::SummaryNode, ) -> Result, QueryPlanError> { use planner_types::post_asap::{ExactKind, SummaryExpr, SummaryFamilyType}; + let node = match &node.expr { + SummaryExpr::ValueOperation { + child, + operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, + .. + } => child.as_ref(), + _ => node, + }; if !matches!( &node.expr, SummaryExpr::SummaryAgg { @@ -1229,6 +1257,7 @@ pub fn eligible_materialization_keys( visit(original, left, keys)?; visit(original, right, keys)?; } + SummaryExpr::ValueOperation { child, .. } => visit(original, child, keys)?, SummaryExpr::SummaryAgg { child, .. } => visit(original, child, keys)?, SummaryExpr::SummaryEstimate { summary_input, .. } @@ -1342,7 +1371,7 @@ pub fn externalize_residuals(entry: &mut QueryPlanEntry) -> Result<(), QueryPlan ) || matches!( node, QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { .. }, + operator: ResidualQueryOperator::Limit { .. }, .. } ); @@ -1679,7 +1708,7 @@ mod tests { assert!(matches!( entry.nodes[&entry.root], QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { k: actual, .. }, + operator: ResidualQueryOperator::Limit { n: actual, .. }, .. } if actual == k )); @@ -1698,7 +1727,7 @@ mod tests { assert!(matches!( entry.nodes[&entry.root], QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { k: 3, .. }, + operator: ResidualQueryOperator::Limit { n: 3, .. }, .. } )); @@ -1727,7 +1756,7 @@ mod tests { assert!(matches!( &entry.nodes[&entry.root], QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { grouping, .. }, + operator: ResidualQueryOperator::Limit { grouping, .. }, .. } if grouping.labels == labels && grouping.without == without )); diff --git a/control_plane/tests/offline_evidence.rs b/control_plane/tests/offline_evidence.rs index ded77eed3..f24dddeec 100644 --- a/control_plane/tests/offline_evidence.rs +++ b/control_plane/tests/offline_evidence.rs @@ -276,6 +276,19 @@ fn sketch(node: &SummaryNode) -> (&SketchAlgorithm, &SketchParams) { fn assert_exact_count(node: &SummaryNode) { match &node.expr { + SummaryExpr::ValueOperation { + child, + operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, + } => { + assert!(node + .schema + .fields + .iter() + .all(|field| matches!(field.dtype, SummaryFamilyType::Plain(_)))); + assert_exact_count(child); + } + SummaryExpr::SummaryEstimate { summary_input, .. } => assert_exact_count(summary_input), SummaryExpr::SummaryAgg { family: diff --git a/crates/asap_sketch_codec/Cargo.toml b/crates/asap_sketch_codec/Cargo.toml deleted file mode 100644 index c2d728895..000000000 --- a/crates/asap_sketch_codec/Cargo.toml +++ /dev/null @@ -1,8 +0,0 @@ -[package] -name = "asap_sketch_codec" -version.workspace = true -edition.workspace = true - -[dependencies] -asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", branch = "main" } -prost = "0.13" diff --git a/crates/asap_sketch_codec/src/lib.rs b/crates/asap_sketch_codec/src/lib.rs deleted file mode 100644 index 279efe168..000000000 --- a/crates/asap_sketch_codec/src/lib.rs +++ /dev/null @@ -1,84 +0,0 @@ -//! 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(DdSketchState { - alpha: sketch.wire_alpha(), - store_counts: sketch.store_counts.clone(), - store_offset: sketch.store_offset, - })), - }; - 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_types/Cargo.toml b/crates/asap_types/Cargo.toml index 2ee4b61fc..d5c83ef67 100644 --- a/crates/asap_types/Cargo.toml +++ b/crates/asap_types/Cargo.toml @@ -21,3 +21,5 @@ planner-types.workspace = true # Accuracy metadata projects the authoritative Planner guarantees. asap-aware-mapping.workspace = true + +asap-physical-operators.workspace = true diff --git a/crates/asap_types/src/aggregation_type.rs b/crates/asap_types/src/aggregation_type.rs index 647f7604f..7ff2ce7ea 100644 --- a/crates/asap_types/src/aggregation_type.rs +++ b/crates/asap_types/src/aggregation_type.rs @@ -1,215 +1,2 @@ -//! 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() - ); - } -} +//! Kernel identity is owned by the shared physical operator library. +pub use asap_physical_operators::AggregationType; diff --git a/crates/asap_types/src/enums.rs b/crates/asap_types/src/enums.rs index 96cf0a0c0..ebf03534b 100644 --- a/crates/asap_types/src/enums.rs +++ b/crates/asap_types/src/enums.rs @@ -1,87 +1,7 @@ use std::fmt; use std::str::FromStr; -use tracing::debug; -/// The scalar value a serving-time query wants out of an already-built -/// accumulator: "given a live `AggregateCore` implementation, which -/// number do you want?" Every accumulator's `AggregateCore::query_statistic` -/// dispatches on this. Distinct from L3's `AggIntent` (a planning-time -/// IR node carrying accuracy targets, column refs, φ, k) — nothing at -/// L3/L4 reaches down to a live Rust struct's fields, so `Statistic` has -/// no upstream ASAPController equivalent; it's this workspace's own -/// serving-time vocabulary. -/// -/// Formerly `promql_utilities::query_logics::enums::Statistic` — moved -/// here because its real center of gravity (`compatible_agg_types`, -/// `QueryRequirements`, capability matching) already lived in this -/// crate, and `asap_types` — not `data_plane` — is the shared foundation -/// both `control_plane`'s ecosystem and `data_plane` can depend on -/// without a cycle. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)] -pub enum Statistic { - Count, - Sum, - Cardinality, - FrequencyL2, - FrequencyEntropy, - Increase, - Rate, - Min, - Max, - Quantile, - Topk, -} - -impl fmt::Display for Statistic { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - debug!("Formatting Statistic: {:?}", self); - match self { - Statistic::Count => write!(f, "count"), - Statistic::Sum => write!(f, "sum"), - Statistic::Cardinality => write!(f, "cardinality"), - Statistic::FrequencyL2 => write!(f, "frequency_l2"), - Statistic::FrequencyEntropy => write!(f, "frequency_entropy"), - Statistic::Increase => write!(f, "increase"), - Statistic::Rate => write!(f, "rate"), - Statistic::Min => write!(f, "min"), - Statistic::Max => write!(f, "max"), - Statistic::Quantile => write!(f, "quantile"), - Statistic::Topk => write!(f, "topk"), - } - } -} - -#[allow(clippy::should_implement_trait)] -impl Statistic { - pub fn from_str(s: &str) -> Option { - debug!("Parsing Statistic from string: {}", s); - match s.to_lowercase().as_str() { - "count" => Some(Statistic::Count), - "sum" => Some(Statistic::Sum), - "cardinality" => Some(Statistic::Cardinality), - "frequency_l2" => Some(Statistic::FrequencyL2), - "frequency_entropy" => Some(Statistic::FrequencyEntropy), - "increase" => Some(Statistic::Increase), - "rate" => Some(Statistic::Rate), - "min" => Some(Statistic::Min), - "max" => Some(Statistic::Max), - "quantile" => Some(Statistic::Quantile), - "topk" => Some(Statistic::Topk), - _ => None, - } - } -} - -impl FromStr for Statistic { - type Err = (); - - /// Parse a statistic from a string (case-insensitive). - /// Use `s.parse::()` or `Statistic::from_str(s)`. - fn from_str(s: &str) -> Result { - debug!("FromStr trait parsing Statistic: {}", s); - Statistic::from_str(s).ok_or(()) - } -} +pub use asap_physical_operators::Statistic; #[derive( clap::ValueEnum, diff --git a/crates/asap_types/src/executable_plan.rs b/crates/asap_types/src/executable_plan.rs index 03e2193f9..c475584a6 100644 --- a/crates/asap_types/src/executable_plan.rs +++ b/crates/asap_types/src/executable_plan.rs @@ -20,8 +20,8 @@ use serde::{Deserialize, Serialize}; #[serde(transparent)] pub struct QueryNodeId(pub u64); -pub const OWNED_POST_ASAP_DAG_SCHEMA_VERSION: u32 = 3; -pub const MAINTENANCE_DAG_SCHEMA_VERSION: u32 = 4; +pub const OWNED_POST_ASAP_DAG_SCHEMA_VERSION: u32 = 5; +pub const MAINTENANCE_DAG_SCHEMA_VERSION: u32 = 6; /// Versioned, language-neutral Planner DAG persisted with an installed plan. /// Plan lifecycle belongs to the enclosing `PrecomputePlan`; this document diff --git a/crates/asap_types/src/query_plan.rs b/crates/asap_types/src/query_plan.rs index f4eb6e76d..452943cca 100644 --- a/crates/asap_types/src/query_plan.rs +++ b/crates/asap_types/src/query_plan.rs @@ -384,6 +384,84 @@ impl QueryPlanEntry { ))); } for (id, node) in &self.nodes { + if let QueryPlanNode::PhysicalFragment { + inputs, + dag, + row_input, + pruning, + } = node + { + let compiled = + asap_physical_operators::physical_planner::CompiledPhysicalDag::decode(dag) + .map_err(|e| QueryPlanError::Invalid(e.to_string()))?; + for (_, contract) in compiled.input_contracts() { + let mut samples = 0; + for (index, field) in contract.schema.fields.iter().enumerate() { + use planner_types::{post_asap::SummaryFamilyType, pre_asap::DataType}; + match &field.dtype { + SummaryFamilyType::Plain(DataType::Float64 | DataType::Int64) => { + samples += 1 + } + SummaryFamilyType::Plain(DataType::Utf8) => {} + SummaryFamilyType::Plain(DataType::Timestamp) + if contract.schema.time_index == Some(index) => {} + _ => { + return Err(QueryPlanError::Invalid(format!( + "PromQL input binding cannot supply field {}", + field.name + ))) + } + } + } + if samples > 1 { + return Err(QueryPlanError::Invalid( + "PromQL vector input has only one numeric sample per row".into(), + )); + } + } + let source = compiled.input_contracts().nth(*row_input).map(|(id, _)| id); + if compiled.roots().len() != 1 + || source.is_none() + || compiled.row_source(compiled.roots()[0]) != source + { + return Err(QueryPlanError::Invalid( + "physical vector output must preserve its bound input rows".into(), + )); + } + if let Some(pruning) = pruning { + if matches!(&pruning.completeness, CandidateCompleteness::Certified { guarantee } + if guarantee.metric != planner_types::post_asap::ErrorMetric::TopKMembership || guarantee.bound.evaluate().is_none() || guarantee.failure_probability.evaluate().is_none()) + { + return Err(QueryPlanError::Invalid( + "invalid physical pruning certificate".into(), + )); + } + let contracts = compiled.input_contracts().collect::>(); + let left = contracts + .get(*row_input) + .ok_or_else(|| QueryPlanError::Invalid("invalid row input".into()))? + .1; + let right = contracts + .get(pruning.candidate_input) + .ok_or_else(|| QueryPlanError::Invalid("invalid candidate input".into()))? + .1; + asap_physical_operators::operators::Operator::semi_join( + left.schema.clone(), + right.schema.clone(), + pruning.keys.clone(), + ) + .map_err(|e| QueryPlanError::Invalid(e.to_string()))?; + } + if compiled.input_contracts().count() != inputs.len() + || compiled.roots().len() != 1 + || *row_input >= inputs.len() + { + return Err(QueryPlanError::Invalid( + "physical input/root binding mismatch".into(), + )); + } + } + if let QueryPlanNode::Logical { operator, inputs } = node { operator.validate(inputs.len())?; } @@ -409,7 +487,17 @@ impl QueryPlanEntry { )); } } - if let QueryPlanNode::MembershipFilter { completeness, .. } = node { + if let QueryPlanNode::RelationalJoin { + pruning: Some(completeness), + join_kind, + .. + } = node + { + if *join_kind != planner_types::pre_asap::JoinKind::Semi { + return Err(QueryPlanError::Invalid( + "pruning evidence requires a semi-join".into(), + )); + } if matches!( completeness, CandidateCompleteness::Certified { guarantee } @@ -419,7 +507,7 @@ impl QueryPlanEntry { || guarantee.failure_probability.evaluate().is_none() ) { return Err(QueryPlanError::Invalid( - "invalid MembershipFilter completeness certificate".into(), + "invalid semi-join pruning certificate".into(), )); } } @@ -568,12 +656,30 @@ pub struct ExternalExactRequest { pub input_contracts: Vec, } +/// Required candidate rows must have authoritative values at the bound source. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(deny_unknown_fields)] +pub struct PruningInputContract { + pub candidate_input: usize, + pub keys: Vec<(usize, usize)>, + pub completeness: CandidateCompleteness, +} + #[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] #[serde(tag = "op", rename_all = "snake_case", deny_unknown_fields)] pub enum QueryPlanNode { + /// Planner-compiled computation. Input order follows the physical input contracts. + PhysicalFragment { + inputs: Vec, + dag: Vec, + row_input: usize, + pruning: Option, + }, + RelationalJoin { inputs: [QueryNodeId; 2], join_kind: planner_types::pre_asap::JoinKind, + pruning: Option, pred: serde_json::Value, left_schema: planner_types::post_asap::SummarySchema, right_schema: planner_types::post_asap::SummarySchema, @@ -616,13 +722,6 @@ pub enum QueryPlanNode { SummaryMerge { inputs: Vec, }, - /// Semijoin value rows against membership identities, preserving their values - /// and order. Inputs are membership and authoritative values respectively. - /// Ranking, grouping and limiting are separate downstream operators. - MembershipFilter { - inputs: [QueryNodeId; 2], - completeness: CandidateCompleteness, - }, /// An exact subtree evaluated outside ASAP. Its results enter the query DAG /// like any other node output and may depend on summary-produced inputs. ExternalExact { @@ -645,10 +744,10 @@ impl QueryPlanNode { | Self::Relational { input, .. } | Self::SummaryEstimate { input, .. } | Self::ExactReadout { input, .. } => std::slice::from_ref(input), - Self::SummaryMerge { inputs } + Self::PhysicalFragment { inputs, .. } + | Self::SummaryMerge { inputs } | Self::Logical { inputs, .. } | Self::ExternalExact { inputs, .. } => inputs, - Self::MembershipFilter { inputs, .. } => inputs, } } } diff --git a/crates/asap_types/src/query_plan/residual.rs b/crates/asap_types/src/query_plan/residual.rs index aa7260a9e..002f98dab 100644 --- a/crates/asap_types/src/query_plan/residual.rs +++ b/crates/asap_types/src/query_plan/residual.rs @@ -36,11 +36,10 @@ pub enum ResidualQueryOperator { operation: Aggregation, grouping: Grouping, }, - /// PromQL `topk(k, vector)` selection over values produced by the child. - /// This is distinct from a frequency-sketch TopK readout: any exact or - /// summary-backed instant-vector child may feed this query-time operator. - TopKSelection { - k: u64, + /// Select an ordered slice independently within each group. + Limit { + n: u64, + offset: u64, grouping: Grouping, }, Binary { @@ -52,6 +51,7 @@ pub enum ResidualQueryOperator { }, Sort { descending: bool, + grouping: Grouping, }, HistogramQuantile, Subquery { diff --git a/crates/asap_types/src/traits.rs b/crates/asap_types/src/traits.rs index 196d90813..51afd904f 100644 --- a/crates/asap_types/src/traits.rs +++ b/crates/asap_types/src/traits.rs @@ -1,7 +1 @@ -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; -} +pub use asap_physical_operators::SerializableToSink; diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index 7482c5210..10bd74c10 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -6,7 +6,8 @@ edition.workspace = true [dependencies] # Internal crates (workspace) asap_types.workspace = true -asap_sketch_codec = { path = "../crates/asap_sketch_codec" } +asap-physical-operators.workspace = true +asap_sketch_codec.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 @@ -64,7 +65,7 @@ prometheus = "0.13" lazy_static = "1.4" reqwest.workspace = true tracing-appender = "0.2" -asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", branch = "main" } +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "5f03ccbd798ed5fec62bdd839bcb331123cab369" } # Generic MessagePack codec for decoding the heap-bearing CountSketch # DELTA-HEAP wire frame (`MSGPACK_DELTA`) directly in the backend, WITHOUT # adding a delta API to the public `asap_sketchlib`: the frame is decoded @@ -103,4 +104,4 @@ default = [] # Enable lock profiling instrumentation lock_profiling = [] # Enable extra debugging output -extra_debugging = [] +extra_debugging = ["asap-physical-operators/extra_debugging"] diff --git a/data_plane/benches/sketch_db.rs b/data_plane/benches/sketch_db.rs index 11cbfa4a0..23aedab50 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 data_plane::precompute_engine::operators::SumAccumulator; +use asap_physical_operators::summary_kernels::SumAccumulator; use data_plane::storage_engines::sketch_db::data::{ AccuracyBound, AggKind, AggregationType, Capability, SketchAlgorithm, SketchConfig, SketchEncoding, @@ -52,6 +52,7 @@ fn encode_ddsketch(values: &[f64], alpha: f64) -> Vec { 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)), diff --git a/data_plane/examples/sketch_db_diag.rs b/data_plane/examples/sketch_db_diag.rs index 325d72b0c..90555f5f6 100644 --- a/data_plane/examples/sketch_db_diag.rs +++ b/data_plane/examples/sketch_db_diag.rs @@ -38,6 +38,7 @@ fn ddsketch_payload() -> Vec { 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)), diff --git a/data_plane/examples/univmon_erp_artifact.rs b/data_plane/examples/univmon_erp_artifact.rs index aa37ec137..c20687a1b 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 data_plane::precompute_engine::operators::hll_sketch_accumulator::HllSketchAccumulator; -use data_plane::precompute_engine::operators::univmon_accumulator::UnivMonAccumulator; +use asap_physical_operators::summary_kernels::hll_sketch::HllSketchAccumulator; +use asap_physical_operators::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 fea49a0a2..2b23439a7 100644 --- a/data_plane/src/drivers/ingest/otel.rs +++ b/data_plane/src/drivers/ingest/otel.rs @@ -24,7 +24,6 @@ use std::collections::HashMap; use std::io::Read; -use crate::precompute_engine::operators::sketch_envelope_accumulator::SketchEnvelopeAccumulator; use crate::precompute_engine::series_router::WorkerMessage; use crate::precompute_engine::IngestState; use crate::query_engines::routing::FreshnessProbeCache; @@ -35,6 +34,7 @@ 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 axum::{body::Bytes, extract::State, routing::post, Json, Router}; @@ -2144,7 +2144,7 @@ fn dp_carries_heap(dp: &ModifiedOtlpSketchDp) -> bool { .unwrap_or(false) } ENCODING_MSGPACK_DELTA => { - use crate::precompute_engine::operators::CountMinSketchWithHeapAccumulator; + use asap_physical_operators::summary_kernels::CountMinSketchWithHeapAccumulator; CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes(&dp.sketch) .map(|acc| !acc.inner.topk_heap_items().is_empty()) .unwrap_or(false) @@ -2573,7 +2573,7 @@ fn decode_modified_otlp_sketch_bytes( encoding: i32, bytes: &[u8], ) -> Result, Box> { - use crate::precompute_engine::operators::{ + use asap_physical_operators::summary_kernels::{ CountMinSketchAccumulator, CountSketchAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, HllSketchAccumulator, }; @@ -2644,7 +2644,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 crate::precompute_engine::operators::CountSketchWithHeapAccumulator; + use asap_physical_operators::summary_kernels::CountSketchWithHeapAccumulator; return Ok(Box::new( CountSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(bytes)?, )); @@ -2709,11 +2709,11 @@ fn empty_accumulator_for_delta_bootstrap( config: &crate::storage_engines::sketch_db::index::SketchConfig, encoding: i32, ) -> Option> { - use crate::precompute_engine::operators::{ + use crate::storage_engines::sketch_db::index::SketchConfig; + use asap_physical_operators::summary_kernels::{ CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, HllSketchAccumulator, }; - use crate::storage_engines::sketch_db::index::SketchConfig; match (algorithm, config) { (SketchAlgorithm::Hll, SketchConfig::Hll { precision }) => { @@ -2783,7 +2783,7 @@ pub(crate) fn apply_modified_otlp_delta_bytes( existing: &mut Box, bytes: &[u8], ) -> Result<(), Box> { - use crate::precompute_engine::operators::{ + use asap_physical_operators::summary_kernels::{ CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, DDSketchAccumulator, HllSketchAccumulator, }; @@ -3029,7 +3029,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 crate::precompute_engine::operators::sum_accumulator::SumAccumulator::from_sum_bytes(&dp.sketch) { + let value = match asap_physical_operators::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}"); @@ -3469,8 +3469,8 @@ mod policy_fp_lookup_tests { #[cfg(test)] mod dispatcher_tests { use super::*; - use crate::precompute_engine::operators::{DDSketchAccumulator, HllSketchAccumulator}; use crate::storage_engines::types::AggregateCore; + use asap_physical_operators::summary_kernels::{DDSketchAccumulator, HllSketchAccumulator}; use asap_sketchlib::DdSketch; use asap_sketchlib::HllVariant; @@ -3822,8 +3822,8 @@ mod sid_resolution_tests { /// directly observable on the bucket counts. #[tokio::test] async fn delta_apply_rotates_per_series_base_at_window_boundary() { - use crate::precompute_engine::operators::DDSketchAccumulator; 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 prost::Message; @@ -3861,6 +3861,7 @@ mod sid_resolution_tests { alpha: 0.01, store_counts: vec![10, 0, 5], store_offset: 0, + ..Default::default() })), ..Default::default() } @@ -4154,8 +4155,8 @@ mod sid_resolution_tests { /// recover after a backend restart. #[tokio::test] async fn leading_cms_delta_bootstraps_onto_empty_base() { - use crate::precompute_engine::operators::CountMinSketchAccumulator; use asap_otel_proto::sketchlib::v1::CountMinDelta as PbDelta; + use asap_physical_operators::summary_kernels::CountMinSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; @@ -4240,8 +4241,8 @@ mod sid_resolution_tests { /// the register-max updates. #[tokio::test] async fn leading_hll_delta_bootstraps_onto_empty_base() { - use crate::precompute_engine::operators::HllSketchAccumulator; use asap_otel_proto::sketchlib::v1::HllDelta as PbDelta; + use asap_physical_operators::summary_kernels::HllSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; diff --git a/data_plane/src/drivers/query/servers/http.rs b/data_plane/src/drivers/query/servers/http.rs index 8c03df23f..3a33f3c83 100644 --- a/data_plane/src/drivers/query/servers/http.rs +++ b/data_plane/src/drivers/query/servers/http.rs @@ -1354,6 +1354,10 @@ async fn process_via_named_engine( ) .into_response() } + Err(error @ EngineError::Physical(_)) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"status": "error", "errorType": "execution", "error": error.to_string()})), + ).into_response(), Err(EngineError::Backend { .. }) => { warn!( data_source_id = data_source_id, @@ -1483,6 +1487,10 @@ async fn process_via_router( EngineError::CapabilityMiss { .. } => { forward_instant_to_fallback(state, parsed_request, headers).await } + EngineError::Physical(_) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"status": "error", "errorType": "execution", "error": last.to_string()})), + ).into_response(), EngineError::Backend { .. } => ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ @@ -2297,6 +2305,10 @@ async fn process_range_query_request( } } } + EngineError::Physical(_) => ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(serde_json::json!({"status": "error", "errorType": "execution", "error": last.to_string()})), + ).into_response(), EngineError::Backend { .. } => ( StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ @@ -3576,6 +3588,7 @@ mod tests { #[derive(Clone)] enum MockOutcome { + Physical(asap_physical_operators::Error), /// Empty instant vector — the Prometheus adapter still /// produces `status=success` with `data.result=[]`. OkEmpty, @@ -3604,7 +3617,8 @@ mod tests { impl QueryEngine for MockQueryEngine { async fn execute(&self, _query: &str) -> Result { self.calls.fetch_add(1, Ordering::SeqCst); - match self.outcome { + match &self.outcome { + MockOutcome::Physical(error) => Err(EngineError::Physical(error.clone())), MockOutcome::OkEmpty => Ok(QueryResult::vector(Vec::new(), 0)), MockOutcome::Backend => Err(EngineError::backend( self.caps.data_source_id, @@ -3624,7 +3638,8 @@ mod tests { _step_ms: u64, ) -> Result { self.calls.fetch_add(1, Ordering::SeqCst); - match self.outcome { + match &self.outcome { + MockOutcome::Physical(error) => Err(EngineError::Physical(error.clone())), MockOutcome::OkEmpty => Ok(QueryResult::matrix(Vec::new())), MockOutcome::Backend => Err(EngineError::backend( self.caps.data_source_id, @@ -4074,6 +4089,54 @@ mod tests { assert_eq!(accuracy["delta"], 0.0); } + // Execution failure must survive the router and HTTP boundary without failover. + #[tokio::test] + async fn physical_resource_errors_never_forward_instant_or_range_queries() { + for cause in [ + asap_physical_operators::Error::MemoryLimit, + asap_physical_operators::Error::Cancelled, + ] { + let error = asap_physical_operators::Error::AtNode { + node: 7, + operation: "sort".into(), + source: Box::new(cause.clone()), + }; + let (warm, warm_calls) = + MockQueryEngine::new(StorageBackend::SketchStore, MockOutcome::Physical(error)); + let (other, other_calls) = + MockQueryEngine::new(StorageBackend::DoubleWrite, MockOutcome::OkEmpty); + let port = + setup_test_server_with_router(StorageBackend::DoubleWrite, vec![warm, other]).await; + for endpoint in ["query", "query_range"] { + let response = Client::new() + .get(format!("http://127.0.0.1:{port}/api/v1/{endpoint}")) + .query(&[ + ("query", "sum_over_time(foo[5m])"), + ("time", "1700000000"), + ("start", "1700000000"), + ("end", "1700000060"), + ("step", "15"), + ]) + .send() + .await + .unwrap(); + assert_eq!( + response.status(), + reqwest::StatusCode::INTERNAL_SERVER_ERROR + ); + let body: serde_json::Value = response.json().await.unwrap(); + assert_eq!(body["status"], "error"); + assert_eq!(body["errorType"], "execution"); + assert!( + body["error"].as_str().unwrap().contains(&cause.to_string()), + "{body}" + ); + } + assert_eq!(warm_calls.load(Ordering::SeqCst), 2); + assert_eq!(other_calls.load(Ordering::SeqCst), 0); + } + } + #[tokio::test] async fn http_router_serves_double_write_via_warm_head() { // Step-1 of the JSONL deprecation deleted the diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index 8c3ac483c..2968fed2a 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 precompute_engine::operators::{ +pub use asap_physical_operators::summary_kernels::{ IncreaseAccumulator, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, }; diff --git a/data_plane/src/main.rs b/data_plane/src/main.rs index 1448cc337..e8db1fd1e 100644 --- a/data_plane/src/main.rs +++ b/data_plane/src/main.rs @@ -746,6 +746,8 @@ async fn main() -> Result<()> { capability_snapshot_id: "bootstrap".into(), }, ingest: asap_types::precompute_plan::IngestContract { + // Bootstrap envelope: no plan is installed yet, so there is no + // dataset to bind to. A published plan supplies the identity. dataset_identity: None, protocol: asap_types::precompute_plan::IngestProtocol::ModifiedOtlpMetricsV1, endpoint_path: "/v1/metrics".into(), diff --git a/data_plane/src/precompute_engine/accumulator_factory.rs b/data_plane/src/precompute_engine/accumulator_factory.rs deleted file mode 100644 index 3346f6059..000000000 --- a/data_plane/src/precompute_engine/accumulator_factory.rs +++ /dev/null @@ -1,2025 +0,0 @@ -use crate::precompute_engine::operators::{ - CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, - CountSketchWithHeapAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, - HydraKllSketchAccumulator, IncreaseAccumulator, KeyedCounterState, KeyedMaxState, - KeyedMinState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, -}; -use crate::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, Measurement, -}; -use asap_types::aggregation_config::PrecomputeMaterialization; -// Production dispatch consumes Planner SummaryAgg payloads directly. The -// config adapter below is compiled only for isolated historical kernel tests. -use super::operators::hll_sketch_accumulator::HllSketchAccumulator; -use super::operators::univmon_accumulator::UnivMonAccumulator; -#[cfg(test)] -use asap_types::accumulator_spec::cms_params; -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) - } - }; -} - -/// Trait for feeding samples into accumulators in the precompute engine. -/// -/// 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 maintenance 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, - weight_scale: f64, -} - -impl CmsHeapAccumulatorUpdater { - pub fn new(row_num: usize, col_num: usize, heap_size: usize, weight: TopkWeight) -> Self { - Self::with_weight_scale(row_num, col_num, heap_size, weight, 1.0) - } - - pub fn with_weight_scale( - row_num: usize, - col_num: usize, - heap_size: usize, - weight: TopkWeight, - weight_scale: f64, - ) -> Self { - Self { - acc: CountMinSketchWithHeapAccumulator::new(row_num, col_num, heap_size), - row_num, - col_num, - heap_size, - weight, - weight_scale, - } - } -} - -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 * self.weight_scale, - // Σ 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, - weight_scale: f64, -} - -impl CountSketchWithHeapAccumulatorUpdater { - pub fn new(row_num: usize, col_num: usize, heap_size: usize, weight: TopkWeight) -> Self { - Self::with_weight_scale(row_num, col_num, heap_size, weight, 1.0) - } - - pub fn with_weight_scale( - row_num: usize, - col_num: usize, - heap_size: usize, - weight: TopkWeight, - weight_scale: f64, - ) -> Self { - Self { - acc: CountSketchWithHeapAccumulator::new(row_num, col_num, heap_size), - row_num, - col_num, - heap_size, - weight, - weight_scale, - } - } -} - -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 * self.weight_scale, - 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 -// --------------------------------------------------------------------------- - -#[cfg(test)] -/// Return `true` if `config` produces a keyed (MultipleSubpopulation) updater, -/// without allocating an updater object. -/// -/// **Contract:** this must agree with every concrete `AccumulatorUpdater::is_keyed()` -/// implementation. When a new accumulator type is added, update both here and -/// in the corresponding struct. -pub fn config_is_keyed(config: &PrecomputeMaterialization) -> bool { - config - .accumulator_spec() - .expect("valid fixture") - .grouping - .is_some() -} - -/// Top-k ranking quantity, selected by `weight_mode` or its alias `topk_weight`. -/// -/// * `value` / `sum`: sum values per key (default). -/// * `count` / `frequency` / `freq`: count occurrences per key. -#[cfg(test)] -fn topk_weight_param(config: &PrecomputeMaterialization) -> TopkWeight { - match config.sample_update_rule() { - asap_types::SampleUpdateRule::Count => TopkWeight::Count, - asap_types::SampleUpdateRule::Value { .. } - | asap_types::SampleUpdateRule::CounterDelta { .. } => TopkWeight::Value, - } -} - -#[cfg(test)] -fn topk_weight_scale_param(config: &PrecomputeMaterialization) -> f64 { - match config.sample_update_rule() { - asap_types::SampleUpdateRule::Value { scale } => scale, - asap_types::SampleUpdateRule::CounterDelta { scale } => scale, - asap_types::SampleUpdateRule::Count => 1.0, - } -} - -// --------------------------------------------------------------------------- -// Factory function -// --------------------------------------------------------------------------- - -/// Read the KLL `k` out of `SketchParams::Kll`. `accumulator_spec()` -/// always builds a `SketchKind` whose `SketchAlgorithm::Kll` is paired with -/// `SketchParams::Kll`, so the -/// other arm is unreachable from a `spec` this module builds itself. -#[cfg(test)] -fn kll_k(params: &SketchParams) -> u16 { - match params { - // Lossless: `accumulator_spec()` only ever stores a value that - // already fit in `u16` (via `kll_k_param`'s own `u16::try_from` - // fallback) widened to `u32`. - SketchParams::Kll { k } => *k as u16, - other => unreachable!( - "accumulator_spec() paired SketchAlgorithm::Kll with non-Kll params: {other:?}" - ), - } -} - -/// Read `(rows = depth, columns = width)` out of `SketchParams::Cms` or `::CountSketch` -/// — same shape, different variant per bare-sketch identity. -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:?}" - ), - } -} - -/// Read the DDSketch relative-accuracy `alpha` out of `SketchParams::DDSketch`. -#[cfg(test)] -fn ddsketch_alpha(params: &SketchParams) -> f64 { - match params { - SketchParams::DDSketch { alpha } => *alpha, - other => unreachable!( - "accumulator_spec() paired SketchAlgorithm::DDSketch with non-DDSketch params: {other:?}" - ), - } -} - -/// Construct isolated payload fixtures for kernel/storage unit tests. -/// Production execution requires a validated Planner DAG program. -#[cfg(test)] -pub fn create_fixture_accumulator( - config: &PrecomputeMaterialization, -) -> Box { - let spec = config - .accumulator_spec() - .expect("invalid isolated kernel fixture"); - - let keyed = spec.grouping.is_some(); - - match (&spec.family, keyed) { - (SummaryFamilyType::ExactAggregate(ExactKind::Sum | ExactKind::Count, _), false) => { - Box::new(SumAccumulatorUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), true) => { - Box::new(KeyedSumCountAccumulatorUpdater::for_family(ExactKind::Sum)) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Count, _), true) => Box::new( - KeyedSumCountAccumulatorUpdater::for_family(ExactKind::Count), - ), - - // Direction comes off the family itself now. It used to be read - // back out of `aggregation_sub_type` because Planner had one - // `MinMax` accumulator for both directions, which meant a config - // whose sub_type was lost or misspelled silently built the wrong - // extremum. - (SummaryFamilyType::ExactAggregate(ExactKind::Min, _), false) => { - Box::new(MinAccumulatorUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Min, _), true) => { - Box::new(KeyedMinStateUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Max, _), false) => { - Box::new(MaxAccumulatorUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Max, _), true) => { - Box::new(KeyedMaxStateUpdater::new()) - } - - (SummaryFamilyType::ExactAggregate(ExactKind::Increase | ExactKind::Rate, _), false) => { - Box::new(IncreaseAccumulatorUpdater::new()) - } - (SummaryFamilyType::ExactAggregate(ExactKind::Increase | ExactKind::Rate, _), true) => { - Box::new(KeyedCounterStateUpdater::new()) - } - - (SummaryFamilyType::Sketch(kind, _), false) - if kind.algorithm() == &SketchAlgorithm::Kll => - { - Box::new(KllAccumulatorUpdater::new(kll_k(kind.params()))) - } - // HydraKLL: `k` comes off the typed params like the unkeyed case, - // but the `(row, col)` tiling grid has no `SketchParams::Kll` - // field to live in (see `asap_types::accumulator_spec`'s module - // doc) — read it the same way bare CMS does, via `cms_params`. - (SummaryFamilyType::Sketch(kind, _), true) if kind.algorithm() == &SketchAlgorithm::Kll => { - let (row_num, col_num) = cms_params(config); - Box::new(HydraKllAccumulatorUpdater::new( - row_num, - col_num, - kll_k(kind.params()), - )) - } - - // Bare CMS: point-frequency only, min-of-rows estimator. `keyed=false` - // can't actually arise here today (no `AggregationType` resolves to - // bare Cms unkeyed — see accumulator_spec.rs), matched anyway as a - // safe default. - (SummaryFamilyType::Sketch(kind, _), _) if kind.algorithm() == &SketchAlgorithm::Cms => { - let (row_num, col_num) = cms_dims(kind.params()); - Box::new(CmsAccumulatorUpdater::new(row_num, col_num)) - } - - // CountSketch uses the median-of-signed-rows estimator. - (SummaryFamilyType::Sketch(kind, _), _) - if kind.algorithm() == &SketchAlgorithm::CountSketch => - { - let (row_num, col_num) = cms_dims(kind.params()); - Box::new(CountSketchAccumulatorUpdater::new(row_num, col_num)) - } - - // Heap-bearing top-k variant (raw-input ingest path): route to the - // real `CmsHeapAccumulatorUpdater` so the per-policy top-k heap is - // BUILT (heap-less CMS could not answer `topk(...)` — recall 0). - // Keyed by the configured group-by `aggregated_labels` (e.g. `host`), - // ranked by Σ value per key by default (`weight_mode: value`), or Σ - // count for genuine frequency-top-k (`weight_mode: count`). The OTLP - // modified-sketch path builds the heap agent-side and uses - // `SketchEnvelope` ingest, not this raw arm. - (SummaryFamilyType::Sketch(kind, _), _) - if kind.algorithm() == &SketchAlgorithm::CmsWithHeap => - { - let (row_num, col_num, heap_size) = cms_heap_dims(kind.params()); - Box::new(CmsHeapAccumulatorUpdater::with_weight_scale( - row_num, - col_num, - heap_size, - topk_weight_param(config), - topk_weight_scale_param(config), - )) - } - - // Heap-bearing CountSketch retains CountSketch estimation semantics. - (SummaryFamilyType::Sketch(kind, _), _) - if kind.algorithm() == &SketchAlgorithm::CountSketchWithHeap => - { - let (row_num, col_num, heap_size) = cms_heap_dims(kind.params()); - Box::new(CountSketchWithHeapAccumulatorUpdater::with_weight_scale( - row_num, - col_num, - heap_size, - topk_weight_param(config), - topk_weight_scale_param(config), - )) - } - - (SummaryFamilyType::Sketch(kind, _), _) - if kind.algorithm() == &SketchAlgorithm::DDSketch => - { - Box::new(DDSketchAccumulatorUpdater::new(ddsketch_alpha( - kind.params(), - ))) - } - - (SummaryFamilyType::Sketch(kind, _), false) - if kind.algorithm() == &SketchAlgorithm::UnivMon => - { - let SketchParams::UnivMon { - heap_size, - sketch_rows, - sketch_cols, - layers, - } = kind.params() - else { - unreachable!("validated UnivMon family parameters") - }; - Box::new(UnivMonUpdater { - acc: UnivMonAccumulator::new( - *heap_size as usize, - *sketch_rows as usize, - *sketch_cols as usize, - *layers as usize, - ) - .expect("validated UnivMon dimensions"), - }) - } - - (SummaryFamilyType::Sketch(kind, _), false) - if kind.algorithm() == &SketchAlgorithm::Hll => - { - let SketchParams::Hll { precision } = kind.params() else { - unreachable!("validated HLL family parameters") - }; - Box::new(HllUpdater { - acc: HllSketchAccumulator::new( - asap_sketchlib::HllVariant::Regular, - u32::from(*precision), - ), - }) - } - - (other_family, keyed) => { - panic!("unsupported isolated kernel fixture {other_family:?}, keyed={keyed}") - } - } -} - -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 tests { - use super::*; - use asap_types::enums::WindowKind; - use asap_types::AggregationType; - - #[test] - fn immutable_dds_inputs_reject_nonpositive_and_unrepresentable_values() { - let updater = DDSketchAccumulatorUpdater::new(0.01); - for value in [-20.0, -0.0, 0.0, f64::NAN, f64::INFINITY, f64::MAX] { - assert!(updater.validate_single_input(value).is_err()); - } - for value in [0.5, 20.0, 40.0] { - assert!(updater.validate_single_input(value).is_ok()); - } - } - - /// Both cardinality implementations consume values, with a single signed-zero identity. - #[test] - fn hll_and_univmon_raw_updates_share_value_identity() { - for family in [AggregationType::HLL, AggregationType::UnivMon] { - let config = PrecomputeMaterialization::new( - family, - String::new(), - Default::default(), - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - String::new(), - 60, - 60, - WindowKind::Tumbling, - "m".into(), - "m".into(), - None, - None, - None, - ); - let mut updater = create_fixture_accumulator(&config); - for value in [0.0, -0.0, 2.0, 2.0, f64::NAN] { - updater.update_single(value, 1000); - } - let state = updater.take_accumulator(); - assert_eq!(state.get_accumulator_type(), family); - let estimate = state - .query_statistic( - asap_types::Statistic::Cardinality, - &None, - &Default::default(), - ) - .unwrap(); - assert!((estimate - 2.0).abs() < 0.05, "{family:?}: {estimate}"); - assert!(updater.memory_usage_bytes() >= 4096); - let empty = updater - .snapshot_accumulator() - .query_statistic( - asap_types::Statistic::Cardinality, - &None, - &Default::default(), - ) - .unwrap(); - assert_eq!(empty, 0.0); - } - } - - #[test] - fn test_sum_updater() { - let mut updater = SumAccumulatorUpdater::new(); - assert!(!updater.is_keyed()); - - updater.update_single(1.0, 1000); - updater.update_single(2.0, 2000); - updater.update_single(3.0, 3000); - - let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "SumAccumulator"); - } - - #[test] - fn test_minmax_updater() { - let mut updater = MaxAccumulatorUpdater::new(); - updater.update_single(5.0, 1000); - updater.update_single(3.0, 2000); - updater.update_single(7.0, 3000); - - let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "MaxAccumulator"); - } - - #[test] - fn test_increase_updater() { - let mut updater = IncreaseAccumulatorUpdater::new(); - updater.update_single(10.0, 1000); - updater.update_single(15.0, 2000); - - let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "IncreaseAccumulator"); - } - - #[test] - fn test_kll_updater() { - let mut updater = KllAccumulatorUpdater::new(200); - for i in 1..=10 { - updater.update_single(i as f64, i * 1000); - } - - let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "DatasketchesKLLAccumulator"); - } - - #[test] - fn test_multiple_sum_updater() { - let mut updater = KeyedSumCountAccumulatorUpdater::new(); - assert!(updater.is_keyed()); - - let key_a = KeyByLabelValues::new_with_labels(vec!["a".to_string()]); - let key_b = KeyByLabelValues::new_with_labels(vec!["b".to_string()]); - - updater.update_keyed(&key_a, 1.0, 1000); - updater.update_keyed(&key_b, 2.0, 2000); - - let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "KeyedSumCountAccumulator"); - } - - #[test] - fn bare_cms_adds_sample_values() { - let mut updater = CmsAccumulatorUpdater::new(4, 256); - let key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - updater.update_keyed(&key, 2.0, 1000); - updater.update_keyed(&key, 3.0, 2000); - updater.update_keyed(&key, 5.0, 3000); - - let acc = updater.snapshot_accumulator(); - let cms = acc - .as_any() - .downcast_ref::() - .expect("should be a CountMinSketchAccumulator"); - assert_eq!(cms.query_key(&key), 10.0); - } - - #[test] - fn bare_count_sketch_adds_sample_values() { - let mut updater = CountSketchAccumulatorUpdater::new(5, 256); - let key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - updater.update_keyed(&key, 2.0, 1000); - updater.update_keyed(&key, 3.0, 2000); - updater.update_keyed(&key, 5.0, 3000); - - let acc = updater.snapshot_accumulator(); - let count_sketch = acc - .as_any() - .downcast_ref::() - .expect("should be a CountSketchAccumulator"); - assert_eq!(count_sketch.query_key(&key), 10.0); - } - - #[test] - fn test_reset_clears_state() { - let mut updater = SumAccumulatorUpdater::new(); - updater.update_single(100.0, 1000); - updater.reset(); - // After reset, should produce a fresh accumulator - let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "SumAccumulator"); - } - - #[test] - fn test_config_is_keyed() { - use std::collections::HashMap; - - let make_config = |agg_type: AggregationType, sub_type: &str| { - PrecomputeMaterialization::new( - agg_type, - sub_type.to_string(), - HashMap::new(), - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - String::new(), - 60, - 0, - WindowKind::Tumbling, - "m".to_string(), - "m".to_string(), - None, - None, - None, - ) - }; - - // Non-keyed types - assert!(!config_is_keyed(&make_config( - AggregationType::SingleSubpopulation, - "Sum" - ))); - assert!(!config_is_keyed(&make_config(AggregationType::Sum, ""))); - assert!(!config_is_keyed(&make_config( - AggregationType::DatasketchesKLL, - "" - ))); - assert!(!config_is_keyed(&make_config( - AggregationType::Increase, - "" - ))); - - // Keyed types - assert!(config_is_keyed(&make_config( - AggregationType::MultipleSubpopulation, - "Sum" - ))); - let mut keyed = make_config(AggregationType::Sum, ""); - keyed.aggregated_labels = asap_types::KeyByLabelNames::new(vec!["host".into()]); - assert!(config_is_keyed(&keyed)); - let mut keyed = make_config(AggregationType::Increase, ""); - keyed.aggregated_labels = asap_types::KeyByLabelNames::new(vec!["host".into()]); - assert!(config_is_keyed(&keyed)); - let mut keyed = make_config(AggregationType::Max, ""); - keyed.aggregated_labels = asap_types::KeyByLabelNames::new(vec!["host".into()]); - assert!(config_is_keyed(&keyed)); - assert!(config_is_keyed(&make_config( - AggregationType::CountMinSketch, - "" - ))); - assert!(config_is_keyed(&make_config( - AggregationType::CountMinSketchWithHeap, - "" - ))); - assert!(config_is_keyed(&make_config( - AggregationType::CountSketch, - "" - ))); - assert!(config_is_keyed(&make_config( - AggregationType::CountSketchWithHeap, - "" - ))); - assert!(config_is_keyed(&make_config(AggregationType::HydraKLL, ""))); - - // Verify agreement with updater.is_keyed() - for (agg_type, sub_type) in &[ - (AggregationType::SingleSubpopulation, "Sum"), - (AggregationType::MultipleSubpopulation, "Sum"), - (AggregationType::Sum, ""), - (AggregationType::DatasketchesKLL, ""), - (AggregationType::CountMinSketch, ""), - ] { - let config = make_config(*agg_type, sub_type); - let updater = create_fixture_accumulator(&config); - assert_eq!( - config_is_keyed(&config), - updater.is_keyed(), - "config_is_keyed disagrees with updater.is_keyed() for type={:?}", - agg_type - ); - } - } - - #[test] - fn test_kll_k_param_capital_k() { - // SingleSubpopulation/KLL with capital "K" param should use it (not default to 200) - use std::collections::HashMap; - let mut params = HashMap::new(); - params.insert("K".to_string(), serde_json::Value::from(50_u64)); - let config = PrecomputeMaterialization::new( - AggregationType::SingleSubpopulation, - "DatasketchesKLL".to_string(), - params, - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - String::new(), - 60, - 0, - WindowKind::Tumbling, - "m".to_string(), - "m".to_string(), - None, - None, - None, - ); - let updater = create_fixture_accumulator(&config); - let acc = updater.snapshot_accumulator(); - let kll = acc - .as_any() - .downcast_ref::() - .expect("should be KLL"); - assert_eq!(kll.inner.k, 50, "k should be 50 from capital-K param"); - } - - #[test] - fn cms_params_reads_canonical_w_d_keys() { - use std::collections::HashMap; - // Canonical `w`/`d` form — what the control plane's - // `sketch_params_to_json` emits and what asapcollector - // streaming-config YAMLs ship (asapcollector PR - // `sync-config-canonical-w-d` migrated them in lock-step - // with the legacy-fallback removal). - let mut params = HashMap::new(); - params.insert("d".to_string(), serde_json::Value::from(7_u64)); - params.insert("w".to_string(), serde_json::Value::from(2048_u64)); - let config = PrecomputeMaterialization::new( - AggregationType::CountMinSketch, - String::new(), - params, - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - String::new(), - 60, - 0, - WindowKind::Tumbling, - "m".to_string(), - "m".to_string(), - None, - None, - None, - ); - assert_eq!(super::cms_params(&config), (7, 2048)); - - // Empty params — defaults `(4, 1000)`. - let empty_config = PrecomputeMaterialization::new( - AggregationType::CountMinSketch, - String::new(), - HashMap::new(), - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - asap_types::KeyByLabelNames::new(vec![]), - String::new(), - 60, - 0, - WindowKind::Tumbling, - "m".to_string(), - "m".to_string(), - None, - None, - None, - ); - assert_eq!(super::cms_params(&empty_config), (4, 1000)); - } - - // ----------------------------------------------------------------- - // value-weighted vs count-weighted top-k (fix/value-weighted-topk) - // ----------------------------------------------------------------- - - /// Build a `*WithHeap` config keyed by group-by label `host`, with the - /// given `weight_mode` param (None → default = value-weighted). - fn topk_config( - agg_type: AggregationType, - weight_mode: Option<&str>, - ) -> PrecomputeMaterialization { - use std::collections::HashMap; - let mut params = HashMap::new(); - // Small, deterministic geometry; heap big enough to hold all hosts. - params.insert("d".to_string(), serde_json::Value::from(4_u64)); - params.insert("w".to_string(), serde_json::Value::from(256_u64)); - params.insert("heap_size".to_string(), serde_json::Value::from(8_u64)); - if let Some(m) = weight_mode { - params.insert("weight_mode".to_string(), serde_json::Value::from(m)); - } - PrecomputeMaterialization::new( - agg_type, - String::new(), - params, - asap_types::KeyByLabelNames::new(vec![]), - // group-by = `host` (NOT the metric label `item`). - asap_types::KeyByLabelNames::new(vec!["host".to_string()]), - asap_types::KeyByLabelNames::new(vec![]), - String::new(), - 60, - 0, - WindowKind::Tumbling, - "cpu".to_string(), - "cpu".to_string(), - None, - None, - None, - ) - } - - /// Read the heap as a sorted-descending `(host, value)` list from a - /// finished accumulator — mirrors the read-side reducer's - /// `topk_heap_items()` + sort-by-value-desc. - fn ranked_topk(acc: &dyn AggregateCore) -> Vec<(String, f64)> { - let heap = acc - .as_any() - .downcast_ref::() - .expect("WithHeap config must build a heap accumulator"); - let mut items = heap.inner.topk_heap_items(); - items.sort_by(|a, b| { - b.value - .partial_cmp(&a.value) - .unwrap_or(std::cmp::Ordering::Equal) - }); - items.into_iter().map(|i| (i.key, i.value)).collect() - } - - /// Same as `ranked_topk`, but for the real `CountSketchWithHeapAccumulator` - /// (median-of-signed-rows) built by `SketchAlgorithm::CountSketchWithHeap` — - /// no longer conflated with the CMS-family accumulator above. - fn ranked_topk_cs(acc: &dyn AggregateCore) -> Vec<(String, f64)> { - let heap = acc - .as_any() - .downcast_ref::() - .expect("CountSketchWithHeap config must build a CountSketchWithHeapAccumulator"); - let mut items = heap.inner.topk_heap_items(); - items.sort_by(|a, b| { - b.value - .partial_cmp(&a.value) - .unwrap_or(std::cmp::Ordering::Equal) - }); - items.into_iter().map(|i| (i.key, i.value)).collect() - } - - fn host_key(h: &str) -> KeyByLabelValues { - KeyByLabelValues::new_with_labels(vec![h.to_string()]) - } - - /// A multi-host CPU stream where value-rank and count-rank DISAGREE, - /// so the test distinguishes a correct value-weighted answer from the - /// (buggy) count-weighted one. - /// - /// host-a: ONE big sample -> value 100, count 1 - /// host-b: TWO mid samples -> value 60, count 2 - /// host-c: FOUR tiny ones -> value 20, count 4 - /// - /// By Σ VALUE: a(100) > b(60) > c(20) → top-2 = [a, b] - /// By Σ COUNT: c(4) > b(2) > a(1) → top-2 = [c, b] - const STREAM: &[(&str, f64)] = &[ - ("host-a", 100.0), - ("host-b", 30.0), - ("host-b", 30.0), - ("host-c", 5.0), - ("host-c", 5.0), - ("host-c", 5.0), - ("host-c", 5.0), - ]; - - fn feed_stream(updater: &mut dyn AccumulatorUpdater) { - for (i, (host, val)) in STREAM.iter().enumerate() { - updater.update_keyed(&host_key(host), *val, 1_000 + i as i64); - } - } - - #[test] - fn value_weighted_topk_ranks_hosts_by_sum_of_value() { - // DEFAULT mode (no weight_mode param) must be value-weighted. - let config = topk_config(AggregationType::CountMinSketchWithHeap, None); - let mut updater = create_fixture_accumulator(&config); - assert!(updater.is_keyed()); - - feed_stream(&mut *updater); - let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "CountMinSketchWithHeapAccumulator"); - - let ranked = ranked_topk(&*acc); - // Σ value: host-a=100, host-b=60, host-c=20. - assert_eq!(ranked[0].0, "host-a", "top host by Σ value"); - assert_eq!(ranked[0].1, 100.0); - assert_eq!(ranked[1].0, "host-b"); - assert_eq!(ranked[1].1, 60.0); - assert_eq!(ranked[2].0, "host-c"); - assert_eq!(ranked[2].1, 20.0); - - // Recall of value-weighted top-2 against ground truth {host-a, host-b}. - let truth: std::collections::HashSet<&str> = ["host-a", "host-b"].into_iter().collect(); - let got: std::collections::HashSet<&str> = - ranked.iter().take(2).map(|(h, _)| h.as_str()).collect(); - let recall = got.intersection(&truth).count() as f64 / truth.len() as f64; - assert_eq!(recall, 1.0, "value-weighted top-2 recall must be 1.0"); - } - - #[test] - fn counter_delta_scale_preserves_sub_unit_membership_weights() { - let mut config = topk_config( - AggregationType::CountMinSketchWithHeap, - Some("counter_delta"), - ); - config - .parameters - .insert("weight_scale".into(), serde_json::json!(1_000_000)); - let mut updater = create_fixture_accumulator(&config); - updater.update_keyed(&host_key("payment"), 0.004, 1_000); - updater.update_keyed(&host_key("order"), 0.002, 1_000); - let ranked = ranked_topk(&*updater.take_accumulator()); - assert_eq!(ranked[0], ("payment".into(), 4_000.0)); - assert_eq!(ranked[1], ("order".into(), 2_000.0)); - } - - #[test] - fn count_weighted_topk_still_ranks_by_occurrence_frequency() { - // Opt-in frequency-top-k: weight_mode=count must rank by event count. - let config = topk_config(AggregationType::CountMinSketchWithHeap, Some("count")); - let mut updater = create_fixture_accumulator(&config); - feed_stream(&mut *updater); - let acc = updater.take_accumulator(); - - let ranked = ranked_topk(&*acc); - // Σ count: host-c=4, host-b=2, host-a=1. - assert_eq!(ranked[0].0, "host-c", "top host by Σ count"); - assert_eq!(ranked[0].1, 4.0); - assert_eq!(ranked[1].0, "host-b"); - assert_eq!(ranked[1].1, 2.0); - assert_eq!(ranked[2].0, "host-a"); - assert_eq!(ranked[2].1, 1.0); - } - - #[test] - fn countsketch_with_heap_also_routes_to_value_weighted_heap() { - // CountSketchWithHeap gets its OWN dedicated updater/accumulator - // (real median-of-signed-rows math) — same value-weighted default - // as the CMS-family heap path, but no longer conflated with it. - let config = topk_config(AggregationType::CountSketchWithHeap, None); - let mut updater = create_fixture_accumulator(&config); - feed_stream(&mut *updater); - let acc = updater.take_accumulator(); - assert_eq!(acc.type_name(), "CountSketchWithHeapAccumulator"); - let ranked = ranked_topk_cs(&*acc); - assert_eq!(ranked[0].0, "host-a"); - assert_eq!(ranked[0].1, 100.0); - } - - #[test] - fn topk_weight_param_parses_modes() { - assert_eq!( - super::topk_weight_param(&topk_config(AggregationType::CountMinSketchWithHeap, None)), - TopkWeight::Value, - "unset defaults to value-weighted" - ); - for m in ["value", "sum", "VALUE"] { - assert_eq!( - super::topk_weight_param(&topk_config( - AggregationType::CountMinSketchWithHeap, - Some(m) - )), - TopkWeight::Value, - ); - } - for m in ["count", "frequency", "freq", "COUNT"] { - assert_eq!( - super::topk_weight_param(&topk_config( - AggregationType::CountMinSketchWithHeap, - Some(m) - )), - TopkWeight::Count, - ); - } - } -} - -#[cfg(test)] -mod planner_family_regression { - use super::*; - use asap_types::{enums::WindowKind, KeyByLabelNames}; - - // Every installed exact producer must retain its family in runtime state. - #[test] - fn exact_state_identity_survives_factory_and_reset() { - for kind in [ - AggregationType::Sum, - AggregationType::Count, - AggregationType::Rate, - AggregationType::Increase, - AggregationType::Min, - AggregationType::Max, - ] { - let config = PrecomputeMaterialization::new( - kind, - String::new(), - Default::default(), - KeyByLabelNames::empty(), - KeyByLabelNames::empty(), - KeyByLabelNames::empty(), - String::new(), - 60, - 60, - WindowKind::Tumbling, - String::new(), - "metric".into(), - None, - None, - None, - ); - let mut updater = create_planner_accumulator( - &config.accumulator_spec().unwrap().family, - &planner_types::post_asap::SummaryUpdate::column( - planner_types::pre_asap::ColumnRef::SampleValue, - ), - &Default::default(), - ) - .unwrap(); - updater.update_single(4.0, 1000); - updater.update_single(7.0, 2000); - assert_eq!(updater.take_accumulator().get_accumulator_type(), kind); - assert_eq!(updater.snapshot_accumulator().get_accumulator_type(), kind); - } - } -} - -/// 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> { - 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: super::operators::exact_accumulator::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 weight_scale = 1.0; - 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::with_weight_scale( - r, - c, - h, - TopkWeight::Value, - weight_scale, - )) - } - ( - SketchAlgorithm::CountSketchWithHeap, - params @ SketchParams::CountSketchWithHeap { .. }, - ) => { - let (r, c, h) = cms_heap_dims(params); - Box::new(CountSketchWithHeapAccumulatorUpdater::with_weight_scale( - r, - c, - h, - TopkWeight::Value, - weight_scale, - )) - } - (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() - && !asap_types::accumulator_spec::is_unit_sample_frequency(input) - { - return Err("Planner item expression does not match the selected kernel layout".into()); - } - Ok(updater) -} - -struct PlannerExactUpdater { - acc: super::operators::exact_accumulator::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 = super::operators::exact_accumulator::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() - } -} - -#[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/data_plane/src/precompute_engine/ingest_handler.rs b/data_plane/src/precompute_engine/ingest_handler.rs index 67940478b..a531f3c0c 100644 --- a/data_plane/src/precompute_engine/ingest_handler.rs +++ b/data_plane/src/precompute_engine/ingest_handler.rs @@ -366,8 +366,8 @@ mod tests { #[tokio::test] async fn delta_path_reconstitutes_cumulative_state() { use crate::drivers::ingest::otel::apply_modified_otlp_delta_bytes; - use crate::precompute_engine::operators::DDSketchAccumulator; use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_sketchlib::DdSketch; 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 crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::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 683ed5c06..bbed2b039 100644 --- a/data_plane/src/precompute_engine/maintenance_runtime.rs +++ b/data_plane/src/precompute_engine/maintenance_runtime.rs @@ -167,7 +167,7 @@ impl PrecomputeOperatorRegistry for OperatorAdapter<'_> { inputs: &[Arc], ) -> Result { if node.output_state.timing != planner_types::post_asap::ExecutionTiming::IngestionTime { - return Err("maintenance runtime requires ingestion-time nodes".into()); + return Err("ingestion executor received a query-time node".into()); } match &node.payload { ExecutableOperatorPayload::SummaryMerge => merge_inputs(inputs), @@ -253,7 +253,7 @@ impl PrecomputeOperatorRegistry for OperatorAdapter<'_> { "keyed maintenance updates require explicit row identity routing".into(), ); } - let mut updater = super::accumulator_factory::create_planner_accumulator( + let mut updater = asap_physical_operators::factory::create_planner_accumulator( family, input, grouping, )?; if updater.is_keyed() { @@ -466,8 +466,9 @@ fn evaluate_aligned_binary( let right = right_rows .get(×tamp) .ok_or("maintenance binary requires matching timestamp sets")?; - let value = - crate::utils::arithmetic::evaluate_float64_arithmetic(arithmetic, left, *right); + let value = asap_physical_operators::arithmetic::evaluate_float64_arithmetic( + arithmetic, left, *right, + ); if !value.is_finite() { return Err("maintenance binary produced a non-finite update".into()); } @@ -2148,7 +2149,7 @@ pub(crate) fn affected_materializations( #[cfg(test)] mod tests { use super::*; - use crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::summary_kernels::SumAccumulator; use planner_types::post_asap::{ EdgeRole, ExecutableDag, ExecutableDagEdge, GroupingEdgeCompatibility, SummarySchema, WindowEdgeCompatibility, @@ -2182,7 +2183,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 = crate::precompute_engine::operators::SumAccumulator::new(); + let mut state = asap_physical_operators::summary_kernels::SumAccumulator::new(); state.update(value); FrozenExactWindows { sid, @@ -3411,6 +3412,7 @@ mod tests { assert!(frozen .execute(&operation, &[left.clone(), right.clone()]) .is_err()); + operation.output_state = planner_types::post_asap::ExecutionDataState::INGESTION_ROWS; for invalid in [ rows(vec![]), diff --git a/data_plane/src/precompute_engine/mod.rs b/data_plane/src/precompute_engine/mod.rs index 3f744870a..7c8176b50 100644 --- a/data_plane/src/precompute_engine/mod.rs +++ b/data_plane/src/precompute_engine/mod.rs @@ -1,4 +1,3 @@ -pub mod accumulator_factory; pub mod config; pub mod coordination_checkpoint; mod engine; @@ -9,7 +8,6 @@ pub mod ingest_handler; pub mod maintenance_runtime; pub(crate) mod metrics; pub mod multisource_coordinator; -pub mod operators; pub mod output_sink; pub mod raw_dag; pub mod series_buffer; diff --git a/data_plane/src/precompute_engine/operators/count_min_sketch_accumulator.rs b/data_plane/src/precompute_engine/operators/count_min_sketch_accumulator.rs deleted file mode 100644 index b3307fdf6..000000000 --- a/data_plane/src/precompute_engine/operators/count_min_sketch_accumulator.rs +++ /dev/null @@ -1,1323 +0,0 @@ -use crate::precompute_engine::operators::dd_sketch_accumulator::normalize_sample_p; -use crate::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CountMinSketch, CountMinSketchDelta, MessagePackCodec}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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_otel_proto::sketchlib::v1::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: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::MultipleSubpopulationAggregate; - use asap_types::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::precompute_engine::operators::sum_accumulator::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_otel_proto::sketchlib::v1::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![], - } - .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 asap_types::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 asap_sketchlib::proto::sketchlib::{ - sketch_envelope, CountMinState, CounterType, SketchEnvelope, - }; - use asap_types::Statistic; - 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/data_plane/src/precompute_engine/operators/count_min_sketch_with_heap_accumulator.rs b/data_plane/src/precompute_engine/operators/count_min_sketch_with_heap_accumulator.rs deleted file mode 100644 index 3eea0afdb..000000000 --- a/data_plane/src/precompute_engine/operators/count_min_sketch_with_heap_accumulator.rs +++ /dev/null @@ -1,832 +0,0 @@ -use crate::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CmsHeapItem, CountMinSketchWithHeap, MessagePackCodec}; -use serde::Deserialize; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs b/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs deleted file mode 100644 index 9a353a33e..000000000 --- a/data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs +++ /dev/null @@ -1,686 +0,0 @@ -//! 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::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CountSketch, CountSketchDelta, MessagePackCodec}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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::precompute_engine::operators::count_min_sketch_accumulator::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_otel_proto::sketchlib::v1::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(); - // Proto-schema-divergence-tracker: the Go-side - // `CountSketchDelta` proto carries an `hh_keys` field - // (heavy-hitter candidate keys forwarded by the upstream - // Space-Saving tracker). The Rust wire-format struct now - // models it (`asap_sketchlib::CountSketchDelta::hh_keys`), - // but the vendored Rust proto bindings in - // `asap_otel_proto::sketchlib::v1` haven't been regenerated - // against the latest `.proto` yet, so no `hh_keys` arrive on - // the wire from Go producers. Sending an empty `hh_keys` - // disables the TopK rebuild path; it'll start firing once the - // proto-schema sync PR lands. - 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: asap_types::Statistic, - key: &Option, - query_kwargs: &HashMap, - ) -> Result> { - use asap_types::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::precompute_engine::operators::count_min_sketch_accumulator::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_otel_proto::sketchlib::v1::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![], - } - .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/data_plane/src/precompute_engine/operators/count_sketch_with_heap_accumulator.rs b/data_plane/src/precompute_engine/operators/count_sketch_with_heap_accumulator.rs deleted file mode 100644 index 7c8de44a9..000000000 --- a/data_plane/src/precompute_engine/operators/count_sketch_with_heap_accumulator.rs +++ /dev/null @@ -1,575 +0,0 @@ -//! Count Sketch with Heap accumulator — wraps -//! `asap_sketchlib::CountSketchWithHeap`. -//! -//! Port of `count_min_sketch_with_heap_accumulator.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::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CountSketchWithHeap, CsHeapItem, MessagePackCodec}; -use serde::Deserialize; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::Statistic; - -/// Local serde view of the DELTA-HEAP wire frame (encoding `MSGPACK_DELTA`). -/// Identical shape to `count_min_sketch_with_heap_accumulator.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: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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::precompute_engine::operators::count_min_sketch_with_heap_accumulator::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/data_plane/src/precompute_engine/operators/datasketches_kll_accumulator.rs b/data_plane/src/precompute_engine/operators/datasketches_kll_accumulator.rs deleted file mode 100644 index 2874f5102..000000000 --- a/data_plane/src/precompute_engine/operators/datasketches_kll_accumulator.rs +++ /dev/null @@ -1,733 +0,0 @@ -use crate::storage_engines::types::{ - 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 asap_types::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: asap_types::Statistic, - _key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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::*; - - fn encode_state(state: asap_sketchlib::proto::sketchlib::KllState) -> Vec { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; - use prost::Message; - 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::precompute_engine::operators::sum_accumulator::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; - use prost::Message; - - 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}; - use prost::Message; - 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}; - use prost::Message; - - 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}; - use prost::Message; - - 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; - use prost::Message; - 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; - use prost::Message; - // 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/data_plane/src/precompute_engine/operators/dd_sketch_accumulator.rs b/data_plane/src/precompute_engine/operators/dd_sketch_accumulator.rs deleted file mode 100644 index 0f63348b1..000000000 --- a/data_plane/src/precompute_engine/operators/dd_sketch_accumulator.rs +++ /dev/null @@ -1,667 +0,0 @@ -//! 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::storage_engines::types::{ - 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()); - } - // The DataPoint-level METRIC scalars (count/sum/min/max) were - // dropped from `DDSketchState` (ProjectASAP/sketchlib-go#243 / - // asap_sketchlib#57). Reconstruct from the bucket store only: - // `DdSketch::from_raw` now takes just (alpha, store_counts, - // store_offset) and recovers `count` by summing the bucket - // counts via `total_count()`. - let inner = DdSketch::from_raw(state.alpha, state.store_counts.clone(), state.store_offset); - 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_otel_proto::sketchlib::v1::DdSketchDelta` message. - pub fn apply_proto_delta_bytes( - &mut self, - buffer: &[u8], - ) -> Result<(), Box> { - use asap_otel_proto::sketchlib::v1::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, - ..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: asap_types::Statistic, - _key: &Option, - query_kwargs: &HashMap, - ) -> Result> { - use asap_types::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, - }; - 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, - }; - 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::precompute_engine::operators::count_sketch_accumulator::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_otel_proto::sketchlib::v1::{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, - }, - ], - } - .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_otel_proto::sketchlib::v1::{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, - }], - } - .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 asap_types::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 asap_types::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 asap_types::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 asap_types::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 asap_types::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 asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; - use asap_types::Statistic; - 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() - }; - 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, - } - .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); - } -} diff --git a/data_plane/src/precompute_engine/operators/exact_accumulator.rs b/data_plane/src/precompute_engine/operators/exact_accumulator.rs deleted file mode 100644 index b7fcead12..000000000 --- a/data_plane/src/precompute_engine/operators/exact_accumulator.rs +++ /dev/null @@ -1,327 +0,0 @@ -//! Exact summary state identified by Planner family, independent of keyed layout. -use super::increase_accumulator::IncreaseAccumulator; -use crate::storage_engines::types::{ - AggregateCore, AggregationType, AuxStats, KeyByLabelValues, Measurement, SerializableToSink, -}; -use asap_types::Statistic; -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 { - 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 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( - >::merge_accumulators(vec![a.clone(), b.clone()])?, - ), - (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")?; - if self.family != other.family || self.is_keyed() != other.is_keyed() { - return Err("cannot merge different Planner families or layouts".into()); - } - let mut merged = self.clone(); - if let (Some(target), Some(source)) = (&mut merged.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 { - merged.scalar = merge_scalar(&self.scalar, &other.scalar)?; - } - 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/data_plane/src/precompute_engine/operators/hll_sketch_accumulator.rs b/data_plane/src/precompute_engine/operators/hll_sketch_accumulator.rs deleted file mode 100644 index b4c3b4d0c..000000000 --- a/data_plane/src/precompute_engine/operators/hll_sketch_accumulator.rs +++ /dev/null @@ -1,790 +0,0 @@ -//! 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::precompute_engine::operators::dd_sketch_accumulator::normalize_sample_p; -use crate::storage_engines::types::{ - 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_otel_proto::sketchlib::v1::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: asap_types::Statistic, - _key: &Option, - _query_kwargs: &HashMap, - ) -> Result> { - use asap_types::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::precompute_engine::operators::count_sketch_accumulator::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_otel_proto::sketchlib::v1::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 asap_types::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 asap_types::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 asap_sketchlib::proto::sketchlib::{ - sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, - }; - use asap_types::Statistic; - 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/data_plane/src/precompute_engine/operators/hydra_kll_accumulator.rs b/data_plane/src/precompute_engine/operators/hydra_kll_accumulator.rs deleted file mode 100644 index c3793584b..000000000 --- a/data_plane/src/precompute_engine/operators/hydra_kll_accumulator.rs +++ /dev/null @@ -1,168 +0,0 @@ -use crate::{ - storage_engines::types::{ - AggregateCore, AggregationType, MergeableAccumulator, MultipleSubpopulationAggregate, - SerializableToSink, - }, - KeyByLabelValues, -}; -use asap_sketchlib::{HydraKllSketch, MessagePackCodec}; -use base64::{engine::general_purpose, Engine as _}; -use std::collections::HashMap; - -use asap_types::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: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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/data_plane/src/precompute_engine/operators/increase_accumulator.rs b/data_plane/src/precompute_engine/operators/increase_accumulator.rs deleted file mode 100644 index 421feaa4d..000000000 --- a/data_plane/src/precompute_engine/operators/increase_accumulator.rs +++ /dev/null @@ -1,742 +0,0 @@ -use crate::storage_engines::types::{ - AggregateCore, AggregationType, Measurement, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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 { - /// 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..] { - if acc.starting_timestamp > result.last_seen_timestamp { - result.total_increase += - if acc.starting_measurement.value >= result.last_seen_measurement.value { - acc.starting_measurement.value - result.last_seen_measurement.value - } else { - acc.starting_measurement.value - }; - } - result.total_increase += acc.total_increase; - result.sample_count = result.sample_count.saturating_add(acc.sample_count); - if acc.last_seen_timestamp > result.last_seen_timestamp { - result.last_seen_measurement = acc.last_seen_measurement.clone(); - result.last_seen_timestamp = acc.last_seen_timestamp; - } - } - - 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 (first, second) = if self.starting_timestamp <= other_increase.starting_timestamp { - (self, other_increase) - } else { - (other_increase, self) - }; - 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; - } - - 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: asap_types::Statistic, - _key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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/data_plane/src/precompute_engine/operators/keyed_counter_state.rs b/data_plane/src/precompute_engine/operators/keyed_counter_state.rs deleted file mode 100644 index b94d2faba..000000000 --- a/data_plane/src/precompute_engine/operators/keyed_counter_state.rs +++ /dev/null @@ -1,529 +0,0 @@ -use crate::precompute_engine::operators::IncreaseAccumulator; -use crate::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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::storage_engines::types::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/data_plane/src/precompute_engine/operators/keyed_max_state.rs b/data_plane/src/precompute_engine/operators/keyed_max_state.rs deleted file mode 100644 index 4309c04a0..000000000 --- a/data_plane/src/precompute_engine/operators/keyed_max_state.rs +++ /dev/null @@ -1,335 +0,0 @@ -use crate::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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/data_plane/src/precompute_engine/operators/keyed_min_state.rs b/data_plane/src/precompute_engine/operators/keyed_min_state.rs deleted file mode 100644 index 5be698f50..000000000 --- a/data_plane/src/precompute_engine/operators/keyed_min_state.rs +++ /dev/null @@ -1,335 +0,0 @@ -use crate::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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/data_plane/src/precompute_engine/operators/keyed_sum_count_accumulator.rs b/data_plane/src/precompute_engine/operators/keyed_sum_count_accumulator.rs deleted file mode 100644 index c39d55831..000000000 --- a/data_plane/src/precompute_engine/operators/keyed_sum_count_accumulator.rs +++ /dev/null @@ -1,558 +0,0 @@ -use crate::storage_engines::types::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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/data_plane/src/precompute_engine/operators/max_accumulator.rs b/data_plane/src/precompute_engine/operators/max_accumulator.rs deleted file mode 100644 index 733c2028a..000000000 --- a/data_plane/src/precompute_engine/operators/max_accumulator.rs +++ /dev/null @@ -1,248 +0,0 @@ -use crate::storage_engines::types::{ - AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::Statistic; - -/// Exact maximum over one population, mergeable by comparison. -/// -/// See [`MinAccumulator`](super::min_accumulator::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: asap_types::Statistic, - _key: &Option, - _query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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_accumulator::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/data_plane/src/precompute_engine/operators/min_accumulator.rs b/data_plane/src/precompute_engine/operators/min_accumulator.rs deleted file mode 100644 index e69cda830..000000000 --- a/data_plane/src/precompute_engine/operators/min_accumulator.rs +++ /dev/null @@ -1,253 +0,0 @@ -use crate::storage_engines::types::{ - AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::Statistic; - -/// Exact minimum over one population, mergeable by comparison. -/// -/// The sibling [`MaxAccumulator`](super::max_accumulator::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: asap_types::Statistic, - _key: &Option, - _query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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_accumulator::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/data_plane/src/precompute_engine/operators/mod.rs b/data_plane/src/precompute_engine/operators/mod.rs deleted file mode 100644 index 073db6e82..000000000 --- a/data_plane/src/precompute_engine/operators/mod.rs +++ /dev/null @@ -1,37 +0,0 @@ -pub mod count_min_sketch_accumulator; -pub mod count_min_sketch_with_heap_accumulator; -pub mod count_sketch_accumulator; -pub mod count_sketch_with_heap_accumulator; -pub mod datasketches_kll_accumulator; -pub mod dd_sketch_accumulator; -pub mod exact_accumulator; -pub mod hll_sketch_accumulator; -pub mod hydra_kll_accumulator; -pub mod increase_accumulator; -pub mod keyed_counter_state; -pub mod keyed_max_state; -pub mod keyed_min_state; -pub mod keyed_sum_count_accumulator; -pub mod max_accumulator; -pub mod min_accumulator; -pub mod sketch_envelope_accumulator; -pub mod sum_accumulator; -pub mod univmon_accumulator; - -pub use count_min_sketch_accumulator::*; -pub use count_min_sketch_with_heap_accumulator::*; -pub use count_sketch_accumulator::*; -pub use count_sketch_with_heap_accumulator::*; -pub use datasketches_kll_accumulator::*; -pub use dd_sketch_accumulator::*; -pub use hll_sketch_accumulator::*; -pub use hydra_kll_accumulator::*; -pub use increase_accumulator::*; -pub use keyed_counter_state::*; -pub use keyed_max_state::*; -pub use keyed_min_state::*; -pub use keyed_sum_count_accumulator::*; -pub use max_accumulator::*; -pub use min_accumulator::*; -pub use sketch_envelope_accumulator::*; -pub use sum_accumulator::*; diff --git a/data_plane/src/precompute_engine/operators/sketch_envelope_accumulator.rs b/data_plane/src/precompute_engine/operators/sketch_envelope_accumulator.rs deleted file mode 100644 index 08956e0ad..000000000 --- a/data_plane/src/precompute_engine/operators/sketch_envelope_accumulator.rs +++ /dev/null @@ -1,156 +0,0 @@ -//! 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::storage_engines::types::{AggregateCore, KeyByLabelValues, SerializableToSink}; -use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; -use prost::Message; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::AggregationType; -use asap_types::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::storage_engines::types::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/data_plane/src/precompute_engine/operators/sum_accumulator.rs b/data_plane/src/precompute_engine/operators/sum_accumulator.rs deleted file mode 100644 index d5ff3b02c..000000000 --- a/data_plane/src/precompute_engine/operators/sum_accumulator.rs +++ /dev/null @@ -1,413 +0,0 @@ -use crate::storage_engines::types::{ - AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use std::collections::HashMap; - -use asap_types::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: asap_types::Statistic, - _key: &Option, - _query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::storage_engines::types::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 asap_types::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/data_plane/src/precompute_engine/operators/univmon_accumulator.rs b/data_plane/src/precompute_engine/operators/univmon_accumulator.rs deleted file mode 100644 index b216f8f54..000000000 --- a/data_plane/src/precompute_engine/operators/univmon_accumulator.rs +++ /dev/null @@ -1,236 +0,0 @@ -//! One frequency state shared by count, distinct, L2 and entropy readouts. - -use crate::storage_engines::types::{ - AggregateCore, AuxStats, KeyByLabelValues, SerializableToSink, -}; -use asap_sketchlib::{DataInput, UnivMon}; -use asap_types::{AggregationType, Statistic}; -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/data_plane/src/precompute_engine/output_sink.rs b/data_plane/src/precompute_engine/output_sink.rs index 940df03e0..d32388b78 100644 --- a/data_plane/src/precompute_engine/output_sink.rs +++ b/data_plane/src/precompute_engine/output_sink.rs @@ -335,9 +335,9 @@ impl OutputSink for NoopOutputSink { #[cfg(test)] mod tests { use super::*; - use crate::precompute_engine::operators::{DDSketchAccumulator, SumAccumulator}; use crate::storage_engines::sketch_db::index::{AggKind, SeriesLookup}; use crate::storage_engines::types::{KeyByLabelValues, StreamingConfig}; + use asap_physical_operators::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 75bdad529..5d94407a9 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 super::accumulator_factory::{create_planner_accumulator, AccumulatorUpdater}; use crate::storage_engines::types::KeyByLabelValues; +use asap_physical_operators::factory::{create_planner_accumulator, AccumulatorUpdater}; use asap_types::{executable_plan::BackendNodeBinding, PrecomputeMaterialization}; use planner_types::post_asap::{ EdgeRole, ExecutableOperatorPayload, GroupingStrategy, PostAsapNodeId, SummaryFamilyType, @@ -224,6 +224,7 @@ impl RawDagProgram { } pub fn uses_counter_delta(&self) -> bool { + // Planner represents rate computation as an explicit upstream operator. false } diff --git a/data_plane/src/precompute_engine/worker.rs b/data_plane/src/precompute_engine/worker.rs index 847245e13..ae9db38fe 100644 --- a/data_plane/src/precompute_engine/worker.rs +++ b/data_plane/src/precompute_engine/worker.rs @@ -1,16 +1,16 @@ -#[cfg(test)] -use crate::precompute_engine::accumulator_factory::create_fixture_accumulator; -use crate::precompute_engine::accumulator_factory::AccumulatorUpdater; use crate::precompute_engine::config::LateDataPolicy; use crate::precompute_engine::group_key::GroupKey; use crate::precompute_engine::metrics::record_late_input; -use crate::precompute_engine::operators::sum_accumulator::SumAccumulator; use crate::precompute_engine::output_sink::OutputSink; use crate::precompute_engine::series_router::WorkerMessage; use crate::precompute_engine::window_manager::WindowManager; use crate::storage_engines::types::{ AggregateCore, KeyByLabelValues, PrecomputedOutput, StreamingConfigHandle, }; +#[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_types::aggregation_config::PrecomputeMaterialization; use asap_types::PolicyFingerprint; use asap_types::SampleUpdateRule; @@ -1874,11 +1874,11 @@ mod tests { // ----------------------------------------------------------------------- use crate::precompute_engine::config::LateDataPolicy; - use crate::precompute_engine::operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; - use crate::precompute_engine::operators::keyed_sum_count_accumulator::KeyedSumCountAccumulator; - use crate::precompute_engine::operators::sum_accumulator::SumAccumulator; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::StreamingConfig; + 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_types::enums::WindowKind; use asap_types::AggregationType; @@ -3169,7 +3169,7 @@ mod tests { // OTLP ingest dispatch builds via `decode_modified_otlp_sketch_bytes`. // ----------------------------------------------------------------------- - use crate::precompute_engine::operators::DDSketchAccumulator; + use asap_physical_operators::summary_kernels::DDSketchAccumulator; use asap_sketchlib::DdSketch; /// Build a fresh DDSketch holding `vals` so each test has a real, @@ -3934,7 +3934,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 crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::summary_kernels::SumAccumulator; let config = make_agg_config( 1, "gauge", @@ -3988,7 +3988,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 crate::precompute_engine::operators::IncreaseAccumulator; + use asap_physical_operators::summary_kernels::IncreaseAccumulator; let config = make_agg_config( 1, "requests_total", @@ -4280,9 +4280,9 @@ mod tests { #[cfg(test)] mod dag_execution_tests { use super::*; - use crate::precompute_engine::operators::exact_accumulator::ExactAccumulator; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::StreamingConfig; + use asap_physical_operators::summary_kernels::exact::ExactAccumulator; use asap_types::query_plan::ExactReadout; fn plan(query: &str) -> control_plane::physical::compiler::CompiledPhysicalPlan { @@ -4451,9 +4451,6 @@ mod dag_execution_tests { let error = StreamingConfig::from_precompute_plan(plan) .unwrap_err() .to_string(); - assert!( - error.contains("update") || error.contains("semantic definition differs"), - "{error}" - ); + assert!(error.contains("update"), "{error}"); } } 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 0ef02672d..1544a67b1 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 @@ -361,10 +361,8 @@ mod tests { } } - use crate::{ - precompute_engine::operators::SumAccumulator, - storage_engines::sketch_db::index::{AggKind, Capability, SummarySeriesMetadata}, - }; + use crate::storage_engines::sketch_db::index::{AggKind, Capability, SummarySeriesMetadata}; + use asap_physical_operators::summary_kernels::SumAccumulator; use asap_types::query_plan::{ ClickHousePlanningContext, ExactReadout, ExternalExactOutput, ExternalExactRequest, FallbackPolicy, FixedEvaluationRange, InstantExecution, MaterializationBinding, @@ -458,6 +456,7 @@ mod tests { QueryPlanNode::RelationalJoin { inputs: [left, external], join_kind: planner_types::pre_asap::JoinKind::Inner, + pruning: None, pred: serde_json::to_value(Predicate(Rc::new(QueryExpr::Compare { left: Rc::new(QueryExpr::Column(0)), op: CompareOpKind::Eq, @@ -668,7 +667,7 @@ mod tests { root, QueryPlanNode::Relational { input: sort, - operation: serde_json::to_value(ValueOperation::Limit { n: 1, offset: 0, partition_by: planner_types::pre_asap::GroupKeys::by(vec![]) }) + operation: serde_json::to_value(ValueOperation::Limit { n: 1, offset: 0, partition_by: planner_types::pre_asap::GroupKeys::none() }) .unwrap(), input_schema: projected_schema.clone(), output_schema: projected_schema, 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 65b43e033..3563a60d3 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 @@ -104,7 +104,11 @@ fn execute_relation_subtree( left_schema, right_schema, output_schema, + pruning, }) => { + if pruning.is_some() { + return Err("candidate pruning requires a certified semi-join binding".into()); + } if !matches!(join_kind, planner_types::pre_asap::JoinKind::Inner) { return Err("only inner relational joins are executable".into()); } diff --git a/data_plane/src/query_engines/asap_clickhouse_query_engine/relational_adapter.rs b/data_plane/src/query_engines/asap_clickhouse_query_engine/relational_adapter.rs index 8d673fff3..c0acadd8e 100644 --- a/data_plane/src/query_engines/asap_clickhouse_query_engine/relational_adapter.rs +++ b/data_plane/src/query_engines/asap_clickhouse_query_engine/relational_adapter.rs @@ -491,7 +491,7 @@ impl ClickHouseRelationalAdapter { } => { if !partition_by.keys().is_empty() || partition_by.is_without() { return Err(ClickHouseRelationalError::Unsupported( - "partitioned Limit".into(), + "partitioned relation Limit is not bound".into(), )); } input.rows = input.rows.into_iter().skip(*offset).take(*n).collect(); @@ -1616,7 +1616,7 @@ mod tests { ValueOperation::Limit { n: 1, offset: 0, - partition_by: planner_types::pre_asap::GroupKeys::by(vec![]), + partition_by: planner_types::pre_asap::GroupKeys::none(), }, ] { relation = adapter 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 cc6407ee8..05d1e8c26 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -286,7 +286,7 @@ impl ASAPQueryEngine { // Candidate-filtered exact cuts have a data dependency: read the // installed membership subtree once, then use that vector to build the // Prometheus selector. Keeping the result as a prepared leaf also means - // MembershipFilter reuses the same membership readout during composition. + // semi-join reuses the same membership readout during composition. let dependencies = super::exact_subqueries::external_dependencies(entry, times)?; let mut prepared = super::logical_dag::PreparedLeaves::new(); let unique_inputs = dependencies @@ -751,7 +751,11 @@ impl ASAPQueryEngine { if let Some(physical) = self.active_physical_plan_snapshot() { if let Ok(entry) = physical.query_plan.lookup(query) { if entry.nodes.values().any(|node| { - matches!(node, asap_types::query_plan::QueryPlanNode::Logical { .. }) + matches!( + node, + asap_types::query_plan::QueryPlanNode::Logical { .. } + | asap_types::query_plan::QueryPlanNode::PhysicalFragment { .. } + ) }) { return self .execute_logical_range(&physical, entry, start_ms, end_ms, step_ms) @@ -1394,11 +1398,10 @@ mod sketch_query_tests { #[cfg(test)] mod aux_pushdown_tests { use super::*; - use crate::precompute_engine::operators::{ - max_accumulator::MaxAccumulator, min_accumulator::MinAccumulator, - sum_accumulator::SumAccumulator, - }; use crate::storage_engines::types::AggregationType; + use asap_physical_operators::summary_kernels::{ + max::MaxAccumulator, min::MinAccumulator, sum::SumAccumulator, + }; use asap_types::Statistic; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; @@ -1613,9 +1616,9 @@ mod asap_tier_classify_tests { /// results instead of a CapabilityMiss. #[tokio::test] async fn execute_sum_by_zone_dispatches_to_exact_agg_reducer() { - use crate::precompute_engine::operators::sum_accumulator::SumAccumulator; use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; + use asap_physical_operators::summary_kernels::sum::SumAccumulator; let idx = Arc::new(SketchStore::new()); // Mirror the acceptance-test setup: four ExactAgg(Sum) sids, one @@ -2263,9 +2266,9 @@ mod asap_tier_classify_tests { /// `OuterFn::Plain` instant sums. #[tokio::test] async fn execute_instant_sum_accumulates_all_windows_not_last() { - use crate::precompute_engine::operators::sum_accumulator::SumAccumulator; use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; + use asap_physical_operators::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 b24908b45..aed8dd8b9 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 @@ -85,9 +85,12 @@ fn leaves( } _ => pending.extend(inputs.iter().map(|input| (*input, at))), }, - // MembershipFilter is a typed composition node rather than a Logical + // A join is a typed composition node rather than a Logical // wrapper, but its value input can still be a Prometheus leaf. - QueryPlanNode::MembershipFilter { inputs, .. } => { + QueryPlanNode::PhysicalFragment { inputs, .. } => { + pending.extend(inputs.iter().map(|input| (*input, at))); + } + QueryPlanNode::RelationalJoin { inputs, .. } => { pending.extend(inputs.iter().map(|input| (*input, at))); } QueryPlanNode::ExternalExact { request, inputs } => { @@ -650,21 +653,43 @@ mod tests { } #[tokio::test] - async fn candidate_exact_is_discovered_and_prepared_behind_membership_filter_root() { + async fn candidate_exact_is_discovered_and_prepared_behind_semi_join_root() { use asap_types::query_plan::CandidateCompleteness; let mut entry = candidate_entry("sum by (job) (rate(m[5m]))"); - entry.nodes.insert( - QueryNodeId(2), - QueryPlanNode::MembershipFilter { - inputs: [QueryNodeId(1), QueryNodeId(0)], - completeness: CandidateCompleteness::BestEffort { guarantee: None }, - }, - ); + entry.nodes.insert(QueryNodeId(2), { + let schema = planner_types::post_asap::SummarySchema { + fields: vec![planner_types::post_asap::SummaryField { + name: "job".into(), + dtype: planner_types::post_asap::SummaryFamilyType::Plain( + planner_types::pre_asap::DataType::Utf8, + ), + nullable: false, + }], + time_index: None, + }; + QueryPlanNode::RelationalJoin { + inputs: [QueryNodeId(0), QueryNodeId(1)], + join_kind: planner_types::pre_asap::JoinKind::Semi, + pred: serde_json::to_value(planner_types::pre_asap::Predicate(std::rc::Rc::new( + planner_types::pre_asap::QueryExpr::Compare { + left: std::rc::Rc::new(planner_types::pre_asap::QueryExpr::Column(0)), + op: planner_types::pre_asap::CompareOpKind::Eq, + right: std::rc::Rc::new(planner_types::pre_asap::QueryExpr::Column(1)), + }, + ))) + .unwrap(), + pruning: Some(CandidateCompleteness::BestEffort { guarantee: None }), + left_schema: schema.clone(), + right_schema: schema.clone(), + output_schema: schema, + } + }); entry.nodes.insert( QueryNodeId(3), QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { - k: 2, + operator: ResidualQueryOperator::Limit { + offset: 0, + n: 2, grouping: asap_types::query_plan::residual::Grouping { labels: vec![], without: false, @@ -888,13 +913,13 @@ mod tests { #[tokio::test] async fn five_minute_error_ratio_combines_prometheus_cut_with_summary_store() { - use crate::precompute_engine::operators::IncreaseAccumulator; use crate::query_engines::query_result::{InstantVectorElement, QueryResult}; use crate::storage_engines::sketch_db::{ data::AggKind, index::{Capability, SummarySeriesMetadata}, }; use crate::storage_engines::types::{KeyByLabelValues, Measurement}; + use asap_physical_operators::summary_kernels::IncreaseAccumulator; use asap_types::query_plan::{ residual::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 103247fdf..42f0974f8 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(crate::precompute_engine::operators::SumAccumulator::with_sum(value)), + Box::new(asap_physical_operators::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 31ef9f7c8..b14fe9d17 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,5 @@ //! Executes the installed typed logical DAG. No serving-time PromQL parsing. +mod native_values; use crate::query_engines::{ query_result::{InstantVectorElement, QueryResult}, EngineError, @@ -112,6 +113,14 @@ where memo: BTreeMap::new(), active: BTreeSet::new(), warnings: Vec::new(), + context: asap_physical_operators::dag::RunContext::new( + asap_physical_operators::dag::Scope::Query { + evaluation_time_ms: i64::try_from(at) + .map_err(|_| miss("evaluation timestamp overflow"))?, + revision: 0, + }, + asap_physical_operators::dag::Limits::default(), + )?, }; let at_signed = i64::try_from(at).map_err(|_| miss("evaluation timestamp overflow"))?; let evaluated = evaluator.eval(entry.root, at_signed)?; @@ -149,9 +158,13 @@ struct Evaluator<'a, F> { memo: BTreeMap<(QueryNodeId, i64), Value>, active: BTreeSet<(QueryNodeId, i64)>, warnings: Vec, + context: asap_physical_operators::dag::RunContext, } impl Result> Evaluator<'_, F> { fn eval(&mut self, id: QueryNodeId, at: i64) -> Result { + if self.context.is_cancelled() { + return Err(asap_physical_operators::Error::Cancelled.into()); + } if let Some(value) = self.memo.get(&(id, at)) { self.stats.memo_hits += 1; return Ok(value.clone()); @@ -168,11 +181,16 @@ impl Result> Evaluator<' self.memo.insert((id, at), value.clone()); return Ok(value); } - if self.active.len() >= 256 || !self.active.insert((id, at)) { - return Err(miss("cyclic or excessively deep installed DAG")); + if self.active.len() >= 256 || self.memo.len() >= 200_000 { + return Err(asap_physical_operators::Error::Operator( + "installed DAG evaluation budget exceeded".into(), + ) + .into()); } - if self.memo.len() >= 200_000 { - return Err(miss("installed DAG evaluation budget exceeded")); + if !self.active.insert((id, at)) { + return Err( + asap_physical_operators::Error::Invalid("cyclic installed DAG".into()).into(), + ); } let node = self .entry @@ -181,6 +199,38 @@ impl Result> Evaluator<' .ok_or_else(|| miss("missing installed node"))? .clone(); let value = match node { + QueryPlanNode::PhysicalFragment { + inputs, + dag, + row_input, + pruning, + } => { + let values = inputs + .iter() + .map(|id| self.eval(*id, at).and_then(vector)) + .collect::, _>>()?; + if let Some(contract) = &pruning { + native_values::validate_pruning( + &dag, + &values, + row_input, + contract, + at, + self.context.clone(), + )?; + if let Some(warning) = pruning_warning(Some(&contract.completeness)) { + self.warnings.push(warning); + } + } + Value::Vector(native_values::physical( + &dag, + values, + row_input, + at, + self.context.clone(), + )?) + } + QueryPlanNode::Scalar { value } => Value::Scalar(value), QueryPlanNode::Logical { operator: ResidualQueryOperator::CurrentSeries { .. }, @@ -205,17 +255,33 @@ impl Result> Evaluator<' } self.logical(operator, &inputs, at)? } - QueryPlanNode::MembershipFilter { + QueryPlanNode::RelationalJoin { inputs, - completeness, + join_kind: planner_types::pre_asap::JoinKind::Semi, + pred, + pruning, + left_schema, + right_schema, + output_schema, } => { - let candidates = vector(self.eval(inputs[0], at)?)?; - let values = vector(self.eval(inputs[1], at)?)?; - let (selected, warning) = membership_filter(candidates, values, &completeness)?; - if let Some(warning) = warning { + let values = vector(self.eval(inputs[0], at)?)?; + let candidates = vector(self.eval(inputs[1], at)?)?; + let predicate = serde_json::from_value(pred) + .map_err(|_| miss("invalid semi-join predicate"))?; + if let Some(warning) = pruning_warning(pruning.as_ref()) { self.warnings.push(warning); } - Value::Vector(selected) + Value::Vector(native_values::relation( + values, + candidates, + predicate, + std::sync::Arc::new(left_schema), + std::sync::Arc::new(right_schema), + std::sync::Arc::new(output_schema), + pruning, + at, + self.context.clone(), + )?) } _ => { self.stats.summary_readout_evaluations += 1; @@ -277,9 +343,19 @@ impl Result> Evaluator<' let values = vector(self.eval(input(0)?, at)?)?; Ok(Value::Vector(aggregate(operation, &grouping, values))) } - ResidualQueryOperator::TopKSelection { k, grouping } => { + ResidualQueryOperator::Limit { + n, + offset, + grouping, + } => { let values = vector(self.eval(input(0)?, at)?)?; - Ok(Value::Vector(topk_selection(k, &grouping, values))) + Ok(Value::Vector(native_values::limit( + values, + &grouping, + n, + offset, + self.context.clone(), + )?)) } ResidualQueryOperator::Binary { operation, @@ -350,22 +426,17 @@ impl Result> Evaluator<' .collect(), )) } - ResidualQueryOperator::Sort { descending } => { - let mut values = vector(self.eval(input(0)?, at)?)?; - values.sort_by(|a, b| { - if a.1.is_nan() && b.1.is_nan() { - std::cmp::Ordering::Equal - } else if a.1.is_nan() { - std::cmp::Ordering::Greater - } else if b.1.is_nan() { - std::cmp::Ordering::Less - } else if descending { - b.1.total_cmp(&a.1) - } else { - a.1.total_cmp(&b.1) - } - }); - Ok(Value::Vector(values)) + ResidualQueryOperator::Sort { + descending, + grouping, + } => { + let values = vector(self.eval(input(0)?, at)?)?; + Ok(Value::Vector(native_values::sort( + values, + &grouping, + descending, + self.context.clone(), + )?)) } ResidualQueryOperator::HistogramQuantile => { let Value::Scalar(quantile) = self.eval(input(0)?, at)? else { @@ -420,40 +491,81 @@ impl Result> Evaluator<' } } -fn membership_filter( - candidates: Vector, - values: Vector, - completeness: &CandidateCompleteness, -) -> Result<(Vector, Option), EngineError> { - let identity = |labels: &Labels| { - let mut labels = labels.clone(); - labels.remove("__name__"); - labels - }; - let candidate_ids: BTreeSet<_> = candidates - .iter() - .map(|(labels, _)| identity(labels)) - .collect(); - let value_ids: BTreeSet<_> = values.iter().map(|(labels, _)| identity(labels)).collect(); - let missing: BTreeSet<_> = candidate_ids.difference(&value_ids).collect(); - let selected = values - .into_iter() - .filter(|(labels, _)| candidate_ids.contains(&identity(labels))) - .collect(); - if !missing.is_empty() && matches!(completeness, CandidateCompleteness::Certified { .. }) { - return Err(miss("certified membership key has no authoritative value")); - } - let warning = match completeness { - CandidateCompleteness::Certified { .. } => None, - CandidateCompleteness::BestEffort { guarantee } => Some(match guarantee { +fn pruning_warning(completeness: Option<&CandidateCompleteness>) -> Option { + match completeness { + None | Some(CandidateCompleteness::Certified { .. }) => None, + Some(CandidateCompleteness::BestEffort { guarantee }) => Some(match guarantee { Some(guarantee) => format!( "ASAP membership pruning is approximate: {:?}", guarantee.metric ), None => "ASAP membership pruning is approximate and uncertified".into(), }), + } +} + +#[cfg(test)] +fn semi_join( + candidates: Vector, + values: Vector, + keys: &[(String, String)], + completeness: Option<&CandidateCompleteness>, +) -> Result<(Vector, Option), EngineError> { + use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::{CompareOpKind, DataType, Predicate, QueryExpr}, }; - Ok((selected, warning)) + use std::{rc::Rc, sync::Arc}; + let schema = |right: bool| { + Arc::new(SummarySchema { + fields: keys + .iter() + .map(|(l, r)| { + let name = if right { r } else { l }; + SummaryField { + name: name.clone(), + dtype: SummaryFamilyType::Plain(if name == "value" { + DataType::Float64 + } else { + DataType::Utf8 + }), + nullable: false, + } + }) + .collect(), + time_index: None, + }) + }; + let left = schema(false); + let right = schema(true); + let predicate = Predicate(Rc::new(QueryExpr::BoolAnd( + (0..keys.len()) + .map(|i| QueryExpr::Compare { + left: Rc::new(QueryExpr::Column(i)), + op: CompareOpKind::Eq, + right: Rc::new(QueryExpr::Column(keys.len() + i)), + }) + .collect(), + ))); + let context = asap_physical_operators::dag::RunContext::new( + asap_physical_operators::dag::Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Default::default(), + )?; + let rows = native_values::relation( + values, + candidates, + predicate, + left.clone(), + right, + left, + completeness.cloned(), + 0, + context, + )?; + Ok((rows, pruning_warning(completeness))) } fn aggregate(operation: Aggregation, grouping: &Grouping, values: Vector) -> Vector { @@ -511,31 +623,24 @@ fn grouping_key(labels: &Labels, grouping: &Grouping) -> Labels { /// Select by the child sample value while retaining every selected series' /// labels. NaN ranks below every numeric value, matching Prometheus' TOPK heap. /// Stable sorting also leaves equal-valued series in the child's order. +#[cfg(test)] fn topk_selection(k: u64, grouping: &Grouping, values: Vector) -> Vector { - if k == 0 { - return Vec::new(); - } - let mut groups: BTreeMap = BTreeMap::new(); - for (labels, value) in values { - groups - .entry(grouping_key(&labels, grouping)) - .or_default() - .push((labels, value)); - } - let limit = usize::try_from(k).unwrap_or(usize::MAX); - groups - .into_values() - .flat_map(|mut group| { - group.sort_by(|a, b| match (a.1.is_nan(), b.1.is_nan()) { - (true, true) => std::cmp::Ordering::Equal, - (true, false) => std::cmp::Ordering::Greater, - (false, true) => std::cmp::Ordering::Less, - (false, false) => b.1.total_cmp(&a.1), - }); - group.truncate(limit); - group - }) - .collect() + let context = asap_physical_operators::dag::RunContext::new( + asap_physical_operators::dag::Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Default::default(), + ) + .unwrap(); + native_values::limit( + native_values::sort(values, grouping, true, context.clone()).unwrap(), + grouping, + k, + 0, + context, + ) + .unwrap() } fn binary( @@ -782,7 +887,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 = crate::precompute_engine::operators::sum_accumulator::SumAccumulator::new(); + let mut sum = asap_physical_operators::summary_kernels::sum::SumAccumulator::new(); sum.update(1e308); sum.update(1e308); assert!(binary( @@ -1093,8 +1198,22 @@ mod topk_tests { ( root, QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { - k: 2, + operator: ResidualQueryOperator::Limit { + offset: 0, + n: 2, + grouping: Grouping { + labels: vec![], + without: false, + }, + }, + inputs: vec![QueryNodeId(98)], + }, + ), + ( + QueryNodeId(98), + QueryPlanNode::Logical { + operator: ResidualQueryOperator::Sort { + descending: true, grouping: Grouping { labels: vec![], without: false, @@ -1166,6 +1285,52 @@ mod topk_tests { } } + // Numeric boundary fields must reach Planner as numbers, never absent labels. + #[test] + fn semi_join_does_not_match_unequal_numeric_samples() { + let (rows, _) = semi_join( + vec![(Labels::new(), 2.0)], + vec![(Labels::new(), 1.0)], + &[("value".into(), "value".into())], + None, + ) + .unwrap(); + assert!(rows.is_empty(), "unequal numeric samples matched: {rows:?}"); + } + + #[test] + fn native_join_preserves_renamed_and_multiple_typed_keys() { + let left = vec![ + ( + labels(&[("instance", "a"), ("zone", "east"), ("extra", "kept")]), + 1., + ), + (labels(&[("instance", "a"), ("zone", "west")]), 2.), + ]; + let right = vec![(labels(&[("pod", "a"), ("region", "east")]), 99.)]; + let (selected, _) = semi_join( + right, + left.clone(), + &[ + ("instance".into(), "pod".into()), + ("zone".into(), "region".into()), + ], + None, + ) + .unwrap(); + assert_eq!(selected, vec![left[0].clone()]); + for (a, b, matched) in [(f64::NAN, f64::NAN, false), (-0., 0., true), (1., 1., true)] { + let (rows, _) = semi_join( + vec![(Labels::new(), b)], + vec![(Labels::new(), a)], + &[("value".into(), "value".into())], + None, + ) + .unwrap(); + assert_eq!(!rows.is_empty(), matched); + } + } + #[test] fn candidate_sidecar_intersects_then_reranks_exact_values() { let candidates = vec![ @@ -1177,12 +1342,13 @@ mod topk_tests { (labels(&[("pod", "b")]), 8.0), (labels(&[("pod", "c")]), 9.0), ]; - let (selected, warning) = membership_filter( + let (selected, warning) = semi_join( candidates, exact, - &CandidateCompleteness::Certified { + &[("pod".into(), "pod".into())], + Some(&CandidateCompleteness::Certified { guarantee: topk_membership_guarantee(), - }, + }), ) .unwrap(); let selected = topk_selection( @@ -1228,20 +1394,59 @@ mod topk_tests { reason: "prepared exact counter readout".into(), }, ), - ( - filter, - QueryPlanNode::MembershipFilter { - inputs: [candidate_id, value_id], - completeness: CandidateCompleteness::Certified { + (filter, { + let schema = planner_types::post_asap::SummarySchema { + fields: vec![planner_types::post_asap::SummaryField { + name: "pod".into(), + dtype: planner_types::post_asap::SummaryFamilyType::Plain( + planner_types::pre_asap::DataType::Utf8, + ), + nullable: false, + }], + time_index: None, + }; + QueryPlanNode::RelationalJoin { + inputs: [value_id, candidate_id], + join_kind: planner_types::pre_asap::JoinKind::Semi, + pred: serde_json::to_value(planner_types::pre_asap::Predicate( + std::rc::Rc::new(planner_types::pre_asap::QueryExpr::Compare { + left: std::rc::Rc::new(planner_types::pre_asap::QueryExpr::Column( + 0, + )), + op: planner_types::pre_asap::CompareOpKind::Eq, + right: std::rc::Rc::new( + planner_types::pre_asap::QueryExpr::Column(1), + ), + }), + )) + .unwrap(), + pruning: Some(CandidateCompleteness::Certified { guarantee: topk_membership_guarantee(), + }), + left_schema: schema.clone(), + right_schema: schema.clone(), + output_schema: schema, + } + }), + ( + root, + QueryPlanNode::Logical { + operator: ResidualQueryOperator::Limit { + offset: 0, + n: 1, + grouping: Grouping { + labels: vec![], + without: false, + }, }, + inputs: vec![QueryNodeId(98)], }, ), ( - root, + QueryNodeId(98), QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { - k: 1, + operator: ResidualQueryOperator::Sort { + descending: true, grouping: Grouping { labels: vec![], without: false, @@ -1304,23 +1509,25 @@ mod topk_tests { fn uncertified_candidate_sidecar_warns_or_falls_back_explicitly() { let candidates = vec![(labels(&[("pod", "a")]), 1.0)]; let exact = vec![(labels(&[("pod", "a")]), 2.0)]; - let (_, warning) = membership_filter( + let (_, warning) = semi_join( candidates.clone(), exact.clone(), - &CandidateCompleteness::BestEffort { guarantee: None }, + &[("pod".into(), "pod".into())], + Some(&CandidateCompleteness::BestEffort { guarantee: None }), ) .unwrap(); assert!(warning.unwrap().contains("approximate")); - // Exact queries never lower an uncertified MembershipFilter. The Planner + // Exact queries never lower an uncertified pruning semi-join. The Planner // emits its ordinary exact fallback instead; this runtime node is only // valid for certified or explicitly approximate plans. let certified = CandidateCompleteness::Certified { guarantee: topk_membership_guarantee(), }; - assert!(membership_filter( + assert!(semi_join( vec![(labels(&[("pod", "missing")]), 1.0)], exact, - &certified, + &[("pod".into(), "pod".into())], + Some(&certified), ) .is_err()); } 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 new file mode 100644 index 000000000..621562ef6 --- /dev/null +++ b/data_plane/src/query_engines/asap_query_engine/logical_dag/native_values.rs @@ -0,0 +1,568 @@ +//! Bind protocol vectors to native batch operators; computation stays in Planner. +use super::{grouping_key, miss, EngineError, Grouping, Labels, Vector}; +use asap_physical_operators::dag::{ + self, batch_execution, + operators::{Operator, SortKey}, + values::{Batch, Schema, Value}, +}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, +}; +use std::sync::Arc; +fn schema(fields: &[(&str, DataType)]) -> Schema { + Arc::new(SummarySchema { + fields: fields + .iter() + .map(|(name, dtype)| SummaryField { + name: (*name).into(), + dtype: SummaryFamilyType::Plain(dtype.clone()), + nullable: false, + }) + .collect(), + time_index: None, + }) +} +fn key(value: &T) -> Value { + Value::Utf8( + serde_json::to_string(value) + .expect("string keys serialize") + .into(), + ) +} +fn ranked_batch(values: &Vector, grouping: &Grouping) -> Result { + let schema = schema(&[ + ("index", DataType::Int64), + ("group", DataType::Utf8), + ("value", DataType::Float64), + ]); + Batch::try_new( + schema, + values + .iter() + .enumerate() + .map(|(i, (labels, v))| { + vec![ + Value::Int64(i as i64), + key(&grouping_key(labels, grouping)), + Value::Float64(*v), + ] + }) + .collect(), + ) + .map_err(EngineError::from) +} +fn output(values: Vector, batches: Vec>) -> Result { + batches + .iter() + .flat_map(|b| b.rows()) + .map(|row| match row.first() { + Some(Value::Int64(index)) => values + .get(*index as usize) + .cloned() + .ok_or_else(|| miss("native result index outside input")), + _ => Err(miss("native result has no row identity")), + }) + .collect() +} +pub(super) fn sort( + values: Vector, + grouping: &Grouping, + descending: bool, + context: dag::RunContext, +) -> Result { + let batch = ranked_batch(&values, grouping)?; + let op = Operator::sort( + batch.schema().clone(), + vec![SortKey { + column: 2, + descending, + nulls_first: false, + }], + vec![1], + ) + .map_err(EngineError::from)?; + let result = + batch_execution::evaluate_batch(batch, vec![op], context).map_err(EngineError::from)?; + output(values, result) +} +pub(super) fn limit( + values: Vector, + grouping: &Grouping, + n: u64, + offset: u64, + context: dag::RunContext, +) -> Result { + let batch = ranked_batch(&values, grouping)?; + let op = + Operator::limit(batch.schema().clone(), n, offset, vec![1]).map_err(EngineError::from)?; + let result = + batch_execution::evaluate_batch(batch, vec![op], context).map_err(EngineError::from)?; + output(values, result) +} +/// Relational boundary used by explicit row plans. Planner owns predicate lowering. +pub(super) fn relation( + values: Vector, + candidates: Vector, + predicate: planner_types::pre_asap::Predicate, + left: Schema, + right: Schema, + output_schema: Schema, + completeness: Option, + at: i64, + context: dag::RunContext, +) -> Result { + use asap_physical_operators::physical_planner::{ + compile_node, CompiledPhysicalDag, InputContract, + }; + use planner_types::post_asap::{ + ExecutableDagNode, ExecutableOperatorPayload, ExecutionDataState, PostAsapNodeId, + }; + let pruning = completeness + .as_ref() + .map(|completeness| { + Ok::<_, EngineError>(asap_types::query_plan::PruningInputContract { + candidate_input: 1, + keys: asap_physical_operators::physical_planner::equijoin_keys( + &predicate, &left, &right, + )?, + completeness: completeness.clone(), + }) + }) + .transpose()?; + let node = ExecutableDagNode { + id: PostAsapNodeId(2), + output_state: ExecutionDataState::QUERY_ROWS, + output_schema: (*output_schema).clone(), + guarantee: None, + payload: ExecutableOperatorPayload::RelationalJoin { + join_kind: planner_types::pre_asap::JoinKind::Semi, + pred: predicate, + pruning: completeness, + }, + }; + let operator = compile_node(&node, &[left.clone(), right.clone()])?; + let compiled = CompiledPhysicalDag::from_operators( + [ + (0, InputContract::bounded(left)), + (1, InputContract::bounded(right)), + ] + .into(), + [(2, (vec![0, 1], operator))].into(), + vec![2], + )?; + let encoded = compiled.encode()?; + let inputs = vec![values, candidates]; + if let Some(pruning) = pruning { + validate_pruning(&encoded, &inputs, 0, &pruning, at, context.clone())?; + } + physical(&encoded, inputs, 0, at, context) +} + +// Equality keys canonicalize signed zero and NaNs; row transport must preserve their bits. +fn identity_key(value: &Value) -> Result, asap_physical_operators::Error> { + if let Value::Float64(number) = value { + let mut key = vec![0xff]; + key.extend_from_slice(&number.to_bits().to_be_bytes()); + Ok(key) + } else { + value.key() + } +} + +/// Bind protocol values to the selected physical input contracts; no operator lowering. +pub(super) fn physical( + encoded: &[u8], + inputs: Vec, + row_input: usize, + at: i64, + context: dag::RunContext, +) -> Result { + use asap_physical_operators::physical_planner::{CompiledPhysicalDag, Source}; + use futures::{FutureExt, StreamExt}; + use std::collections::{BTreeMap, VecDeque}; + let compiled = CompiledPhysicalDag::decode(encoded)?; + let contracts = compiled.input_contracts().collect::>(); + if contracts.len() != inputs.len() || row_input >= inputs.len() { + return Err(asap_physical_operators::Error::Invalid( + "physical input arity mismatch".into(), + ) + .into()); + } + let mut identities: BTreeMap>, VecDeque<(Labels, f64)>> = BTreeMap::new(); + let mut sources = BTreeMap::new(); + for (position, ((id, contract), values)) in contracts.iter().zip(inputs).enumerate() { + let rows = values + .iter() + .map(|(labels, sample)| { + contract + .schema + .fields + .iter() + .enumerate() + .map(|(column, field)| { + if Some(column) == contract.schema.time_index { + return Ok(Value::Timestamp(at)); + } + match &field.dtype { + SummaryFamilyType::Plain(DataType::Float64) => { + Ok(Value::Float64(*sample)) + } + SummaryFamilyType::Plain(DataType::Int64) => { + if sample.is_finite() && sample.fract() == 0. && sample.abs() <= (1_u64 << 53) as f64 { + Ok(Value::Int64(*sample as i64)) + } else { + Err(asap_physical_operators::Error::Invalid("protocol sample cannot represent the required Int64 input exactly".into()).into()) + } + } + SummaryFamilyType::Plain(DataType::Utf8) => { + Ok(labels.get(&field.name).map_or_else( + || { + if field.nullable { + Value::Null + } else { + Value::Utf8("".into()) + } + }, + |value| Value::Utf8(value.clone().into()), + )) + } + _ => Err(miss(format!( + "PromQL input cannot supply field {} of type {:?}", + field.name, field.dtype + ))), + } + }) + .collect::, EngineError>>() + }) + .collect::, _>>()?; + if position == row_input { + for (row, original) in rows.iter().zip(values) { + let key = row + .iter() + .map(identity_key) + .collect::, _>>()?; + identities.entry(key).or_default().push_back(original); + } + } + let batch = Batch::try_new(contract.schema.clone(), rows)?; + sources.insert( + *id, + Box::new(Operator::source(contract.schema.clone(), vec![batch])?) as Source<'_>, + ); + } + let graph = compiled.instantiate(sources)?; + let mut streams = graph.execute(compiled.roots(), context)?; + if streams.len() != 1 { + return Err( + asap_physical_operators::Error::Invalid("expected one physical output".into()).into(), + ); + } + let mut stream = streams.remove(0); + let mut result = Vec::new(); + loop { + match stream.next().now_or_never() { + Some(Some(batch)) => { + for row in batch?.rows() { + let key = row + .iter() + .map(identity_key) + .collect::, _>>()?; + let original = identities + .get_mut(&key) + .and_then(VecDeque::pop_front) + .ok_or_else(|| { + asap_physical_operators::Error::Invalid( + "physical row has no protocol identity".into(), + ) + })?; + result.push(original); + } + } + Some(None) => return Ok(result), + None => continue, + } + } +} + +pub(super) fn validate_pruning( + encoded: &[u8], + inputs: &[Vector], + row_input: usize, + contract: &asap_types::query_plan::PruningInputContract, + at: i64, + context: dag::RunContext, +) -> Result<(), EngineError> { + use asap_physical_operators::physical_planner::{CompiledPhysicalDag, InputContract}; + use planner_types::post_asap::CandidateCompleteness; + if !matches!( + contract.completeness, + CandidateCompleteness::Certified { .. } + ) { + return Ok(()); + } + let compiled = CompiledPhysicalDag::decode(encoded)?; + let schemas = compiled + .input_contracts() + .map(|(_, c)| c.schema.clone()) + .collect::>(); + let candidates = inputs + .get(contract.candidate_input) + .ok_or_else(|| miss("missing candidate input"))?; + let values = inputs + .get(row_input) + .ok_or_else(|| miss("missing authoritative input"))?; + let coverage = Operator::semi_join( + schemas[contract.candidate_input].clone(), + schemas[row_input].clone(), + contract.keys.iter().map(|&(l, r)| (r, l)).collect(), + )?; + let check = CompiledPhysicalDag::from_operators( + [ + ( + 0, + InputContract::bounded(schemas[contract.candidate_input].clone()), + ), + (1, InputContract::bounded(schemas[row_input].clone())), + ] + .into(), + [(2, (vec![0, 1], coverage))].into(), + vec![2], + )?; + let matched = physical( + &check.encode()?, + vec![candidates.clone(), values.clone()], + 0, + at, + context, + )?; + if matched.len() != candidates.len() { + return Err(miss("certified pruning key has no authoritative value")); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use asap_physical_operators::{ + physical_planner::{CompiledPhysicalDag, InputContract}, + Error, + }; + + // An available label is bound by Backend; Planner evaluates its predicate. + #[test] + fn planner_filter_compiles_and_executes_bound_labels() { + use asap_types::query_plan::{FallbackPolicy, InstantExecution, QueryPlanNode}; + use planner_types::{ + post_asap::{ + execution_data_state::lift_plain, ExecutionTiming, SummaryExpr, SummaryNode, + ValueOperation, + }, + pre_asap::QueryExpr, + }; + let query = "m{instance=\"pod\"}"; + let logical = control_plane::query_parser::parse_query_expr_canonical( + query, + planner_types::types::AccuracyTarget::Exact, + ) + .unwrap(); + let QueryExpr::TimeRange { child, .. } = logical else { + panic!("expected instant source") + }; + let QueryExpr::Scan { + source, + predicates, + schema, + } = child.as_ref() + else { + panic!("expected source selector") + }; + let pred = predicates[0].clone(); + let child = std::rc::Rc::new(QueryExpr::Scan { + source: source.clone(), + predicates: vec![], + schema: schema.clone(), + }); + let input = std::rc::Rc::new(SummaryNode { + schema: lift_plain(&child.output_schema().unwrap()), + guarantee: None, + expr: SummaryExpr::KeepPreAsap(child), + }); + let root = std::rc::Rc::new(SummaryNode { + schema: input.schema.clone(), + guarantee: None, + expr: SummaryExpr::ValueOperation { + child: input, + operation: ValueOperation::Filter { pred }, + timing: ExecutionTiming::QueryTime, + }, + }); + let entry = control_plane::query_plan::compile_bound_mapped( + "filter".into(), + query.into(), + &root, + InstantExecution { + lookback_ms: 0, + full_history: false, + cumulative_readout: false, + }, + FallbackPolicy::Reject, + |_, _| panic!("filter requires no materialization"), + |_, _| {}, + ) + .unwrap(); + let QueryPlanNode::PhysicalFragment { dag, row_input, .. } = &entry.nodes[&entry.root] + else { + panic!("filter was rejected: {:?}", entry.nodes) + }; + let values = vec![ + ( + [ + ("instance".into(), "pod".into()), + ("extra".into(), "kept".into()), + ] + .into(), + 7., + ), + ([("instance".into(), "other".into())].into(), 9.), + ]; + let expected = values[0].clone(); + assert_eq!( + physical(dag, vec![values], *row_input, 42, context(1 << 20)).unwrap(), + vec![expected] + ); + } + + fn sorted() -> Vec { + let input = schema(&[("value", DataType::Float64)]); + CompiledPhysicalDag::from_operators( + [(0, InputContract::bounded(input.clone()))].into(), + [( + 1, + ( + vec![0], + Operator::sort( + input, + vec![SortKey { + column: 0, + descending: false, + nulls_first: false, + }], + vec![], + ) + .unwrap(), + ), + )] + .into(), + vec![1], + ) + .unwrap() + .encode() + .unwrap() + } + fn context(max_bytes: usize) -> dag::RunContext { + dag::RunContext::new( + dag::Scope::Query { + evaluation_time_ms: 42, + revision: 7, + }, + dag::Limits { + max_bytes, + ..dag::Limits::default() + }, + ) + .unwrap() + } + fn cause(error: &Error) -> &Error { + match error { + Error::AtNode { source, .. } => cause(source), + other => other, + } + } + + // A real native execution failure must remain typed through the protocol adapter. + #[test] + fn physical_budget_and_cancellation_are_terminal() { + for cancelled in [false, true] { + let run = context(if cancelled { 4096 } else { 1 }); + if cancelled { + run.cancel(); + } + let error = physical( + &sorted(), + vec![vec![(Labels::new(), 2.), (Labels::new(), 1.)]], + 0, + 42, + run.clone(), + ) + .unwrap_err(); + let EngineError::Physical(error) = error else { + panic!("physical failure lost its type") + }; + assert!( + if cancelled { + matches!(cause(&error), Error::Cancelled) + } else { + matches!(cause(&error), Error::MemoryLimit) + }, + "{error}" + ); + assert_eq!(run.retained_bytes(), 0); + } + } + + // Count-like values are bound as integers only when the protocol sample is exact. + #[test] + fn integer_input_binding_preserves_type_and_rejects_rounding() { + let input = schema(&[("count", DataType::Int64)]); + let compiled = CompiledPhysicalDag::from_operators( + [(0, InputContract::bounded(input.clone()))].into(), + [(1, (vec![0], Operator::limit(input, 1, 0, vec![]).unwrap()))].into(), + vec![1], + ) + .unwrap() + .encode() + .unwrap(); + for (sample, valid) in [ + (3., true), + (-4., true), + (0.5, false), + (f64::INFINITY, false), + (9_007_199_254_740_994., false), + ] { + let result = physical( + &compiled, + vec![vec![(Labels::new(), sample)]], + 0, + 42, + context(4096), + ); + if valid { + assert_eq!(result.unwrap()[0].1, sample); + } else { + assert!(matches!( + result, + Err(EngineError::Physical(Error::Invalid(_))) + )); + } + } + } + + // Transport identity preserves the exact value selected by native total-order sorting. + #[test] + fn native_sort_preserves_signed_zero_bits() { + let rows = physical( + &sorted(), + vec![vec![(Labels::new(), 0.), (Labels::new(), -0.)]], + 0, + 42, + context(4096), + ) + .unwrap(); + assert_eq!( + rows.iter().map(|row| row.1.to_bits()).collect::>(), + vec![0.0_f64.to_bits(), (-0.0_f64).to_bits()] + ); + } +} 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 2ea8c6229..35f6e22c6 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 @@ -2,7 +2,7 @@ use std::collections::BTreeMap; -use crate::utils::arithmetic::evaluate_float64_arithmetic as arithmetic; +use asap_physical_operators::arithmetic::evaluate_float64_arithmetic as arithmetic; use asap_types::query_plan::{QueryNodeId, QueryPlanNode}; @@ -305,9 +305,9 @@ impl QueryNodeRuntime for PhysicalQueryRuntime<'_> { }) } QueryPlanNode::Logical { .. } - | QueryPlanNode::MembershipFilter { .. } | QueryPlanNode::Relational { .. } | QueryPlanNode::ExternalExact { .. } + | QueryPlanNode::PhysicalFragment { .. } | QueryPlanNode::RelationalJoin { .. } => Err(PhysicalNodeError::Fallback( "logical node requires installed logical runtime".into(), )), @@ -1050,7 +1050,7 @@ mod tests { 1, BTreeMap::new(), (1_000, 2_000), - Box::new(crate::precompute_engine::operators::SumAccumulator::with_sum(42.0)), + Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(42.0)), ); let config = test_plan::materialization("bytes_total", "Sum", serde_json::json!({}), &[], 1000); @@ -1148,7 +1148,9 @@ mod tests { BTreeMap::new(), bounds, Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum(sum), + asap_physical_operators::summary_kernels::SumAccumulator::with_sum( + sum, + ), ), ); } @@ -1216,7 +1218,7 @@ mod tests { BTreeMap::new(), (pane * 60_000, (pane + 1) * 60_000), Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum( + asap_physical_operators::summary_kernels::SumAccumulator::with_sum( (pane + 1) as f64, ), ), @@ -1291,7 +1293,7 @@ mod tests { BTreeMap::new(), (pane * 10_000, (pane + 1) * 10_000), Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum( + asap_physical_operators::summary_kernels::SumAccumulator::with_sum( (pane + 1) as f64, ), ), @@ -1358,7 +1360,7 @@ mod tests { 7, BTreeMap::new(), (pane * 10_000, (pane + 1) * 10_000), - Box::new(crate::precompute_engine::operators::SumAccumulator::with_sum(1.0)), + Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(1.0)), ); } assert!( @@ -1388,7 +1390,7 @@ mod tests { policy_fp: policy, }); use crate::storage_engines::types::Measurement; - let mut accumulator = crate::precompute_engine::operators::IncreaseAccumulator::new( + let mut accumulator = asap_physical_operators::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/summary_exec.rs b/data_plane/src/query_engines/asap_query_engine/summary_exec.rs index 803653718..9c174d95c 100644 --- a/data_plane/src/query_engines/asap_query_engine/summary_exec.rs +++ b/data_plane/src/query_engines/asap_query_engine/summary_exec.rs @@ -223,7 +223,7 @@ pub fn execute( SummaryExpr::SummaryJoin { .. } => Err(ExecError::NotYetSupported("SummaryJoin")), SummaryExpr::RelationalJoin { .. } => Err(ExecError::NotYetSupported("RelationalJoin")), - // MembershipFilter is lowered to the deployed QueryPlan DAG, where both + // semi-join is lowered to the deployed QueryPlan DAG, where both // row inputs retain labels for intersection and exact reranking. This // legacy generic adapter exposes opaque GroupKey values and cannot // implement that contract without losing label identity. 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 e25ea4fcd..7aa36ac0c 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 @@ -70,9 +70,6 @@ use planner_types::post_asap::{ }; use planner_types::pre_asap::{ColumnId, ColumnRef, QueryExpr, Reduction, Source}; -use crate::precompute_engine::operators::increase_accumulator::IncreaseAccumulator; -use crate::precompute_engine::operators::max_accumulator::MaxAccumulator; -use crate::precompute_engine::operators::min_accumulator::MinAccumulator; use crate::storage_engines::sketch_db::data::{AggKind, SketchConfig, SketchTimeSeries}; use crate::storage_engines::sketch_db::index::{SketchSampleState, SketchStore}; use crate::storage_engines::sketch_db::query::delta_apply::{ @@ -81,6 +78,9 @@ use crate::storage_engines::sketch_db::query::delta_apply::{ use crate::storage_engines::types::{ AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, }; +use asap_physical_operators::summary_kernels::increase::IncreaseAccumulator; +use asap_physical_operators::summary_kernels::max::MaxAccumulator; +use asap_physical_operators::summary_kernels::min::MinAccumulator; /// Per-query, per-call execution context — constructed fresh for each /// incoming query (never shared across concurrent queries, never @@ -251,7 +251,7 @@ impl GroupState { let planner_state = entries.iter().flat_map(|w| w.values()).any(|a| { a.as_any() - .is::() + .is::() }); // Temporal exact summaries are the hot path for long-window // dashboards. Merge their concrete, fixed-size states in one batch @@ -1456,7 +1456,7 @@ mod tests { #[test] fn keyed_count_state_follows_planner_family_and_query_readout() { - use crate::precompute_engine::operators::KeyedSumCountAccumulator; + use asap_physical_operators::summary_kernels::KeyedSumCountAccumulator; use asap_types::query_plan::ExactReadout; let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); @@ -1971,9 +1971,9 @@ mod tests { /// One installed frequency summary merges panes before all four readouts. #[test] fn bound_univmon_merges_panes_for_four_readouts() { - use crate::precompute_engine::operators::univmon_accumulator::UnivMonAccumulator; use crate::storage_engines::sketch_db::index::SketchEncoding; use crate::storage_engines::types::SerializableToSink; + use asap_physical_operators::summary_kernels::univmon::UnivMonAccumulator; use asap_types::query_plan::{MaterializationBinding, PhysicalGrouping}; let index = SketchStore::new(); let fp = asap_types::PolicyFingerprint(701); @@ -3201,13 +3201,13 @@ mod tests { sid, BTreeMap::new(), (T0, T0 + 1000), - Box::new(crate::precompute_engine::operators::SumAccumulator::with_sum(10.0)), + Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(10.0)), ); idx.append_precompute( sid, BTreeMap::new(), (T0 + 1000, T0 + 2000), - Box::new(crate::precompute_engine::operators::SumAccumulator::with_sum(15.0)), + Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(15.0)), ); let child = scan_node("bytes_total", None); @@ -3307,13 +3307,13 @@ mod tests { 1, BTreeMap::new(), (T0, T0 + 1000), - Box::new(crate::precompute_engine::operators::SumAccumulator::with_sum(30.0)), + Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(30.0)), ); idx.append_precompute( 2, BTreeMap::new(), (T0, T0 + 1000), - Box::new(crate::precompute_engine::operators::SumAccumulator::with_sum(12.0)), + Box::new(asap_physical_operators::summary_kernels::SumAccumulator::with_sum(12.0)), ); let child = scan_node("bytes_total", None); @@ -3353,7 +3353,7 @@ mod tests { sid, BTreeMap::new(), (T0, T0 + 1000), - Box::new(crate::precompute_engine::operators::MaxAccumulator::new()), + Box::new(asap_physical_operators::summary_kernels::MaxAccumulator::new()), ); let child = scan_node("latency_max_ms", None); diff --git a/data_plane/src/query_engines/mod.rs b/data_plane/src/query_engines/mod.rs index 0470cb7db..1672e1095 100644 --- a/data_plane/src/query_engines/mod.rs +++ b/data_plane/src/query_engines/mod.rs @@ -55,6 +55,10 @@ use thiserror::Error; /// `CapabilityMiss` does not — the caller should escalate). #[derive(Debug, Error)] pub enum EngineError { + /// A failed physical execution is terminal; it must never select another engine. + #[error(transparent)] + Physical(#[from] asap_physical_operators::Error), + /// The engine has no aggregation that can answer this query. Mirrors /// `ASAPQueryEngine::handle_query` returning `None`. The router treats /// this as a "hard miss" and falls through to the next backend in diff --git a/data_plane/src/query_engines/routing/query_engine_routing.rs b/data_plane/src/query_engines/routing/query_engine_routing.rs index d0b6bc16f..89b6cb1d8 100644 --- a/data_plane/src/query_engines/routing/query_engine_routing.rs +++ b/data_plane/src/query_engines/routing/query_engine_routing.rs @@ -250,6 +250,9 @@ impl EngineRouter { ); return Ok((result, id)); } + Err(e @ EngineError::Physical(_)) => { + return Err(EngineRouterError::AllFailed { last: e }); + } Err(e) => { warn!( backend = ?backend, @@ -343,6 +346,9 @@ impl EngineRouter { ); return Ok((result, id)); } + Err(e @ EngineError::Physical(_)) => { + return Err(EngineRouterError::AllFailed { last: e }); + } Err(e) => { warn!( backend = ?backend, 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 8bfcf0cf0..b38df3268 100644 --- a/data_plane/src/storage_engines/sketch_db/backfill/processor.rs +++ b/data_plane/src/storage_engines/sketch_db/backfill/processor.rs @@ -572,7 +572,7 @@ mod tests { /// when given the same ordered samples. #[test] fn backfill_builds_bit_identical_sum_accumulator_to_live() { - use crate::precompute_engine::accumulator_factory::create_fixture_accumulator; + use crate::tests::accumulator_fixture::create_fixture_accumulator; let cfg = sum_config(1, "m", vec![]); @@ -931,8 +931,7 @@ mod tests { ), cfg.policy_fingerprint(), ); - let acc = - crate::precompute_engine::operators::sum_accumulator::SumAccumulator::with_sum(1.0); + let acc = asap_physical_operators::summary_kernels::sum::SumAccumulator::with_sum(1.0); let live_sid = store .ingest_precompute_for_agg_config( |metric, attrs, kind| resolver.resolve(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 6bf028c9c..f6dc241ac 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 @@ -4,15 +4,15 @@ //! It shares the pure accumulator factory and update primitives with live ingest //! so both paths use the same sketch semantics. -#[cfg(test)] -use crate::precompute_engine::accumulator_factory::{ - create_fixture_accumulator, AccumulatorUpdater, -}; #[cfg(test)] use crate::precompute_engine::worker::apply_sample; use crate::storage_engines::sketch_db::backfill::raw_sample_reader::RawSample; use crate::storage_engines::types::AggregateCore; #[cfg(test)] +use crate::tests::accumulator_fixture::create_fixture_accumulator; +#[cfg(test)] +use asap_physical_operators::factory::AccumulatorUpdater; +#[cfg(test)] use asap_types::aggregation_config::PrecomputeMaterialization; /// Construct the accumulator for one `(agg_id, window)` pair by @@ -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 crate::precompute_engine::operators::{ + use asap_physical_operators::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 e696fe525..475fc6b14 100644 --- a/data_plane/src/storage_engines/sketch_db/data/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/data/mod.rs @@ -472,20 +472,7 @@ impl AccuracyBound { /// Per-sample sketch state. Stored as the payload column inside the /// per-sid `SidStoreData` columnar storage. -#[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, -} +pub use asap_physical_operators::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 199b643d5..1d59bb02c 100644 --- a/data_plane/src/storage_engines/sketch_db/index/maintenance.rs +++ b/data_plane/src/storage_engines/sketch_db/index/maintenance.rs @@ -729,8 +729,8 @@ impl SketchStore { #[cfg(test)] mod tests { use super::*; - use crate::precompute_engine::operators::SumAccumulator; use crate::storage_engines::types::PrecomputedOutput; + use asap_physical_operators::summary_kernels::SumAccumulator; use asap_types::traits::SerializableToSink; #[test] @@ -1237,8 +1237,8 @@ mod tests { )]) ); let assert_complete_output = |store: &SketchStore| { - use crate::precompute_engine::operators::DDSketchAccumulator; use crate::storage_engines::sketch_db::data::SketchEncoding; + use asap_physical_operators::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 57fc7c953..3c9dc4812 100644 --- a/data_plane/src/storage_engines/sketch_db/index/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/index/mod.rs @@ -59,6 +59,7 @@ fn encoding_to_tag(enc: SketchEncoding) -> u8 { SketchEncoding::ProtoDelta => t::PROTO_DELTA, SketchEncoding::MsgpackFull => t::MSGPACK_FULL, SketchEncoding::MsgpackDelta => t::MSGPACK_DELTA, + SketchEncoding::NativeBatchV1 => t::NATIVE_BATCH_V1, } } @@ -72,6 +73,7 @@ fn tag_to_encoding(tag: u8) -> SketchEncoding { t::PROTO_DELTA => SketchEncoding::ProtoDelta, t::MSGPACK_FULL => SketchEncoding::MsgpackFull, t::MSGPACK_DELTA => SketchEncoding::MsgpackDelta, + t::NATIVE_BATCH_V1 => SketchEncoding::NativeBatchV1, // t::PROTO_FULL and t::UNKNOWN (legacy) both → Full. _ => SketchEncoding::ProtoFull, } @@ -91,13 +93,13 @@ fn reconstruct_exact_agg( type_name: &str, bytes: &[u8], ) -> Option> { - use crate::precompute_engine::operators::{ + use crate::storage_engines::types::AggregateCore; + use asap_physical_operators::summary_kernels::{ IncreaseAccumulator, KeyedCounterState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, }; - use crate::storage_engines::types::AggregateCore; match type_name { - "PlannerExactAccumulatorV1" => crate::precompute_engine::operators::exact_accumulator::ExactAccumulator::deserialize_from_bytes(bytes).ok().map(|a|Box::new(a) as Box), + "PlannerExactAccumulatorV1" => asap_physical_operators::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), @@ -1522,12 +1524,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 @@ -4413,7 +4415,7 @@ mod tests { #[test] fn precompute_payload_round_trips_through_storage() { - use crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::summary_kernels::SumAccumulator; let idx = SketchStore::new(); let cfg = SketchConfig::DDSketch { @@ -4454,7 +4456,7 @@ mod tests { #[test] fn query_precomputes_by_agg_returns_data_grouped_by_label_values() { - use crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::summary_kernels::SumAccumulator; let idx = SketchStore::new(); let cfg = SketchConfig::DDSketch { @@ -4608,7 +4610,7 @@ mod tests { assert!(sketch.as_sketch().is_some()); assert!(sketch.as_exact_agg().is_none()); - use crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::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()); @@ -5188,7 +5190,7 @@ mod tests { ) .unwrap(); assert!(inventory.instances.is_empty()); - store.register(meta_with_policy(510, fingerprint)); + store.register(meta_for_config(510, &state_config)); store.append_sample(510, BTreeMap::new(), (0, 10_000), sample(2)); assert_eq!(store.series_ids_for_policy(fingerprint), vec![510]); let incompatible = asap_types::summary_catalog::SummaryCatalog::from_materializations( @@ -5420,7 +5422,7 @@ mod tests { 850, BTreeMap::new(), (0, 30_000), - Box::new(crate::precompute_engine::operators::SumAccumulator::new()) + Box::new(asap_physical_operators::summary_kernels::SumAccumulator::new()) )); // A flusher that captured metadata before completion cannot reopen it. writer.upsert_all(&[stale_record]).unwrap(); @@ -5882,7 +5884,7 @@ mod tests { lv_zone("z0"), (s, s + 30_000), Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum( + asap_physical_operators::summary_kernels::SumAccumulator::with_sum( (i + 1) as f64, ), ), @@ -6101,7 +6103,9 @@ mod tests { lv_zone("z0"), (s, s + 30_000), Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum((i + 1) as f64), + asap_physical_operators::summary_kernels::SumAccumulator::with_sum( + (i + 1) as f64, + ), ), ); } @@ -6165,7 +6169,7 @@ mod tests { lv_zone("z0"), (s, s + 30_000), Box::new({ - let mut acc = crate::precompute_engine::operators::SumAccumulator::new(); + let mut acc = asap_physical_operators::summary_kernels::SumAccumulator::new(); acc.update((i + 1) as f64); acc.update(10.0); acc @@ -6358,8 +6362,8 @@ mod tests { // Flush and reopen must preserve Planner family rather than reconstructing Rate as Increase. #[test] fn planner_exact_families_survive_disk_eviction_and_restart() { - use crate::precompute_engine::operators::exact_accumulator::ExactAccumulator; use crate::storage_engines::types::{AggregateCore, AggregationType}; + use asap_physical_operators::summary_kernels::exact::ExactAccumulator; let kinds = [ AggregationType::Sum, AggregationType::Count, 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 f82f06a17..c988369ab 100644 --- a/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs +++ b/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs @@ -193,8 +193,8 @@ pub fn warn_if_retention_inverted( #[cfg(test)] mod tests { use super::*; - use crate::precompute_engine::operators::SumAccumulator; use crate::storage_engines::types::{AggregationType, StreamingConfig}; + use asap_physical_operators::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/part.rs b/data_plane/src/storage_engines/sketch_db/persistence/part.rs index 52078104c..c760668cb 100644 --- a/data_plane/src/storage_engines/sketch_db/persistence/part.rs +++ b/data_plane/src/storage_engines/sketch_db/persistence/part.rs @@ -105,6 +105,7 @@ pub mod encoding_tag { pub const PROTO_DELTA: u8 = 2; pub const MSGPACK_FULL: u8 = 3; pub const MSGPACK_DELTA: u8 = 4; + pub const NATIVE_BATCH_V1: u8 = 5; } /// One entry inside a decoded part. The `start_ts`/`end_ts`/`label` 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 f1698c2f8..8f5da9ef7 100644 --- a/data_plane/src/storage_engines/sketch_db/query/decoders.rs +++ b/data_plane/src/storage_engines/sketch_db/query/decoders.rs @@ -24,7 +24,7 @@ use asap_sketchlib::CountSketchWithHeap; use asap_sketchlib::CsHeapItem; use asap_sketchlib::MessagePackCodec; -use crate::precompute_engine::operators::count_min_sketch_with_heap_accumulator::CountMinSketchWithHeapAccumulator; +use asap_physical_operators::summary_kernels::count_min_sketch_with_heap::CountMinSketchWithHeapAccumulator; /// Decode a `CountMinSketch` from the modified-OTLP wire bytes. /// MSGPACK path round-trips `CountMinSketch::deserialize_msgpack`; 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 894b15ee6..05656df78 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 @@ -115,7 +115,7 @@ impl DeltaSketchKind { sketch_cols, layers, } => SummaryState::UnivMon( - crate::precompute_engine::operators::univmon_accumulator::UnivMonAccumulator::new( + asap_physical_operators::summary_kernels::univmon::UnivMonAccumulator::new( *heap_size as usize, *sketch_rows as usize, *sketch_cols as usize, @@ -167,7 +167,10 @@ fn decode_full( }, SketchEncoding::MsgpackFull, ) => { - let state = crate::precompute_engine::operators::univmon_accumulator::UnivMonAccumulator::from_bytes(bytes) + let state = + asap_physical_operators::summary_kernels::univmon::UnivMonAccumulator::from_bytes( + bytes, + ) .map_err(|e| e.to_string())?; if state.dimensions() != ( @@ -244,7 +247,7 @@ fn decode_full( /// folded across a window (or several) via delta application, or merged /// in from another sid's own reconstruction. pub enum SummaryState { - UnivMon(crate::precompute_engine::operators::univmon_accumulator::UnivMonAccumulator), + UnivMon(asap_physical_operators::summary_kernels::univmon::UnivMonAccumulator), Dd(DdSketch), Hll(HllSketch), Kll(KllSketch), @@ -321,7 +324,7 @@ impl SummaryState { } // Shape (2): bucket-delta proto → additive apply via the // SAME decoder the ingest delta path uses. - use crate::precompute_engine::operators::dd_sketch_accumulator::DDSketchAccumulator; + use asap_physical_operators::summary_kernels::dd_sketch::DDSketchAccumulator; let mut acc = DDSketchAccumulator { inner: std::mem::replace(sk, DdSketch::new(sk.alpha)), sample_p: 1.0, @@ -583,6 +586,9 @@ pub fn cumulative_summary_state( let mut rolling: Option = None; for (_window_end, state) in samples { match state.encoding { + SketchEncoding::NativeBatchV1 => { + return Err("native batch requires the physical batch reader".into()); + } SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull => { let new_state = decode_full(&kind, &state.bytes, state.encoding)?; rolling = Some(match rolling.take() { @@ -698,6 +704,9 @@ pub fn per_window_summary_states( } match state.encoding { + SketchEncoding::NativeBatchV1 => { + return Err("native batch requires the physical batch reader".into()); + } SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull => { // A Full (re)sets this window's base. rolling = Some(decode_full(&kind, &state.bytes, state.encoding)?); @@ -740,21 +749,21 @@ pub fn per_window_summary_states( // --------------------------------------------------------------------------- fn dd_from_proto(buffer: &[u8]) -> Result { - use crate::precompute_engine::operators::dd_sketch_accumulator::DDSketchAccumulator; + use asap_physical_operators::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::precompute_engine::operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; + use asap_physical_operators::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::precompute_engine::operators::hll_sketch_accumulator::HllSketchAccumulator; + use asap_physical_operators::summary_kernels::hll_sketch::HllSketchAccumulator; HllSketchAccumulator::from_sketchlib_proto_bytes(buffer) .map(|acc| acc.inner) .map_err(|e| e.to_string()) @@ -781,6 +790,17 @@ mod tests { use super::*; use asap_sketchlib::HllVariant; + #[test] + fn native_batches_are_not_legacy_sketch_frames() { + let state = SketchSampleState { + bytes: vec![], + encoding: SketchEncoding::NativeBatchV1, + }; + let samples = [(1000, &state)]; + assert!(cumulative_summary_state(&samples, DeltaSketchKind::Kll { k: 200 }).is_err()); + assert!(per_window_summary_states(&samples, DeltaSketchKind::Kll { k: 200 }).is_err()); + } + fn encode_dd(sk: &DdSketch) -> Vec { use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; use prost::Message; @@ -788,6 +808,7 @@ mod tests { 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)), @@ -910,7 +931,7 @@ mod tests { 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::precompute_engine::operators::hll_sketch_accumulator::HllSketchAccumulator; + use asap_physical_operators::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()); @@ -929,7 +950,7 @@ mod tests { #[test] fn dd_from_proto_matches_accumulator_decoder() { - use crate::precompute_engine::operators::dd_sketch_accumulator::DDSketchAccumulator; + use asap_physical_operators::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); @@ -946,7 +967,7 @@ mod tests { #[test] fn kll_from_proto_matches_accumulator_decoder() { - use crate::precompute_engine::operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; + use asap_physical_operators::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"); diff --git a/data_plane/src/storage_engines/types/key_by_label_values.rs b/data_plane/src/storage_engines/types/key_by_label_values.rs deleted file mode 100644 index 34bc84899..000000000 --- a/data_plane/src/storage_engines/types/key_by_label_values.rs +++ /dev/null @@ -1,164 +0,0 @@ -use serde::{Deserialize, Serialize}; -// use std::collections::HashMap; -use std::hash::{Hash, Hasher}; - -#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] -pub struct KeyByLabelValues { - // pub labels: HashMap, - pub labels: Vec, -} - -impl KeyByLabelValues { - pub fn new() -> Self { - Self { labels: Vec::new() } - } - - pub fn new_with_labels(labels: Vec) -> Self { - Self { labels } - } - - pub fn insert(&mut self, value: String) { - self.labels.push(value); - } - - pub fn get(&self, index: usize) -> Option<&String> { - self.labels.get(index) - } - - pub fn serialize_to_json(&self) -> serde_json::Value { - serde_json::to_value(&self.labels).unwrap_or(serde_json::Value::Null) - } - - pub fn deserialize_from_json(data: &serde_json::Value) -> Result { - let labels: Vec = serde_json::from_value(data.clone())?; - Ok(Self { labels }) - } - - pub fn serialize_to_bytes(&self) -> Vec { - bincode::serialize(&self.labels).unwrap_or_default() - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - let labels: Vec = bincode::deserialize(buffer)?; - Ok(Self { labels }) - } - - /// Encode labels as a semicolon-joined string — the canonical key format used - /// for all sketch hashing (CountMinSketch, HydraKLL, SetAggregator, DeltaSet). - pub fn to_semicolon_str(&self) -> String { - self.labels.join(";") - } - - #[cfg(test)] - /// Decode a semicolon-joined string back into a KeyByLabelValues. - pub fn from_semicolon_str(s: &str) -> Self { - Self { - labels: s.split(';').map(|s| s.to_string()).collect(), - } - } - - pub fn is_empty(&self) -> bool { - self.labels.is_empty() - } - - pub fn len(&self) -> usize { - self.labels.len() - } -} - -impl Hash for KeyByLabelValues { - fn hash(&self, state: &mut H) { - // Create a sorted vector of key-value pairs for consistent hashing - let mut sorted_pairs: Vec<_> = self.labels.iter().collect(); - sorted_pairs.sort(); - - for value in sorted_pairs { - value.hash(state); - } - } -} - -impl Default for KeyByLabelValues { - fn default() -> Self { - Self::new() - } -} - -impl std::fmt::Display for KeyByLabelValues { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "{{")?; - let mut first = true; - for value in &self.labels { - if !first { - write!(f, ", ")?; - } - write!(f, "{value}")?; - first = false; - } - write!(f, "}}") - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_key_by_label_values() { - let mut key = KeyByLabelValues::new(); - key.insert("localhost:8080".to_string()); - key.insert("prometheus".to_string()); - - assert_eq!(key.len(), 2); - assert_eq!(key.get(0), Some(&"localhost:8080".to_string())); - assert_eq!(key.get(1), Some(&"prometheus".to_string())); - } - - #[test] - fn test_serialization() { - let mut key = KeyByLabelValues::new(); - key.insert("test".to_string()); - - let json = key.serialize_to_json(); - let deserialized = KeyByLabelValues::deserialize_from_json(&json).unwrap(); - assert_eq!(key, deserialized); - } - - #[test] - fn test_byte_serialization() { - let mut key = KeyByLabelValues::new(); - key.insert("test".to_string()); - - let bytes = key.serialize_to_bytes(); - let deserialized = KeyByLabelValues::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!(key, deserialized); - } - - #[test] - fn test_semicolon_roundtrip() { - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string(), "prod".to_string()]); - assert_eq!(key.to_semicolon_str(), "web;prod"); - let roundtripped = KeyByLabelValues::from_semicolon_str("web;prod"); - assert_eq!(roundtripped, key); - } - - #[test] - fn test_hash_consistency() { - let mut key1 = KeyByLabelValues::new(); - key1.insert("a".to_string()); - key1.insert("b".to_string()); - - let mut key2 = KeyByLabelValues::new(); - key2.insert("b".to_string()); - key2.insert("a".to_string()); - - // Should hash to the same value regardless of insertion order - let mut hasher1 = std::collections::hash_map::DefaultHasher::new(); - let mut hasher2 = std::collections::hash_map::DefaultHasher::new(); - - key1.hash(&mut hasher1); - key2.hash(&mut hasher2); - - assert_eq!(hasher1.finish(), hasher2.finish()); - } -} diff --git a/data_plane/src/storage_engines/types/measurement.rs b/data_plane/src/storage_engines/types/measurement.rs deleted file mode 100644 index 0fe1abc0d..000000000 --- a/data_plane/src/storage_engines/types/measurement.rs +++ /dev/null @@ -1,94 +0,0 @@ -use serde::{Deserialize, Serialize}; -use std::ops::Add; - -#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] -pub struct Measurement { - pub value: f64, -} - -impl Measurement { - pub fn new(value: f64) -> Self { - Self { value } - } - - pub fn serialize_to_bytes(&self) -> Vec { - self.value.to_le_bytes().to_vec() - } - - pub fn serialize_to_json(&self) -> serde_json::Value { - serde_json::json!({ - "value": self.value - }) - } - - pub 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)) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - if buffer.len() < 8 { - return Err("Buffer too short for f64".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::new(value)) - } -} - -impl Add for Measurement { - type Output = Measurement; - - fn add(self, other: Measurement) -> Measurement { - Measurement::new(self.value + other.value) - } -} - -impl Add for &Measurement { - type Output = Measurement; - - fn add(self, other: &Measurement) -> Measurement { - Measurement::new(self.value + other.value) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_measurement_creation() { - let measurement = Measurement::new(42.5); - assert_eq!(measurement.value, 42.5); - } - - #[test] - fn test_measurement_addition() { - let m1 = Measurement::new(10.0); - let m2 = Measurement::new(20.0); - let result = m1 + m2; - assert_eq!(result.value, 30.0); - } - - #[test] - fn test_serialization() { - let measurement = Measurement::new(42.5); - let json = measurement.serialize_to_json(); - let deserialized = Measurement::deserialize_from_json(&json).unwrap(); - assert_eq!(measurement, deserialized); - } - - #[test] - fn test_byte_serialization() { - let measurement = Measurement::new(42.5); - let bytes = measurement.serialize_to_bytes(); - let deserialized = Measurement::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!(measurement, deserialized); - } -} diff --git a/data_plane/src/storage_engines/types/mod.rs b/data_plane/src/storage_engines/types/mod.rs index a91e83148..562d08d7f 100644 --- a/data_plane/src/storage_engines/types/mod.rs +++ b/data_plane/src/storage_engines/types/mod.rs @@ -8,21 +8,18 @@ pub mod enums; pub mod hot_reload_config; -pub mod key_by_label_values; -pub mod measurement; pub mod precomputed_output; pub mod storage_backend; pub mod streaming_config; -pub mod traits; +pub use asap_physical_operators::key_by_label_values::*; +pub use asap_physical_operators::measurement::*; +pub use asap_physical_operators::traits::*; pub use enums::*; pub use hot_reload_config::*; -pub use key_by_label_values::*; -pub use measurement::*; pub use precomputed_output::*; pub use storage_backend::*; pub use streaming_config::*; -pub use traits::*; // Cross-module re-export of asap_types data types so callers can // write `crate::storage_engines::types::PrecomputeMaterialization` instead of diff --git a/data_plane/src/storage_engines/types/traits.rs b/data_plane/src/storage_engines/types/traits.rs deleted file mode 100644 index 97f2c96df..000000000 --- a/data_plane/src/storage_engines/types/traits.rs +++ /dev/null @@ -1,351 +0,0 @@ -use crate::storage_engines::types::KeyByLabelValues; -use std::collections::HashMap; - -use asap_types::AggregationType; -use asap_types::Statistic; - -pub use asap_types::traits::SerializableToSink; - -/// 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/data_plane/src/tests/accumulator_fixture.rs b/data_plane/src/tests/accumulator_fixture.rs new file mode 100644 index 000000000..e8346cba4 --- /dev/null +++ b/data_plane/src/tests/accumulator_fixture.rs @@ -0,0 +1,265 @@ +//! Config fixtures for backend integration tests; production binds Planner payloads. +use asap_physical_operators::factory::*; +use asap_physical_operators::{AggregateCore, AggregationType}; +use asap_types::{accumulator_spec::cms_params, PrecomputeMaterialization}; +use planner_types::post_asap::{ExactKind, SketchAlgorithm, SketchParams, SummaryFamilyType}; +#[cfg(test)] +/// Return `true` if `config` produces a keyed (MultipleSubpopulation) updater, +/// without allocating an updater object. +/// +/// **Contract:** this must agree with every concrete `AccumulatorUpdater::is_keyed()` +/// implementation. When a new accumulator type is added, update both here and +/// in the corresponding struct. +pub fn config_is_keyed(config: &PrecomputeMaterialization) -> bool { + config + .accumulator_spec() + .expect("valid fixture") + .grouping + .is_some() +} + +/// Top-k ranking quantity, selected by `weight_mode` or its alias `topk_weight`. +/// +/// * `value` / `sum`: sum values per key (default). +/// * `count` / `frequency` / `freq`: count occurrences per key. +#[cfg(test)] +fn topk_weight_param(config: &PrecomputeMaterialization) -> TopkWeight { + match config.sample_update_rule() { + asap_types::SampleUpdateRule::Count => TopkWeight::Count, + asap_types::SampleUpdateRule::Value { .. } + | asap_types::SampleUpdateRule::CounterDelta { .. } => TopkWeight::Value, + } +} + +// --------------------------------------------------------------------------- +// Factory function +// --------------------------------------------------------------------------- + +/// Read the KLL `k` out of `SketchParams::Kll`. `accumulator_spec()` +/// always builds a `SketchKind` whose `SketchAlgorithm::Kll` is paired with +/// `SketchParams::Kll`, so the +/// other arm is unreachable from a `spec` this module builds itself. +#[cfg(test)] +fn kll_k(params: &SketchParams) -> u16 { + match params { + // Lossless: `accumulator_spec()` only ever stores a value that + // already fit in `u16` (via `kll_k_param`'s own `u16::try_from` + // fallback) widened to `u32`. + SketchParams::Kll { k } => *k as u16, + other => unreachable!( + "accumulator_spec() paired SketchAlgorithm::Kll with non-Kll params: {other:?}" + ), + } +} + +/// Read `(rows = depth, columns = width)` out of `SketchParams::Cms` or `::CountSketch` +/// — same shape, different variant per bare-sketch identity. +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:?}" + ), + } +} + +/// Read the DDSketch relative-accuracy `alpha` out of `SketchParams::DDSketch`. +#[cfg(test)] +fn ddsketch_alpha(params: &SketchParams) -> f64 { + match params { + SketchParams::DDSketch { alpha } => *alpha, + other => unreachable!( + "accumulator_spec() paired SketchAlgorithm::DDSketch with non-DDSketch params: {other:?}" + ), + } +} + +/// Construct isolated payload fixtures for kernel/storage unit tests. +/// Production execution requires a validated Planner DAG program. +#[cfg(test)] +pub fn create_fixture_accumulator( + config: &PrecomputeMaterialization, +) -> Box { + let spec = config + .accumulator_spec() + .expect("invalid isolated kernel fixture"); + + let keyed = spec.grouping.is_some(); + + match (&spec.family, keyed) { + (SummaryFamilyType::ExactAggregate(ExactKind::Sum | ExactKind::Count, _), false) => { + Box::new(SumAccumulatorUpdater::new()) + } + (SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), true) => { + Box::new(KeyedSumCountAccumulatorUpdater::for_family(ExactKind::Sum)) + } + (SummaryFamilyType::ExactAggregate(ExactKind::Count, _), true) => Box::new( + KeyedSumCountAccumulatorUpdater::for_family(ExactKind::Count), + ), + + // Direction comes off the family itself now. It used to be read + // back out of `aggregation_sub_type` because Planner had one + // `MinMax` accumulator for both directions, which meant a config + // whose sub_type was lost or misspelled silently built the wrong + // extremum. + (SummaryFamilyType::ExactAggregate(ExactKind::Min, _), false) => { + Box::new(MinAccumulatorUpdater::new()) + } + (SummaryFamilyType::ExactAggregate(ExactKind::Min, _), true) => { + Box::new(KeyedMinStateUpdater::new()) + } + (SummaryFamilyType::ExactAggregate(ExactKind::Max, _), false) => { + Box::new(MaxAccumulatorUpdater::new()) + } + (SummaryFamilyType::ExactAggregate(ExactKind::Max, _), true) => { + Box::new(KeyedMaxStateUpdater::new()) + } + + (SummaryFamilyType::ExactAggregate(ExactKind::Increase | ExactKind::Rate, _), false) => { + Box::new(IncreaseAccumulatorUpdater::new()) + } + (SummaryFamilyType::ExactAggregate(ExactKind::Increase | ExactKind::Rate, _), true) => { + Box::new(KeyedCounterStateUpdater::new()) + } + + (SummaryFamilyType::Sketch(kind, _), false) + if kind.algorithm() == &SketchAlgorithm::Kll => + { + Box::new(KllAccumulatorUpdater::new(kll_k(kind.params()))) + } + // HydraKLL: `k` comes off the typed params like the unkeyed case, + // but the `(row, col)` tiling grid has no `SketchParams::Kll` + // field to live in (see `asap_types::accumulator_spec`'s module + // doc) — read it the same way bare CMS does, via `cms_params`. + (SummaryFamilyType::Sketch(kind, _), true) if kind.algorithm() == &SketchAlgorithm::Kll => { + let (row_num, col_num) = cms_params(config); + Box::new(HydraKllAccumulatorUpdater::new( + row_num, + col_num, + kll_k(kind.params()), + )) + } + + // Bare CMS: point-frequency only, min-of-rows estimator. `keyed=false` + // can't actually arise here today (no `AggregationType` resolves to + // bare Cms unkeyed — see accumulator_spec.rs), matched anyway as a + // safe default. + (SummaryFamilyType::Sketch(kind, _), _) if kind.algorithm() == &SketchAlgorithm::Cms => { + let (row_num, col_num) = cms_dims(kind.params()); + Box::new(CmsAccumulatorUpdater::new(row_num, col_num)) + } + + // CountSketch uses the median-of-signed-rows estimator. + (SummaryFamilyType::Sketch(kind, _), _) + if kind.algorithm() == &SketchAlgorithm::CountSketch => + { + let (row_num, col_num) = cms_dims(kind.params()); + Box::new(CountSketchAccumulatorUpdater::new(row_num, col_num)) + } + + // Heap-bearing top-k variant (raw-input ingest path): route to the + // real `CmsHeapAccumulatorUpdater` so the per-policy top-k heap is + // BUILT (heap-less CMS could not answer `topk(...)` — recall 0). + // Keyed by the configured group-by `aggregated_labels` (e.g. `host`), + // ranked by Σ value per key by default (`weight_mode: value`), or Σ + // count for genuine frequency-top-k (`weight_mode: count`). The OTLP + // modified-sketch path builds the heap agent-side and uses + // `SketchEnvelope` ingest, not this raw arm. + (SummaryFamilyType::Sketch(kind, _), _) + if kind.algorithm() == &SketchAlgorithm::CmsWithHeap => + { + let (row_num, col_num, heap_size) = cms_heap_dims(kind.params()); + Box::new(CmsHeapAccumulatorUpdater::new( + row_num, + col_num, + heap_size, + topk_weight_param(config), + )) + } + + // Heap-bearing CountSketch retains CountSketch estimation semantics. + (SummaryFamilyType::Sketch(kind, _), _) + if kind.algorithm() == &SketchAlgorithm::CountSketchWithHeap => + { + let (row_num, col_num, heap_size) = cms_heap_dims(kind.params()); + Box::new(CountSketchWithHeapAccumulatorUpdater::new( + row_num, + col_num, + heap_size, + topk_weight_param(config), + )) + } + + (SummaryFamilyType::Sketch(kind, _), _) + if kind.algorithm() == &SketchAlgorithm::DDSketch => + { + Box::new(DDSketchAccumulatorUpdater::new(ddsketch_alpha( + kind.params(), + ))) + } + + (SummaryFamilyType::Sketch(kind, _), false) + if kind.algorithm() == &SketchAlgorithm::UnivMon => + { + let SketchParams::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + } = kind.params() + else { + unreachable!("validated UnivMon family parameters") + }; + asap_physical_operators::factory::create_planner_accumulator( + &spec.family, + &planner_types::post_asap::SummaryUpdate::column( + planner_types::pre_asap::ColumnRef::SampleValue, + ), + &Default::default(), + ) + .unwrap() + } + + (SummaryFamilyType::Sketch(kind, _), false) + if kind.algorithm() == &SketchAlgorithm::Hll => + { + let SketchParams::Hll { precision } = kind.params() else { + unreachable!("validated HLL family parameters") + }; + asap_physical_operators::factory::create_planner_accumulator( + &spec.family, + &planner_types::post_asap::SummaryUpdate::column( + planner_types::pre_asap::ColumnRef::SampleValue, + ), + &Default::default(), + ) + .unwrap() + } + + (other_family, keyed) => { + panic!("unsupported isolated kernel fixture {other_family:?}, keyed={keyed}") + } + } +} diff --git a/data_plane/src/tests/mod.rs b/data_plane/src/tests/mod.rs index a6d769261..3a179e8ba 100644 --- a/data_plane/src/tests/mod.rs +++ b/data_plane/src/tests/mod.rs @@ -5,3 +5,5 @@ pub mod trait_design_tests; #[cfg(test)] pub mod test_utilities; + +pub(crate) mod accumulator_fixture; diff --git a/data_plane/src/tests/trait_design_tests.rs b/data_plane/src/tests/trait_design_tests.rs index a10cd3d14..f8f5d6f04 100644 --- a/data_plane/src/tests/trait_design_tests.rs +++ b/data_plane/src/tests/trait_design_tests.rs @@ -1,8 +1,8 @@ -use crate::precompute_engine::operators::{KeyedSumCountAccumulator, SumAccumulator}; #[cfg(test)] use crate::storage_engines::types::{ KeyByLabelValues, MultipleSubpopulationAggregate, SingleSubpopulationAggregate, }; +use asap_physical_operators::summary_kernels::{KeyedSumCountAccumulator, SumAccumulator}; use asap_types::Statistic; #[test] diff --git a/data_plane/src/utils/arithmetic.rs b/data_plane/src/utils/arithmetic.rs deleted file mode 100644 index 94aba683f..000000000 --- a/data_plane/src/utils/arithmetic.rs +++ /dev/null @@ -1,19 +0,0 @@ -//! Float64 arithmetic shared by data-plane execution engines. -//! Preserve IEEE non-finite results; callers own their output policies. - -pub(crate) fn evaluate_float64_arithmetic( - operator: &planner_types::pre_asap::ArithmeticOpKind, - left: f64, - right: f64, -) -> f64 { - use planner_types::pre_asap::ArithmeticOpKind::*; - match operator { - Add => left + right, - Sub => left - right, - Mul => left * right, - Div => left / right, - Mod => left % right, - Pow => left.powf(right), - Atan2 => left.atan2(right), - } -} diff --git a/data_plane/src/utils/mod.rs b/data_plane/src/utils/mod.rs index 72f331c5d..5d620636b 100644 --- a/data_plane/src/utils/mod.rs +++ b/data_plane/src/utils/mod.rs @@ -1,4 +1,3 @@ -pub(crate) mod arithmetic; pub mod file_io; pub mod http; diff --git a/data_plane/tests/component_process_e2e.rs b/data_plane/tests/component_process_e2e.rs index 8f0d89dc6..945a82aec 100644 --- a/data_plane/tests/component_process_e2e.rs +++ b/data_plane/tests/component_process_e2e.rs @@ -42,6 +42,7 @@ fn ddsketch_export(metric: &str, timestamp_ns: u64, counts: Vec) -> Vec alpha, store_counts: counts, store_offset: -1, + ..Default::default() }; let point = DdSketchDataPoint { attributes: vec![KeyValue { diff --git a/data_plane/tests/e2e_controller_plans_and_backend_serves.rs b/data_plane/tests/e2e_controller_plans_and_backend_serves.rs index 8c078ec49..21dc9d6cf 100644 --- a/data_plane/tests/e2e_controller_plans_and_backend_serves.rs +++ b/data_plane/tests/e2e_controller_plans_and_backend_serves.rs @@ -370,6 +370,7 @@ fn build_dd_sketch_state(alpha: f64, store_counts: Vec, store_offset: i32) alpha, store_counts, store_offset, + ..Default::default() } } @@ -1400,6 +1401,7 @@ fn encode_dd_full_envelope(sk: &asap_sketchlib::DdSketch) -> Vec { 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)), diff --git a/data_plane/tests/edge_sketch_codec.rs b/data_plane/tests/edge_sketch_codec.rs index 4f94d9fee..03f4e40f1 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 = data_plane::precompute_engine::operators::DDSketchAccumulator { + let accumulator = asap_physical_operators::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 = data_plane::precompute_engine::operators::DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&snapshot_bytes).unwrap(); + let accumulator = asap_physical_operators::summary_kernels::DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&snapshot_bytes).unwrap(); assert!(accumulator.get_quantile(0.5).is_finite()); } diff --git a/data_plane/tests/promql_differential_process_e2e.rs b/data_plane/tests/promql_differential_process_e2e.rs index 59fc84659..fdda417eb 100644 --- a/data_plane/tests/promql_differential_process_e2e.rs +++ b/data_plane/tests/promql_differential_process_e2e.rs @@ -69,6 +69,7 @@ fn ddsketch_export(metric: &str, timestamp_ns: u64, values: &[f64]) -> Vec { alpha: sketch.wire_alpha(), store_counts: sketch.store_counts, store_offset: sketch.store_offset, + ..Default::default() }; let point = DdSketchDataPoint { attributes: vec![KeyValue { diff --git a/data_plane/tests/support/univmon_erp_process.rs b/data_plane/tests/support/univmon_erp_process.rs index 28bb7b01c..52a43ae79 100644 --- a/data_plane/tests/support/univmon_erp_process.rs +++ b/data_plane/tests/support/univmon_erp_process.rs @@ -1,6 +1,6 @@ use super::*; +use asap_physical_operators::summary_kernels::univmon::UnivMonAccumulator; use control_plane::physical::erp::ErpShapeObserver; -use data_plane::precompute_engine::operators::univmon_accumulator::UnivMonAccumulator; use data_plane::storage_engines::types::{AggregateCore, SerializableToSink}; fn values(offset: usize) -> Vec { diff --git a/docs/design_docs/physical-operators.md b/docs/design_docs/physical-operators.md new file mode 100644 index 000000000..ade079abf --- /dev/null +++ b/docs/design_docs/physical-operators.md @@ -0,0 +1,70 @@ +# Consuming Planner physical operators + +## Ownership and contract + +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. + +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 +phase belongs to the physical node's data state, not the operator payload. +Computation has the same semantics at ingestion time and query time. + +Candidate pruning uses a general semi-join with explicit matching keys, followed +by grouped Sort and grouped Limit. The completeness certificate belongs to the +pruning step; sorting exact scores does not prove completeness. There is no +`MembershipFilter` physical operator or compatibility dispatch. + +## Acceptance + +Planner validates operation, expression, family and parameter support when it +compiles the physical DAG. Backend validates deployment input and storage +contracts before execution. It must not reject a supported computation because +its adapter discarded a type, predicate or join key. An implicit external fallback is not an implementation. +Planner tests cover shared producers, phase assignment, typed batches and +composed candidate pruning. Backend tests cover source binding, installed plan +validation, storage compatibility and query responses. + +The migration does not supply a local raw Scan. That deployment capability +remains deferred. Library tests supplied with raw batches are not evidence that +the backend can execute arbitrary raw-only installed plans. + +## Stack integration + +Planner PR #462 owns the library and depends on Planner #461, including its +composed candidate-pruning API. Backend #774 consumes the pinned library; +#763 integrates ingestion DAG execution and #765 integrates query DAG execution. +The remaining backend stack builds on those integrations. #728 checks candidate +structure, #742 checks selection with synthetic costs, and #775 checks installed +plans against data-plane results. Production evidence and performance validation +remain separate from these correctness tests. + +The library also owns stored-summary decoding, delta reconstruction and +family-specific readout kernels. Deployment adapters select compatible panes +and translate inputs and outputs; they do not copy those computations. + +The shared `physical_planner::compile` API validates concrete operators and typed +input contracts before opening sources. Deployment resolves those inputs and +instantiates the compiled DAG. `CompiledPhysicalDag::from_operators` supports +adapters that already have a concrete physical fragment. Unsupported deployment +input frontiers are reported to Planner as feasibility evidence, before selection. + +Selected Filter, Sort, Limit and semi-join fragments are compiled by Planner and persisted +with typed input contracts. Complete query candidates include Planner readouts +that turn accumulator state into values; internal shared and stored edges retain +their state types. Runtime binds protocol vectors to these contracts; +renamed or multiple join keys retain their original types and positions. External +Prometheus bindings fetch the selected authoritative subquery without rewriting +its labels from the candidate side. The native join performs the comparison. + +Physical execution errors retain their original cause. Memory exhaustion and +cancellation terminate both instant and range requests; routing does not try a +second engine or exact fallback. Physical fragments in one query evaluation share +one run context. This does not yet account for every protocol-buffer allocation +or provide an HTTP-disconnect cancellation mechanism. diff --git a/docs/design_docs/precompute-dag-execution.md b/docs/design_docs/precompute-dag-execution.md new file mode 100644 index 000000000..a56eb003b --- /dev/null +++ b/docs/design_docs/precompute-dag-execution.md @@ -0,0 +1,24 @@ +# Precompute execution from post-ASAP IR + +Audience: backend designers and developers. + +The execution installation is `PrecomputePlan`: selected Planner DAGs, their node bindings, and physical window/storage placement. The former standalone `AggregationConfig` type is removed. `PrecomputeMaterialization` describes storage and routing; it is not independently executable. The streaming configuration serializes the DAG plan and derives its routing index after validation. Flat `aggregations` / `aggregation_configs` documents are rejected. Publish the complete physical plan through `/api/v1/physical-plan` and activate its generation; partial streaming configuration updates are removed. + +```mermaid +flowchart LR + P[Selected Planner post-ASAP DAG] --> I[Validate DAG and physical bindings] + I --> R[Raw source → SummaryAgg streaming kernel] + I --> M[Maintenance dependency scheduler] + R --> S[Stored summary frontier] + S --> M + S --> Q[Query projection and readout] + M --> S +``` + +For a raw producer, installation checks its `SummaryAgg` payload, input edge, source selection, reduction, family, and supported update expressions. The worker executes that validated projection with Planner-owned family and update parameters. Ingestion retains physical window management and routes populations using the validated binding. Shared producers have one installed program and one state per population/window. An unsupported raw path fails installation; the worker cannot choose Sum as a fallback. Backfill uses the same program and update evaluator. Derived summaries continue through the production maintenance scheduler, which observes stored frontiers, dependency roles and shared-node memoization. + +`SummaryAgg` is the operator; Sum, Count, Min, Max, Rate and Increase are its exact families. `ExactAccumulator` retains the family and population layout across updates, reset, merge and serialization. Counter arithmetic can be shared internally, while a Rate state still rejects Increase readout or merge. Keyed layout does not introduce `MultipleX` Planner families. Config-based dispatch remains only in isolated kernel test fixtures and cannot execute in a production build. + +Catalog schema version 6 separates semantic definitions from deployed output identities and carries the Planner family in SDS. Writer and reader bindings identify both the deployed output and its semantic definition; installation and recovery validate their agreement. Installation rejects disagreement between DAG and storage descriptors; storage admission rejects wrong exact families. The persisted `PlannerExactAccumulatorV1` encoding includes family and population layout. Tests cover a real Planner-selected DAG through worker execution and query readout, all six exact families through disk eviction/restart, invalid installations, and the native backend process Remote Write/HTTP query suite. + +The runtime supports explicit subsets of Planner operators. Shared Hydra grouping and unsupported raw input programs are rejected rather than silently assigned another algorithm. Existing imported collector state and isolated payload kernels are not alternate executable configuration formats. diff --git a/docs/developer_docs/control-plane/physical-compiler.md b/docs/developer_docs/control-plane/physical-compiler.md index 7f029b8b5..3885be978 100644 --- a/docs/developer_docs/control-plane/physical-compiler.md +++ b/docs/developer_docs/control-plane/physical-compiler.md @@ -3,9 +3,35 @@ > Interface status: implemented MVP API in > `control_plane::physical::compiler`. +## Dataset-bound planning input + +Planning snapshots use version 3 and require an explicit identity: + +```json +{ + "snapshot_version": 3, + "environment": { + "dataset_identity": {"namespace": "tenant-a", "dataset": "requests"} + } +} +``` + +This is an excerpt; the normal workload, capabilities and policy fields remain +required. All input source names in this deployment resolve within that logical +dataset. The source authority must supply the identity; the compiler does not +infer a tenant from a metric name or endpoint. Separate datasets require separate +bindings rather than multiplexing indistinguishable inputs through this channel. + +Planner exports dataset-bound semantic fragments (version 2). The precompute +input contract carries the same identity and validates it against stored-output +semantics. Catalog schema 6 marks the changed contract; old planning snapshots +and catalog schemas fail validation. Changing an endpoint or replica for the same +dataset does not change the semantic definition. Dataset identity participates +in workload cost manifests so quotes cannot cross dataset scopes. + ## Current implementation boundary -The compiler consumes ASAPPlanner types from `main`, with the resolved revision +The compiler consumes pinned ASAPPlanner types, with the resolved revision exposed as `physical::compiler::PLANNER_REVISION`, and selects from Planner's legal candidate space with backend-owned cost and evidence inputs, and emits one `CompiledPhysicalPlan`. The plan contains one SummaryCatalog plus @@ -15,44 +41,39 @@ for every target collector. Legacy `StageAllocator`/`ThreeStageEmitter` paths remain for older publication flows; they are not a second semantic planner. -ASAPPlanner owns abstract semantics and selection: summary family and -parameters, summary-maintenance lifecycle, and summary-window framework. The -backend enumerates executor-feasible concrete implementations and supplies -complete workload-scoped cost evidence to Planner. It then retains the -concrete identity corresponding to Planner's selected abstract framework. +ASAPPlanner owns semantics and physical candidate construction, including summary +families, maintenance lifecycles and physical boundaries. Backend evaluates +candidate feasibility and complete workload-scoped costs, then selects and binds +the deployment without changing its computation. Missing or stale implementation evidence makes the candidate unavailable; the compiler never invents a framework or assigns it an optimistic zero cost. ## 1. Code architecture -The control plane has three public layers: +The planning and deployment boundary is: ```text -PhysicalCompilationRequest + DataWorkload + concrete implementation evidence - | ^ - | abstract candidates | complete physical costs - v | -ASAPPlanner selection <---------- DeploymentPlanCompiler - | - v -DeploymentPlanCompiler -------> CompiledPhysicalPlan - | | | | - v v v v - Collector Precompute Backend Query - Plan Plan Plan DAG +Query workload + accuracy requirements + ↓ +ASAPPlanner: supported, semantically legal physical candidates + ↓ +Backend: feasibility checks + workload-scoped candidate costs + ↓ +DeploymentPlanCompiler: select and bind a physical candidate + ↓ +CompiledPhysicalPlan + CollectorPlan + PrecomputePlan + TransmissionPlan + QueryPlan + ↓ +Publication and execution ``` -- **Planner selection boundary** is - `planner_selection::select_summary_with_evidence`. It enumerates Planner's - candidates and commits only a legal candidate. -- **Physical compiler** enumerates concrete window/pane/state-layout, - placement, transport, and runtime implementations without changing the - Planner-owned abstract framework. -- **CompiledPhysicalPlan** is the only output passed to publication. Its SummaryCatalog, - CollectorPlan, PrecomputePlan, TransmissionPlan, and QueryPlan are created together and share identities. - -Logical query parsing, summary alternatives, guarantees, and candidate search -remain public ASAPPlanner interfaces. Runtime publication is documented in +Planner owns computation, operator choices, sharing and materialization +boundaries. Backend supplies deployment capabilities and evidence, selects among +admitted physical candidates, and binds the chosen graph to concrete inputs and +stored outputs. It does not reinterpret summary operations into another logical +plan. All published projections share the same catalog and identities. + +Runtime publication is documented in [Runtime plan publication](plan-publication.md). ## 2. Public interfaces and definitions diff --git a/scripts/e2e.sh b/scripts/e2e.sh index 2600999f9..b495ade61 100755 --- a/scripts/e2e.sh +++ b/scripts/e2e.sh @@ -81,6 +81,9 @@ contracts() { say "contracts: shared policy and routing types" rust_test asap_types + # The physical library's independent tests run in the ASAPPlanner workspace. + # Backend type/control-plane/data-plane tests cover its deployment bindings. + CURRENT_STAGE="contracts/asap_otel_proto" say "contracts: modified OTLP and monitor protobuf compatibility" rust_test asap_otel_proto --tests