From 06d6b6d35c151f63e22ddd1ad379b30895e5afe9 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 18:56:02 +0000 Subject: [PATCH 01/15] Introduce independent physical operator library before engine integration --- Cargo.lock | 32 +- Cargo.toml | 10 +- .../examples/calibration_candidates.rs | 10 +- .../examples/offline_planner_replay.rs | 4 +- control_plane/src/emit/mod.rs | 4 +- control_plane/src/physical/compiler.rs | 49 +- .../src/physical/executable_binding.rs | 2 +- control_plane/src/physical/plan_dot.rs | 2 +- .../src/physical/post_asap/cost_model.rs | 4 +- control_plane/src/physical/post_asap/tests.rs | 4 +- control_plane/src/query_plan.rs | 73 +- control_plane/src/query_plan/residual.rs | 56 +- control_plane/tests/offline_evidence.rs | 2 +- crates/asap-physical-operators/Cargo.toml | 27 + crates/asap-physical-operators/README.md | 80 + .../count_min_sketch_accumulator.rs | 1323 +++++++++++ .../count_min_sketch_with_heap_accumulator.rs | 832 +++++++ .../accumulators/count_sketch_accumulator.rs | 678 ++++++ .../count_sketch_with_heap_accumulator.rs | 575 +++++ .../datasketches_kll_accumulator.rs | 727 ++++++ .../src/accumulators/dd_sketch_accumulator.rs | 665 ++++++ .../src/accumulators/exact_accumulator.rs | 326 +++ .../accumulators/hll_sketch_accumulator.rs | 788 +++++++ .../src/accumulators/hydra_kll_accumulator.rs | 165 ++ .../src/accumulators/increase_accumulator.rs | 742 ++++++ .../src/accumulators/keyed_counter_state.rs | 529 +++++ .../src/accumulators/keyed_max_state.rs | 335 +++ .../src/accumulators/keyed_min_state.rs | 335 +++ .../keyed_sum_count_accumulator.rs | 558 +++++ .../src/accumulators/max_accumulator.rs | 248 ++ .../src/accumulators/min_accumulator.rs | 253 ++ .../src/accumulators/mod.rs | 37 + .../sketch_envelope_accumulator.rs | 154 ++ .../src/accumulators/sum_accumulator.rs | 413 ++++ .../src/accumulators/univmon_accumulator.rs | 234 ++ .../asap-physical-operators/src/arithmetic.rs | 19 + .../asap-physical-operators/src/capability.rs | 115 + crates/asap-physical-operators/src/dag/mod.rs | 515 +++++ .../src/dag/operators.rs | 1170 ++++++++++ .../src/dag/planner.rs | 479 ++++ .../asap-physical-operators/src/dag/tests.rs | 260 +++ .../asap-physical-operators/src/dag/values.rs | 318 +++ crates/asap-physical-operators/src/factory.rs | 2046 +++++++++++++++++ .../src/key_by_label_values.rs | 164 ++ crates/asap-physical-operators/src/lib.rs | 22 + .../src/measurement.rs | 94 + crates/asap-physical-operators/src/rows.rs | 88 + crates/asap-physical-operators/src/traits.rs | 351 +++ .../tests/deployment.rs | 96 + .../tests/physical_dag.rs | 663 ++++++ crates/asap_types/src/derived_input.rs | 4 +- crates/asap_types/src/executable_plan.rs | 15 +- crates/asap_types/src/precompute_plan.rs | 12 +- crates/asap_types/src/query_plan.rs | 24 +- data_plane/Cargo.toml | 1 + data_plane/src/drivers/query/servers/http.rs | 6 +- .../precompute_engine/maintenance_runtime.rs | 43 +- .../src/precompute_engine/subdag_scheduler.rs | 21 +- data_plane/src/precompute_engine/worker.rs | 8 +- .../query_engines/asap_query_engine/engine.rs | 2 +- .../asap_query_engine/exact_subqueries.rs | 30 +- .../asap_query_engine/logical_dag.rs | 141 +- .../asap_query_engine/post_asap_readout.rs | 2 +- .../asap_query_engine/summary_exec.rs | 13 +- .../asap_query_engine/summary_executor.rs | 2 +- docs/design_docs/physical-operators.md | 80 + tools/o11y-execution/calibrate_runtime.py | 42 +- .../o11y-execution/test_calibrate_runtime.py | 42 +- 68 files changed, 16848 insertions(+), 316 deletions(-) create mode 100644 crates/asap-physical-operators/Cargo.toml create mode 100644 crates/asap-physical-operators/README.md create mode 100644 crates/asap-physical-operators/src/accumulators/count_min_sketch_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/count_min_sketch_with_heap_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/count_sketch_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/count_sketch_with_heap_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/datasketches_kll_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/dd_sketch_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/exact_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/hll_sketch_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/hydra_kll_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/increase_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/keyed_counter_state.rs create mode 100644 crates/asap-physical-operators/src/accumulators/keyed_max_state.rs create mode 100644 crates/asap-physical-operators/src/accumulators/keyed_min_state.rs create mode 100644 crates/asap-physical-operators/src/accumulators/keyed_sum_count_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/max_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/min_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/mod.rs create mode 100644 crates/asap-physical-operators/src/accumulators/sketch_envelope_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/sum_accumulator.rs create mode 100644 crates/asap-physical-operators/src/accumulators/univmon_accumulator.rs create mode 100644 crates/asap-physical-operators/src/arithmetic.rs create mode 100644 crates/asap-physical-operators/src/capability.rs create mode 100644 crates/asap-physical-operators/src/dag/mod.rs create mode 100644 crates/asap-physical-operators/src/dag/operators.rs create mode 100644 crates/asap-physical-operators/src/dag/planner.rs create mode 100644 crates/asap-physical-operators/src/dag/tests.rs create mode 100644 crates/asap-physical-operators/src/dag/values.rs create mode 100644 crates/asap-physical-operators/src/factory.rs create mode 100644 crates/asap-physical-operators/src/key_by_label_values.rs create mode 100644 crates/asap-physical-operators/src/lib.rs create mode 100644 crates/asap-physical-operators/src/measurement.rs create mode 100644 crates/asap-physical-operators/src/rows.rs create mode 100644 crates/asap-physical-operators/src/traits.rs create mode 100644 crates/asap-physical-operators/tests/deployment.rs create mode 100644 crates/asap-physical-operators/tests/physical_dag.rs create mode 100644 docs/design_docs/physical-operators.md diff --git a/Cargo.lock b/Cargo.lock index d4dfd33fa..bc5762e36 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=cd7e9e0f710816d49190dabd6c789359067a208f#cd7e9e0f710816d49190dabd6c789359067a208f" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" 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=cd7e9e0f710816d49190dabd6c789359067a208f#cd7e9e0f710816d49190dabd6c789359067a208f" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" 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=cd7e9e0f710816d49190dabd6c789359067a208f#cd7e9e0f710816d49190dabd6c789359067a208f" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" dependencies = [ "asap-sql-function-catalog", "asap-types", @@ -393,15 +393,36 @@ dependencies = [ "serde_json", ] +[[package]] +name = "asap-physical-operators" +version = "0.1.0" +dependencies = [ + "asap-types", + "asap_sketch_codec", + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?branch=main)", + "asap_types", + "base64 0.21.7", + "bincode", + "futures", + "hex", + "prost", + "rmp-serde", + "serde", + "serde_json", + "thiserror 1.0.69", + "tracing", + "xxhash-rust", +] + [[package]] name = "asap-sql-function-catalog" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=cd7e9e0f710816d49190dabd6c789359067a208f#cd7e9e0f710816d49190dabd6c789359067a208f" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" [[package]] name = "asap-types" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=cd7e9e0f710816d49190dabd6c789359067a208f#cd7e9e0f710816d49190dabd6c789359067a208f" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" dependencies = [ "serde", "serde_json", @@ -1155,6 +1176,7 @@ dependencies = [ "arrow", "asap-aware-mapping", "asap-frontend-promql", + "asap-physical-operators", "asap-types", "asap_otel_proto", "asap_sketch_codec", diff --git a/Cargo.toml b/Cargo.toml index 5bce59217..852c7a476 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -4,6 +4,7 @@ members = [ "crates/asap_otel_proto", "crates/asap_types", "crates/asap_sketch_codec", + "crates/asap-physical-operators", "data_plane", "control_plane", ] @@ -15,10 +16,10 @@ version = "0.1.0" [workspace.dependencies] # Keep Planner frontends, selection, and IR on the same immutable revision. # Alias upstream asap-types because this workspace also defines asap_types. -planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "cd7e9e0f710816d49190dabd6c789359067a208f" } -asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "cd7e9e0f710816d49190dabd6c789359067a208f" } -asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "cd7e9e0f710816d49190dabd6c789359067a208f" } -asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "cd7e9e0f710816d49190dabd6c789359067a208f" } +planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "09129074c1894f313b98764dd0400ecd73334a2d" } +asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "09129074c1894f313b98764dd0400ecd73334a2d" } +asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "09129074c1894f313b98764dd0400ecd73334a2d" } +asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "09129074c1894f313b98764dd0400ecd73334a2d" } # Shared external deps (used by 2+ crates) serde = { version = "1.0", features = ["derive"] } @@ -38,6 +39,7 @@ arc-swap = "1.7" reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } # Internal crates +asap-physical-operators = { path = "crates/asap-physical-operators" } 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/examples/calibration_candidates.rs b/control_plane/examples/calibration_candidates.rs index f2894c4fb..d83c50ad9 100644 --- a/control_plane/examples/calibration_candidates.rs +++ b/control_plane/examples/calibration_candidates.rs @@ -37,16 +37,14 @@ fn planner_forest(queries: &[control_plane::physical::compiler::QueryCompilation vec![lhs, rhs], json!({"operator_debug":format!("{operator:?}"),"timing_debug":format!("{timing:?}")}), ), - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, - k, - grouping, completeness, } => ( - "CandidateTopK", + "MembershipFilter", vec![candidates, values], - json!({"k":k,"grouping_debug":format!("{grouping:?}"),"completeness_debug":format!("{completeness:?}")}), + json!({"completeness_debug":format!("{completeness:?}")}), ), SummaryExpr::ValueOperation { child, @@ -105,7 +103,7 @@ fn planner_forest(queries: &[control_plane::physical::compiler::QueryCompilation vec![summary_input], json!({"query_debug":format!("{query:?}")}), ), - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { ("SummaryMerge", children.iter().collect(), json!({})) } }; diff --git a/control_plane/examples/offline_planner_replay.rs b/control_plane/examples/offline_planner_replay.rs index 7a703ce18..f87f8ab8b 100644 --- a/control_plane/examples/offline_planner_replay.rs +++ b/control_plane/examples/offline_planner_replay.rs @@ -59,7 +59,7 @@ fn inspect( inspect(summary_input, model, seen, states, raw) } SummaryExpr::ValueOperation { child, .. } => inspect(child, model, seen, states, raw), - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { for child in children { inspect(child, model, seen, states, raw); } @@ -82,7 +82,7 @@ fn inspect( inspect(lhs, model, seen, states, raw); inspect(rhs, model, seen, states, raw); } - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, .. } => { inspect(candidates, model, seen, states, raw); diff --git a/control_plane/src/emit/mod.rs b/control_plane/src/emit/mod.rs index d9d13f75a..f1bd209c0 100644 --- a/control_plane/src/emit/mod.rs +++ b/control_plane/src/emit/mod.rs @@ -54,9 +54,9 @@ fn extract_from_node(node: &Rc) -> Option { // `ExactAgg` case. SummaryExpr::SummaryAgg { .. } => None, SummaryExpr::SummaryEstimate { summary_input, .. } => extract_from_node(summary_input), - SummaryExpr::SummaryMerge { children } => children.iter().find_map(extract_from_node), + SummaryExpr::SummaryMerge { children, .. } => children.iter().find_map(extract_from_node), SummaryExpr::ValueOperation { child, .. } => extract_from_node(child), - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, .. } => extract_from_node(candidates).or_else(|| extract_from_node(values)), // Not surfaced by any `Bind*` path yet (gated on rules that diff --git a/control_plane/src/physical/compiler.rs b/control_plane/src/physical/compiler.rs index 0da9f9033..fdb12922c 100644 --- a/control_plane/src/physical/compiler.rs +++ b/control_plane/src/physical/compiler.rs @@ -807,7 +807,7 @@ fn has_unsafe_raw_entity_leaf( SummaryExpr::SummaryEstimate { summary_input, .. } => { has_unsafe_raw_entity_leaf(summary_input, selected, false) } - SummaryExpr::SummaryMerge { children } => children + SummaryExpr::SummaryMerge { children, .. } => children .iter() .any(|child| has_unsafe_raw_entity_leaf(child, selected, false)), _ => false, @@ -2056,7 +2056,7 @@ fn summary_agg_metric(node: &SummaryNode) -> Option { } } SummaryExpr::SummaryAgg { child, .. } => walk(child, metrics), - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, .. } => { walk(candidates, metrics); @@ -2064,7 +2064,7 @@ fn summary_agg_metric(node: &SummaryNode) -> Option { } SummaryExpr::ValueOperation { child, .. } => walk(child, metrics), SummaryExpr::SummaryEstimate { summary_input, .. } => walk(summary_input, metrics), - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { for child in children { walk(child, metrics); } @@ -2321,7 +2321,7 @@ fn requires_exact_erp_fallback( child: summary_input, .. } => walk(summary_input, out), - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { children.iter().for_each(|child| walk(child, out)) } SummaryExpr::SummaryJoin { outer, inner, .. } => { @@ -2338,7 +2338,7 @@ fn requires_exact_erp_fallback( walk(left, out); walk(right, out); } - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, .. } => { walk(candidates, out); @@ -3206,7 +3206,7 @@ fn immutable_materialization_sources(node: &SummaryNode) -> Option Option @@ -3535,7 +3535,7 @@ fn collect_selected_materializations( } } match &node.expr { - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, .. } => { walk(candidates, readout, composable, grouping.clone(), selected)?; @@ -3589,7 +3589,7 @@ fn collect_selected_materializations( grouping.clone(), selected, )?, - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { for child in children { walk(child, readout, composable, grouping.clone(), selected)?; } @@ -4441,10 +4441,17 @@ pub(crate) mod tests { .compile_promql(request, environment(10_000)) .unwrap(); let entry = plan.query_plan.entries.values().next().unwrap(); - let crate::query_plan::QueryPlanNode::CandidateTopK { inputs, .. } = - &entry.nodes[&entry.root] + let crate::query_plan::QueryPlanNode::Logical { + operator: asap_types::query_plan::residual::ResidualQueryOperator::TopKSelection { .. }, + inputs, + } = &entry.nodes[&entry.root] else { - panic!("Planner weighted TopK must lower to CandidateTopK: {entry:#?}"); + panic!("expected ordinary TopK root") + }; + let crate::query_plan::QueryPlanNode::MembershipFilter { inputs, .. } = + &entry.nodes[&inputs[0]] + else { + panic!("Planner weighted TopK must lower to MembershipFilter: {entry:#?}"); }; assert!(matches!( entry.nodes[&inputs[0]], @@ -4540,7 +4547,14 @@ pub(crate) mod tests { asap_types::AggregationType::CountMinSketchWithHeap ); let entry = plan.query_plan.lookup(query).unwrap(); - let QueryPlanNode::CandidateTopK { inputs, .. } = &entry.nodes[&entry.root] else { + let QueryPlanNode::Logical { + operator: ResidualQueryOperator::TopKSelection { .. }, + inputs, + } = &entry.nodes[&entry.root] + else { + panic!("expected ordinary TopK root") + }; + let QueryPlanNode::MembershipFilter { inputs, .. } = &entry.nodes[&inputs[0]] else { panic!("expected candidate TopK: {entry:#?}"); }; assert!(matches!( @@ -4581,7 +4595,7 @@ pub(crate) mod tests { .nodes .iter() .all(|node| node.output_state.timing - == planner_types::post_asap::ExecutionTiming::MaintenanceTime)); + == planner_types::post_asap::ExecutionTiming::IngestionTime)); assert_eq!(installed.binding.query_plan_sink, entry.root); let mut mismatched = plan.to_publication_artifact().unwrap(); let projected = mismatched @@ -5766,7 +5780,7 @@ pub(crate) mod tests { }; request.queries[0].selected_plan_root = Rc::new(SummaryNode { expr: SummaryExpr::BinaryOp { - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, lhs: selected.clone(), rhs: selected.clone(), operator: planner_types::post_asap::BinaryOperator { @@ -5967,7 +5981,7 @@ pub(crate) mod tests { let right = right.queries[0].selected_plan_root.clone(); let right = Rc::new(SummaryNode { expr: SummaryExpr::ValueOperation { - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, child: right.clone(), }, @@ -5976,7 +5990,7 @@ pub(crate) mod tests { }); request.queries[0].selected_plan_root = Rc::new(SummaryNode { expr: SummaryExpr::BinaryOp { - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, lhs: left.clone(), rhs: right, operator: planner_types::post_asap::BinaryOperator { @@ -7598,6 +7612,7 @@ pub(crate) mod tests { }; let merge = Rc::new(SummaryNode { expr: SummaryExpr::SummaryMerge { + timing: planner_types::post_asap::ExecutionTiming::QueryTime, children: vec![left.clone(), right.clone()], }, schema: left.schema.clone(), diff --git a/control_plane/src/physical/executable_binding.rs b/control_plane/src/physical/executable_binding.rs index 187235955..ef0609e6a 100644 --- a/control_plane/src/physical/executable_binding.rs +++ b/control_plane/src/physical/executable_binding.rs @@ -19,7 +19,7 @@ pub fn install_selected_dag( precompute_sinks.push(node.id); BackendNodeBinding::Materialization { summary_definition } } else if node.output_state.timing - == planner_types::post_asap::ExecutionTiming::MaintenanceTime + == planner_types::post_asap::ExecutionTiming::IngestionTime { BackendNodeBinding::MaintenanceInput } else { diff --git a/control_plane/src/physical/plan_dot.rs b/control_plane/src/physical/plan_dot.rs index 67bad1b79..19d602413 100644 --- a/control_plane/src/physical/plan_dot.rs +++ b/control_plane/src/physical/plan_dot.rs @@ -154,7 +154,7 @@ 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::CandidateTopK { k, .. } => format!("CandidateTopK\nk={k}"), + QueryPlanNode::MembershipFilter { .. } => "MembershipFilter".into(), QueryPlanNode::ExternalExact { .. } => "ExternalExact".into(), QueryPlanNode::ExactFallback { reason } => format!("ExactFallback\n{reason}"), } diff --git a/control_plane/src/physical/post_asap/cost_model.rs b/control_plane/src/physical/post_asap/cost_model.rs index 571c054c3..5fa6c75d9 100644 --- a/control_plane/src/physical/post_asap/cost_model.rs +++ b/control_plane/src/physical/post_asap/cost_model.rs @@ -521,8 +521,8 @@ impl CostModel for ControlPlaneCostModel { fn value_operation_capabilities(&self) -> ValueOperationCapabilities { ValueOperationCapabilities { - read_time: true, - maintenance_time: false, + query_time: true, + ingestion_time: false, } } diff --git a/control_plane/src/physical/post_asap/tests.rs b/control_plane/src/physical/post_asap/tests.rs index 070b6ea62..74bb605a7 100644 --- a/control_plane/src/physical/post_asap/tests.rs +++ b/control_plane/src/physical/post_asap/tests.rs @@ -112,11 +112,11 @@ fn node_is_archive(node: &Rc) -> bool { SummaryExpr::SummaryAgg { child, .. } => node_is_archive(child), SummaryExpr::ValueOperation { child, .. } => node_is_archive(child), SummaryExpr::SummaryEstimate { summary_input, .. } => node_is_archive(summary_input), - SummaryExpr::SummaryMerge { children } => children.iter().any(node_is_archive), + SummaryExpr::SummaryMerge { children, .. } => children.iter().any(node_is_archive), SummaryExpr::SummaryJoin { outer, inner, .. } => { node_is_archive(outer) || node_is_archive(inner) } - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, .. } => node_is_archive(candidates) || node_is_archive(values), SummaryExpr::SummarySubtract { left, right } diff --git a/control_plane/src/query_plan.rs b/control_plane/src/query_plan.rs index 19c12c689..40a7e8b38 100644 --- a/control_plane/src/query_plan.rs +++ b/control_plane/src/query_plan.rs @@ -179,7 +179,7 @@ where *input = remap[input]; } } - QueryPlanNode::CandidateTopK { inputs, .. } + QueryPlanNode::MembershipFilter { inputs, .. } | QueryPlanNode::Binary { inputs, .. } | QueryPlanNode::RelationalJoin { inputs, .. } => { for input in inputs { @@ -308,7 +308,7 @@ where .. }, ), - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, } if measures.len() == 1 => { use planner_types::pre_asap::AggIntent; let operation = match &measures[0] { @@ -368,12 +368,12 @@ where SummaryExpr::ValueOperation { child: sort, operation: planner_types::post_asap::ValueOperation::Limit { n, offset: 0 }, - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + 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::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, } = &sort.expr else { return Err(QueryPlanError::Invalid( @@ -434,7 +434,7 @@ where SummaryExpr::ValueOperation { child, operation: planner_types::post_asap::ValueOperation::Sort { keys, .. }, - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, } if keys.len() == 1 => QueryPlanNode::Logical { operator: residual::ResidualQueryOperator::Sort { descending: !keys[0].ascending, @@ -444,43 +444,14 @@ where SummaryExpr::ValueOperation { .. } => QueryPlanNode::ExactFallback { reason: "unsupported post-ASAP value operation".into(), }, - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, - k, - grouping, completeness, } => { - let labels = grouping - .keys() - .iter() - .map(|&column| { - values - .schema - .fields - .get(column) - .map(|field| field.name.clone()) - .ok_or_else(|| { - QueryPlanError::Invalid( - "unresolved CandidateTopK grouping column".into(), - ) - }) - }) - .collect::, _>>()?; let candidate_input = self.lower(candidates)?; let value_input = if let Some(original) = &self.logical_source { - let parsed = promql_parser::parser::parse(original) - .map_err(|error| QueryPlanError::Invalid(error.to_string()))?; - let promql_parser::parser::Expr::Aggregate(aggregate) = parsed else { - return Err(QueryPlanError::Invalid( - "CandidateTopK requires a top-level PromQL aggregate".into(), - )); - }; - if aggregate.op.to_string() != "topk" { - return Err(QueryPlanError::Invalid( - "CandidateTopK requires a topk source expression".into(), - )); - } + let exact_expression = residual::selected_native_expression(original, values)?; fn item_label(node: &SummaryNode) -> Option { match &node.expr { SummaryExpr::SummaryEstimate { summary_input, .. } => { @@ -500,7 +471,7 @@ where } let item_label = item_label(candidates).ok_or_else(|| { QueryPlanError::Invalid( - "CandidateTopK membership has no named item label".into(), + "MembershipFilter membership has no named item label".into(), ) })?; let value_id = QueryNodeId(self.next_id); @@ -510,7 +481,7 @@ where QueryPlanNode::ExternalExact { request: ExternalExactRequest { language: QueryLanguage::PromQl, - expression: aggregate.expr.to_string(), + expression: exact_expression.to_string(), output: ExternalExactOutput::InstantVector, parameters: BTreeMap::new(), start_parameter: None, @@ -526,15 +497,8 @@ where } else { self.lower(values)? }; - QueryPlanNode::CandidateTopK { + QueryPlanNode::MembershipFilter { inputs: [candidate_input, value_input], - k: u64::try_from(*k).map_err(|_| { - QueryPlanError::Invalid("CandidateTopK k exceeds u64".into()) - })?, - grouping: residual::Grouping { - labels, - without: grouping.is_without(), - }, completeness: completeness.clone(), } } @@ -542,7 +506,7 @@ where lhs, rhs, operator, - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, } if self.logical_source.is_some() || operator.checked_relative_division || operator.checked_finite_division => @@ -624,7 +588,7 @@ where lhs, rhs, operator, - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, } if exact_value_executable(node) => { let planner_types::pre_asap::BinaryOpKind::Arithmetic(operator) = &operator.kind else { @@ -782,7 +746,7 @@ where input: self.lower(summary_input)?, query: query.clone().into(), }, - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { if children.is_empty() { QueryPlanNode::ExactFallback { reason: "empty summary_merge".into(), @@ -881,7 +845,7 @@ pub(crate) fn exact_value_executable(node: &SummaryNode) -> bool { lhs, rhs, operator, - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, } => { matches!( operator.kind, @@ -1393,7 +1357,7 @@ mod tests { } #[test] - fn candidate_topk_rejects_invalid_completeness_contract() { + fn membership_filter_rejects_invalid_completeness_contract() { let leaf = QueryPlanNode::ExactFallback { reason: "prepared".into(), }; @@ -1408,13 +1372,8 @@ mod tests { (QueryNodeId(1), leaf), ( QueryNodeId(2), - QueryPlanNode::CandidateTopK { + QueryPlanNode::MembershipFilter { inputs: [QueryNodeId(0), QueryNodeId(1)], - k: 2, - grouping: residual::Grouping { - labels: vec![], - without: false, - }, completeness: CandidateCompleteness::Certified { guarantee: planner_types::post_asap::ResultGuarantee { metric: planner_types::post_asap::ErrorMetric::Frequency, diff --git a/control_plane/src/query_plan/residual.rs b/control_plane/src/query_plan/residual.rs index 3727e7742..5722d50c2 100644 --- a/control_plane/src/query_plan/residual.rs +++ b/control_plane/src/query_plan/residual.rs @@ -442,11 +442,34 @@ pub(crate) fn selected_residual_nodes( original: &str, selected: &planner_types::post_asap::SummaryNode, ) -> Result<(QueryNodeId, BTreeMap), QueryPlanError> { + let expression = selected_native_expression(original, selected)?; + let mut lower = Lower { + nodes: BTreeMap::new(), + seen: BTreeMap::new(), + }; + let root = lower.lower(&expression)?; + Ok((root, lower.nodes)) +} + +/// Resolve the selected exact subtree to a verified native expression before +/// binding an external input. Never substitute the top-level query's child. +pub(super) fn selected_native_expression( + original: &str, + selected: &planner_types::post_asap::SummaryNode, +) -> Result { if !selected.guarantee.as_ref().is_some_and(|g| g.is_exact()) { return Err(invalid( "native residual substitution requires an exact selected value", )); } + let selected = match &selected.expr { + planner_types::post_asap::SummaryExpr::ValueOperation { + child, + operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, + .. + } => child.as_ref(), + _ => selected, + }; fn visit<'a>(expr: &'a Expr, output: &mut Vec<&'a Expr>) { output.push(expr); match expr { @@ -485,7 +508,7 @@ pub(crate) fn selected_residual_nodes( rhs: right, .. } - | SummaryExpr::CandidateTopK { + | SummaryExpr::MembershipFilter { candidates: left, values: right, .. @@ -500,7 +523,7 @@ pub(crate) fn selected_residual_nodes( selected_horizons(left, out); selected_horizons(right, out); } - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { for child in children { selected_horizons(child, out); } @@ -530,12 +553,7 @@ pub(crate) fn selected_residual_nodes( let candidates = SketchAlgorithmStrategy::new(&asap_aware_mapping::DefaultCostModel) .replacements(&TargetSubDAG::new(&root)); if candidates.iter().any(|candidate| matches!(&candidate.replacement, Replacement::Summary(node) if node.as_ref() == selected)) { - let mut lower = Lower { - nodes: BTreeMap::new(), - seen: BTreeMap::new(), - }; - let root = lower.lower(expression)?; - let candidate = (root, lower.nodes); + let candidate = expression.clone(); if matched .as_ref() .is_some_and(|previous| previous != &candidate) @@ -575,6 +593,24 @@ pub(super) fn selected_aggregate_operator( mod hybrid_tests { use super::*; use crate::query_plan::{MaterializationBinding, PhysicalGrouping}; + #[test] + fn external_binding_rejects_an_unrelated_selected_exact_subtree() { + let exact = crate::query_parser::parse_query_expr_with_interval( + "sum_over_time(other_metric[5m])", + planner_types::types::AccuracyTarget::Exact, + 1_000, + ) + .unwrap(); + let selected = crate::planner_selection::plan_test_query(&exact).unwrap(); + assert!(selected_native_expression("topk(2, sum_over_time(m[5m]))", &selected).is_err()); + assert_eq!( + selected_native_expression("sum_over_time(other_metric[5m])", &selected) + .unwrap() + .to_string(), + "sum_over_time(other_metric[5m])" + ); + } + #[test] fn selected_summary_and_filtered_residual_share_installed_binary() { // Both filtered and unfiltered leaves bind independently. @@ -1086,7 +1122,7 @@ pub fn eligible_materialization_keys( visit(original, left, keys)?; visit(original, right, keys)?; } - SummaryExpr::CandidateTopK { + SummaryExpr::MembershipFilter { candidates, values, .. } => { visit(original, candidates, keys)?; @@ -1098,7 +1134,7 @@ pub fn eligible_materialization_keys( | SummaryExpr::SummaryDelete { summary_input, .. } => { visit(original, summary_input, keys)? } - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { for child in children { visit(original, child, keys)?; } diff --git a/control_plane/tests/offline_evidence.rs b/control_plane/tests/offline_evidence.rs index 5e334d261..a07a9a163 100644 --- a/control_plane/tests/offline_evidence.rs +++ b/control_plane/tests/offline_evidence.rs @@ -351,7 +351,7 @@ fn binary_summary_has_explicit_warm_tier_fallback() { let child = bound(&model()); let root = std::rc::Rc::new(SummaryNode { expr: SummaryExpr::BinaryOp { - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, lhs: child.clone(), rhs: child.clone(), operator: BinaryOperator { diff --git a/crates/asap-physical-operators/Cargo.toml b/crates/asap-physical-operators/Cargo.toml new file mode 100644 index 000000000..601271713 --- /dev/null +++ b/crates/asap-physical-operators/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "asap-physical-operators" +version.workspace = true +edition.workspace = true + +[dependencies] +futures = "0.3" +asap_types.workspace = true +planner-types.workspace = true +asap_sketch_codec = { path = "../asap_sketch_codec" } +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", branch = "main" } +serde.workspace = true +serde_json.workspace = true +tracing.workspace = true +thiserror.workspace = true +base64 = "0.21" +bincode = "1.3" +rmp-serde = "1.3" +prost = "0.13" +xxhash-rust = { version = "0.8", features = ["xxh32", "xxh64"] } + +[features] +default = [] +extra_debugging = [] + +[dev-dependencies] +hex = "0.4" diff --git a/crates/asap-physical-operators/README.md b/crates/asap-physical-operators/README.md new file mode 100644 index 000000000..8bde41a54 --- /dev/null +++ b/crates/asap-physical-operators/README.md @@ -0,0 +1,80 @@ +# ASAP physical operators + +An independent Rust physical operator DAG runtime shared by ingestion time and +query time execution. The library requires neither backend engine, a server, +a storage implementation, Arrow nor DataFusion. DataFusion informed the design; +it is not the execution framework. + +`dag::PhysicalDag` binds typed operator inputs to node IDs. Each execution starts +one producer per reachable node, shares output batches among its consumers, and +bounds buffering. Dropping one consumer does not cancel other consumers. A +`RunContext` carries query or ingestion scope, cancellation and byte accounting. +Executions use the caller's worker and worker-local streams, with no internal +thread pool. Poll multiple root streams concurrently when they share inputs. + +`dag::operators::Operator` implements native batch sources, scalar values, +projection, filtering, grouped exact aggregation, semi-join, grouped Sort and +Limit, vector-to-scalar conversion, Union, and summary construction/merge/readout. +Sort followed by Limit implements grouped ranking; no dedicated TopK physical +operator is needed. Summary construction updates state batch by batch. End of +input means the supplied query range or ingestion window is complete. + +```rust +use asap_physical_operators::dag::{ + operators::{Expression, Operator}, + values::Value, + Limits, PhysicalDag, RunContext, Scope, +}; +use asap_physical_operators::planner::pre_asap::DataType; +use futures::{executor::block_on, StreamExt}; + +let source = Operator::scalar(Value::Int64(7), DataType::Int64)?; +let negate = Operator::project(source.schema(), vec![ + ("value".into(), Expression::Negate(Box::new(Expression::Column(0)))), +])?; +let mut plan = PhysicalDag::default(); +plan.add(0, vec![], source)?; +plan.add(1, vec![0], negate)?; +let run = RunContext::new( + Scope::Query { evaluation_time_ms: 1000, revision: 1 }, + Limits::default(), +)?; +let mut output = plan.execute(&[1], run)?.remove(0); +let batch = block_on(output.next()).unwrap()?; +assert!(matches!(batch.rows()[0][0], Value::Int64(-7))); +# Ok::<(), asap_physical_operators::dag::Error>(()) +``` + +`dag::planner::bind` accepts a post-ASAP DAG and explicit source bindings for +installed ingestion/storage frontiers. It rejects unsupported operations and +schema mismatches before starting a source. Implement `PhysicalOperator` for a +deployment source, including asynchronous I/O; computation operators remain in +the library. The public `planner` export identifies the exact Planner types used +by the crate. The native binder currently supports a subset of those types and +operations; it does not interpret an unknown node as external fallback. + +Plain values preserve Planner scalar/collection types and nullability. Numeric +arithmetic uses matching Int64 or Float64 inputs; integer overflow is an error. +Boolean predicates use three-valued logic. Native summary states currently cover +exact Sum/Count/Min/Max/Rate/Increase, KLL, DDSketch and HLL. Binding checks family, +parameters and readout compatibility; source batches also validate state payloads. +Existing accumulator algorithms are reused as kernels behind these operators. + +Backend ingestion integration is delivered in #763 and query integration in +#765, after this foundation. Installed value/storage adapters provide deployment-specific +computation; they have not all been replaced by native batch bindings. Local raw +Scan remains deferred. See the [design and coverage table](../../docs/design_docs/query-dag-execution.md) +for the distinction between native operator support and backend integration. + +The default limits are eight buffered batches per producer and 64 MiB of estimated +retained execution data. Callers can set both through `Limits`. Accounting includes +consumer-held outputs and reserved operator state, but is not a hard RSS cap or an +allocator hook. Source-owned data and temporary allocation peaks are excluded. +Blocking operators have no spill support. Plan depth is limited to 128. No execution +state is shared between runs, and no implicit fallback or legacy traversal API is +provided. + +Run `cargo test -p asap-physical-operators --locked` for the independent library +acceptance tests, including shared producers, backpressure, cancellation, grouping, +state restoration and raw/partial/fully precomputed DAG examples. These examples +supply in-memory batches; they do not establish backend local raw-Scan support. diff --git a/crates/asap-physical-operators/src/accumulators/count_min_sketch_accumulator.rs b/crates/asap-physical-operators/src/accumulators/count_min_sketch_accumulator.rs new file mode 100644 index 000000000..a1fc7e694 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/count_min_sketch_accumulator.rs @@ -0,0 +1,1323 @@ +use crate::accumulators::dd_sketch_accumulator::normalize_sample_p; +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountMinSketch, CountMinSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +use 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_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; + + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountMinDelta packed-array length mismatch: \ + cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + ) + .into()); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + let delta = CountMinSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l1: pb.l1, + l2: pb.l2, + // The Go-side CountMinDelta proto now carries an hh_keys field + // (heavy-hitter candidates), mirrored on asap_sketchlib's + // CountMinSketchDelta. The vendored Rust proto bindings here don't + // decode it yet, and CountMin has no TopK to rebuild, so pass an + // empty set — same handling as CountSketch's hh_keys. + hh_keys: Vec::new(), + }; + self.inner + .apply_delta(&delta) + .map_err(|e| format!("apply CountMinDelta: {e}"))?; + Ok(()) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 8 { + return Err("Buffer too short for row_num and col_num".into()); + } + + // TODO: this logic will need to be checked for i32 -> f64 + // Github Issue #11 + + let row_num = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; + let col_num = u32::from_le_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]) as usize; + + let expected_size = 8 + (row_num * col_num * 4); + if buffer.len() < expected_size { + return Err("Buffer too short for sketch data".into()); + } + + let mut sketch = Vec::new(); + let mut offset = 8; + + for _ in 0..row_num { + let mut row = Vec::new(); + for _ in 0..col_num { + let value = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + row.push(value); + offset += 8; + } + sketch.push(row); + } + + Ok(Self { + inner: CountMinSketch::from_legacy_matrix(sketch, row_num, col_num), + sample_p: 1.0, + }) + } + + /// Merge multiple accumulators efficiently without cloning all of them. + pub fn merge_multiple( + accumulators: &[Box], + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut cms_accumulators = Vec::with_capacity(accumulators.len()); + for acc in accumulators { + if acc.get_accumulator_type() != AggregationType::CountMinSketch { + return Err(format!( + "Cannot merge CountMinSketchAccumulator with {:?}", + acc.get_accumulator_type() + ) + .into()); + } + let cms_acc = acc + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountMinSketchAccumulator")?; + cms_accumulators.push(cms_acc); + } + + // Check dimensions are consistent + let rows = cms_accumulators[0].inner.rows(); + let cols = cms_accumulators[0].inner.cols(); + for acc in &cms_accumulators { + if acc.inner.rows() != rows || acc.inner.cols() != cols { + return Err( + "Cannot merge CountMinSketch accumulators with different dimensions".into(), + ); + } + } + + let inner_refs: Vec<&CountMinSketch> = + cms_accumulators.iter().map(|acc| &acc.inner).collect(); + let merged_inner = CountMinSketch::merge_refs(&inner_refs)?; + // sample_p is a per-series config constant, so all operands carry the + // same value in practice. Mirror DDSketch's merge policy: prefer a + // sampled factor (< 1.0) over the no-sampling default so a merge with + // a freshly-reset (1.0) base keeps the series' sampling rate. + let sample_p = cms_accumulators + .iter() + .map(|acc| acc.sample_p) + .find(|&p| p < 1.0) + .unwrap_or(cms_accumulators[0].sample_p); + Ok(Self { + inner: merged_inner, + sample_p, + }) + } +} + +/// Defensive upper bound on the number of matrix cells (`rows * cols`) +/// we'll reconstruct from an inbound wire-declared CMS / CountSketch +/// dimension pair. A malformed / hostile payload could declare absurd +/// dims (e.g. `rows = cols = u32::MAX`) and trick the decoder into a +/// huge `Vec` allocation before the `counts_*.len() != rows*cols` +/// check ever runs. Realistic sketches are at most a few hundred rows +/// by tens-of-thousands of columns, so 8M cells (~64 MiB of f64) is a +/// generous ceiling that no legitimate producer reaches. +pub(crate) const MAX_SKETCH_CELLS: usize = 8 * 1024 * 1024; + +/// Validate an inbound, wire-declared `(rows, cols)` pair for a +/// matrix-backed frequency sketch (CMS / CountSketch) BEFORE any matrix +/// is reconstructed from it. Returns `Ok(())` for dimensions a +/// legitimate producer could have emitted, and an `Err` (never a panic) +/// for malformed / degenerate ones so the ingest path can skip the data +/// point and fall through to its existing decode-failure accounting. +/// +/// Rejections: +/// 1. `rows < 1` or `cols < 1` — a zero-dim matrix has no cells. +/// 2. Narrow-hash-budget violation. The cross-language wire hasher +/// (`sketchlib`'s `MatrixHashType::Packed64`) derives every row's +/// column index from disjoint bit-fields of a single 64-bit hash +/// word: row `r` reads `mask_bits = ceil(log2(cols))` bits at offset +/// `r * mask_bits`. Once `rows * mask_bits > 64` the per-row column +/// slices overflow / alias the 64-bit word and the matrix-cell +/// layout is no longer the one the producer hashed into — the sketch +/// is internally degenerate. This mirrors sketchlib's own +/// `MatrixFastHash::assert_compatible` budget (`rows * (mask_bits + 1) <= 64`); we check the column-index bits alone so realistic +/// configs (5x2048, 5x4096, 5x2000) — for which the sign bits share +/// the top of the word without affecting the cell layout — still +/// pass. +/// 3. Obviously-oversized dims: `rows * cols > MAX_SKETCH_CELLS`, +/// guarding against a huge allocation from a malformed payload. +/// +/// `what` names the wire struct for the error message (e.g. +/// `"CountMinState"`). +pub(crate) fn validate_sketch_dims(what: &str, rows: usize, cols: usize) -> Result<(), String> { + if rows < 1 || cols < 1 { + return Err(format!( + "{what} has degenerate dims (rows={rows}, cols={cols}); rejecting" + )); + } + // mask_bits = ceil(log2(cols)); cols >= 1 here. ilog2 is floor(log2). + let mask_bits = if cols.is_power_of_two() { + cols.ilog2() as usize + } else { + cols.ilog2() as usize + 1 + }; + if rows.saturating_mul(mask_bits) > 64 { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) exceed the 64-bit \ + packed-hash column budget (rows * ceil(log2(cols)) = {} > 64); \ + the sketch's matrix-cell layout is degenerate, rejecting", + rows.saturating_mul(mask_bits) + )); + } + if rows.saturating_mul(cols) > MAX_SKETCH_CELLS { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) declare {} cells, \ + exceeding the {MAX_SKETCH_CELLS}-cell ingest cap; rejecting to \ + avoid a huge allocation from a malformed payload", + rows.saturating_mul(cols) + )); + } + Ok(()) +} + +impl SerializableToSink for CountMinSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), + "sketch": self.inner.sketch() + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountMinSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountMinSketchAccumulator" + } + + /// Per-window base rotation: rebuild an empty counter matrix with + /// the same (rows, cols) so the next window's additive cell deltas + /// align to the identical hash geometry. `sample_p` is a per-series + /// config constant (not per-window data), so it is intentionally + /// preserved across the rotation — mirrors `DDSketchAccumulator`. + fn reset_to_empty(&mut self) { + self.inner = CountMinSketch::new(self.inner.rows(), self.inner.cols()); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountMinSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_cms = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountMinSketchAccumulator")?; + + let merged_inner = CountMinSketch::merge_refs(&[&self.inner, &other_cms.inner])?; + // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a + // per-series config constant, so both operands carry the same value + // in practice. Prefer a sampled factor over the no-sampling default + // so a merge with a freshly-reset (1.0) base keeps the series' + // sampling rate. + let sample_p = if self.sample_p < 1.0 { + self.sample_p + } else { + other_cms.sample_p + }; + Ok(Box::new(Self { + inner: merged_inner, + sample_p, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountMinSketch + } + + fn approx_memory_bytes(&self) -> usize { + // Conservative constant for the CountMinSketch counter matrix. + // Real per-instance sizing would require exposing rows/cols on + // the inner sketch; 16 KiB is a reasonable v1 default. + 16 * 1024 + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: asap_types::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::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::accumulators::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_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + use prost::Message; + + let mut acc = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let bytes = PbDelta { + rows: 2, + cols: 3, + cell_rows: vec![0, 1], + cell_cols: vec![0, 2], + d_counts: vec![10, 100], + l1: vec![], + l2: vec![], + ..Default::default() + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!( + acc.inner.sketch(), + vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 106.0]] + ); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = CountMinSketchAccumulator::new(2, 3); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ---------------------------------------------------------------- + // Statistic::Rate / Statistic::Increase — PR #111 honest-gap closure. + // CMS records insert counts but not timestamps. The Rate readout + // requires the engine to push `range_ms` via query_kwargs; without + // it the accumulator falls back to the raw event count (units of + // events/window) so the answer is at least non-empty. + // ---------------------------------------------------------------- + + #[test] + fn test_query_statistic_rate_with_range_ms() { + // Build a CMS whose min-row-sum is 100 events. With a 5-minute + // (300_000 ms) range, the per-second rate is 100 / 300 ≈ 0.333. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![100.0, 0.0], vec![100.0, 0.0]], + 2, + 2, + ), + sample_p: 1.0, + }; + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "300000".to_string()); + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect("Rate with range_ms is supported"); + assert!( + (v - (100.0 / 300.0)).abs() < 1e-9, + "expected 100/300 = {}, got {v}", + 100.0 / 300.0, + ); + } + + #[test] + fn test_query_statistic_rate_without_range_ms_falls_back_to_count() { + // Without `range_ms` in kwargs the accumulator returns the raw + // event volume (events/window units). Caller is responsible for + // surfacing that caveat to the user; this avoids `status=error` + // for instant rate-shape queries that bypass the matrix-selector + // code path. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix(vec![vec![42.0, 0.0], vec![42.0, 0.0]], 2, 2), + sample_p: 1.0, + }; + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Rate, &None, &HashMap::new()) + .expect("Rate without range_ms still answers (fallback)"); + assert_eq!(v, 42.0); + } + + #[test] + fn test_query_statistic_increase_returns_total_count() { + // Increase semantics on CMS: total events in the window — the + // same min-row-sum as Sum / Count. Differs from Rate only in + // that it never divides by range. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix(vec![vec![5.0, 7.0], vec![3.0, 9.0]], 2, 2), + sample_p: 1.0, + }; + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Increase, &None, &HashMap::new()) + .expect("Increase is supported"); + // min-row-sum: row0 = 12, row1 = 12, min = 12. + assert_eq!(v, 12.0); + } + + // ---------------------------------------------------------------- + // Defensive inbound-dimension validation (harden/sketch-dim-validation). + // Malformed / degenerate / narrow-hash-budget-violating CMS dims must + // be rejected gracefully (Err, never a panic); valid configs the + // backend actually uses (5x2048, 5x4096, 5x2000) must still decode. + // ---------------------------------------------------------------- + + /// Build a bare `CountMinState` proto carrying the given dims and a + /// row-major INT64 counts vector sized to `rows*cols` so that, IF the + /// dims pass validation, the reshape also succeeds. Used to prove a + /// malformed-dim payload is rejected at the dim gate, not later. + fn cms_state_bytes(rows: u32, cols: u32) -> Vec { + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + let n = (rows as usize).saturating_mul(cols as usize); + let state = CountMinState { + rows, + cols, + counter_type: CounterType::Int64 as i32, + counts_int: vec![0i64; n], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + state.encode_to_vec() + } + + #[test] + fn test_validate_sketch_dims_accepts_valid_configs() { + // The realistic configs the backend uses must pass unchanged. + for (r, c) in [(5usize, 2048usize), (5, 4096), (5, 2000), (4, 1000), (2, 3)] { + assert!( + validate_sketch_dims("CountMinState", r, c).is_ok(), + "valid config {r}x{c} was wrongly rejected" + ); + } + } + + #[test] + fn test_validate_sketch_dims_rejects_malformed() { + // Zero dims. + assert!(validate_sketch_dims("CountMinState", 0, 2048).is_err()); + assert!(validate_sketch_dims("CountMinState", 5, 0).is_err()); + // Narrow-hash-budget violation: 5 * ceil(log2(8192))=5*13=65 > 64. + let err = validate_sketch_dims("CountMinState", 5, 8192).unwrap_err(); + assert!(err.contains("budget"), "expected budget error, got: {err}"); + // Absurdly oversized: 1 x 16,777,216 = 16M cells > 8M cap. (1 row + // keeps the hash budget tiny — 1*24=24 — so the cap check, not the + // budget check, is what fires here.) + let err = validate_sketch_dims("CountMinState", 1, 16_777_216).unwrap_err(); + assert!(err.contains("cap"), "expected cell-cap error, got: {err}"); + // No panic on extreme dims (saturating_mul guards the products). + assert!(validate_sketch_dims("CountMinState", usize::MAX, usize::MAX).is_err()); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { + // A data point declaring narrow-hash-budget-violating dims must be + // skipped (Err returned, NOT a panic). The ingest caller turns + // this Err into a dropped data point + WARN log. + let bytes = cms_state_bytes(5, 8192); + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "budget-violating dims should be rejected"); + assert!(result.unwrap_err().to_string().contains("rejecting")); + + // A valid neighbour (5x4096) on the same path still decodes fine. + let ok_bytes = cms_state_bytes(5, 4096); + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) + .expect("valid 5x4096 CMS should still decode"); + assert_eq!(acc.inner.rows(), 5); + assert_eq!(acc.inner.cols(), 4096); + } + + #[test] + fn test_query_statistic_rate_rejects_invalid_range_ms() { + let cms = CountMinSketchAccumulator::new(2, 2); + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "0".to_string()); + let trait_obj: &dyn AggregateCore = &cms; + let err = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect_err("range_ms=0 should error"); + assert!(err.to_string().contains("positive")); + + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "not-a-number".to_string()); + let err = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect_err("non-numeric range_ms should error"); + assert!(err.to_string().contains("bad range_ms")); + } + + // ---------------------------------------------------------------- + // sample_p rescale. The edge admits each insert with probability `p`, + // so every stored cell is ~p× the true count. CMS is L1/additive and + // linear, so BOTH the point-frequency (query_key) and the aggregate + // total-event statistics (Count/Sum/Increase/Rate) rescale by 1/p. + // ---------------------------------------------------------------- + + #[test] + fn test_query_key_rescaled_by_sample_p() { + // Same stored cell counts, two sample_p values: the p=0.25 sketch + // must report 4× the point-frequency of the unsampled one. + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let mut unsampled = CountMinSketchAccumulator::new(4, 1000); + unsampled._update(&key, 10.0); + let mut sampled = CountMinSketchAccumulator::new(4, 1000); + sampled._update(&key, 10.0); + sampled.sample_p = 0.25; + + let raw = unsampled.query_key(&key); + let rescaled = sampled.query_key(&key); + assert!( + raw >= 10.0, + "raw estimate should be >= inserted 10, got {raw}" + ); + assert!( + (rescaled - raw * 4.0).abs() < 1e-9, + "expected point-frequency rescaled ≈ 4×raw ({}), got {rescaled}", + raw * 4.0 + ); + } + + #[test] + fn test_aggregate_statistics_rescaled_by_sample_p() { + use 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/crates/asap-physical-operators/src/accumulators/count_min_sketch_with_heap_accumulator.rs b/crates/asap-physical-operators/src/accumulators/count_min_sketch_with_heap_accumulator.rs new file mode 100644 index 000000000..259ab9d21 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/count_min_sketch_with_heap_accumulator.rs @@ -0,0 +1,832 @@ +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CmsHeapItem, CountMinSketchWithHeap, MessagePackCodec}; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; + +use 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::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for CountMinSketchWithHeapAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for CountMinSketchWithHeapAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountMinSketchWithHeapAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_min_sketch_with_heap_creation() { + let cms = CountMinSketchWithHeapAccumulator::new(4, 1000, 20); + assert_eq!(cms.inner.rows(), 4); + assert_eq!(cms.inner.cols(), 1000); + assert_eq!(cms.inner.heap_size, 20); + assert_eq!(cms.inner.topk_heap_items().len(), 0); + } + + #[test] + fn test_count_min_sketch_with_heap_query() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 10, 5); + let key = KeyByLabelValues::new(); + assert_eq!(cms.query_key(&key), 0.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); + } + + #[test] + fn test_count_min_sketch_with_heap_merge() { + // Build controlled state via from_legacy_matrix (works regardless of backend config). + let sketch1 = vec![ + vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap1 = vec![ + CmsHeapItem { + key: "key1".to_string(), + value: 100.0, + }, + CmsHeapItem { + key: "key2".to_string(), + value: 50.0, + }, + ]; + let sketch2 = vec![ + vec![5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap2 = vec![ + CmsHeapItem { + key: "key3".to_string(), + value: 75.0, + }, + CmsHeapItem { + key: "key1".to_string(), + value: 80.0, + }, + ]; + + let cms1 = CountMinSketchWithHeapAccumulator { + inner: CountMinSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), + }; + let cms2 = CountMinSketchWithHeapAccumulator { + inner: CountMinSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), + }; + + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.sketch_matrix()[0][0], 15.0); + assert_eq!(merged.inner.sketch_matrix()[1][1], 35.0); + assert_eq!(merged.inner.heap_size, 3); + assert!(merged.inner.topk_heap_items().len() <= 3); + } + + #[test] + fn test_count_min_sketch_with_heap_merge_single() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms.clone()]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.rows(), cms.inner.rows()); + assert_eq!(merged.inner.cols(), cms.inner.cols()); + assert_eq!(merged.inner.heap_size, cms.inner.heap_size); + } + + #[test] + fn test_count_min_sketch_with_heap_merge_dimension_mismatch() { + let cms1 = CountMinSketchWithHeapAccumulator::new(2, 10, 5); + let cms2 = CountMinSketchWithHeapAccumulator::new(3, 10, 5); + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("dimension")); + } + + #[test] + fn test_count_min_sketch_with_heap_as_aggregate_core() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + assert_eq!(cms.type_name(), "CountMinSketchWithHeapAccumulator"); + } + + #[test] + fn test_get_topk_keys() { + let mut cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + cms.inner.update("label1;label2", 100.0); + cms.inner.update("label3;label4", 50.0); + + let keys = cms.get_topk_keys(); + assert_eq!(keys.len(), 2); + // Top-k order can differ between Legacy and Sketchlib backends (heap ordering / estimates). + let label_sets: std::collections::HashSet<_> = + keys.iter().map(|k| k.labels.clone()).collect(); + assert!(label_sets.contains(&vec!["label1".to_string(), "label2".to_string()])); + assert!(label_sets.contains(&vec!["label3".to_string(), "label4".to_string()])); + } + + #[test] + fn test_multiple_subpopulation_aggregate() { + let cms = CountMinSketchWithHeapAccumulator::new(3, 50, 10); + let key = KeyByLabelValues::new(); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, 0.0); + + let keys = multi_trait.get_keys(); + assert!(keys.is_some()); + assert_eq!(keys.unwrap().len(), 0); + } + + // ---------------------------------------------------------------- + // DELTA-HEAP wire form (encoding MSGPACK_DELTA): apply a sparse matrix + // delta + replace the heap, decoded generically (rmp_serde) WITHOUT any + // asap_sketchlib delta API. The first test feeds a frame produced by the + // Go encoder (sketchlib-go `MarshalCountSketchWithHeapDelta`) to prove + // cross-language byte parity — mirrors how the full-heap parity is + // proven. The second proves PWR full -> delta -> delta reconstruction. + // ---------------------------------------------------------------- + + /// Cross-language byte-parity: this hex is the exact output of + /// sketchlib-go's `asapmsgpack.MarshalCountSketchWithHeapDelta(5, 1024, + /// cells=[(0,1,50),(1,3,-4),(4,1023,1_000_000)], + /// heap=[("/checkout",50),("/cart",20)], heap_size=20)` (captured via a + /// throw-away Go print test, identical methodology to the full-heap + /// golden in `sketchlib-go/.../count_sketch_with_heap_test.go`). If the + /// Go encoder or the rmp_serde layout ever shifts, this decode fails + /// loudly. + const GO_DELTA_HEAP_GOLDEN_HEX: &str = "94c39305cd04009393000132930103fc9304cd03ffce000f42409292a92f636865636b6f7574cb404900000000000092a52f63617274cb403400000000000014"; + + #[test] + fn test_apply_go_produced_delta_heap_frame_matrix_and_heap() { + let bytes = hex::decode(GO_DELTA_HEAP_GOLDEN_HEX).expect("hex"); + + // Base = empty heap accumulator with the frame's dims (what the + // ingest caller holds after the per-window base rotation). + let mut acc = CountMinSketchWithHeapAccumulator::new(5, 1024, 20); + acc.apply_msgpack_heap_delta_bytes(&bytes) + .expect("apply Go delta-heap frame"); + + // Matrix: the three sparse cells landed onto the empty base. + let m = acc.inner.sketch_matrix(); + assert_eq!(m.len(), 5); + assert_eq!(m[0].len(), 1024); + assert_eq!(m[0][1], 50.0, "cell (0,1)"); + assert_eq!(m[1][3], -4.0, "cell (1,3)"); + assert_eq!(m[4][1023], 1_000_000.0, "cell (4,1023)"); + // Everything else stays zero. + assert_eq!(m[2][2], 0.0); + assert_eq!(m[0][0], 0.0); + + // Heap: the frame's full heap, with /checkout ranked above /cart. + let mut items = acc.inner.topk_heap_items(); + items.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap()); + assert_eq!(items.len(), 2); + assert_eq!(items[0].key, "/checkout"); + assert_eq!(items[0].value, 50.0); + assert_eq!(items[1].key, "/cart"); + assert_eq!(items[1].value, 20.0); + } + + #[test] + fn test_pwr_full_then_delta_then_delta_reconstructs_per_window() { + use asap_sketchlib::MessagePackCodec; + + // Window 1 (full frame): build a heap-bearing CountSketch with mass + // and serialize the FULL `{sketch,topk_heap,heap_size}` frame, then + // decode it into a heap accumulator (the cached per-series base). + let w1 = CountMinSketchWithHeap::from_legacy_matrix( + vec![vec![300.0; 4]; 5], + vec![CmsHeapItem { + key: "k".into(), + value: 300.0, + }], + 5, + 4, + 20, + ); + let w1_bytes = w1.to_msgpack().expect("w1 full msgpack"); + let mut base = CountMinSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(&w1_bytes) + .expect("decode w1 full frame as heap accumulator"); + assert_eq!(base.inner.sketch_matrix()[0][0], 300.0); + + // Window 2 delta: this window's own state is matrix cells of value 50 + // against an EMPTY base + heap {k:50}. The DELTA-HEAP frame is encoded + // the same way the Go producer does (4-array, is_delta, sparse cells). + let w2_frame = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + // PWR: rotate base to empty at the window boundary, then apply. + base.reset_to_empty(); + assert_eq!( + base.inner.sketch_matrix()[0][0], + 0.0, + "reset_to_empty cleared matrix" + ); + base.apply_msgpack_heap_delta_bytes(&w2_frame) + .expect("apply w2 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 50.0, "window-2 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 50.0); + // No cross-window leakage from window 1's 300s. + assert_eq!(base.inner.sketch_matrix()[2][2], 0.0); + let h2: Vec<_> = base.inner.topk_heap_items(); + assert_eq!(h2.len(), 1); + assert_eq!(h2[0].key, "k"); + assert_eq!(h2[0].value, 50.0); + + // Window 3 delta: 80s against empty + heap {k:80}. + let w3_frame = encode_delta_heap(5, 4, &[(0, 0, 80)], &[("k", 80.0)], 20); + base.reset_to_empty(); + base.apply_msgpack_heap_delta_bytes(&w3_frame) + .expect("apply w3 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 80.0, "window-3 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 0.0, "no window-2 leakage"); + let h3 = base.inner.topk_heap_items(); + assert_eq!(h3.len(), 1); + assert_eq!(h3[0].value, 80.0); + } + + #[test] + fn test_rmp_serde_layout_is_byte_identical_to_go_encoder() { + // The rmp_serde positional encoding of the delta-heap frame must be + // BYTE-IDENTICAL to sketchlib-go's hand-rolled + // `MarshalCountSketchWithHeapDelta`. This hex is the Go encoder's + // output for (5, 4, cells=[(0,0,50),(1,1,50)], heap=[("k",50)], + // heap_size=20) — the same inputs `encode_delta_heap` uses below. + // Equality here proves both encode AND decode are cross-language + // byte-compatible (the decode path is exercised by the Go-golden + // test above). + const GO_PARITY_HEX: &str = "94c39305049293000032930101329192a16bcb404900000000000014"; + let rust_bytes = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + assert_eq!(hex::encode(&rust_bytes), GO_PARITY_HEX); + } + + #[test] + fn test_apply_delta_rejects_full_frame_and_garbage() { + use asap_sketchlib::MessagePackCodec; + let mut acc = CountMinSketchWithHeapAccumulator::new(2, 4, 5); + // A FULL frame (3-array, no is_delta marker) must NOT decode as a + // delta — the routing relies on the two shapes being distinct. + let full = CountMinSketchWithHeap::from_legacy_matrix( + vec![vec![1.0; 4]; 2], + vec![CmsHeapItem { + key: "a".into(), + value: 1.0, + }], + 2, + 4, + 5, + ) + .to_msgpack() + .unwrap(); + assert!(acc.apply_msgpack_heap_delta_bytes(&full).is_err()); + assert!(acc.apply_msgpack_heap_delta_bytes(b"not msgpack").is_err()); + } + + /// Encode a DELTA-HEAP frame the same way sketchlib-go's + /// `MarshalCountSketchWithHeapDelta` does (rmp_serde positional layout), + /// so the test exercises the real decode path. Tuple structs serialize + /// as msgpack fixed arrays — byte-identical to the Go hand-rolled writer. + fn encode_delta_heap( + rows: u32, + cols: u32, + cells: &[(u32, u32, i64)], + heap: &[(&str, f64)], + heap_size: u64, + ) -> Vec { + #[derive(serde::Serialize)] + struct W<'a>( + bool, + (u32, u32, &'a [(u32, u32, i64)]), + Vec<(String, f64)>, + u64, + ); + let heap_owned: Vec<(String, f64)> = + heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); + let w = W(true, (rows, cols, cells), heap_owned, heap_size); + rmp_serde::to_vec(&w).expect("encode delta-heap") + } + + // ---------------------------------------------------------------- + // FIX 1 — VALUE-WEIGHTED top-k (recall 0 → correct). + // + // `topk(k, sum by (host) (cpu_load))` asks for the top-k hosts by + // SUM OF VALUE. The heavy-hitter heap built by the default `+1`-per- + // occurrence update ranks by COUNT keyed by `item`, so its recall + // against the value-weighted ground truth is 0 when the busiest host + // (most samples) is NOT the heaviest host (largest Σvalue). + // `insert_value(group_label, value)` adds the sample VALUE keyed by the + // GROUP LABEL, so `topk_by_value` ranks by Σvalue — correct recall. + // ---------------------------------------------------------------- + + /// Crafted adversarial dataset: the host with the MOST samples + /// (`h_chatty`, 100 tiny samples) is NOT the host with the largest + /// value-sum (`h_heavy`, a handful of huge samples). A COUNT-ranked + /// heap would surface `h_chatty`; the value-weighted top-k must surface + /// the true heavy hitters by Σvalue, giving recall 1.0 against the + /// ground-truth top-k-by-value-sum. + #[test] + fn value_weighted_topk_has_full_recall_vs_count_topk() { + // (host, per-sample value, sample count) → true Σvalue: + // h_heavy : 1000 × 3 = 3000 (few samples, huge value) + // h_mid : 200 × 5 = 1000 + // h_small : 50 × 6 = 300 + // h_chatty: 1 × 100 = 100 (MOST samples, tiny value) + let data: &[(&str, f64, usize)] = &[ + ("h_heavy", 1000.0, 3), + ("h_mid", 200.0, 5), + ("h_small", 50.0, 6), + ("h_chatty", 1.0, 100), + ]; + + // Wide CMS + heap large enough to hold every group exactly (4 groups) + // so the estimate equals the true Σvalue with no hash collisions. + let mut acc = CountMinSketchWithHeapAccumulator::new(5, 4096, 16); + let mut truth: std::collections::HashMap<&str, f64> = std::collections::HashMap::new(); + for (host, value, count) in data { + for _ in 0..*count { + acc.insert_value(host, *value); + } + *truth.entry(*host).or_insert(0.0) += value * (*count as f64); + } + + // Ground-truth top-2 by value-sum: h_heavy (3000), h_mid (1000). + let mut truth_ranked: Vec<(&str, f64)> = truth.into_iter().collect(); + truth_ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + let truth_top2: std::collections::HashSet<&str> = + truth_ranked.iter().take(2).map(|(k, _)| *k).collect(); + assert!( + truth_top2.contains("h_heavy") && truth_top2.contains("h_mid"), + "ground-truth top-2 by value-sum should be h_heavy + h_mid" + ); + + // Value-weighted top-2 from the heap. + let got = acc.topk_by_value(2); + assert_eq!(got.len(), 2, "k=2 → two groups: {got:?}"); + let got_keys: std::collections::HashSet<&str> = + got.iter().map(|(k, _)| k.as_str()).collect(); + + // RECALL = |got ∩ truth| / |truth| must be 1.0. + let hits = got_keys.intersection(&truth_top2).count(); + let recall = hits as f64 / truth_top2.len() as f64; + assert_eq!( + recall, 1.0, + "value-weighted top-k recall must be 1.0 (count-ranked heap would \ + surface h_chatty and miss h_heavy → recall < 1): got={got:?}" + ); + + // The busiest-by-count host (h_chatty) must NOT be in the top-2, + // proving we rank by value-sum, not occurrence count. + assert!( + !got_keys.contains("h_chatty"), + "h_chatty (most samples, smallest value-sum) must be excluded: {got:?}" + ); + + // Estimates are exact here (no collisions, heap holds all groups): + // top-1 must be h_heavy with Σvalue 3000. + assert_eq!(got[0].0, "h_heavy"); + assert!( + (got[0].1 - 3000.0).abs() < 1e-6, + "h_heavy value-sum estimate ≈ 3000, got {}", + got[0].1 + ); + assert_eq!(got[1].0, "h_mid"); + assert!( + (got[1].1 - 1000.0).abs() < 1e-6, + "h_mid value-sum estimate ≈ 1000, got {}", + got[1].1 + ); + } + + /// A single value-weighted insert must put the full value (not +1) into + /// the heap, and repeated inserts for the same group must accumulate. + #[test] + fn insert_value_accumulates_summed_value_in_heap() { + let mut acc = CountMinSketchWithHeapAccumulator::new(4, 1024, 8); + acc.insert_value("g", 10.0); + acc.insert_value("g", 25.0); + let top = acc.topk_by_value(1); + assert_eq!(top.len(), 1); + assert_eq!(top[0].0, "g"); + assert!( + (top[0].1 - 35.0).abs() < 1e-6, + "summed value should be 35 (10+25), got {}", + top[0].1 + ); + } +} diff --git a/crates/asap-physical-operators/src/accumulators/count_sketch_accumulator.rs b/crates/asap-physical-operators/src/accumulators/count_sketch_accumulator.rs new file mode 100644 index 000000000..c12eda0c3 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/count_sketch_accumulator.rs @@ -0,0 +1,678 @@ +//! CountSketch accumulator backed by `asap_sketchlib::CountSketch`. +//! +//! Supports worker merge, persistence serialization, and modified-OTLP proto +//! decoding. Per-key queries delegate to sketchlib's median-of-signed-rows +//! estimator so query and ingest use the same hash specification. Top-k +//! requires the separate heap-bearing accumulator. + +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountSketch, CountSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +use 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::accumulators::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_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; + + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountSketchDelta packed-array length mismatch: \ + cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + ) + .into()); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + // This is the heap-less matrix kernel; ranked membership is handled + // by the explicit heap-bearing operator, not inferred from delta keys. + let delta = CountSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l2: pb.l2, + hh_keys: Vec::new(), + }; + self.inner + .apply_delta(&delta) + .map_err(|e| format!("apply CountSketchDelta: {e}"))?; + Ok(()) + } +} + +impl SerializableToSink for CountSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "row_num": self.inner.rows, + "col_num": self.inner.cols, + "sketch": self.inner.sketch(), + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountSketchAccumulator" + } + + /// Per-window base rotation: rebuild an empty signed-counter matrix + /// with the same (rows, cols) so the next window's additive cell + /// deltas align to the identical hash geometry. + fn reset_to_empty(&mut self) { + self.inner = CountSketch::new(self.inner.rows, self.inner.cols); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_cs = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountSketchAccumulator")?; + + let merged_inner = CountSketch::merge_refs(&[&self.inner, &other_cs.inner])?; + Ok(Box::new(Self { + inner: merged_inner, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountSketch + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: 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::accumulators::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_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + use prost::Message; + + let mut acc = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix( + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], + 2, + 3, + ), + }; + let bytes = PbDelta { + rows: 2, + cols: 3, + cell_rows: vec![0, 1], + cell_cols: vec![0, 2], + d_counts: vec![10, -6], + l2: vec![], + ..Default::default() + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!( + acc.inner.sketch(), + &vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 0.0]] + ); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = CountSketchAccumulator::new(2, 3); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ---------------------------------------------------------------- + // Defensive inbound-dimension validation (harden/sketch-dim-validation). + // Malformed / narrow-hash-budget-violating CountSketch dims must be + // rejected gracefully (Err, never a panic); valid configs the backend + // actually uses (5x2048, 5x4096, 5x2000) must still decode. + // ---------------------------------------------------------------- + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { + use asap_sketchlib::proto::sketchlib::CounterType; + // 5 * ceil(log2(8192))=5*13=65 > 64 — narrow-hash-budget violation. + // counts sized to rows*cols so rejection is on dims, not length. + let n = 5usize * 8192usize; + let bytes = encode_state( + 5, + 8192, + CounterType::Int64 as i32, + vec![0i64; n], + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "budget-violating dims should be rejected"); + assert!(result.unwrap_err().to_string().contains("rejecting")); + + // A valid neighbour (5x4096) on the same path still decodes fine. + let n_ok = 5usize * 4096usize; + let ok_bytes = encode_state( + 5, + 4096, + CounterType::Int64 as i32, + vec![0i64; n_ok], + Vec::new(), + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) + .expect("valid 5x4096 CountSketch should still decode"); + assert_eq!(acc.inner.rows, 5); + assert_eq!(acc.inner.cols, 4096); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_oversized_dims() { + use asap_sketchlib::proto::sketchlib::CounterType; + // Declare 1 x 16,777,216 = 16M cells (> 8M cap) but send an empty + // counts vector: validation must reject on the dim cap BEFORE the + // decoder tries to allocate/reshape a 16M-entry matrix. (1 row keeps + // the hash budget tiny so the cap check, not the budget check, fires.) + let bytes = encode_state( + 1, + 16_777_216, + CounterType::Int64 as i32, + Vec::new(), + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "oversized dims should be rejected"); + let msg = result.unwrap_err().to_string(); + assert!(msg.contains("cap"), "expected cell-cap error, got: {msg}"); + } +} diff --git a/crates/asap-physical-operators/src/accumulators/count_sketch_with_heap_accumulator.rs b/crates/asap-physical-operators/src/accumulators/count_sketch_with_heap_accumulator.rs new file mode 100644 index 000000000..6b8c1b24d --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/count_sketch_with_heap_accumulator.rs @@ -0,0 +1,575 @@ +//! 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::{ + 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::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::accumulators::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/crates/asap-physical-operators/src/accumulators/datasketches_kll_accumulator.rs b/crates/asap-physical-operators/src/accumulators/datasketches_kll_accumulator.rs new file mode 100644 index 000000000..7a29e84b5 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/datasketches_kll_accumulator.rs @@ -0,0 +1,727 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use asap_sketchlib::{KllSketch, MessagePackCodec}; +use base64::{engine::general_purpose, Engine as _}; +use serde_json::Value; +use std::collections::HashMap; +#[cfg(feature = "extra_debugging")] +use std::time::Instant; +use tracing::debug; + +use 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::SingleSubpopulationAggregate; + self.query(statistic, Some(query_kwargs)) + } +} + +impl SingleSubpopulationAggregate for DatasketchesKLLAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Quantile => { + debug!( + "Querying DatasketchesKLLAccumulator for quantile with kwargs: {:?}", + query_kwargs + ); + let quantile = query_kwargs + .and_then(|kwargs| kwargs.get("quantile")) + .ok_or("Missing quantile parameter for quantile query")? + .parse::() + .map_err(|_| "Invalid quantile parameter format")?; + + if !(0.0..=1.0).contains(&quantile) { + return Err("Quantile must be between 0.0 and 1.0".into()); + } + + Ok(self.get_quantile(quantile)) + } + _ => Err( + format!("Unsupported statistic in DatasketchesKLLAccumulator: {statistic:?}") + .into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for DatasketchesKLLAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use prost::Message; + + fn encode_state(state: asap_sketchlib::proto::sketchlib::KllState) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + } + .encode_to_vec() + } + + #[test] + fn test_datasketches_kll_creation() { + let kll = DatasketchesKLLAccumulator::new(200); + assert!(kll.inner.count() == 0); + assert_eq!(kll.inner.k, 200); + } + + #[test] + fn test_datasketches_kll_update() { + let mut kll = DatasketchesKLLAccumulator::new(200); + kll.update(10.0); + kll.update(20.0); + kll.update(15.0); + assert_eq!(kll.inner.count(), 3); + } + + #[test] + fn test_datasketches_kll_quantile() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + assert_eq!(kll.get_quantile(0.0), 1.0); + assert_eq!(kll.get_quantile(1.0), 10.0); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + let q50 = kll.get_quantile(0.5); + assert!((q50 - 6.0).abs() <= 1.0, "expected median ~6, got {q50}"); + } + + #[test] + fn test_datasketches_kll_query() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + + let mut query_kwargs = HashMap::new(); + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + assert!( + (result - 6.0).abs() <= 1.0, + "expected median ~6, got {result}" + ); + + assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); + } + + #[test] + fn test_datasketches_kll_merge() { + let mut kll1 = DatasketchesKLLAccumulator::new(200); + let mut kll2 = DatasketchesKLLAccumulator::new(200); + + for i in 1..=5 { + kll1.update(i as f64); + } + for i in 6..=10 { + kll2.update(i as f64); + } + + let merged = DatasketchesKLLAccumulator::merge_accumulators(vec![kll1, kll2]).unwrap(); + assert_eq!(merged.inner.count(), 10); + assert_eq!(merged.get_quantile(0.0), 1.0); + assert_eq!(merged.get_quantile(1.0), 10.0); + } + + #[test] + fn test_datasketches_kll_get_keys() { + let kll = DatasketchesKLLAccumulator::new(200); + assert_eq!(kll.type_name(), "DatasketchesKLLAccumulator"); + } + + #[test] + fn test_trait_object() { + let mut kll = DatasketchesKLLAccumulator::new(200); + kll.update(5.0); + let trait_obj: Box = Box::new(kll); + assert_eq!(trait_obj.type_name(), "DatasketchesKLLAccumulator"); + } + + #[test] + fn test_datasketches_kll_query_with_kwargs() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + + let mut query_kwargs = HashMap::new(); + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + assert!( + (result - 6.0).abs() <= 1.0, + "expected median ~6, got {result}" + ); + + query_kwargs.insert("quantile".to_string(), "0.9".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.9 quantile of 1..10 may be 9 or 10. + assert!( + (9.0..=10.0).contains(&result), + "expected 0.9 quantile in [9,10], got {result}" + ); + + query_kwargs.insert("quantile".to_string(), "0.0".to_string()); + assert_eq!( + kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), + 1.0 + ); + + query_kwargs.insert("quantile".to_string(), "1.0".to_string()); + assert_eq!( + kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), + 10.0 + ); + + assert!(kll.query(Statistic::Quantile, None).is_err()); + + query_kwargs.insert("quantile".to_string(), "invalid".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "1.5".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "-0.1".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); + } + + #[test] + fn test_datasketches_kll_merge_multiple() { + let mut kll1 = DatasketchesKLLAccumulator::new(200); + let mut kll2 = DatasketchesKLLAccumulator::new(200); + let mut kll3 = DatasketchesKLLAccumulator::new(200); + + for i in 1..=5 { + kll1.update(i as f64); + } + for i in 6..=10 { + kll2.update(i as f64); + } + for i in 11..=15 { + kll3.update(i as f64); + } + + let boxed_accs: Vec> = + vec![Box::new(kll1), Box::new(kll2), Box::new(kll3)]; + + let merged = DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).unwrap(); + assert_eq!(merged.inner.count(), 15); + assert_eq!(merged.get_quantile(0.0), 1.0); + assert_eq!(merged.get_quantile(1.0), 15.0); + assert_eq!(merged.get_quantile(0.5), 8.0); + } + + #[test] + fn test_datasketches_kll_merge_multiple_error_cases() { + let empty: Vec> = vec![]; + assert!(DatasketchesKLLAccumulator::merge_multiple(&empty).is_err()); + + let kll1 = DatasketchesKLLAccumulator::new(200); + let kll2 = DatasketchesKLLAccumulator::new(100); + let boxed_accs: Vec> = vec![Box::new(kll1), Box::new(kll2)]; + assert!(DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).is_err()); + + use crate::accumulators::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; + + let items: Vec = (0..64).map(|i| i as f64).collect(); + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 64], + items: items.clone(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + + let acc = + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.count(), 64); + // For 64 values 0..63, the true median is 31.5 and quantile + // error is ~1% × range = 0.63. KLL's own point query can + // legally be off by up to ε × N ~= 0.01 × 64 = 0.64. Allow a + // generous tolerance since the important invariant is "the + // decoded sketch is queryable and returns a sensible value". + let median = acc.get_quantile(0.5); + assert!( + (median - 31.5).abs() <= 10.0, + "reconstructed median {median} is outside tolerance of true median 31.5" + ); + let q01 = acc.get_quantile(0.01); + let q99 = acc.get_quantile(0.99); + assert!( + q01 <= q99, + "quantile monotonicity violated: q01={q01}, q99={q99}" + ); + } + + // Compacted portable state is highest-level first, unlike the runtime buffer. + #[test] + fn compacted_wire_state_preserves_count_and_quantiles() { + use asap_sketchlib::{proto::sketchlib::KllState, sketches::KLL}; + let mut source = KLL::::init_kll_with_seed(32, 123); + for i in 0..1000 { + source.update(&(((i * 7919 + 17) % 1009) as f64 / 1009.0)); + } + assert!(source.wire_num_levels() > 1); + let state = KllState { + k: 32, + m: source.wire_m(), + num_levels: source.wire_num_levels(), + levels: source.wire_levels(), + items: source.wire_items(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: vec![], + }; + let decoded = + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&encode_state(state)).unwrap(); + assert_eq!(decoded.inner.count(), source.count() as u64); + for q in [0.0, 0.1, 0.5, 0.9, 1.0] { + assert_eq!(decoded.inner.quantile(q), source.quantile(q), "q={q}"); + } + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's kllprocessor emits: the state + // wrapped in a `SketchEnvelope{kll: ...}` via sketchlib-go's + // `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + + let items: Vec = (0..64).map(|i| i as f64).collect(); + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 64], + items, + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + assert_eq!(acc.inner.count(), 64); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, CountMinState, SketchEnvelope}; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::CountMin( + CountMinState::default(), + )), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_small_k() { + use asap_sketchlib::proto::sketchlib::KllState; + let state = KllState { + k: 4, // < minimum of 8 + m: 2, + num_levels: 0, + levels: Vec::new(), + items: Vec::new(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("k must be >= 8")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_inconsistent_levels() { + use asap_sketchlib::proto::sketchlib::KllState; + // num_levels=1 but levels array has 3 entries instead of 2 + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 5, 10], + items: vec![1.0, 2.0, 3.0, 4.0, 5.0], + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("levels length")); + } + + #[test] + fn aux_stats_exposes_count_via_kll_n() { + let mut acc = DatasketchesKLLAccumulator::new(200); + for i in 0..50 { + acc.update(i as f64); + } + let aux = acc.aux_stats(); + assert_eq!(aux.count, Some(50)); + // KLL doesn't natively expose min/max cheaply and doesn't + // track sum at all — those fields must be None so callers + // fall through to query_statistic. + assert_eq!(aux.sum, None); + assert_eq!(aux.min, None); + assert_eq!(aux.max, None); + } + + #[test] + fn aux_stats_empty_kll_has_zero_count() { + let acc = DatasketchesKLLAccumulator::new(200); + assert_eq!(acc.aux_stats().count, Some(0)); + } +} diff --git a/crates/asap-physical-operators/src/accumulators/dd_sketch_accumulator.rs b/crates/asap-physical-operators/src/accumulators/dd_sketch_accumulator.rs new file mode 100644 index 000000000..9aa1132e6 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/dd_sketch_accumulator.rs @@ -0,0 +1,665 @@ +//! DDSketch accumulator — wraps `asap_sketchlib::DdSketch`. +//! +//! Concrete accumulator reached from the modified-OTLP +//! `Metric.data = DDSketch{…}` hot path (PR C-CountSketch follow-up). +//! Merge via bucket-index alignment on the inner sketch, serialize as +//! MessagePack for the sink, and decode from the sketchlib +//! `DDSketchState` proto. +//! +//! Query semantics follow the STRICT policy after the DataPoint-level +//! METRIC scalars were dropped from the wire format +//! (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57): the sketch serves +//! Quantile (log-bucket estimation) and Count (sum of bucket counts). +//! Sum/Min/Max are no longer derivable from the wire bytes and are +//! served by controller-provisioned exact aggregations — `query_statistic` +//! returns the unavailable-statistic error for them. + +use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::{DdSketch, DdSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +/// DDSketch accumulator — inner log-bucketed sketch. +#[derive(Debug, Clone)] +pub struct DDSketchAccumulator { + pub inner: DdSketch, + /// Edge sampling probability `p ∈ (0,1]` carried on the producer's + /// `SketchEnvelope.sample_p`. The edge admits each value with probability + /// `p` (NitroSketch geometric skip), so `inner.total_count()` is ~`p`× the + /// true count and a `Count` query must rescale by `1/p`. Quantiles are + /// rank-preserving and need NO rescale. `1.0` (and the proto3 default `0.0`, + /// dual-read as `1.0`) means no sampling, so the rescale is a no-op and the + /// behaviour is identical to before. The factor is a per-series config + /// constant: it is set from the first (always-full, otel.rs ingest + /// contract) frame and preserved across delta applies, window-boundary + /// `reset_to_empty`, and `merge_with`. + pub sample_p: f64, +} + +/// Normalize a wire `sample_p` to a usable rescale denominator. `0.0` (proto3 +/// default), `>= 1.0`, and non-finite all collapse to `1.0` (no sampling), so a +/// `Count` rescale by `1/p` is a no-op on unsampled / legacy frames. +pub(crate) fn normalize_sample_p(p: f64) -> f64 { + if p.is_finite() && p > 0.0 && p < 1.0 { + p + } else { + 1.0 + } +} + +impl DDSketchAccumulator { + pub fn new(alpha: f64) -> Self { + Self { + inner: DdSketch::new(alpha), + sample_p: 1.0, + } + } + + /// Read the normalized edge sampling probability from a full-frame + /// `SketchEnvelope`'s `sample_p`. Returns `1.0` (no sampling) for bare + /// `DdSketchState` bytes or any decode failure — the primary production + /// decode path (`reconstruct_via_runtime`) discards the envelope's + /// `sample_p`, so the ingest call site re-reads it from the same bytes. + pub fn sample_p_from_envelope_bytes(buffer: &[u8]) -> f64 { + use asap_sketchlib::proto::sketchlib::SketchEnvelope; + use prost::Message; + SketchEnvelope::decode(buffer) + .map(|env| normalize_sample_p(env.sample_p)) + .unwrap_or(1.0) + } + + /// Decode from the modified OTLP wire format's + /// `DDSketchDataPoint.sketch` bytes when + /// `encoding = DDSKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `DdSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: DdSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize DdSketch msgpack: {e}"))?, + // The msgpack DdSketch struct carries no envelope/sample_p; the + // msgpack path is parity/test-only and is never edge-sampled. + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `DDSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::DDSketchState` message that + /// DataCollector's `ddsketchprocessor` emits when + /// `encoding = DD_SKETCH_ENCODING_PROTO`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + let (state, sample_p) = asap_sketch_codec::ddsketch_state(buffer)?; + if !(state.alpha > 0.0 && state.alpha < 1.0) { + return Err(format!( + "DDSketchState alpha {} out of range (expected 0 < alpha < 1)", + state.alpha + ) + .into()); + } + // 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_sketchlib::proto::sketchlib::DdSketchDelta` message. + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + use asap_sketchlib::proto::sketchlib::DdSketchDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode DDSketchDelta: {e}"))?; + + // The delta no longer carries d_count/d_sum/min/max + // (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). Apply the + // bucket deltas only; `DdSketch` recomputes its total count from + // the merged bucket counts (`total_count()`). + let buckets = pb + .buckets + .into_iter() + .map(|b| (b.index, b.d_count)) + .collect(); + let delta = DdSketchDelta { + buckets, + ..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::accumulators::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_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use prost::Message; + + let mut acc = DDSketchAccumulator::new(0.01); + acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); + + // The wire delta now carries only bucket deltas (tags 2-7 + // reserved); `DdSketchBucketDelta` has just `index` + `d_count`. + let bytes = PbDelta { + buckets: vec![ + DdSketchBucketDelta { + index: 0, + d_count: 10, + }, + DdSketchBucketDelta { + index: 2, + d_count: 20, + }, + ], + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!(acc.inner.store_counts, vec![11, 2, 23]); + // `count` recomputed from the merged buckets: 11 + 2 + 23 = 36. + assert_eq!(acc.inner.total_count(), 36); + } + + /// A valid protobuf with an inadmissible span must not acknowledge a dropped update. + #[test] + fn test_apply_proto_delta_rejects_span_without_mutating_state() { + use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use prost::Message; + let mut acc = DDSketchAccumulator::new(0.01); + acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); + let bytes = PbDelta { + buckets: vec![DdSketchBucketDelta { + index: i32::MAX, + d_count: 1, + }], + } + .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/crates/asap-physical-operators/src/accumulators/exact_accumulator.rs b/crates/asap-physical-operators/src/accumulators/exact_accumulator.rs new file mode 100644 index 000000000..5d23acd87 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/exact_accumulator.rs @@ -0,0 +1,326 @@ +//! Exact summary state identified by Planner family, independent of keyed layout. +use super::increase_accumulator::IncreaseAccumulator; +use crate::{ + 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/crates/asap-physical-operators/src/accumulators/hll_sketch_accumulator.rs b/crates/asap-physical-operators/src/accumulators/hll_sketch_accumulator.rs new file mode 100644 index 000000000..254eee38e --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/hll_sketch_accumulator.rs @@ -0,0 +1,788 @@ +//! HLL accumulator — wraps `asap_sketchlib::HllSketch`. +//! +//! Concrete accumulator reached from the modified-OTLP +//! `Metric.data = HLLSketch{…}` hot path (PR C-CountSketch follow-up). +//! Mirrors the CountSketch accumulator's shape: merge via register-wise +//! max on the inner sketch, serialize as MessagePack for the sink, and +//! decode from the sketchlib `HyperLogLogState` proto. +//! +//! Query semantics (cardinality estimation via the three HLL variants' +//! estimators) are intentionally deferred — the wire format carries the +//! registers + variant + HIP accumulators losslessly, so the merge + +//! store round-trip works end-to-end without that richer query surface. + +use crate::accumulators::dd_sketch_accumulator::normalize_sample_p; +use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::{HllSketch, HllVariant, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +/// Decode one protobuf base-128 varint (LEB128) from the front of `buf`. +/// Returns `(value, bytes_consumed)`, or `None` if the buffer is truncated +/// or the varint overflows u64. +pub(crate) fn read_uvarint(buf: &[u8]) -> Option<(u64, usize)> { + let mut result: u64 = 0; + let mut shift: u32 = 0; + for (i, &b) in buf.iter().enumerate() { + if shift >= 64 { + return None; + } + result |= u64::from(b & 0x7f) << shift; + if b & 0x80 == 0 { + return Some((result, i + 1)); + } + shift += 7; + } + None +} + +/// Expand sketchlib-go's sparse HLL register encoding +/// (`HLLSparseRegisters.packed`) into the dense `num_registers`-byte array. +/// +/// Layout (sketchlib-go `proto/hll/hll.proto`): varint-packed +/// `(index_delta, value)` pairs in ascending index order; `prev_index` +/// starts at 0, so each register's absolute index is the running sum of the +/// deltas. Mirrors the Go encoder in `sketches/HLL/sparse.go` +/// (`encodeSparseRegisters`). The reconstructed array is byte-identical to +/// the dense `registers` field a high-cardinality producer would have sent. +pub(crate) fn expand_sparse_hll_registers( + packed: &[u8], + num_registers: usize, +) -> Result, Box> { + let mut regs = vec![0u8; num_registers]; + let mut prev: u64 = 0; + let mut pos = 0usize; + while pos < packed.len() { + let (delta, n1) = read_uvarint(&packed[pos..]) + .ok_or("HLLSparseRegisters.packed: truncated index_delta varint")?; + pos += n1; + let (value, n2) = read_uvarint(&packed[pos..]) + .ok_or("HLLSparseRegisters.packed: truncated value varint")?; + pos += n2; + let idx = prev + delta; + let i = usize::try_from(idx) + .map_err(|_| format!("HLLSparseRegisters: index {idx} overflows usize"))?; + if i >= num_registers { + return Err(format!( + "HLLSparseRegisters: register index {i} >= num_registers {num_registers}" + ) + .into()); + } + regs[i] = u8::try_from(value) + .map_err(|_| format!("HLLSparseRegisters: register value {value} > 255"))?; + prev = idx; + } + Ok(regs) +} + +/// HLL accumulator — inner register array + variant metadata. +#[derive(Debug, Clone)] +pub struct HllSketchAccumulator { + pub inner: HllSketch, + /// Edge sampling probability `p ∈ (0,1]` carried on the producer's + /// `SketchEnvelope.sample_p`. HLL uses HASH-THRESHOLD sampling — each + /// DISTINCT key is admitted into the sketch with probability `p`, so the + /// register-derived distinct-count estimate is ~`p`× the true + /// cardinality and a `Cardinality`/`Count` query must rescale by `1/p`. + /// `1.0` (and the proto3 default `0.0`, dual-read as `1.0`) means no + /// sampling, so the rescale is a no-op and the behaviour is identical to + /// before. Mirrors `DDSketchAccumulator::sample_p`; set from the envelope + /// at the `from_sketchlib_proto_bytes` decode site and preserved across + /// `reset_to_empty` and `merge_with`. + /// + /// NOTE: HLL edge sampling is currently force-disabled in the edge + /// (`warm_sketch.go` HLL case always emits `sample_p = 1.0`), so in + /// practice `p = 1.0` today and this is a latent-correctness fix that + /// activates if HLL sampling is ever enabled. + pub sample_p: f64, +} + +impl HllSketchAccumulator { + pub fn new(variant: HllVariant, precision: u32) -> Self { + Self { + inner: HllSketch::new(variant, precision), + sample_p: 1.0, + } + } + + /// Decode from the modified OTLP wire format's + /// `HLLSketchDataPoint.sketch` bytes when + /// `encoding = HLL_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `HllSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: HllSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?, + // The msgpack HllSketch struct carries no envelope/sample_p; the + // msgpack path is parity/test-only and is never edge-sampled. + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `HLLSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::HyperLogLogState` message + /// that DataCollector's `hllprocessor` emits when + /// `encoding = HLL_SKETCH_ENCODING_PROTO`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + + // DataCollector's hllprocessor wraps the state in a + // `SketchEnvelope{hll: HyperLogLogState}` via sketchlib-go's + // `SerializePortableFO` + `proto.Marshal`. Try envelope first, + // fall back to bare `HyperLogLogState` for callers (e.g. unit + // tests) that encode the state directly. Mirrors the PR #14 + // fix on `CountMinSketchAccumulator::from_sketchlib_proto_bytes`. + // Capture the envelope's `sample_p` alongside the state so a + // Cardinality query can rescale the distinct-count estimate by + // `1/p`. Bare `HyperLogLogState` bytes (no envelope) carry no + // sampling info → `sample_p` 1.0 (no rescale). Mirrors + // `DDSketchAccumulator`. + let (state, sample_p) = match SketchEnvelope::decode(buffer) { + Ok(env) => { + let sp = env.sample_p; + match env.sketch_state { + Some(sketch_envelope::SketchState::Hll(st)) => (st, sp), + Some(other) => { + return Err(format!( + "SketchEnvelope contains non-HLL sketch: {:?}", + std::mem::discriminant(&other) + ) + .into()); + } + None => ( + HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?, + 1.0, + ), + } + } + Err(_) => ( + HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?, + 1.0, + ), + }; + if state.precision == 0 || state.precision > 20 { + return Err(format!( + "HyperLogLogState precision {} out of range (expected 1..=20)", + state.precision + ) + .into()); + } + let expected_len = 1usize << state.precision; + // Register resolution. sketchlib-go emits the SPARSE + // `registers_sparse` (proto tag 7) form below its dense/sparse + // crossover (~6000 non-zero registers — see + // sketchlib-go/sketches/HLL/sparse.go); low-cardinality producers + // (the common case) therefore leave the dense `registers` (tag 3) + // field empty. The proto contract (hll.proto) is: read whichever of + // `registers` / `registers_sparse` is present; if both are empty the + // sketch is all-zero. Reconstruct the dense 2^precision array in all + // three cases so the inner `HllSketch` always gets a full register + // vector. + let dense_registers: Vec = if state.registers.len() == expected_len { + state.registers.clone() + } else if !state.registers.is_empty() { + // A non-empty dense field of the wrong length is a malformed frame. + return Err(format!( + "HyperLogLogState registers has {} bytes, expected 2^precision = {}", + state.registers.len(), + expected_len + ) + .into()); + } else if let Some(sparse) = state.registers_sparse.as_ref() { + expand_sparse_hll_registers(&sparse.packed, expected_len)? + } else { + // Neither representation populated → all-zero register array. + vec![0u8; expected_len] + }; + let proto_variant = ProtoVariant::try_from(state.variant) + .map_err(|_| format!("HyperLogLogState has unknown variant tag {}", state.variant))?; + let variant = match proto_variant { + ProtoVariant::Unspecified => HllVariant::Unspecified, + ProtoVariant::Regular => HllVariant::Regular, + ProtoVariant::ErtlMle => HllVariant::Datafusion, + ProtoVariant::Hip => HllVariant::Hip, + }; + let inner = HllSketch::from_raw( + variant, + state.precision, + dense_registers, + state.hip_kxq0, + state.hip_kxq1, + state.hip_est, + ); + Ok(Self { + inner, + sample_p: normalize_sample_p(sample_p), + }) + } + + /// Apply a proto-encoded `HLLDelta` frame to this accumulator's + /// inner sketch — the decode path for + /// `HLL_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + /// + /// Called against an accumulator that already carries the base + /// sketch state; the caller is the per-series snapshot cache in + /// the ingest path. Bytes are the + /// `asap_sketchlib::proto::sketchlib::HllDelta` message. + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + // The HLLDelta wire format is a varint-packed (index_delta, value) blob; + // decode + apply (register-wise max) via the shared sketch library so + // the unpacking stays a single source of truth. + self.inner + .apply_delta_bytes(buffer) + .map_err(|e| format!("apply HLLDelta: {e}"))?; + Ok(()) + } +} + +impl SerializableToSink for HllSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "variant": format!("{:?}", self.inner.variant), + "precision": self.inner.precision, + "register_bytes": self.inner.registers.len(), + "hip_kxq0": self.inner.hip_kxq0, + "hip_kxq1": self.inner.hip_kxq1, + "hip_est": self.inner.hip_est, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for HllSketchAccumulator { + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::().saturating_add(self.inner.registers.capacity()) + } + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "HllSketchAccumulator" + } + + /// Per-window base rotation: zero the registers but keep the variant + /// and precision. Critical for HLL — its register-wise `max` merge + /// has no inverse, so a never-reset base accumulates the all-time-max + /// across windows (`docs/delta-baseline-contract.md` §1.5); rotating + /// to an empty register array makes per-window cardinality correct. + /// `sample_p` is a per-series config constant (not per-window data), so + /// it is intentionally preserved across the rotation — mirrors + /// `DDSketchAccumulator`. + fn reset_to_empty(&mut self) { + self.inner = HllSketch::new(self.inner.variant, self.inner.precision); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge HllSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_hll = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to HllSketchAccumulator")?; + let merged_inner = HllSketch::merge_refs(&[&self.inner, &other_hll.inner])?; + // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a + // per-series config constant, so both operands carry the same value + // in practice. Prefer a sampled factor over the no-sampling default + // so a merge with a freshly-reset (1.0) base keeps the series' + // sampling rate. + let sample_p = if self.sample_p < 1.0 { + self.sample_p + } else { + other_hll.sample_p + }; + Ok(Box::new(Self { + inner: merged_inner, + sample_p, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::HLL + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: 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::accumulators::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_sketchlib::proto::sketchlib::HllDelta as PbDelta; + use prost::Message; + + let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); + acc.inner.registers = vec![1, 5, 3, 7]; + + // Packed (index_delta, value) blob for updates {0:4, 2:6}: + // varint(0),varint(4),varint(2),varint(6). + let delta_bytes = PbDelta { + packed_updates: vec![0, 4, 2, 6], + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&delta_bytes).expect("apply ok"); + // Max semantics: reg[0]=max(1,4)=4, reg[2]=max(3,6)=6; others unchanged. + assert_eq!(acc.inner.registers, vec![4, 5, 6, 7]); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ----- sample_p cardinality rescale ----- + // + // HLL uses hash-threshold sampling: each distinct key is admitted into + // the sketch with probability `p`, so the register-derived cardinality + // estimate is ~p× the true distinct count and must be rescaled by 1/p. + + #[test] + fn test_cardinality_is_rescaled_by_sample_p() { + use 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/crates/asap-physical-operators/src/accumulators/hydra_kll_accumulator.rs b/crates/asap-physical-operators/src/accumulators/hydra_kll_accumulator.rs new file mode 100644 index 000000000..a5cc64dc1 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/hydra_kll_accumulator.rs @@ -0,0 +1,165 @@ +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{HydraKllSketch, MessagePackCodec}; +use base64::{engine::general_purpose, Engine as _}; +use std::collections::HashMap; + +use 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::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for HydraKllSketchAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for HydraKllSketchAccumulator { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Quantile => { + let quantile = query_kwargs + .and_then(|kwargs| kwargs.get("quantile")) + .ok_or("Missing quantile parameter for quantile query")? + .parse::() + .map_err(|_| "Invalid quantile parameter format")?; + + if !(0.0..=1.0).contains(&quantile) { + return Err("Quantile must be between 0.0 and 1.0".into()); + } + + Ok(self.query_key(key, quantile)) + } + _ => Err( + format!("Unsupported statistic in HydraKllSketchAccumulator: {statistic:?}").into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} diff --git a/crates/asap-physical-operators/src/accumulators/increase_accumulator.rs b/crates/asap-physical-operators/src/accumulators/increase_accumulator.rs new file mode 100644 index 000000000..7407800b0 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/increase_accumulator.rs @@ -0,0 +1,742 @@ +use crate::{ + 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::SingleSubpopulationAggregate; + self.query( + statistic, + (!query_kwargs.is_empty()).then_some(query_kwargs), + ) + } +} + +impl SingleSubpopulationAggregate for IncreaseAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Increase => Ok(self.extrapolated_value(query_kwargs, false)?), + Statistic::Rate => Ok(self.extrapolated_value(query_kwargs, true)?), + // For instant `sum [by (...)] (counter_metric)` Prometheus + // sums the latest cumulative value of each matching series. + // The IncreaseAccumulator already tracks that latest value + // in `last_seen_measurement`, so per-series Sum is just + // that scalar; the engine's outer aggregation groups by the + // `by` labels and adds the per-series totals across keys. + // + // See PR #108 audit conclusion (commit 4359e10) and issue + // ProjectASAP/ASAPCollector#46: pre-fix the ASAP tier ingested + // counters as IncreaseAccumulator and bare `sum by (...) ()` + // capability-missed because this trait did not answer Sum. + Statistic::Sum => Ok(self.last_seen_measurement.value), + _ => Err(format!("Unsupported statistic in IncreaseAccumulator: {statistic:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl IncreaseAccumulator { + fn extrapolated_value( + &self, + query_kwargs: Option<&HashMap>, + is_rate: bool, + ) -> Result> { + if self.sample_count < 2 || self.last_seen_timestamp <= self.starting_timestamp { + return Err("at least two ordered counter samples are required".into()); + } + let sampled_interval = (self.last_seen_timestamp - self.starting_timestamp) as f64 / 1000.0; + let Some(kwargs) = query_kwargs else { + return Ok(if is_rate { + self.total_increase / sampled_interval + } else { + self.total_increase + }); + }; + let range_start = kwargs + .get("range_start_ms") + .ok_or("missing range_start_ms")? + .parse::()?; + let range_end = kwargs + .get("range_end_ms") + .ok_or("missing range_end_ms")? + .parse::()?; + if range_end <= range_start { + return Err("invalid counter evaluation range".into()); + } + + let mut duration_to_start = + (self.starting_timestamp.saturating_sub(range_start)) as f64 / 1000.0; + let duration_to_end = (range_end.saturating_sub(self.last_seen_timestamp)) as f64 / 1000.0; + let average_sample_interval = sampled_interval / (self.sample_count - 1) as f64; + let extrapolation_threshold = average_sample_interval * 1.1; + + if self.total_increase > 0.0 && self.starting_measurement.value >= 0.0 { + let duration_to_zero = + sampled_interval * (self.starting_measurement.value / self.total_increase); + duration_to_start = duration_to_start.min(duration_to_zero); + } + let mut extrapolate_to = sampled_interval; + extrapolate_to += if duration_to_start < extrapolation_threshold { + duration_to_start.max(0.0) + } else { + average_sample_interval / 2.0 + }; + extrapolate_to += if duration_to_end < extrapolation_threshold { + duration_to_end.max(0.0) + } else { + average_sample_interval / 2.0 + }; + let mut factor = extrapolate_to / sampled_interval; + if is_rate { + factor /= (range_end - range_start) as f64 / 1000.0; + } + Ok(self.total_increase * factor) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_increase_accumulator_creation() { + let starting_measurement = Measurement::new(10.0); + let last_seen_measurement = Measurement::new(25.0); + let acc = IncreaseAccumulator::new( + starting_measurement.clone(), + 1000, + last_seen_measurement.clone(), + 2000, + ); + + assert_eq!(acc.starting_measurement.value, 10.0); + assert_eq!(acc.starting_timestamp, 1000); + assert_eq!(acc.last_seen_measurement.value, 25.0); + assert_eq!(acc.last_seen_timestamp, 2000); + } + + #[test] + fn test_increase_accumulator_update() { + let starting_measurement = Measurement::new(10.0); + let mut acc = IncreaseAccumulator::new( + starting_measurement.clone(), + 1000, + starting_measurement.clone(), + 1000, + ); + + let new_measurement = Measurement::new(25.0); + acc.update(new_measurement.clone(), 2000); + + assert_eq!(acc.last_seen_measurement.value, 25.0); + assert_eq!(acc.last_seen_timestamp, 2000); + assert_eq!(acc.starting_measurement.value, 10.0); // Should remain unchanged + } + + #[test] + fn test_increase_accumulator_query() { + let starting_measurement = Measurement::new(10.0); + let last_seen_measurement = Measurement::new(25.0); + let acc = IncreaseAccumulator::new( + starting_measurement, + 1000, + last_seen_measurement, + 3000, // 2 second difference + ); + + // Test increase calculation + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, None).unwrap(), + 15.0 + ); + + // Test rate calculation (per second) + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, None).unwrap(), + 7.5 + ); // 15.0 / 2.0 + + // Statistic::Sum returns the latest cumulative counter value, + // matching Prometheus semantics for instant `sum()`. + // (Issue ProjectASAP/ASAPCollector#46, PR #108 diagnosis.) + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), + 25.0 + ); + + // Unsupported statistics still error. + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); + } + + #[test] + fn prometheus_counter_reset_and_boundary_extrapolation() { + let mut acc = IncreaseAccumulator::new( + Measurement::new(10.0), + 10_000, + Measurement::new(10.0), + 10_000, + ); + acc.update(Measurement::new(20.0), 20_000); + acc.update(Measurement::new(3.0), 30_000); + acc.update(Measurement::new(13.0), 50_000); + assert_eq!(acc.total_increase, 23.0); + assert_eq!(acc.sample_count, 4); + + let kwargs = HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + let increase = + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, Some(&kwargs)) + .unwrap(); + let rate = crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, Some(&kwargs)) + .unwrap(); + assert!((increase - 34.5).abs() < 1e-12); + assert!((rate - 0.575).abs() < 1e-12); + } + + #[test] + fn pane_merge_preserves_resets_and_prometheus_extrapolation() { + let mut left = IncreaseAccumulator::new( + Measurement::new(10.0), + 10_000, + Measurement::new(10.0), + 10_000, + ); + left.update(Measurement::new(20.0), 20_000); + let mut right = + IncreaseAccumulator::new(Measurement::new(3.0), 30_000, Measurement::new(3.0), 30_000); + right.update(Measurement::new(13.0), 50_000); + let merged = IncreaseAccumulator::merge_accumulators(vec![right, left]).unwrap(); + assert_eq!(merged.total_increase, 23.0); + assert_eq!(merged.sample_count, 4); + let kwargs = HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&merged, Statistic::Increase, Some(&kwargs)) + .unwrap(), + 34.5 + ); + } + + #[test] + fn counter_sds_state_is_constant_size_per_pane() { + let mut acc = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(0.0), 0); + let initial = acc.serialize_to_bytes().len(); + for second in 1..=86_400 { + acc.update(Measurement::new(second as f64), second * 1_000); + } + assert_eq!(acc.serialize_to_bytes().len(), initial); + assert_eq!(acc.sample_count, 86_401); + assert_eq!( + acc.approx_memory_bytes(), + std::mem::size_of::() + ); + } + + #[test] + fn test_increase_accumulator_sum_is_latest_cumulative_value() { + // Instant `sum ()` semantics: the per-series summand is + // the latest cumulative counter value. Two series with latest + // values 100 and 50 (started at 10 and 5 respectively) should + // each report Sum = 100 and Sum = 50 — the engine's `sum by` + // outer aggregation does the cross-series total. + let acc_a = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000); + let acc_b = + IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc_a, Statistic::Sum, None).unwrap(), + 100.0 + ); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc_b, Statistic::Sum, None).unwrap(), + 50.0 + ); + } + + #[test] + fn test_increase_accumulator_merge() { + let acc1 = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(20.0), 2000); + let acc2 = IncreaseAccumulator::new( + Measurement::new(5.0), + 500, // Earlier start + Measurement::new(15.0), + 1500, + ); + let acc3 = IncreaseAccumulator::new( + Measurement::new(20.0), + 2000, + Measurement::new(30.0), + 3000, // Later end + ); + + let merged = + >::merge_accumulators( + vec![acc1, acc2, acc3], + ) + .unwrap(); + + // Should use earliest start and latest end + assert_eq!(merged.starting_measurement.value, 5.0); + assert_eq!(merged.starting_timestamp, 500); + assert_eq!(merged.last_seen_measurement.value, 30.0); + assert_eq!(merged.last_seen_timestamp, 3000); + } + + #[test] + fn test_increase_accumulator_serialization() { + let acc = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(25.0), 2000); + + // Test JSON serialization + let json = acc.serialize_to_json(); + let deserialized = IncreaseAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!( + acc.starting_measurement.value, + deserialized.starting_measurement.value + ); + assert_eq!(acc.starting_timestamp, deserialized.starting_timestamp); + assert_eq!( + acc.last_seen_measurement.value, + deserialized.last_seen_measurement.value + ); + assert_eq!(acc.last_seen_timestamp, deserialized.last_seen_timestamp); + + // Test byte serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = IncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + acc.starting_measurement.value, + deserialized_bytes.starting_measurement.value + ); + assert_eq!( + acc.starting_timestamp, + deserialized_bytes.starting_timestamp + ); + assert_eq!( + acc.last_seen_measurement.value, + deserialized_bytes.last_seen_measurement.value + ); + assert_eq!( + acc.last_seen_timestamp, + deserialized_bytes.last_seen_timestamp + ); + assert_eq!(acc.total_increase, deserialized_bytes.total_increase); + assert_eq!(acc.sample_count, deserialized_bytes.sample_count); + + let legacy = &bytes[..bytes.len() - RESET_AWARE_WIRE_EXTENSION_LEN]; + let legacy_value = IncreaseAccumulator::deserialize_from_bytes(legacy).unwrap(); + assert_eq!(legacy_value.total_increase, 15.0); + assert_eq!(legacy_value.sample_count, 2); + } + + #[test] + fn test_trait_object() { + let acc: Box = Box::new(IncreaseAccumulator::new( + Measurement::new(10.0), + 1000, + Measurement::new(25.0), + 2000, + )); + + assert_eq!(acc.type_name(), "IncreaseAccumulator"); + } +} diff --git a/crates/asap-physical-operators/src/accumulators/keyed_counter_state.rs b/crates/asap-physical-operators/src/accumulators/keyed_counter_state.rs new file mode 100644 index 000000000..1d4cf1c7f --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/keyed_counter_state.rs @@ -0,0 +1,529 @@ +use crate::accumulators::IncreaseAccumulator; +use crate::{ + 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::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedCounterState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedCounterState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + query_kwargs: Option<&HashMap>, + ) -> Result> { + let data = self + .increases + .get(key) + .ok_or_else(|| format!("Key {key} not found in KeyedCounterState"))?; + + data.query(statistic, query_kwargs) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedCounterState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedCounterState::new(); + + for accumulator in accumulators { + for (key, data) in accumulator.increases { + if let Some(existing_data) = result.increases.get_mut(&key) { + *existing_data = + IncreaseAccumulator::merge_accumulators(vec![existing_data.clone(), data])?; + } else { + result.increases.insert(key, data); + } + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Measurement; + + fn create_test_increase_accumulator(start_val: f64, end_val: f64) -> IncreaseAccumulator { + IncreaseAccumulator::new( + Measurement::new(start_val), + 1000, + Measurement::new(end_val), + 2000, + ) + } + + fn create_test_increase_accumulator_with_time( + start_val: f64, + start_time: i64, + end_val: f64, + end_time: i64, + ) -> IncreaseAccumulator { + IncreaseAccumulator::new( + Measurement::new(start_val), + start_time, + Measurement::new(end_val), + end_time, + ) + } + + #[test] + fn test_keyed_counter_state_creation() { + let acc = KeyedCounterState::new(); + assert!(acc.increases.is_empty()); + } + + #[test] + fn test_keyed_counter_state_update() { + let mut acc = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + let increase1 = create_test_increase_accumulator(10.0, 25.0); + let increase2 = create_test_increase_accumulator(5.0, 15.0); + + acc.update(key1.clone(), increase1); + acc.update(key2.clone(), increase2); + + assert_eq!(acc.increases.len(), 2); + assert!(acc.increases.contains_key(&key1)); + assert!(acc.increases.contains_key(&key2)); + } + + #[test] + fn test_keyed_counter_state_query() { + let mut acc = KeyedCounterState::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let increase_acc = create_test_increase_accumulator(10.0, 25.0); + acc.update(key.clone(), increase_acc); + + // Test increase query + assert_eq!(acc.query(Statistic::Increase, &key, None).unwrap(), 15.0); + + // Test rate query (15.0 increase over 1 second = 15.0 per second) + assert_eq!(acc.query(Statistic::Rate, &key, None).unwrap(), 15.0); + + // Sum returns the latest cumulative counter value for the + // queried key (per-series Prometheus `sum()` semantics; + // see issue ProjectASAP/ASAPCollector#46 and PR #108 diagnosis). + // The series here was created with last_seen=25.0. + assert_eq!(acc.query(Statistic::Sum, &key, None).unwrap(), 25.0); + + // Unsupported statistic still errors. + assert!(acc.query(Statistic::Min, &key, None).is_err()); + + let unknown_key = KeyByLabelValues::new(); + assert!(acc.query(Statistic::Increase, &unknown_key, None).is_err()); + } + + #[test] + fn test_keyed_counter_state_sum_per_key() { + // `sum by (zone) (counter)` reaches KeyedCounterState + // only when the ASAP-tier ingest groups multiple series under + // a single accumulator (the `Multiple*` variant). In that case + // each per-key Sum should be the series' latest cumulative + // value; the engine's outer `by` aggregation does the cross-key + // grouping. (Issue ProjectASAP/ASAPCollector#46.) + let mut acc = KeyedCounterState::new(); + let east = KeyByLabelValues::new_with_labels(vec!["us-east-1".to_string()]); + let west = KeyByLabelValues::new_with_labels(vec!["us-west-2".to_string()]); + + acc.update( + east.clone(), + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000), + ); + acc.update( + west.clone(), + IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000), + ); + + assert_eq!(acc.query(Statistic::Sum, &east, None).unwrap(), 100.0); + assert_eq!(acc.query(Statistic::Sum, &west, None).unwrap(), 50.0); + } + + #[test] + fn test_keyed_counter_state_merge() { + let mut acc1 = KeyedCounterState::new(); + let mut acc2 = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + // Add different keys to each accumulator + acc1.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); + acc2.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); + + // Also add overlapping key with different time ranges (later timestamps) + acc2.update( + key1.clone(), + create_test_increase_accumulator_with_time(15.0, 2000, 30.0, 3000), + ); // Later time range + + let merged = KeyedCounterState::merge_accumulators(vec![acc1, acc2]).unwrap(); + + assert_eq!(merged.increases.len(), 2); + assert!(merged.increases.contains_key(&key1)); + assert!(merged.increases.contains_key(&key2)); + + // The merged key1 should have the full range (earliest start to latest end) + let merged_key1 = merged.increases.get(&key1).unwrap(); + assert_eq!(merged_key1.starting_measurement.value, 10.0); // Earlier start + assert_eq!(merged_key1.last_seen_measurement.value, 30.0); // Later end + } + + #[test] + fn test_keyed_counter_state_serialization() { + let mut acc = KeyedCounterState::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let second_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + let mut reset_aware = create_test_increase_accumulator(10.0, 25.0); + reset_aware.update(Measurement::new(3.0), 3000); + acc.update(key.clone(), reset_aware); + acc.update( + second_key.clone(), + create_test_increase_accumulator(4.0, 9.0), + ); + + // Test JSON serialization + let json_value = acc.serialize_to_json(); + let deserialized = KeyedCounterState::deserialize_from_json(&json_value).unwrap(); + + assert_eq!(deserialized.increases.len(), 2); + let deserialized_acc = deserialized.increases.get(&key).unwrap(); + assert_eq!(deserialized_acc.starting_measurement.value, 10.0); + assert_eq!(deserialized_acc.last_seen_measurement.value, 3.0); + assert_eq!(deserialized_acc.total_increase, 18.0); + + // Test binary serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = KeyedCounterState::deserialize_from_bytes(&bytes).unwrap(); + + assert_eq!(deserialized_bytes.increases.len(), 2); + let deserialized_acc_bytes = deserialized_bytes.increases.get(&key).unwrap(); + assert_eq!(deserialized_acc_bytes.starting_measurement.value, 10.0); + assert_eq!(deserialized_acc_bytes.last_seen_measurement.value, 3.0); + assert_eq!(deserialized_acc_bytes.total_increase, 18.0); + assert_eq!( + deserialized_bytes + .increases + .get(&second_key) + .unwrap() + .last_seen_measurement + .value, + 9.0 + ); + } + + #[test] + fn test_keyed_counter_state_get_keys() { + let mut acc = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); + acc.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); + + let keys = acc.get_keys().unwrap(); + assert_eq!(keys.len(), 2); + assert!(keys.contains(&key1)); + assert!(keys.contains(&key2)); + } + + #[test] + fn test_trait_object() { + let mut acc = KeyedCounterState::new(); + let key = KeyByLabelValues::new(); + acc.update(key.clone(), create_test_increase_accumulator(10.0, 25.0)); + + let trait_obj: Box = Box::new(acc); + assert_eq!( + trait_obj.query(Statistic::Increase, &key, None).unwrap(), + 15.0 + ); + + let keys = trait_obj.get_keys().unwrap(); + assert_eq!(keys.len(), 1); + } + + // #[test] + // fn test_keyed_counter_state_arroyo_deserialization() { + // // Create test data in Arroyo MessagePack format + // // Format: {key: [starting_value, starting_timestamp, last_seen_value, last_seen_timestamp]} + // let mut test_data = std::collections::HashMap::new(); + // test_data.insert("web;service".to_string(), vec![10.0, 1000.0, 25.0, 2000.0]); + // test_data.insert("api;service".to_string(), vec![5.0, 1500.0, 15.0, 2500.0]); + + // // Serialize to MessagePack + // let arroyo_buffer = rmp_serde::to_vec(&test_data).unwrap(); + + // // Test Arroyo deserialization + // let deserialized_acc = + // KeyedCounterState::deserialize_from_bytes_arroyo(&arroyo_buffer).unwrap(); + + // // Verify the deserialized accumulator has the correct data + // assert_eq!(deserialized_acc.increases.len(), 2); + + // // Check first key (web;service) + // let keys: Vec<_> = deserialized_acc.increases.keys().collect(); + // let key1 = keys + // .iter() + // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "web")) + // .unwrap(); + + // let increase1 = deserialized_acc.increases.get(key1).unwrap(); + // assert_eq!(increase1.starting_measurement.value, 10.0); + // assert_eq!(increase1.starting_timestamp, 1000); + // assert_eq!(increase1.last_seen_measurement.value, 25.0); + // assert_eq!(increase1.last_seen_timestamp, 2000); + + // // Check second key (api;service) + // let key2 = keys + // .iter() + // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "api")) + // .unwrap(); + + // let increase2 = deserialized_acc.increases.get(key2).unwrap(); + // assert_eq!(increase2.starting_measurement.value, 5.0); + // assert_eq!(increase2.starting_timestamp, 1500); + // assert_eq!(increase2.last_seen_measurement.value, 15.0); + // assert_eq!(increase2.last_seen_timestamp, 2500); + + // // Test querying + // assert_eq!( + // deserialized_acc.query(Statistic::Increase, key1).unwrap(), + // 15.0 + // ); // 25.0 - 10.0 + // assert_eq!( + // deserialized_acc.query(Statistic::Increase, key2).unwrap(), + // 10.0 + // ); // 15.0 - 5.0 + // } +} diff --git a/crates/asap-physical-operators/src/accumulators/keyed_max_state.rs b/crates/asap-physical-operators/src/accumulators/keyed_max_state.rs new file mode 100644 index 000000000..30a4a666b --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/keyed_max_state.rs @@ -0,0 +1,335 @@ +use crate::{ + 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::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedMaxState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedMaxState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Max => self + .values + .get(key) + .copied() + .ok_or_else(|| format!("Key {key} not found in KeyedMaxState").into()), + other => Err(format!("Unsupported statistic in KeyedMaxState: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedMaxState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedMaxState::new(); + + for acc in accumulators { + for (key, value) in acc.values { + result.update(key, value); + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(value: &str) -> KeyByLabelValues { + KeyByLabelValues::new_with_labels(vec![value.to_string()]) + } + + #[test] + fn keeps_the_largest_per_key() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 10.0); + acc.update(key("a"), 5.0); + acc.update(key("a"), 15.0); + acc.update(key("b"), 7.0); + + assert_eq!(acc.query(Statistic::Max, &key("a"), None).unwrap(), 15.0); + assert_eq!(acc.query(Statistic::Max, &key("b"), None).unwrap(), 7.0); + } + + #[test] + fn refuses_the_opposite_statistic_and_unknown_keys() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 1.0); + assert!(acc.query(Statistic::Min, &key("a"), None).is_err()); + assert!(acc.query(Statistic::Max, &key("missing"), None).is_err()); + } + + #[test] + fn merges_per_key() { + let mut left = KeyedMaxState::new(); + left.update(key("a"), 10.0); + let mut right = KeyedMaxState::new(); + right.update(key("a"), 5.0); + right.update(key("b"), 3.0); + + let merged = + >::merge_accumulators(vec![ + left, right, + ]) + .unwrap(); + + assert_eq!(merged.query(Statistic::Max, &key("a"), None).unwrap(), 10.0); + assert_eq!(merged.query(Statistic::Max, &key("b"), None).unwrap(), 3.0); + } + + #[test] + fn refuses_to_merge_with_the_opposite_direction() { + use super::super::keyed_min_state::KeyedMinState; + let mine = KeyedMaxState::new(); + let theirs = KeyedMinState::new(); + assert!(mine.merge_with(&theirs).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 4.0); + + let json = acc.serialize_to_json(); + let from_json = KeyedMaxState::deserialize_from_json(&json).unwrap(); + assert_eq!( + from_json.query(Statistic::Max, &key("a"), None).unwrap(), + 4.0 + ); + + let bytes = acc.serialize_to_bytes(); + let from_bytes = KeyedMaxState::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + from_bytes.query(Statistic::Max, &key("a"), None).unwrap(), + 4.0 + ); + } +} diff --git a/crates/asap-physical-operators/src/accumulators/keyed_min_state.rs b/crates/asap-physical-operators/src/accumulators/keyed_min_state.rs new file mode 100644 index 000000000..f6bbf2be9 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/keyed_min_state.rs @@ -0,0 +1,335 @@ +use crate::{ + 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::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedMinState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedMinState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Min => self + .values + .get(key) + .copied() + .ok_or_else(|| format!("Key {key} not found in KeyedMinState").into()), + other => Err(format!("Unsupported statistic in KeyedMinState: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedMinState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedMinState::new(); + + for acc in accumulators { + for (key, value) in acc.values { + result.update(key, value); + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(value: &str) -> KeyByLabelValues { + KeyByLabelValues::new_with_labels(vec![value.to_string()]) + } + + #[test] + fn keeps_the_smallest_per_key() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 10.0); + acc.update(key("a"), 5.0); + acc.update(key("a"), 15.0); + acc.update(key("b"), 7.0); + + assert_eq!(acc.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); + assert_eq!(acc.query(Statistic::Min, &key("b"), None).unwrap(), 7.0); + } + + #[test] + fn refuses_the_opposite_statistic_and_unknown_keys() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 1.0); + assert!(acc.query(Statistic::Max, &key("a"), None).is_err()); + assert!(acc.query(Statistic::Min, &key("missing"), None).is_err()); + } + + #[test] + fn merges_per_key() { + let mut left = KeyedMinState::new(); + left.update(key("a"), 10.0); + let mut right = KeyedMinState::new(); + right.update(key("a"), 5.0); + right.update(key("b"), 3.0); + + let merged = + >::merge_accumulators(vec![ + left, right, + ]) + .unwrap(); + + assert_eq!(merged.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); + assert_eq!(merged.query(Statistic::Min, &key("b"), None).unwrap(), 3.0); + } + + #[test] + fn refuses_to_merge_with_the_opposite_direction() { + use super::super::keyed_max_state::KeyedMaxState; + let mine = KeyedMinState::new(); + let theirs = KeyedMaxState::new(); + assert!(mine.merge_with(&theirs).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 4.0); + + let json = acc.serialize_to_json(); + let from_json = KeyedMinState::deserialize_from_json(&json).unwrap(); + assert_eq!( + from_json.query(Statistic::Min, &key("a"), None).unwrap(), + 4.0 + ); + + let bytes = acc.serialize_to_bytes(); + let from_bytes = KeyedMinState::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + from_bytes.query(Statistic::Min, &key("a"), None).unwrap(), + 4.0 + ); + } +} diff --git a/crates/asap-physical-operators/src/accumulators/keyed_sum_count_accumulator.rs b/crates/asap-physical-operators/src/accumulators/keyed_sum_count_accumulator.rs new file mode 100644 index 000000000..486b3fee3 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/keyed_sum_count_accumulator.rs @@ -0,0 +1,558 @@ +use crate::{ + 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::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for KeyedSumCountAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedSumCountAccumulator { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match (&self.family, statistic) { + (ExactKind::Sum, Statistic::Sum) => self.sums.get(key).copied().ok_or_else(|| { + "Key not found in KeyedSumCountAccumulator" + .to_string() + .into() + }), + (ExactKind::Count, Statistic::Count) => self + .counts + .get(key) + .map(|count| *count as f64) + .ok_or_else(|| { + "Sample count unavailable in KeyedSumCountAccumulator" + .to_string() + .into() + }), + _ => Err( + format!("Unsupported statistic in KeyedSumCountAccumulator: {statistic:?}").into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedSumCountAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let family = accumulators[0].family.clone(); + if accumulators.iter().any(|acc| acc.family != family) { + return Err("Cannot merge different keyed additive families".into()); + } + let mut result = KeyedSumCountAccumulator::for_family(family); + + for acc in accumulators { + for key in acc.sums.keys() { + match ( + result.counts.get(key).copied(), + acc.counts.get(key).copied(), + ) { + (None, Some(count)) if !result.sums.contains_key(key) => { + result.counts.insert(key.clone(), count); + } + (Some(existing), Some(count)) => { + if let Some(total) = existing.checked_add(count) { + result.counts.insert(key.clone(), total); + } else { + result.counts.remove(key); + } + } + _ => { + result.counts.remove(key); + } + } + } + for (key, sum) in acc.sums { + *result.sums.entry(key).or_insert(0.0) += sum; + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use std::vec; + + use super::*; + + #[test] + fn test_keyed_sum_count_accumulator_creation() { + let acc = KeyedSumCountAccumulator::new(); + assert!(acc.sums.is_empty()); + } + + #[test] + fn test_keyed_sum_count_accumulator_update() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.update(key1.clone(), 10.0); + acc.update(key2.clone(), 20.0); + acc.update(key1.clone(), 5.0); // Should add to existing + + assert_eq!(acc.sums.get(&key1), Some(&15.0)); + assert_eq!(acc.sums.get(&key2), Some(&20.0)); + } + + #[test] + fn grouped_count_reads_sample_count_and_survives_merge_and_round_trip() { + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let mut first = KeyedSumCountAccumulator::for_family(ExactKind::Count); + first.update(key.clone(), 10.0); + first.update(key.clone(), 20.0); + let mut second = KeyedSumCountAccumulator::for_family(ExactKind::Count); + second.update(key.clone(), 7.0); + let merged = KeyedSumCountAccumulator::merge_accumulators(vec![first, second]).unwrap(); + for acc in [ + merged.clone(), + KeyedSumCountAccumulator::deserialize_from_json(&merged.serialize_to_json()).unwrap(), + KeyedSumCountAccumulator::deserialize_from_bytes(&merged.serialize_to_bytes()).unwrap(), + ] { + assert_eq!(acc.family, ExactKind::Count); + assert!(acc.query(Statistic::Sum, &key, None).is_err()); + assert_eq!(acc.query(Statistic::Count, &key, None).unwrap(), 3.0); + } + } + + #[test] + fn keyed_additive_merge_rejects_different_planner_families() { + assert!(KeyedSumCountAccumulator::merge_accumulators(vec![ + KeyedSumCountAccumulator::for_family(ExactKind::Sum), + KeyedSumCountAccumulator::for_family(ExactKind::Count), + ]) + .is_err()); + } + + #[test] + fn test_keyed_sum_count_accumulator_query() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); + + acc.add_sum(key.clone(), 42.0); + + // Test total queries (querying with the specific key) + assert_eq!( + crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Sum, &key, None).unwrap(), + 42.0 + ); + + // Test error cases + assert!( + crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Min, &key, None).is_err() + ); + } + + #[test] + fn test_keyed_sum_count_accumulator_get_keys() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.add_sum(key1.clone(), 10.0); + acc.add_sum(key2.clone(), 20.0); + + let keys = crate::AggregateCore::get_keys(&acc).unwrap(); + assert_eq!(keys.len(), 2); + assert!(keys.contains(&key1)); + assert!(keys.contains(&key2)); + } + + #[test] + fn test_keyed_sum_count_accumulator_merge() { + let mut acc1 = KeyedSumCountAccumulator::new(); + let mut acc2 = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc1.add_sum(key1.clone(), 10.0); + acc1.add_sum(key2.clone(), 20.0); + + acc2.add_sum(key1.clone(), 5.0); // Same key, different accumulator + + let merged = >::merge_accumulators(vec![acc1, acc2]).unwrap(); + + assert_eq!(merged.sums.get(&key1), Some(&15.0)); // Should be merged + assert_eq!(merged.sums.get(&key2), Some(&20.0)); // Should be preserved + } + + #[test] + fn test_keyed_sum_count_accumulator_serialization() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); + + acc.add_sum(key.clone(), 42.5); + + // Test JSON serialization + let json = acc.serialize_to_json(); + let deserialized = KeyedSumCountAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!(deserialized.sums.get(&key), Some(&42.5)); + + // Test byte serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = KeyedSumCountAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(deserialized_bytes.sums.get(&key), Some(&42.5)); + } + + #[test] + fn test_trait_object() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + acc.add_sum(key.clone(), 42.0); + + let trait_obj: Box = Box::new(acc); + + // Test type name through trait object + assert_eq!(trait_obj.type_name(), "KeyedSumCountAccumulator"); + } +} diff --git a/crates/asap-physical-operators/src/accumulators/max_accumulator.rs b/crates/asap-physical-operators/src/accumulators/max_accumulator.rs new file mode 100644 index 000000000..235b5fc98 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/max_accumulator.rs @@ -0,0 +1,248 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use 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::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/crates/asap-physical-operators/src/accumulators/min_accumulator.rs b/crates/asap-physical-operators/src/accumulators/min_accumulator.rs new file mode 100644 index 000000000..fff2fa0e9 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/min_accumulator.rs @@ -0,0 +1,253 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use 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::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/crates/asap-physical-operators/src/accumulators/mod.rs b/crates/asap-physical-operators/src/accumulators/mod.rs new file mode 100644 index 000000000..073db6e82 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/mod.rs @@ -0,0 +1,37 @@ +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/crates/asap-physical-operators/src/accumulators/sketch_envelope_accumulator.rs b/crates/asap-physical-operators/src/accumulators/sketch_envelope_accumulator.rs new file mode 100644 index 000000000..930dbe6d9 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/sketch_envelope_accumulator.rs @@ -0,0 +1,154 @@ +//! SketchEnvelopeAccumulator — wraps a raw SketchEnvelope protobuf payload +//! received via OTLP ingest so it can be stored through the `Store` trait. +//! +//! The accumulator preserves the opaque proto bytes and decodes them lazily +//! (via `SketchEnvelope::decode`) only when merge or query operations need +//! the inner sketch type. + +use crate::{AggregateCore, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; +use prost::Message; +use serde_json::Value; +use std::collections::HashMap; + +use 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::MultipleSubpopulationAggregate for SketchEnvelopeAccumulator { + fn query( + &self, + _statistic: Statistic, + _key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Err( + "SketchEnvelopeAccumulator: direct query not supported; use native sketch query path" + .into(), + ) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} diff --git a/crates/asap-physical-operators/src/accumulators/sum_accumulator.rs b/crates/asap-physical-operators/src/accumulators/sum_accumulator.rs new file mode 100644 index 000000000..c5293911d --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/sum_accumulator.rs @@ -0,0 +1,413 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use 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::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/crates/asap-physical-operators/src/accumulators/univmon_accumulator.rs b/crates/asap-physical-operators/src/accumulators/univmon_accumulator.rs new file mode 100644 index 000000000..2d18895e0 --- /dev/null +++ b/crates/asap-physical-operators/src/accumulators/univmon_accumulator.rs @@ -0,0 +1,234 @@ +//! One frequency state shared by count, distinct, L2 and entropy readouts. + +use crate::{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/crates/asap-physical-operators/src/arithmetic.rs b/crates/asap-physical-operators/src/arithmetic.rs new file mode 100644 index 000000000..bfc50694b --- /dev/null +++ b/crates/asap-physical-operators/src/arithmetic.rs @@ -0,0 +1,19 @@ +//! Float64 arithmetic shared by ASAP execution engines. +//! Preserve IEEE non-finite results; callers own their output policies. + +pub 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/crates/asap-physical-operators/src/capability.rs b/crates/asap-physical-operators/src/capability.rs new file mode 100644 index 000000000..fdd2fee7a --- /dev/null +++ b/crates/asap-physical-operators/src/capability.rs @@ -0,0 +1,115 @@ +//! Allocation-free checks for the concrete summary kernels in this crate. +use planner_types::post_asap::{ + ExactKind, ExactParams, GroupingStrategy, SketchAlgorithm, SketchParams, SummaryFamilyType, + SummaryUpdate, +}; + +/// Check the same contract used by `create_planner_accumulator` before a plan +/// is accepted. Execution timing is deliberately not a kernel property. +pub fn validate_summary_kernel( + family: &SummaryFamilyType, + input: &SummaryUpdate, + grouping: &GroupingStrategy, +) -> Result<(), String> { + if grouping != &GroupingStrategy::PerSubpopulationInstance { + return Err("shared summary grouping has no registered kernel".into()); + } + let keyed = match family { + SummaryFamilyType::ExactAggregate(kind, params) => { + use ExactKind as K; + use ExactParams as P; + if !matches!( + (kind, params), + (K::Sum, P::Sum) + | (K::Count, P::Count) + | (K::Min, P::Min) + | (K::Max, P::Max) + | (K::Rate, P::Rate) + | (K::Increase, P::Increase) + ) { + return Err(format!("unsupported exact kernel {family:?}")); + } + input.item.is_some() + } + SummaryFamilyType::Sketch(kind, layout) => { + if layout != grouping { + return Err("Planner family and operator grouping disagree".into()); + } + use SketchAlgorithm as A; + use SketchParams as P; + match (kind.algorithm(), kind.params()) { + (A::Kll, P::Kll { k }) if (8..=u16::MAX as u32).contains(k) => false, + (A::DDSketch, P::DDSketch { alpha }) + if alpha.is_finite() && *alpha > 0.0 && *alpha < 1.0 => + { + false + } + (A::Hll, P::Hll { precision }) if (4..=18).contains(precision) => false, + (A::Cms, P::Cms { width, depth }) + | (A::CountSketch, P::CountSketch { width, depth }) + if valid_matrix(*width, *depth) => + { + true + } + ( + A::CmsWithHeap, + P::CmsWithHeap { + width, + depth, + heap_size, + }, + ) + | ( + A::CountSketchWithHeap, + P::CountSketchWithHeap { + width, + depth, + heap_size, + }, + ) if valid_matrix(*width, *depth) && *heap_size > 0 => true, + ( + A::UnivMon, + P::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + }, + ) if *heap_size > 0 + && *sketch_cols > 0 + && (1..=20).contains(sketch_rows) + && (1..=64).contains(layers) + && (*sketch_rows as usize) + .checked_mul(*sketch_cols as usize) + .and_then(|n| n.checked_mul(*layers as usize)) + .is_some() => + { + false + } + _ => { + return Err(format!( + "unsupported kernel or invalid parameters: {kind:?}" + )) + } + } + } + _ => return Err(format!("unsupported summary kernel {family:?}")), + }; + if keyed != input.item.is_some() + && !asap_types::accumulator_spec::is_unit_sample_frequency(input) + { + return Err("Planner item expression does not match kernel layout".into()); + } + Ok(()) +} + +fn valid_matrix(width: u32, depth: u32) -> bool { + // Construction uses the kernel's native row hashing. Packed-wire decoder + // limits describe a different representation and must not reject it here. + width > 0 + && depth > 0 + && (width as usize) + .checked_mul(depth as usize) + .and_then(|n| n.checked_mul(std::mem::size_of::())) + .is_some() +} diff --git a/crates/asap-physical-operators/src/dag/mod.rs b/crates/asap-physical-operators/src/dag/mod.rs new file mode 100644 index 000000000..b5a2e7b7d --- /dev/null +++ b/crates/asap-physical-operators/src/dag/mod.rs @@ -0,0 +1,515 @@ +//! Independent operator DAG execution. No backend plan or engine types are used. +//! +//! Each run creates one stream per reachable node. Consumers subscribe to that +//! stream independently; retained outputs are released after the last consumer. +use futures::{stream::LocalBoxStream, Stream}; +use std::{ + cell::{Cell, RefCell}, + collections::{BTreeMap, BTreeSet, VecDeque}, + fmt::Debug, + pin::Pin, + rc::Rc, + sync::Arc, + task::{Context, Poll, Waker}, +}; + +pub type NodeId = u64; +pub type OutputStream<'a, V> = LocalBoxStream<'a, Result>; +#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] +pub enum Error { + #[error("invalid DAG: {0}")] + Invalid(String), + #[error("operator failed: {0}")] + Operator(String), + #[error("node {node} ({operation}) failed: {source}")] + AtNode { + node: NodeId, + operation: String, + source: Box, + }, + #[error("execution memory limit exceeded")] + MemoryLimit, + #[error("execution cancelled")] + Cancelled, +} + +/// Scope is part of an execution instance, never mutable state in a reusable plan. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Scope { + Ingestion { + window_start_ms: i64, + window_end_ms: i64, + revision: u64, + }, + Query { + evaluation_time_ms: i64, + revision: u64, + }, +} +#[derive(Clone, Debug)] +pub struct Limits { + pub max_buffered_batches: usize, + pub max_bytes: usize, +} +impl Default for Limits { + fn default() -> Self { + Self { + max_buffered_batches: 8, + max_bytes: 64 * 1024 * 1024, + } + } +} +struct Control { + cancelled: Cell, + bytes: Cell, + peak: Cell, + limits: Limits, + waiters: RefCell>, +} +#[derive(Clone)] +pub struct RunContext { + pub scope: Scope, + control: Rc, +} +impl RunContext { + pub fn new(scope: Scope, limits: Limits) -> Result { + if limits.max_buffered_batches == 0 || limits.max_bytes == 0 { + return Err(Error::Invalid("execution limits must be positive".into())); + } + if matches!(&scope, Scope::Ingestion { window_start_ms, window_end_ms, .. } if window_start_ms > window_end_ms) + { + return Err(Error::Invalid("inverted ingestion window".into())); + } + Ok(Self { + scope, + control: Rc::new(Control { + cancelled: Cell::new(false), + bytes: Cell::new(0), + peak: Cell::new(0), + limits, + waiters: RefCell::new(Vec::new()), + }), + }) + } + pub fn cancel(&self) { + self.control.cancelled.set(true); + for waiter in self.control.waiters.borrow_mut().drain(..) { + waiter.wake(); + } + } + pub fn is_cancelled(&self) -> bool { + self.control.cancelled.get() + } + pub fn retained_bytes(&self) -> usize { + self.control.bytes.get() + } + pub fn peak_bytes(&self) -> usize { + self.control.peak.get() + } + pub fn reserve(&self, bytes: usize) -> Result { + let total = self + .control + .bytes + .get() + .checked_add(bytes) + .ok_or(Error::MemoryLimit)?; + if total > self.control.limits.max_bytes { + return Err(Error::MemoryLimit); + } + self.control.bytes.set(total); + self.control.peak.set(self.control.peak.get().max(total)); + Ok(Reservation { + bytes, + control: Rc::clone(&self.control), + }) + } + fn register(&self, waker: &Waker) { + let mut waiters = self.control.waiters.borrow_mut(); + if !waiters.iter().any(|old| old.will_wake(waker)) { + waiters.push(waker.clone()); + } + } +} +pub struct Reservation { + bytes: usize, + control: Rc, +} +impl Reservation { + /// Adjust an operator-owned allocation without accumulating bookkeeping entries. + pub fn resize(&mut self, bytes: usize) -> Result<(), Error> { + let total = self + .control + .bytes + .get() + .checked_sub(self.bytes) + .and_then(|total| total.checked_add(bytes)) + .ok_or(Error::MemoryLimit)?; + if total > self.control.limits.max_bytes { + return Err(Error::MemoryLimit); + } + self.control.bytes.set(total); + self.control.peak.set(self.control.peak.get().max(total)); + self.bytes = bytes; + Ok(()) + } +} +impl Drop for Reservation { + fn drop(&mut self) { + self.control + .bytes + .set(self.control.bytes.get().saturating_sub(self.bytes)); + } +} + +/// An output owns its memory reservation even after it leaves the DAG's queue. +pub struct SharedValue { + value: Arc, + _reservation: Rc, +} +impl Clone for SharedValue { + fn clone(&self) -> Self { + Self { + value: Arc::clone(&self.value), + _reservation: Rc::clone(&self._reservation), + } + } +} +impl std::ops::Deref for SharedValue { + type Target = V; + fn deref(&self) -> &V { + &self.value + } +} +impl SharedValue { + pub fn value(&self) -> &V { + &self.value + } +} + +/// Operators own computation. The runtime provides already-connected inputs; +/// an operator must not recursively execute another plan node itself. +pub trait PhysicalOperator { + fn name(&self) -> &str; + fn input_schemas(&self) -> Vec; + fn output_schema(&self) -> S; + fn start<'a>( + &'a self, + inputs: Vec>, + context: RunContext, + ) -> Result, Error>; + fn output_bytes(&self, value: &V) -> usize; +} +struct Node<'a, V, S> { + inputs: Vec, + operator: Box + 'a>, +} +pub struct PhysicalDag<'a, V, S> { + nodes: BTreeMap>, +} +impl Default for PhysicalDag<'_, V, S> { + fn default() -> Self { + Self { + nodes: BTreeMap::new(), + } + } +} +impl<'a, V: 'a, S: Clone + PartialEq + Debug + 'a> PhysicalDag<'a, V, S> { + pub fn add( + &mut self, + id: NodeId, + inputs: Vec, + operator: impl PhysicalOperator + 'a, + ) -> Result<(), Error> { + self.add_boxed(id, inputs, Box::new(operator)) + } + pub fn add_boxed( + &mut self, + id: NodeId, + inputs: Vec, + operator: Box + 'a>, + ) -> Result<(), Error> { + if self.nodes.contains_key(&id) { + return Err(Error::Invalid(format!("duplicate node {id}"))); + } + self.nodes.insert(id, Node { inputs, operator }); + Ok(()) + } + pub fn validate(&self, roots: &[NodeId]) -> Result<(), Error> { + fn visit( + dag: &PhysicalDag<'_, V, S>, + id: NodeId, + active: &mut BTreeSet, + done: &mut BTreeMap, + ) -> Result { + if let Some(depth) = done.get(&id) { + return Ok(*depth); + } + if active.len() >= 128 { + return Err(Error::Invalid( + "DAG exceeds the supported execution depth of 128".into(), + )); + } + if !active.insert(id) { + return Err(Error::Invalid(format!("cycle at node {id}"))); + } + let node = dag + .nodes + .get(&id) + .ok_or_else(|| Error::Invalid(format!("missing node {id}")))?; + let expected = node.operator.input_schemas(); + if expected.len() != node.inputs.len() { + return Err(Error::Invalid(format!("node {id} input arity mismatch"))); + } + let mut depth = 1; + for (input, schema) in node.inputs.iter().zip(expected) { + depth = depth.max(1 + visit(dag, *input, active, done)?); + let actual = dag.nodes[input].operator.output_schema(); + if actual != schema { + return Err(Error::Invalid(format!( + "node {id} input {input} schema mismatch: {actual:?} vs {schema:?}" + ))); + } + } + if depth > 128 { + return Err(Error::Invalid( + "DAG exceeds the supported execution depth of 128".into(), + )); + } + active.remove(&id); + done.insert(id, depth); + Ok(depth) + } + if roots.is_empty() { + return Err(Error::Invalid("execution needs a root".into())); + } + let mut done = BTreeMap::new(); + for &root in roots { + visit(self, root, &mut BTreeSet::new(), &mut done)?; + } + Ok(()) + } + pub fn execute<'r>( + &'r self, + roots: &[NodeId], + context: RunContext, + ) -> Result>, Error> + where + 'a: 'r, + { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + self.validate(roots)?; + fn build<'r, V: 'r, S: 'r>( + dag: &'r PhysicalDag<'_, V, S>, + id: NodeId, + context: &RunContext, + states: &mut BTreeMap>>>, + ) -> Result>>, Error> { + if let Some(state) = states.get(&id) { + return Ok(Rc::clone(state)); + } + let node = &dag.nodes[&id]; + let mut inputs = Vec::new(); + for &child in &node.inputs { + inputs.push(Input::subscribe(build(dag, child, context, states)?)); + } + let stream = node + .operator + .start(inputs, context.clone()) + .map_err(|source| Error::AtNode { + node: id, + operation: node.operator.name().into(), + source: Box::new(source), + })?; + let op = node.operator.as_ref(); + let state = Rc::new(RefCell::new(Producer { + stream: Some(stream), + node: id, + operation: node.operator.name().into(), + size: Box::new(move |value| op.output_bytes(value)), + context: context.clone(), + queue: VecDeque::new(), + base: 0, + next_reader: 0, + batches_polled: 0, + readers: BTreeMap::new(), + waiters: BTreeMap::new(), + finished: false, + failure: None, + })); + states.insert(id, Rc::clone(&state)); + Ok(state) + } + let mut states = BTreeMap::new(); + roots + .iter() + .map(|&id| build(self, id, &context, &mut states).map(Input::subscribe)) + .collect() + } +} +struct Producer<'a, V> { + node: NodeId, + operation: String, + stream: Option>, + size: Box usize + 'a>, + context: RunContext, + queue: VecDeque>, + base: u64, + next_reader: u64, + batches_polled: usize, + readers: BTreeMap, + waiters: BTreeMap, + finished: bool, + failure: Option, +} +impl Producer<'_, V> { + fn trim(&mut self) { + let minimum = self + .readers + .values() + .copied() + .min() + .unwrap_or(self.base + self.queue.len() as u64); + while self.base < minimum { + self.queue.pop_front(); + self.base += 1; + } + for (_, waker) in std::mem::take(&mut self.waiters) { + waker.wake(); + } + if self.readers.is_empty() { + self.stream = None; + self.queue.clear(); + } + } +} +pub struct Input<'a, V> { + producer: Rc>>, + reader: u64, + done: bool, +} +impl<'a, V> Input<'a, V> { + fn subscribe(producer: Rc>>) -> Self { + let reader = { + let mut state = producer.borrow_mut(); + let id = state.next_reader; + state.next_reader += 1; + let base = state.base; + state.readers.insert(id, base); + id + }; + Self { + producer, + reader, + done: false, + } + } +} +impl Drop for Input<'_, V> { + fn drop(&mut self) { + let mut state = self.producer.borrow_mut(); + state.readers.remove(&self.reader); + state.waiters.remove(&self.reader); + state.trim(); + } +} +impl Stream for Input<'_, V> { + type Item = Result, Error>; + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); + if this.done { + return Poll::Ready(None); + } + let mut state = this.producer.borrow_mut(); + state.context.register(cx.waker()); + if state.context.is_cancelled() { + state.failure = Some(Error::Cancelled); + state.finished = true; + state.stream = None; + state.queue.clear(); + } + let position = state.readers[&this.reader]; + let index = (position - state.base) as usize; + if let Some(value) = state.queue.get(index).cloned() { + state.readers.insert(this.reader, position + 1); + state.trim(); + return Poll::Ready(Some(Ok(value))); + } + if state.finished { + this.done = true; + state.readers.remove(&this.reader); + let failure = state.failure.clone(); + state.trim(); + return Poll::Ready(failure.map(Err)); + } + state.waiters.insert(this.reader, cx.waker().clone()); + if state.queue.len() >= state.context.control.limits.max_buffered_batches { + return Poll::Pending; + } + // Always-ready sources must still give cancellation and other roots a turn. + if state.batches_polled >= 32 { + state.batches_polled = 0; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + let polled = state + .stream + .as_mut() + .expect("unfinished producer") + .as_mut() + .poll_next(cx); + if matches!(&polled, Poll::Ready(Some(Ok(_)))) { + state.batches_polled += 1; + } + match polled { + Poll::Pending => Poll::Pending, + Poll::Ready(Some(Ok(value))) => match state.context.reserve((state.size)(&value)) { + Ok(reservation) => { + let value = SharedValue { + value: Arc::new(value), + _reservation: Rc::new(reservation), + }; + state.queue.push_back(value.clone()); + state.readers.insert(this.reader, position + 1); + state.trim(); + Poll::Ready(Some(Ok(value))) + } + Err(error) => { + state.failure = Some(error.clone()); + state.finished = true; + state.stream = None; + this.done = true; + state.readers.remove(&this.reader); + state.trim(); + Poll::Ready(Some(Err(error))) + } + }, + Poll::Ready(result) => { + let error = result.and_then(Result::err).map(|source| match source { + Error::AtNode { .. } | Error::Cancelled | Error::MemoryLimit => source, + source => Error::AtNode { + node: state.node, + operation: state.operation.clone(), + source: Box::new(source), + }, + }); + state.failure = error.clone(); + state.finished = true; + state.stream = None; + this.done = true; + state.readers.remove(&this.reader); + state.trim(); + Poll::Ready(error.map(Err)) + } + } + } +} + +pub mod operators; +pub mod values; + +#[cfg(test)] +mod tests; + +pub mod planner; diff --git a/crates/asap-physical-operators/src/dag/operators.rs b/crates/asap-physical-operators/src/dag/operators.rs new file mode 100644 index 000000000..eeb1a922e --- /dev/null +++ b/crates/asap-physical-operators/src/dag/operators.rs @@ -0,0 +1,1170 @@ +//! Native DAG operators. Engines bind sources; computation lives here. +use super::{ + values::{group_key, Batch, Schema, Value}, + Error, Input, OutputStream, PhysicalOperator, Reservation, RunContext, +}; +use futures::StreamExt; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema, SummaryUpdate}, + pre_asap::{ArithmeticOpKind, ColumnRef, DataType}, +}; +use std::{collections::BTreeMap, sync::Arc}; + +fn invalid(message: &str) -> Error { + Error::Invalid(message.into()) +} +fn field(schema: &Schema, column: usize) -> Result<&SummaryField, Error> { + schema + .fields + .get(column) + .ok_or_else(|| invalid("column out of range")) +} +fn plain(schema: &Schema, column: usize) -> Result<(&DataType, bool), Error> { + let f = field(schema, column)?; + let SummaryFamilyType::Plain(dtype) = &f.dtype else { + return Err(invalid("plain value required")); + }; + Ok((dtype, f.nullable)) +} +fn schema(fields: Vec) -> Schema { + Arc::new(SummarySchema { + fields, + time_index: None, + }) +} +fn result_field(name: &str, dtype: DataType, nullable: bool) -> SummaryField { + SummaryField { + name: name.into(), + dtype: SummaryFamilyType::Plain(dtype), + nullable, + } +} + +#[derive(Clone, Debug)] +pub enum Expression { + Column(usize), + Literal { + value: Value, + dtype: DataType, + }, + Negate(Box), + Arithmetic { + op: ArithmeticOpKind, + left: Box, + right: Box, + }, + Equal(Box, Box), + Less(Box, Box), + And(Box, Box), + Or(Box, Box), + Not(Box), + IsNull(Box), +} +impl Expression { + fn dtype(&self, input: &Schema) -> Result<(DataType, bool), Error> { + use Expression::*; + match self { + Column(i) => { + let (t, n) = plain(input, *i)?; + Ok((t.clone(), n)) + } + Literal { value, dtype } => { + if value.matches(dtype, true) { + Ok((dtype.clone(), matches!(value, Value::Null))) + } else { + Err(invalid("literal type mismatch")) + } + } + Negate(v) => { + let (t, n) = v.dtype(input)?; + if matches!(t, DataType::Int64 | DataType::Float64) { + Ok((t, n)) + } else { + Err(invalid("numeric negation required")) + } + } + Arithmetic { op, left, right } => { + let (a, n) = left.dtype(input)?; + let (b, m) = right.dtype(input)?; + if a == b + && matches!(a, DataType::Int64 | DataType::Float64) + && !(a == DataType::Int64 && *op == ArithmeticOpKind::Atan2) + { + Ok((a, n || m)) + } else { + Err(invalid("arithmetic requires matching numeric types")) + } + } + Equal(a, b) | Less(a, b) => { + let (a, n) = a.dtype(input)?; + let (b, m) = b.dtype(input)?; + if a == b && ordered(&a) { + Ok((DataType::Bool, n || m)) + } else { + Err(invalid("comparison requires matching ordered types")) + } + } + And(a, b) | Or(a, b) => { + let (a, n) = a.dtype(input)?; + let (b, m) = b.dtype(input)?; + if a == DataType::Bool && b == DataType::Bool { + Ok((DataType::Bool, n || m)) + } else { + Err(invalid("boolean operands required")) + } + } + Not(v) => { + let (t, n) = v.dtype(input)?; + if t == DataType::Bool { + Ok((t, n)) + } else { + Err(invalid("boolean operand required")) + } + } + IsNull(v) => { + v.dtype(input)?; + Ok((DataType::Bool, false)) + } + } + } + fn evaluate(&self, row: &[Value]) -> Result { + use Expression::*; + Ok(match self { + Column(i) => row[*i].clone(), + Literal { value, .. } => value.clone(), + Negate(v) => match v.evaluate(row)? { + Value::Int64(v) => Value::Int64( + v.checked_neg() + .ok_or_else(|| invalid("integer negation overflow"))?, + ), + Value::Float64(v) => Value::Float64(-v), + Value::Null => Value::Null, + _ => return Err(invalid("numeric negation required")), + }, + Arithmetic { op, left, right } => { + numeric(op, left.evaluate(row)?, right.evaluate(row)?)? + } + Equal(a, b) | Less(a, b) => { + let (a, b) = (a.evaluate(row)?, b.evaluate(row)?); + if matches!(a, Value::Null) || matches!(b, Value::Null) { + Value::Null + } else if matches!((&a,&b),(Value::Float64(a),Value::Float64(b)) if a.is_nan() || b.is_nan()) + { + Value::Bool(false) + } else { + let c = a.compare(&b)?; + Value::Bool(if matches!(self, Equal(..)) { + c.is_eq() + } else { + c.is_lt() + }) + } + } + And(a, b) | Or(a, b) => { + let (a, b) = (a.evaluate(row)?, b.evaluate(row)?); + match (a, b, matches!(self, And(..))) { + (Value::Bool(false), _, true) | (_, Value::Bool(false), true) => { + Value::Bool(false) + } + (Value::Bool(true), _, false) | (_, Value::Bool(true), false) => { + Value::Bool(true) + } + (Value::Null, _, _) | (_, Value::Null, _) => Value::Null, + (Value::Bool(a), Value::Bool(b), true) => Value::Bool(a && b), + (Value::Bool(a), Value::Bool(b), false) => Value::Bool(a || b), + _ => return Err(invalid("boolean operands required")), + } + } + Not(v) => match v.evaluate(row)? { + Value::Bool(v) => Value::Bool(!v), + Value::Null => Value::Null, + _ => return Err(invalid("boolean operand required")), + }, + IsNull(v) => Value::Bool(matches!(v.evaluate(row)?, Value::Null)), + }) + } +} +fn ordered(dtype: &DataType) -> bool { + matches!( + dtype, + DataType::Int64 + | DataType::Float64 + | DataType::Utf8 + | DataType::Bool + | DataType::Timestamp + | DataType::Date + ) +} +fn numeric(op: &ArithmeticOpKind, a: Value, b: Value) -> Result { + use ArithmeticOpKind::*; + Ok(match (a, b) { + (Value::Null, _) | (_, Value::Null) => Value::Null, + (Value::Float64(a), Value::Float64(b)) => { + Value::Float64(crate::arithmetic::evaluate_float64_arithmetic(op, a, b)) + } + (Value::Int64(a), Value::Int64(b)) => Value::Int64( + match op { + Add => a.checked_add(b), + Sub => a.checked_sub(b), + Mul => a.checked_mul(b), + Div => a.checked_div(b), + Mod => a.checked_rem(b), + Pow => u32::try_from(b).ok().and_then(|b| a.checked_pow(b)), + Atan2 => None, + } + .ok_or_else(|| invalid("invalid integer arithmetic or overflow"))?, + ), + _ => return Err(invalid("arithmetic type mismatch")), + }) +} +#[derive(Clone, Debug)] +pub struct SortKey { + pub column: usize, + pub descending: bool, + pub nulls_first: bool, +} +#[derive(Clone, Debug)] +pub enum Reduction { + Count, + Sum(usize), + Avg(usize), + Min(usize), + Max(usize), +} +#[derive(Clone)] +enum Kind { + Source(Vec), + Union, + VectorToScalar { + column: usize, + }, + Project(Vec), + Filter(Expression), + Limit { + n: u64, + offset: u64, + groups: Vec, + }, + Sort { + keys: Vec, + groups: Vec, + }, + Aggregate { + groups: Vec, + measures: Vec, + }, + SemiJoin { + keys: Vec<(usize, usize)>, + }, + SummaryBuild { + family: SummaryFamilyType, + value: usize, + time: Option, + groups: Vec, + }, + SummaryMerge { + state: usize, + groups: Vec, + }, + Readout { + state: usize, + statistic: crate::Statistic, + parameters: std::collections::HashMap, + }, +} +/// A bound operation has a fully checked input/output contract before execution. +#[derive(Clone)] +pub struct Operator { + kind: Kind, + inputs: Vec, + output: Schema, +} +impl Operator { + pub fn source(output: Schema, batches: Vec) -> Result { + super::values::validate_schema(&output)?; + if batches.iter().any(|b| b.schema() != &output) { + return Err(invalid("source schema mismatch")); + } + Ok(Self { + kind: Kind::Source(batches), + inputs: vec![], + output, + }) + } + /// Union polls every input fairly, including branches sharing a producer. + pub fn union(input: Schema, arity: usize) -> Result { + if arity == 0 { + return Err(invalid("union needs at least one input")); + } + Ok(Self { + kind: Kind::Union, + inputs: vec![input.clone(); arity], + output: input, + }) + } + pub fn scalar(value: Value, dtype: DataType) -> Result { + let schema = schema(vec![result_field( + "value", + dtype, + matches!(value, Value::Null), + )]); + Self::source( + schema.clone(), + vec![Batch::try_new(schema, vec![vec![value]])?], + ) + } + /// PromQL scalar conversion: zero or multiple elements produce NaN. + pub fn vector_to_scalar(input: Schema, column: usize) -> Result { + if plain(&input, column)? != (&DataType::Float64, false) { + return Err(invalid("scalar conversion requires non-null Float64")); + } + Ok(Self { + kind: Kind::VectorToScalar { column }, + inputs: vec![input], + output: schema(vec![result_field("value", DataType::Float64, false)]), + }) + } + pub fn project(input: Schema, columns: Vec<(String, Expression)>) -> Result { + let fields = columns + .iter() + .map(|(name, e)| { + let (t, n) = e.dtype(&input)?; + Ok(result_field(name, t, n)) + }) + .collect::>()?; + Ok(Self { + kind: Kind::Project(columns.into_iter().map(|(_, e)| e).collect()), + inputs: vec![input], + output: schema(fields), + }) + } + pub fn filter(input: Schema, predicate: Expression) -> Result { + if predicate.dtype(&input)?.0 != DataType::Bool { + return Err(invalid("filter predicate must be boolean")); + } + Ok(Self { + kind: Kind::Filter(predicate), + inputs: vec![input.clone()], + output: input, + }) + } + pub fn limit(input: Schema, n: u64, offset: u64, groups: Vec) -> Result { + validate_groups(&input, &groups)?; + Ok(Self { + kind: Kind::Limit { n, offset, groups }, + inputs: vec![input.clone()], + output: input, + }) + } + pub fn sort(input: Schema, keys: Vec, groups: Vec) -> Result { + validate_groups(&input, &groups)?; + for key in &keys { + if !ordered(plain(&input, key.column)?.0) { + return Err(invalid("unsupported sort type")); + } + } + Ok(Self { + kind: Kind::Sort { keys, groups }, + inputs: vec![input.clone()], + output: input, + }) + } + pub fn aggregate( + input: Schema, + groups: Vec, + measures: Vec<(String, Reduction)>, + ) -> Result { + validate_groups(&input, &groups)?; + let mut fields = groups + .iter() + .map(|&i| input.fields[i].clone()) + .collect::>(); + for (name, reduction) in &measures { + let (t, n) = match reduction { + Reduction::Count => (DataType::Int64, false), + Reduction::Sum(i) | Reduction::Avg(i) => { + let (t, _) = plain(&input, *i)?; + if !matches!(t, DataType::Int64 | DataType::Float64) { + return Err(invalid("numeric aggregate input required")); + } + ( + if matches!(reduction, Reduction::Avg(_)) { + DataType::Float64 + } else { + t.clone() + }, + false, + ) + } + Reduction::Min(i) | Reduction::Max(i) => { + let (t, _) = plain(&input, *i)?; + if !ordered(t) { + return Err(invalid("ordered aggregate input required")); + } + (t.clone(), true) + } + }; + fields.push(result_field(name, t, n)); + } + Ok(Self { + kind: Kind::Aggregate { + groups, + measures: measures.into_iter().map(|(_, r)| r).collect(), + }, + inputs: vec![input], + output: schema(fields), + }) + } + pub fn semi_join( + left: Schema, + right: Schema, + keys: Vec<(usize, usize)>, + ) -> Result { + if keys.is_empty() { + return Err(invalid("semi-join needs matching keys")); + } + for &(l, r) in &keys { + if plain(&left, l)?.0 != plain(&right, r)?.0 { + return Err(invalid("join key types differ")); + } + } + Ok(Self { + kind: Kind::SemiJoin { keys }, + inputs: vec![left.clone(), right], + output: left, + }) + } + pub fn summary_build( + input: Schema, + family: SummaryFamilyType, + value: usize, + time: Option, + groups: Vec, + ) -> Result { + super::values::validate_family(&family)?; + validate_groups(&input, &groups)?; + if plain(&input, value)? != (&DataType::Float64, false) { + return Err(invalid("summary numeric update requires non-null Float64")); + } + if let Some(time) = time { + if plain(&input, time)? != (&DataType::Timestamp, false) { + return Err(invalid("summary time column must be a timestamp")); + } + } + if time.is_none() + && matches!( + family, + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Rate + | planner_types::post_asap::ExactKind::Increase, + _ + ) + ) + { + return Err(invalid("counter summary requires a timestamp column")); + } + crate::capability::validate_summary_kernel( + &family, + &SummaryUpdate::column(ColumnRef::SampleValue), + &Default::default(), + ) + .map_err(Error::Invalid)?; + let mut fields = groups + .iter() + .map(|&i| input.fields[i].clone()) + .collect::>(); + fields.push(SummaryField { + name: "state".into(), + dtype: family.clone(), + nullable: false, + }); + Ok(Self { + kind: Kind::SummaryBuild { + family, + value, + time, + groups, + }, + inputs: vec![input], + output: schema(fields), + }) + } + pub fn summary_merge(input: Schema, state: usize, groups: Vec) -> Result { + validate_groups(&input, &groups)?; + super::values::validate_family(&field(&input, state)?.dtype)?; + if matches!(field(&input, state)?.dtype, SummaryFamilyType::Plain(_)) { + return Err(invalid("summary state required")); + } + let mut fields = groups + .iter() + .map(|&i| input.fields[i].clone()) + .collect::>(); + fields.push(input.fields[state].clone()); + Ok(Self { + kind: Kind::SummaryMerge { state, groups }, + inputs: vec![input], + output: schema(fields), + }) + } + pub fn readout( + input: Schema, + state: usize, + statistic: crate::Statistic, + parameters: std::collections::HashMap, + ) -> Result { + super::values::validate_family(&field(&input, state)?.dtype)?; + if matches!(field(&input, state)?.dtype, SummaryFamilyType::Plain(_)) { + return Err(invalid("summary state required")); + } + validate_readout(&field(&input, state)?.dtype, statistic, ¶meters)?; + let mut fields = input.fields.clone(); + let result_type = if matches!( + fields[state].dtype, + SummaryFamilyType::ExactAggregate(planner_types::post_asap::ExactKind::Count, _) + ) { + DataType::Int64 + } else { + DataType::Float64 + }; + fields[state] = result_field("value", result_type, false); + Ok(Self { + kind: Kind::Readout { + state, + statistic, + parameters, + }, + inputs: vec![input], + output: schema(fields), + }) + } + pub(crate) fn with_output_schema(mut self, output: Schema) -> Result { + if self.output.fields.len() != output.fields.len() + || self + .output + .fields + .iter() + .zip(&output.fields) + .any(|(actual, declared)| { + actual.dtype != declared.dtype || (actual.nullable && !declared.nullable) + }) + { + return Err(invalid("native output type differs from Planner output")); + } + if output.time_index.is_some_and(|i| { + i >= output.fields.len() + || output.fields[i].dtype != SummaryFamilyType::Plain(DataType::Timestamp) + }) { + return Err(invalid("invalid output time column")); + } + self.output = output; + Ok(self) + } + pub fn schema(&self) -> Schema { + self.output.clone() + } +} +fn validate_groups(input: &Schema, groups: &[usize]) -> Result<(), Error> { + for &i in groups { + plain(input, i)?; + } + if groups + .iter() + .collect::>() + .len() + != groups.len() + { + return Err(invalid("duplicate group columns")); + } + Ok(()) +} +async fn collect_rows( + mut input: Input<'_, Batch>, + context: &RunContext, +) -> Result<(Vec>, Vec), Error> { + let mut rows = Vec::new(); + let mut reservations = Vec::new(); + while let Some(batch) = input.next().await { + let batch = batch?; + reservations.push(context.reserve(batch.bytes())?); + rows.extend(batch.rows().iter().cloned()); + } + Ok((rows, reservations)) +} +impl PhysicalOperator for Operator { + fn name(&self) -> &str { + match self.kind { + Kind::Source(_) => "Source", + Kind::Union => "Union", + Kind::VectorToScalar { .. } => "VectorToScalar", + Kind::Project(_) => "Project", + Kind::Filter(_) => "Filter", + Kind::Limit { .. } => "Limit", + Kind::Sort { .. } => "Sort", + Kind::Aggregate { .. } => "Aggregate", + Kind::SemiJoin { .. } => "SemiJoin", + Kind::SummaryBuild { .. } => "SummaryAgg", + Kind::SummaryMerge { .. } => "SummaryMerge", + Kind::Readout { .. } => "SummaryReadout", + } + } + fn input_schemas(&self) -> Vec { + self.inputs.clone() + } + fn output_schema(&self) -> Schema { + self.output.clone() + } + fn output_bytes(&self, value: &Batch) -> usize { + value.bytes() + } + fn start<'a>( + &'a self, + mut inputs: Vec>, + context: RunContext, + ) -> Result, Error> { + let output = self.output.clone(); + if let Kind::Source(batches) = &self.kind { + return Ok(futures::stream::iter(batches.iter().cloned().map(Ok)).boxed_local()); + } + if matches!(self.kind, Kind::Union) { + return Ok(futures::stream::select_all(inputs) + .map(|batch| batch.map(|batch| batch.value().clone())) + .boxed_local()); + } + if let Kind::SemiJoin { keys } = &self.kind { + let right = inputs.pop().ok_or_else(|| invalid("right input missing"))?; + let left = inputs.pop().ok_or_else(|| invalid("left input missing"))?; + return Ok(futures::stream::once(async move { + // Poll both branches together: either may depend on a common producer. + let ((left, _left_memory), (right, _right_memory)) = futures::try_join!( + collect_rows(left, &context), + collect_rows(right, &context) + )?; + let right_cols = keys.iter().map(|(_, r)| *r).collect::>(); + let left_cols = keys.iter().map(|(l, _)| *l).collect::>(); + let members = right + .iter() + .filter(|row| right_cols.iter().all(|&i| !matches!(row[i], Value::Null))) + .map(|r| group_key(r, &right_cols)) + .collect::, _>>()?; + let rows = left + .into_iter() + .filter_map(|r| match group_key(&r, &left_cols) { + Ok(k) + if left_cols.iter().all(|&i| !matches!(r[i], Value::Null)) + && members.contains(&k) => + { + Some(Ok(r)) + } + Ok(_) => None, + Err(e) => Some(Err(e)), + }) + .collect::, _>>()?; + Batch::try_new(output, rows) + }) + .boxed_local()); + } + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &self.kind { + Kind::VectorToScalar { column } => Ok(futures::stream::once(async move { + let mut input = input; + let mut value = f64::NAN; + let mut count = 0usize; + while let Some(batch) = input.next().await { + for row in batch?.rows() { + count = count.saturating_add(1); + if let Value::Float64(v) = row[*column] { + value = v; + } + } + } + Batch::try_new( + output, + vec![vec![Value::Float64(if count == 1 { + value + } else { + f64::NAN + })]], + ) + }) + .boxed_local()), + Kind::Project(expressions) => Ok(input + .map(move |batch| { + let batch = batch?; + let rows = batch + .rows() + .iter() + .map(|r| { + expressions + .iter() + .map(|e| e.evaluate(r)) + .collect::, _>>() + }) + .collect::, _>>()?; + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + Kind::Filter(predicate) => Ok(input + .map(move |batch| { + let batch = batch?; + let mut rows = Vec::new(); + for row in batch.rows() { + if matches!(predicate.evaluate(row)?, Value::Bool(true)) { + rows.push(row.clone()); + } + } + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + Kind::Limit { n, offset, groups } => { + let counts = BTreeMap::>, u64>::new(); + Ok(futures::stream::try_unfold( + (input, counts, Vec::::new(), false), + move |(mut input, mut counts, mut memory, done)| { + let output = output.clone(); + let context = context.clone(); + async move { + if done || *n == 0 { + return Ok(None); + } + let Some(batch) = input.next().await else { + return Ok(None); + }; + let batch = batch?; + let mut rows = Vec::new(); + for row in batch.rows() { + let key = group_key(row, groups)?; + if !counts.contains_key(&key) { + memory.push( + context.reserve( + key.iter() + .map(|part| { + part.len() + std::mem::size_of::>() + }) + .sum::() + + 64, + )?, + ); + } + let count = counts.entry(key).or_default(); + if *count >= *offset && count.saturating_sub(*offset) < *n { + rows.push(row.clone()); + } + *count = count.saturating_add(1); + } + let done = groups.is_empty() + && counts + .get(&vec![]) + .is_some_and(|count| count.saturating_sub(*offset) >= *n); + Ok(Some(( + Batch::try_new(output, rows)?, + (input, counts, memory, done), + ))) + } + }, + ) + .boxed_local()) + } + Kind::SummaryBuild { + family, + value, + time, + groups, + } => Ok(futures::stream::once(async move { + Batch::try_new( + output, + build_summary(input, family, *value, *time, groups, &context).await?, + ) + }) + .boxed_local()), + Kind::Readout { + state, + statistic, + parameters, + } => Ok(input + .map(move |batch| { + let batch = batch?; + let mut rows = batch.rows().to_vec(); + for row in &mut rows { + let Value::Summary { state: summary, .. } = &row[*state] else { + return Err(invalid("summary value required")); + }; + row[*state] = if output.fields[*state].dtype + == SummaryFamilyType::Plain(DataType::Int64) + { + let count = summary.aux_stats().count.ok_or_else(|| { + Error::Operator("exact count state lacks an integer count".into()) + })?; + Value::Int64( + i64::try_from(count).map_err(|_| { + Error::Operator("exact count exceeds Int64".into()) + })?, + ) + } else { + Value::Float64( + summary + .query_statistic(*statistic, &None, parameters) + .map_err(|e| Error::Operator(e.to_string()))?, + ) + }; + } + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + _ => Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let result = match &self.kind { + Kind::Sort { keys, groups } => { + let mut grouped = BTreeMap::>, Vec>>::new(); + for row in rows { + grouped + .entry(group_key(&row, groups)?) + .or_default() + .push(row); + } + let mut result = Vec::new(); + for mut rows in grouped.into_values() { + rows.sort_by(|a, b| compare_rows(a, b, keys)); + result.extend(rows); + } + result + } + Kind::Aggregate { groups, measures } => { + reduce(rows, groups, measures, &self.inputs[0])? + } + Kind::SummaryMerge { state, groups } => merge_summary(rows, *state, groups)?, + _ => return Err(invalid("unexpected blocking operation")), + }; + Batch::try_new(output, result) + }) + .boxed_local()), + } + } +} +fn compare_rows(a: &[Value], b: &[Value], keys: &[SortKey]) -> std::cmp::Ordering { + use std::cmp::Ordering::*; + for key in keys { + let (a, b) = (&a[key.column], &b[key.column]); + let order = match (a, b) { + (Value::Null, Value::Null) => Equal, + (Value::Null, _) => { + if key.nulls_first { + Less + } else { + Greater + } + } + (_, Value::Null) => { + if key.nulls_first { + Greater + } else { + Less + } + } + (Value::Float64(a), Value::Float64(b)) if a.is_nan() || b.is_nan() => { + match (a.is_nan(), b.is_nan()) { + (true, true) => Equal, + (true, false) => Greater, + _ => Less, + } + } + _ => { + let order = a.compare(b).expect("bound ordered types"); + if key.descending { + order.reverse() + } else { + order + } + } + }; + if order != Equal { + return order; + } + } + Equal +} +fn reduce( + rows: Vec>, + groups: &[usize], + measures: &[Reduction], + input: &Schema, +) -> Result>, Error> { + let mut grouped = BTreeMap::>, Vec>>::new(); + if rows.is_empty() && groups.is_empty() { + grouped.insert(vec![], vec![]); + } + for row in rows { + grouped + .entry(group_key(&row, groups)?) + .or_default() + .push(row); + } + grouped + .into_values() + .map(|rows| { + let mut result = groups + .iter() + .map(|&i| rows[0][i].clone()) + .collect::>(); + for measure in measures { + result.push(reduce_one(&rows, measure, input)?); + } + Ok(result) + }) + .collect() +} +fn reduce_one(rows: &[Vec], measure: &Reduction, input: &Schema) -> Result { + let column = match measure { + Reduction::Count => { + return Ok(Value::Int64( + i64::try_from(rows.len()).map_err(|_| invalid("count overflow"))?, + )) + } + Reduction::Sum(i) | Reduction::Avg(i) | Reduction::Min(i) | Reduction::Max(i) => *i, + }; + let values = rows + .iter() + .map(|r| &r[column]) + .filter(|v| !matches!(v, Value::Null)) + .collect::>(); + if matches!(measure, Reduction::Min(_) | Reduction::Max(_)) { + if plain(input, column)?.0 == &DataType::Float64 { + // Match exact-state kernels: ignore NaN when a numeric value exists. + let mut best: Option = None; + for value in values { + let Value::Float64(value) = value else { + return Err(invalid("floating aggregate value required")); + }; + best = Some(best.map_or(*value, |old| { + if matches!(measure, Reduction::Min(_)) { + old.min(*value) + } else { + old.max(*value) + } + })); + } + return Ok(best.map(Value::Float64).unwrap_or(Value::Null)); + } + let mut best: Option<&Value> = None; + for value in values { + if best + .map(|b| value.compare(b)) + .transpose()? + .is_none_or(|order| { + if matches!(measure, Reduction::Min(_)) { + order.is_lt() + } else { + order.is_gt() + } + }) + { + best = Some(value); + } + } + return Ok(best.cloned().unwrap_or(Value::Null)); + } + let count = values.len(); + let dtype = plain(input, column)?.0; + if dtype == &DataType::Int64 { + let sum = values.into_iter().try_fold(0i128, |sum, v| { + let Value::Int64(v) = v else { + return Err(invalid("integer aggregate value required")); + }; + sum.checked_add(i128::from(*v)) + .ok_or_else(|| invalid("integer aggregate overflow")) + })?; + return if matches!(measure, Reduction::Avg(_)) { + Ok(Value::Float64(sum as f64 / count as f64)) + } else { + Ok(Value::Int64( + i64::try_from(sum).map_err(|_| invalid("integer sum overflow"))?, + )) + }; + } + let sum = values + .into_iter() + .map(|v| { + if let Value::Float64(v) = v { + *v + } else { + unreachable!() + } + }) + .sum::(); + Ok(Value::Float64(if matches!(measure, Reduction::Avg(_)) { + sum / count as f64 + } else { + sum + })) +} +async fn build_summary( + mut input: Input<'_, Batch>, + family: &SummaryFamilyType, + value: usize, + time: Option, + groups: &[usize], + context: &RunContext, +) -> Result>, Error> { + type State = ( + Vec, + Box, + Reservation, + usize, + Option, + ); + let create = |labels: Vec, key_bytes: usize| -> Result { + let updater = crate::factory::create_planner_accumulator( + family, + &SummaryUpdate::column(ColumnRef::SampleValue), + &Default::default(), + ) + .map_err(Error::Operator)?; + let overhead = labels.iter().map(Value::bytes).sum::() + key_bytes + 64; + let memory = context.reserve(updater.memory_usage_bytes() + overhead)?; + Ok((labels, updater, memory, overhead, None)) + }; + let mut states = BTreeMap::>, State>::new(); + if groups.is_empty() { + states.insert(vec![], create(vec![], 0)?); + } + let ordered_time = matches!( + family, + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Rate + | planner_types::post_asap::ExactKind::Increase, + _ + ) + ); + while let Some(batch) = input.next().await { + let batch = batch?; + for row in batch.rows() { + let key = group_key(row, groups)?; + if !states.contains_key(&key) { + let labels = groups.iter().map(|&i| row[i].clone()).collect(); + let state = create( + labels, + key.iter() + .map(|v| v.len() + std::mem::size_of::>()) + .sum(), + )?; + states.insert(key.clone(), state); + } + let (_, updater, memory, overhead, previous) = + states.get_mut(&key).expect("inserted group"); + let Value::Float64(value) = row[value] else { + return Err(invalid("summary update type")); + }; + let timestamp = if let Some(time) = time { + let Value::Timestamp(time) = row[time] else { + return Err(invalid("summary time type")); + }; + time + } else { + 0 + }; + if ordered_time && previous.is_some_and(|prior| timestamp <= prior) { + return Err(Error::Operator( + "counter samples must have strictly increasing timestamps within each group" + .into(), + )); + } + updater + .validate_single_input(value) + .map_err(Error::Operator)?; + updater.update_single(value, timestamp); + *previous = Some(timestamp); + memory.resize(updater.memory_usage_bytes() + *overhead)?; + } + } + Ok(states + .into_values() + .map(|(mut labels, updater, _memory, _, _)| { + labels.push(Value::Summary { + family: family.clone(), + state: Arc::from(updater.into_accumulator()), + }); + labels + }) + .collect()) +} + +fn merge_summary( + rows: Vec>, + state_column: usize, + groups: &[usize], +) -> Result>, Error> { + type GroupState = (Vec, SummaryFamilyType, Arc); + let mut states: BTreeMap>, GroupState> = BTreeMap::new(); + for row in rows { + let Value::Summary { family, state } = &row[state_column] else { + return Err(invalid("summary state required")); + }; + let key = group_key(&row, groups)?; + if let Some((_, expected, existing)) = states.get_mut(&key) { + if expected != family { + return Err(invalid("incompatible summary family")); + } + *existing = Arc::from( + existing + .merge_with(state.as_ref()) + .map_err(|e| Error::Operator(e.to_string()))?, + ); + } else { + states.insert( + key, + ( + groups.iter().map(|&i| row[i].clone()).collect(), + family.clone(), + state.clone(), + ), + ); + } + } + Ok(states + .into_values() + .map(|(mut keys, family, state)| { + keys.push(Value::Summary { family, state }); + keys + }) + .collect()) +} + +fn validate_readout( + family: &SummaryFamilyType, + statistic: crate::Statistic, + parameters: &std::collections::HashMap, +) -> Result<(), Error> { + use crate::Statistic as S; + use planner_types::post_asap::{ExactKind as E, SketchAlgorithm as A}; + let supported = match family { + SummaryFamilyType::ExactAggregate(kind, _) => matches!( + (kind, statistic), + (E::Sum, S::Sum) + | (E::Count, S::Count) + | (E::Min, S::Min) + | (E::Max, S::Max) + | (E::Rate, S::Rate) + | (E::Increase, S::Increase) + ), + SummaryFamilyType::Sketch(kind, _) => match kind.algorithm() { + A::Kll => statistic == S::Quantile, + A::DDSketch => matches!(statistic, S::Quantile | S::Count), + A::Hll => matches!(statistic, S::Cardinality | S::Count), + _ => false, + }, + _ => false, + }; + if !supported { + return Err(invalid( + "readout is not implemented for this summary family", + )); + } + if statistic == S::Quantile + && !parameters + .get("quantile") + .and_then(|s| s.parse::().ok()) + .is_some_and(|q| (0.0..=1.0).contains(&q)) + { + return Err(invalid("quantile readout requires quantile in [0,1]")); + } + Ok(()) +} diff --git a/crates/asap-physical-operators/src/dag/planner.rs b/crates/asap-physical-operators/src/dag/planner.rs new file mode 100644 index 000000000..1952f1790 --- /dev/null +++ b/crates/asap-physical-operators/src/dag/planner.rs @@ -0,0 +1,479 @@ +//! Bind a post-ASAP DAG to native operators. Sources are explicit execution +//! frontiers supplied by the deployment; unsupported computation is an error. +use super::{ + operators::{Expression, Operator, Reduction, SortKey}, + values::{Batch, Schema, Value}, + Error, NodeId, PhysicalDag, PhysicalOperator, +}; +use planner_types::{ + post_asap::{ + ExactOperation, ExecutableDag, ExecutableDagNode, ExecutableOperatorPayload as Payload, + SketchQuery, SummaryFamilyType, SummaryInputExpr, ValueOperation, + }, + pre_asap::{ + AggIntent, ColumnRef, CompareOpKind, DataType, GroupKeys, QueryExpr, + Reduction as PlannerReduction, ScalarValue, + }, +}; +use std::{ + collections::{BTreeMap, BTreeSet}, + sync::Arc, +}; +fn invalid(message: impl Into) -> Error { + Error::Invalid(message.into()) +} + +/// Source nodes cut the DAG at an installed storage/ingestion frontier. The +/// binding must have exactly the declared schema and no upstream dependencies. +/// A deployment must authorize these frontiers before calling this function. +pub type Source<'a> = Box + 'a>; + +pub fn bind<'a>( + dag: &ExecutableDag, + mut sources: BTreeMap>, + roots: &[NodeId], +) -> Result, Error> { + preflight_depth(dag)?; + dag.validate().map_err(|e| invalid(e.to_string()))?; + let nodes = dag + .nodes + .iter() + .map(|node| (u64::from(node.id.0), node)) + .collect::>(); + let mut dependencies = BTreeMap::>::new(); + for edge in &dag.edges { + dependencies + .entry(u64::from(edge.consumer.0)) + .or_default() + .push(u64::from(edge.producer.0)); + } + if sources.keys().any(|id| !nodes.contains_key(id)) { + return Err(invalid("source binding names an unknown node")); + } + let mut ordered = Vec::new(); + let mut seen = BTreeSet::new(); + let mut pending = roots.iter().map(|&id| (id, false)).collect::>(); + while let Some((id, expanded)) = pending.pop() { + if expanded { + ordered.push(id); + continue; + } + if !seen.insert(id) { + continue; + } + if !nodes.contains_key(&id) { + return Err(invalid(format!("missing root {id}"))); + } + pending.push((id, true)); + if !sources.contains_key(&id) { + for &input in dependencies.get(&id).into_iter().flatten() { + pending.push((input, false)); + } + } + } + let mut graph = PhysicalDag::default(); + let mut auxiliary = u64::MAX; + for id in ordered { + let node = nodes[&id]; + let output = Arc::new(node.output_schema.clone()); + super::values::validate_schema(&output)?; + let (operator, inputs) = if let Some(source) = sources.remove(&id) { + if !source.input_schemas().is_empty() || source.output_schema() != output { + return Err(invalid("frontier is not a source with the declared schema")); + } + ( + Box::new(CheckedSource { source, output }) as Source<'a>, + vec![], + ) + } else { + let mut inputs = dependencies.get(&id).cloned().unwrap_or_default(); + let mut schemas = inputs + .iter() + .map(|id| Arc::new(nodes[id].output_schema.clone())) + .collect::>(); + if matches!(node.payload, Payload::SummaryMerge { .. }) && inputs.len() > 1 { + if schemas.iter().any(|s| s != &schemas[0]) { + return Err(invalid("summary merge inputs have different schemas")); + } + graph.add( + auxiliary, + inputs, + Operator::union(schemas[0].clone(), schemas.len())?, + )?; + inputs = vec![auxiliary]; + auxiliary -= 1; + schemas.truncate(1); + } + let operator = bind_operation(node, &schemas) + .map_err(|error| invalid(format!("node {id}: {error}")))? + .with_output_schema(output)?; + (Box::new(operator) as Source<'a>, inputs) + }; + graph.add_boxed(id, inputs, operator)?; + } + graph.validate(roots)?; + Ok(graph) +} + +fn bind_operation(node: &ExecutableDagNode, inputs: &[Schema]) -> Result { + let [input] = inputs else { + return Err(invalid( + "native Planner binding currently requires a unary operation or an explicit source", + )); + }; + match &node.payload { + Payload::Value { operation, .. } => match operation { + ValueOperation::Project { cols, .. } => Operator::project( + input.clone(), + cols.iter() + .enumerate() + .map(|(i, col)| { + Ok(( + node.output_schema + .fields + .get(i) + .ok_or_else(|| invalid("projection width mismatch"))? + .name + .clone(), + expression(&col.expr)?, + )) + }) + .collect::>()?, + ), + ValueOperation::Filter { pred } => { + Operator::filter(input.clone(), expression(&pred.0)?) + } + ValueOperation::Sort { keys, partition_by } => Operator::sort( + input.clone(), + keys.iter() + .map(|key| { + let QueryExpr::Column(column) = key.expr else { + return Err(invalid( + "sort expression must be projected before sorting", + )); + }; + Ok(SortKey { + column, + descending: !key.ascending, + nulls_first: key.nulls_first, + }) + }) + .collect::>()?, + groups(input, partition_by)?, + ), + ValueOperation::Limit { n, offset } => { + Operator::limit(input.clone(), *n as u64, *offset as u64, vec![]) + } + ValueOperation::Exact(ExactOperation::Aggregate { + reduction, + measures, + output_names, + having: None, + }) => { + if measures.len() != output_names.len() { + return Err(invalid("aggregate output names differ from measures")); + } + let PlannerReduction::Reduce(keys) = reduction else { + return Err(invalid( + "per-entity aggregate requires an explicit entity binding", + )); + }; + let measures = measures + .iter() + .zip(output_names) + .map(|(m, name)| { + let column = |col: Option| { + col.map(Ok) + .unwrap_or_else(|| named_column(input, &ColumnRef::SampleValue)) + }; + let m = match m { + AggIntent::Count { .. } => Reduction::Count, + AggIntent::Sum { col } => Reduction::Sum(column(*col)?), + AggIntent::Avg { col } => Reduction::Avg(column(*col)?), + AggIntent::Min { col } => Reduction::Min(column(*col)?), + AggIntent::Max { col } => Reduction::Max(column(*col)?), + _ => { + return Err(invalid( + "aggregate intent has no native implementation", + )) + } + }; + Ok((name.clone(), m)) + }) + .collect::>()?; + Operator::aggregate(input.clone(), groups(input, keys)?, measures) + } + ValueOperation::FinalizeExactAccumulator => { + let state = summary_column(input)?; + use crate::Statistic as S; + use planner_types::post_asap::ExactKind as E; + let statistic = match &input.fields[state].dtype { + SummaryFamilyType::ExactAggregate(kind, _) => match kind { + E::Sum => S::Sum, + E::Count => S::Count, + E::Min => S::Min, + E::Max => S::Max, + E::Rate => S::Rate, + E::Increase => S::Increase, + _ => return Err(invalid("exact family readout is unsupported")), + }, + _ => return Err(invalid("exact finalization requires exact state")), + }; + Operator::readout(input.clone(), state, statistic, Default::default()) + } + _ => Err(invalid("value operation has no native implementation")), + }, + Payload::SummaryAgg { + family, + input: update, + reduction, + grouping, + } => { + if update.item.is_some() { + return Err(invalid("keyed summary update binding is not implemented")); + } + crate::capability::validate_summary_kernel(family, update, grouping) + .map_err(Error::Invalid)?; + let SummaryInputExpr::Column(column) = &update.weight else { + return Err(invalid( + "summary update expression must be projected to a column", + )); + }; + let PlannerReduction::Reduce(keys) = reduction else { + return Err(invalid( + "summary construction requires explicit grouping columns", + )); + }; + Operator::summary_build( + input.clone(), + family.clone(), + named_column(input, column)?, + input.time_index, + groups(input, keys)?, + ) + } + Payload::SummaryMerge { .. } => { + let state = summary_column(input)?; + Operator::summary_merge( + input.clone(), + state, + (0..input.fields.len()) + .filter(|&i| i != state && Some(i) != input.time_index) + .collect(), + ) + } + Payload::SummaryEstimate { query } => { + let mut params = std::collections::HashMap::new(); + let statistic = match query { + SketchQuery::Quantile { q } => { + params.insert("quantile".into(), q.to_string()); + crate::Statistic::Quantile + } + SketchQuery::Cardinality => crate::Statistic::Cardinality, + SketchQuery::PointCount { value: None, .. } => crate::Statistic::Count, + _ => return Err(invalid("summary readout is not implemented")), + }; + Operator::readout(input.clone(), summary_column(input)?, statistic, params) + } + _ => Err(invalid( + "physical operation has no native binding; no fallback is installed", + )), + } +} +fn summary_column(input: &Schema) -> Result { + let columns = input + .fields + .iter() + .enumerate() + .filter(|(_, f)| !matches!(f.dtype, SummaryFamilyType::Plain(_))) + .map(|(i, _)| i) + .collect::>(); + match columns.as_slice() { + [column] => Ok(*column), + _ => Err(invalid("one summary state column required")), + } +} +fn named_column(input: &Schema, column: &ColumnRef) -> Result { + let name = match column { + ColumnRef::Named(name) => name.as_str(), + ColumnRef::SampleValue => "value", + _ => { + return Err(invalid( + "summary update requires an unambiguous bound column", + )) + } + }; + let matches = input + .fields + .iter() + .enumerate() + .filter(|(_, field)| field.name == name) + .map(|(i, _)| i) + .collect::>(); + match matches.as_slice() { + [column] => Ok(*column), + _ => Err(invalid("summary update column missing or ambiguous")), + } +} +fn groups(input: &Schema, groups: &GroupKeys) -> Result, Error> { + if groups.is_without() { + return Err(invalid("grouping without requires resolved label columns")); + } + if groups.keys().iter().any(|&i| i >= input.fields.len()) { + return Err(invalid("grouping column out of range")); + } + Ok(groups.keys().to_vec()) +} +fn expression(expr: &QueryExpr) -> Result { + let bind = |e: &QueryExpr| expression(e).map(Box::new); + Ok(match expr { + QueryExpr::Column(i) => Expression::Column(*i), + QueryExpr::Literal(value) => { + let (value, dtype) = match value { + ScalarValue::Int64(v) => (Value::Int64(*v), DataType::Int64), + ScalarValue::Float64(v) => (Value::Float64(*v), DataType::Float64), + ScalarValue::Utf8(v) => (Value::Utf8(v.as_str().into()), DataType::Utf8), + ScalarValue::Boolean(v) => (Value::Bool(*v), DataType::Bool), + ScalarValue::Null => (Value::Null, DataType::Null), + ScalarValue::Interval { + months, + days, + nanos, + } => ( + Value::Interval { + months: *months, + days: *days, + nanos: *nanos, + }, + DataType::Interval, + ), + }; + Expression::Literal { value, dtype } + } + QueryExpr::Arithmetic { op, left, right } => Expression::Arithmetic { + op: op.clone(), + left: bind(left)?, + right: bind(right)?, + }, + QueryExpr::Compare { + left, + op: CompareOpKind::Eq, + right, + } => Expression::Equal(bind(left)?, bind(right)?), + QueryExpr::Compare { + left, + op: CompareOpKind::Lt, + right, + } => Expression::Less(bind(left)?, bind(right)?), + QueryExpr::Not(v) => Expression::Not(bind(v)?), + QueryExpr::IsNull(v) => Expression::IsNull(bind(v)?), + QueryExpr::IsNotNull(v) => Expression::Not(Box::new(Expression::IsNull(bind(v)?))), + QueryExpr::BoolAnd(items) | QueryExpr::BoolOr(items) => { + let and = matches!(expr, QueryExpr::BoolAnd(_)); + let mut result = Expression::Literal { + value: Value::Bool(and), + dtype: DataType::Bool, + }; + for item in items { + result = if and { + Expression::And(Box::new(result), bind(item)?) + } else { + Expression::Or(Box::new(result), bind(item)?) + }; + } + result + } + _ => return Err(invalid("expression has no native implementation")), + }) +} + +// Source adapters may perform I/O, but their actual batches must honor the +// schema accepted by the binder before a downstream expression sees a row. +struct CheckedSource<'a> { + source: Source<'a>, + output: Schema, +} +impl PhysicalOperator for CheckedSource<'_> { + fn name(&self) -> &str { + self.source.name() + } + fn input_schemas(&self) -> Vec { + vec![] + } + fn output_schema(&self) -> Schema { + self.output.clone() + } + fn output_bytes(&self, batch: &Batch) -> usize { + self.source.output_bytes(batch) + } + fn start<'a>( + &'a self, + inputs: Vec>, + context: super::RunContext, + ) -> Result, Error> { + use futures::StreamExt; + Ok(self + .source + .start(inputs, context)? + .map(|batch| { + let batch = batch?; + if batch.schema() != &self.output { + return Err(invalid("source batch differs from its bound schema")); + } + Ok(batch) + }) + .boxed_local()) + } +} + +// Bound recursion before invoking the upstream recursive provenance validator. +fn preflight_depth(dag: &ExecutableDag) -> Result<(), Error> { + let mut remaining = dag + .nodes + .iter() + .map(|node| (node.id, 0usize)) + .collect::>(); + if remaining.len() != dag.nodes.len() { + return Err(invalid("duplicate Planner node")); + } + let mut consumers = BTreeMap::<_, Vec<_>>::new(); + for edge in &dag.edges { + if !remaining.contains_key(&edge.producer) { + return Err(invalid("missing Planner edge producer")); + } + *remaining + .get_mut(&edge.consumer) + .ok_or_else(|| invalid("missing Planner edge consumer"))? += 1; + consumers + .entry(edge.producer) + .or_default() + .push(edge.consumer); + } + let mut ready = remaining + .iter() + .filter(|(_, n)| **n == 0) + .map(|(id, _)| *id) + .collect::>(); + let mut depths = BTreeMap::new(); + let mut visited = 0; + while let Some(id) = ready.pop_front() { + visited += 1; + let depth = *depths.get(&id).unwrap_or(&1usize); + if depth > 128 { + return Err(invalid("DAG exceeds the supported execution depth of 128")); + } + for &consumer in consumers.get(&id).into_iter().flatten() { + let next = depths.entry(consumer).or_insert(1); + *next = (*next).max(depth + 1); + let count = remaining.get_mut(&consumer).expect("validated endpoint"); + *count -= 1; + if *count == 0 { + ready.push_back(consumer); + } + } + } + if visited != dag.nodes.len() { + return Err(invalid("Planner DAG contains a cycle")); + } + Ok(()) +} diff --git a/crates/asap-physical-operators/src/dag/tests.rs b/crates/asap-physical-operators/src/dag/tests.rs new file mode 100644 index 000000000..e051fa313 --- /dev/null +++ b/crates/asap-physical-operators/src/dag/tests.rs @@ -0,0 +1,260 @@ +use super::*; +use futures::{executor::block_on, stream, StreamExt}; + +struct Source { + starts: Rc>, + polls: Rc>, + fail: bool, + end: u64, +} +impl PhysicalOperator for Source { + fn name(&self) -> &str { + "CountingSource" + } + fn input_schemas(&self) -> Vec<()> { + vec![] + } + fn output_schema(&self) {} + fn output_bytes(&self, _: &u64) -> usize { + 8 + } + fn start<'a>( + &'a self, + _: Vec>, + _: RunContext, + ) -> Result, Error> { + self.starts.set(self.starts.get() + 1); + Ok(stream::iter(0..self.end) + .map(move |n| { + self.polls.set(self.polls.get() + 1); + if self.fail && n == 1 { + Err(Error::Operator("source failure".into())) + } else { + Ok(n) + } + }) + .boxed_local()) + } +} +struct Identity; +impl PhysicalOperator for Identity { + fn name(&self) -> &str { + "Identity" + } + fn input_schemas(&self) -> Vec<()> { + vec![()] + } + fn output_schema(&self) {} + fn output_bytes(&self, _: &u64) -> usize { + 8 + } + fn start<'a>( + &'a self, + mut inputs: Vec>, + _: RunContext, + ) -> Result, Error> { + Ok(inputs + .remove(0) + .map(|value| value.map(|v| *v)) + .boxed_local()) + } +} +fn context() -> RunContext { + RunContext::new( + Scope::Query { + evaluation_time_ms: 100, + revision: 1, + }, + Limits { + max_buffered_batches: 1, + max_bytes: 1024, + }, + ) + .unwrap() +} +fn source(fail: bool) -> (Source, Rc>, Rc>) { + let starts = Rc::new(Cell::new(0)); + let polls = Rc::new(Cell::new(0)); + ( + Source { + starts: starts.clone(), + polls: polls.clone(), + fail, + end: 4, + }, + starts, + polls, + ) +} + +// A shared producer runs once, and the slow reader bounds producer progress. +#[test] +fn shared_source_backpressure_and_reader_drop() { + let (source, starts, polls) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + let context = context(); + let mut readers = dag.execute(&[0, 0], context.clone()).unwrap(); + let mut slow = readers.pop().unwrap(); + let mut fast = readers.pop().unwrap(); + assert_eq!(starts.get(), 1); + let first = block_on(fast.next()).unwrap().unwrap(); + assert_eq!(*first, 0); + let mut cx = Context::from_waker(futures::task::noop_waker_ref()); + assert!(Pin::new(&mut fast).poll_next(&mut cx).is_pending()); + assert_eq!(polls.get(), 1); + let same = block_on(slow.next()).unwrap().unwrap(); + assert!(Arc::ptr_eq(&first.value, &same.value)); + drop(same); + drop(first); + assert_eq!(context.retained_bytes(), 0); + assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 1); + drop(slow); + assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 2); + assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 3); + assert!(block_on(fast.next()).is_none()); + assert_eq!(polls.get(), 4); + drop(fast); + assert_eq!(context.retained_bytes(), 0); +} + +// Independent branches consume a common node concurrently without duplicate work. +#[test] +fn diamond_and_run_isolation() { + let (source, starts, polls) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + dag.add(1, vec![0], Identity).unwrap(); + dag.add(2, vec![0], Identity).unwrap(); + for _ in 0..2 { + let mut outputs = dag.execute(&[1, 2], context()).unwrap(); + let a = outputs.pop().unwrap(); + let b = outputs.pop().unwrap(); + let (a, b) = + block_on(async { futures::join!(a.collect::>(), b.collect::>()) }); + assert_eq!( + a.iter().map(|v| **v.as_ref().unwrap()).collect::>(), + vec![0, 1, 2, 3] + ); + assert_eq!( + b.iter().map(|v| **v.as_ref().unwrap()).collect::>(), + vec![0, 1, 2, 3] + ); + } + assert_eq!(starts.get(), 2); + assert_eq!(polls.get(), 8); +} + +// Failure reaches every subscriber; cancellation stops further producer work. +#[test] +fn broadcast_error_and_cancel() { + let (source, _, polls) = source(true); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + let mut outputs = dag.execute(&[0, 0], context()).unwrap(); + let a = outputs.pop().unwrap(); + let b = outputs.pop().unwrap(); + let (a, b) = block_on(async { futures::join!(a.collect::>(), b.collect::>()) }); + for values in [a, b] { + assert_eq!(values.len(), 2); + assert!(matches!(values[1], Err(Error::AtNode { node: 0, .. }))); + } + assert_eq!(polls.get(), 2); + let run = context(); + let mut output = dag.execute(&[0], run.clone()).unwrap().remove(0); + run.cancel(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::Cancelled)) + )); + assert!(block_on(output.next()).is_none()); + assert_eq!(polls.get(), 2); +} + +// Retaining a consumer output retains its budget lease after queue eviction. +#[test] +fn retained_outputs_count_against_budget() { + let (source, _, _) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + let run = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits { + max_buffered_batches: 1, + max_bytes: 8, + }, + ) + .unwrap(); + let mut input = dag.execute(&[0], run.clone()).unwrap().remove(0); + let held = block_on(input.next()).unwrap().unwrap(); + assert_eq!(run.retained_bytes(), 8); + assert!(matches!( + block_on(input.next()), + Some(Err(Error::MemoryLimit)) + )); + drop(input); + assert_eq!(run.retained_bytes(), 8); + drop(held); + assert_eq!(run.retained_bytes(), 0); +} + +// Invalid graphs fail before even starting a source. +#[test] +fn invalid_graphs_do_not_start_sources() { + let (source, starts, _) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + dag.add(1, vec![2], Identity).unwrap(); + dag.add(2, vec![1], Identity).unwrap(); + assert!(dag.execute(&[0, 1], context()).is_err()); + assert_eq!(starts.get(), 0); + let mut missing = PhysicalDag::default(); + missing.add(1, vec![9], Identity).unwrap(); + assert!(missing.validate(&[1]).is_err()); + let mut arity = PhysicalDag::default(); + arity.add(1, vec![], Identity).unwrap(); + assert!(arity.validate(&[1]).is_err()); +} + +// An always-ready source must yield so cancellation can be polled on this worker. +#[test] +fn ready_sources_cooperate_with_cancellation() { + let (mut source, _, polls) = source(false); + source.end = 10_000; + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + let context = context(); + let mut input = dag.execute(&[0], context.clone()).unwrap().remove(0); + block_on(async { + let drain = async { + while let Some(result) = input.next().await { + if let Err(error) = result { + assert_eq!(error, Error::Cancelled); + return; + } + } + panic!("source completed without yielding"); + }; + let cancel = async { + context.cancel(); + }; + futures::join!(drain, cancel); + }); + assert_eq!(polls.get(), 32); + assert_eq!(context.retained_bytes(), 0); +} + +// Cached shorter paths must not hide an over-deep path through shared nodes. +#[test] +fn depth_limit_covers_shared_paths() { + let (source, _, _) = source(false); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], source).unwrap(); + for id in 1..129 { + dag.add(id, vec![id - 1], Identity).unwrap(); + } + assert!(dag.validate(&(0..129).collect::>()).is_err()); +} diff --git a/crates/asap-physical-operators/src/dag/values.rs b/crates/asap-physical-operators/src/dag/values.rs new file mode 100644 index 000000000..5b4e043f7 --- /dev/null +++ b/crates/asap-physical-operators/src/dag/values.rs @@ -0,0 +1,318 @@ +//! Runtime values preserve Planner schemas; summary states are typed values too. +use super::Error; +use crate::AggregateCore; +use planner_types::{ + post_asap::{SummaryFamilyType, SummarySchema}, + pre_asap::DataType, +}; +use std::{cmp::Ordering, sync::Arc}; +pub type Schema = Arc; +#[derive(Clone)] +pub enum Value { + Null, + Bool(bool), + Int64(i64), + Float64(f64), + Utf8(Arc), + Timestamp(i64), + Date(i32), + Interval { + months: i32, + days: i32, + nanos: i64, + }, + List(Arc<[Value]>), + Struct(Arc<[Value]>), + Map(Arc<[(Value, Value)]>), + Summary { + family: SummaryFamilyType, + state: Arc, + }, +} +impl std::fmt::Debug for Value { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Summary { family, .. } => f.debug_tuple("Summary").field(family).finish(), + _ => write!(f, "{:?}", self.key()), + } + } +} +impl Value { + pub fn bytes(&self) -> usize { + std::mem::size_of::() + + match self { + Self::Utf8(s) => s.len(), + Self::List(v) | Self::Struct(v) => v.iter().map(Self::bytes).sum(), + Self::Map(v) => v.iter().map(|(k, v)| k.bytes() + v.bytes()).sum(), + Self::Summary { state, .. } => state.approx_memory_bytes(), + _ => 0, + } + } + pub fn matches(&self, dtype: &DataType, nullable: bool) -> bool { + if matches!(self, Self::Null) { + return nullable || matches!(dtype, DataType::Null); + } + match (self, dtype) { + (Self::Bool(_), DataType::Bool) + | (Self::Int64(_), DataType::Int64) + | (Self::Float64(_), DataType::Float64) + | (Self::Utf8(_), DataType::Utf8) + | (Self::Timestamp(_), DataType::Timestamp) + | (Self::Date(_), DataType::Date) + | (Self::Interval { .. }, DataType::Interval) => true, + (Self::List(v), DataType::List { element }) => v + .iter() + .all(|v| v.matches(&element.dtype, element.nullable)), + (Self::Struct(v), DataType::Struct { fields }) => { + v.len() == fields.len() + && v.iter() + .zip(fields) + .all(|(v, f)| v.matches(&f.dtype, f.nullable)) + } + ( + Self::Map(v), + DataType::Map { + key, + value, + value_nullable, + }, + ) => v + .iter() + .all(|(k, v)| k.matches(key, false) && v.matches(value, *value_nullable)), + _ => false, + } + } + /// Stable typed equality key. Zero signs and NaN payloads form one group. + pub fn key(&self) -> Result, Error> { + let mut out = Vec::new(); + macro_rules! number { + ($tag:expr,$v:expr) => {{ + out.push($tag); + out.extend_from_slice(&$v.to_le_bytes()); + }}; + } + match self { + Self::Null => out.push(0), + Self::Bool(v) => out.extend([1, *v as u8]), + Self::Int64(v) => number!(2, v), + Self::Float64(v) => { + let bits = if *v == 0. { + 0 + } else if v.is_nan() { + f64::NAN.to_bits() + } else { + v.to_bits() + }; + number!(3, bits); + } + Self::Utf8(v) => { + out.push(4); + out.extend(v.as_bytes()); + } + Self::Timestamp(v) => number!(5, v), + Self::Date(v) => number!(6, v), + Self::Interval { + months, + days, + nanos, + } => { + number!(7, months); + number!(8, days); + number!(9, nanos); + } + Self::List(v) | Self::Struct(v) => { + out.push(if matches!(self, Self::List(_)) { + 10 + } else { + 11 + }); + for v in v.iter() { + let key = v.key()?; + out.extend((key.len() as u64).to_le_bytes()); + out.extend(key); + } + } + Self::Map(v) => { + out.push(12); + for (k, v) in v.iter() { + for value in [k, v] { + let key = value.key()?; + out.extend((key.len() as u64).to_le_bytes()); + out.extend(key); + } + } + } + Self::Summary { .. } => { + return Err(Error::Invalid( + "summary states cannot be grouping keys".into(), + )) + } + } + Ok(out) + } + pub fn compare(&self, other: &Self) -> Result { + Ok(match (self, other) { + (Self::Null, Self::Null) => Ordering::Equal, + (Self::Int64(a), Self::Int64(b)) | (Self::Timestamp(a), Self::Timestamp(b)) => a.cmp(b), + (Self::Float64(a), Self::Float64(b)) => { + if a == b { + Ordering::Equal + } else { + a.total_cmp(b) + } + } + (Self::Utf8(a), Self::Utf8(b)) => a.cmp(b), + (Self::Bool(a), Self::Bool(b)) => a.cmp(b), + (Self::Date(a), Self::Date(b)) => a.cmp(b), + _ => { + return Err(Error::Operator( + "values do not have a supported common ordering".into(), + )) + } + }) + } +} +#[derive(Clone, Debug)] +pub struct Batch { + schema: Schema, + rows: Vec>, +} +impl Batch { + pub fn try_new(schema: Schema, rows: Vec>) -> Result { + validate_schema(&schema)?; + for row in &rows { + if row.len() != schema.fields.len() { + return Err(Error::Invalid( + "row width differs from Planner schema".into(), + )); + } + for (value, field) in row.iter().zip(&schema.fields) { + let matches = match (&field.dtype, value) { + (SummaryFamilyType::Plain(dtype), value) => { + value.matches(dtype, field.nullable) + } + (expected, Value::Summary { family, state }) => { + expected == family && validate_state(family, state.as_ref()).is_ok() + } + _ => false, + }; + if !matches { + return Err(Error::Invalid(format!( + "value differs from type of {}", + field.name + ))); + } + } + } + Ok(Self { schema, rows }) + } + pub fn schema(&self) -> &Schema { + &self.schema + } + pub fn rows(&self) -> &[Vec] { + &self.rows + } + pub fn bytes(&self) -> usize { + std::mem::size_of::() + + self + .rows + .iter() + .flat_map(|r| r.iter()) + .map(Value::bytes) + .sum::() + } +} +pub(crate) fn group_key(row: &[Value], columns: &[usize]) -> Result>, Error> { + columns + .iter() + .map(|&i| { + row.get(i) + .ok_or_else(|| Error::Invalid("group column out of range".into()))? + .key() + }) + .collect() +} + +pub(crate) fn validate_family(family: &SummaryFamilyType) -> Result<(), Error> { + use planner_types::post_asap::SketchAlgorithm as A; + match family { + SummaryFamilyType::ExactAggregate(..) => {} + SummaryFamilyType::Sketch(kind, _) + if matches!(kind.algorithm(), A::Kll | A::DDSketch | A::Hll) => {} + _ => { + return Err(Error::Invalid( + "summary family has no native DAG state implementation".into(), + )) + } + } + crate::capability::validate_summary_kernel( + family, + &planner_types::post_asap::SummaryUpdate::column( + planner_types::pre_asap::ColumnRef::SampleValue, + ), + &Default::default(), + ) + .map_err(Error::Invalid) +} +fn validate_state(family: &SummaryFamilyType, state: &dyn AggregateCore) -> Result<(), Error> { + use crate::accumulators::{ + datasketches_kll_accumulator::DatasketchesKLLAccumulator, + dd_sketch_accumulator::DDSketchAccumulator, exact_accumulator::ExactAccumulator, + hll_sketch_accumulator::HllSketchAccumulator, + }; + use planner_types::post_asap::SketchParams; + validate_family(family)?; + let valid = match family { + SummaryFamilyType::ExactAggregate(..) => state + .as_any() + .downcast_ref::() + .is_some_and(|s| s.family() == family && !s.is_keyed()), + SummaryFamilyType::Sketch(kind, _) => match kind.params() { + SketchParams::Kll { k } => state + .as_any() + .downcast_ref::() + .is_some_and(|s| u32::from(s.inner.k()) == *k), + SketchParams::DDSketch { alpha } => state + .as_any() + .downcast_ref::() + .is_some_and(|s| s.inner.alpha == *alpha && s.sample_p == 1.0), + SketchParams::Hll { precision } => state + .as_any() + .downcast_ref::() + .is_some_and(|s| s.inner.precision == u32::from(*precision) && s.sample_p == 1.0), + _ => false, + }, + _ => false, + }; + if valid { + Ok(()) + } else { + Err(Error::Invalid( + "state payload differs from declared family, parameters or population layout".into(), + )) + } +} + +pub(crate) fn validate_schema(schema: &Schema) -> Result<(), Error> { + if schema.time_index.is_some_and(|index| { + schema + .fields + .get(index) + .is_none_or(|field| field.dtype != SummaryFamilyType::Plain(DataType::Timestamp)) + }) { + return Err(Error::Invalid( + "time index must name a Timestamp column".into(), + )); + } + for field in &schema.fields { + if !matches!(field.dtype, SummaryFamilyType::Plain(_)) { + validate_family(&field.dtype)?; + if field.nullable { + return Err(Error::Invalid( + "nullable summary states are not supported".into(), + )); + } + } + } + Ok(()) +} diff --git a/crates/asap-physical-operators/src/factory.rs b/crates/asap-physical-operators/src/factory.rs new file mode 100644 index 000000000..fe6bed87f --- /dev/null +++ b/crates/asap-physical-operators/src/factory.rs @@ -0,0 +1,2046 @@ +use crate::accumulators::{ + CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, + CountSketchWithHeapAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, + HydraKllSketchAccumulator, IncreaseAccumulator, KeyedCounterState, KeyedMaxState, + KeyedMinState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, +}; +#[cfg(test)] +use crate::AggregationType; +use crate::{AggregateCore, KeyByLabelValues, Measurement}; +#[cfg(test)] +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 crate::accumulators::hll_sketch_accumulator::HllSketchAccumulator; +use crate::accumulators::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) + } + }; +} + +/// Shared update interface for query-time and maintenance-time accumulation. +/// +/// This provides a uniform interface over all accumulator types so that the +/// worker loop doesn't need to know which concrete type it's dealing with. +pub trait AccumulatorUpdater: Send { + /// Validate an immutable 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() { + use planner_types::post_asap::{ + EntityIdentity, NonNegativeWeightProof, SummaryInputExpr, SummaryUpdate, WeightDomain, + }; + let config = topk_config(AggregationType::CountMinSketchWithHeap, None); + let family = config.accumulator_spec().unwrap().family; + let input = SummaryUpdate { + item: Some(SummaryInputExpr::Column( + planner_types::pre_asap::ColumnRef::Named("host".into()), + )), + weight: SummaryInputExpr::ResetAwareCounterDelta { + value: planner_types::pre_asap::ColumnRef::SampleValue, + series: EntityIdentity::PromqlLabelSet { excluding: vec![] }, + }, + weight_domain: WeightDomain::NonNegative { + proof: NonNegativeWeightProof::ResetAwareCounterDerivative, + }, + }; + let mut updater = create_planner_accumulator(&family, &input, &Default::default()).unwrap(); + 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> { + crate::capability::validate_summary_kernel(family, input, grouping)?; + use planner_types::post_asap::GroupingStrategy; + if grouping != &GroupingStrategy::PerSubpopulationInstance { + return Err("shared summary grouping requires a supported Planner Hydra kernel".into()); + } + if matches!(family, SummaryFamilyType::ExactAggregate(..)) { + return Ok(Box::new(PlannerExactUpdater { + acc: crate::accumulators::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()); + } + // Heap counters use fixed-point storage for fractional counter deltas. + // This encodes the selected update; it does not choose another family. + let weight_scale = if matches!( + input.weight, + planner_types::post_asap::SummaryInputExpr::ResetAwareCounterDelta { .. } + ) { + 1_000_000.0 + } else { + 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: crate::accumulators::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 = crate::accumulators::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/crates/asap-physical-operators/src/key_by_label_values.rs b/crates/asap-physical-operators/src/key_by_label_values.rs new file mode 100644 index 000000000..34bc84899 --- /dev/null +++ b/crates/asap-physical-operators/src/key_by_label_values.rs @@ -0,0 +1,164 @@ +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/crates/asap-physical-operators/src/lib.rs b/crates/asap-physical-operators/src/lib.rs new file mode 100644 index 000000000..266a516fa --- /dev/null +++ b/crates/asap-physical-operators/src/lib.rs @@ -0,0 +1,22 @@ +#![doc = include_str!("../README.md")] + +pub mod accumulators; +pub mod key_by_label_values; +pub mod measurement; +pub mod traits; + +pub use asap_types::{AggregationType, Statistic}; +pub use key_by_label_values::KeyByLabelValues; +pub use measurement::Measurement; +pub use traits::*; + +pub mod arithmetic; +pub mod capability; +pub mod factory; + +/// The exact Planner contract used by these kernels. +pub use planner_types as planner; + +pub mod rows; + +pub mod dag; diff --git a/crates/asap-physical-operators/src/measurement.rs b/crates/asap-physical-operators/src/measurement.rs new file mode 100644 index 000000000..0fe1abc0d --- /dev/null +++ b/crates/asap-physical-operators/src/measurement.rs @@ -0,0 +1,94 @@ +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/crates/asap-physical-operators/src/rows.rs b/crates/asap-physical-operators/src/rows.rs new file mode 100644 index 000000000..869ac0c30 --- /dev/null +++ b/crates/asap-physical-operators/src/rows.rs @@ -0,0 +1,88 @@ +//! Composable row operators, independent of storage, query language and sketches. +use std::collections::{BTreeMap, BTreeSet}; + +/// A semijoin preserves value-row order and multiplicity; duplicate membership +/// keys never multiply rows. Missing membership keys are reported separately so +/// the deployment can enforce the pruning proof attached to its plan. +pub fn membership_filter( + members: impl IntoIterator, + values: Vec, + identity: impl Fn(&T) -> K, +) -> (Vec, BTreeSet) { + let members: BTreeSet = members.into_iter().collect(); + let mut missing = members.clone(); + let rows = values + .into_iter() + .filter(|row| { + let key = identity(row); + missing.remove(&key); + members.contains(&key) + }) + .collect(); + (rows, missing) +} + +/// Stable descending TopK per group. NaN sorts after numeric values; ties keep +/// input order. This operator does not know how its input was filtered or built. +pub fn grouped_topk( + values: Vec, + k: usize, + group_key: impl Fn(&T) -> K, + score: impl Fn(&T) -> f64, +) -> Vec { + let mut groups: BTreeMap> = BTreeMap::new(); + for row in values { + groups.entry(group_key(&row)).or_default().push(row); + } + groups + .into_values() + .flat_map(|mut rows| { + rows.sort_by(|a, b| { + let (a, b) = (score(a), score(b)); + match (a.is_nan(), b.is_nan()) { + (true, true) => std::cmp::Ordering::Equal, + (true, false) => std::cmp::Ordering::Greater, + (false, true) => std::cmp::Ordering::Less, + (false, false) => b.total_cmp(&a), + } + }); + rows.truncate(k); + rows + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn semijoin_preserves_values_order_and_duplicates_without_ranking() { + let (rows, missing) = membership_filter( + ["b", "c", "b", "missing"], + vec![("a", 100.), ("b", 2.), ("c", 9.), ("b", 3.)], + |row| row.0, + ); + assert_eq!(rows, vec![("b", 2.), ("c", 9.), ("b", 3.)]); + assert_eq!(missing, BTreeSet::from(["missing"])); + assert_eq!( + grouped_topk(rows, 2, |_| (), |r| r.1), + vec![("c", 9.), ("b", 3.)] + ); + } + + #[test] + fn grouped_ranking_preserves_ties_and_places_nan_last() { + let rows = vec![("x", 0, f64::NAN), ("x", 1, 2.), ("y", 2, 8.), ("x", 3, 2.)]; + let ranked = grouped_topk(rows, 2, |r| r.0, |r| r.2); + assert_eq!(ranked, vec![("x", 1, 2.), ("x", 3, 2.), ("y", 2, 8.)]); + assert!(grouped_topk(vec![1], 0, |_| (), |r| *r as f64).is_empty()); + } + + #[test] + fn empty_membership_removes_all_rows() { + let (rows, missing) = membership_filter([], vec![1, 2], |r| *r); + assert!(rows.is_empty()); + assert!(missing.is_empty()); + } +} diff --git a/crates/asap-physical-operators/src/traits.rs b/crates/asap-physical-operators/src/traits.rs new file mode 100644 index 000000000..2b99da789 --- /dev/null +++ b/crates/asap-physical-operators/src/traits.rs @@ -0,0 +1,351 @@ +use crate::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/crates/asap-physical-operators/tests/deployment.rs b/crates/asap-physical-operators/tests/deployment.rs new file mode 100644 index 000000000..ffdd3aea5 --- /dev/null +++ b/crates/asap-physical-operators/tests/deployment.rs @@ -0,0 +1,96 @@ +//! Exercise the public library without a backend server, store, or scheduler. +use asap_physical_operators::planner::{ + post_asap::{ + GroupingStrategy, SketchAlgorithm, SketchKind, SketchParams, SummaryFamilyType, + SummaryUpdate, + }, + pre_asap::ColumnRef, +}; +use asap_physical_operators::{factory::create_planner_accumulator, AggregateCore, Statistic}; +use std::collections::HashMap; + +fn family(k: u32) -> SummaryFamilyType { + SummaryFamilyType::Sketch( + SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k }), + GroupingStrategy::PerSubpopulationInstance, + ) +} +fn build(values: &[f64]) -> Box { + let mut operator = create_planner_accumulator( + &family(512), + &SummaryUpdate::column(ColumnRef::SampleValue), + &Default::default(), + ) + .unwrap(); + for (at, value) in values.iter().enumerate() { + operator.validate_single_input(*value).unwrap(); + operator.update_single(*value, at as i64); + } + operator.into_accumulator() +} +fn read(state: &dyn AggregateCore) -> f64 { + state + .query_statistic( + Statistic::Quantile, + &None, + &HashMap::from([("quantile".into(), "0.5".into())]), + ) + .unwrap() +} + +// The same kernels work when every build is query-time, when only a prefix +// was precomputed, and when all state was precomputed before the readout. +#[test] +fn raw_partial_and_fully_precomputed_use_the_same_kernels() { + let raw: Vec = (0..128).map(f64::from).collect(); + let raw_only = build(&raw); + let stored_prefix = build(&raw[..64]); + let query_time_suffix = build(&raw[64..]); + let partial = stored_prefix.merge_with(&*query_time_suffix).unwrap(); + let stored_complete = build(&raw); + assert_eq!(read(&*raw_only), read(&*partial)); + assert_eq!(read(&*partial), read(&*stored_complete)); + assert!((read(&*raw_only) - 64.0).abs() <= 1.0); +} + +// A compiler must reject invalid physical parameters before starting execution. +#[test] +fn invalid_kll_parameters_are_rejected_at_binding() { + let result = create_planner_accumulator( + &family(0), + &SummaryUpdate::column(ColumnRef::SampleValue), + &Default::default(), + ); + assert!(result.is_err()); +} + +// Native CountSketch supports the confidence-sized depth used by the backend; +// a packed-wire column-bit budget must not be imposed on this constructor. +#[test] +fn native_count_sketch_dimensions_are_not_packed_wire_dimensions() { + use asap_physical_operators::planner::post_asap::SummaryInputExpr; + use asap_physical_operators::KeyByLabelValues; + let family = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::CountSketchWithHeap, + SketchParams::CountSketchWithHeap { + width: 1200, + depth: 55, + heap_size: 3, + }, + ), + Default::default(), + ); + let mut update = SummaryUpdate::column(ColumnRef::SampleValue); + update.item = Some(SummaryInputExpr::Column(ColumnRef::Named("host".into()))); + let mut operator = create_planner_accumulator(&family, &update, &Default::default()).unwrap(); + let key = KeyByLabelValues::new_with_labels(vec!["a".into()]); + operator.update_keyed(&key, 7.0, 1000); + let state = operator.into_accumulator(); + assert_eq!( + state + .query_statistic(Statistic::Sum, &Some(key), &Default::default()) + .unwrap(), + 7.0 + ); +} diff --git a/crates/asap-physical-operators/tests/physical_dag.rs b/crates/asap-physical-operators/tests/physical_dag.rs new file mode 100644 index 000000000..a46cd2d01 --- /dev/null +++ b/crates/asap-physical-operators/tests/physical_dag.rs @@ -0,0 +1,663 @@ +//! Acceptance tests use the library directly, without either backend engine. +use asap_physical_operators::{ + dag::{ + operators::{Expression, Operator, Reduction, SortKey}, + values::{Batch, Schema, Value}, + Limits, PhysicalDag, RunContext, Scope, + }, + Statistic, +}; +use futures::{executor::block_on, StreamExt}; +use planner_types::{ + post_asap::{ExactKind, ExactParams, SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, +}; +use std::sync::Arc; +fn schema(fields: &[(&str, DataType, bool)]) -> Schema { + Arc::new(SummarySchema { + fields: fields + .iter() + .map(|(name, dtype, nullable)| SummaryField { + name: (*name).into(), + dtype: SummaryFamilyType::Plain(dtype.clone()), + nullable: *nullable, + }) + .collect(), + time_index: None, + }) +} +fn run(dag: &PhysicalDag<'_, Batch, Schema>, root: u64, scope: Scope) -> Vec> { + let context = RunContext::new( + scope, + Limits { + max_buffered_batches: 1, + ..Limits::default() + }, + ) + .unwrap(); + block_on(async { + let mut stream = dag.execute(&[root], context.clone()).unwrap().remove(0); + let mut rows = vec![]; + while let Some(batch) = stream.next().await { + rows.extend(batch.unwrap().rows().iter().cloned()); + } + assert_eq!(context.retained_bytes(), 0); + rows + }) +} +fn query() -> Scope { + Scope::Query { + evaluation_time_ms: 1000, + revision: 2, + } +} +fn floats(rows: &[Vec], column: usize) -> Vec { + rows.iter() + .map(|r| { + if let Value::Float64(v) = r[column] { + v + } else { + panic!("not Float64") + } + }) + .collect() +} + +// Sort followed by partitioned Limit implements ranking independently per group. +#[test] +fn grouped_sort_limit_across_batches() { + let schema = schema(&[ + ("group", DataType::Int64, false), + ("score", DataType::Float64, false), + ]); + let batches = [ + vec![(1, 1.), (2, 4.), (1, 9.)], + vec![(2, 8.), (1, 5.), (2, 2.)], + ] + .into_iter() + .map(|rows| { + Batch::try_new( + schema.clone(), + rows.into_iter() + .map(|(g, v)| vec![Value::Int64(g), Value::Float64(v)]) + .collect(), + ) + .unwrap() + }) + .collect(); + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source(schema.clone(), batches).unwrap(), + ) + .unwrap(); + dag.add( + 1, + vec![0], + Operator::sort( + schema.clone(), + vec![SortKey { + column: 1, + descending: true, + nulls_first: false, + }], + vec![0], + ) + .unwrap(), + ) + .unwrap(); + dag.add(2, vec![1], Operator::limit(schema, 1, 1, vec![0]).unwrap()) + .unwrap(); + assert_eq!(floats(&run(&dag, 2, query()), 1), vec![5., 4.]); +} + +// The same computation runs in either engine scope with fresh per-run state. +#[test] +fn summary_construction_merge_and_readout_at_both_phases() { + let schema = schema(&[("v", DataType::Float64, false)]); + let batches = (1..=20) + .map(|v| Batch::try_new(schema.clone(), vec![vec![Value::Float64(v as f64)]]).unwrap()) + .collect(); + let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + let build = Operator::summary_build(schema.clone(), family, 0, None, vec![]).unwrap(); + let state = build.schema(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], Operator::source(schema, batches).unwrap()) + .unwrap(); + dag.add(1, vec![0], build).unwrap(); + dag.add(2, vec![1, 1], Operator::union(state.clone(), 2).unwrap()) + .unwrap(); + dag.add( + 3, + vec![2], + Operator::summary_merge(state.clone(), 0, vec![]).unwrap(), + ) + .unwrap(); + dag.add( + 4, + vec![3], + Operator::readout(state, 0, Statistic::Sum, Default::default()).unwrap(), + ) + .unwrap(); + for scope in [ + query(), + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 1000, + revision: 2, + }, + ] { + assert_eq!(floats(&run(&dag, 4, scope), 0), vec![420.]); + } +} + +// A semi-join can consume two branches of one producer with a one-batch buffer. +#[test] +fn diamond_semijoin_preserves_left_values_and_multiplicity() { + let schema = schema(&[("key", DataType::Int64, false)]); + let batches = [1, 2, 2, 3] + .into_iter() + .map(|v| Batch::try_new(schema.clone(), vec![vec![Value::Int64(v)]]).unwrap()) + .collect(); + let filter = Operator::filter( + schema.clone(), + Expression::Equal( + Box::new(Expression::Column(0)), + Box::new(Expression::Literal { + value: Value::Int64(2), + dtype: DataType::Int64, + }), + ), + ) + .unwrap(); + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source(schema.clone(), batches).unwrap(), + ) + .unwrap(); + dag.add(1, vec![0], filter).unwrap(); + dag.add( + 2, + vec![0, 1], + Operator::semi_join(schema.clone(), schema, vec![(0, 0)]).unwrap(), + ) + .unwrap(); + let rows = run(&dag, 2, query()); + assert_eq!(rows.len(), 2); + assert!(rows.iter().all(|r| matches!(r[0], Value::Int64(2)))); +} + +// Integer aggregation must not silently lose precision through Float64. +#[test] +fn exact_integer_and_empty_extrema() { + let schema = schema(&[("v", DataType::Int64, false)]); + let aggregate = Operator::aggregate( + schema.clone(), + vec![], + vec![("sum".into(), Reduction::Sum(0))], + ) + .unwrap(); + let mut dag = PhysicalDag::default(); + let value = 9_007_199_254_740_993; + dag.add( + 0, + vec![], + Operator::source( + schema.clone(), + vec![Batch::try_new( + schema.clone(), + vec![vec![Value::Int64(value)], vec![Value::Int64(2)]], + ) + .unwrap()], + ) + .unwrap(), + ) + .unwrap(); + dag.add(1, vec![0], aggregate).unwrap(); + assert!(matches!(run(&dag,1,query())[0][0],Value::Int64(v) if v==value+2)); + let mut empty = PhysicalDag::default(); + empty + .add(0, vec![], Operator::source(schema.clone(), vec![]).unwrap()) + .unwrap(); + empty + .add( + 1, + vec![0], + Operator::aggregate(schema, vec![], vec![("min".into(), Reduction::Min(0))]).unwrap(), + ) + .unwrap(); + assert!(matches!(run(&empty, 1, query())[0][0], Value::Null)); +} + +// Plain value operators are library implementations, including NaN comparison. +#[test] +fn scalar_negation_and_vector_conversion() { + let scalar = Operator::scalar(Value::Float64(7.), DataType::Float64).unwrap(); + let project = Operator::project( + scalar.schema(), + vec![( + "v".into(), + Expression::Negate(Box::new(Expression::Column(0))), + )], + ) + .unwrap(); + let convert = Operator::vector_to_scalar(project.schema(), 0).unwrap(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], scalar).unwrap(); + dag.add(1, vec![0], project).unwrap(); + dag.add(2, vec![1], convert).unwrap(); + assert_eq!(floats(&run(&dag, 2, query()), 0), vec![-7.]); + let scalar = Operator::scalar(Value::Float64(f64::NAN), DataType::Float64).unwrap(); + let predicate = Expression::Equal( + Box::new(Expression::Column(0)), + Box::new(Expression::Column(0)), + ); + let filter = Operator::filter(scalar.schema(), predicate).unwrap(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], scalar).unwrap(); + dag.add(1, vec![0], filter).unwrap(); + assert!(run(&dag, 1, query()).is_empty()); +} + +// Invalid operations fail at binding rather than becoming external fallbacks. +#[test] +fn binding_rejects_unsupported_operations() { + let schema = schema(&[("v", DataType::Float64, false)]); + assert!(Operator::summary_build( + schema.clone(), + SummaryFamilyType::ExactAggregate(ExactKind::Rate, ExactParams::Rate), + 0, + None, + vec![] + ) + .is_err()); + let sum = Operator::summary_build( + schema.clone(), + SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum), + 0, + None, + vec![], + ) + .unwrap(); + assert!(Operator::readout(sum.schema(), 0, Statistic::Quantile, Default::default()).is_err()); + assert!(Operator::filter(schema, Expression::Column(0)).is_err()); +} + +// KLL is one family example: precomputation changes input sources, not operators. +#[test] +fn kll_raw_partial_and_precomputed_are_native_dags() { + use planner_types::post_asap::{GroupingStrategy, SketchAlgorithm, SketchKind, SketchParams}; + let input = schema(&[("value", DataType::Float64, false)]); + let family = SummaryFamilyType::Sketch( + SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k: 512 }), + GroupingStrategy::PerSubpopulationInstance, + ); + let build = Operator::summary_build(input.clone(), family, 0, None, vec![]).unwrap(); + let state = build.schema(); + let build_range = |start: u32, end: u32| { + let mut dag = PhysicalDag::default(); + let batch = Batch::try_new( + input.clone(), + (start..end) + .map(|v| vec![Value::Float64(f64::from(v))]) + .collect(), + ) + .unwrap(); + dag.add( + 0, + vec![], + Operator::source(input.clone(), vec![batch]).unwrap(), + ) + .unwrap(); + dag.add(1, vec![0], build.clone()).unwrap(); + run( + &dag, + 1, + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 1000, + revision: 1, + }, + ) + }; + let prefix = build_range(0, 64); + let complete = build_range(0, 128); + let query_plan = |stored: Option>>, raw_start: Option| { + let mut dag = PhysicalDag::default(); + let mut states = vec![]; + if let Some(rows) = stored { + dag.add( + 0, + vec![], + Operator::source( + state.clone(), + vec![Batch::try_new(state.clone(), rows).unwrap()], + ) + .unwrap(), + ) + .unwrap(); + states.push(0); + } + if let Some(start) = raw_start { + dag.add( + 1, + vec![], + Operator::source( + input.clone(), + vec![Batch::try_new( + input.clone(), + (start..128) + .map(|v| vec![Value::Float64(f64::from(v))]) + .collect(), + ) + .unwrap()], + ) + .unwrap(), + ) + .unwrap(); + dag.add(2, vec![1], build.clone()).unwrap(); + states.push(2); + } + dag.add( + 3, + states.clone(), + Operator::union(state.clone(), states.len()).unwrap(), + ) + .unwrap(); + dag.add( + 4, + vec![3], + Operator::summary_merge(state.clone(), 0, vec![]).unwrap(), + ) + .unwrap(); + dag.add( + 5, + vec![4], + Operator::readout( + state.clone(), + 0, + Statistic::Quantile, + std::collections::HashMap::from([("quantile".into(), "0.5".into())]), + ) + .unwrap(), + ) + .unwrap(); + floats(&run(&dag, 5, query()), 0)[0] + }; + let raw = query_plan(None, Some(0)); + let partial = query_plan(Some(prefix), Some(64)); + let full = query_plan(Some(complete), None); + assert_eq!(raw, partial); + assert_eq!(partial, full); + assert!((raw - 64.).abs() <= 1.); +} + +// Restored state must retain its family; a mislabeled state is rejected. +#[test] +fn restored_exact_state_and_family_validation() { + use asap_physical_operators::{ + accumulators::exact_accumulator::ExactAccumulator, SerializableToSink, + }; + let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + let mut acc = ExactAccumulator::new(family.clone(), false).unwrap(); + acc.update(None, 7., 0); + let acc = ExactAccumulator::deserialize_from_bytes(&acc.serialize_to_bytes()).unwrap(); + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "state".into(), + dtype: family.clone(), + nullable: false, + }], + time_index: None, + }); + let value = Value::Summary { + family: family.clone(), + state: Arc::new(acc), + }; + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source( + schema.clone(), + vec![Batch::try_new(schema.clone(), vec![vec![value]]).unwrap()], + ) + .unwrap(), + ) + .unwrap(); + dag.add( + 1, + vec![0], + Operator::readout(schema.clone(), 0, Statistic::Sum, Default::default()).unwrap(), + ) + .unwrap(); + assert_eq!(floats(&run(&dag, 1, query()), 0), vec![7.]); + let wrong = ExactAccumulator::new( + SummaryFamilyType::ExactAggregate(ExactKind::Max, ExactParams::Max), + false, + ) + .unwrap(); + assert!(Batch::try_new( + schema, + vec![vec![Value::Summary { + family, + state: Arc::new(wrong) + }]] + ) + .is_err()); +} + +// Planner binding rejects unknown computation instead of accepting a fallback. +#[test] +fn bind_post_asap_before_execution() { + use asap_physical_operators::dag::planner::bind; + use planner_types::{ + post_asap::{ + EdgeRole, ExecutableDag, ExecutableDagEdge, ExecutableDagNode, + ExecutableOperatorPayload, ExecutionDataState, ExecutionTiming, + GroupingEdgeCompatibility, PostAsapNodeId, ValueOperation, WindowEdgeCompatibility, + }, + pre_asap::{ArithmeticOpKind, ProjectItem, QueryExpr, ScalarValue}, + }; + use std::{collections::BTreeMap, rc::Rc}; + let schema = schema(&[("value", DataType::Float64, false)]); + let node = |id, payload| ExecutableDagNode { + id: PostAsapNodeId(id), + payload, + output_state: ExecutionDataState::QUERY_ROWS, + output_schema: (*schema).clone(), + guarantee: None, + }; + let mut dag = ExecutableDag { + nodes: vec![ + node( + 0, + ExecutableOperatorPayload::Fallback { + expression: QueryExpr::promql_scalar(1.), + }, + ), + node( + 1, + ExecutableOperatorPayload::Value { + timing: ExecutionTiming::QueryTime, + operation: ValueOperation::Project { + cols: vec![ProjectItem { + alias: None, + expr: QueryExpr::Arithmetic { + op: ArithmeticOpKind::Add, + left: Rc::new(QueryExpr::Column(0)), + right: Rc::new(QueryExpr::Literal(ScalarValue::Float64(2.))), + }, + }], + qualifier: None, + }, + }, + ), + ], + edges: vec![ExecutableDagEdge { + producer: PostAsapNodeId(0), + consumer: PostAsapNodeId(1), + role: EdgeRole::Input, + intermediate_schema: (*schema).clone(), + data_state: ExecutionDataState::QUERY_ROWS, + grouping: GroupingEdgeCompatibility::NotApplicable, + window: WindowEdgeCompatibility::NotApplicable, + }], + root: PostAsapNodeId(1), + }; + let sources = || -> BTreeMap> { + BTreeMap::from([( + 0, + Box::new( + Operator::source( + schema.clone(), + vec![Batch::try_new(schema.clone(), vec![vec![Value::Float64(1.)]]).unwrap()], + ) + .unwrap(), + ) as asap_physical_operators::dag::planner::Source<'static>, + )]) + }; + let native = bind(&dag, sources(), &[1]).unwrap(); + assert_eq!(floats(&run(&native, 1, query()), 0), vec![3.]); + assert!(bind(&dag, BTreeMap::new(), &[1]).is_err()); + dag.nodes[1].payload = ExecutableOperatorPayload::Value { + timing: ExecutionTiming::QueryTime, + operation: ValueOperation::Extension { + name: "unknown".into(), + }, + }; + assert!(bind(&dag, sources(), &[1]).is_err()); +} + +// A completed empty population has an exact zero count, with integer output. +#[test] +fn empty_exact_count_is_an_integer_state_readout() { + let input = schema(&[("value", DataType::Float64, false)]); + let build = Operator::summary_build( + input.clone(), + SummaryFamilyType::ExactAggregate(ExactKind::Count, ExactParams::Count), + 0, + None, + vec![], + ) + .unwrap(); + let read = Operator::readout(build.schema(), 0, Statistic::Count, Default::default()).unwrap(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], Operator::source(input, vec![]).unwrap()) + .unwrap(); + dag.add(1, vec![0], build).unwrap(); + dag.add(2, vec![1], read).unwrap(); + assert!(matches!(run(&dag, 2, query())[0][0], Value::Int64(0))); +} + +// A deployment source cannot pass a different row shape to bound expressions. +#[test] +fn source_batches_must_match_the_bound_schema() { + use asap_physical_operators::dag::{self, PhysicalOperator}; + use planner_types::{ + post_asap::{ + ExecutableDag, ExecutableDagNode, ExecutableOperatorPayload, ExecutionDataState, + PostAsapNodeId, + }, + pre_asap::QueryExpr, + }; + use std::{cell::Cell, collections::BTreeMap, rc::Rc}; + struct WrongSource { + schema: Schema, + starts: Rc>, + } + impl PhysicalOperator for WrongSource { + fn name(&self) -> &str { + "ExternalSource" + } + fn input_schemas(&self) -> Vec { + vec![] + } + fn output_schema(&self) -> Schema { + self.schema.clone() + } + fn output_bytes(&self, value: &Batch) -> usize { + value.bytes() + } + fn start<'a>( + &'a self, + _: Vec>, + _: RunContext, + ) -> Result, dag::Error> { + self.starts.set(self.starts.get() + 1); + Ok( + futures::stream::once(async { Batch::try_new(schema(&[]), vec![vec![]]) }) + .boxed_local(), + ) + } + } + let expected = schema(&[("value", DataType::Float64, false)]); + let starts = Rc::new(Cell::new(0)); + let plan = ExecutableDag { + nodes: vec![ExecutableDagNode { + id: PostAsapNodeId(0), + payload: ExecutableOperatorPayload::Fallback { + expression: QueryExpr::promql_scalar(1.), + }, + output_state: ExecutionDataState::QUERY_ROWS, + output_schema: (*expected).clone(), + guarantee: None, + }], + edges: vec![], + root: PostAsapNodeId(0), + }; + let source = Box::new(WrongSource { + schema: expected, + starts: starts.clone(), + }) as dag::planner::Source<'static>; + let native = dag::planner::bind(&plan, BTreeMap::from([(0, source)]), &[0]).unwrap(); + assert_eq!(starts.get(), 0); + let context = RunContext::new(query(), Limits::default()).unwrap(); + let mut output = native.execute(&[0], context).unwrap().remove(0); + assert!(matches!( + block_on(output.next()), + Some(Err(dag::Error::AtNode { node: 0, .. })) + )); + assert_eq!(starts.get(), 1); +} + +// Float extrema have the same NaN behavior as the exact summary kernels. +#[test] +fn extrema_preserve_numeric_values_in_the_presence_of_nan() { + let input = schema(&[("v", DataType::Float64, false)]); + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source( + input.clone(), + vec![Batch::try_new( + input.clone(), + vec![vec![Value::Float64(-f64::NAN)], vec![Value::Float64(5.)]], + ) + .unwrap()], + ) + .unwrap(), + ) + .unwrap(); + dag.add( + 1, + vec![0], + Operator::aggregate( + input, + vec![], + vec![ + ("min".into(), Reduction::Min(0)), + ("max".into(), Reduction::Max(0)), + ], + ) + .unwrap(), + ) + .unwrap(); + let rows = run(&dag, 1, query()); + assert_eq!(floats(&rows, 0), vec![5.]); + assert_eq!(floats(&rows, 1), vec![5.]); +} diff --git a/crates/asap_types/src/derived_input.rs b/crates/asap_types/src/derived_input.rs index d53dd0807..e4596b7ca 100644 --- a/crates/asap_types/src/derived_input.rs +++ b/crates/asap_types/src/derived_input.rs @@ -215,7 +215,7 @@ mod tests { use planner_types::post_asap::{ EdgeRole, ExecutionDataState, GroupingEdgeCompatibility, WindowEdgeCompatibility, }; - let state = ExecutionDataState::MAINTENANCE_SUMMARY; + let state = ExecutionDataState::INGESTION_SUMMARY; OwnedPostAsapDag { schema_version: crate::executable_plan::OWNED_POST_ASAP_DAG_SCHEMA_VERSION, query_id: "query-a".into(), @@ -224,7 +224,7 @@ mod tests { .into_iter() .map(|id| OwnedPostAsapNode { id: PostAsapNodeId(id), - payload: serde_json::json!({"kind":"summary_merge"}), + payload: serde_json::json!({"kind":"summary_merge", "timing":"ingestion_time"}), output_state: state, output_schema: serde_json::json!({"fields":[],"time_index":null}), guarantee: None, diff --git a/crates/asap_types/src/executable_plan.rs b/crates/asap_types/src/executable_plan.rs index d84764dce..7c1b947a5 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 = 2; -pub const MAINTENANCE_DAG_SCHEMA_VERSION: u32 = 3; +pub const OWNED_POST_ASAP_DAG_SCHEMA_VERSION: u32 = 3; +pub const MAINTENANCE_DAG_SCHEMA_VERSION: u32 = 4; /// Versioned, language-neutral Planner DAG persisted with an installed plan. /// Plan lifecycle belongs to the enclosing `PrecomputePlan`; this document @@ -293,7 +293,7 @@ impl BackendExecutableBinding { for node in &dag.nodes { match (node.output_state.timing, self.node(node.id)) { ( - ExecutionTiming::MaintenanceTime, + ExecutionTiming::IngestionTime, Some( BackendNodeBinding::MaintenanceInput | BackendNodeBinding::Materialization { .. }, @@ -331,11 +331,10 @@ impl BackendExecutableBinding { } for node in &dag.nodes { match (node.output_state.timing, self.node(node.id).unwrap()) { - (ExecutionTiming::ReadTime, BackendNodeBinding::Query { .. }) - | (ExecutionTiming::ReadTime, BackendNodeBinding::QueryInput) - | (ExecutionTiming::MaintenanceTime, BackendNodeBinding::MaintenanceInput) - | (ExecutionTiming::MaintenanceTime, BackendNodeBinding::Materialization { .. }) => { - } + (ExecutionTiming::QueryTime, BackendNodeBinding::Query { .. }) + | (ExecutionTiming::QueryTime, BackendNodeBinding::QueryInput) + | (ExecutionTiming::IngestionTime, BackendNodeBinding::MaintenanceInput) + | (ExecutionTiming::IngestionTime, BackendNodeBinding::Materialization { .. }) => {} _ => { return Err(format!( "backend placement disagrees with node {} mode", diff --git a/crates/asap_types/src/precompute_plan.rs b/crates/asap_types/src/precompute_plan.rs index 47f30192a..12d9f421f 100644 --- a/crates/asap_types/src/precompute_plan.rs +++ b/crates/asap_types/src/precompute_plan.rs @@ -626,19 +626,19 @@ impl PrecomputePlan { ExecutableOperatorPayload as Payload, ExecutionTiming, ValueOperation, }; if node.output_state - != planner_types::post_asap::ExecutionDataState::MAINTENANCE_ROWS + != planner_types::post_asap::ExecutionDataState::INGESTION_ROWS { return Err(invalid()); } match &node.payload { Payload::Value { operation: ValueOperation::FinalizeExactAccumulator, - timing: ExecutionTiming::MaintenanceTime, + timing: ExecutionTiming::IngestionTime, } if children.len() == 1 && frontiers.contains_key(&children[0].producer) => {} Payload::Binary { operator, - timing: ExecutionTiming::MaintenanceTime, + timing: ExecutionTiming::IngestionTime, } if children.len() == 2 && children .iter() @@ -1056,7 +1056,7 @@ mod source_window_cohort_tests { reduction: Reduction::by(vec![]), grouping: Default::default(), }, - output_state: ExecutionDataState::MAINTENANCE_SUMMARY, + output_state: ExecutionDataState::INGESTION_SUMMARY, output_schema: SummarySchema { fields: vec![], time_index: None, @@ -1077,7 +1077,9 @@ mod source_window_cohort_tests { assert!(validate_maintenance_reduction(&config, &node).is_err()); config.partitioning = None; assert!(validate_maintenance_reduction(&config, &node).is_err()); - node.payload = ExecutableOperatorPayload::SummaryMerge; + node.payload = ExecutableOperatorPayload::SummaryMerge { + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, + }; assert!(validate_maintenance_reduction(&config, &node).is_err()); } diff --git a/crates/asap_types/src/query_plan.rs b/crates/asap_types/src/query_plan.rs index f7071c7b0..fc73f58c8 100644 --- a/crates/asap_types/src/query_plan.rs +++ b/crates/asap_types/src/query_plan.rs @@ -391,15 +391,7 @@ impl QueryPlanEntry { )); } } - if let QueryPlanNode::CandidateTopK { - k, completeness, .. - } = node - { - if *k == 0 { - return Err(QueryPlanError::Invalid( - "CandidateTopK requires k > 0".into(), - )); - } + if let QueryPlanNode::MembershipFilter { completeness, .. } = node { if matches!( completeness, CandidateCompleteness::Certified { guarantee } @@ -409,7 +401,7 @@ impl QueryPlanEntry { || guarantee.failure_probability.evaluate().is_none() ) { return Err(QueryPlanError::Invalid( - "invalid CandidateTopK completeness certificate".into(), + "invalid MembershipFilter completeness certificate".into(), )); } } @@ -613,13 +605,11 @@ pub enum QueryPlanNode { SummaryMerge { inputs: Vec, }, - /// Use an approximate heap only as a membership sidecar, then rerank the - /// matching exact counter readouts. `inputs[0]` is candidate membership; - /// `inputs[1]` is the authoritative exact value vector. - CandidateTopK { + /// 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], - k: u64, - grouping: residual::Grouping, completeness: CandidateCompleteness, }, /// An exact subtree evaluated outside ASAP. Its results enter the query DAG @@ -647,7 +637,7 @@ impl QueryPlanNode { Self::SummaryMerge { inputs } | Self::Logical { inputs, .. } | Self::ExternalExact { inputs, .. } => inputs, - Self::CandidateTopK { inputs, .. } => inputs, + Self::MembershipFilter { inputs, .. } => inputs, } } } diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index 7482c5210..3f999f0e3 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -6,6 +6,7 @@ edition.workspace = true [dependencies] # Internal crates (workspace) asap_types.workspace = true +asap-physical-operators.workspace = true asap_sketch_codec = { path = "../crates/asap_sketch_codec" } # Phase 9: the control plane is now an in-process library inside the # backend binary. Wiring up the in-process OpAMP server + capability-map diff --git a/data_plane/src/drivers/query/servers/http.rs b/data_plane/src/drivers/query/servers/http.rs index 776e91976..6c4c71514 100644 --- a/data_plane/src/drivers/query/servers/http.rs +++ b/data_plane/src/drivers/query/servers/http.rs @@ -6650,7 +6650,11 @@ mod catalog_install_tests { }) .expect("demo has maintained summaries"); binding.window_ms += 1; - assert!(install(request).unwrap_err().contains("query pane differs")); + let error = install(request).unwrap_err(); + assert!( + error.contains("query physical pane duration differs"), + "{error}" + ); } #[test] diff --git a/data_plane/src/precompute_engine/maintenance_runtime.rs b/data_plane/src/precompute_engine/maintenance_runtime.rs index 428a42af1..78c58dea5 100644 --- a/data_plane/src/precompute_engine/maintenance_runtime.rs +++ b/data_plane/src/precompute_engine/maintenance_runtime.rs @@ -167,14 +167,16 @@ impl PrecomputeOperatorRegistry for OperatorAdapter<'_> { inputs: &[Arc], ) -> Result { match &node.payload { - ExecutableOperatorPayload::SummaryMerge => merge_inputs(inputs), + ExecutableOperatorPayload::SummaryMerge { + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, + } => merge_inputs(inputs), ExecutableOperatorPayload::Binary { operator, - timing: planner_types::post_asap::ExecutionTiming::MaintenanceTime, + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, } => { if !self.inputs.frozen_inputs().is_some() || node.output_state - != planner_types::post_asap::ExecutionDataState::MAINTENANCE_ROWS + != planner_types::post_asap::ExecutionDataState::INGESTION_ROWS { return Err("maintenance binary requires immutable completed row inputs".into()); } @@ -183,7 +185,7 @@ impl PrecomputeOperatorRegistry for OperatorAdapter<'_> { ExecutableOperatorPayload::Value { operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, - timing: planner_types::post_asap::ExecutionTiming::MaintenanceTime, + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, } => { if !self.inputs.frozen_inputs().is_some() { return Err( @@ -2171,8 +2173,7 @@ mod tests { .nodes .iter() .filter(|node| { - node.output_state.timing - == planner_types::post_asap::ExecutionTiming::MaintenanceTime + node.output_state.timing == planner_types::post_asap::ExecutionTiming::IngestionTime }) .map(|node| node.id) .collect::>(); @@ -2249,8 +2250,10 @@ mod tests { fn node(id: u32) -> ExecutableDagNode { ExecutableDagNode { id: PostAsapNodeId(id), - payload: ExecutableOperatorPayload::SummaryMerge, - output_state: planner_types::post_asap::ExecutionDataState::MAINTENANCE_SUMMARY, + payload: ExecutableOperatorPayload::SummaryMerge { + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, + }, + output_state: planner_types::post_asap::ExecutionDataState::INGESTION_SUMMARY, output_schema: SummarySchema { fields: vec![], time_index: None, @@ -2268,7 +2271,7 @@ mod tests { fields: vec![], time_index: None, }, - data_state: planner_types::post_asap::ExecutionDataState::MAINTENANCE_SUMMARY, + data_state: planner_types::post_asap::ExecutionDataState::INGESTION_SUMMARY, grouping: GroupingEdgeCompatibility::Identical, window: WindowEdgeCompatibility::NotApplicable, } @@ -2444,7 +2447,7 @@ mod tests { let mut read = node(2); read.payload = ExecutableOperatorPayload::Value { operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, - timing: planner_types::post_asap::ExecutionTiming::MaintenanceTime, + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, }; read.output_schema.fields = vec![SummaryField { name: "value".into(), @@ -2489,14 +2492,14 @@ mod tests { dtype: SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum), nullable: false, }]; - read.output_state = planner_types::post_asap::ExecutionDataState::MAINTENANCE_ROWS; + read.output_state = planner_types::post_asap::ExecutionDataState::INGESTION_ROWS; aggregate.output_schema.fields = vec![SummaryField { name: "state".into(), dtype: configs[1].accumulator_spec().unwrap().family, nullable: false, }]; let mut query = node(4); - query.output_state = planner_types::post_asap::ExecutionDataState::READ_ROWS; + query.output_state = planner_types::post_asap::ExecutionDataState::QUERY_ROWS; let mut first_edge = edge(1, 2); first_edge.intermediate_schema = source_node.output_schema.clone(); let mut second_edge = edge(2, 3); @@ -2907,7 +2910,9 @@ mod tests { second_node.id = PostAsapNodeId(5); let mut merge = second_node.clone(); merge.id = PostAsapNodeId(6); - merge.payload = ExecutableOperatorPayload::SummaryMerge; + merge.payload = ExecutableOperatorPayload::SummaryMerge { + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, + }; dag.nodes.extend([second_node, merge]); let original = dag .edges @@ -3395,9 +3400,9 @@ mod tests { }; operation.payload = ExecutableOperatorPayload::Binary { operator: operator.clone(), - timing: planner_types::post_asap::ExecutionTiming::MaintenanceTime, + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, }; - operation.output_state = planner_types::post_asap::ExecutionDataState::MAINTENANCE_ROWS; + operation.output_state = planner_types::post_asap::ExecutionDataState::INGESTION_ROWS; assert!(frozen .execute(&operation, &[left.clone(), right.clone()]) .is_ok()); @@ -3414,7 +3419,7 @@ mod tests { .is_err()); operation.payload = ExecutableOperatorPayload::Binary { operator: operator.clone(), - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, }; assert!(frozen .execute(&operation, &[left.clone(), right.clone()]) @@ -3870,7 +3875,7 @@ mod tests { .policy_fingerprint() .into(); let mut query = node(2); - query.output_state = planner_types::post_asap::ExecutionDataState::READ_ROWS; + query.output_state = planner_types::post_asap::ExecutionDataState::QUERY_ROWS; let dag = ExecutableDag { nodes: vec![node(0), node(1), query], edges: vec![edge(0, 1), edge(1, 2)], @@ -4010,7 +4015,7 @@ mod tests { // source 0 is shared by both branches; root therefore contains two // copies of its value while node 0 itself is evaluated once. let mut query = node(4); - query.output_state = planner_types::post_asap::ExecutionDataState::READ_ROWS; + query.output_state = planner_types::post_asap::ExecutionDataState::QUERY_ROWS; let dag = ExecutableDag { nodes: (0..4).map(node).chain([query]).collect(), edges: vec![edge(0, 1), edge(0, 2), edge(1, 3), edge(2, 3), edge(3, 4)], @@ -4082,7 +4087,7 @@ mod tests { let mut unsupported = node(1); unsupported.payload = ExecutableOperatorPayload::SummarySubtract; let mut query = node(2); - query.output_state = planner_types::post_asap::ExecutionDataState::READ_ROWS; + query.output_state = planner_types::post_asap::ExecutionDataState::QUERY_ROWS; let dag = ExecutableDag { nodes: vec![node(0), unsupported, query], edges: vec![edge(0, 1), edge(1, 2)], diff --git a/data_plane/src/precompute_engine/subdag_scheduler.rs b/data_plane/src/precompute_engine/subdag_scheduler.rs index 8d20c07f5..a9d10fec3 100644 --- a/data_plane/src/precompute_engine/subdag_scheduler.rs +++ b/data_plane/src/precompute_engine/subdag_scheduler.rs @@ -153,7 +153,7 @@ where let node = nodes .get(&id) .ok_or_else(|| ScheduleError::Invalid(format!("missing node {id}")))?; - if node.output_state == ExecutionDataState::READ_ROWS { + if node.output_state == ExecutionDataState::QUERY_ROWS { return Err(ScheduleError::Invalid(format!( "query-time node {id} in precompute dependency path" ))); @@ -237,8 +237,7 @@ mod tests { .nodes .iter() .filter(|node| { - node.output_state.timing - == planner_types::post_asap::ExecutionTiming::MaintenanceTime + node.output_state.timing == planner_types::post_asap::ExecutionTiming::IngestionTime }) .map(|node| node.id) .collect::>(); @@ -254,7 +253,7 @@ mod tests { ExecutableDagNode { id: PostAsapNodeId(id), payload: ExecutableOperatorPayload::SummarySubtract, - output_state: ExecutionDataState::MAINTENANCE_SUMMARY, + output_state: ExecutionDataState::INGESTION_SUMMARY, output_schema: SummarySchema { fields: Vec::new(), time_index: None, @@ -272,7 +271,7 @@ mod tests { fields: Vec::new(), time_index: None, }, - data_state: ExecutionDataState::MAINTENANCE_SUMMARY, + data_state: ExecutionDataState::INGESTION_SUMMARY, grouping: GroupingEdgeCompatibility::Identical, window: WindowEdgeCompatibility::NotApplicable, } @@ -348,7 +347,7 @@ mod tests { } let mut binary = node(3); binary.payload = ExecutableOperatorPayload::Binary { - timing: planner_types::post_asap::ExecutionTiming::MaintenanceTime, + timing: planner_types::post_asap::ExecutionTiming::IngestionTime, operator: BinaryOperator { checked_relative_division: false, checked_finite_division: false, @@ -363,7 +362,7 @@ mod tests { let mut dag = ExecutableDag { nodes: vec![node(0), node(1), node(2), binary, { let mut query = node(4); - query.output_state = ExecutionDataState::READ_ROWS; + query.output_state = ExecutionDataState::QUERY_ROWS; query }], edges: vec![right, left], @@ -399,7 +398,7 @@ mod tests { .map(node) .chain([{ let mut query = node(4); - query.output_state = ExecutionDataState::READ_ROWS; + query.output_state = ExecutionDataState::QUERY_ROWS; query }]) .collect(), @@ -444,9 +443,9 @@ mod tests { } } let mut raw = node(0); - raw.output_state = ExecutionDataState::READ_ROWS; + raw.output_state = ExecutionDataState::QUERY_ROWS; let mut query = node(4); - query.output_state = ExecutionDataState::READ_ROWS; + query.output_state = ExecutionDataState::QUERY_ROWS; let dag = ExecutableDag { nodes: vec![raw, node(1), node(2), node(3), query], edges: vec![edge(0, 1), edge(1, 2), edge(1, 3), edge(2, 3), edge(3, 4)], @@ -472,7 +471,7 @@ mod tests { #[test] fn rejects_query_node_in_precompute_path_and_mismatched_lineage_key() { let mut query_child = node(0); - query_child.output_state = ExecutionDataState::READ_ROWS; + query_child.output_state = ExecutionDataState::QUERY_ROWS; let dag = ExecutableDag { nodes: vec![query_child, node(1)], edges: vec![edge(0, 1)], diff --git a/data_plane/src/precompute_engine/worker.rs b/data_plane/src/precompute_engine/worker.rs index c6aac060a..cda3524fd 100644 --- a/data_plane/src/precompute_engine/worker.rs +++ b/data_plane/src/precompute_engine/worker.rs @@ -4446,9 +4446,11 @@ mod dag_execution_tests { &dag, ) .unwrap(); - assert!(StreamingConfig::from_precompute_plan(plan) + installed.document.schema_version = + asap_types::executable_plan::MAINTENANCE_DAG_SCHEMA_VERSION; + let error = StreamingConfig::from_precompute_plan(plan) .unwrap_err() - .to_string() - .contains("update")); + .to_string(); + assert!(error.contains("update"), "{error}"); } } 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 cae8ff229..cc6407ee8 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 - // CandidateTopK reuses the same membership readout during composition. + // MembershipFilter 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 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 9a992ca25..0e1752951 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,9 @@ fn leaves( } _ => pending.extend(inputs.iter().map(|input| (*input, at))), }, - // CandidateTopK is a typed composition node rather than a Logical + // MembershipFilter is a typed composition node rather than a Logical // wrapper, but its value input can still be a Prometheus leaf. - QueryPlanNode::CandidateTopK { inputs, .. } => { + QueryPlanNode::MembershipFilter { inputs, .. } => { pending.extend(inputs.iter().map(|input| (*input, at))); } QueryPlanNode::ExternalExact { request, inputs } => { @@ -650,22 +650,30 @@ mod tests { } #[tokio::test] - async fn candidate_exact_is_discovered_and_prepared_behind_candidate_topk_root() { - use asap_types::query_plan::{residual::Grouping, CandidateCompleteness}; + async fn candidate_exact_is_discovered_and_prepared_behind_membership_filter_root() { + use asap_types::query_plan::CandidateCompleteness; let mut entry = candidate_entry("sum by (job) (rate(m[5m]))"); entry.nodes.insert( QueryNodeId(2), - QueryPlanNode::CandidateTopK { + QueryPlanNode::MembershipFilter { inputs: [QueryNodeId(1), QueryNodeId(0)], - k: 2, - grouping: Grouping { - labels: vec![], - without: false, - }, completeness: CandidateCompleteness::BestEffort { guarantee: None }, }, ); - entry.root = QueryNodeId(2); + entry.nodes.insert( + QueryNodeId(3), + QueryPlanNode::Logical { + operator: ResidualQueryOperator::TopKSelection { + k: 2, + grouping: asap_types::query_plan::residual::Grouping { + labels: vec![], + without: false, + }, + }, + inputs: vec![QueryNodeId(2)], + }, + ); + entry.root = QueryNodeId(3); let dependencies = external_dependencies(&entry, &[1_000]).unwrap(); assert_eq!(dependencies, vec![(QueryNodeId(0), QueryNodeId(1), 1_000)]); let prepared = prepare_external( 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 4b77a5a8f..a9e4d8460 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 @@ -205,16 +205,13 @@ impl Result> Evaluator<' } self.logical(operator, &inputs, at)? } - QueryPlanNode::CandidateTopK { + QueryPlanNode::MembershipFilter { inputs, - k, - grouping, completeness, } => { let candidates = vector(self.eval(inputs[0], at)?)?; let values = vector(self.eval(inputs[1], at)?)?; - let (selected, warning) = - candidate_topk(k, &grouping, candidates, values, &completeness)?; + let (selected, warning) = membership_filter(candidates, values, &completeness)?; if let Some(warning) = warning { self.warnings.push(warning); } @@ -423,9 +420,7 @@ impl Result> Evaluator<' } } -fn candidate_topk( - k: u64, - grouping: &Grouping, +fn membership_filter( candidates: Vector, values: Vector, completeness: &CandidateCompleteness, @@ -435,30 +430,22 @@ fn candidate_topk( 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 dangling = candidate_ids - .iter() - .any(|candidate| !value_ids.contains(candidate)); - if dangling && matches!(completeness, CandidateCompleteness::Certified { .. }) { - return Err(miss("certified TopK candidate has no exact counter value")); + let (selected, missing) = asap_physical_operators::rows::membership_filter( + candidates.iter().map(|(labels, _)| identity(labels)), + values, + |(labels, _)| identity(labels), + ); + if !missing.is_empty() && matches!(completeness, CandidateCompleteness::Certified { .. }) { + return Err(miss("certified membership key has no authoritative value")); } - let matched = values - .into_iter() - .filter(|(labels, _)| candidate_ids.contains(&identity(labels))) - .collect(); - let selected = topk_selection(k, grouping, matched); let warning = match completeness { CandidateCompleteness::Certified { .. } => None, CandidateCompleteness::BestEffort { guarantee } => Some(match guarantee { Some(guarantee) => format!( - "ASAP TopK candidate membership is approximate: {:?}", + "ASAP membership pruning is approximate: {:?}", guarantee.metric ), - None => "ASAP TopK candidate membership is approximate and uncertified".into(), + None => "ASAP membership pruning is approximate and uncertified".into(), }), }; Ok((selected, warning)) @@ -520,30 +507,12 @@ fn grouping_key(labels: &Labels, grouping: &Grouping) -> Labels { /// labels. NaN ranks below every numeric value, matching Prometheus' TOPK heap. /// Stable sorting also leaves equal-valued series in the child's order. 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() + asap_physical_operators::rows::grouped_topk( + values, + usize::try_from(k).unwrap_or(usize::MAX), + |(labels, _)| grouping_key(labels, grouping), + |(_, value)| *value, + ) } fn binary( @@ -908,6 +877,14 @@ mod topk_tests { (labels(&[("series", "high")]), 3.0), ], ); + let selected = topk_selection( + 2, + &Grouping { + labels: vec![], + without: false, + }, + selected, + ); assert_eq!( selected .iter() @@ -1177,12 +1154,7 @@ mod topk_tests { (labels(&[("pod", "b")]), 8.0), (labels(&[("pod", "c")]), 9.0), ]; - let (selected, warning) = candidate_topk( - 2, - &Grouping { - labels: vec![], - without: false, - }, + let (selected, warning) = membership_filter( candidates, exact, &CandidateCompleteness::Certified { @@ -1190,6 +1162,14 @@ mod topk_tests { }, ) .unwrap(); + let selected = topk_selection( + 2, + &Grouping { + labels: vec![], + without: false, + }, + selected, + ); assert_eq!( selected .iter() @@ -1204,7 +1184,8 @@ mod topk_tests { fn installed_candidate_sidecar_reads_both_summary_inputs() { let candidate_id = QueryNodeId(0); let value_id = QueryNodeId(1); - let root = QueryNodeId(2); + let filter = QueryNodeId(2); + let root = QueryNodeId(3); let entry = QueryPlanEntry { language: asap_types::query_plan::QueryLanguage::PromQl, query_id: "candidate-topk".into(), @@ -1225,19 +1206,27 @@ mod topk_tests { }, ), ( - root, - QueryPlanNode::CandidateTopK { + filter, + QueryPlanNode::MembershipFilter { inputs: [candidate_id, value_id], - k: 1, - grouping: Grouping { - labels: vec![], - without: false, - }, completeness: CandidateCompleteness::Certified { guarantee: topk_membership_guarantee(), }, }, ), + ( + root, + QueryPlanNode::Logical { + operator: ResidualQueryOperator::TopKSelection { + k: 1, + grouping: Grouping { + labels: vec![], + without: false, + }, + }, + inputs: vec![filter], + }, + ), ]), instant: InstantExecution { lookback_ms: 300_000, @@ -1251,7 +1240,10 @@ mod topk_tests { ( (candidate_id, at), PreparedLeaf { - value: Value::Vector(vec![(labels(&[("pod", "b")]), 100.0)]), + value: Value::Vector(vec![ + (labels(&[("pod", "b")]), 100.0), + (labels(&[("pod", "c")]), 1.0), + ]), remote: false, remote_evaluations: 0, remote_rpcs: 0, @@ -1263,6 +1255,7 @@ mod topk_tests { value: Value::Vector(vec![ (labels(&[("pod", "a")]), 2.0), (labels(&[("pod", "b")]), 1.0), + (labels(&[("pod", "c")]), 3.0), ]), remote: false, remote_evaluations: 0, @@ -1278,8 +1271,8 @@ mod topk_tests { panic!("vector expected") }; assert_eq!(result.values.len(), 1); - assert_eq!(result.values[0].value, 1.0, "exact value is authoritative"); - assert_eq!(result.values[0].labels.labels, vec!["b"]); + assert_eq!(result.values[0].value, 3.0, "exact value is authoritative"); + assert_eq!(result.values[0].labels.labels, vec!["c"]); assert_eq!(stats.summary_readout_evaluations, 2); assert!(result.warnings.is_empty()); } @@ -1288,30 +1281,20 @@ 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) = candidate_topk( - 1, - &Grouping { - labels: vec![], - without: false, - }, + let (_, warning) = membership_filter( candidates.clone(), exact.clone(), &CandidateCompleteness::BestEffort { guarantee: None }, ) .unwrap(); assert!(warning.unwrap().contains("approximate")); - // Exact queries never lower an uncertified CandidateTopK. The Planner + // Exact queries never lower an uncertified MembershipFilter. 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!(candidate_topk( - 1, - &Grouping { - labels: vec![], - without: false - }, + assert!(membership_filter( vec![(labels(&[("pod", "missing")]), 1.0)], exact, &certified, 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 73d234ee5..1da37ae39 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 @@ -305,7 +305,7 @@ impl QueryNodeRuntime for PhysicalQueryRuntime<'_> { }) } QueryPlanNode::Logical { .. } - | QueryPlanNode::CandidateTopK { .. } + | QueryPlanNode::MembershipFilter { .. } | QueryPlanNode::Relational { .. } | QueryPlanNode::ExternalExact { .. } | QueryPlanNode::RelationalJoin { .. } => Err(PhysicalNodeError::Fallback( 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 fb75ae6bc..02c2adac6 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 @@ -184,7 +184,7 @@ pub fn execute( Ok(ExecOutcome::Value(out)) } - SummaryExpr::SummaryMerge { children } => { + SummaryExpr::SummaryMerge { children, .. } => { if children.is_empty() { return Err(ExecError::EmptyMerge); } @@ -223,11 +223,11 @@ pub fn execute( SummaryExpr::SummaryJoin { .. } => Err(ExecError::NotYetSupported("SummaryJoin")), SummaryExpr::RelationalJoin { .. } => Err(ExecError::NotYetSupported("RelationalJoin")), - // CandidateTopK is lowered to the deployed QueryPlan DAG, where both + // MembershipFilter 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. - SummaryExpr::CandidateTopK { .. } => Err(ExecError::NotYetSupported("CandidateTopK")), + SummaryExpr::MembershipFilter { .. } => Err(ExecError::NotYetSupported("MembershipFilter")), SummaryExpr::BinaryOp { .. } => Err(ExecError::NotYetSupported("BinaryOp")), SummaryExpr::ValueOperation { .. } => Err(ExecError::NotYetSupported("ValueOperation")), SummaryExpr::SummarySubtract { .. } => Err(ExecError::NotYetSupported("SummarySubtract")), @@ -350,7 +350,10 @@ mod tests { fn merge_node(children: Vec>) -> Rc { Rc::new(SummaryNode { - expr: SummaryExpr::SummaryMerge { children }, + expr: SummaryExpr::SummaryMerge { + children, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, + }, schema: lift(vec!["value"]), guarantee: None, }) @@ -522,7 +525,7 @@ mod tests { let child = logical_node(); let tree = SummaryNode { expr: SummaryExpr::BinaryOp { - timing: planner_types::post_asap::ExecutionTiming::ReadTime, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, lhs: child.clone(), rhs: child.clone(), operator: planner_types::post_asap::BinaryOperator { 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 058cd8d82..dfcf32f6b 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 @@ -1389,7 +1389,7 @@ fn find_metric(node: &SummaryNode) -> Option { SummaryExpr::KeepPreAsap(qe) => find_metric_in_query_expr(qe), SummaryExpr::SummaryAgg { child, .. } => find_metric(child), SummaryExpr::SummaryEstimate { summary_input, .. } => find_metric(summary_input), - SummaryExpr::SummaryMerge { children } => children.first().and_then(|c| find_metric(c)), + SummaryExpr::SummaryMerge { children, .. } => children.first().and_then(|c| find_metric(c)), _ => None, } } diff --git a/docs/design_docs/physical-operators.md b/docs/design_docs/physical-operators.md new file mode 100644 index 000000000..02d321a19 --- /dev/null +++ b/docs/design_docs/physical-operators.md @@ -0,0 +1,80 @@ +# Shared physical operators and DAG execution + +## Decision + +ASAP owns an independent physical operator library and DAG runtime. Precompute +and query engines bind inputs and consume outputs from the same library. The +operator defines computation; the engine supplies ingestion time or query time, +window boundaries, storage access and publication. There is no second execution +algorithm selected by phase. + +This foundation precedes the precompute integration in #763 and query integration +in #765. Its native operators can execute independently of either backend engine. +Engine integration must use these operators for computation, rather than merely +using the shared scheduler around a second implementation. + +## Execution contract + +An immutable plan describes typed nodes and dependency edges. Each execution +creates its own operator state. One producer may have multiple consumers; the +producer executes once in that run and sends the same outputs to all consumers. +Separate runs, query evaluation times and ingestion windows do not share mutable +state. Request-local caching of intermediate results is scoped to execution. + +The runtime validates dependencies, schemas, arity and cycles before sources +start. Each consumer advances independently. Bounded queues apply backpressure; +dropping one consumer does not cancel other consumers. Whole-run cancellation +wakes readers and releases queued work as streams are polled or dropped. + +Execution runs on the caller's worker without an internal thread pool. Active +streams are worker-local. Deployments poll all consumers concurrently. The byte +budget accounts for retained outputs and native operator state, including outputs +held after queue eviction. It is not an RSS limit: source-owned data, temporary +allocation peaks and allocator overhead remain outside that estimate. Blocking +operators currently have no spill implementation. + +## Operator coverage + +Native operations include scalar sources, typed Project and Filter, arithmetic +and boolean expressions, exact grouped aggregation, semi-join, grouped Sort and +Limit, Union, vector-to-scalar conversion, and summary construction, merge and +readout. Grouped TopK composes Sort and Limit within each group; candidate +completeness is an earlier pruning obligation. + +Values retain Planner types and nullability. Native summary batches currently +support exact Sum/Count/Min/Max/Rate/Increase, KLL, DDSketch and HLL. Other available +low-level kernels do not imply native batch bindings. Unsupported expressions, +state families and parameters must be rejected during binding, without an +implicit external fallback. The installed engine adapters and remaining gaps +are tracked in the query DAG design's unified coverage table. + +Deployments provide explicit storage or ingestion source frontiers. A supplied +batch source is not a backend raw Scan implementation. Local backend raw Scan +is deferred; a raw-only library test does not establish that deployment capability. + +## DataFusion reuse vs independent implementation + +| Decision dimension | Reuse DataFusion | Independent ASAP implementation | +| --- | --- | --- | +| General computation | Reuse mature Arrow operators and expression execution | Implement and test the supported Planner vocabulary explicitly | +| Shared DAG producer | Shared plan references need an explicit execution-sharing and buffering policy | One producer and independent consumer cursors are part of the runtime contract | +| Summary lifecycle | Add custom summary state operators to the framework | Summary construction, merge and readout are native capabilities | +| Engine reuse | Adapt both engines to DataFusion's execution model | Both engines bind the same ASAP interfaces | +| Engineering cost | Less generic operator work; integration and semantic adaptation remain | More operator, typing, scheduling and resource-accounting responsibility | + +DataFusion is a design reference, not this library's execution dependency. This +choice does not claim that DataFusion cannot express shared dependencies. ASAP +chooses direct ownership of execution sharing and summary-state semantics across +both engines. Mathematical sketch kernels remain reusable implementation details. + +## Acceptance + +Independent tests must execute shared-producer diamonds without duplicated work +or deadlock, exercise slow and dropped consumers, propagate cancellation and +errors, retain memory accounting, and isolate separate executions. Operator tests +must cover types, nulls, grouped limits, state compatibility and unsupported +bindings. The same summary pipeline must run at ingestion time and query time. + +#763 and #765 add deployment acceptance for source binding, window and revision +scope, durable publication and query output adaptation. External exact forwarding +does not count as evidence that a local operator was implemented. diff --git a/tools/o11y-execution/calibrate_runtime.py b/tools/o11y-execution/calibrate_runtime.py index 5587dca6e..7b263db7c 100644 --- a/tools/o11y-execution/calibrate_runtime.py +++ b/tools/o11y-execution/calibrate_runtime.py @@ -212,7 +212,7 @@ def launch(name, command): query_phase = phase(folder, "query-" + qid, before, after, time.perf_counter_ns() - start) raw = folder / f"queries-{qid}.json" runner.write_json(raw, records) - validate_candidate_topk_execution(artifact, records) + validate_membership_filter_execution(artifact, records) routes = {record["execution"] for record in records} correct = all(record["comparison"]["equal"] and record["exact"]["http_status"] == 200 for record in records) row["queries"][qid] = {"cpu_ns": query_phase["cpu_ns"], "evaluations": len(records), @@ -267,20 +267,20 @@ def launch(name, command): -def _candidate_topk_inputs(nodes, root): +def _membership_filter_inputs(nodes, root): bindings, visiting, visited = set(), set(), set() def visit(node_id): node_id = str(node_id) if node_id in visiting: - raise ValueError("CandidateTopK input DAG contains a cycle") + raise ValueError("MembershipFilter input DAG contains a cycle") if node_id in visited: return if node_id not in nodes: - raise ValueError(f"CandidateTopK input DAG references missing node {node_id}") + raise ValueError(f"MembershipFilter input DAG references missing node {node_id}") visiting.add(node_id) node = nodes[node_id] if node.get("op") == "exact_fallback": - raise ValueError("CandidateTopK input contains ExactFallback") + raise ValueError("MembershipFilter input contains ExactFallback") if node.get("op") == "read_materialization": bindings.add(str(node["binding"]["materialization"])) children = [str(value) for value in node.get("inputs", [])] @@ -294,21 +294,21 @@ def visit(node_id): return bindings -def validate_candidate_topk_artifact(artifact): - """Reject CandidateTopK plans whose membership sidecar is not locally installed.""" +def validate_membership_filter_artifact(artifact): + """Reject MembershipFilter plans whose membership sidecar is not locally installed.""" request = artifact.get("install_request", {}) schemas = {str(row["materialization"]): row for row in request.get("precompute_plan", {}).get("schemas", [])} modes = set() for entry in request.get("query_plan", {}).get("entries", {}).values(): nodes = entry.get("nodes", {}) for node in nodes.values(): - if node.get("op") != "candidate_top_k": + if node.get("op") != "membership_filter": continue inputs = node.get("inputs", []) if len(inputs) != 2: - raise ValueError("CandidateTopK requires membership and exact-value inputs") - membership_bindings = _candidate_topk_inputs(nodes, inputs[0]) - value_bindings = _candidate_topk_inputs(nodes, inputs[1]) + raise ValueError("MembershipFilter requires membership and exact-value inputs") + membership_bindings = _membership_filter_inputs(nodes, inputs[0]) + value_bindings = _membership_filter_inputs(nodes, inputs[1]) heap_bindings = [] for materialization in membership_bindings: schema = schemas.get(materialization) @@ -316,7 +316,7 @@ def validate_candidate_topk_artifact(artifact): if "CmsWithHeap" in family or "CountSketchWithHeap" in family: heap_bindings.append(materialization) if not heap_bindings: - raise ValueError("CandidateTopK membership input has no installed heap materialization") + raise ValueError("MembershipFilter membership input has no installed heap materialization") value_node = nodes.get(str(inputs[1]), {}) operator = value_node.get("operator", {}) if value_node.get("op") == "logical" else {} if operator.get("kind") == "candidate_exact_subquery": @@ -339,21 +339,21 @@ def validate_candidate_topk_artifact(artifact): and any(kind in json.dumps((schemas.get(mid) or {}).get("family", {})).lower() for kind in ("counter", "rate", "increase")) for mid in value_bindings): - raise ValueError("CandidateTopK value input has no installed ExactCounter materialization") + raise ValueError("MembershipFilter value input has no installed ExactCounter materialization") return modes -def validate_candidate_topk_execution(artifact, records): - modes = validate_candidate_topk_artifact(artifact) +def validate_membership_filter_execution(artifact, records): + modes = validate_membership_filter_artifact(artifact) if not modes: return if len(modes) != 1: - raise ValueError("mixed CandidateTopK execution contracts are not calibratable together") + raise ValueError("mixed MembershipFilter execution contracts are not calibratable together") mode = next(iter(modes)) for record in records: provenance = record.get("execution_provenance", {}) if provenance.get("raw_scan_evaluations", 0) not in (0, None): - raise ValueError("CandidateTopK execution used a forbidden local raw scan") + raise ValueError("MembershipFilter execution used a forbidden local raw scan") if mode == "candidate_filtered_exact": if record.get("execution") != "hybrid" or provenance.get("detail") != "hybrid": raise ValueError("candidate-filtered TopK did not report hybrid execution") @@ -364,12 +364,12 @@ def validate_candidate_topk_execution(artifact, records): raise ValueError(f"candidate-filtered TopK has invalid provenance: {key}") else: if record.get("execution") != "warm" or provenance.get("detail") not in (None, "asap"): - raise ValueError("local CandidateTopK execution was not warm") + raise ValueError("local MembershipFilter execution was not warm") for key in ("exact_subquery_rpcs", "exact_subquery_evaluations", "exact_branch_evaluations"): if provenance.get(key, 0) != 0: - raise ValueError(f"CandidateTopK execution used exact path: {key}") + raise ValueError(f"MembershipFilter execution used exact path: {key}") if provenance.get("summary_readout_evaluations", 0) < 2: - raise ValueError("CandidateTopK execution did not read both summary branches") + raise ValueError("MembershipFilter execution did not read both summary branches") def main(): @@ -413,7 +413,7 @@ def main(): candidates = candidate_document["candidates"] for candidate in candidates: if "manifest" in candidate and "install_request" in candidate: - validate_candidate_topk_artifact(candidate) + validate_membership_filter_artifact(candidate) for index, candidate in enumerate(candidates): if "manifest" not in candidate or "install_request" not in candidate: continue diff --git a/tools/o11y-execution/test_calibrate_runtime.py b/tools/o11y-execution/test_calibrate_runtime.py index 74a0931dc..891829f36 100644 --- a/tools/o11y-execution/test_calibrate_runtime.py +++ b/tools/o11y-execution/test_calibrate_runtime.py @@ -1,14 +1,14 @@ -"""Fail-closed validation for calibrated CandidateTopK artifacts.""" +"""Fail-closed validation for calibrated MembershipFilter artifacts.""" import unittest -from calibrate_runtime import validate_candidate_topk_artifact, validate_candidate_topk_execution +from calibrate_runtime import validate_membership_filter_artifact, validate_membership_filter_execution -class CandidateTopKArtifactTests(unittest.TestCase): +class MembershipFilterArtifactTests(unittest.TestCase): def artifact(self, membership): return {"install_request": { "query_plan": {"entries": {"q": {"nodes": { - "0": {"op": "candidate_top_k", "inputs": [1, 3]}, + "0": {"op": "membership_filter", "inputs": [1, 3]}, "1": {"op": "summary_estimate", "input": 2}, "2": membership, "3": {"op": "exact_readout", "input": 4}, @@ -35,20 +35,20 @@ def candidate_filtered_artifact(self): def test_rejects_exact_membership_fallback(self): with self.assertRaisesRegex(ValueError, "contains ExactFallback"): - validate_candidate_topk_artifact(self.artifact({"op": "exact_fallback", "reason": "unsupported"})) + validate_membership_filter_artifact(self.artifact({"op": "exact_fallback", "reason": "unsupported"})) def test_rejects_uninstalled_heap_membership(self): artifact = self.artifact({"op": "read_materialization", "binding": {"materialization": 9}}) with self.assertRaisesRegex(ValueError, "no installed heap"): - validate_candidate_topk_artifact(artifact) + validate_membership_filter_artifact(artifact) def test_accepts_heap_membership_and_exact_values(self): - validate_candidate_topk_artifact( + validate_membership_filter_artifact( self.artifact({"op": "read_materialization", "binding": {"materialization": 7}}) ) - def test_ignores_plans_without_candidate_topk(self): - validate_candidate_topk_artifact({"install_request": { + def test_ignores_plans_without_membership_filter(self): + validate_membership_filter_artifact({"install_request": { "query_plan": {"entries": {"q": {"nodes": {"0": {"op": "exact_fallback"}}}}}, "precompute_plan": {"schemas": []}, }}) @@ -57,56 +57,56 @@ def test_rejects_exact_value_fallback(self): artifact = self.artifact({"op": "read_materialization", "binding": {"materialization": 7}}) artifact["install_request"]["query_plan"]["entries"]["q"]["nodes"]["3"] = {"op": "exact_fallback"} with self.assertRaisesRegex(ValueError, "contains ExactFallback"): - validate_candidate_topk_artifact(artifact) + validate_membership_filter_artifact(artifact) def test_rejects_missing_or_cyclic_input_nodes(self): artifact = self.artifact({"op": "summary_estimate", "input": 99}) with self.assertRaisesRegex(ValueError, "missing node 99"): - validate_candidate_topk_artifact(artifact) + validate_membership_filter_artifact(artifact) artifact = self.artifact({"op": "summary_estimate", "input": 2}) with self.assertRaisesRegex(ValueError, "contains a cycle"): - validate_candidate_topk_artifact(artifact) + validate_membership_filter_artifact(artifact) def test_requires_two_local_summary_readouts_at_runtime(self): artifact = self.artifact({"op": "read_materialization", "binding": {"materialization": 7}}) provenance = {"summary_readout_evaluations": 2, "exact_subquery_rpcs": 0, "exact_subquery_evaluations": 0, "exact_branch_evaluations": 0} - validate_candidate_topk_execution(artifact, [{"execution": "warm", "execution_provenance": provenance}]) + validate_membership_filter_execution(artifact, [{"execution": "warm", "execution_provenance": provenance}]) with self.assertRaisesRegex(ValueError, "both summary branches"): - validate_candidate_topk_execution(artifact, [{"execution": "warm", "execution_provenance": { + validate_membership_filter_execution(artifact, [{"execution": "warm", "execution_provenance": { **provenance, "summary_readout_evaluations": 1}}]) with self.assertRaisesRegex(ValueError, "used exact path"): - validate_candidate_topk_execution(artifact, [{"execution": "warm", "execution_provenance": { + validate_membership_filter_execution(artifact, [{"execution": "warm", "execution_provenance": { **provenance, "exact_subquery_rpcs": 1}}]) def test_accepts_one_heap_and_candidate_filtered_external_exact(self): - validate_candidate_topk_artifact(self.candidate_filtered_artifact()) + validate_membership_filter_artifact(self.candidate_filtered_artifact()) def test_candidate_filtered_contract_rejects_local_exact_state_or_unshared_input(self): artifact = self.candidate_filtered_artifact() artifact["install_request"]["precompute_plan"]["schemas"].append( {"materialization": 8, "family": {"family": "exact", "kind": "increase"}}) with self.assertRaisesRegex(ValueError, "must not install"): - validate_candidate_topk_artifact(artifact) + validate_membership_filter_artifact(artifact) artifact = self.candidate_filtered_artifact() artifact["install_request"]["query_plan"]["entries"]["q"]["nodes"]["3"]["inputs"] = [2] with self.assertRaisesRegex(ValueError, "shared membership"): - validate_candidate_topk_artifact(artifact) + validate_membership_filter_artifact(artifact) def test_candidate_filtered_execution_requires_hybrid_one_rpc_and_one_summary_read(self): artifact = self.candidate_filtered_artifact() provenance = {"detail": "hybrid", "raw_scan_evaluations": 0, "summary_readout_evaluations": 1, "exact_subquery_rpcs": 1, "exact_subquery_evaluations": 1, "exact_branch_evaluations": 1} - validate_candidate_topk_execution( + validate_membership_filter_execution( artifact, [{"execution": "hybrid", "execution_provenance": provenance}]) for key in ("summary_readout_evaluations", "exact_subquery_rpcs", "exact_subquery_evaluations", "exact_branch_evaluations"): with self.subTest(key=key), self.assertRaisesRegex(ValueError, "invalid provenance"): - validate_candidate_topk_execution(artifact, [{"execution": "hybrid", + validate_membership_filter_execution(artifact, [{"execution": "hybrid", "execution_provenance": {**provenance, key: 0}}]) with self.assertRaisesRegex(ValueError, "hybrid execution"): - validate_candidate_topk_execution( + validate_membership_filter_execution( artifact, [{"execution": "hybrid", "execution_provenance": { **provenance, "detail": "external_exact"}}]) From 425fce6745e611e36c5fca69774c5a7d67a4ae59 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 19:03:06 +0000 Subject: [PATCH 02/15] Use shared kernels exclusively and expose scoped native batch execution --- .../src/dag/batch_execution.rs | 112 + crates/asap-physical-operators/src/dag/mod.rs | 2 + .../asap-physical-operators/src/dag/values.rs | 17 +- data_plane/Cargo.toml | 2 +- data_plane/benches/sketch_db.rs | 2 +- data_plane/examples/univmon_erp_artifact.rs | 4 +- data_plane/src/drivers/ingest/otel.rs | 22 +- data_plane/src/lib.rs | 2 +- .../precompute_engine/accumulator_factory.rs | 2044 ----------------- .../src/precompute_engine/ingest_handler.rs | 4 +- .../precompute_engine/maintenance_runtime.rs | 8 +- data_plane/src/precompute_engine/mod.rs | 2 - .../operators/count_min_sketch_accumulator.rs | 1323 ----------- .../count_min_sketch_with_heap_accumulator.rs | 832 ------- .../operators/count_sketch_accumulator.rs | 686 ------ .../count_sketch_with_heap_accumulator.rs | 575 ----- .../operators/datasketches_kll_accumulator.rs | 733 ------ .../operators/dd_sketch_accumulator.rs | 667 ------ .../operators/exact_accumulator.rs | 327 --- .../operators/hll_sketch_accumulator.rs | 790 ------- .../operators/hydra_kll_accumulator.rs | 168 -- .../operators/increase_accumulator.rs | 742 ------ .../operators/keyed_counter_state.rs | 529 ----- .../operators/keyed_max_state.rs | 335 --- .../operators/keyed_min_state.rs | 335 --- .../operators/keyed_sum_count_accumulator.rs | 558 ----- .../operators/max_accumulator.rs | 248 -- .../operators/min_accumulator.rs | 253 -- .../src/precompute_engine/operators/mod.rs | 37 - .../operators/sketch_envelope_accumulator.rs | 156 -- .../operators/sum_accumulator.rs | 413 ---- .../operators/univmon_accumulator.rs | 236 -- .../src/precompute_engine/output_sink.rs | 2 +- data_plane/src/precompute_engine/raw_dag.rs | 2 +- data_plane/src/precompute_engine/worker.rs | 20 +- .../accelerator.rs | 6 +- .../query_engines/asap_query_engine/engine.rs | 6 +- .../asap_query_engine/exact_subqueries.rs | 2 +- .../asap_query_engine/live_serve.rs | 2 +- .../asap_query_engine/logical_dag.rs | 2 +- .../asap_query_engine/post_asap_readout.rs | 14 +- .../asap_query_engine/summary_executor.rs | 22 +- .../sketch_db/backfill/processor.rs | 4 +- .../sketch_db/backfill/window_builder.rs | 8 +- .../sketch_db/index/maintenance.rs | 4 +- .../storage_engines/sketch_db/index/mod.rs | 24 +- .../sketch_db/lifecycle/eviction.rs | 2 +- .../sketch_db/query/decoders.rs | 2 +- .../sketch_db/query/delta_apply.rs | 20 +- .../types/key_by_label_values.rs | 164 -- .../src/storage_engines/types/measurement.rs | 94 - data_plane/src/storage_engines/types/mod.rs | 9 +- .../src/storage_engines/types/traits.rs | 351 --- data_plane/src/tests/accumulator_fixture.rs | 276 +++ data_plane/src/tests/mod.rs | 2 + data_plane/src/tests/trait_design_tests.rs | 2 +- data_plane/src/utils/arithmetic.rs | 19 - data_plane/src/utils/mod.rs | 1 - data_plane/tests/edge_sketch_codec.rs | 4 +- .../tests/support/univmon_erp_process.rs | 2 +- 60 files changed, 504 insertions(+), 12726 deletions(-) create mode 100644 crates/asap-physical-operators/src/dag/batch_execution.rs delete mode 100644 data_plane/src/precompute_engine/accumulator_factory.rs delete mode 100644 data_plane/src/precompute_engine/operators/count_min_sketch_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/count_min_sketch_with_heap_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/count_sketch_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/count_sketch_with_heap_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/datasketches_kll_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/dd_sketch_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/exact_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/hll_sketch_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/hydra_kll_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/increase_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/keyed_counter_state.rs delete mode 100644 data_plane/src/precompute_engine/operators/keyed_max_state.rs delete mode 100644 data_plane/src/precompute_engine/operators/keyed_min_state.rs delete mode 100644 data_plane/src/precompute_engine/operators/keyed_sum_count_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/max_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/min_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/mod.rs delete mode 100644 data_plane/src/precompute_engine/operators/sketch_envelope_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/sum_accumulator.rs delete mode 100644 data_plane/src/precompute_engine/operators/univmon_accumulator.rs delete mode 100644 data_plane/src/storage_engines/types/key_by_label_values.rs delete mode 100644 data_plane/src/storage_engines/types/measurement.rs delete mode 100644 data_plane/src/storage_engines/types/traits.rs create mode 100644 data_plane/src/tests/accumulator_fixture.rs delete mode 100644 data_plane/src/utils/arithmetic.rs diff --git a/crates/asap-physical-operators/src/dag/batch_execution.rs b/crates/asap-physical-operators/src/dag/batch_execution.rs new file mode 100644 index 000000000..d29a9b0cc --- /dev/null +++ b/crates/asap-physical-operators/src/dag/batch_execution.rs @@ -0,0 +1,112 @@ +//! Execute a bounded in-memory batch through native operators. This is also the +//! bridge for deployments whose boundary values are not yet streaming batches. +use super::{operators::Operator, values::Batch, Error, PhysicalDag, RunContext}; +use futures::{FutureExt, StreamExt}; + +/// Every input is already in memory; the chain contains native operators only. +/// This deliberately does not enter a nested executor when called from a DAG +/// adapter. I/O belongs to source operators in the surrounding execution. +pub fn evaluate_batch( + input: Batch, + operators: Vec, + context: RunContext, +) -> Result, Error> { + let mut graph = PhysicalDag::default(); + graph.add( + 0, + vec![], + Operator::source(input.schema().clone(), vec![input])?, + )?; + let mut root = 0; + for operator in operators { + graph.add(root + 1, vec![root], operator)?; + root += 1; + } + let mut output = graph.execute(&[root], context)?.remove(0); + let mut batches = Vec::new(); + loop { + match output.next().now_or_never() { + Some(Some(Ok(batch))) => batches.push(batch.value().clone()), + Some(Some(Err(error))) => return Err(error), + Some(None) => return Ok(batches), + None => { + return Err(Error::Operator( + "in-memory native batch chain unexpectedly awaited I/O".into(), + )) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dag::{operators::Expression, values::Value, Limits, Scope}; + use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, + }; + use std::sync::Arc; + + // Engine adapters can run the identical native chain from an outer executor. + #[test] + fn same_native_chain_inside_query_and_ingestion_execution() { + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }], + time_index: None, + }); + for scope in [ + Scope::Query { + evaluation_time_ms: 20, + revision: 1, + }, + Scope::Ingestion { + window_start_ms: 10, + window_end_ms: 20, + revision: 1, + }, + ] { + let batch = Batch::try_new(schema.clone(), vec![vec![Value::Float64(7.)]]).unwrap(); + let negate = Operator::project( + schema.clone(), + vec![( + "value".into(), + Expression::Negate(Box::new(Expression::Column(0))), + )], + ) + .unwrap(); + let context = RunContext::new(scope, Limits::default()).unwrap(); + let result = + futures::executor::block_on(async { evaluate_batch(batch, vec![negate], context) }) + .unwrap(); + assert!(matches!(result[0].rows()[0][0], Value::Float64(-7.))); + } + } + + // A cancelled surrounding execution also prevents its native computation. + #[test] + fn cancellation_is_not_bypassed_by_in_memory_execution() { + let schema = Arc::new(SummarySchema { + fields: vec![], + time_index: None, + }); + let batch = Batch::try_new(schema, vec![vec![]]).unwrap(); + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + context.cancel(); + assert!(matches!( + evaluate_batch(batch, vec![], context), + Err(Error::Cancelled) + )); + } +} diff --git a/crates/asap-physical-operators/src/dag/mod.rs b/crates/asap-physical-operators/src/dag/mod.rs index b5a2e7b7d..1c9ee9e86 100644 --- a/crates/asap-physical-operators/src/dag/mod.rs +++ b/crates/asap-physical-operators/src/dag/mod.rs @@ -513,3 +513,5 @@ pub mod values; mod tests; pub mod planner; + +pub mod batch_execution; diff --git a/crates/asap-physical-operators/src/dag/values.rs b/crates/asap-physical-operators/src/dag/values.rs index 5b4e043f7..82b873143 100644 --- a/crates/asap-physical-operators/src/dag/values.rs +++ b/crates/asap-physical-operators/src/dag/values.rs @@ -263,10 +263,19 @@ fn validate_state(family: &SummaryFamilyType, state: &dyn AggregateCore) -> Resu use planner_types::post_asap::SketchParams; validate_family(family)?; let valid = match family { - SummaryFamilyType::ExactAggregate(..) => state - .as_any() - .downcast_ref::() - .is_some_and(|s| s.family() == family && !s.is_keyed()), + SummaryFamilyType::ExactAggregate(..) => { + state + .as_any() + .downcast_ref::() + .is_some_and(|s| s.family() == family && !s.is_keyed()) + || (matches!( + family, + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Sum, + planner_types::post_asap::ExactParams::Sum + ) + ) && state.as_any().is::()) + } SummaryFamilyType::Sketch(kind, _) => match kind.params() { SketchParams::Kll { k } => state .as_any() diff --git a/data_plane/Cargo.toml b/data_plane/Cargo.toml index 3f999f0e3..5e9c27b28 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -104,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..6462d255a 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::accumulators::SumAccumulator; use data_plane::storage_engines::sketch_db::data::{ AccuracyBound, AggKind, AggregationType, Capability, SketchAlgorithm, SketchConfig, SketchEncoding, diff --git a/data_plane/examples/univmon_erp_artifact.rs b/data_plane/examples/univmon_erp_artifact.rs index aa37ec137..109103068 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::accumulators::hll_sketch_accumulator::HllSketchAccumulator; +use asap_physical_operators::accumulators::univmon_accumulator::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 0d9a93048..cdfb1836b 100644 --- a/data_plane/src/drivers/ingest/otel.rs +++ b/data_plane/src/drivers/ingest/otel.rs @@ -24,7 +24,7 @@ use std::collections::HashMap; use std::io::Read; -use crate::precompute_engine::operators::sketch_envelope_accumulator::SketchEnvelopeAccumulator; +use asap_physical_operators::accumulators::sketch_envelope_accumulator::SketchEnvelopeAccumulator; use crate::precompute_engine::series_router::WorkerMessage; use crate::precompute_engine::IngestState; use crate::query_engines::routing::FreshnessProbeCache; @@ -2120,7 +2120,7 @@ fn dp_carries_heap(dp: &ModifiedOtlpSketchDp) -> bool { .unwrap_or(false) } ENCODING_MSGPACK_DELTA => { - use crate::precompute_engine::operators::CountMinSketchWithHeapAccumulator; + use asap_physical_operators::accumulators::CountMinSketchWithHeapAccumulator; CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes(&dp.sketch) .map(|acc| !acc.inner.topk_heap_items().is_empty()) .unwrap_or(false) @@ -2548,7 +2548,7 @@ fn decode_modified_otlp_sketch_bytes( encoding: i32, bytes: &[u8], ) -> Result, Box> { - use crate::precompute_engine::operators::{ + use asap_physical_operators::accumulators::{ CountMinSketchAccumulator, CountSketchAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, HllSketchAccumulator, }; @@ -2619,7 +2619,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::accumulators::CountSketchWithHeapAccumulator; return Ok(Box::new( CountSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(bytes)?, )); @@ -2684,7 +2684,7 @@ fn empty_accumulator_for_delta_bootstrap( config: &crate::storage_engines::sketch_db::index::SketchConfig, encoding: i32, ) -> Option> { - use crate::precompute_engine::operators::{ + use asap_physical_operators::accumulators::{ CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, HllSketchAccumulator, }; @@ -2758,7 +2758,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::accumulators::{ CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, DDSketchAccumulator, HllSketchAccumulator, }; @@ -3004,7 +3004,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::accumulators::sum_accumulator::SumAccumulator::from_sum_bytes(&dp.sketch) { Ok(acc) => acc.sum, Err(e) => { debug!("asap_edge: SumAgg data point decode failed (skipping): {e}"); @@ -3444,7 +3444,7 @@ mod policy_fp_lookup_tests { #[cfg(test)] mod dispatcher_tests { use super::*; - use crate::precompute_engine::operators::{DDSketchAccumulator, HllSketchAccumulator}; + use asap_physical_operators::accumulators::{DDSketchAccumulator, HllSketchAccumulator}; use crate::storage_engines::types::AggregateCore; use asap_sketchlib::DdSketch; use asap_sketchlib::HllVariant; @@ -3797,7 +3797,7 @@ 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_physical_operators::accumulators::DDSketchAccumulator; use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; use prost::Message; @@ -4129,7 +4129,7 @@ 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_physical_operators::accumulators::CountMinSketchAccumulator; use asap_otel_proto::sketchlib::v1::CountMinDelta as PbDelta; use prost::Message; @@ -4215,7 +4215,7 @@ 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_physical_operators::accumulators::HllSketchAccumulator; use asap_otel_proto::sketchlib::v1::HllDelta as PbDelta; use prost::Message; diff --git a/data_plane/src/lib.rs b/data_plane/src/lib.rs index 8c3ac483c..fd8b50149 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::accumulators::{ IncreaseAccumulator, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, }; 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 9f94c8c93..000000000 --- a/data_plane/src/precompute_engine/accumulator_factory.rs +++ /dev/null @@ -1,2044 +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() { - use planner_types::post_asap::{ - EntityIdentity, NonNegativeWeightProof, SummaryInputExpr, SummaryUpdate, WeightDomain, - }; - let config = topk_config(AggregationType::CountMinSketchWithHeap, None); - let family = config.accumulator_spec().unwrap().family; - let input = SummaryUpdate { - item: Some(SummaryInputExpr::Column( - planner_types::pre_asap::ColumnRef::Named("host".into()), - )), - weight: SummaryInputExpr::ResetAwareCounterDelta { - value: planner_types::pre_asap::ColumnRef::SampleValue, - series: EntityIdentity::PromqlLabelSet { excluding: vec![] }, - }, - weight_domain: WeightDomain::NonNegative { - proof: NonNegativeWeightProof::ResetAwareCounterDerivative, - }, - }; - let mut updater = create_planner_accumulator(&family, &input, &Default::default()).unwrap(); - 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()); - } - // Heap counters use fixed-point storage for fractional counter deltas. - // This encodes the selected update; it does not choose another family. - let weight_scale = if matches!( - input.weight, - planner_types::post_asap::SummaryInputExpr::ResetAwareCounterDelta { .. } - ) { - 1_000_000.0 - } else { - 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..393690a97 100644 --- a/data_plane/src/precompute_engine/ingest_handler.rs +++ b/data_plane/src/precompute_engine/ingest_handler.rs @@ -366,7 +366,7 @@ 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_physical_operators::accumulators::DDSketchAccumulator; use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; use asap_sketchlib::DdSketch; use planner_types::post_asap::SketchAlgorithm; @@ -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::accumulators::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 78c58dea5..f504770fb 100644 --- a/data_plane/src/precompute_engine/maintenance_runtime.rs +++ b/data_plane/src/precompute_engine/maintenance_runtime.rs @@ -258,7 +258,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() { @@ -472,7 +472,7 @@ fn evaluate_aligned_binary( .get(×tamp) .ok_or("maintenance binary requires matching timestamp sets")?; let value = - crate::utils::arithmetic::evaluate_float64_arithmetic(arithmetic, left, *right); + asap_physical_operators::arithmetic::evaluate_float64_arithmetic(arithmetic, left, *right); if !value.is_finite() { return Err("maintenance binary produced a non-finite update".into()); } @@ -2155,7 +2155,7 @@ pub(crate) fn affected_materializations( #[cfg(test)] mod tests { use super::*; - use crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::accumulators::SumAccumulator; use planner_types::post_asap::{ EdgeRole, ExecutableDag, ExecutableDagEdge, GroupingEdgeCompatibility, SummarySchema, WindowEdgeCompatibility, @@ -2189,7 +2189,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::accumulators::SumAccumulator::new(); state.update(value); FrozenExactWindows { sid, 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 f9583d608..de328f256 100644 --- a/data_plane/src/precompute_engine/output_sink.rs +++ b/data_plane/src/precompute_engine/output_sink.rs @@ -335,7 +335,7 @@ impl OutputSink for NoopOutputSink { #[cfg(test)] mod tests { use super::*; - use crate::precompute_engine::operators::{DDSketchAccumulator, SumAccumulator}; + use asap_physical_operators::accumulators::{DDSketchAccumulator, SumAccumulator}; use crate::storage_engines::sketch_db::index::{AggKind, SeriesLookup}; use crate::storage_engines::types::{KeyByLabelValues, StreamingConfig}; use asap_types::aggregation_config::PrecomputeMaterialization; diff --git a/data_plane/src/precompute_engine/raw_dag.rs b/data_plane/src/precompute_engine/raw_dag.rs index ee76c8b49..5baa25ac6 100644 --- a/data_plane/src/precompute_engine/raw_dag.rs +++ b/data_plane/src/precompute_engine/raw_dag.rs @@ -1,5 +1,5 @@ //! Bind raw ingestion to a selected Planner producer and its raw dependency edge. -use super::accumulator_factory::{create_planner_accumulator, AccumulatorUpdater}; +use asap_physical_operators::factory::{create_planner_accumulator, AccumulatorUpdater}; use crate::storage_engines::types::KeyByLabelValues; use asap_types::{executable_plan::BackendNodeBinding, PrecomputeMaterialization}; use planner_types::post_asap::{ diff --git a/data_plane/src/precompute_engine/worker.rs b/data_plane/src/precompute_engine/worker.rs index cda3524fd..a21890931 100644 --- a/data_plane/src/precompute_engine/worker.rs +++ b/data_plane/src/precompute_engine/worker.rs @@ -1,10 +1,10 @@ #[cfg(test)] -use crate::precompute_engine::accumulator_factory::create_fixture_accumulator; -use crate::precompute_engine::accumulator_factory::AccumulatorUpdater; +use crate::tests::accumulator_fixture::create_fixture_accumulator; +use asap_physical_operators::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 asap_physical_operators::accumulators::sum_accumulator::SumAccumulator; use crate::precompute_engine::output_sink::OutputSink; use crate::precompute_engine::series_router::WorkerMessage; use crate::precompute_engine::window_manager::WindowManager; @@ -1874,9 +1874,9 @@ 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 asap_physical_operators::accumulators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; + use asap_physical_operators::accumulators::keyed_sum_count_accumulator::KeyedSumCountAccumulator; + use asap_physical_operators::accumulators::sum_accumulator::SumAccumulator; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::StreamingConfig; use asap_sketchlib::KllSketch; @@ -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::accumulators::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::accumulators::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::accumulators::IncreaseAccumulator; let config = make_agg_config( 1, "requests_total", @@ -4280,7 +4280,7 @@ mod tests { #[cfg(test)] mod dag_execution_tests { use super::*; - use crate::precompute_engine::operators::exact_accumulator::ExactAccumulator; + use asap_physical_operators::accumulators::exact_accumulator::ExactAccumulator; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::StreamingConfig; use asap_types::query_plan::ExactReadout; 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 87c751e06..2709c45dd 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 asap_physical_operators::accumulators::SumAccumulator; + use crate::storage_engines::sketch_db::index::{AggKind, Capability, SummarySeriesMetadata}; use asap_types::query_plan::{ ClickHousePlanningContext, ExactReadout, ExternalExactOutput, ExternalExactRequest, FallbackPolicy, FixedEvaluationRange, InstantExecution, MaterializationBinding, 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..3c31156e5 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -1394,7 +1394,7 @@ mod sketch_query_tests { #[cfg(test)] mod aux_pushdown_tests { use super::*; - use crate::precompute_engine::operators::{ + use asap_physical_operators::accumulators::{ max_accumulator::MaxAccumulator, min_accumulator::MinAccumulator, sum_accumulator::SumAccumulator, }; @@ -1613,7 +1613,7 @@ 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 asap_physical_operators::accumulators::sum_accumulator::SumAccumulator; use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; @@ -2263,7 +2263,7 @@ 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 asap_physical_operators::accumulators::sum_accumulator::SumAccumulator; use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; 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 0e1752951..78ae660ad 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 @@ -888,7 +888,7 @@ mod tests { #[tokio::test] async fn five_minute_error_ratio_combines_prometheus_cut_with_summary_store() { - use crate::precompute_engine::operators::IncreaseAccumulator; + use asap_physical_operators::accumulators::IncreaseAccumulator; use crate::query_engines::query_result::{InstantVectorElement, QueryResult}; use crate::storage_engines::sketch_db::{ data::AggKind, 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 59bffd71d..4485cf078 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::accumulators::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 a9e4d8460..86b29c8f6 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 @@ -759,7 +759,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::accumulators::sum_accumulator::SumAccumulator::new(); sum.update(1e308); sum.update(1e308); assert!(binary( 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 1da37ae39..e449560a9 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}; @@ -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::accumulators::SumAccumulator::with_sum(42.0)), ); let config = test_plan::materialization("bytes_total", "Sum", serde_json::json!({}), &[], 1000); @@ -1146,7 +1146,7 @@ mod tests { BTreeMap::new(), bounds, Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum(sum), + asap_physical_operators::accumulators::SumAccumulator::with_sum(sum), ), ); } @@ -1214,7 +1214,7 @@ mod tests { BTreeMap::new(), (pane * 60_000, (pane + 1) * 60_000), Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum( + asap_physical_operators::accumulators::SumAccumulator::with_sum( (pane + 1) as f64, ), ), @@ -1289,7 +1289,7 @@ mod tests { BTreeMap::new(), (pane * 10_000, (pane + 1) * 10_000), Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum( + asap_physical_operators::accumulators::SumAccumulator::with_sum( (pane + 1) as f64, ), ), @@ -1356,7 +1356,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::accumulators::SumAccumulator::with_sum(1.0)), ); } assert!( @@ -1386,7 +1386,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::accumulators::IncreaseAccumulator::new( Measurement::new(10.0), 10_000, Measurement::new(10.0), 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 dfcf32f6b..0a8549d38 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,9 @@ 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 asap_physical_operators::accumulators::increase_accumulator::IncreaseAccumulator; +use asap_physical_operators::accumulators::max_accumulator::MaxAccumulator; +use asap_physical_operators::accumulators::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::{ @@ -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 @@ -1441,7 +1441,7 @@ mod tests { #[test] fn keyed_count_state_follows_planner_family_and_query_readout() { - use crate::precompute_engine::operators::KeyedSumCountAccumulator; + use asap_physical_operators::accumulators::KeyedSumCountAccumulator; use asap_types::query_plan::ExactReadout; let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); @@ -1958,7 +1958,7 @@ 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 asap_physical_operators::accumulators::univmon_accumulator::UnivMonAccumulator; use crate::storage_engines::sketch_db::index::SketchEncoding; use crate::storage_engines::types::SerializableToSink; use asap_types::query_plan::{MaterializationBinding, PhysicalGrouping}; @@ -3190,13 +3190,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::accumulators::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::accumulators::SumAccumulator::with_sum(15.0)), ); let child = scan_node("bytes_total", None); @@ -3296,13 +3296,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::accumulators::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::accumulators::SumAccumulator::with_sum(12.0)), ); let child = scan_node("bytes_total", None); @@ -3342,7 +3342,7 @@ mod tests { sid, BTreeMap::new(), (T0, T0 + 1000), - Box::new(crate::precompute_engine::operators::MaxAccumulator::new()), + Box::new(asap_physical_operators::accumulators::MaxAccumulator::new()), ); let child = scan_node("latency_max_ms", None); 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..30a7f8383 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![]); @@ -932,7 +932,7 @@ mod tests { cfg.policy_fingerprint(), ); let acc = - crate::precompute_engine::operators::sum_accumulator::SumAccumulator::with_sum(1.0); + asap_physical_operators::accumulators::sum_accumulator::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..7c8b09a6d 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 @@ -5,9 +5,9 @@ //! so both paths use the same sketch semantics. #[cfg(test)] -use crate::precompute_engine::accumulator_factory::{ - create_fixture_accumulator, AccumulatorUpdater, -}; +use asap_physical_operators::factory::AccumulatorUpdater; +#[cfg(test)] +use crate::tests::accumulator_fixture::create_fixture_accumulator; #[cfg(test)] use crate::precompute_engine::worker::apply_sample; use crate::storage_engines::sketch_db::backfill::raw_sample_reader::RawSample; @@ -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::accumulators::{ CountMinSketchWithHeapAccumulator, CountSketchWithHeapAccumulator, }; for kind in [ 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 c1bc5f220..6c465d7cb 100644 --- a/data_plane/src/storage_engines/sketch_db/index/maintenance.rs +++ b/data_plane/src/storage_engines/sketch_db/index/maintenance.rs @@ -729,7 +729,7 @@ impl SketchStore { #[cfg(test)] mod tests { use super::*; - use crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::accumulators::SumAccumulator; use crate::storage_engines::types::PrecomputedOutput; use asap_types::traits::SerializableToSink; @@ -1237,7 +1237,7 @@ mod tests { )]) ); let assert_complete_output = |store: &SketchStore| { - use crate::precompute_engine::operators::DDSketchAccumulator; + use asap_physical_operators::accumulators::DDSketchAccumulator; use crate::storage_engines::sketch_db::data::SketchEncoding; let rows = store.query_range(target_sid, 0, 60_000); assert_eq!(rows.len(), 1); 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 098c90777..ed654f0a5 100644 --- a/data_plane/src/storage_engines/sketch_db/index/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/index/mod.rs @@ -91,13 +91,13 @@ fn reconstruct_exact_agg( type_name: &str, bytes: &[u8], ) -> Option> { - use crate::precompute_engine::operators::{ + use asap_physical_operators::accumulators::{ 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::accumulators::exact_accumulator::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), @@ -1571,12 +1571,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 @@ -4403,7 +4403,7 @@ mod tests { #[test] fn precompute_payload_round_trips_through_storage() { - use crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::accumulators::SumAccumulator; let idx = SketchStore::new(); let cfg = SketchConfig::DDSketch { @@ -4444,7 +4444,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::accumulators::SumAccumulator; let idx = SketchStore::new(); let cfg = SketchConfig::DDSketch { @@ -4598,7 +4598,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::accumulators::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()); @@ -5342,7 +5342,7 @@ mod tests { 850, BTreeMap::new(), (0, 30_000), - Box::new(crate::precompute_engine::operators::SumAccumulator::new()) + Box::new(asap_physical_operators::accumulators::SumAccumulator::new()) )); // A flusher that captured metadata before completion cannot reopen it. writer.upsert_all(&[stale_record]).unwrap(); @@ -5767,7 +5767,7 @@ mod tests { lv_zone("z0"), (s, s + 30_000), Box::new( - crate::precompute_engine::operators::SumAccumulator::with_sum( + asap_physical_operators::accumulators::SumAccumulator::with_sum( (i + 1) as f64, ), ), @@ -5986,7 +5986,7 @@ 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::accumulators::SumAccumulator::with_sum((i + 1) as f64), ), ); } @@ -6050,7 +6050,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::accumulators::SumAccumulator::new(); acc.update((i + 1) as f64); acc.update(10.0); acc @@ -6243,7 +6243,7 @@ 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 asap_physical_operators::accumulators::exact_accumulator::ExactAccumulator; use crate::storage_engines::types::{AggregateCore, AggregationType}; let kinds = [ AggregationType::Sum, 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 768f09469..df7c1f7bb 100644 --- a/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs +++ b/data_plane/src/storage_engines/sketch_db/lifecycle/eviction.rs @@ -193,7 +193,7 @@ pub fn warn_if_retention_inverted( #[cfg(test)] mod tests { use super::*; - use crate::precompute_engine::operators::SumAccumulator; + use asap_physical_operators::accumulators::SumAccumulator; use crate::storage_engines::types::{AggregationType, StreamingConfig}; use asap_types::aggregation_config::PrecomputeMaterialization; use asap_types::enums::WindowKind; 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..3d254f2cf 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::accumulators::count_min_sketch_with_heap_accumulator::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..d36b69552 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::accumulators::univmon_accumulator::UnivMonAccumulator::new( *heap_size as usize, *sketch_rows as usize, *sketch_cols as usize, @@ -167,7 +167,7 @@ fn decode_full( }, SketchEncoding::MsgpackFull, ) => { - let state = crate::precompute_engine::operators::univmon_accumulator::UnivMonAccumulator::from_bytes(bytes) + let state = asap_physical_operators::accumulators::univmon_accumulator::UnivMonAccumulator::from_bytes(bytes) .map_err(|e| e.to_string())?; if state.dimensions() != ( @@ -244,7 +244,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::accumulators::univmon_accumulator::UnivMonAccumulator), Dd(DdSketch), Hll(HllSketch), Kll(KllSketch), @@ -321,7 +321,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::accumulators::dd_sketch_accumulator::DDSketchAccumulator; let mut acc = DDSketchAccumulator { inner: std::mem::replace(sk, DdSketch::new(sk.alpha)), sample_p: 1.0, @@ -740,21 +740,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::accumulators::dd_sketch_accumulator::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::accumulators::datasketches_kll_accumulator::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::accumulators::hll_sketch_accumulator::HllSketchAccumulator; HllSketchAccumulator::from_sketchlib_proto_bytes(buffer) .map(|acc| acc.inner) .map_err(|e| e.to_string()) @@ -910,7 +910,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::accumulators::hll_sketch_accumulator::HllSketchAccumulator; let mut sk = HllSketch::new(HllVariant::Regular, 12); for i in 0..500u64 { sk.update(format!("item-{i}").as_bytes()); @@ -929,7 +929,7 @@ mod tests { #[test] fn dd_from_proto_matches_accumulator_decoder() { - use crate::precompute_engine::operators::dd_sketch_accumulator::DDSketchAccumulator; + use asap_physical_operators::accumulators::dd_sketch_accumulator::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 +946,7 @@ mod tests { #[test] fn kll_from_proto_matches_accumulator_decoder() { - use crate::precompute_engine::operators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; + use asap_physical_operators::accumulators::datasketches_kll_accumulator::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..8e18c22d6 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 enums::*; pub use hot_reload_config::*; -pub use key_by_label_values::*; -pub use measurement::*; +pub use asap_physical_operators::key_by_label_values::*; +pub use asap_physical_operators::measurement::*; pub use precomputed_output::*; pub use storage_backend::*; pub use streaming_config::*; -pub use traits::*; +pub use asap_physical_operators::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..8c7eaac56 --- /dev/null +++ b/data_plane/src/tests/accumulator_fixture.rs @@ -0,0 +1,276 @@ +//! 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, + } +} + +#[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") + }; + 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..22ff11181 100644 --- a/data_plane/src/tests/trait_design_tests.rs +++ b/data_plane/src/tests/trait_design_tests.rs @@ -1,4 +1,4 @@ -use crate::precompute_engine::operators::{KeyedSumCountAccumulator, SumAccumulator}; +use asap_physical_operators::accumulators::{KeyedSumCountAccumulator, SumAccumulator}; #[cfg(test)] use crate::storage_engines::types::{ KeyByLabelValues, MultipleSubpopulationAggregate, SingleSubpopulationAggregate, 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/edge_sketch_codec.rs b/data_plane/tests/edge_sketch_codec.rs index 4f94d9fee..0055124e4 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::accumulators::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::accumulators::DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&snapshot_bytes).unwrap(); assert!(accumulator.get_quantile(0.5).is_finite()); } diff --git a/data_plane/tests/support/univmon_erp_process.rs b/data_plane/tests/support/univmon_erp_process.rs index 28bb7b01c..d849405b9 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 control_plane::physical::erp::ErpShapeObserver; -use data_plane::precompute_engine::operators::univmon_accumulator::UnivMonAccumulator; +use asap_physical_operators::accumulators::univmon_accumulator::UnivMonAccumulator; use data_plane::storage_engines::types::{AggregateCore, SerializableToSink}; fn values(offset: usize) -> Vec { From 53b762af5057fe9212c7993d5f614eb0ae8de007 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 12:42:02 +0000 Subject: [PATCH 03/15] test: include shared physical kernels in the contracts suite --- scripts/e2e.sh | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/scripts/e2e.sh b/scripts/e2e.sh index 2600999f9..19878aa6a 100755 --- a/scripts/e2e.sh +++ b/scripts/e2e.sh @@ -81,6 +81,10 @@ contracts() { say "contracts: shared policy and routing types" rust_test asap_types + CURRENT_STAGE="contracts/asap-physical-operators" + say "contracts: shared physical kernels and deployment-independent execution" + rust_test asap-physical-operators + CURRENT_STAGE="contracts/asap_otel_proto" say "contracts: modified OTLP and monitor protobuf compatibility" rust_test asap_otel_proto --tests From d3e09996027198b12dc8c663fbf794f5fbfe0cb8 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 19:08:26 +0000 Subject: [PATCH 04/15] Expose scoped native source execution for engine adapters --- .../src/dag/batch_execution.rs | 25 ++++++++++++++++--- 1 file changed, 22 insertions(+), 3 deletions(-) diff --git a/crates/asap-physical-operators/src/dag/batch_execution.rs b/crates/asap-physical-operators/src/dag/batch_execution.rs index d29a9b0cc..9ae78d74e 100644 --- a/crates/asap-physical-operators/src/dag/batch_execution.rs +++ b/crates/asap-physical-operators/src/dag/batch_execution.rs @@ -22,6 +22,21 @@ pub fn evaluate_batch( graph.add(root + 1, vec![root], operator)?; root += 1; } + evaluate_graph(graph, root, context) +} + +/// Evaluate a native in-memory source, including scalar sources, in the caller's scope. +pub fn evaluate_source(source: Operator, context: RunContext) -> Result, Error> { + let mut graph = PhysicalDag::default(); + graph.add(0, vec![], source)?; + evaluate_graph(graph, 0, context) +} + +fn evaluate_graph( + graph: PhysicalDag<'_, Batch, super::values::Schema>, + root: super::NodeId, + context: RunContext, +) -> Result, Error> { let mut output = graph.execute(&[root], context)?.remove(0); let mut batches = Vec::new(); loop { @@ -80,10 +95,14 @@ mod tests { ) .unwrap(); let context = RunContext::new(scope, Limits::default()).unwrap(); - let result = - futures::executor::block_on(async { evaluate_batch(batch, vec![negate], context) }) - .unwrap(); + let result = futures::executor::block_on(async { + evaluate_batch(batch, vec![negate], context.clone()) + }) + .unwrap(); assert!(matches!(result[0].rows()[0][0], Value::Float64(-7.))); + let source = Operator::scalar(Value::Float64(9.), DataType::Float64).unwrap(); + let scalar = evaluate_source(source, context).unwrap(); + assert!(matches!(scalar[0].rows()[0][0], Value::Float64(9.))); } } From ee976dc5932ac2333f0b963c5f421a5e6d254355 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 19:13:38 +0000 Subject: [PATCH 05/15] Drive native source yields and keep pane rejection assertions semantic --- .../src/dag/batch_execution.rs | 28 +++++++++++++++---- data_plane/src/drivers/query/servers/http.rs | 2 +- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/crates/asap-physical-operators/src/dag/batch_execution.rs b/crates/asap-physical-operators/src/dag/batch_execution.rs index 9ae78d74e..cc43e7d43 100644 --- a/crates/asap-physical-operators/src/dag/batch_execution.rs +++ b/crates/asap-physical-operators/src/dag/batch_execution.rs @@ -44,11 +44,9 @@ fn evaluate_graph( Some(Some(Ok(batch))) => batches.push(batch.value().clone()), Some(Some(Err(error))) => return Err(error), Some(None) => return Ok(batches), - None => { - return Err(Error::Operator( - "in-memory native batch chain unexpectedly awaited I/O".into(), - )) - } + // Native operators have no I/O sources here. Pending is the + // shared runtime's cooperative yield after a batch quantum. + None => continue, } } } @@ -106,6 +104,26 @@ mod tests { } } + // Native sources may cross the runtime's cooperative batch quantum. + #[test] + fn in_memory_source_drives_cooperative_yields() { + let schema = Arc::new(SummarySchema { + fields: vec![], + time_index: None, + }); + let batch = Batch::try_new(schema.clone(), vec![vec![]]).unwrap(); + let source = Operator::source(schema, vec![batch; 65]).unwrap(); + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + assert_eq!(evaluate_source(source, context).unwrap().len(), 65); + } + // A cancelled surrounding execution also prevents its native computation. #[test] fn cancellation_is_not_bypassed_by_in_memory_execution() { diff --git a/data_plane/src/drivers/query/servers/http.rs b/data_plane/src/drivers/query/servers/http.rs index 6c4c71514..ddfaeecd9 100644 --- a/data_plane/src/drivers/query/servers/http.rs +++ b/data_plane/src/drivers/query/servers/http.rs @@ -6652,7 +6652,7 @@ mod catalog_install_tests { binding.window_ms += 1; let error = install(request).unwrap_err(); assert!( - error.contains("query physical pane duration differs"), + error.contains("query") && error.contains("pane") && error.contains("differs"), "{error}" ); } From 704ee1b4b11fa30c7277be4d79d665b0648c6dc4 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 19:16:55 +0000 Subject: [PATCH 06/15] Link standalone library documentation to its foundation design --- crates/asap-physical-operators/README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/asap-physical-operators/README.md b/crates/asap-physical-operators/README.md index 8bde41a54..bc9a439c0 100644 --- a/crates/asap-physical-operators/README.md +++ b/crates/asap-physical-operators/README.md @@ -63,8 +63,8 @@ Existing accumulator algorithms are reused as kernels behind these operators. Backend ingestion integration is delivered in #763 and query integration in #765, after this foundation. Installed value/storage adapters provide deployment-specific computation; they have not all been replaced by native batch bindings. Local raw -Scan remains deferred. See the [design and coverage table](../../docs/design_docs/query-dag-execution.md) -for the distinction between native operator support and backend integration. +Scan remains deferred. See the [shared operator design](../../docs/design_docs/physical-operators.md). +The dependent #765 query DAG design tracks installed engine coverage separately. The default limits are eight buffered batches per producer and 64 MiB of estimated retained execution data. Callers can set both through `Limits`. Accounting includes From 66849f95a3313e5d84b32906852f86467e9b988d Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 19:22:04 +0000 Subject: [PATCH 07/15] Retain native output reservations across engine adapter boundaries --- .../src/dag/batch_execution.rs | 42 ++++++++++++++++--- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/crates/asap-physical-operators/src/dag/batch_execution.rs b/crates/asap-physical-operators/src/dag/batch_execution.rs index cc43e7d43..78b521724 100644 --- a/crates/asap-physical-operators/src/dag/batch_execution.rs +++ b/crates/asap-physical-operators/src/dag/batch_execution.rs @@ -1,6 +1,6 @@ //! Execute a bounded in-memory batch through native operators. This is also the //! bridge for deployments whose boundary values are not yet streaming batches. -use super::{operators::Operator, values::Batch, Error, PhysicalDag, RunContext}; +use super::{operators::Operator, values::Batch, Error, PhysicalDag, RunContext, SharedValue}; use futures::{FutureExt, StreamExt}; /// Every input is already in memory; the chain contains native operators only. @@ -10,7 +10,7 @@ pub fn evaluate_batch( input: Batch, operators: Vec, context: RunContext, -) -> Result, Error> { +) -> Result>, Error> { let mut graph = PhysicalDag::default(); graph.add( 0, @@ -26,7 +26,10 @@ pub fn evaluate_batch( } /// Evaluate a native in-memory source, including scalar sources, in the caller's scope. -pub fn evaluate_source(source: Operator, context: RunContext) -> Result, Error> { +pub fn evaluate_source( + source: Operator, + context: RunContext, +) -> Result>, Error> { let mut graph = PhysicalDag::default(); graph.add(0, vec![], source)?; evaluate_graph(graph, 0, context) @@ -36,12 +39,12 @@ fn evaluate_graph( graph: PhysicalDag<'_, Batch, super::values::Schema>, root: super::NodeId, context: RunContext, -) -> Result, Error> { +) -> Result>, Error> { let mut output = graph.execute(&[root], context)?.remove(0); let mut batches = Vec::new(); loop { match output.next().now_or_never() { - Some(Some(Ok(batch))) => batches.push(batch.value().clone()), + Some(Some(Ok(batch))) => batches.push(batch), Some(Some(Err(error))) => return Err(error), Some(None) => return Ok(batches), // Native operators have no I/O sources here. Pending is the @@ -124,6 +127,35 @@ mod tests { assert_eq!(evaluate_source(source, context).unwrap().len(), 65); } + // An adapter-held output must retain its parent's reservation after execution. + #[test] + fn returned_batches_keep_their_resource_reservation() { + let schema = Arc::new(SummarySchema { + fields: vec![], + time_index: None, + }); + let batch = Batch::try_new(schema.clone(), vec![vec![]]).unwrap(); + let bytes = batch.bytes(); + let source = Operator::source(schema, vec![batch]).unwrap(); + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits { + max_bytes: bytes, + max_buffered_batches: 1, + }, + ) + .unwrap(); + let held = evaluate_source(source.clone(), context.clone()).unwrap(); + assert_eq!(context.retained_bytes(), bytes); + assert!(evaluate_source(source.clone(), context.clone()).is_err()); + drop(held); + assert_eq!(context.retained_bytes(), 0); + assert!(evaluate_source(source, context).is_ok()); + } + // A cancelled surrounding execution also prevents its native computation. #[test] fn cancellation_is_not_bypassed_by_in_memory_execution() { From fb8bc803484839dad28ed0e16b44d7a7d75e1c57 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 19:28:25 +0000 Subject: [PATCH 08/15] Format relocated physical operator imports at the foundation boundary --- data_plane/src/drivers/ingest/otel.rs | 12 ++++++------ data_plane/src/precompute_engine/ingest_handler.rs | 2 +- .../src/precompute_engine/maintenance_runtime.rs | 5 +++-- data_plane/src/precompute_engine/output_sink.rs | 2 +- data_plane/src/precompute_engine/raw_dag.rs | 2 +- data_plane/src/precompute_engine/worker.rs | 14 +++++++------- .../asap_clickhouse_query_engine/accelerator.rs | 2 +- .../src/query_engines/asap_query_engine/engine.rs | 6 +++--- .../asap_query_engine/exact_subqueries.rs | 2 +- .../asap_query_engine/post_asap_readout.rs | 4 +++- .../asap_query_engine/summary_executor.rs | 8 ++++---- .../sketch_db/backfill/window_builder.rs | 8 ++++---- .../storage_engines/sketch_db/index/maintenance.rs | 4 ++-- .../src/storage_engines/sketch_db/index/mod.rs | 4 ++-- .../sketch_db/lifecycle/eviction.rs | 2 +- data_plane/src/storage_engines/types/mod.rs | 6 +++--- data_plane/src/tests/trait_design_tests.rs | 2 +- data_plane/tests/support/univmon_erp_process.rs | 2 +- 18 files changed, 45 insertions(+), 42 deletions(-) diff --git a/data_plane/src/drivers/ingest/otel.rs b/data_plane/src/drivers/ingest/otel.rs index cdfb1836b..985edc6c4 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 asap_physical_operators::accumulators::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::accumulators::sketch_envelope_accumulator::SketchEnvelopeAccumulator; use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; use asap_sketchlib::MessagePackCodec; use axum::{body::Bytes, extract::State, routing::post, Json, Router}; @@ -2684,11 +2684,11 @@ fn empty_accumulator_for_delta_bootstrap( config: &crate::storage_engines::sketch_db::index::SketchConfig, encoding: i32, ) -> Option> { + use crate::storage_engines::sketch_db::index::SketchConfig; use asap_physical_operators::accumulators::{ CountMinSketchAccumulator, CountSketchAccumulator, CountSketchWithHeapAccumulator, HllSketchAccumulator, }; - use crate::storage_engines::sketch_db::index::SketchConfig; match (algorithm, config) { (SketchAlgorithm::Hll, SketchConfig::Hll { precision }) => { @@ -3444,8 +3444,8 @@ mod policy_fp_lookup_tests { #[cfg(test)] mod dispatcher_tests { use super::*; - use asap_physical_operators::accumulators::{DDSketchAccumulator, HllSketchAccumulator}; use crate::storage_engines::types::AggregateCore; + use asap_physical_operators::accumulators::{DDSketchAccumulator, HllSketchAccumulator}; use asap_sketchlib::DdSketch; use asap_sketchlib::HllVariant; @@ -3797,8 +3797,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 asap_physical_operators::accumulators::DDSketchAccumulator; use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use asap_physical_operators::accumulators::DDSketchAccumulator; use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; use prost::Message; @@ -4129,8 +4129,8 @@ mod sid_resolution_tests { /// recover after a backend restart. #[tokio::test] async fn leading_cms_delta_bootstraps_onto_empty_base() { - use asap_physical_operators::accumulators::CountMinSketchAccumulator; use asap_otel_proto::sketchlib::v1::CountMinDelta as PbDelta; + use asap_physical_operators::accumulators::CountMinSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; @@ -4215,8 +4215,8 @@ mod sid_resolution_tests { /// the register-max updates. #[tokio::test] async fn leading_hll_delta_bootstraps_onto_empty_base() { - use asap_physical_operators::accumulators::HllSketchAccumulator; use asap_otel_proto::sketchlib::v1::HllDelta as PbDelta; + use asap_physical_operators::accumulators::HllSketchAccumulator; use prost::Message; let (state, drain) = make_state().await; diff --git a/data_plane/src/precompute_engine/ingest_handler.rs b/data_plane/src/precompute_engine/ingest_handler.rs index 393690a97..77bc6c387 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 asap_physical_operators::accumulators::DDSketchAccumulator; use asap_otel_proto::sketchlib::v1::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use asap_physical_operators::accumulators::DDSketchAccumulator; use asap_sketchlib::DdSketch; use planner_types::post_asap::SketchAlgorithm; use prost::Message; diff --git a/data_plane/src/precompute_engine/maintenance_runtime.rs b/data_plane/src/precompute_engine/maintenance_runtime.rs index f504770fb..f33317a79 100644 --- a/data_plane/src/precompute_engine/maintenance_runtime.rs +++ b/data_plane/src/precompute_engine/maintenance_runtime.rs @@ -471,8 +471,9 @@ fn evaluate_aligned_binary( let right = right_rows .get(×tamp) .ok_or("maintenance binary requires matching timestamp sets")?; - let value = - asap_physical_operators::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()); } diff --git a/data_plane/src/precompute_engine/output_sink.rs b/data_plane/src/precompute_engine/output_sink.rs index de328f256..3713e0b01 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 asap_physical_operators::accumulators::{DDSketchAccumulator, SumAccumulator}; use crate::storage_engines::sketch_db::index::{AggKind, SeriesLookup}; use crate::storage_engines::types::{KeyByLabelValues, StreamingConfig}; + use asap_physical_operators::accumulators::{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 5baa25ac6..a1884e322 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 asap_physical_operators::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, diff --git a/data_plane/src/precompute_engine/worker.rs b/data_plane/src/precompute_engine/worker.rs index a21890931..de14847a6 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::tests::accumulator_fixture::create_fixture_accumulator; -use asap_physical_operators::factory::AccumulatorUpdater; use crate::precompute_engine::config::LateDataPolicy; use crate::precompute_engine::group_key::GroupKey; use crate::precompute_engine::metrics::record_late_input; -use asap_physical_operators::accumulators::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::accumulators::sum_accumulator::SumAccumulator; +use asap_physical_operators::factory::AccumulatorUpdater; 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::output_sink::CapturingOutputSink; + use crate::storage_engines::types::StreamingConfig; use asap_physical_operators::accumulators::datasketches_kll_accumulator::DatasketchesKLLAccumulator; use asap_physical_operators::accumulators::keyed_sum_count_accumulator::KeyedSumCountAccumulator; use asap_physical_operators::accumulators::sum_accumulator::SumAccumulator; - use crate::precompute_engine::output_sink::CapturingOutputSink; - use crate::storage_engines::types::StreamingConfig; use asap_sketchlib::KllSketch; use asap_types::enums::WindowKind; use asap_types::AggregationType; @@ -4280,9 +4280,9 @@ mod tests { #[cfg(test)] mod dag_execution_tests { use super::*; - use asap_physical_operators::accumulators::exact_accumulator::ExactAccumulator; use crate::precompute_engine::output_sink::CapturingOutputSink; use crate::storage_engines::types::StreamingConfig; + use asap_physical_operators::accumulators::exact_accumulator::ExactAccumulator; use asap_types::query_plan::ExactReadout; fn plan(query: &str) -> control_plane::physical::compiler::CompiledPhysicalPlan { diff --git a/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs b/data_plane/src/query_engines/asap_clickhouse_query_engine/accelerator.rs index 2709c45dd..30f39c8bf 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,8 +361,8 @@ mod tests { } } - use asap_physical_operators::accumulators::SumAccumulator; use crate::storage_engines::sketch_db::index::{AggKind, Capability, SummarySeriesMetadata}; + use asap_physical_operators::accumulators::SumAccumulator; use asap_types::query_plan::{ ClickHousePlanningContext, ExactReadout, ExternalExactOutput, ExternalExactRequest, FallbackPolicy, FixedEvaluationRange, InstantExecution, MaterializationBinding, 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 3c31156e5..72cf22f6a 100644 --- a/data_plane/src/query_engines/asap_query_engine/engine.rs +++ b/data_plane/src/query_engines/asap_query_engine/engine.rs @@ -1394,11 +1394,11 @@ mod sketch_query_tests { #[cfg(test)] mod aux_pushdown_tests { use super::*; + use crate::storage_engines::types::AggregationType; use asap_physical_operators::accumulators::{ max_accumulator::MaxAccumulator, min_accumulator::MinAccumulator, sum_accumulator::SumAccumulator, }; - use crate::storage_engines::types::AggregationType; use asap_types::Statistic; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; @@ -1613,9 +1613,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 asap_physical_operators::accumulators::sum_accumulator::SumAccumulator; use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; + use asap_physical_operators::accumulators::sum_accumulator::SumAccumulator; let idx = Arc::new(SketchStore::new()); // Mirror the acceptance-test setup: four ExactAgg(Sum) sids, one @@ -2263,9 +2263,9 @@ mod asap_tier_classify_tests { /// `OuterFn::Plain` instant sums. #[tokio::test] async fn execute_instant_sum_accumulates_all_windows_not_last() { - use asap_physical_operators::accumulators::sum_accumulator::SumAccumulator; use crate::query_engines::query_result::QueryResult; use crate::storage_engines::sketch_db::data::AggregationType; + use asap_physical_operators::accumulators::sum_accumulator::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 78ae660ad..ba8356858 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 @@ -888,13 +888,13 @@ mod tests { #[tokio::test] async fn five_minute_error_ratio_combines_prometheus_cut_with_summary_store() { - use asap_physical_operators::accumulators::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::accumulators::IncreaseAccumulator; use asap_types::query_plan::{ residual::BinaryOperation, ExactReadout, MaterializationBinding, PhysicalGrouping, }; 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 e449560a9..aeeceebc7 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 @@ -1146,7 +1146,9 @@ mod tests { BTreeMap::new(), bounds, Box::new( - asap_physical_operators::accumulators::SumAccumulator::with_sum(sum), + asap_physical_operators::accumulators::SumAccumulator::with_sum( + sum, + ), ), ); } 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 0a8549d38..851e1c1c3 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 asap_physical_operators::accumulators::increase_accumulator::IncreaseAccumulator; -use asap_physical_operators::accumulators::max_accumulator::MaxAccumulator; -use asap_physical_operators::accumulators::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::accumulators::increase_accumulator::IncreaseAccumulator; +use asap_physical_operators::accumulators::max_accumulator::MaxAccumulator; +use asap_physical_operators::accumulators::min_accumulator::MinAccumulator; /// Per-query, per-call execution context — constructed fresh for each /// incoming query (never shared across concurrent queries, never @@ -1958,9 +1958,9 @@ mod tests { /// One installed frequency summary merges panes before all four readouts. #[test] fn bound_univmon_merges_panes_for_four_readouts() { - use asap_physical_operators::accumulators::univmon_accumulator::UnivMonAccumulator; use crate::storage_engines::sketch_db::index::SketchEncoding; use crate::storage_engines::types::SerializableToSink; + use asap_physical_operators::accumulators::univmon_accumulator::UnivMonAccumulator; use asap_types::query_plan::{MaterializationBinding, PhysicalGrouping}; let index = SketchStore::new(); let fp = asap_types::PolicyFingerprint(701); diff --git a/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs b/data_plane/src/storage_engines/sketch_db/backfill/window_builder.rs index 7c8b09a6d..430c59aa4 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 asap_physical_operators::factory::AccumulatorUpdater; -#[cfg(test)] -use crate::tests::accumulator_fixture::create_fixture_accumulator; #[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 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 6c465d7cb..a4410abe9 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 asap_physical_operators::accumulators::SumAccumulator; use crate::storage_engines::types::PrecomputedOutput; + use asap_physical_operators::accumulators::SumAccumulator; use asap_types::traits::SerializableToSink; #[test] @@ -1237,8 +1237,8 @@ mod tests { )]) ); let assert_complete_output = |store: &SketchStore| { - use asap_physical_operators::accumulators::DDSketchAccumulator; use crate::storage_engines::sketch_db::data::SketchEncoding; + use asap_physical_operators::accumulators::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 ed654f0a5..fb9105cef 100644 --- a/data_plane/src/storage_engines/sketch_db/index/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/index/mod.rs @@ -91,11 +91,11 @@ fn reconstruct_exact_agg( type_name: &str, bytes: &[u8], ) -> Option> { + use crate::storage_engines::types::AggregateCore; use asap_physical_operators::accumulators::{ IncreaseAccumulator, KeyedCounterState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, }; - use crate::storage_engines::types::AggregateCore; match type_name { "PlannerExactAccumulatorV1" => asap_physical_operators::accumulators::exact_accumulator::ExactAccumulator::deserialize_from_bytes(bytes).ok().map(|a|Box::new(a) as Box), "SumAccumulator" => SumAccumulator::deserialize_from_bytes(bytes) @@ -6243,8 +6243,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 asap_physical_operators::accumulators::exact_accumulator::ExactAccumulator; use crate::storage_engines::types::{AggregateCore, AggregationType}; + use asap_physical_operators::accumulators::exact_accumulator::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 df7c1f7bb..2d16af788 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 asap_physical_operators::accumulators::SumAccumulator; use crate::storage_engines::types::{AggregationType, StreamingConfig}; + use asap_physical_operators::accumulators::SumAccumulator; use asap_types::aggregation_config::PrecomputeMaterialization; use asap_types::enums::WindowKind; use asap_types::KeyByLabelNames; diff --git a/data_plane/src/storage_engines/types/mod.rs b/data_plane/src/storage_engines/types/mod.rs index 8e18c22d6..562d08d7f 100644 --- a/data_plane/src/storage_engines/types/mod.rs +++ b/data_plane/src/storage_engines/types/mod.rs @@ -12,14 +12,14 @@ pub mod precomputed_output; pub mod storage_backend; pub mod streaming_config; -pub use enums::*; -pub use hot_reload_config::*; 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 precomputed_output::*; pub use storage_backend::*; pub use streaming_config::*; -pub use asap_physical_operators::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/tests/trait_design_tests.rs b/data_plane/src/tests/trait_design_tests.rs index 22ff11181..b5b80a769 100644 --- a/data_plane/src/tests/trait_design_tests.rs +++ b/data_plane/src/tests/trait_design_tests.rs @@ -1,8 +1,8 @@ -use asap_physical_operators::accumulators::{KeyedSumCountAccumulator, SumAccumulator}; #[cfg(test)] use crate::storage_engines::types::{ KeyByLabelValues, MultipleSubpopulationAggregate, SingleSubpopulationAggregate, }; +use asap_physical_operators::accumulators::{KeyedSumCountAccumulator, SumAccumulator}; use asap_types::Statistic; #[test] diff --git a/data_plane/tests/support/univmon_erp_process.rs b/data_plane/tests/support/univmon_erp_process.rs index d849405b9..f9eecb4ef 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 control_plane::physical::erp::ErpShapeObserver; use asap_physical_operators::accumulators::univmon_accumulator::UnivMonAccumulator; +use control_plane::physical::erp::ErpShapeObserver; use data_plane::storage_engines::types::{AggregateCore, SerializableToSink}; fn values(offset: usize) -> Vec { From 3cfc9efa82ca3a02181fca7c6188e8e4b71cc5b6 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 21:05:15 +0000 Subject: [PATCH 09/15] refactor: consume Planner-owned physical operators and composed pruning --- Cargo.lock | 26 +- Cargo.toml | 13 +- control_plane/Cargo.toml | 1 + .../examples/calibration_candidates.rs | 11 +- .../examples/offline_planner_replay.rs | 6 - control_plane/src/emit/mod.rs | 4 +- control_plane/src/physical/compiler.rs | 58 +- control_plane/src/physical/plan_dot.rs | 3 +- control_plane/src/physical/post_asap/tests.rs | 4 +- control_plane/src/query_plan.rs | 237 +- control_plane/src/query_plan/residual.rs | 54 +- crates/asap-physical-operators/Cargo.toml | 27 - crates/asap-physical-operators/README.md | 80 - .../count_min_sketch_accumulator.rs | 1323 ----------- .../count_min_sketch_with_heap_accumulator.rs | 832 ------- .../accumulators/count_sketch_accumulator.rs | 678 ------ .../count_sketch_with_heap_accumulator.rs | 575 ----- .../datasketches_kll_accumulator.rs | 727 ------ .../src/accumulators/dd_sketch_accumulator.rs | 665 ------ .../src/accumulators/exact_accumulator.rs | 326 --- .../accumulators/hll_sketch_accumulator.rs | 788 ------- .../src/accumulators/hydra_kll_accumulator.rs | 165 -- .../src/accumulators/increase_accumulator.rs | 742 ------ .../src/accumulators/keyed_counter_state.rs | 529 ----- .../src/accumulators/keyed_max_state.rs | 335 --- .../src/accumulators/keyed_min_state.rs | 335 --- .../keyed_sum_count_accumulator.rs | 558 ----- .../src/accumulators/max_accumulator.rs | 248 -- .../src/accumulators/min_accumulator.rs | 253 -- .../src/accumulators/mod.rs | 37 - .../sketch_envelope_accumulator.rs | 154 -- .../src/accumulators/sum_accumulator.rs | 413 ---- .../src/accumulators/univmon_accumulator.rs | 234 -- .../asap-physical-operators/src/arithmetic.rs | 19 - .../asap-physical-operators/src/capability.rs | 115 - .../src/dag/batch_execution.rs | 181 -- crates/asap-physical-operators/src/dag/mod.rs | 517 ----- .../src/dag/operators.rs | 1170 ---------- .../src/dag/planner.rs | 479 ---- .../asap-physical-operators/src/dag/tests.rs | 260 --- .../asap-physical-operators/src/dag/values.rs | 327 --- crates/asap-physical-operators/src/factory.rs | 2046 ----------------- .../src/key_by_label_values.rs | 164 -- crates/asap-physical-operators/src/lib.rs | 22 - .../src/measurement.rs | 94 - crates/asap-physical-operators/src/rows.rs | 88 - crates/asap-physical-operators/src/traits.rs | 351 --- .../tests/deployment.rs | 96 - .../tests/physical_dag.rs | 663 ------ crates/asap_sketch_codec/Cargo.toml | 8 - crates/asap_sketch_codec/src/lib.rs | 84 - crates/asap_types/Cargo.toml | 2 + crates/asap_types/src/aggregation_type.rs | 217 +- crates/asap_types/src/derived_input.rs | 2 +- crates/asap_types/src/enums.rs | 82 +- crates/asap_types/src/executable_plan.rs | 4 +- crates/asap_types/src/precompute_plan.rs | 51 +- crates/asap_types/src/query_plan.rs | 23 +- crates/asap_types/src/query_plan/residual.rs | 10 +- crates/asap_types/src/traits.rs | 8 +- data_plane/Cargo.toml | 4 +- .../precompute_engine/maintenance_runtime.rs | 26 +- .../src/precompute_engine/subdag_scheduler.rs | 1 - .../accelerator.rs | 3 +- .../asap_clickhouse_query_engine/execution.rs | 4 + .../relational_adapter.rs | 17 +- .../query_engines/asap_query_engine/engine.rs | 2 +- .../asap_query_engine/exact_subqueries.rs | 46 +- .../asap_query_engine/logical_dag.rs | 212 +- .../logical_dag/native_values.rs | 137 ++ .../asap_query_engine/post_asap_readout.rs | 1 - .../asap_query_engine/summary_exec.rs | 3 +- docs/design_docs/physical-operators.md | 96 +- 73 files changed, 645 insertions(+), 17431 deletions(-) delete mode 100644 crates/asap-physical-operators/Cargo.toml delete mode 100644 crates/asap-physical-operators/README.md delete mode 100644 crates/asap-physical-operators/src/accumulators/count_min_sketch_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/count_min_sketch_with_heap_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/count_sketch_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/count_sketch_with_heap_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/datasketches_kll_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/dd_sketch_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/exact_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/hll_sketch_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/hydra_kll_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/increase_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/keyed_counter_state.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/keyed_max_state.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/keyed_min_state.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/keyed_sum_count_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/max_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/min_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/mod.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/sketch_envelope_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/sum_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/accumulators/univmon_accumulator.rs delete mode 100644 crates/asap-physical-operators/src/arithmetic.rs delete mode 100644 crates/asap-physical-operators/src/capability.rs delete mode 100644 crates/asap-physical-operators/src/dag/batch_execution.rs delete mode 100644 crates/asap-physical-operators/src/dag/mod.rs delete mode 100644 crates/asap-physical-operators/src/dag/operators.rs delete mode 100644 crates/asap-physical-operators/src/dag/planner.rs delete mode 100644 crates/asap-physical-operators/src/dag/tests.rs delete mode 100644 crates/asap-physical-operators/src/dag/values.rs delete mode 100644 crates/asap-physical-operators/src/factory.rs delete mode 100644 crates/asap-physical-operators/src/key_by_label_values.rs delete mode 100644 crates/asap-physical-operators/src/lib.rs delete mode 100644 crates/asap-physical-operators/src/measurement.rs delete mode 100644 crates/asap-physical-operators/src/rows.rs delete mode 100644 crates/asap-physical-operators/src/traits.rs delete mode 100644 crates/asap-physical-operators/tests/deployment.rs delete mode 100644 crates/asap-physical-operators/tests/physical_dag.rs delete mode 100644 crates/asap_sketch_codec/Cargo.toml delete mode 100644 crates/asap_sketch_codec/src/lib.rs create mode 100644 data_plane/src/query_engines/asap_query_engine/logical_dag/native_values.rs diff --git a/Cargo.lock b/Cargo.lock index bc5762e36..3a765ca52 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=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" 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=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" 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=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" dependencies = [ "asap-sql-function-catalog", "asap-types", @@ -396,20 +396,19 @@ dependencies = [ [[package]] name = "asap-physical-operators" version = "0.1.0" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" dependencies = [ "asap-types", "asap_sketch_codec", - "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?branch=main)", - "asap_types", + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=026cd18c7b8c23ae6c46d4d683151ba562b8cd3a)", "base64 0.21.7", "bincode", "futures", - "hex", "prost", "rmp-serde", "serde", "serde_json", - "thiserror 1.0.69", + "thiserror 2.0.20", "tracing", "xxhash-rust", ] @@ -417,12 +416,12 @@ dependencies = [ [[package]] name = "asap-sql-function-catalog" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" [[package]] name = "asap-types" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=09129074c1894f313b98764dd0400ecd73334a2d#09129074c1894f313b98764dd0400ecd73334a2d" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" dependencies = [ "serde", "serde_json", @@ -443,15 +442,16 @@ dependencies = [ [[package]] name = "asap_sketch_codec" version = "0.1.0" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" 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=026cd18c7b8c23ae6c46d4d683151ba562b8cd3a)", "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=026cd18c7b8c23ae6c46d4d683151ba562b8cd3a#026cd18c7b8c23ae6c46d4d683151ba562b8cd3a" dependencies = [ "bytes", "prost", @@ -488,6 +488,7 @@ version = "0.1.0" dependencies = [ "anyhow", "asap-aware-mapping", + "asap-physical-operators", "asap-types", "base64 0.21.7", "clap", @@ -966,6 +967,7 @@ dependencies = [ "asap-aware-mapping", "asap-frontend-promql", "asap-frontend-sql", + "asap-physical-operators", "asap-types", "asap_types", "axum", @@ -1180,7 +1182,7 @@ dependencies = [ "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=026cd18c7b8c23ae6c46d4d683151ba562b8cd3a)", "asap_types", "async-trait", "axum", diff --git a/Cargo.toml b/Cargo.toml index 852c7a476..a4fd489c4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -3,8 +3,6 @@ resolver = "2" members = [ "crates/asap_otel_proto", "crates/asap_types", - "crates/asap_sketch_codec", - "crates/asap-physical-operators", "data_plane", "control_plane", ] @@ -16,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 = "09129074c1894f313b98764dd0400ecd73334a2d" } -asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "09129074c1894f313b98764dd0400ecd73334a2d" } -asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "09129074c1894f313b98764dd0400ecd73334a2d" } -asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "09129074c1894f313b98764dd0400ecd73334a2d" } +planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } +asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } +asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } +asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } # Shared external deps (used by 2+ crates) serde = { version = "1.0", features = ["derive"] } @@ -39,7 +37,8 @@ arc-swap = "1.7" reqwest = { version = "0.12", default-features = false, features = ["json", "rustls-tls"] } # Internal crates -asap-physical-operators = { path = "crates/asap-physical-operators" } +asap-physical-operators = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } +asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } 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 d83c50ad9..1c45cbf57 100644 --- a/control_plane/examples/calibration_candidates.rs +++ b/control_plane/examples/calibration_candidates.rs @@ -37,15 +37,7 @@ fn planner_forest(queries: &[control_plane::physical::compiler::QueryCompilation vec![lhs, rhs], json!({"operator_debug":format!("{operator:?}"),"timing_debug":format!("{timing:?}")}), ), - SummaryExpr::MembershipFilter { - candidates, - values, - completeness, - } => ( - "MembershipFilter", - vec![candidates, values], - json!({"completeness_debug":format!("{completeness:?}")}), - ), + SummaryExpr::ValueOperation { child, operation, @@ -82,6 +74,7 @@ fn planner_forest(queries: &[control_plane::physical::compiler::QueryCompilation right, kind, pred, + .. } => ( "RelationalJoin", vec![left, right], diff --git a/control_plane/examples/offline_planner_replay.rs b/control_plane/examples/offline_planner_replay.rs index f87f8ab8b..78b6e1dd0 100644 --- a/control_plane/examples/offline_planner_replay.rs +++ b/control_plane/examples/offline_planner_replay.rs @@ -82,12 +82,6 @@ fn inspect( inspect(lhs, model, seen, states, raw); inspect(rhs, model, seen, states, raw); } - SummaryExpr::MembershipFilter { - candidates, values, .. - } => { - inspect(candidates, model, seen, states, raw); - inspect(values, model, seen, states, raw); - } } } diff --git a/control_plane/src/emit/mod.rs b/control_plane/src/emit/mod.rs index f1bd209c0..a56536256 100644 --- a/control_plane/src/emit/mod.rs +++ b/control_plane/src/emit/mod.rs @@ -56,9 +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), - SummaryExpr::MembershipFilter { - candidates, values, .. - } => extract_from_node(candidates).or_else(|| extract_from_node(values)), + // 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 fdb12922c..c11eebe47 100644 --- a/control_plane/src/physical/compiler.rs +++ b/control_plane/src/physical/compiler.rs @@ -2056,12 +2056,7 @@ fn summary_agg_metric(node: &SummaryNode) -> Option { } } SummaryExpr::SummaryAgg { child, .. } => walk(child, metrics), - SummaryExpr::MembershipFilter { - candidates, values, .. - } => { - walk(candidates, metrics); - walk(values, metrics); - } + SummaryExpr::ValueOperation { child, .. } => walk(child, metrics), SummaryExpr::SummaryEstimate { summary_input, .. } => walk(summary_input, metrics), SummaryExpr::SummaryMerge { children, .. } => { @@ -2338,12 +2333,7 @@ fn requires_exact_erp_fallback( walk(left, out); walk(right, out); } - SummaryExpr::MembershipFilter { - candidates, values, .. - } => { - walk(candidates, out); - walk(values, out); - } + SummaryExpr::KeepPreAsap(_) => {} } } @@ -3535,13 +3525,15 @@ fn collect_selected_materializations( } } match &node.expr { - SummaryExpr::MembershipFilter { - candidates, values, .. + SummaryExpr::RelationalJoin { + 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)?; } @@ -4442,25 +4434,32 @@ pub(crate) mod tests { .unwrap(); let entry = plan.query_plan.entries.values().next().unwrap(); let crate::query_plan::QueryPlanNode::Logical { - operator: asap_types::query_plan::residual::ResidualQueryOperator::TopKSelection { .. }, + operator: asap_types::query_plan::residual::ResidualQueryOperator::Limit { .. }, inputs, } = &entry.nodes[&entry.root] else { panic!("expected ordinary TopK root") }; - let crate::query_plan::QueryPlanNode::MembershipFilter { inputs, .. } = + let crate::query_plan::QueryPlanNode::Logical { + operator: asap_types::query_plan::residual::ResidualQueryOperator::Sort { .. }, + inputs, + } = &entry.nodes[&inputs[0]] + else { + panic!("expected grouped Sort below Limit") + }; + let crate::query_plan::QueryPlanNode::RelationalJoin { inputs, .. } = &entry.nodes[&inputs[0]] else { - panic!("Planner weighted TopK must lower to MembershipFilter: {entry:#?}"); + panic!("Planner weighted TopK must lower to semi-join: {entry:#?}"); }; assert!(matches!( - entry.nodes[&inputs[0]], + entry.nodes[&inputs[1]], crate::query_plan::QueryPlanNode::SummaryEstimate { query: crate::query_plan::QueryReadout::TopK { .. }, .. } )); - let candidate_read = match &entry.nodes[&inputs[0]] { + let candidate_read = match &entry.nodes[&inputs[1]] { crate::query_plan::QueryPlanNode::SummaryEstimate { input, .. } => *input, _ => unreachable!(), }; @@ -4548,17 +4547,24 @@ pub(crate) mod tests { ); let entry = plan.query_plan.lookup(query).unwrap(); let QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { .. }, + operator: ResidualQueryOperator::Limit { .. }, inputs, } = &entry.nodes[&entry.root] else { panic!("expected ordinary TopK root") }; - let QueryPlanNode::MembershipFilter { inputs, .. } = &entry.nodes[&inputs[0]] else { + let QueryPlanNode::Logical { + operator: ResidualQueryOperator::Sort { .. }, + inputs, + } = &entry.nodes[&inputs[0]] + else { + panic!("expected grouped Sort below Limit") + }; + let QueryPlanNode::RelationalJoin { inputs, .. } = &entry.nodes[&inputs[0]] else { panic!("expected candidate TopK: {entry:#?}"); }; assert!(matches!( - &entry.nodes[&inputs[1]], + &entry.nodes[&inputs[0]], QueryPlanNode::ExternalExact { request, inputs: exact_inputs, @@ -4568,7 +4574,7 @@ pub(crate) mod tests { && request.input_contracts == vec![ExternalExactInput::CandidateMembership { item_label: "job".into(), }] - && exact_inputs == &vec![inputs[0]] + && exact_inputs == &vec![inputs[1]] )); assert!(entry.nodes.values().all(|node| !matches!( node, diff --git a/control_plane/src/physical/plan_dot.rs b/control_plane/src/physical/plan_dot.rs index 19d602413..715aacf7b 100644 --- a/control_plane/src/physical/plan_dot.rs +++ b/control_plane/src/physical/plan_dot.rs @@ -154,7 +154,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 +168,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/tests.rs b/control_plane/src/physical/post_asap/tests.rs index 74bb605a7..deb1b4644 100644 --- a/control_plane/src/physical/post_asap/tests.rs +++ b/control_plane/src/physical/post_asap/tests.rs @@ -116,9 +116,7 @@ fn node_is_archive(node: &Rc) -> bool { SummaryExpr::SummaryJoin { outer, inner, .. } => { node_is_archive(outer) || node_is_archive(inner) } - SummaryExpr::MembershipFilter { - candidates, values, .. - } => node_is_archive(candidates) || node_is_archive(values), + SummaryExpr::SummarySubtract { left, right } | SummaryExpr::RelationalJoin { left, right, .. } | SummaryExpr::BinaryOp { diff --git a/control_plane/src/query_plan.rs b/control_plane/src/query_plan.rs index 40a7e8b38..0a1eb1a56 100644 --- a/control_plane/src/query_plan.rs +++ b/control_plane/src/query_plan.rs @@ -179,8 +179,7 @@ where *input = remap[input]; } } - QueryPlanNode::MembershipFilter { inputs, .. } - | QueryPlanNode::Binary { inputs, .. } + QueryPlanNode::Binary { inputs, .. } | QueryPlanNode::RelationalJoin { inputs, .. } => { for input in inputs { *input = remap[input]; @@ -257,9 +256,11 @@ where right, kind, pred, + 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}" @@ -269,9 +270,6 @@ where right_schema: right.schema.clone(), output_schema: node.schema.clone(), }, - SummaryExpr::RelationalJoin { .. } => QueryPlanNode::ExactFallback { - reason: "read-time relational join requires the relational compiler".into(), - }, SummaryExpr::ValueOperation { child, operation, .. } if self.preserve_relational @@ -317,7 +315,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(), @@ -349,16 +346,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, @@ -366,114 +356,80 @@ where } } SummaryExpr::ValueOperation { - child: sort, - operation: planner_types::post_asap::ValueOperation::Limit { n, offset: 0 }, + child, + operation: + planner_types::post_asap::ValueOperation::Limit { + n, + offset, + partition_by, + }, 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 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 { + } => QueryPlanNode::Logical { + operator: residual::ResidualQueryOperator::Limit { + n: *n as u64, + offset: *offset as u64, + grouping: vector_grouping(partition_by, &child.schema)?, + }, + inputs: vec![self.lower(child)?], + }, + SummaryExpr::ValueOperation { + child, + operation: planner_types::post_asap::ValueOperation::Sort { keys, partition_by }, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, + } if keys.len() == 1 => { + let planner_types::pre_asap::QueryExpr::Column(column) = keys[0].expr else { return Err(QueryPlanError::Invalid( - "TopK sort key must reference the child value column".into(), + "vector Sort requires a value column".into(), )); }; if !matches!( - child - .schema - .fields - .get(*sort_column) - .map(|field| &field.dtype), + child.schema.fields.get(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(), + "vector Sort requires the numeric value column".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(), - }, + operator: residual::ResidualQueryOperator::Sort { + descending: !keys[0].ascending, + grouping: vector_grouping(partition_by, &child.schema)?, }, inputs: vec![self.lower(child)?], } } - 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(), }, - SummaryExpr::MembershipFilter { - candidates, - values, - completeness, + SummaryExpr::RelationalJoin { + right: candidates, + left: values, + kind: planner_types::pre_asap::JoinKind::Semi, + pred, + pruning, } => { + let keys = asap_physical_operators::dag::planner::equijoin_keys( + pred, + &values.schema, + &candidates.schema, + ) + .map_err(|error| QueryPlanError::Invalid(error.to_string()))? + .into_iter() + .map(|(left, right)| { + ( + values.schema.fields[left].name.clone(), + candidates.schema.fields[right].name.clone(), + ) + }) + .collect::>(); 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 item_label = keys[0].1.clone(); let value_id = QueryNodeId(self.next_id); self.next_id += 1; self.nodes.insert( @@ -497,11 +453,20 @@ where } else { self.lower(values)? }; - QueryPlanNode::MembershipFilter { - inputs: [candidate_input, value_input], - completeness: completeness.clone(), + QueryPlanNode::RelationalJoin { + inputs: [value_input, candidate_input], + join_kind: planner_types::pre_asap::JoinKind::Semi, + pred: serde_json::to_value(pred) + .map_err(|error| QueryPlanError::Invalid(error.to_string()))?, + pruning: pruning.clone(), + left_schema: values.schema.clone(), + right_schema: candidates.schema.clone(), + output_schema: node.schema.clone(), } } + SummaryExpr::RelationalJoin { .. } => QueryPlanNode::ExactFallback { + reason: "unsupported join in vector adapter".into(), + }, SummaryExpr::BinaryOp { lhs, rhs, @@ -1357,7 +1322,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(), }; @@ -1370,11 +1335,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 { @@ -1386,9 +1373,12 @@ mod tests { }, provenance: vec![], }, - }, - }, - ), + }), + left_schema: schema.clone(), + right_schema: schema.clone(), + output_schema: schema, + } + }), ]), instant: InstantExecution { lookback_ms: 300_000, @@ -1400,3 +1390,24 @@ mod tests { assert!(entry.validate(&BTreeSet::new()).is_err()); } } + +fn vector_grouping( + keys: &planner_types::pre_asap::GroupKeys, + schema: &planner_types::post_asap::SummarySchema, +) -> Result { + let labels = keys + .keys() + .iter() + .map(|&index| { + schema + .fields + .get(index) + .map(|field| field.name.clone()) + .ok_or_else(|| QueryPlanError::Invalid("unresolved partition column".into())) + }) + .collect::, _>>()?; + Ok(residual::Grouping { + labels, + without: keys.is_without(), + }) +} diff --git a/control_plane/src/query_plan/residual.rs b/control_plane/src/query_plan/residual.rs index 5722d50c2..94e08c7c8 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, @@ -508,11 +528,6 @@ pub(super) fn selected_native_expression( rhs: right, .. } - | SummaryExpr::MembershipFilter { - candidates: left, - values: right, - .. - } | SummaryExpr::RelationalJoin { left, right, .. } | SummaryExpr::SummaryJoin { outer: left, @@ -748,7 +763,7 @@ mod planner_workload_tests { matches!( entry.nodes[&entry.root], QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { .. }, + operator: ResidualQueryOperator::Limit { .. }, .. } ), @@ -776,7 +791,7 @@ mod planner_workload_tests { assert!(matches!( entry.nodes[&entry.root], QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { .. }, + operator: ResidualQueryOperator::Limit { .. }, .. } )); @@ -1122,12 +1137,7 @@ pub fn eligible_materialization_keys( visit(original, left, keys)?; visit(original, right, keys)?; } - SummaryExpr::MembershipFilter { - candidates, values, .. - } => { - visit(original, candidates, keys)?; - visit(original, values, keys)?; - } + SummaryExpr::ValueOperation { child, .. } => visit(original, child, keys)?, SummaryExpr::SummaryAgg { child, .. } => visit(original, child, keys)?, SummaryExpr::SummaryEstimate { summary_input, .. } @@ -1241,7 +1251,7 @@ pub fn externalize_residuals(entry: &mut QueryPlanEntry) -> Result<(), QueryPlan ) || matches!( node, QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { .. }, + operator: ResidualQueryOperator::Limit { .. }, .. } ); @@ -1522,7 +1532,7 @@ mod tests { assert!(matches!( entry.nodes[&entry.root], QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { k: actual, .. }, + operator: ResidualQueryOperator::Limit { n: actual, .. }, .. } if actual == k )); @@ -1541,7 +1551,7 @@ mod tests { assert!(matches!( entry.nodes[&entry.root], QueryPlanNode::Logical { - operator: ResidualQueryOperator::TopKSelection { k: 3, .. }, + operator: ResidualQueryOperator::Limit { n: 3, .. }, .. } )); @@ -1570,7 +1580,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/crates/asap-physical-operators/Cargo.toml b/crates/asap-physical-operators/Cargo.toml deleted file mode 100644 index 601271713..000000000 --- a/crates/asap-physical-operators/Cargo.toml +++ /dev/null @@ -1,27 +0,0 @@ -[package] -name = "asap-physical-operators" -version.workspace = true -edition.workspace = true - -[dependencies] -futures = "0.3" -asap_types.workspace = true -planner-types.workspace = true -asap_sketch_codec = { path = "../asap_sketch_codec" } -asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", branch = "main" } -serde.workspace = true -serde_json.workspace = true -tracing.workspace = true -thiserror.workspace = true -base64 = "0.21" -bincode = "1.3" -rmp-serde = "1.3" -prost = "0.13" -xxhash-rust = { version = "0.8", features = ["xxh32", "xxh64"] } - -[features] -default = [] -extra_debugging = [] - -[dev-dependencies] -hex = "0.4" diff --git a/crates/asap-physical-operators/README.md b/crates/asap-physical-operators/README.md deleted file mode 100644 index bc9a439c0..000000000 --- a/crates/asap-physical-operators/README.md +++ /dev/null @@ -1,80 +0,0 @@ -# ASAP physical operators - -An independent Rust physical operator DAG runtime shared by ingestion time and -query time execution. The library requires neither backend engine, a server, -a storage implementation, Arrow nor DataFusion. DataFusion informed the design; -it is not the execution framework. - -`dag::PhysicalDag` binds typed operator inputs to node IDs. Each execution starts -one producer per reachable node, shares output batches among its consumers, and -bounds buffering. Dropping one consumer does not cancel other consumers. A -`RunContext` carries query or ingestion scope, cancellation and byte accounting. -Executions use the caller's worker and worker-local streams, with no internal -thread pool. Poll multiple root streams concurrently when they share inputs. - -`dag::operators::Operator` implements native batch sources, scalar values, -projection, filtering, grouped exact aggregation, semi-join, grouped Sort and -Limit, vector-to-scalar conversion, Union, and summary construction/merge/readout. -Sort followed by Limit implements grouped ranking; no dedicated TopK physical -operator is needed. Summary construction updates state batch by batch. End of -input means the supplied query range or ingestion window is complete. - -```rust -use asap_physical_operators::dag::{ - operators::{Expression, Operator}, - values::Value, - Limits, PhysicalDag, RunContext, Scope, -}; -use asap_physical_operators::planner::pre_asap::DataType; -use futures::{executor::block_on, StreamExt}; - -let source = Operator::scalar(Value::Int64(7), DataType::Int64)?; -let negate = Operator::project(source.schema(), vec![ - ("value".into(), Expression::Negate(Box::new(Expression::Column(0)))), -])?; -let mut plan = PhysicalDag::default(); -plan.add(0, vec![], source)?; -plan.add(1, vec![0], negate)?; -let run = RunContext::new( - Scope::Query { evaluation_time_ms: 1000, revision: 1 }, - Limits::default(), -)?; -let mut output = plan.execute(&[1], run)?.remove(0); -let batch = block_on(output.next()).unwrap()?; -assert!(matches!(batch.rows()[0][0], Value::Int64(-7))); -# Ok::<(), asap_physical_operators::dag::Error>(()) -``` - -`dag::planner::bind` accepts a post-ASAP DAG and explicit source bindings for -installed ingestion/storage frontiers. It rejects unsupported operations and -schema mismatches before starting a source. Implement `PhysicalOperator` for a -deployment source, including asynchronous I/O; computation operators remain in -the library. The public `planner` export identifies the exact Planner types used -by the crate. The native binder currently supports a subset of those types and -operations; it does not interpret an unknown node as external fallback. - -Plain values preserve Planner scalar/collection types and nullability. Numeric -arithmetic uses matching Int64 or Float64 inputs; integer overflow is an error. -Boolean predicates use three-valued logic. Native summary states currently cover -exact Sum/Count/Min/Max/Rate/Increase, KLL, DDSketch and HLL. Binding checks family, -parameters and readout compatibility; source batches also validate state payloads. -Existing accumulator algorithms are reused as kernels behind these operators. - -Backend ingestion integration is delivered in #763 and query integration in -#765, after this foundation. Installed value/storage adapters provide deployment-specific -computation; they have not all been replaced by native batch bindings. Local raw -Scan remains deferred. See the [shared operator design](../../docs/design_docs/physical-operators.md). -The dependent #765 query DAG design tracks installed engine coverage separately. - -The default limits are eight buffered batches per producer and 64 MiB of estimated -retained execution data. Callers can set both through `Limits`. Accounting includes -consumer-held outputs and reserved operator state, but is not a hard RSS cap or an -allocator hook. Source-owned data and temporary allocation peaks are excluded. -Blocking operators have no spill support. Plan depth is limited to 128. No execution -state is shared between runs, and no implicit fallback or legacy traversal API is -provided. - -Run `cargo test -p asap-physical-operators --locked` for the independent library -acceptance tests, including shared producers, backpressure, cancellation, grouping, -state restoration and raw/partial/fully precomputed DAG examples. These examples -supply in-memory batches; they do not establish backend local raw-Scan support. diff --git a/crates/asap-physical-operators/src/accumulators/count_min_sketch_accumulator.rs b/crates/asap-physical-operators/src/accumulators/count_min_sketch_accumulator.rs deleted file mode 100644 index a1fc7e694..000000000 --- a/crates/asap-physical-operators/src/accumulators/count_min_sketch_accumulator.rs +++ /dev/null @@ -1,1323 +0,0 @@ -use crate::accumulators::dd_sketch_accumulator::normalize_sample_p; -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CountMinSketch, CountMinSketchDelta, MessagePackCodec}; -use serde_json::Value; -use std::collections::HashMap; - -use 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_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; - use prost::Message; - - let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; - - if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { - return Err(format!( - "CountMinDelta packed-array length mismatch: \ - cell_rows={}, cell_cols={}, d_counts={}", - pb.cell_rows.len(), - pb.cell_cols.len(), - pb.d_counts.len() - ) - .into()); - } - let cells = pb - .cell_rows - .iter() - .zip(pb.cell_cols.iter()) - .zip(pb.d_counts.iter()) - .map(|((r, c), dc)| (*r, *c, *dc)) - .collect(); - let delta = CountMinSketchDelta { - rows: pb.rows, - cols: pb.cols, - cells, - l1: pb.l1, - l2: pb.l2, - // The Go-side CountMinDelta proto now carries an hh_keys field - // (heavy-hitter candidates), mirrored on asap_sketchlib's - // CountMinSketchDelta. The vendored Rust proto bindings here don't - // decode it yet, and CountMin has no TopK to rebuild, so pass an - // empty set — same handling as CountSketch's hh_keys. - hh_keys: Vec::new(), - }; - self.inner - .apply_delta(&delta) - .map_err(|e| format!("apply CountMinDelta: {e}"))?; - Ok(()) - } - - pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { - if buffer.len() < 8 { - return Err("Buffer too short for row_num and col_num".into()); - } - - // TODO: this logic will need to be checked for i32 -> f64 - // Github Issue #11 - - let row_num = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; - let col_num = u32::from_le_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]) as usize; - - let expected_size = 8 + (row_num * col_num * 4); - if buffer.len() < expected_size { - return Err("Buffer too short for sketch data".into()); - } - - let mut sketch = Vec::new(); - let mut offset = 8; - - for _ in 0..row_num { - let mut row = Vec::new(); - for _ in 0..col_num { - let value = f64::from_le_bytes([ - buffer[offset], - buffer[offset + 1], - buffer[offset + 2], - buffer[offset + 3], - buffer[offset + 4], - buffer[offset + 5], - buffer[offset + 6], - buffer[offset + 7], - ]); - row.push(value); - offset += 8; - } - sketch.push(row); - } - - Ok(Self { - inner: CountMinSketch::from_legacy_matrix(sketch, row_num, col_num), - sample_p: 1.0, - }) - } - - /// Merge multiple accumulators efficiently without cloning all of them. - pub fn merge_multiple( - accumulators: &[Box], - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut cms_accumulators = Vec::with_capacity(accumulators.len()); - for acc in accumulators { - if acc.get_accumulator_type() != AggregationType::CountMinSketch { - return Err(format!( - "Cannot merge CountMinSketchAccumulator with {:?}", - acc.get_accumulator_type() - ) - .into()); - } - let cms_acc = acc - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to CountMinSketchAccumulator")?; - cms_accumulators.push(cms_acc); - } - - // Check dimensions are consistent - let rows = cms_accumulators[0].inner.rows(); - let cols = cms_accumulators[0].inner.cols(); - for acc in &cms_accumulators { - if acc.inner.rows() != rows || acc.inner.cols() != cols { - return Err( - "Cannot merge CountMinSketch accumulators with different dimensions".into(), - ); - } - } - - let inner_refs: Vec<&CountMinSketch> = - cms_accumulators.iter().map(|acc| &acc.inner).collect(); - let merged_inner = CountMinSketch::merge_refs(&inner_refs)?; - // sample_p is a per-series config constant, so all operands carry the - // same value in practice. Mirror DDSketch's merge policy: prefer a - // sampled factor (< 1.0) over the no-sampling default so a merge with - // a freshly-reset (1.0) base keeps the series' sampling rate. - let sample_p = cms_accumulators - .iter() - .map(|acc| acc.sample_p) - .find(|&p| p < 1.0) - .unwrap_or(cms_accumulators[0].sample_p); - Ok(Self { - inner: merged_inner, - sample_p, - }) - } -} - -/// Defensive upper bound on the number of matrix cells (`rows * cols`) -/// we'll reconstruct from an inbound wire-declared CMS / CountSketch -/// dimension pair. A malformed / hostile payload could declare absurd -/// dims (e.g. `rows = cols = u32::MAX`) and trick the decoder into a -/// huge `Vec` allocation before the `counts_*.len() != rows*cols` -/// check ever runs. Realistic sketches are at most a few hundred rows -/// by tens-of-thousands of columns, so 8M cells (~64 MiB of f64) is a -/// generous ceiling that no legitimate producer reaches. -pub(crate) const MAX_SKETCH_CELLS: usize = 8 * 1024 * 1024; - -/// Validate an inbound, wire-declared `(rows, cols)` pair for a -/// matrix-backed frequency sketch (CMS / CountSketch) BEFORE any matrix -/// is reconstructed from it. Returns `Ok(())` for dimensions a -/// legitimate producer could have emitted, and an `Err` (never a panic) -/// for malformed / degenerate ones so the ingest path can skip the data -/// point and fall through to its existing decode-failure accounting. -/// -/// Rejections: -/// 1. `rows < 1` or `cols < 1` — a zero-dim matrix has no cells. -/// 2. Narrow-hash-budget violation. The cross-language wire hasher -/// (`sketchlib`'s `MatrixHashType::Packed64`) derives every row's -/// column index from disjoint bit-fields of a single 64-bit hash -/// word: row `r` reads `mask_bits = ceil(log2(cols))` bits at offset -/// `r * mask_bits`. Once `rows * mask_bits > 64` the per-row column -/// slices overflow / alias the 64-bit word and the matrix-cell -/// layout is no longer the one the producer hashed into — the sketch -/// is internally degenerate. This mirrors sketchlib's own -/// `MatrixFastHash::assert_compatible` budget (`rows * (mask_bits + 1) <= 64`); we check the column-index bits alone so realistic -/// configs (5x2048, 5x4096, 5x2000) — for which the sign bits share -/// the top of the word without affecting the cell layout — still -/// pass. -/// 3. Obviously-oversized dims: `rows * cols > MAX_SKETCH_CELLS`, -/// guarding against a huge allocation from a malformed payload. -/// -/// `what` names the wire struct for the error message (e.g. -/// `"CountMinState"`). -pub(crate) fn validate_sketch_dims(what: &str, rows: usize, cols: usize) -> Result<(), String> { - if rows < 1 || cols < 1 { - return Err(format!( - "{what} has degenerate dims (rows={rows}, cols={cols}); rejecting" - )); - } - // mask_bits = ceil(log2(cols)); cols >= 1 here. ilog2 is floor(log2). - let mask_bits = if cols.is_power_of_two() { - cols.ilog2() as usize - } else { - cols.ilog2() as usize + 1 - }; - if rows.saturating_mul(mask_bits) > 64 { - return Err(format!( - "{what} dims (rows={rows}, cols={cols}) exceed the 64-bit \ - packed-hash column budget (rows * ceil(log2(cols)) = {} > 64); \ - the sketch's matrix-cell layout is degenerate, rejecting", - rows.saturating_mul(mask_bits) - )); - } - if rows.saturating_mul(cols) > MAX_SKETCH_CELLS { - return Err(format!( - "{what} dims (rows={rows}, cols={cols}) declare {} cells, \ - exceeding the {MAX_SKETCH_CELLS}-cell ingest cap; rejecting to \ - avoid a huge allocation from a malformed payload", - rows.saturating_mul(cols) - )); - } - Ok(()) -} - -impl SerializableToSink for CountMinSketchAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "row_num": self.inner.rows(), - "col_num": self.inner.cols(), - "sketch": self.inner.sketch() - }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl AggregateCore for CountMinSketchAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "CountMinSketchAccumulator" - } - - /// Per-window base rotation: rebuild an empty counter matrix with - /// the same (rows, cols) so the next window's additive cell deltas - /// align to the identical hash geometry. `sample_p` is a per-series - /// config constant (not per-window data), so it is intentionally - /// preserved across the rotation — mirrors `DDSketchAccumulator`. - fn reset_to_empty(&mut self) { - self.inner = CountMinSketch::new(self.inner.rows(), self.inner.cols()); - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } - - fn merge_with( - &self, - other: &dyn AggregateCore, - ) -> Result, Box> { - if other.get_accumulator_type() != self.get_accumulator_type() { - return Err(format!( - "Cannot merge CountMinSketchAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - - let other_cms = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to CountMinSketchAccumulator")?; - - let merged_inner = CountMinSketch::merge_refs(&[&self.inner, &other_cms.inner])?; - // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a - // per-series config constant, so both operands carry the same value - // in practice. Prefer a sampled factor over the no-sampling default - // so a merge with a freshly-reset (1.0) base keeps the series' - // sampling rate. - let sample_p = if self.sample_p < 1.0 { - self.sample_p - } else { - other_cms.sample_p - }; - Ok(Box::new(Self { - inner: merged_inner, - sample_p, - })) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::CountMinSketch - } - - fn approx_memory_bytes(&self) -> usize { - // Conservative constant for the CountMinSketch counter matrix. - // Real per-instance sizing would require exposing rows/cols on - // the inner sketch; 16 KiB is a reasonable v1 default. - 16 * 1024 - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: asap_types::Statistic, - key: &Option, - query_kwargs: &std::collections::HashMap, - ) -> Result> { - use crate::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::accumulators::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_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; - use prost::Message; - - let mut acc = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], - 2, - 3, - ), - sample_p: 1.0, - }; - let bytes = PbDelta { - rows: 2, - cols: 3, - cell_rows: vec![0, 1], - cell_cols: vec![0, 2], - d_counts: vec![10, 100], - l1: vec![], - l2: vec![], - ..Default::default() - } - .encode_to_vec(); - - acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); - assert_eq!( - acc.inner.sketch(), - vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 106.0]] - ); - } - - #[test] - fn test_apply_proto_delta_bytes_rejects_garbage() { - let mut acc = CountMinSketchAccumulator::new(2, 3); - assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); - } - - // ---------------------------------------------------------------- - // Statistic::Rate / Statistic::Increase — PR #111 honest-gap closure. - // CMS records insert counts but not timestamps. The Rate readout - // requires the engine to push `range_ms` via query_kwargs; without - // it the accumulator falls back to the raw event count (units of - // events/window) so the answer is at least non-empty. - // ---------------------------------------------------------------- - - #[test] - fn test_query_statistic_rate_with_range_ms() { - // Build a CMS whose min-row-sum is 100 events. With a 5-minute - // (300_000 ms) range, the per-second rate is 100 / 300 ≈ 0.333. - let cms = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix( - vec![vec![100.0, 0.0], vec![100.0, 0.0]], - 2, - 2, - ), - sample_p: 1.0, - }; - let mut kwargs = HashMap::new(); - kwargs.insert("range_ms".to_string(), "300000".to_string()); - let trait_obj: &dyn AggregateCore = &cms; - let v = trait_obj - .query_statistic(Statistic::Rate, &None, &kwargs) - .expect("Rate with range_ms is supported"); - assert!( - (v - (100.0 / 300.0)).abs() < 1e-9, - "expected 100/300 = {}, got {v}", - 100.0 / 300.0, - ); - } - - #[test] - fn test_query_statistic_rate_without_range_ms_falls_back_to_count() { - // Without `range_ms` in kwargs the accumulator returns the raw - // event volume (events/window units). Caller is responsible for - // surfacing that caveat to the user; this avoids `status=error` - // for instant rate-shape queries that bypass the matrix-selector - // code path. - let cms = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix(vec![vec![42.0, 0.0], vec![42.0, 0.0]], 2, 2), - sample_p: 1.0, - }; - let trait_obj: &dyn AggregateCore = &cms; - let v = trait_obj - .query_statistic(Statistic::Rate, &None, &HashMap::new()) - .expect("Rate without range_ms still answers (fallback)"); - assert_eq!(v, 42.0); - } - - #[test] - fn test_query_statistic_increase_returns_total_count() { - // Increase semantics on CMS: total events in the window — the - // same min-row-sum as Sum / Count. Differs from Rate only in - // that it never divides by range. - let cms = CountMinSketchAccumulator { - inner: CountMinSketch::from_legacy_matrix(vec![vec![5.0, 7.0], vec![3.0, 9.0]], 2, 2), - sample_p: 1.0, - }; - let trait_obj: &dyn AggregateCore = &cms; - let v = trait_obj - .query_statistic(Statistic::Increase, &None, &HashMap::new()) - .expect("Increase is supported"); - // min-row-sum: row0 = 12, row1 = 12, min = 12. - assert_eq!(v, 12.0); - } - - // ---------------------------------------------------------------- - // Defensive inbound-dimension validation (harden/sketch-dim-validation). - // Malformed / degenerate / narrow-hash-budget-violating CMS dims must - // be rejected gracefully (Err, never a panic); valid configs the - // backend actually uses (5x2048, 5x4096, 5x2000) must still decode. - // ---------------------------------------------------------------- - - /// Build a bare `CountMinState` proto carrying the given dims and a - /// row-major INT64 counts vector sized to `rows*cols` so that, IF the - /// dims pass validation, the reshape also succeeds. Used to prove a - /// malformed-dim payload is rejected at the dim gate, not later. - fn cms_state_bytes(rows: u32, cols: u32) -> Vec { - use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; - use prost::Message; - let n = (rows as usize).saturating_mul(cols as usize); - let state = CountMinState { - rows, - cols, - counter_type: CounterType::Int64 as i32, - counts_int: vec![0i64; n], - counts_float: Vec::new(), - sum_counts: Vec::new(), - sum2_counts: Vec::new(), - l1: Vec::new(), - l2: Vec::new(), - }; - state.encode_to_vec() - } - - #[test] - fn test_validate_sketch_dims_accepts_valid_configs() { - // The realistic configs the backend uses must pass unchanged. - for (r, c) in [(5usize, 2048usize), (5, 4096), (5, 2000), (4, 1000), (2, 3)] { - assert!( - validate_sketch_dims("CountMinState", r, c).is_ok(), - "valid config {r}x{c} was wrongly rejected" - ); - } - } - - #[test] - fn test_validate_sketch_dims_rejects_malformed() { - // Zero dims. - assert!(validate_sketch_dims("CountMinState", 0, 2048).is_err()); - assert!(validate_sketch_dims("CountMinState", 5, 0).is_err()); - // Narrow-hash-budget violation: 5 * ceil(log2(8192))=5*13=65 > 64. - let err = validate_sketch_dims("CountMinState", 5, 8192).unwrap_err(); - assert!(err.contains("budget"), "expected budget error, got: {err}"); - // Absurdly oversized: 1 x 16,777,216 = 16M cells > 8M cap. (1 row - // keeps the hash budget tiny — 1*24=24 — so the cap check, not the - // budget check, is what fires here.) - let err = validate_sketch_dims("CountMinState", 1, 16_777_216).unwrap_err(); - assert!(err.contains("cap"), "expected cell-cap error, got: {err}"); - // No panic on extreme dims (saturating_mul guards the products). - assert!(validate_sketch_dims("CountMinState", usize::MAX, usize::MAX).is_err()); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { - // A data point declaring narrow-hash-budget-violating dims must be - // skipped (Err returned, NOT a panic). The ingest caller turns - // this Err into a dropped data point + WARN log. - let bytes = cms_state_bytes(5, 8192); - let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "budget-violating dims should be rejected"); - assert!(result.unwrap_err().to_string().contains("rejecting")); - - // A valid neighbour (5x4096) on the same path still decodes fine. - let ok_bytes = cms_state_bytes(5, 4096); - let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) - .expect("valid 5x4096 CMS should still decode"); - assert_eq!(acc.inner.rows(), 5); - assert_eq!(acc.inner.cols(), 4096); - } - - #[test] - fn test_query_statistic_rate_rejects_invalid_range_ms() { - let cms = CountMinSketchAccumulator::new(2, 2); - let mut kwargs = HashMap::new(); - kwargs.insert("range_ms".to_string(), "0".to_string()); - let trait_obj: &dyn AggregateCore = &cms; - let err = trait_obj - .query_statistic(Statistic::Rate, &None, &kwargs) - .expect_err("range_ms=0 should error"); - assert!(err.to_string().contains("positive")); - - let mut kwargs = HashMap::new(); - kwargs.insert("range_ms".to_string(), "not-a-number".to_string()); - let err = trait_obj - .query_statistic(Statistic::Rate, &None, &kwargs) - .expect_err("non-numeric range_ms should error"); - assert!(err.to_string().contains("bad range_ms")); - } - - // ---------------------------------------------------------------- - // sample_p rescale. The edge admits each insert with probability `p`, - // so every stored cell is ~p× the true count. CMS is L1/additive and - // linear, so BOTH the point-frequency (query_key) and the aggregate - // total-event statistics (Count/Sum/Increase/Rate) rescale by 1/p. - // ---------------------------------------------------------------- - - #[test] - fn test_query_key_rescaled_by_sample_p() { - // Same stored cell counts, two sample_p values: the p=0.25 sketch - // must report 4× the point-frequency of the unsampled one. - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let mut unsampled = CountMinSketchAccumulator::new(4, 1000); - unsampled._update(&key, 10.0); - let mut sampled = CountMinSketchAccumulator::new(4, 1000); - sampled._update(&key, 10.0); - sampled.sample_p = 0.25; - - let raw = unsampled.query_key(&key); - let rescaled = sampled.query_key(&key); - assert!( - raw >= 10.0, - "raw estimate should be >= inserted 10, got {raw}" - ); - assert!( - (rescaled - raw * 4.0).abs() < 1e-9, - "expected point-frequency rescaled ≈ 4×raw ({}), got {rescaled}", - raw * 4.0 - ); - } - - #[test] - fn test_aggregate_statistics_rescaled_by_sample_p() { - use 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/crates/asap-physical-operators/src/accumulators/count_min_sketch_with_heap_accumulator.rs b/crates/asap-physical-operators/src/accumulators/count_min_sketch_with_heap_accumulator.rs deleted file mode 100644 index 259ab9d21..000000000 --- a/crates/asap-physical-operators/src/accumulators/count_min_sketch_with_heap_accumulator.rs +++ /dev/null @@ -1,832 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{CmsHeapItem, CountMinSketchWithHeap, MessagePackCodec}; -use serde::Deserialize; -use serde_json::Value; -use std::collections::HashMap; - -use 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::MultipleSubpopulationAggregate; - let key_val = key - .as_ref() - .ok_or("Key required for CountMinSketchWithHeapAccumulator")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for CountMinSketchWithHeapAccumulator { - fn query( - &self, - _statistic: Statistic, - key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - Ok(self.query_key(key)) - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for CountMinSketchWithHeapAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - let mut iter = accumulators.into_iter(); - let mut merged = iter.next().unwrap(); - for acc in iter { - merged.inner.merge(&acc.inner)?; - } - Ok(merged) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_count_min_sketch_with_heap_creation() { - let cms = CountMinSketchWithHeapAccumulator::new(4, 1000, 20); - assert_eq!(cms.inner.rows(), 4); - assert_eq!(cms.inner.cols(), 1000); - assert_eq!(cms.inner.heap_size, 20); - assert_eq!(cms.inner.topk_heap_items().len(), 0); - } - - #[test] - fn test_count_min_sketch_with_heap_query() { - let cms = CountMinSketchWithHeapAccumulator::new(2, 10, 5); - let key = KeyByLabelValues::new(); - assert_eq!(cms.query_key(&key), 0.0); - - let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; - assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); - } - - #[test] - fn test_count_min_sketch_with_heap_merge() { - // Build controlled state via from_legacy_matrix (works regardless of backend config). - let sketch1 = vec![ - vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - ]; - let heap1 = vec![ - CmsHeapItem { - key: "key1".to_string(), - value: 100.0, - }, - CmsHeapItem { - key: "key2".to_string(), - value: 50.0, - }, - ]; - let sketch2 = vec![ - vec![5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], - ]; - let heap2 = vec![ - CmsHeapItem { - key: "key3".to_string(), - value: 75.0, - }, - CmsHeapItem { - key: "key1".to_string(), - value: 80.0, - }, - ]; - - let cms1 = CountMinSketchWithHeapAccumulator { - inner: CountMinSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), - }; - let cms2 = CountMinSketchWithHeapAccumulator { - inner: CountMinSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), - }; - - let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); - assert!(result.is_ok()); - let merged = result.unwrap(); - assert_eq!(merged.inner.sketch_matrix()[0][0], 15.0); - assert_eq!(merged.inner.sketch_matrix()[1][1], 35.0); - assert_eq!(merged.inner.heap_size, 3); - assert!(merged.inner.topk_heap_items().len() <= 3); - } - - #[test] - fn test_count_min_sketch_with_heap_merge_single() { - let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); - let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms.clone()]); - assert!(result.is_ok()); - let merged = result.unwrap(); - assert_eq!(merged.inner.rows(), cms.inner.rows()); - assert_eq!(merged.inner.cols(), cms.inner.cols()); - assert_eq!(merged.inner.heap_size, cms.inner.heap_size); - } - - #[test] - fn test_count_min_sketch_with_heap_merge_dimension_mismatch() { - let cms1 = CountMinSketchWithHeapAccumulator::new(2, 10, 5); - let cms2 = CountMinSketchWithHeapAccumulator::new(3, 10, 5); - let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("dimension")); - } - - #[test] - fn test_count_min_sketch_with_heap_as_aggregate_core() { - let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); - assert_eq!(cms.type_name(), "CountMinSketchWithHeapAccumulator"); - } - - #[test] - fn test_get_topk_keys() { - let mut cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); - cms.inner.update("label1;label2", 100.0); - cms.inner.update("label3;label4", 50.0); - - let keys = cms.get_topk_keys(); - assert_eq!(keys.len(), 2); - // Top-k order can differ between Legacy and Sketchlib backends (heap ordering / estimates). - let label_sets: std::collections::HashSet<_> = - keys.iter().map(|k| k.labels.clone()).collect(); - assert!(label_sets.contains(&vec!["label1".to_string(), "label2".to_string()])); - assert!(label_sets.contains(&vec!["label3".to_string(), "label4".to_string()])); - } - - #[test] - fn test_multiple_subpopulation_aggregate() { - let cms = CountMinSketchWithHeapAccumulator::new(3, 50, 10); - let key = KeyByLabelValues::new(); - - let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; - let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); - assert_eq!(result, 0.0); - - let keys = multi_trait.get_keys(); - assert!(keys.is_some()); - assert_eq!(keys.unwrap().len(), 0); - } - - // ---------------------------------------------------------------- - // DELTA-HEAP wire form (encoding MSGPACK_DELTA): apply a sparse matrix - // delta + replace the heap, decoded generically (rmp_serde) WITHOUT any - // asap_sketchlib delta API. The first test feeds a frame produced by the - // Go encoder (sketchlib-go `MarshalCountSketchWithHeapDelta`) to prove - // cross-language byte parity — mirrors how the full-heap parity is - // proven. The second proves PWR full -> delta -> delta reconstruction. - // ---------------------------------------------------------------- - - /// Cross-language byte-parity: this hex is the exact output of - /// sketchlib-go's `asapmsgpack.MarshalCountSketchWithHeapDelta(5, 1024, - /// cells=[(0,1,50),(1,3,-4),(4,1023,1_000_000)], - /// heap=[("/checkout",50),("/cart",20)], heap_size=20)` (captured via a - /// throw-away Go print test, identical methodology to the full-heap - /// golden in `sketchlib-go/.../count_sketch_with_heap_test.go`). If the - /// Go encoder or the rmp_serde layout ever shifts, this decode fails - /// loudly. - const GO_DELTA_HEAP_GOLDEN_HEX: &str = "94c39305cd04009393000132930103fc9304cd03ffce000f42409292a92f636865636b6f7574cb404900000000000092a52f63617274cb403400000000000014"; - - #[test] - fn test_apply_go_produced_delta_heap_frame_matrix_and_heap() { - let bytes = hex::decode(GO_DELTA_HEAP_GOLDEN_HEX).expect("hex"); - - // Base = empty heap accumulator with the frame's dims (what the - // ingest caller holds after the per-window base rotation). - let mut acc = CountMinSketchWithHeapAccumulator::new(5, 1024, 20); - acc.apply_msgpack_heap_delta_bytes(&bytes) - .expect("apply Go delta-heap frame"); - - // Matrix: the three sparse cells landed onto the empty base. - let m = acc.inner.sketch_matrix(); - assert_eq!(m.len(), 5); - assert_eq!(m[0].len(), 1024); - assert_eq!(m[0][1], 50.0, "cell (0,1)"); - assert_eq!(m[1][3], -4.0, "cell (1,3)"); - assert_eq!(m[4][1023], 1_000_000.0, "cell (4,1023)"); - // Everything else stays zero. - assert_eq!(m[2][2], 0.0); - assert_eq!(m[0][0], 0.0); - - // Heap: the frame's full heap, with /checkout ranked above /cart. - let mut items = acc.inner.topk_heap_items(); - items.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap()); - assert_eq!(items.len(), 2); - assert_eq!(items[0].key, "/checkout"); - assert_eq!(items[0].value, 50.0); - assert_eq!(items[1].key, "/cart"); - assert_eq!(items[1].value, 20.0); - } - - #[test] - fn test_pwr_full_then_delta_then_delta_reconstructs_per_window() { - use asap_sketchlib::MessagePackCodec; - - // Window 1 (full frame): build a heap-bearing CountSketch with mass - // and serialize the FULL `{sketch,topk_heap,heap_size}` frame, then - // decode it into a heap accumulator (the cached per-series base). - let w1 = CountMinSketchWithHeap::from_legacy_matrix( - vec![vec![300.0; 4]; 5], - vec![CmsHeapItem { - key: "k".into(), - value: 300.0, - }], - 5, - 4, - 20, - ); - let w1_bytes = w1.to_msgpack().expect("w1 full msgpack"); - let mut base = CountMinSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(&w1_bytes) - .expect("decode w1 full frame as heap accumulator"); - assert_eq!(base.inner.sketch_matrix()[0][0], 300.0); - - // Window 2 delta: this window's own state is matrix cells of value 50 - // against an EMPTY base + heap {k:50}. The DELTA-HEAP frame is encoded - // the same way the Go producer does (4-array, is_delta, sparse cells). - let w2_frame = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); - // PWR: rotate base to empty at the window boundary, then apply. - base.reset_to_empty(); - assert_eq!( - base.inner.sketch_matrix()[0][0], - 0.0, - "reset_to_empty cleared matrix" - ); - base.apply_msgpack_heap_delta_bytes(&w2_frame) - .expect("apply w2 delta"); - assert_eq!(base.inner.sketch_matrix()[0][0], 50.0, "window-2 cell"); - assert_eq!(base.inner.sketch_matrix()[1][1], 50.0); - // No cross-window leakage from window 1's 300s. - assert_eq!(base.inner.sketch_matrix()[2][2], 0.0); - let h2: Vec<_> = base.inner.topk_heap_items(); - assert_eq!(h2.len(), 1); - assert_eq!(h2[0].key, "k"); - assert_eq!(h2[0].value, 50.0); - - // Window 3 delta: 80s against empty + heap {k:80}. - let w3_frame = encode_delta_heap(5, 4, &[(0, 0, 80)], &[("k", 80.0)], 20); - base.reset_to_empty(); - base.apply_msgpack_heap_delta_bytes(&w3_frame) - .expect("apply w3 delta"); - assert_eq!(base.inner.sketch_matrix()[0][0], 80.0, "window-3 cell"); - assert_eq!(base.inner.sketch_matrix()[1][1], 0.0, "no window-2 leakage"); - let h3 = base.inner.topk_heap_items(); - assert_eq!(h3.len(), 1); - assert_eq!(h3[0].value, 80.0); - } - - #[test] - fn test_rmp_serde_layout_is_byte_identical_to_go_encoder() { - // The rmp_serde positional encoding of the delta-heap frame must be - // BYTE-IDENTICAL to sketchlib-go's hand-rolled - // `MarshalCountSketchWithHeapDelta`. This hex is the Go encoder's - // output for (5, 4, cells=[(0,0,50),(1,1,50)], heap=[("k",50)], - // heap_size=20) — the same inputs `encode_delta_heap` uses below. - // Equality here proves both encode AND decode are cross-language - // byte-compatible (the decode path is exercised by the Go-golden - // test above). - const GO_PARITY_HEX: &str = "94c39305049293000032930101329192a16bcb404900000000000014"; - let rust_bytes = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); - assert_eq!(hex::encode(&rust_bytes), GO_PARITY_HEX); - } - - #[test] - fn test_apply_delta_rejects_full_frame_and_garbage() { - use asap_sketchlib::MessagePackCodec; - let mut acc = CountMinSketchWithHeapAccumulator::new(2, 4, 5); - // A FULL frame (3-array, no is_delta marker) must NOT decode as a - // delta — the routing relies on the two shapes being distinct. - let full = CountMinSketchWithHeap::from_legacy_matrix( - vec![vec![1.0; 4]; 2], - vec![CmsHeapItem { - key: "a".into(), - value: 1.0, - }], - 2, - 4, - 5, - ) - .to_msgpack() - .unwrap(); - assert!(acc.apply_msgpack_heap_delta_bytes(&full).is_err()); - assert!(acc.apply_msgpack_heap_delta_bytes(b"not msgpack").is_err()); - } - - /// Encode a DELTA-HEAP frame the same way sketchlib-go's - /// `MarshalCountSketchWithHeapDelta` does (rmp_serde positional layout), - /// so the test exercises the real decode path. Tuple structs serialize - /// as msgpack fixed arrays — byte-identical to the Go hand-rolled writer. - fn encode_delta_heap( - rows: u32, - cols: u32, - cells: &[(u32, u32, i64)], - heap: &[(&str, f64)], - heap_size: u64, - ) -> Vec { - #[derive(serde::Serialize)] - struct W<'a>( - bool, - (u32, u32, &'a [(u32, u32, i64)]), - Vec<(String, f64)>, - u64, - ); - let heap_owned: Vec<(String, f64)> = - heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); - let w = W(true, (rows, cols, cells), heap_owned, heap_size); - rmp_serde::to_vec(&w).expect("encode delta-heap") - } - - // ---------------------------------------------------------------- - // FIX 1 — VALUE-WEIGHTED top-k (recall 0 → correct). - // - // `topk(k, sum by (host) (cpu_load))` asks for the top-k hosts by - // SUM OF VALUE. The heavy-hitter heap built by the default `+1`-per- - // occurrence update ranks by COUNT keyed by `item`, so its recall - // against the value-weighted ground truth is 0 when the busiest host - // (most samples) is NOT the heaviest host (largest Σvalue). - // `insert_value(group_label, value)` adds the sample VALUE keyed by the - // GROUP LABEL, so `topk_by_value` ranks by Σvalue — correct recall. - // ---------------------------------------------------------------- - - /// Crafted adversarial dataset: the host with the MOST samples - /// (`h_chatty`, 100 tiny samples) is NOT the host with the largest - /// value-sum (`h_heavy`, a handful of huge samples). A COUNT-ranked - /// heap would surface `h_chatty`; the value-weighted top-k must surface - /// the true heavy hitters by Σvalue, giving recall 1.0 against the - /// ground-truth top-k-by-value-sum. - #[test] - fn value_weighted_topk_has_full_recall_vs_count_topk() { - // (host, per-sample value, sample count) → true Σvalue: - // h_heavy : 1000 × 3 = 3000 (few samples, huge value) - // h_mid : 200 × 5 = 1000 - // h_small : 50 × 6 = 300 - // h_chatty: 1 × 100 = 100 (MOST samples, tiny value) - let data: &[(&str, f64, usize)] = &[ - ("h_heavy", 1000.0, 3), - ("h_mid", 200.0, 5), - ("h_small", 50.0, 6), - ("h_chatty", 1.0, 100), - ]; - - // Wide CMS + heap large enough to hold every group exactly (4 groups) - // so the estimate equals the true Σvalue with no hash collisions. - let mut acc = CountMinSketchWithHeapAccumulator::new(5, 4096, 16); - let mut truth: std::collections::HashMap<&str, f64> = std::collections::HashMap::new(); - for (host, value, count) in data { - for _ in 0..*count { - acc.insert_value(host, *value); - } - *truth.entry(*host).or_insert(0.0) += value * (*count as f64); - } - - // Ground-truth top-2 by value-sum: h_heavy (3000), h_mid (1000). - let mut truth_ranked: Vec<(&str, f64)> = truth.into_iter().collect(); - truth_ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); - let truth_top2: std::collections::HashSet<&str> = - truth_ranked.iter().take(2).map(|(k, _)| *k).collect(); - assert!( - truth_top2.contains("h_heavy") && truth_top2.contains("h_mid"), - "ground-truth top-2 by value-sum should be h_heavy + h_mid" - ); - - // Value-weighted top-2 from the heap. - let got = acc.topk_by_value(2); - assert_eq!(got.len(), 2, "k=2 → two groups: {got:?}"); - let got_keys: std::collections::HashSet<&str> = - got.iter().map(|(k, _)| k.as_str()).collect(); - - // RECALL = |got ∩ truth| / |truth| must be 1.0. - let hits = got_keys.intersection(&truth_top2).count(); - let recall = hits as f64 / truth_top2.len() as f64; - assert_eq!( - recall, 1.0, - "value-weighted top-k recall must be 1.0 (count-ranked heap would \ - surface h_chatty and miss h_heavy → recall < 1): got={got:?}" - ); - - // The busiest-by-count host (h_chatty) must NOT be in the top-2, - // proving we rank by value-sum, not occurrence count. - assert!( - !got_keys.contains("h_chatty"), - "h_chatty (most samples, smallest value-sum) must be excluded: {got:?}" - ); - - // Estimates are exact here (no collisions, heap holds all groups): - // top-1 must be h_heavy with Σvalue 3000. - assert_eq!(got[0].0, "h_heavy"); - assert!( - (got[0].1 - 3000.0).abs() < 1e-6, - "h_heavy value-sum estimate ≈ 3000, got {}", - got[0].1 - ); - assert_eq!(got[1].0, "h_mid"); - assert!( - (got[1].1 - 1000.0).abs() < 1e-6, - "h_mid value-sum estimate ≈ 1000, got {}", - got[1].1 - ); - } - - /// A single value-weighted insert must put the full value (not +1) into - /// the heap, and repeated inserts for the same group must accumulate. - #[test] - fn insert_value_accumulates_summed_value_in_heap() { - let mut acc = CountMinSketchWithHeapAccumulator::new(4, 1024, 8); - acc.insert_value("g", 10.0); - acc.insert_value("g", 25.0); - let top = acc.topk_by_value(1); - assert_eq!(top.len(), 1); - assert_eq!(top[0].0, "g"); - assert!( - (top[0].1 - 35.0).abs() < 1e-6, - "summed value should be 35 (10+25), got {}", - top[0].1 - ); - } -} diff --git a/crates/asap-physical-operators/src/accumulators/count_sketch_accumulator.rs b/crates/asap-physical-operators/src/accumulators/count_sketch_accumulator.rs deleted file mode 100644 index c12eda0c3..000000000 --- a/crates/asap-physical-operators/src/accumulators/count_sketch_accumulator.rs +++ /dev/null @@ -1,678 +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::{ - 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::accumulators::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_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; - use prost::Message; - - let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; - - if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { - return Err(format!( - "CountSketchDelta packed-array length mismatch: \ - cell_rows={}, cell_cols={}, d_counts={}", - pb.cell_rows.len(), - pb.cell_cols.len(), - pb.d_counts.len() - ) - .into()); - } - let cells = pb - .cell_rows - .iter() - .zip(pb.cell_cols.iter()) - .zip(pb.d_counts.iter()) - .map(|((r, c), dc)| (*r, *c, *dc)) - .collect(); - // This is the heap-less matrix kernel; ranked membership is handled - // by the explicit heap-bearing operator, not inferred from delta keys. - let delta = CountSketchDelta { - rows: pb.rows, - cols: pb.cols, - cells, - l2: pb.l2, - hh_keys: Vec::new(), - }; - self.inner - .apply_delta(&delta) - .map_err(|e| format!("apply CountSketchDelta: {e}"))?; - Ok(()) - } -} - -impl SerializableToSink for CountSketchAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "row_num": self.inner.rows, - "col_num": self.inner.cols, - "sketch": self.inner.sketch(), - }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl AggregateCore for CountSketchAccumulator { - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "CountSketchAccumulator" - } - - /// Per-window base rotation: rebuild an empty signed-counter matrix - /// with the same (rows, cols) so the next window's additive cell - /// deltas align to the identical hash geometry. - fn reset_to_empty(&mut self) { - self.inner = CountSketch::new(self.inner.rows, self.inner.cols); - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } - - fn merge_with( - &self, - other: &dyn AggregateCore, - ) -> Result, Box> { - if other.get_accumulator_type() != self.get_accumulator_type() { - return Err(format!( - "Cannot merge CountSketchAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - let other_cs = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to CountSketchAccumulator")?; - - let merged_inner = CountSketch::merge_refs(&[&self.inner, &other_cs.inner])?; - Ok(Box::new(Self { - inner: merged_inner, - })) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::CountSketch - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: 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::accumulators::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_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; - use prost::Message; - - let mut acc = CountSketchAccumulator { - inner: CountSketch::from_legacy_matrix( - vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], - 2, - 3, - ), - }; - let bytes = PbDelta { - rows: 2, - cols: 3, - cell_rows: vec![0, 1], - cell_cols: vec![0, 2], - d_counts: vec![10, -6], - l2: vec![], - ..Default::default() - } - .encode_to_vec(); - - acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); - assert_eq!( - acc.inner.sketch(), - &vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 0.0]] - ); - } - - #[test] - fn test_apply_proto_delta_bytes_rejects_garbage() { - let mut acc = CountSketchAccumulator::new(2, 3); - assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); - } - - // ---------------------------------------------------------------- - // Defensive inbound-dimension validation (harden/sketch-dim-validation). - // Malformed / narrow-hash-budget-violating CountSketch dims must be - // rejected gracefully (Err, never a panic); valid configs the backend - // actually uses (5x2048, 5x4096, 5x2000) must still decode. - // ---------------------------------------------------------------- - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { - use asap_sketchlib::proto::sketchlib::CounterType; - // 5 * ceil(log2(8192))=5*13=65 > 64 — narrow-hash-budget violation. - // counts sized to rows*cols so rejection is on dims, not length. - let n = 5usize * 8192usize; - let bytes = encode_state( - 5, - 8192, - CounterType::Int64 as i32, - vec![0i64; n], - Vec::new(), - ); - let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "budget-violating dims should be rejected"); - assert!(result.unwrap_err().to_string().contains("rejecting")); - - // A valid neighbour (5x4096) on the same path still decodes fine. - let n_ok = 5usize * 4096usize; - let ok_bytes = encode_state( - 5, - 4096, - CounterType::Int64 as i32, - vec![0i64; n_ok], - Vec::new(), - ); - let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) - .expect("valid 5x4096 CountSketch should still decode"); - assert_eq!(acc.inner.rows, 5); - assert_eq!(acc.inner.cols, 4096); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_oversized_dims() { - use asap_sketchlib::proto::sketchlib::CounterType; - // Declare 1 x 16,777,216 = 16M cells (> 8M cap) but send an empty - // counts vector: validation must reject on the dim cap BEFORE the - // decoder tries to allocate/reshape a 16M-entry matrix. (1 row keeps - // the hash budget tiny so the cap check, not the budget check, fires.) - let bytes = encode_state( - 1, - 16_777_216, - CounterType::Int64 as i32, - Vec::new(), - Vec::new(), - ); - let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "oversized dims should be rejected"); - let msg = result.unwrap_err().to_string(); - assert!(msg.contains("cap"), "expected cell-cap error, got: {msg}"); - } -} diff --git a/crates/asap-physical-operators/src/accumulators/count_sketch_with_heap_accumulator.rs b/crates/asap-physical-operators/src/accumulators/count_sketch_with_heap_accumulator.rs deleted file mode 100644 index 6b8c1b24d..000000000 --- a/crates/asap-physical-operators/src/accumulators/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::{ - 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::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::accumulators::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/crates/asap-physical-operators/src/accumulators/datasketches_kll_accumulator.rs b/crates/asap-physical-operators/src/accumulators/datasketches_kll_accumulator.rs deleted file mode 100644 index 7a29e84b5..000000000 --- a/crates/asap-physical-operators/src/accumulators/datasketches_kll_accumulator.rs +++ /dev/null @@ -1,727 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, - SingleSubpopulationAggregate, -}; -use asap_sketchlib::{KllSketch, MessagePackCodec}; -use base64::{engine::general_purpose, Engine as _}; -use serde_json::Value; -use std::collections::HashMap; -#[cfg(feature = "extra_debugging")] -use std::time::Instant; -use tracing::debug; - -use 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::SingleSubpopulationAggregate; - self.query(statistic, Some(query_kwargs)) - } -} - -impl SingleSubpopulationAggregate for DatasketchesKLLAccumulator { - fn query( - &self, - statistic: Statistic, - query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Quantile => { - debug!( - "Querying DatasketchesKLLAccumulator for quantile with kwargs: {:?}", - query_kwargs - ); - let quantile = query_kwargs - .and_then(|kwargs| kwargs.get("quantile")) - .ok_or("Missing quantile parameter for quantile query")? - .parse::() - .map_err(|_| "Invalid quantile parameter format")?; - - if !(0.0..=1.0).contains(&quantile) { - return Err("Quantile must be between 0.0 and 1.0".into()); - } - - Ok(self.get_quantile(quantile)) - } - _ => Err( - format!("Unsupported statistic in DatasketchesKLLAccumulator: {statistic:?}") - .into(), - ), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for DatasketchesKLLAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - let mut iter = accumulators.into_iter(); - let mut merged = iter.next().unwrap(); - for acc in iter { - merged.inner.merge(&acc.inner)?; - } - Ok(merged) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use prost::Message; - - fn encode_state(state: asap_sketchlib::proto::sketchlib::KllState) -> Vec { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; - SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(state)), - ..Default::default() - } - .encode_to_vec() - } - - #[test] - fn test_datasketches_kll_creation() { - let kll = DatasketchesKLLAccumulator::new(200); - assert!(kll.inner.count() == 0); - assert_eq!(kll.inner.k, 200); - } - - #[test] - fn test_datasketches_kll_update() { - let mut kll = DatasketchesKLLAccumulator::new(200); - kll.update(10.0); - kll.update(20.0); - kll.update(15.0); - assert_eq!(kll.inner.count(), 3); - } - - #[test] - fn test_datasketches_kll_quantile() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for i in 1..=10 { - kll.update(i as f64); - } - assert_eq!(kll.get_quantile(0.0), 1.0); - assert_eq!(kll.get_quantile(1.0), 10.0); - // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. - let q50 = kll.get_quantile(0.5); - assert!((q50 - 6.0).abs() <= 1.0, "expected median ~6, got {q50}"); - } - - #[test] - fn test_datasketches_kll_query() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for i in 1..=10 { - kll.update(i as f64); - } - - let mut query_kwargs = HashMap::new(); - query_kwargs.insert("quantile".to_string(), "0.5".to_string()); - let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); - // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. - assert!( - (result - 6.0).abs() <= 1.0, - "expected median ~6, got {result}" - ); - - assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); - } - - #[test] - fn test_datasketches_kll_merge() { - let mut kll1 = DatasketchesKLLAccumulator::new(200); - let mut kll2 = DatasketchesKLLAccumulator::new(200); - - for i in 1..=5 { - kll1.update(i as f64); - } - for i in 6..=10 { - kll2.update(i as f64); - } - - let merged = DatasketchesKLLAccumulator::merge_accumulators(vec![kll1, kll2]).unwrap(); - assert_eq!(merged.inner.count(), 10); - assert_eq!(merged.get_quantile(0.0), 1.0); - assert_eq!(merged.get_quantile(1.0), 10.0); - } - - #[test] - fn test_datasketches_kll_get_keys() { - let kll = DatasketchesKLLAccumulator::new(200); - assert_eq!(kll.type_name(), "DatasketchesKLLAccumulator"); - } - - #[test] - fn test_trait_object() { - let mut kll = DatasketchesKLLAccumulator::new(200); - kll.update(5.0); - let trait_obj: Box = Box::new(kll); - assert_eq!(trait_obj.type_name(), "DatasketchesKLLAccumulator"); - } - - #[test] - fn test_datasketches_kll_query_with_kwargs() { - let mut kll = DatasketchesKLLAccumulator::new(200); - for i in 1..=10 { - kll.update(i as f64); - } - - let mut query_kwargs = HashMap::new(); - query_kwargs.insert("quantile".to_string(), "0.5".to_string()); - let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); - // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. - assert!( - (result - 6.0).abs() <= 1.0, - "expected median ~6, got {result}" - ); - - query_kwargs.insert("quantile".to_string(), "0.9".to_string()); - let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); - // Sketchlib KLL is approximate; 0.9 quantile of 1..10 may be 9 or 10. - assert!( - (9.0..=10.0).contains(&result), - "expected 0.9 quantile in [9,10], got {result}" - ); - - query_kwargs.insert("quantile".to_string(), "0.0".to_string()); - assert_eq!( - kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), - 1.0 - ); - - query_kwargs.insert("quantile".to_string(), "1.0".to_string()); - assert_eq!( - kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), - 10.0 - ); - - assert!(kll.query(Statistic::Quantile, None).is_err()); - - query_kwargs.insert("quantile".to_string(), "invalid".to_string()); - assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); - - query_kwargs.insert("quantile".to_string(), "1.5".to_string()); - assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); - - query_kwargs.insert("quantile".to_string(), "-0.1".to_string()); - assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); - - query_kwargs.insert("quantile".to_string(), "0.5".to_string()); - assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); - } - - #[test] - fn test_datasketches_kll_merge_multiple() { - let mut kll1 = DatasketchesKLLAccumulator::new(200); - let mut kll2 = DatasketchesKLLAccumulator::new(200); - let mut kll3 = DatasketchesKLLAccumulator::new(200); - - for i in 1..=5 { - kll1.update(i as f64); - } - for i in 6..=10 { - kll2.update(i as f64); - } - for i in 11..=15 { - kll3.update(i as f64); - } - - let boxed_accs: Vec> = - vec![Box::new(kll1), Box::new(kll2), Box::new(kll3)]; - - let merged = DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).unwrap(); - assert_eq!(merged.inner.count(), 15); - assert_eq!(merged.get_quantile(0.0), 1.0); - assert_eq!(merged.get_quantile(1.0), 15.0); - assert_eq!(merged.get_quantile(0.5), 8.0); - } - - #[test] - fn test_datasketches_kll_merge_multiple_error_cases() { - let empty: Vec> = vec![]; - assert!(DatasketchesKLLAccumulator::merge_multiple(&empty).is_err()); - - let kll1 = DatasketchesKLLAccumulator::new(200); - let kll2 = DatasketchesKLLAccumulator::new(100); - let boxed_accs: Vec> = vec![Box::new(kll1), Box::new(kll2)]; - assert!(DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).is_err()); - - use crate::accumulators::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; - - let items: Vec = (0..64).map(|i| i as f64).collect(); - let state = KllState { - k: 200, - m: 8, - num_levels: 1, - levels: vec![0, 64], - items: items.clone(), - coin: None, - offset: 0.0, - value_scale: 0, - residuals: Vec::new(), - }; - let bytes = encode_state(state); - - let acc = - DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); - assert_eq!(acc.inner.count(), 64); - // For 64 values 0..63, the true median is 31.5 and quantile - // error is ~1% × range = 0.63. KLL's own point query can - // legally be off by up to ε × N ~= 0.01 × 64 = 0.64. Allow a - // generous tolerance since the important invariant is "the - // decoded sketch is queryable and returns a sensible value". - let median = acc.get_quantile(0.5); - assert!( - (median - 31.5).abs() <= 10.0, - "reconstructed median {median} is outside tolerance of true median 31.5" - ); - let q01 = acc.get_quantile(0.01); - let q99 = acc.get_quantile(0.99); - assert!( - q01 <= q99, - "quantile monotonicity violated: q01={q01}, q99={q99}" - ); - } - - // Compacted portable state is highest-level first, unlike the runtime buffer. - #[test] - fn compacted_wire_state_preserves_count_and_quantiles() { - use asap_sketchlib::{proto::sketchlib::KllState, sketches::KLL}; - let mut source = KLL::::init_kll_with_seed(32, 123); - for i in 0..1000 { - source.update(&(((i * 7919 + 17) % 1009) as f64 / 1009.0)); - } - assert!(source.wire_num_levels() > 1); - let state = KllState { - k: 32, - m: source.wire_m(), - num_levels: source.wire_num_levels(), - levels: source.wire_levels(), - items: source.wire_items(), - coin: None, - offset: 0.0, - value_scale: 0, - residuals: vec![], - }; - let decoded = - DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&encode_state(state)).unwrap(); - assert_eq!(decoded.inner.count(), source.count() as u64); - for q in [0.0, 0.1, 0.5, 0.9, 1.0] { - assert_eq!(decoded.inner.quantile(q), source.quantile(q), "q={q}"); - } - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrapped() { - // Mirrors what DataCollector's kllprocessor emits: the state - // wrapped in a `SketchEnvelope{kll: ...}` via sketchlib-go's - // `SerializePortableFO` + `proto.Marshal`. - use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; - - let items: Vec = (0..64).map(|i| i as f64).collect(); - let state = KllState { - k: 200, - m: 8, - num_levels: 1, - levels: vec![0, 64], - items, - coin: None, - offset: 0.0, - value_scale: 0, - residuals: Vec::new(), - }; - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::Kll(state)), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes) - .expect("envelope-wrapped decode should succeed"); - assert_eq!(acc.inner.count(), 64); - } - - #[test] - fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { - use asap_sketchlib::proto::sketchlib::{sketch_envelope, CountMinState, SketchEnvelope}; - - let env = SketchEnvelope { - sketch_state: Some(sketch_envelope::SketchState::CountMin( - CountMinState::default(), - )), - ..Default::default() - }; - let bytes = env.encode_to_vec(); - - let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err(), "wrong-sketch envelope should error"); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_small_k() { - use asap_sketchlib::proto::sketchlib::KllState; - let state = KllState { - k: 4, // < minimum of 8 - m: 2, - num_levels: 0, - levels: Vec::new(), - items: Vec::new(), - coin: None, - offset: 0.0, - value_scale: 0, - residuals: Vec::new(), - }; - let bytes = encode_state(state); - let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("k must be >= 8")); - } - - #[test] - fn test_from_sketchlib_proto_bytes_rejects_inconsistent_levels() { - use asap_sketchlib::proto::sketchlib::KllState; - // num_levels=1 but levels array has 3 entries instead of 2 - let state = KllState { - k: 200, - m: 8, - num_levels: 1, - levels: vec![0, 5, 10], - items: vec![1.0, 2.0, 3.0, 4.0, 5.0], - coin: None, - offset: 0.0, - value_scale: 0, - residuals: Vec::new(), - }; - let bytes = encode_state(state); - let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); - assert!(result.is_err()); - assert!(result.unwrap_err().to_string().contains("levels length")); - } - - #[test] - fn aux_stats_exposes_count_via_kll_n() { - let mut acc = DatasketchesKLLAccumulator::new(200); - for i in 0..50 { - acc.update(i as f64); - } - let aux = acc.aux_stats(); - assert_eq!(aux.count, Some(50)); - // KLL doesn't natively expose min/max cheaply and doesn't - // track sum at all — those fields must be None so callers - // fall through to query_statistic. - assert_eq!(aux.sum, None); - assert_eq!(aux.min, None); - assert_eq!(aux.max, None); - } - - #[test] - fn aux_stats_empty_kll_has_zero_count() { - let acc = DatasketchesKLLAccumulator::new(200); - assert_eq!(acc.aux_stats().count, Some(0)); - } -} diff --git a/crates/asap-physical-operators/src/accumulators/dd_sketch_accumulator.rs b/crates/asap-physical-operators/src/accumulators/dd_sketch_accumulator.rs deleted file mode 100644 index 9aa1132e6..000000000 --- a/crates/asap-physical-operators/src/accumulators/dd_sketch_accumulator.rs +++ /dev/null @@ -1,665 +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::{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_sketchlib::proto::sketchlib::DdSketchDelta` message. - pub fn apply_proto_delta_bytes( - &mut self, - buffer: &[u8], - ) -> Result<(), Box> { - use asap_sketchlib::proto::sketchlib::DdSketchDelta as PbDelta; - use prost::Message; - - let pb = PbDelta::decode(buffer).map_err(|e| format!("decode DDSketchDelta: {e}"))?; - - // The delta no longer carries d_count/d_sum/min/max - // (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). Apply the - // bucket deltas only; `DdSketch` recomputes its total count from - // the merged bucket counts (`total_count()`). - let buckets = pb - .buckets - .into_iter() - .map(|b| (b.index, b.d_count)) - .collect(); - let delta = DdSketchDelta { - buckets, - ..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::accumulators::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_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; - use prost::Message; - - let mut acc = DDSketchAccumulator::new(0.01); - acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); - - // The wire delta now carries only bucket deltas (tags 2-7 - // reserved); `DdSketchBucketDelta` has just `index` + `d_count`. - let bytes = PbDelta { - buckets: vec![ - DdSketchBucketDelta { - index: 0, - d_count: 10, - }, - DdSketchBucketDelta { - index: 2, - d_count: 20, - }, - ], - } - .encode_to_vec(); - - acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); - assert_eq!(acc.inner.store_counts, vec![11, 2, 23]); - // `count` recomputed from the merged buckets: 11 + 2 + 23 = 36. - assert_eq!(acc.inner.total_count(), 36); - } - - /// A valid protobuf with an inadmissible span must not acknowledge a dropped update. - #[test] - fn test_apply_proto_delta_rejects_span_without_mutating_state() { - use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; - use prost::Message; - let mut acc = DDSketchAccumulator::new(0.01); - acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); - let bytes = PbDelta { - buckets: vec![DdSketchBucketDelta { - index: i32::MAX, - d_count: 1, - }], - } - .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/crates/asap-physical-operators/src/accumulators/exact_accumulator.rs b/crates/asap-physical-operators/src/accumulators/exact_accumulator.rs deleted file mode 100644 index 5d23acd87..000000000 --- a/crates/asap-physical-operators/src/accumulators/exact_accumulator.rs +++ /dev/null @@ -1,326 +0,0 @@ -//! Exact summary state identified by Planner family, independent of keyed layout. -use super::increase_accumulator::IncreaseAccumulator; -use crate::{ - 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/crates/asap-physical-operators/src/accumulators/hll_sketch_accumulator.rs b/crates/asap-physical-operators/src/accumulators/hll_sketch_accumulator.rs deleted file mode 100644 index 254eee38e..000000000 --- a/crates/asap-physical-operators/src/accumulators/hll_sketch_accumulator.rs +++ /dev/null @@ -1,788 +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::accumulators::dd_sketch_accumulator::normalize_sample_p; -use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; -use asap_sketchlib::{HllSketch, HllVariant, MessagePackCodec}; -use serde_json::Value; -use std::collections::HashMap; - -/// Decode one protobuf base-128 varint (LEB128) from the front of `buf`. -/// Returns `(value, bytes_consumed)`, or `None` if the buffer is truncated -/// or the varint overflows u64. -pub(crate) fn read_uvarint(buf: &[u8]) -> Option<(u64, usize)> { - let mut result: u64 = 0; - let mut shift: u32 = 0; - for (i, &b) in buf.iter().enumerate() { - if shift >= 64 { - return None; - } - result |= u64::from(b & 0x7f) << shift; - if b & 0x80 == 0 { - return Some((result, i + 1)); - } - shift += 7; - } - None -} - -/// Expand sketchlib-go's sparse HLL register encoding -/// (`HLLSparseRegisters.packed`) into the dense `num_registers`-byte array. -/// -/// Layout (sketchlib-go `proto/hll/hll.proto`): varint-packed -/// `(index_delta, value)` pairs in ascending index order; `prev_index` -/// starts at 0, so each register's absolute index is the running sum of the -/// deltas. Mirrors the Go encoder in `sketches/HLL/sparse.go` -/// (`encodeSparseRegisters`). The reconstructed array is byte-identical to -/// the dense `registers` field a high-cardinality producer would have sent. -pub(crate) fn expand_sparse_hll_registers( - packed: &[u8], - num_registers: usize, -) -> Result, Box> { - let mut regs = vec![0u8; num_registers]; - let mut prev: u64 = 0; - let mut pos = 0usize; - while pos < packed.len() { - let (delta, n1) = read_uvarint(&packed[pos..]) - .ok_or("HLLSparseRegisters.packed: truncated index_delta varint")?; - pos += n1; - let (value, n2) = read_uvarint(&packed[pos..]) - .ok_or("HLLSparseRegisters.packed: truncated value varint")?; - pos += n2; - let idx = prev + delta; - let i = usize::try_from(idx) - .map_err(|_| format!("HLLSparseRegisters: index {idx} overflows usize"))?; - if i >= num_registers { - return Err(format!( - "HLLSparseRegisters: register index {i} >= num_registers {num_registers}" - ) - .into()); - } - regs[i] = u8::try_from(value) - .map_err(|_| format!("HLLSparseRegisters: register value {value} > 255"))?; - prev = idx; - } - Ok(regs) -} - -/// HLL accumulator — inner register array + variant metadata. -#[derive(Debug, Clone)] -pub struct HllSketchAccumulator { - pub inner: HllSketch, - /// Edge sampling probability `p ∈ (0,1]` carried on the producer's - /// `SketchEnvelope.sample_p`. HLL uses HASH-THRESHOLD sampling — each - /// DISTINCT key is admitted into the sketch with probability `p`, so the - /// register-derived distinct-count estimate is ~`p`× the true - /// cardinality and a `Cardinality`/`Count` query must rescale by `1/p`. - /// `1.0` (and the proto3 default `0.0`, dual-read as `1.0`) means no - /// sampling, so the rescale is a no-op and the behaviour is identical to - /// before. Mirrors `DDSketchAccumulator::sample_p`; set from the envelope - /// at the `from_sketchlib_proto_bytes` decode site and preserved across - /// `reset_to_empty` and `merge_with`. - /// - /// NOTE: HLL edge sampling is currently force-disabled in the edge - /// (`warm_sketch.go` HLL case always emits `sample_p = 1.0`), so in - /// practice `p = 1.0` today and this is a latent-correctness fix that - /// activates if HLL sampling is ever enabled. - pub sample_p: f64, -} - -impl HllSketchAccumulator { - pub fn new(variant: HllVariant, precision: u32) -> Self { - Self { - inner: HllSketch::new(variant, precision), - sample_p: 1.0, - } - } - - /// Decode from the modified OTLP wire format's - /// `HLLSketchDataPoint.sketch` bytes when - /// `encoding = HLL_SKETCH_ENCODING_MSGPACK`. The bytes are the - /// MessagePack serialization of the cross-language sketch-core - /// `HllSketch` struct — PR I parity entrypoint. - pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { - Ok(Self { - inner: HllSketch::from_msgpack(buffer) - .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?, - // The msgpack HllSketch struct carries no envelope/sample_p; the - // msgpack path is parity/test-only and is never edge-sampled. - sample_p: 1.0, - }) - } - - /// Decode from the modified OTLP wire format's - /// `HLLSketchDataPoint.sketch` bytes — the protobuf-encoded - /// `asap_sketchlib::proto::sketchlib::HyperLogLogState` message - /// that DataCollector's `hllprocessor` emits when - /// `encoding = HLL_SKETCH_ENCODING_PROTO`. - pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { - use asap_sketchlib::proto::sketchlib::{ - sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, - }; - use prost::Message; - - // DataCollector's hllprocessor wraps the state in a - // `SketchEnvelope{hll: HyperLogLogState}` via sketchlib-go's - // `SerializePortableFO` + `proto.Marshal`. Try envelope first, - // fall back to bare `HyperLogLogState` for callers (e.g. unit - // tests) that encode the state directly. Mirrors the PR #14 - // fix on `CountMinSketchAccumulator::from_sketchlib_proto_bytes`. - // Capture the envelope's `sample_p` alongside the state so a - // Cardinality query can rescale the distinct-count estimate by - // `1/p`. Bare `HyperLogLogState` bytes (no envelope) carry no - // sampling info → `sample_p` 1.0 (no rescale). Mirrors - // `DDSketchAccumulator`. - let (state, sample_p) = match SketchEnvelope::decode(buffer) { - Ok(env) => { - let sp = env.sample_p; - match env.sketch_state { - Some(sketch_envelope::SketchState::Hll(st)) => (st, sp), - Some(other) => { - return Err(format!( - "SketchEnvelope contains non-HLL sketch: {:?}", - std::mem::discriminant(&other) - ) - .into()); - } - None => ( - HyperLogLogState::decode(buffer) - .map_err(|e| format!("decode HyperLogLogState: {e}"))?, - 1.0, - ), - } - } - Err(_) => ( - HyperLogLogState::decode(buffer) - .map_err(|e| format!("decode HyperLogLogState: {e}"))?, - 1.0, - ), - }; - if state.precision == 0 || state.precision > 20 { - return Err(format!( - "HyperLogLogState precision {} out of range (expected 1..=20)", - state.precision - ) - .into()); - } - let expected_len = 1usize << state.precision; - // Register resolution. sketchlib-go emits the SPARSE - // `registers_sparse` (proto tag 7) form below its dense/sparse - // crossover (~6000 non-zero registers — see - // sketchlib-go/sketches/HLL/sparse.go); low-cardinality producers - // (the common case) therefore leave the dense `registers` (tag 3) - // field empty. The proto contract (hll.proto) is: read whichever of - // `registers` / `registers_sparse` is present; if both are empty the - // sketch is all-zero. Reconstruct the dense 2^precision array in all - // three cases so the inner `HllSketch` always gets a full register - // vector. - let dense_registers: Vec = if state.registers.len() == expected_len { - state.registers.clone() - } else if !state.registers.is_empty() { - // A non-empty dense field of the wrong length is a malformed frame. - return Err(format!( - "HyperLogLogState registers has {} bytes, expected 2^precision = {}", - state.registers.len(), - expected_len - ) - .into()); - } else if let Some(sparse) = state.registers_sparse.as_ref() { - expand_sparse_hll_registers(&sparse.packed, expected_len)? - } else { - // Neither representation populated → all-zero register array. - vec![0u8; expected_len] - }; - let proto_variant = ProtoVariant::try_from(state.variant) - .map_err(|_| format!("HyperLogLogState has unknown variant tag {}", state.variant))?; - let variant = match proto_variant { - ProtoVariant::Unspecified => HllVariant::Unspecified, - ProtoVariant::Regular => HllVariant::Regular, - ProtoVariant::ErtlMle => HllVariant::Datafusion, - ProtoVariant::Hip => HllVariant::Hip, - }; - let inner = HllSketch::from_raw( - variant, - state.precision, - dense_registers, - state.hip_kxq0, - state.hip_kxq1, - state.hip_est, - ); - Ok(Self { - inner, - sample_p: normalize_sample_p(sample_p), - }) - } - - /// Apply a proto-encoded `HLLDelta` frame to this accumulator's - /// inner sketch — the decode path for - /// `HLL_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). - /// - /// Called against an accumulator that already carries the base - /// sketch state; the caller is the per-series snapshot cache in - /// the ingest path. Bytes are the - /// `asap_sketchlib::proto::sketchlib::HllDelta` message. - pub fn apply_proto_delta_bytes( - &mut self, - buffer: &[u8], - ) -> Result<(), Box> { - // The HLLDelta wire format is a varint-packed (index_delta, value) blob; - // decode + apply (register-wise max) via the shared sketch library so - // the unpacking stays a single source of truth. - self.inner - .apply_delta_bytes(buffer) - .map_err(|e| format!("apply HLLDelta: {e}"))?; - Ok(()) - } -} - -impl SerializableToSink for HllSketchAccumulator { - fn serialize_to_json(&self) -> Value { - serde_json::json!({ - "variant": format!("{:?}", self.inner.variant), - "precision": self.inner.precision, - "register_bytes": self.inner.registers.len(), - "hip_kxq0": self.inner.hip_kxq0, - "hip_kxq1": self.inner.hip_kxq1, - "hip_est": self.inner.hip_est, - }) - } - - fn serialize_to_bytes(&self) -> Vec { - self.inner.to_msgpack().unwrap_or_default() - } -} - -impl AggregateCore for HllSketchAccumulator { - fn approx_memory_bytes(&self) -> usize { - std::mem::size_of::().saturating_add(self.inner.registers.capacity()) - } - fn clone_boxed_core(&self) -> Box { - Box::new(self.clone()) - } - - fn type_name(&self) -> &'static str { - "HllSketchAccumulator" - } - - /// Per-window base rotation: zero the registers but keep the variant - /// and precision. Critical for HLL — its register-wise `max` merge - /// has no inverse, so a never-reset base accumulates the all-time-max - /// across windows (`docs/delta-baseline-contract.md` §1.5); rotating - /// to an empty register array makes per-window cardinality correct. - /// `sample_p` is a per-series config constant (not per-window data), so - /// it is intentionally preserved across the rotation — mirrors - /// `DDSketchAccumulator`. - fn reset_to_empty(&mut self) { - self.inner = HllSketch::new(self.inner.variant, self.inner.precision); - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn as_any_mut(&mut self) -> &mut dyn std::any::Any { - self - } - - fn merge_with( - &self, - other: &dyn AggregateCore, - ) -> Result, Box> { - if other.get_accumulator_type() != self.get_accumulator_type() { - return Err(format!( - "Cannot merge HllSketchAccumulator with {}", - other.get_accumulator_type() - ) - .into()); - } - let other_hll = other - .as_any() - .downcast_ref::() - .ok_or("Failed to downcast to HllSketchAccumulator")?; - let merged_inner = HllSketch::merge_refs(&[&self.inner, &other_hll.inner])?; - // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a - // per-series config constant, so both operands carry the same value - // in practice. Prefer a sampled factor over the no-sampling default - // so a merge with a freshly-reset (1.0) base keeps the series' - // sampling rate. - let sample_p = if self.sample_p < 1.0 { - self.sample_p - } else { - other_hll.sample_p - }; - Ok(Box::new(Self { - inner: merged_inner, - sample_p, - })) - } - - fn get_accumulator_type(&self) -> AggregationType { - AggregationType::HLL - } - - fn get_keys(&self) -> Option> { - None - } - - fn query_statistic( - &self, - statistic: 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::accumulators::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_sketchlib::proto::sketchlib::HllDelta as PbDelta; - use prost::Message; - - let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); - acc.inner.registers = vec![1, 5, 3, 7]; - - // Packed (index_delta, value) blob for updates {0:4, 2:6}: - // varint(0),varint(4),varint(2),varint(6). - let delta_bytes = PbDelta { - packed_updates: vec![0, 4, 2, 6], - } - .encode_to_vec(); - - acc.apply_proto_delta_bytes(&delta_bytes).expect("apply ok"); - // Max semantics: reg[0]=max(1,4)=4, reg[2]=max(3,6)=6; others unchanged. - assert_eq!(acc.inner.registers, vec![4, 5, 6, 7]); - } - - #[test] - fn test_apply_proto_delta_bytes_rejects_garbage() { - let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); - assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); - } - - // ----- sample_p cardinality rescale ----- - // - // HLL uses hash-threshold sampling: each distinct key is admitted into - // the sketch with probability `p`, so the register-derived cardinality - // estimate is ~p× the true distinct count and must be rescaled by 1/p. - - #[test] - fn test_cardinality_is_rescaled_by_sample_p() { - use 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/crates/asap-physical-operators/src/accumulators/hydra_kll_accumulator.rs b/crates/asap-physical-operators/src/accumulators/hydra_kll_accumulator.rs deleted file mode 100644 index a5cc64dc1..000000000 --- a/crates/asap-physical-operators/src/accumulators/hydra_kll_accumulator.rs +++ /dev/null @@ -1,165 +0,0 @@ -use crate::{ - AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, - MultipleSubpopulationAggregate, SerializableToSink, -}; -use asap_sketchlib::{HydraKllSketch, MessagePackCodec}; -use base64::{engine::general_purpose, Engine as _}; -use std::collections::HashMap; - -use 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::MultipleSubpopulationAggregate; - let key_val = key - .as_ref() - .ok_or("Key required for HydraKllSketchAccumulator")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for HydraKllSketchAccumulator { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Quantile => { - let quantile = query_kwargs - .and_then(|kwargs| kwargs.get("quantile")) - .ok_or("Missing quantile parameter for quantile query")? - .parse::() - .map_err(|_| "Invalid quantile parameter format")?; - - if !(0.0..=1.0).contains(&quantile) { - return Err("Quantile must be between 0.0 and 1.0".into()); - } - - Ok(self.query_key(key, quantile)) - } - _ => Err( - format!("Unsupported statistic in HydraKllSketchAccumulator: {statistic:?}").into(), - ), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} diff --git a/crates/asap-physical-operators/src/accumulators/increase_accumulator.rs b/crates/asap-physical-operators/src/accumulators/increase_accumulator.rs deleted file mode 100644 index 7407800b0..000000000 --- a/crates/asap-physical-operators/src/accumulators/increase_accumulator.rs +++ /dev/null @@ -1,742 +0,0 @@ -use crate::{ - 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::SingleSubpopulationAggregate; - self.query( - statistic, - (!query_kwargs.is_empty()).then_some(query_kwargs), - ) - } -} - -impl SingleSubpopulationAggregate for IncreaseAccumulator { - fn query( - &self, - statistic: Statistic, - query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Increase => Ok(self.extrapolated_value(query_kwargs, false)?), - Statistic::Rate => Ok(self.extrapolated_value(query_kwargs, true)?), - // For instant `sum [by (...)] (counter_metric)` Prometheus - // sums the latest cumulative value of each matching series. - // The IncreaseAccumulator already tracks that latest value - // in `last_seen_measurement`, so per-series Sum is just - // that scalar; the engine's outer aggregation groups by the - // `by` labels and adds the per-series totals across keys. - // - // See PR #108 audit conclusion (commit 4359e10) and issue - // ProjectASAP/ASAPCollector#46: pre-fix the ASAP tier ingested - // counters as IncreaseAccumulator and bare `sum by (...) ()` - // capability-missed because this trait did not answer Sum. - Statistic::Sum => Ok(self.last_seen_measurement.value), - _ => Err(format!("Unsupported statistic in IncreaseAccumulator: {statistic:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl IncreaseAccumulator { - fn extrapolated_value( - &self, - query_kwargs: Option<&HashMap>, - is_rate: bool, - ) -> Result> { - if self.sample_count < 2 || self.last_seen_timestamp <= self.starting_timestamp { - return Err("at least two ordered counter samples are required".into()); - } - let sampled_interval = (self.last_seen_timestamp - self.starting_timestamp) as f64 / 1000.0; - let Some(kwargs) = query_kwargs else { - return Ok(if is_rate { - self.total_increase / sampled_interval - } else { - self.total_increase - }); - }; - let range_start = kwargs - .get("range_start_ms") - .ok_or("missing range_start_ms")? - .parse::()?; - let range_end = kwargs - .get("range_end_ms") - .ok_or("missing range_end_ms")? - .parse::()?; - if range_end <= range_start { - return Err("invalid counter evaluation range".into()); - } - - let mut duration_to_start = - (self.starting_timestamp.saturating_sub(range_start)) as f64 / 1000.0; - let duration_to_end = (range_end.saturating_sub(self.last_seen_timestamp)) as f64 / 1000.0; - let average_sample_interval = sampled_interval / (self.sample_count - 1) as f64; - let extrapolation_threshold = average_sample_interval * 1.1; - - if self.total_increase > 0.0 && self.starting_measurement.value >= 0.0 { - let duration_to_zero = - sampled_interval * (self.starting_measurement.value / self.total_increase); - duration_to_start = duration_to_start.min(duration_to_zero); - } - let mut extrapolate_to = sampled_interval; - extrapolate_to += if duration_to_start < extrapolation_threshold { - duration_to_start.max(0.0) - } else { - average_sample_interval / 2.0 - }; - extrapolate_to += if duration_to_end < extrapolation_threshold { - duration_to_end.max(0.0) - } else { - average_sample_interval / 2.0 - }; - let mut factor = extrapolate_to / sampled_interval; - if is_rate { - factor /= (range_end - range_start) as f64 / 1000.0; - } - Ok(self.total_increase * factor) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn test_increase_accumulator_creation() { - let starting_measurement = Measurement::new(10.0); - let last_seen_measurement = Measurement::new(25.0); - let acc = IncreaseAccumulator::new( - starting_measurement.clone(), - 1000, - last_seen_measurement.clone(), - 2000, - ); - - assert_eq!(acc.starting_measurement.value, 10.0); - assert_eq!(acc.starting_timestamp, 1000); - assert_eq!(acc.last_seen_measurement.value, 25.0); - assert_eq!(acc.last_seen_timestamp, 2000); - } - - #[test] - fn test_increase_accumulator_update() { - let starting_measurement = Measurement::new(10.0); - let mut acc = IncreaseAccumulator::new( - starting_measurement.clone(), - 1000, - starting_measurement.clone(), - 1000, - ); - - let new_measurement = Measurement::new(25.0); - acc.update(new_measurement.clone(), 2000); - - assert_eq!(acc.last_seen_measurement.value, 25.0); - assert_eq!(acc.last_seen_timestamp, 2000); - assert_eq!(acc.starting_measurement.value, 10.0); // Should remain unchanged - } - - #[test] - fn test_increase_accumulator_query() { - let starting_measurement = Measurement::new(10.0); - let last_seen_measurement = Measurement::new(25.0); - let acc = IncreaseAccumulator::new( - starting_measurement, - 1000, - last_seen_measurement, - 3000, // 2 second difference - ); - - // Test increase calculation - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, None).unwrap(), - 15.0 - ); - - // Test rate calculation (per second) - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, None).unwrap(), - 7.5 - ); // 15.0 / 2.0 - - // Statistic::Sum returns the latest cumulative counter value, - // matching Prometheus semantics for instant `sum()`. - // (Issue ProjectASAP/ASAPCollector#46, PR #108 diagnosis.) - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), - 25.0 - ); - - // Unsupported statistics still error. - assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); - } - - #[test] - fn prometheus_counter_reset_and_boundary_extrapolation() { - let mut acc = IncreaseAccumulator::new( - Measurement::new(10.0), - 10_000, - Measurement::new(10.0), - 10_000, - ); - acc.update(Measurement::new(20.0), 20_000); - acc.update(Measurement::new(3.0), 30_000); - acc.update(Measurement::new(13.0), 50_000); - assert_eq!(acc.total_increase, 23.0); - assert_eq!(acc.sample_count, 4); - - let kwargs = HashMap::from([ - ("range_start_ms".into(), "0".into()), - ("range_end_ms".into(), "60000".into()), - ]); - let increase = - crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, Some(&kwargs)) - .unwrap(); - let rate = crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, Some(&kwargs)) - .unwrap(); - assert!((increase - 34.5).abs() < 1e-12); - assert!((rate - 0.575).abs() < 1e-12); - } - - #[test] - fn pane_merge_preserves_resets_and_prometheus_extrapolation() { - let mut left = IncreaseAccumulator::new( - Measurement::new(10.0), - 10_000, - Measurement::new(10.0), - 10_000, - ); - left.update(Measurement::new(20.0), 20_000); - let mut right = - IncreaseAccumulator::new(Measurement::new(3.0), 30_000, Measurement::new(3.0), 30_000); - right.update(Measurement::new(13.0), 50_000); - let merged = IncreaseAccumulator::merge_accumulators(vec![right, left]).unwrap(); - assert_eq!(merged.total_increase, 23.0); - assert_eq!(merged.sample_count, 4); - let kwargs = HashMap::from([ - ("range_start_ms".into(), "0".into()), - ("range_end_ms".into(), "60000".into()), - ]); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&merged, Statistic::Increase, Some(&kwargs)) - .unwrap(), - 34.5 - ); - } - - #[test] - fn counter_sds_state_is_constant_size_per_pane() { - let mut acc = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(0.0), 0); - let initial = acc.serialize_to_bytes().len(); - for second in 1..=86_400 { - acc.update(Measurement::new(second as f64), second * 1_000); - } - assert_eq!(acc.serialize_to_bytes().len(), initial); - assert_eq!(acc.sample_count, 86_401); - assert_eq!( - acc.approx_memory_bytes(), - std::mem::size_of::() - ); - } - - #[test] - fn test_increase_accumulator_sum_is_latest_cumulative_value() { - // Instant `sum ()` semantics: the per-series summand is - // the latest cumulative counter value. Two series with latest - // values 100 and 50 (started at 10 and 5 respectively) should - // each report Sum = 100 and Sum = 50 — the engine's `sum by` - // outer aggregation does the cross-series total. - let acc_a = - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000); - let acc_b = - IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc_a, Statistic::Sum, None).unwrap(), - 100.0 - ); - assert_eq!( - crate::SingleSubpopulationAggregate::query(&acc_b, Statistic::Sum, None).unwrap(), - 50.0 - ); - } - - #[test] - fn test_increase_accumulator_merge() { - let acc1 = - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(20.0), 2000); - let acc2 = IncreaseAccumulator::new( - Measurement::new(5.0), - 500, // Earlier start - Measurement::new(15.0), - 1500, - ); - let acc3 = IncreaseAccumulator::new( - Measurement::new(20.0), - 2000, - Measurement::new(30.0), - 3000, // Later end - ); - - let merged = - >::merge_accumulators( - vec![acc1, acc2, acc3], - ) - .unwrap(); - - // Should use earliest start and latest end - assert_eq!(merged.starting_measurement.value, 5.0); - assert_eq!(merged.starting_timestamp, 500); - assert_eq!(merged.last_seen_measurement.value, 30.0); - assert_eq!(merged.last_seen_timestamp, 3000); - } - - #[test] - fn test_increase_accumulator_serialization() { - let acc = - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(25.0), 2000); - - // Test JSON serialization - let json = acc.serialize_to_json(); - let deserialized = IncreaseAccumulator::deserialize_from_json(&json).unwrap(); - assert_eq!( - acc.starting_measurement.value, - deserialized.starting_measurement.value - ); - assert_eq!(acc.starting_timestamp, deserialized.starting_timestamp); - assert_eq!( - acc.last_seen_measurement.value, - deserialized.last_seen_measurement.value - ); - assert_eq!(acc.last_seen_timestamp, deserialized.last_seen_timestamp); - - // Test byte serialization - let bytes = acc.serialize_to_bytes(); - let deserialized_bytes = IncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!( - acc.starting_measurement.value, - deserialized_bytes.starting_measurement.value - ); - assert_eq!( - acc.starting_timestamp, - deserialized_bytes.starting_timestamp - ); - assert_eq!( - acc.last_seen_measurement.value, - deserialized_bytes.last_seen_measurement.value - ); - assert_eq!( - acc.last_seen_timestamp, - deserialized_bytes.last_seen_timestamp - ); - assert_eq!(acc.total_increase, deserialized_bytes.total_increase); - assert_eq!(acc.sample_count, deserialized_bytes.sample_count); - - let legacy = &bytes[..bytes.len() - RESET_AWARE_WIRE_EXTENSION_LEN]; - let legacy_value = IncreaseAccumulator::deserialize_from_bytes(legacy).unwrap(); - assert_eq!(legacy_value.total_increase, 15.0); - assert_eq!(legacy_value.sample_count, 2); - } - - #[test] - fn test_trait_object() { - let acc: Box = Box::new(IncreaseAccumulator::new( - Measurement::new(10.0), - 1000, - Measurement::new(25.0), - 2000, - )); - - assert_eq!(acc.type_name(), "IncreaseAccumulator"); - } -} diff --git a/crates/asap-physical-operators/src/accumulators/keyed_counter_state.rs b/crates/asap-physical-operators/src/accumulators/keyed_counter_state.rs deleted file mode 100644 index 1d4cf1c7f..000000000 --- a/crates/asap-physical-operators/src/accumulators/keyed_counter_state.rs +++ /dev/null @@ -1,529 +0,0 @@ -use crate::accumulators::IncreaseAccumulator; -use crate::{ - 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::MultipleSubpopulationAggregate; - let key_val = key.as_ref().ok_or("Key required for KeyedCounterState")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for KeyedCounterState { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - query_kwargs: Option<&HashMap>, - ) -> Result> { - let data = self - .increases - .get(key) - .ok_or_else(|| format!("Key {key} not found in KeyedCounterState"))?; - - data.query(statistic, query_kwargs) - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for KeyedCounterState { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut result = KeyedCounterState::new(); - - for accumulator in accumulators { - for (key, data) in accumulator.increases { - if let Some(existing_data) = result.increases.get_mut(&key) { - *existing_data = - IncreaseAccumulator::merge_accumulators(vec![existing_data.clone(), data])?; - } else { - result.increases.insert(key, data); - } - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::Measurement; - - fn create_test_increase_accumulator(start_val: f64, end_val: f64) -> IncreaseAccumulator { - IncreaseAccumulator::new( - Measurement::new(start_val), - 1000, - Measurement::new(end_val), - 2000, - ) - } - - fn create_test_increase_accumulator_with_time( - start_val: f64, - start_time: i64, - end_val: f64, - end_time: i64, - ) -> IncreaseAccumulator { - IncreaseAccumulator::new( - Measurement::new(start_val), - start_time, - Measurement::new(end_val), - end_time, - ) - } - - #[test] - fn test_keyed_counter_state_creation() { - let acc = KeyedCounterState::new(); - assert!(acc.increases.is_empty()); - } - - #[test] - fn test_keyed_counter_state_update() { - let mut acc = KeyedCounterState::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - let increase1 = create_test_increase_accumulator(10.0, 25.0); - let increase2 = create_test_increase_accumulator(5.0, 15.0); - - acc.update(key1.clone(), increase1); - acc.update(key2.clone(), increase2); - - assert_eq!(acc.increases.len(), 2); - assert!(acc.increases.contains_key(&key1)); - assert!(acc.increases.contains_key(&key2)); - } - - #[test] - fn test_keyed_counter_state_query() { - let mut acc = KeyedCounterState::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let increase_acc = create_test_increase_accumulator(10.0, 25.0); - acc.update(key.clone(), increase_acc); - - // Test increase query - assert_eq!(acc.query(Statistic::Increase, &key, None).unwrap(), 15.0); - - // Test rate query (15.0 increase over 1 second = 15.0 per second) - assert_eq!(acc.query(Statistic::Rate, &key, None).unwrap(), 15.0); - - // Sum returns the latest cumulative counter value for the - // queried key (per-series Prometheus `sum()` semantics; - // see issue ProjectASAP/ASAPCollector#46 and PR #108 diagnosis). - // The series here was created with last_seen=25.0. - assert_eq!(acc.query(Statistic::Sum, &key, None).unwrap(), 25.0); - - // Unsupported statistic still errors. - assert!(acc.query(Statistic::Min, &key, None).is_err()); - - let unknown_key = KeyByLabelValues::new(); - assert!(acc.query(Statistic::Increase, &unknown_key, None).is_err()); - } - - #[test] - fn test_keyed_counter_state_sum_per_key() { - // `sum by (zone) (counter)` reaches KeyedCounterState - // only when the ASAP-tier ingest groups multiple series under - // a single accumulator (the `Multiple*` variant). In that case - // each per-key Sum should be the series' latest cumulative - // value; the engine's outer `by` aggregation does the cross-key - // grouping. (Issue ProjectASAP/ASAPCollector#46.) - let mut acc = KeyedCounterState::new(); - let east = KeyByLabelValues::new_with_labels(vec!["us-east-1".to_string()]); - let west = KeyByLabelValues::new_with_labels(vec!["us-west-2".to_string()]); - - acc.update( - east.clone(), - IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000), - ); - acc.update( - west.clone(), - IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000), - ); - - assert_eq!(acc.query(Statistic::Sum, &east, None).unwrap(), 100.0); - assert_eq!(acc.query(Statistic::Sum, &west, None).unwrap(), 50.0); - } - - #[test] - fn test_keyed_counter_state_merge() { - let mut acc1 = KeyedCounterState::new(); - let mut acc2 = KeyedCounterState::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - // Add different keys to each accumulator - acc1.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); - acc2.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); - - // Also add overlapping key with different time ranges (later timestamps) - acc2.update( - key1.clone(), - create_test_increase_accumulator_with_time(15.0, 2000, 30.0, 3000), - ); // Later time range - - let merged = KeyedCounterState::merge_accumulators(vec![acc1, acc2]).unwrap(); - - assert_eq!(merged.increases.len(), 2); - assert!(merged.increases.contains_key(&key1)); - assert!(merged.increases.contains_key(&key2)); - - // The merged key1 should have the full range (earliest start to latest end) - let merged_key1 = merged.increases.get(&key1).unwrap(); - assert_eq!(merged_key1.starting_measurement.value, 10.0); // Earlier start - assert_eq!(merged_key1.last_seen_measurement.value, 30.0); // Later end - } - - #[test] - fn test_keyed_counter_state_serialization() { - let mut acc = KeyedCounterState::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let second_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - let mut reset_aware = create_test_increase_accumulator(10.0, 25.0); - reset_aware.update(Measurement::new(3.0), 3000); - acc.update(key.clone(), reset_aware); - acc.update( - second_key.clone(), - create_test_increase_accumulator(4.0, 9.0), - ); - - // Test JSON serialization - let json_value = acc.serialize_to_json(); - let deserialized = KeyedCounterState::deserialize_from_json(&json_value).unwrap(); - - assert_eq!(deserialized.increases.len(), 2); - let deserialized_acc = deserialized.increases.get(&key).unwrap(); - assert_eq!(deserialized_acc.starting_measurement.value, 10.0); - assert_eq!(deserialized_acc.last_seen_measurement.value, 3.0); - assert_eq!(deserialized_acc.total_increase, 18.0); - - // Test binary serialization - let bytes = acc.serialize_to_bytes(); - let deserialized_bytes = KeyedCounterState::deserialize_from_bytes(&bytes).unwrap(); - - assert_eq!(deserialized_bytes.increases.len(), 2); - let deserialized_acc_bytes = deserialized_bytes.increases.get(&key).unwrap(); - assert_eq!(deserialized_acc_bytes.starting_measurement.value, 10.0); - assert_eq!(deserialized_acc_bytes.last_seen_measurement.value, 3.0); - assert_eq!(deserialized_acc_bytes.total_increase, 18.0); - assert_eq!( - deserialized_bytes - .increases - .get(&second_key) - .unwrap() - .last_seen_measurement - .value, - 9.0 - ); - } - - #[test] - fn test_keyed_counter_state_get_keys() { - let mut acc = KeyedCounterState::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - acc.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); - acc.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); - - let keys = acc.get_keys().unwrap(); - assert_eq!(keys.len(), 2); - assert!(keys.contains(&key1)); - assert!(keys.contains(&key2)); - } - - #[test] - fn test_trait_object() { - let mut acc = KeyedCounterState::new(); - let key = KeyByLabelValues::new(); - acc.update(key.clone(), create_test_increase_accumulator(10.0, 25.0)); - - let trait_obj: Box = Box::new(acc); - assert_eq!( - trait_obj.query(Statistic::Increase, &key, None).unwrap(), - 15.0 - ); - - let keys = trait_obj.get_keys().unwrap(); - assert_eq!(keys.len(), 1); - } - - // #[test] - // fn test_keyed_counter_state_arroyo_deserialization() { - // // Create test data in Arroyo MessagePack format - // // Format: {key: [starting_value, starting_timestamp, last_seen_value, last_seen_timestamp]} - // let mut test_data = std::collections::HashMap::new(); - // test_data.insert("web;service".to_string(), vec![10.0, 1000.0, 25.0, 2000.0]); - // test_data.insert("api;service".to_string(), vec![5.0, 1500.0, 15.0, 2500.0]); - - // // Serialize to MessagePack - // let arroyo_buffer = rmp_serde::to_vec(&test_data).unwrap(); - - // // Test Arroyo deserialization - // let deserialized_acc = - // KeyedCounterState::deserialize_from_bytes_arroyo(&arroyo_buffer).unwrap(); - - // // Verify the deserialized accumulator has the correct data - // assert_eq!(deserialized_acc.increases.len(), 2); - - // // Check first key (web;service) - // let keys: Vec<_> = deserialized_acc.increases.keys().collect(); - // let key1 = keys - // .iter() - // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "web")) - // .unwrap(); - - // let increase1 = deserialized_acc.increases.get(key1).unwrap(); - // assert_eq!(increase1.starting_measurement.value, 10.0); - // assert_eq!(increase1.starting_timestamp, 1000); - // assert_eq!(increase1.last_seen_measurement.value, 25.0); - // assert_eq!(increase1.last_seen_timestamp, 2000); - - // // Check second key (api;service) - // let key2 = keys - // .iter() - // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "api")) - // .unwrap(); - - // let increase2 = deserialized_acc.increases.get(key2).unwrap(); - // assert_eq!(increase2.starting_measurement.value, 5.0); - // assert_eq!(increase2.starting_timestamp, 1500); - // assert_eq!(increase2.last_seen_measurement.value, 15.0); - // assert_eq!(increase2.last_seen_timestamp, 2500); - - // // Test querying - // assert_eq!( - // deserialized_acc.query(Statistic::Increase, key1).unwrap(), - // 15.0 - // ); // 25.0 - 10.0 - // assert_eq!( - // deserialized_acc.query(Statistic::Increase, key2).unwrap(), - // 10.0 - // ); // 15.0 - 5.0 - // } -} diff --git a/crates/asap-physical-operators/src/accumulators/keyed_max_state.rs b/crates/asap-physical-operators/src/accumulators/keyed_max_state.rs deleted file mode 100644 index 30a4a666b..000000000 --- a/crates/asap-physical-operators/src/accumulators/keyed_max_state.rs +++ /dev/null @@ -1,335 +0,0 @@ -use crate::{ - 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::MultipleSubpopulationAggregate; - let key_val = key.as_ref().ok_or("Key required for KeyedMaxState")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for KeyedMaxState { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Max => self - .values - .get(key) - .copied() - .ok_or_else(|| format!("Key {key} not found in KeyedMaxState").into()), - other => Err(format!("Unsupported statistic in KeyedMaxState: {other:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for KeyedMaxState { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut result = KeyedMaxState::new(); - - for acc in accumulators { - for (key, value) in acc.values { - result.update(key, value); - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn key(value: &str) -> KeyByLabelValues { - KeyByLabelValues::new_with_labels(vec![value.to_string()]) - } - - #[test] - fn keeps_the_largest_per_key() { - let mut acc = KeyedMaxState::new(); - acc.update(key("a"), 10.0); - acc.update(key("a"), 5.0); - acc.update(key("a"), 15.0); - acc.update(key("b"), 7.0); - - assert_eq!(acc.query(Statistic::Max, &key("a"), None).unwrap(), 15.0); - assert_eq!(acc.query(Statistic::Max, &key("b"), None).unwrap(), 7.0); - } - - #[test] - fn refuses_the_opposite_statistic_and_unknown_keys() { - let mut acc = KeyedMaxState::new(); - acc.update(key("a"), 1.0); - assert!(acc.query(Statistic::Min, &key("a"), None).is_err()); - assert!(acc.query(Statistic::Max, &key("missing"), None).is_err()); - } - - #[test] - fn merges_per_key() { - let mut left = KeyedMaxState::new(); - left.update(key("a"), 10.0); - let mut right = KeyedMaxState::new(); - right.update(key("a"), 5.0); - right.update(key("b"), 3.0); - - let merged = - >::merge_accumulators(vec![ - left, right, - ]) - .unwrap(); - - assert_eq!(merged.query(Statistic::Max, &key("a"), None).unwrap(), 10.0); - assert_eq!(merged.query(Statistic::Max, &key("b"), None).unwrap(), 3.0); - } - - #[test] - fn refuses_to_merge_with_the_opposite_direction() { - use super::super::keyed_min_state::KeyedMinState; - let mine = KeyedMaxState::new(); - let theirs = KeyedMinState::new(); - assert!(mine.merge_with(&theirs).is_err()); - } - - #[test] - fn round_trips_through_both_serializations() { - let mut acc = KeyedMaxState::new(); - acc.update(key("a"), 4.0); - - let json = acc.serialize_to_json(); - let from_json = KeyedMaxState::deserialize_from_json(&json).unwrap(); - assert_eq!( - from_json.query(Statistic::Max, &key("a"), None).unwrap(), - 4.0 - ); - - let bytes = acc.serialize_to_bytes(); - let from_bytes = KeyedMaxState::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!( - from_bytes.query(Statistic::Max, &key("a"), None).unwrap(), - 4.0 - ); - } -} diff --git a/crates/asap-physical-operators/src/accumulators/keyed_min_state.rs b/crates/asap-physical-operators/src/accumulators/keyed_min_state.rs deleted file mode 100644 index f6bbf2be9..000000000 --- a/crates/asap-physical-operators/src/accumulators/keyed_min_state.rs +++ /dev/null @@ -1,335 +0,0 @@ -use crate::{ - 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::MultipleSubpopulationAggregate; - let key_val = key.as_ref().ok_or("Key required for KeyedMinState")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for KeyedMinState { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - match statistic { - Statistic::Min => self - .values - .get(key) - .copied() - .ok_or_else(|| format!("Key {key} not found in KeyedMinState").into()), - other => Err(format!("Unsupported statistic in KeyedMinState: {other:?}").into()), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for KeyedMinState { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let mut result = KeyedMinState::new(); - - for acc in accumulators { - for (key, value) in acc.values { - result.update(key, value); - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use super::*; - - fn key(value: &str) -> KeyByLabelValues { - KeyByLabelValues::new_with_labels(vec![value.to_string()]) - } - - #[test] - fn keeps_the_smallest_per_key() { - let mut acc = KeyedMinState::new(); - acc.update(key("a"), 10.0); - acc.update(key("a"), 5.0); - acc.update(key("a"), 15.0); - acc.update(key("b"), 7.0); - - assert_eq!(acc.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); - assert_eq!(acc.query(Statistic::Min, &key("b"), None).unwrap(), 7.0); - } - - #[test] - fn refuses_the_opposite_statistic_and_unknown_keys() { - let mut acc = KeyedMinState::new(); - acc.update(key("a"), 1.0); - assert!(acc.query(Statistic::Max, &key("a"), None).is_err()); - assert!(acc.query(Statistic::Min, &key("missing"), None).is_err()); - } - - #[test] - fn merges_per_key() { - let mut left = KeyedMinState::new(); - left.update(key("a"), 10.0); - let mut right = KeyedMinState::new(); - right.update(key("a"), 5.0); - right.update(key("b"), 3.0); - - let merged = - >::merge_accumulators(vec![ - left, right, - ]) - .unwrap(); - - assert_eq!(merged.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); - assert_eq!(merged.query(Statistic::Min, &key("b"), None).unwrap(), 3.0); - } - - #[test] - fn refuses_to_merge_with_the_opposite_direction() { - use super::super::keyed_max_state::KeyedMaxState; - let mine = KeyedMinState::new(); - let theirs = KeyedMaxState::new(); - assert!(mine.merge_with(&theirs).is_err()); - } - - #[test] - fn round_trips_through_both_serializations() { - let mut acc = KeyedMinState::new(); - acc.update(key("a"), 4.0); - - let json = acc.serialize_to_json(); - let from_json = KeyedMinState::deserialize_from_json(&json).unwrap(); - assert_eq!( - from_json.query(Statistic::Min, &key("a"), None).unwrap(), - 4.0 - ); - - let bytes = acc.serialize_to_bytes(); - let from_bytes = KeyedMinState::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!( - from_bytes.query(Statistic::Min, &key("a"), None).unwrap(), - 4.0 - ); - } -} diff --git a/crates/asap-physical-operators/src/accumulators/keyed_sum_count_accumulator.rs b/crates/asap-physical-operators/src/accumulators/keyed_sum_count_accumulator.rs deleted file mode 100644 index 486b3fee3..000000000 --- a/crates/asap-physical-operators/src/accumulators/keyed_sum_count_accumulator.rs +++ /dev/null @@ -1,558 +0,0 @@ -use crate::{ - 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::MultipleSubpopulationAggregate; - let key_val = key - .as_ref() - .ok_or("Key required for KeyedSumCountAccumulator")?; - self.query(statistic, key_val, Some(query_kwargs)) - } -} - -impl MultipleSubpopulationAggregate for KeyedSumCountAccumulator { - fn query( - &self, - statistic: Statistic, - key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - match (&self.family, statistic) { - (ExactKind::Sum, Statistic::Sum) => self.sums.get(key).copied().ok_or_else(|| { - "Key not found in KeyedSumCountAccumulator" - .to_string() - .into() - }), - (ExactKind::Count, Statistic::Count) => self - .counts - .get(key) - .map(|count| *count as f64) - .ok_or_else(|| { - "Sample count unavailable in KeyedSumCountAccumulator" - .to_string() - .into() - }), - _ => Err( - format!("Unsupported statistic in KeyedSumCountAccumulator: {statistic:?}").into(), - ), - } - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} - -impl MergeableAccumulator for KeyedSumCountAccumulator { - fn merge_accumulators( - accumulators: Vec, - ) -> Result> { - if accumulators.is_empty() { - return Err("No accumulators to merge".into()); - } - - let family = accumulators[0].family.clone(); - if accumulators.iter().any(|acc| acc.family != family) { - return Err("Cannot merge different keyed additive families".into()); - } - let mut result = KeyedSumCountAccumulator::for_family(family); - - for acc in accumulators { - for key in acc.sums.keys() { - match ( - result.counts.get(key).copied(), - acc.counts.get(key).copied(), - ) { - (None, Some(count)) if !result.sums.contains_key(key) => { - result.counts.insert(key.clone(), count); - } - (Some(existing), Some(count)) => { - if let Some(total) = existing.checked_add(count) { - result.counts.insert(key.clone(), total); - } else { - result.counts.remove(key); - } - } - _ => { - result.counts.remove(key); - } - } - } - for (key, sum) in acc.sums { - *result.sums.entry(key).or_insert(0.0) += sum; - } - } - - Ok(result) - } -} - -#[cfg(test)] -mod tests { - use std::vec; - - use super::*; - - #[test] - fn test_keyed_sum_count_accumulator_creation() { - let acc = KeyedSumCountAccumulator::new(); - assert!(acc.sums.is_empty()); - } - - #[test] - fn test_keyed_sum_count_accumulator_update() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - acc.update(key1.clone(), 10.0); - acc.update(key2.clone(), 20.0); - acc.update(key1.clone(), 5.0); // Should add to existing - - assert_eq!(acc.sums.get(&key1), Some(&15.0)); - assert_eq!(acc.sums.get(&key2), Some(&20.0)); - } - - #[test] - fn grouped_count_reads_sample_count_and_survives_merge_and_round_trip() { - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - let mut first = KeyedSumCountAccumulator::for_family(ExactKind::Count); - first.update(key.clone(), 10.0); - first.update(key.clone(), 20.0); - let mut second = KeyedSumCountAccumulator::for_family(ExactKind::Count); - second.update(key.clone(), 7.0); - let merged = KeyedSumCountAccumulator::merge_accumulators(vec![first, second]).unwrap(); - for acc in [ - merged.clone(), - KeyedSumCountAccumulator::deserialize_from_json(&merged.serialize_to_json()).unwrap(), - KeyedSumCountAccumulator::deserialize_from_bytes(&merged.serialize_to_bytes()).unwrap(), - ] { - assert_eq!(acc.family, ExactKind::Count); - assert!(acc.query(Statistic::Sum, &key, None).is_err()); - assert_eq!(acc.query(Statistic::Count, &key, None).unwrap(), 3.0); - } - } - - #[test] - fn keyed_additive_merge_rejects_different_planner_families() { - assert!(KeyedSumCountAccumulator::merge_accumulators(vec![ - KeyedSumCountAccumulator::for_family(ExactKind::Sum), - KeyedSumCountAccumulator::for_family(ExactKind::Count), - ]) - .is_err()); - } - - #[test] - fn test_keyed_sum_count_accumulator_query() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); - - acc.add_sum(key.clone(), 42.0); - - // Test total queries (querying with the specific key) - assert_eq!( - crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Sum, &key, None).unwrap(), - 42.0 - ); - - // Test error cases - assert!( - crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Min, &key, None).is_err() - ); - } - - #[test] - fn test_keyed_sum_count_accumulator_get_keys() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - acc.add_sum(key1.clone(), 10.0); - acc.add_sum(key2.clone(), 20.0); - - let keys = crate::AggregateCore::get_keys(&acc).unwrap(); - assert_eq!(keys.len(), 2); - assert!(keys.contains(&key1)); - assert!(keys.contains(&key2)); - } - - #[test] - fn test_keyed_sum_count_accumulator_merge() { - let mut acc1 = KeyedSumCountAccumulator::new(); - let mut acc2 = KeyedSumCountAccumulator::new(); - - let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); - - acc1.add_sum(key1.clone(), 10.0); - acc1.add_sum(key2.clone(), 20.0); - - acc2.add_sum(key1.clone(), 5.0); // Same key, different accumulator - - let merged = >::merge_accumulators(vec![acc1, acc2]).unwrap(); - - assert_eq!(merged.sums.get(&key1), Some(&15.0)); // Should be merged - assert_eq!(merged.sums.get(&key2), Some(&20.0)); // Should be preserved - } - - #[test] - fn test_keyed_sum_count_accumulator_serialization() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); - - acc.add_sum(key.clone(), 42.5); - - // Test JSON serialization - let json = acc.serialize_to_json(); - let deserialized = KeyedSumCountAccumulator::deserialize_from_json(&json).unwrap(); - assert_eq!(deserialized.sums.get(&key), Some(&42.5)); - - // Test byte serialization - let bytes = acc.serialize_to_bytes(); - let deserialized_bytes = KeyedSumCountAccumulator::deserialize_from_bytes(&bytes).unwrap(); - assert_eq!(deserialized_bytes.sums.get(&key), Some(&42.5)); - } - - #[test] - fn test_trait_object() { - let mut acc = KeyedSumCountAccumulator::new(); - - let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); - - acc.add_sum(key.clone(), 42.0); - - let trait_obj: Box = Box::new(acc); - - // Test type name through trait object - assert_eq!(trait_obj.type_name(), "KeyedSumCountAccumulator"); - } -} diff --git a/crates/asap-physical-operators/src/accumulators/max_accumulator.rs b/crates/asap-physical-operators/src/accumulators/max_accumulator.rs deleted file mode 100644 index 235b5fc98..000000000 --- a/crates/asap-physical-operators/src/accumulators/max_accumulator.rs +++ /dev/null @@ -1,248 +0,0 @@ -use crate::{ - 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::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/crates/asap-physical-operators/src/accumulators/min_accumulator.rs b/crates/asap-physical-operators/src/accumulators/min_accumulator.rs deleted file mode 100644 index fff2fa0e9..000000000 --- a/crates/asap-physical-operators/src/accumulators/min_accumulator.rs +++ /dev/null @@ -1,253 +0,0 @@ -use crate::{ - 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::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/crates/asap-physical-operators/src/accumulators/mod.rs b/crates/asap-physical-operators/src/accumulators/mod.rs deleted file mode 100644 index 073db6e82..000000000 --- a/crates/asap-physical-operators/src/accumulators/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/crates/asap-physical-operators/src/accumulators/sketch_envelope_accumulator.rs b/crates/asap-physical-operators/src/accumulators/sketch_envelope_accumulator.rs deleted file mode 100644 index 930dbe6d9..000000000 --- a/crates/asap-physical-operators/src/accumulators/sketch_envelope_accumulator.rs +++ /dev/null @@ -1,154 +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::{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::MultipleSubpopulationAggregate for SketchEnvelopeAccumulator { - fn query( - &self, - _statistic: Statistic, - _key: &KeyByLabelValues, - _query_kwargs: Option<&HashMap>, - ) -> Result> { - Err( - "SketchEnvelopeAccumulator: direct query not supported; use native sketch query path" - .into(), - ) - } - - fn clone_boxed(&self) -> Box { - Box::new(self.clone()) - } -} diff --git a/crates/asap-physical-operators/src/accumulators/sum_accumulator.rs b/crates/asap-physical-operators/src/accumulators/sum_accumulator.rs deleted file mode 100644 index c5293911d..000000000 --- a/crates/asap-physical-operators/src/accumulators/sum_accumulator.rs +++ /dev/null @@ -1,413 +0,0 @@ -use crate::{ - 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::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/crates/asap-physical-operators/src/accumulators/univmon_accumulator.rs b/crates/asap-physical-operators/src/accumulators/univmon_accumulator.rs deleted file mode 100644 index 2d18895e0..000000000 --- a/crates/asap-physical-operators/src/accumulators/univmon_accumulator.rs +++ /dev/null @@ -1,234 +0,0 @@ -//! One frequency state shared by count, distinct, L2 and entropy readouts. - -use crate::{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/crates/asap-physical-operators/src/arithmetic.rs b/crates/asap-physical-operators/src/arithmetic.rs deleted file mode 100644 index bfc50694b..000000000 --- a/crates/asap-physical-operators/src/arithmetic.rs +++ /dev/null @@ -1,19 +0,0 @@ -//! Float64 arithmetic shared by ASAP execution engines. -//! Preserve IEEE non-finite results; callers own their output policies. - -pub 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/crates/asap-physical-operators/src/capability.rs b/crates/asap-physical-operators/src/capability.rs deleted file mode 100644 index fdd2fee7a..000000000 --- a/crates/asap-physical-operators/src/capability.rs +++ /dev/null @@ -1,115 +0,0 @@ -//! Allocation-free checks for the concrete summary kernels in this crate. -use planner_types::post_asap::{ - ExactKind, ExactParams, GroupingStrategy, SketchAlgorithm, SketchParams, SummaryFamilyType, - SummaryUpdate, -}; - -/// Check the same contract used by `create_planner_accumulator` before a plan -/// is accepted. Execution timing is deliberately not a kernel property. -pub fn validate_summary_kernel( - family: &SummaryFamilyType, - input: &SummaryUpdate, - grouping: &GroupingStrategy, -) -> Result<(), String> { - if grouping != &GroupingStrategy::PerSubpopulationInstance { - return Err("shared summary grouping has no registered kernel".into()); - } - let keyed = match family { - SummaryFamilyType::ExactAggregate(kind, params) => { - use ExactKind as K; - use ExactParams as P; - if !matches!( - (kind, params), - (K::Sum, P::Sum) - | (K::Count, P::Count) - | (K::Min, P::Min) - | (K::Max, P::Max) - | (K::Rate, P::Rate) - | (K::Increase, P::Increase) - ) { - return Err(format!("unsupported exact kernel {family:?}")); - } - input.item.is_some() - } - SummaryFamilyType::Sketch(kind, layout) => { - if layout != grouping { - return Err("Planner family and operator grouping disagree".into()); - } - use SketchAlgorithm as A; - use SketchParams as P; - match (kind.algorithm(), kind.params()) { - (A::Kll, P::Kll { k }) if (8..=u16::MAX as u32).contains(k) => false, - (A::DDSketch, P::DDSketch { alpha }) - if alpha.is_finite() && *alpha > 0.0 && *alpha < 1.0 => - { - false - } - (A::Hll, P::Hll { precision }) if (4..=18).contains(precision) => false, - (A::Cms, P::Cms { width, depth }) - | (A::CountSketch, P::CountSketch { width, depth }) - if valid_matrix(*width, *depth) => - { - true - } - ( - A::CmsWithHeap, - P::CmsWithHeap { - width, - depth, - heap_size, - }, - ) - | ( - A::CountSketchWithHeap, - P::CountSketchWithHeap { - width, - depth, - heap_size, - }, - ) if valid_matrix(*width, *depth) && *heap_size > 0 => true, - ( - A::UnivMon, - P::UnivMon { - heap_size, - sketch_rows, - sketch_cols, - layers, - }, - ) if *heap_size > 0 - && *sketch_cols > 0 - && (1..=20).contains(sketch_rows) - && (1..=64).contains(layers) - && (*sketch_rows as usize) - .checked_mul(*sketch_cols as usize) - .and_then(|n| n.checked_mul(*layers as usize)) - .is_some() => - { - false - } - _ => { - return Err(format!( - "unsupported kernel or invalid parameters: {kind:?}" - )) - } - } - } - _ => return Err(format!("unsupported summary kernel {family:?}")), - }; - if keyed != input.item.is_some() - && !asap_types::accumulator_spec::is_unit_sample_frequency(input) - { - return Err("Planner item expression does not match kernel layout".into()); - } - Ok(()) -} - -fn valid_matrix(width: u32, depth: u32) -> bool { - // Construction uses the kernel's native row hashing. Packed-wire decoder - // limits describe a different representation and must not reject it here. - width > 0 - && depth > 0 - && (width as usize) - .checked_mul(depth as usize) - .and_then(|n| n.checked_mul(std::mem::size_of::())) - .is_some() -} diff --git a/crates/asap-physical-operators/src/dag/batch_execution.rs b/crates/asap-physical-operators/src/dag/batch_execution.rs deleted file mode 100644 index 78b521724..000000000 --- a/crates/asap-physical-operators/src/dag/batch_execution.rs +++ /dev/null @@ -1,181 +0,0 @@ -//! Execute a bounded in-memory batch through native operators. This is also the -//! bridge for deployments whose boundary values are not yet streaming batches. -use super::{operators::Operator, values::Batch, Error, PhysicalDag, RunContext, SharedValue}; -use futures::{FutureExt, StreamExt}; - -/// Every input is already in memory; the chain contains native operators only. -/// This deliberately does not enter a nested executor when called from a DAG -/// adapter. I/O belongs to source operators in the surrounding execution. -pub fn evaluate_batch( - input: Batch, - operators: Vec, - context: RunContext, -) -> Result>, Error> { - let mut graph = PhysicalDag::default(); - graph.add( - 0, - vec![], - Operator::source(input.schema().clone(), vec![input])?, - )?; - let mut root = 0; - for operator in operators { - graph.add(root + 1, vec![root], operator)?; - root += 1; - } - evaluate_graph(graph, root, context) -} - -/// Evaluate a native in-memory source, including scalar sources, in the caller's scope. -pub fn evaluate_source( - source: Operator, - context: RunContext, -) -> Result>, Error> { - let mut graph = PhysicalDag::default(); - graph.add(0, vec![], source)?; - evaluate_graph(graph, 0, context) -} - -fn evaluate_graph( - graph: PhysicalDag<'_, Batch, super::values::Schema>, - root: super::NodeId, - context: RunContext, -) -> Result>, Error> { - let mut output = graph.execute(&[root], context)?.remove(0); - let mut batches = Vec::new(); - loop { - match output.next().now_or_never() { - Some(Some(Ok(batch))) => batches.push(batch), - Some(Some(Err(error))) => return Err(error), - Some(None) => return Ok(batches), - // Native operators have no I/O sources here. Pending is the - // shared runtime's cooperative yield after a batch quantum. - None => continue, - } - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::dag::{operators::Expression, values::Value, Limits, Scope}; - use planner_types::{ - post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, - pre_asap::DataType, - }; - use std::sync::Arc; - - // Engine adapters can run the identical native chain from an outer executor. - #[test] - fn same_native_chain_inside_query_and_ingestion_execution() { - let schema = Arc::new(SummarySchema { - fields: vec![SummaryField { - name: "value".into(), - dtype: SummaryFamilyType::Plain(DataType::Float64), - nullable: false, - }], - time_index: None, - }); - for scope in [ - Scope::Query { - evaluation_time_ms: 20, - revision: 1, - }, - Scope::Ingestion { - window_start_ms: 10, - window_end_ms: 20, - revision: 1, - }, - ] { - let batch = Batch::try_new(schema.clone(), vec![vec![Value::Float64(7.)]]).unwrap(); - let negate = Operator::project( - schema.clone(), - vec![( - "value".into(), - Expression::Negate(Box::new(Expression::Column(0))), - )], - ) - .unwrap(); - let context = RunContext::new(scope, Limits::default()).unwrap(); - let result = futures::executor::block_on(async { - evaluate_batch(batch, vec![negate], context.clone()) - }) - .unwrap(); - assert!(matches!(result[0].rows()[0][0], Value::Float64(-7.))); - let source = Operator::scalar(Value::Float64(9.), DataType::Float64).unwrap(); - let scalar = evaluate_source(source, context).unwrap(); - assert!(matches!(scalar[0].rows()[0][0], Value::Float64(9.))); - } - } - - // Native sources may cross the runtime's cooperative batch quantum. - #[test] - fn in_memory_source_drives_cooperative_yields() { - let schema = Arc::new(SummarySchema { - fields: vec![], - time_index: None, - }); - let batch = Batch::try_new(schema.clone(), vec![vec![]]).unwrap(); - let source = Operator::source(schema, vec![batch; 65]).unwrap(); - let context = RunContext::new( - Scope::Query { - evaluation_time_ms: 0, - revision: 0, - }, - Limits::default(), - ) - .unwrap(); - assert_eq!(evaluate_source(source, context).unwrap().len(), 65); - } - - // An adapter-held output must retain its parent's reservation after execution. - #[test] - fn returned_batches_keep_their_resource_reservation() { - let schema = Arc::new(SummarySchema { - fields: vec![], - time_index: None, - }); - let batch = Batch::try_new(schema.clone(), vec![vec![]]).unwrap(); - let bytes = batch.bytes(); - let source = Operator::source(schema, vec![batch]).unwrap(); - let context = RunContext::new( - Scope::Query { - evaluation_time_ms: 0, - revision: 0, - }, - Limits { - max_bytes: bytes, - max_buffered_batches: 1, - }, - ) - .unwrap(); - let held = evaluate_source(source.clone(), context.clone()).unwrap(); - assert_eq!(context.retained_bytes(), bytes); - assert!(evaluate_source(source.clone(), context.clone()).is_err()); - drop(held); - assert_eq!(context.retained_bytes(), 0); - assert!(evaluate_source(source, context).is_ok()); - } - - // A cancelled surrounding execution also prevents its native computation. - #[test] - fn cancellation_is_not_bypassed_by_in_memory_execution() { - let schema = Arc::new(SummarySchema { - fields: vec![], - time_index: None, - }); - let batch = Batch::try_new(schema, vec![vec![]]).unwrap(); - let context = RunContext::new( - Scope::Query { - evaluation_time_ms: 0, - revision: 0, - }, - Limits::default(), - ) - .unwrap(); - context.cancel(); - assert!(matches!( - evaluate_batch(batch, vec![], context), - Err(Error::Cancelled) - )); - } -} diff --git a/crates/asap-physical-operators/src/dag/mod.rs b/crates/asap-physical-operators/src/dag/mod.rs deleted file mode 100644 index 1c9ee9e86..000000000 --- a/crates/asap-physical-operators/src/dag/mod.rs +++ /dev/null @@ -1,517 +0,0 @@ -//! Independent operator DAG execution. No backend plan or engine types are used. -//! -//! Each run creates one stream per reachable node. Consumers subscribe to that -//! stream independently; retained outputs are released after the last consumer. -use futures::{stream::LocalBoxStream, Stream}; -use std::{ - cell::{Cell, RefCell}, - collections::{BTreeMap, BTreeSet, VecDeque}, - fmt::Debug, - pin::Pin, - rc::Rc, - sync::Arc, - task::{Context, Poll, Waker}, -}; - -pub type NodeId = u64; -pub type OutputStream<'a, V> = LocalBoxStream<'a, Result>; -#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)] -pub enum Error { - #[error("invalid DAG: {0}")] - Invalid(String), - #[error("operator failed: {0}")] - Operator(String), - #[error("node {node} ({operation}) failed: {source}")] - AtNode { - node: NodeId, - operation: String, - source: Box, - }, - #[error("execution memory limit exceeded")] - MemoryLimit, - #[error("execution cancelled")] - Cancelled, -} - -/// Scope is part of an execution instance, never mutable state in a reusable plan. -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum Scope { - Ingestion { - window_start_ms: i64, - window_end_ms: i64, - revision: u64, - }, - Query { - evaluation_time_ms: i64, - revision: u64, - }, -} -#[derive(Clone, Debug)] -pub struct Limits { - pub max_buffered_batches: usize, - pub max_bytes: usize, -} -impl Default for Limits { - fn default() -> Self { - Self { - max_buffered_batches: 8, - max_bytes: 64 * 1024 * 1024, - } - } -} -struct Control { - cancelled: Cell, - bytes: Cell, - peak: Cell, - limits: Limits, - waiters: RefCell>, -} -#[derive(Clone)] -pub struct RunContext { - pub scope: Scope, - control: Rc, -} -impl RunContext { - pub fn new(scope: Scope, limits: Limits) -> Result { - if limits.max_buffered_batches == 0 || limits.max_bytes == 0 { - return Err(Error::Invalid("execution limits must be positive".into())); - } - if matches!(&scope, Scope::Ingestion { window_start_ms, window_end_ms, .. } if window_start_ms > window_end_ms) - { - return Err(Error::Invalid("inverted ingestion window".into())); - } - Ok(Self { - scope, - control: Rc::new(Control { - cancelled: Cell::new(false), - bytes: Cell::new(0), - peak: Cell::new(0), - limits, - waiters: RefCell::new(Vec::new()), - }), - }) - } - pub fn cancel(&self) { - self.control.cancelled.set(true); - for waiter in self.control.waiters.borrow_mut().drain(..) { - waiter.wake(); - } - } - pub fn is_cancelled(&self) -> bool { - self.control.cancelled.get() - } - pub fn retained_bytes(&self) -> usize { - self.control.bytes.get() - } - pub fn peak_bytes(&self) -> usize { - self.control.peak.get() - } - pub fn reserve(&self, bytes: usize) -> Result { - let total = self - .control - .bytes - .get() - .checked_add(bytes) - .ok_or(Error::MemoryLimit)?; - if total > self.control.limits.max_bytes { - return Err(Error::MemoryLimit); - } - self.control.bytes.set(total); - self.control.peak.set(self.control.peak.get().max(total)); - Ok(Reservation { - bytes, - control: Rc::clone(&self.control), - }) - } - fn register(&self, waker: &Waker) { - let mut waiters = self.control.waiters.borrow_mut(); - if !waiters.iter().any(|old| old.will_wake(waker)) { - waiters.push(waker.clone()); - } - } -} -pub struct Reservation { - bytes: usize, - control: Rc, -} -impl Reservation { - /// Adjust an operator-owned allocation without accumulating bookkeeping entries. - pub fn resize(&mut self, bytes: usize) -> Result<(), Error> { - let total = self - .control - .bytes - .get() - .checked_sub(self.bytes) - .and_then(|total| total.checked_add(bytes)) - .ok_or(Error::MemoryLimit)?; - if total > self.control.limits.max_bytes { - return Err(Error::MemoryLimit); - } - self.control.bytes.set(total); - self.control.peak.set(self.control.peak.get().max(total)); - self.bytes = bytes; - Ok(()) - } -} -impl Drop for Reservation { - fn drop(&mut self) { - self.control - .bytes - .set(self.control.bytes.get().saturating_sub(self.bytes)); - } -} - -/// An output owns its memory reservation even after it leaves the DAG's queue. -pub struct SharedValue { - value: Arc, - _reservation: Rc, -} -impl Clone for SharedValue { - fn clone(&self) -> Self { - Self { - value: Arc::clone(&self.value), - _reservation: Rc::clone(&self._reservation), - } - } -} -impl std::ops::Deref for SharedValue { - type Target = V; - fn deref(&self) -> &V { - &self.value - } -} -impl SharedValue { - pub fn value(&self) -> &V { - &self.value - } -} - -/// Operators own computation. The runtime provides already-connected inputs; -/// an operator must not recursively execute another plan node itself. -pub trait PhysicalOperator { - fn name(&self) -> &str; - fn input_schemas(&self) -> Vec; - fn output_schema(&self) -> S; - fn start<'a>( - &'a self, - inputs: Vec>, - context: RunContext, - ) -> Result, Error>; - fn output_bytes(&self, value: &V) -> usize; -} -struct Node<'a, V, S> { - inputs: Vec, - operator: Box + 'a>, -} -pub struct PhysicalDag<'a, V, S> { - nodes: BTreeMap>, -} -impl Default for PhysicalDag<'_, V, S> { - fn default() -> Self { - Self { - nodes: BTreeMap::new(), - } - } -} -impl<'a, V: 'a, S: Clone + PartialEq + Debug + 'a> PhysicalDag<'a, V, S> { - pub fn add( - &mut self, - id: NodeId, - inputs: Vec, - operator: impl PhysicalOperator + 'a, - ) -> Result<(), Error> { - self.add_boxed(id, inputs, Box::new(operator)) - } - pub fn add_boxed( - &mut self, - id: NodeId, - inputs: Vec, - operator: Box + 'a>, - ) -> Result<(), Error> { - if self.nodes.contains_key(&id) { - return Err(Error::Invalid(format!("duplicate node {id}"))); - } - self.nodes.insert(id, Node { inputs, operator }); - Ok(()) - } - pub fn validate(&self, roots: &[NodeId]) -> Result<(), Error> { - fn visit( - dag: &PhysicalDag<'_, V, S>, - id: NodeId, - active: &mut BTreeSet, - done: &mut BTreeMap, - ) -> Result { - if let Some(depth) = done.get(&id) { - return Ok(*depth); - } - if active.len() >= 128 { - return Err(Error::Invalid( - "DAG exceeds the supported execution depth of 128".into(), - )); - } - if !active.insert(id) { - return Err(Error::Invalid(format!("cycle at node {id}"))); - } - let node = dag - .nodes - .get(&id) - .ok_or_else(|| Error::Invalid(format!("missing node {id}")))?; - let expected = node.operator.input_schemas(); - if expected.len() != node.inputs.len() { - return Err(Error::Invalid(format!("node {id} input arity mismatch"))); - } - let mut depth = 1; - for (input, schema) in node.inputs.iter().zip(expected) { - depth = depth.max(1 + visit(dag, *input, active, done)?); - let actual = dag.nodes[input].operator.output_schema(); - if actual != schema { - return Err(Error::Invalid(format!( - "node {id} input {input} schema mismatch: {actual:?} vs {schema:?}" - ))); - } - } - if depth > 128 { - return Err(Error::Invalid( - "DAG exceeds the supported execution depth of 128".into(), - )); - } - active.remove(&id); - done.insert(id, depth); - Ok(depth) - } - if roots.is_empty() { - return Err(Error::Invalid("execution needs a root".into())); - } - let mut done = BTreeMap::new(); - for &root in roots { - visit(self, root, &mut BTreeSet::new(), &mut done)?; - } - Ok(()) - } - pub fn execute<'r>( - &'r self, - roots: &[NodeId], - context: RunContext, - ) -> Result>, Error> - where - 'a: 'r, - { - if context.is_cancelled() { - return Err(Error::Cancelled); - } - self.validate(roots)?; - fn build<'r, V: 'r, S: 'r>( - dag: &'r PhysicalDag<'_, V, S>, - id: NodeId, - context: &RunContext, - states: &mut BTreeMap>>>, - ) -> Result>>, Error> { - if let Some(state) = states.get(&id) { - return Ok(Rc::clone(state)); - } - let node = &dag.nodes[&id]; - let mut inputs = Vec::new(); - for &child in &node.inputs { - inputs.push(Input::subscribe(build(dag, child, context, states)?)); - } - let stream = node - .operator - .start(inputs, context.clone()) - .map_err(|source| Error::AtNode { - node: id, - operation: node.operator.name().into(), - source: Box::new(source), - })?; - let op = node.operator.as_ref(); - let state = Rc::new(RefCell::new(Producer { - stream: Some(stream), - node: id, - operation: node.operator.name().into(), - size: Box::new(move |value| op.output_bytes(value)), - context: context.clone(), - queue: VecDeque::new(), - base: 0, - next_reader: 0, - batches_polled: 0, - readers: BTreeMap::new(), - waiters: BTreeMap::new(), - finished: false, - failure: None, - })); - states.insert(id, Rc::clone(&state)); - Ok(state) - } - let mut states = BTreeMap::new(); - roots - .iter() - .map(|&id| build(self, id, &context, &mut states).map(Input::subscribe)) - .collect() - } -} -struct Producer<'a, V> { - node: NodeId, - operation: String, - stream: Option>, - size: Box usize + 'a>, - context: RunContext, - queue: VecDeque>, - base: u64, - next_reader: u64, - batches_polled: usize, - readers: BTreeMap, - waiters: BTreeMap, - finished: bool, - failure: Option, -} -impl Producer<'_, V> { - fn trim(&mut self) { - let minimum = self - .readers - .values() - .copied() - .min() - .unwrap_or(self.base + self.queue.len() as u64); - while self.base < minimum { - self.queue.pop_front(); - self.base += 1; - } - for (_, waker) in std::mem::take(&mut self.waiters) { - waker.wake(); - } - if self.readers.is_empty() { - self.stream = None; - self.queue.clear(); - } - } -} -pub struct Input<'a, V> { - producer: Rc>>, - reader: u64, - done: bool, -} -impl<'a, V> Input<'a, V> { - fn subscribe(producer: Rc>>) -> Self { - let reader = { - let mut state = producer.borrow_mut(); - let id = state.next_reader; - state.next_reader += 1; - let base = state.base; - state.readers.insert(id, base); - id - }; - Self { - producer, - reader, - done: false, - } - } -} -impl Drop for Input<'_, V> { - fn drop(&mut self) { - let mut state = self.producer.borrow_mut(); - state.readers.remove(&self.reader); - state.waiters.remove(&self.reader); - state.trim(); - } -} -impl Stream for Input<'_, V> { - type Item = Result, Error>; - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - let this = self.get_mut(); - if this.done { - return Poll::Ready(None); - } - let mut state = this.producer.borrow_mut(); - state.context.register(cx.waker()); - if state.context.is_cancelled() { - state.failure = Some(Error::Cancelled); - state.finished = true; - state.stream = None; - state.queue.clear(); - } - let position = state.readers[&this.reader]; - let index = (position - state.base) as usize; - if let Some(value) = state.queue.get(index).cloned() { - state.readers.insert(this.reader, position + 1); - state.trim(); - return Poll::Ready(Some(Ok(value))); - } - if state.finished { - this.done = true; - state.readers.remove(&this.reader); - let failure = state.failure.clone(); - state.trim(); - return Poll::Ready(failure.map(Err)); - } - state.waiters.insert(this.reader, cx.waker().clone()); - if state.queue.len() >= state.context.control.limits.max_buffered_batches { - return Poll::Pending; - } - // Always-ready sources must still give cancellation and other roots a turn. - if state.batches_polled >= 32 { - state.batches_polled = 0; - cx.waker().wake_by_ref(); - return Poll::Pending; - } - let polled = state - .stream - .as_mut() - .expect("unfinished producer") - .as_mut() - .poll_next(cx); - if matches!(&polled, Poll::Ready(Some(Ok(_)))) { - state.batches_polled += 1; - } - match polled { - Poll::Pending => Poll::Pending, - Poll::Ready(Some(Ok(value))) => match state.context.reserve((state.size)(&value)) { - Ok(reservation) => { - let value = SharedValue { - value: Arc::new(value), - _reservation: Rc::new(reservation), - }; - state.queue.push_back(value.clone()); - state.readers.insert(this.reader, position + 1); - state.trim(); - Poll::Ready(Some(Ok(value))) - } - Err(error) => { - state.failure = Some(error.clone()); - state.finished = true; - state.stream = None; - this.done = true; - state.readers.remove(&this.reader); - state.trim(); - Poll::Ready(Some(Err(error))) - } - }, - Poll::Ready(result) => { - let error = result.and_then(Result::err).map(|source| match source { - Error::AtNode { .. } | Error::Cancelled | Error::MemoryLimit => source, - source => Error::AtNode { - node: state.node, - operation: state.operation.clone(), - source: Box::new(source), - }, - }); - state.failure = error.clone(); - state.finished = true; - state.stream = None; - this.done = true; - state.readers.remove(&this.reader); - state.trim(); - Poll::Ready(error.map(Err)) - } - } - } -} - -pub mod operators; -pub mod values; - -#[cfg(test)] -mod tests; - -pub mod planner; - -pub mod batch_execution; diff --git a/crates/asap-physical-operators/src/dag/operators.rs b/crates/asap-physical-operators/src/dag/operators.rs deleted file mode 100644 index eeb1a922e..000000000 --- a/crates/asap-physical-operators/src/dag/operators.rs +++ /dev/null @@ -1,1170 +0,0 @@ -//! Native DAG operators. Engines bind sources; computation lives here. -use super::{ - values::{group_key, Batch, Schema, Value}, - Error, Input, OutputStream, PhysicalOperator, Reservation, RunContext, -}; -use futures::StreamExt; -use planner_types::{ - post_asap::{SummaryFamilyType, SummaryField, SummarySchema, SummaryUpdate}, - pre_asap::{ArithmeticOpKind, ColumnRef, DataType}, -}; -use std::{collections::BTreeMap, sync::Arc}; - -fn invalid(message: &str) -> Error { - Error::Invalid(message.into()) -} -fn field(schema: &Schema, column: usize) -> Result<&SummaryField, Error> { - schema - .fields - .get(column) - .ok_or_else(|| invalid("column out of range")) -} -fn plain(schema: &Schema, column: usize) -> Result<(&DataType, bool), Error> { - let f = field(schema, column)?; - let SummaryFamilyType::Plain(dtype) = &f.dtype else { - return Err(invalid("plain value required")); - }; - Ok((dtype, f.nullable)) -} -fn schema(fields: Vec) -> Schema { - Arc::new(SummarySchema { - fields, - time_index: None, - }) -} -fn result_field(name: &str, dtype: DataType, nullable: bool) -> SummaryField { - SummaryField { - name: name.into(), - dtype: SummaryFamilyType::Plain(dtype), - nullable, - } -} - -#[derive(Clone, Debug)] -pub enum Expression { - Column(usize), - Literal { - value: Value, - dtype: DataType, - }, - Negate(Box), - Arithmetic { - op: ArithmeticOpKind, - left: Box, - right: Box, - }, - Equal(Box, Box), - Less(Box, Box), - And(Box, Box), - Or(Box, Box), - Not(Box), - IsNull(Box), -} -impl Expression { - fn dtype(&self, input: &Schema) -> Result<(DataType, bool), Error> { - use Expression::*; - match self { - Column(i) => { - let (t, n) = plain(input, *i)?; - Ok((t.clone(), n)) - } - Literal { value, dtype } => { - if value.matches(dtype, true) { - Ok((dtype.clone(), matches!(value, Value::Null))) - } else { - Err(invalid("literal type mismatch")) - } - } - Negate(v) => { - let (t, n) = v.dtype(input)?; - if matches!(t, DataType::Int64 | DataType::Float64) { - Ok((t, n)) - } else { - Err(invalid("numeric negation required")) - } - } - Arithmetic { op, left, right } => { - let (a, n) = left.dtype(input)?; - let (b, m) = right.dtype(input)?; - if a == b - && matches!(a, DataType::Int64 | DataType::Float64) - && !(a == DataType::Int64 && *op == ArithmeticOpKind::Atan2) - { - Ok((a, n || m)) - } else { - Err(invalid("arithmetic requires matching numeric types")) - } - } - Equal(a, b) | Less(a, b) => { - let (a, n) = a.dtype(input)?; - let (b, m) = b.dtype(input)?; - if a == b && ordered(&a) { - Ok((DataType::Bool, n || m)) - } else { - Err(invalid("comparison requires matching ordered types")) - } - } - And(a, b) | Or(a, b) => { - let (a, n) = a.dtype(input)?; - let (b, m) = b.dtype(input)?; - if a == DataType::Bool && b == DataType::Bool { - Ok((DataType::Bool, n || m)) - } else { - Err(invalid("boolean operands required")) - } - } - Not(v) => { - let (t, n) = v.dtype(input)?; - if t == DataType::Bool { - Ok((t, n)) - } else { - Err(invalid("boolean operand required")) - } - } - IsNull(v) => { - v.dtype(input)?; - Ok((DataType::Bool, false)) - } - } - } - fn evaluate(&self, row: &[Value]) -> Result { - use Expression::*; - Ok(match self { - Column(i) => row[*i].clone(), - Literal { value, .. } => value.clone(), - Negate(v) => match v.evaluate(row)? { - Value::Int64(v) => Value::Int64( - v.checked_neg() - .ok_or_else(|| invalid("integer negation overflow"))?, - ), - Value::Float64(v) => Value::Float64(-v), - Value::Null => Value::Null, - _ => return Err(invalid("numeric negation required")), - }, - Arithmetic { op, left, right } => { - numeric(op, left.evaluate(row)?, right.evaluate(row)?)? - } - Equal(a, b) | Less(a, b) => { - let (a, b) = (a.evaluate(row)?, b.evaluate(row)?); - if matches!(a, Value::Null) || matches!(b, Value::Null) { - Value::Null - } else if matches!((&a,&b),(Value::Float64(a),Value::Float64(b)) if a.is_nan() || b.is_nan()) - { - Value::Bool(false) - } else { - let c = a.compare(&b)?; - Value::Bool(if matches!(self, Equal(..)) { - c.is_eq() - } else { - c.is_lt() - }) - } - } - And(a, b) | Or(a, b) => { - let (a, b) = (a.evaluate(row)?, b.evaluate(row)?); - match (a, b, matches!(self, And(..))) { - (Value::Bool(false), _, true) | (_, Value::Bool(false), true) => { - Value::Bool(false) - } - (Value::Bool(true), _, false) | (_, Value::Bool(true), false) => { - Value::Bool(true) - } - (Value::Null, _, _) | (_, Value::Null, _) => Value::Null, - (Value::Bool(a), Value::Bool(b), true) => Value::Bool(a && b), - (Value::Bool(a), Value::Bool(b), false) => Value::Bool(a || b), - _ => return Err(invalid("boolean operands required")), - } - } - Not(v) => match v.evaluate(row)? { - Value::Bool(v) => Value::Bool(!v), - Value::Null => Value::Null, - _ => return Err(invalid("boolean operand required")), - }, - IsNull(v) => Value::Bool(matches!(v.evaluate(row)?, Value::Null)), - }) - } -} -fn ordered(dtype: &DataType) -> bool { - matches!( - dtype, - DataType::Int64 - | DataType::Float64 - | DataType::Utf8 - | DataType::Bool - | DataType::Timestamp - | DataType::Date - ) -} -fn numeric(op: &ArithmeticOpKind, a: Value, b: Value) -> Result { - use ArithmeticOpKind::*; - Ok(match (a, b) { - (Value::Null, _) | (_, Value::Null) => Value::Null, - (Value::Float64(a), Value::Float64(b)) => { - Value::Float64(crate::arithmetic::evaluate_float64_arithmetic(op, a, b)) - } - (Value::Int64(a), Value::Int64(b)) => Value::Int64( - match op { - Add => a.checked_add(b), - Sub => a.checked_sub(b), - Mul => a.checked_mul(b), - Div => a.checked_div(b), - Mod => a.checked_rem(b), - Pow => u32::try_from(b).ok().and_then(|b| a.checked_pow(b)), - Atan2 => None, - } - .ok_or_else(|| invalid("invalid integer arithmetic or overflow"))?, - ), - _ => return Err(invalid("arithmetic type mismatch")), - }) -} -#[derive(Clone, Debug)] -pub struct SortKey { - pub column: usize, - pub descending: bool, - pub nulls_first: bool, -} -#[derive(Clone, Debug)] -pub enum Reduction { - Count, - Sum(usize), - Avg(usize), - Min(usize), - Max(usize), -} -#[derive(Clone)] -enum Kind { - Source(Vec), - Union, - VectorToScalar { - column: usize, - }, - Project(Vec), - Filter(Expression), - Limit { - n: u64, - offset: u64, - groups: Vec, - }, - Sort { - keys: Vec, - groups: Vec, - }, - Aggregate { - groups: Vec, - measures: Vec, - }, - SemiJoin { - keys: Vec<(usize, usize)>, - }, - SummaryBuild { - family: SummaryFamilyType, - value: usize, - time: Option, - groups: Vec, - }, - SummaryMerge { - state: usize, - groups: Vec, - }, - Readout { - state: usize, - statistic: crate::Statistic, - parameters: std::collections::HashMap, - }, -} -/// A bound operation has a fully checked input/output contract before execution. -#[derive(Clone)] -pub struct Operator { - kind: Kind, - inputs: Vec, - output: Schema, -} -impl Operator { - pub fn source(output: Schema, batches: Vec) -> Result { - super::values::validate_schema(&output)?; - if batches.iter().any(|b| b.schema() != &output) { - return Err(invalid("source schema mismatch")); - } - Ok(Self { - kind: Kind::Source(batches), - inputs: vec![], - output, - }) - } - /// Union polls every input fairly, including branches sharing a producer. - pub fn union(input: Schema, arity: usize) -> Result { - if arity == 0 { - return Err(invalid("union needs at least one input")); - } - Ok(Self { - kind: Kind::Union, - inputs: vec![input.clone(); arity], - output: input, - }) - } - pub fn scalar(value: Value, dtype: DataType) -> Result { - let schema = schema(vec![result_field( - "value", - dtype, - matches!(value, Value::Null), - )]); - Self::source( - schema.clone(), - vec![Batch::try_new(schema, vec![vec![value]])?], - ) - } - /// PromQL scalar conversion: zero or multiple elements produce NaN. - pub fn vector_to_scalar(input: Schema, column: usize) -> Result { - if plain(&input, column)? != (&DataType::Float64, false) { - return Err(invalid("scalar conversion requires non-null Float64")); - } - Ok(Self { - kind: Kind::VectorToScalar { column }, - inputs: vec![input], - output: schema(vec![result_field("value", DataType::Float64, false)]), - }) - } - pub fn project(input: Schema, columns: Vec<(String, Expression)>) -> Result { - let fields = columns - .iter() - .map(|(name, e)| { - let (t, n) = e.dtype(&input)?; - Ok(result_field(name, t, n)) - }) - .collect::>()?; - Ok(Self { - kind: Kind::Project(columns.into_iter().map(|(_, e)| e).collect()), - inputs: vec![input], - output: schema(fields), - }) - } - pub fn filter(input: Schema, predicate: Expression) -> Result { - if predicate.dtype(&input)?.0 != DataType::Bool { - return Err(invalid("filter predicate must be boolean")); - } - Ok(Self { - kind: Kind::Filter(predicate), - inputs: vec![input.clone()], - output: input, - }) - } - pub fn limit(input: Schema, n: u64, offset: u64, groups: Vec) -> Result { - validate_groups(&input, &groups)?; - Ok(Self { - kind: Kind::Limit { n, offset, groups }, - inputs: vec![input.clone()], - output: input, - }) - } - pub fn sort(input: Schema, keys: Vec, groups: Vec) -> Result { - validate_groups(&input, &groups)?; - for key in &keys { - if !ordered(plain(&input, key.column)?.0) { - return Err(invalid("unsupported sort type")); - } - } - Ok(Self { - kind: Kind::Sort { keys, groups }, - inputs: vec![input.clone()], - output: input, - }) - } - pub fn aggregate( - input: Schema, - groups: Vec, - measures: Vec<(String, Reduction)>, - ) -> Result { - validate_groups(&input, &groups)?; - let mut fields = groups - .iter() - .map(|&i| input.fields[i].clone()) - .collect::>(); - for (name, reduction) in &measures { - let (t, n) = match reduction { - Reduction::Count => (DataType::Int64, false), - Reduction::Sum(i) | Reduction::Avg(i) => { - let (t, _) = plain(&input, *i)?; - if !matches!(t, DataType::Int64 | DataType::Float64) { - return Err(invalid("numeric aggregate input required")); - } - ( - if matches!(reduction, Reduction::Avg(_)) { - DataType::Float64 - } else { - t.clone() - }, - false, - ) - } - Reduction::Min(i) | Reduction::Max(i) => { - let (t, _) = plain(&input, *i)?; - if !ordered(t) { - return Err(invalid("ordered aggregate input required")); - } - (t.clone(), true) - } - }; - fields.push(result_field(name, t, n)); - } - Ok(Self { - kind: Kind::Aggregate { - groups, - measures: measures.into_iter().map(|(_, r)| r).collect(), - }, - inputs: vec![input], - output: schema(fields), - }) - } - pub fn semi_join( - left: Schema, - right: Schema, - keys: Vec<(usize, usize)>, - ) -> Result { - if keys.is_empty() { - return Err(invalid("semi-join needs matching keys")); - } - for &(l, r) in &keys { - if plain(&left, l)?.0 != plain(&right, r)?.0 { - return Err(invalid("join key types differ")); - } - } - Ok(Self { - kind: Kind::SemiJoin { keys }, - inputs: vec![left.clone(), right], - output: left, - }) - } - pub fn summary_build( - input: Schema, - family: SummaryFamilyType, - value: usize, - time: Option, - groups: Vec, - ) -> Result { - super::values::validate_family(&family)?; - validate_groups(&input, &groups)?; - if plain(&input, value)? != (&DataType::Float64, false) { - return Err(invalid("summary numeric update requires non-null Float64")); - } - if let Some(time) = time { - if plain(&input, time)? != (&DataType::Timestamp, false) { - return Err(invalid("summary time column must be a timestamp")); - } - } - if time.is_none() - && matches!( - family, - SummaryFamilyType::ExactAggregate( - planner_types::post_asap::ExactKind::Rate - | planner_types::post_asap::ExactKind::Increase, - _ - ) - ) - { - return Err(invalid("counter summary requires a timestamp column")); - } - crate::capability::validate_summary_kernel( - &family, - &SummaryUpdate::column(ColumnRef::SampleValue), - &Default::default(), - ) - .map_err(Error::Invalid)?; - let mut fields = groups - .iter() - .map(|&i| input.fields[i].clone()) - .collect::>(); - fields.push(SummaryField { - name: "state".into(), - dtype: family.clone(), - nullable: false, - }); - Ok(Self { - kind: Kind::SummaryBuild { - family, - value, - time, - groups, - }, - inputs: vec![input], - output: schema(fields), - }) - } - pub fn summary_merge(input: Schema, state: usize, groups: Vec) -> Result { - validate_groups(&input, &groups)?; - super::values::validate_family(&field(&input, state)?.dtype)?; - if matches!(field(&input, state)?.dtype, SummaryFamilyType::Plain(_)) { - return Err(invalid("summary state required")); - } - let mut fields = groups - .iter() - .map(|&i| input.fields[i].clone()) - .collect::>(); - fields.push(input.fields[state].clone()); - Ok(Self { - kind: Kind::SummaryMerge { state, groups }, - inputs: vec![input], - output: schema(fields), - }) - } - pub fn readout( - input: Schema, - state: usize, - statistic: crate::Statistic, - parameters: std::collections::HashMap, - ) -> Result { - super::values::validate_family(&field(&input, state)?.dtype)?; - if matches!(field(&input, state)?.dtype, SummaryFamilyType::Plain(_)) { - return Err(invalid("summary state required")); - } - validate_readout(&field(&input, state)?.dtype, statistic, ¶meters)?; - let mut fields = input.fields.clone(); - let result_type = if matches!( - fields[state].dtype, - SummaryFamilyType::ExactAggregate(planner_types::post_asap::ExactKind::Count, _) - ) { - DataType::Int64 - } else { - DataType::Float64 - }; - fields[state] = result_field("value", result_type, false); - Ok(Self { - kind: Kind::Readout { - state, - statistic, - parameters, - }, - inputs: vec![input], - output: schema(fields), - }) - } - pub(crate) fn with_output_schema(mut self, output: Schema) -> Result { - if self.output.fields.len() != output.fields.len() - || self - .output - .fields - .iter() - .zip(&output.fields) - .any(|(actual, declared)| { - actual.dtype != declared.dtype || (actual.nullable && !declared.nullable) - }) - { - return Err(invalid("native output type differs from Planner output")); - } - if output.time_index.is_some_and(|i| { - i >= output.fields.len() - || output.fields[i].dtype != SummaryFamilyType::Plain(DataType::Timestamp) - }) { - return Err(invalid("invalid output time column")); - } - self.output = output; - Ok(self) - } - pub fn schema(&self) -> Schema { - self.output.clone() - } -} -fn validate_groups(input: &Schema, groups: &[usize]) -> Result<(), Error> { - for &i in groups { - plain(input, i)?; - } - if groups - .iter() - .collect::>() - .len() - != groups.len() - { - return Err(invalid("duplicate group columns")); - } - Ok(()) -} -async fn collect_rows( - mut input: Input<'_, Batch>, - context: &RunContext, -) -> Result<(Vec>, Vec), Error> { - let mut rows = Vec::new(); - let mut reservations = Vec::new(); - while let Some(batch) = input.next().await { - let batch = batch?; - reservations.push(context.reserve(batch.bytes())?); - rows.extend(batch.rows().iter().cloned()); - } - Ok((rows, reservations)) -} -impl PhysicalOperator for Operator { - fn name(&self) -> &str { - match self.kind { - Kind::Source(_) => "Source", - Kind::Union => "Union", - Kind::VectorToScalar { .. } => "VectorToScalar", - Kind::Project(_) => "Project", - Kind::Filter(_) => "Filter", - Kind::Limit { .. } => "Limit", - Kind::Sort { .. } => "Sort", - Kind::Aggregate { .. } => "Aggregate", - Kind::SemiJoin { .. } => "SemiJoin", - Kind::SummaryBuild { .. } => "SummaryAgg", - Kind::SummaryMerge { .. } => "SummaryMerge", - Kind::Readout { .. } => "SummaryReadout", - } - } - fn input_schemas(&self) -> Vec { - self.inputs.clone() - } - fn output_schema(&self) -> Schema { - self.output.clone() - } - fn output_bytes(&self, value: &Batch) -> usize { - value.bytes() - } - fn start<'a>( - &'a self, - mut inputs: Vec>, - context: RunContext, - ) -> Result, Error> { - let output = self.output.clone(); - if let Kind::Source(batches) = &self.kind { - return Ok(futures::stream::iter(batches.iter().cloned().map(Ok)).boxed_local()); - } - if matches!(self.kind, Kind::Union) { - return Ok(futures::stream::select_all(inputs) - .map(|batch| batch.map(|batch| batch.value().clone())) - .boxed_local()); - } - if let Kind::SemiJoin { keys } = &self.kind { - let right = inputs.pop().ok_or_else(|| invalid("right input missing"))?; - let left = inputs.pop().ok_or_else(|| invalid("left input missing"))?; - return Ok(futures::stream::once(async move { - // Poll both branches together: either may depend on a common producer. - let ((left, _left_memory), (right, _right_memory)) = futures::try_join!( - collect_rows(left, &context), - collect_rows(right, &context) - )?; - let right_cols = keys.iter().map(|(_, r)| *r).collect::>(); - let left_cols = keys.iter().map(|(l, _)| *l).collect::>(); - let members = right - .iter() - .filter(|row| right_cols.iter().all(|&i| !matches!(row[i], Value::Null))) - .map(|r| group_key(r, &right_cols)) - .collect::, _>>()?; - let rows = left - .into_iter() - .filter_map(|r| match group_key(&r, &left_cols) { - Ok(k) - if left_cols.iter().all(|&i| !matches!(r[i], Value::Null)) - && members.contains(&k) => - { - Some(Ok(r)) - } - Ok(_) => None, - Err(e) => Some(Err(e)), - }) - .collect::, _>>()?; - Batch::try_new(output, rows) - }) - .boxed_local()); - } - let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; - match &self.kind { - Kind::VectorToScalar { column } => Ok(futures::stream::once(async move { - let mut input = input; - let mut value = f64::NAN; - let mut count = 0usize; - while let Some(batch) = input.next().await { - for row in batch?.rows() { - count = count.saturating_add(1); - if let Value::Float64(v) = row[*column] { - value = v; - } - } - } - Batch::try_new( - output, - vec![vec![Value::Float64(if count == 1 { - value - } else { - f64::NAN - })]], - ) - }) - .boxed_local()), - Kind::Project(expressions) => Ok(input - .map(move |batch| { - let batch = batch?; - let rows = batch - .rows() - .iter() - .map(|r| { - expressions - .iter() - .map(|e| e.evaluate(r)) - .collect::, _>>() - }) - .collect::, _>>()?; - Batch::try_new(output.clone(), rows) - }) - .boxed_local()), - Kind::Filter(predicate) => Ok(input - .map(move |batch| { - let batch = batch?; - let mut rows = Vec::new(); - for row in batch.rows() { - if matches!(predicate.evaluate(row)?, Value::Bool(true)) { - rows.push(row.clone()); - } - } - Batch::try_new(output.clone(), rows) - }) - .boxed_local()), - Kind::Limit { n, offset, groups } => { - let counts = BTreeMap::>, u64>::new(); - Ok(futures::stream::try_unfold( - (input, counts, Vec::::new(), false), - move |(mut input, mut counts, mut memory, done)| { - let output = output.clone(); - let context = context.clone(); - async move { - if done || *n == 0 { - return Ok(None); - } - let Some(batch) = input.next().await else { - return Ok(None); - }; - let batch = batch?; - let mut rows = Vec::new(); - for row in batch.rows() { - let key = group_key(row, groups)?; - if !counts.contains_key(&key) { - memory.push( - context.reserve( - key.iter() - .map(|part| { - part.len() + std::mem::size_of::>() - }) - .sum::() - + 64, - )?, - ); - } - let count = counts.entry(key).or_default(); - if *count >= *offset && count.saturating_sub(*offset) < *n { - rows.push(row.clone()); - } - *count = count.saturating_add(1); - } - let done = groups.is_empty() - && counts - .get(&vec![]) - .is_some_and(|count| count.saturating_sub(*offset) >= *n); - Ok(Some(( - Batch::try_new(output, rows)?, - (input, counts, memory, done), - ))) - } - }, - ) - .boxed_local()) - } - Kind::SummaryBuild { - family, - value, - time, - groups, - } => Ok(futures::stream::once(async move { - Batch::try_new( - output, - build_summary(input, family, *value, *time, groups, &context).await?, - ) - }) - .boxed_local()), - Kind::Readout { - state, - statistic, - parameters, - } => Ok(input - .map(move |batch| { - let batch = batch?; - let mut rows = batch.rows().to_vec(); - for row in &mut rows { - let Value::Summary { state: summary, .. } = &row[*state] else { - return Err(invalid("summary value required")); - }; - row[*state] = if output.fields[*state].dtype - == SummaryFamilyType::Plain(DataType::Int64) - { - let count = summary.aux_stats().count.ok_or_else(|| { - Error::Operator("exact count state lacks an integer count".into()) - })?; - Value::Int64( - i64::try_from(count).map_err(|_| { - Error::Operator("exact count exceeds Int64".into()) - })?, - ) - } else { - Value::Float64( - summary - .query_statistic(*statistic, &None, parameters) - .map_err(|e| Error::Operator(e.to_string()))?, - ) - }; - } - Batch::try_new(output.clone(), rows) - }) - .boxed_local()), - _ => Ok(futures::stream::once(async move { - let (rows, _memory) = collect_rows(input, &context).await?; - let result = match &self.kind { - Kind::Sort { keys, groups } => { - let mut grouped = BTreeMap::>, Vec>>::new(); - for row in rows { - grouped - .entry(group_key(&row, groups)?) - .or_default() - .push(row); - } - let mut result = Vec::new(); - for mut rows in grouped.into_values() { - rows.sort_by(|a, b| compare_rows(a, b, keys)); - result.extend(rows); - } - result - } - Kind::Aggregate { groups, measures } => { - reduce(rows, groups, measures, &self.inputs[0])? - } - Kind::SummaryMerge { state, groups } => merge_summary(rows, *state, groups)?, - _ => return Err(invalid("unexpected blocking operation")), - }; - Batch::try_new(output, result) - }) - .boxed_local()), - } - } -} -fn compare_rows(a: &[Value], b: &[Value], keys: &[SortKey]) -> std::cmp::Ordering { - use std::cmp::Ordering::*; - for key in keys { - let (a, b) = (&a[key.column], &b[key.column]); - let order = match (a, b) { - (Value::Null, Value::Null) => Equal, - (Value::Null, _) => { - if key.nulls_first { - Less - } else { - Greater - } - } - (_, Value::Null) => { - if key.nulls_first { - Greater - } else { - Less - } - } - (Value::Float64(a), Value::Float64(b)) if a.is_nan() || b.is_nan() => { - match (a.is_nan(), b.is_nan()) { - (true, true) => Equal, - (true, false) => Greater, - _ => Less, - } - } - _ => { - let order = a.compare(b).expect("bound ordered types"); - if key.descending { - order.reverse() - } else { - order - } - } - }; - if order != Equal { - return order; - } - } - Equal -} -fn reduce( - rows: Vec>, - groups: &[usize], - measures: &[Reduction], - input: &Schema, -) -> Result>, Error> { - let mut grouped = BTreeMap::>, Vec>>::new(); - if rows.is_empty() && groups.is_empty() { - grouped.insert(vec![], vec![]); - } - for row in rows { - grouped - .entry(group_key(&row, groups)?) - .or_default() - .push(row); - } - grouped - .into_values() - .map(|rows| { - let mut result = groups - .iter() - .map(|&i| rows[0][i].clone()) - .collect::>(); - for measure in measures { - result.push(reduce_one(&rows, measure, input)?); - } - Ok(result) - }) - .collect() -} -fn reduce_one(rows: &[Vec], measure: &Reduction, input: &Schema) -> Result { - let column = match measure { - Reduction::Count => { - return Ok(Value::Int64( - i64::try_from(rows.len()).map_err(|_| invalid("count overflow"))?, - )) - } - Reduction::Sum(i) | Reduction::Avg(i) | Reduction::Min(i) | Reduction::Max(i) => *i, - }; - let values = rows - .iter() - .map(|r| &r[column]) - .filter(|v| !matches!(v, Value::Null)) - .collect::>(); - if matches!(measure, Reduction::Min(_) | Reduction::Max(_)) { - if plain(input, column)?.0 == &DataType::Float64 { - // Match exact-state kernels: ignore NaN when a numeric value exists. - let mut best: Option = None; - for value in values { - let Value::Float64(value) = value else { - return Err(invalid("floating aggregate value required")); - }; - best = Some(best.map_or(*value, |old| { - if matches!(measure, Reduction::Min(_)) { - old.min(*value) - } else { - old.max(*value) - } - })); - } - return Ok(best.map(Value::Float64).unwrap_or(Value::Null)); - } - let mut best: Option<&Value> = None; - for value in values { - if best - .map(|b| value.compare(b)) - .transpose()? - .is_none_or(|order| { - if matches!(measure, Reduction::Min(_)) { - order.is_lt() - } else { - order.is_gt() - } - }) - { - best = Some(value); - } - } - return Ok(best.cloned().unwrap_or(Value::Null)); - } - let count = values.len(); - let dtype = plain(input, column)?.0; - if dtype == &DataType::Int64 { - let sum = values.into_iter().try_fold(0i128, |sum, v| { - let Value::Int64(v) = v else { - return Err(invalid("integer aggregate value required")); - }; - sum.checked_add(i128::from(*v)) - .ok_or_else(|| invalid("integer aggregate overflow")) - })?; - return if matches!(measure, Reduction::Avg(_)) { - Ok(Value::Float64(sum as f64 / count as f64)) - } else { - Ok(Value::Int64( - i64::try_from(sum).map_err(|_| invalid("integer sum overflow"))?, - )) - }; - } - let sum = values - .into_iter() - .map(|v| { - if let Value::Float64(v) = v { - *v - } else { - unreachable!() - } - }) - .sum::(); - Ok(Value::Float64(if matches!(measure, Reduction::Avg(_)) { - sum / count as f64 - } else { - sum - })) -} -async fn build_summary( - mut input: Input<'_, Batch>, - family: &SummaryFamilyType, - value: usize, - time: Option, - groups: &[usize], - context: &RunContext, -) -> Result>, Error> { - type State = ( - Vec, - Box, - Reservation, - usize, - Option, - ); - let create = |labels: Vec, key_bytes: usize| -> Result { - let updater = crate::factory::create_planner_accumulator( - family, - &SummaryUpdate::column(ColumnRef::SampleValue), - &Default::default(), - ) - .map_err(Error::Operator)?; - let overhead = labels.iter().map(Value::bytes).sum::() + key_bytes + 64; - let memory = context.reserve(updater.memory_usage_bytes() + overhead)?; - Ok((labels, updater, memory, overhead, None)) - }; - let mut states = BTreeMap::>, State>::new(); - if groups.is_empty() { - states.insert(vec![], create(vec![], 0)?); - } - let ordered_time = matches!( - family, - SummaryFamilyType::ExactAggregate( - planner_types::post_asap::ExactKind::Rate - | planner_types::post_asap::ExactKind::Increase, - _ - ) - ); - while let Some(batch) = input.next().await { - let batch = batch?; - for row in batch.rows() { - let key = group_key(row, groups)?; - if !states.contains_key(&key) { - let labels = groups.iter().map(|&i| row[i].clone()).collect(); - let state = create( - labels, - key.iter() - .map(|v| v.len() + std::mem::size_of::>()) - .sum(), - )?; - states.insert(key.clone(), state); - } - let (_, updater, memory, overhead, previous) = - states.get_mut(&key).expect("inserted group"); - let Value::Float64(value) = row[value] else { - return Err(invalid("summary update type")); - }; - let timestamp = if let Some(time) = time { - let Value::Timestamp(time) = row[time] else { - return Err(invalid("summary time type")); - }; - time - } else { - 0 - }; - if ordered_time && previous.is_some_and(|prior| timestamp <= prior) { - return Err(Error::Operator( - "counter samples must have strictly increasing timestamps within each group" - .into(), - )); - } - updater - .validate_single_input(value) - .map_err(Error::Operator)?; - updater.update_single(value, timestamp); - *previous = Some(timestamp); - memory.resize(updater.memory_usage_bytes() + *overhead)?; - } - } - Ok(states - .into_values() - .map(|(mut labels, updater, _memory, _, _)| { - labels.push(Value::Summary { - family: family.clone(), - state: Arc::from(updater.into_accumulator()), - }); - labels - }) - .collect()) -} - -fn merge_summary( - rows: Vec>, - state_column: usize, - groups: &[usize], -) -> Result>, Error> { - type GroupState = (Vec, SummaryFamilyType, Arc); - let mut states: BTreeMap>, GroupState> = BTreeMap::new(); - for row in rows { - let Value::Summary { family, state } = &row[state_column] else { - return Err(invalid("summary state required")); - }; - let key = group_key(&row, groups)?; - if let Some((_, expected, existing)) = states.get_mut(&key) { - if expected != family { - return Err(invalid("incompatible summary family")); - } - *existing = Arc::from( - existing - .merge_with(state.as_ref()) - .map_err(|e| Error::Operator(e.to_string()))?, - ); - } else { - states.insert( - key, - ( - groups.iter().map(|&i| row[i].clone()).collect(), - family.clone(), - state.clone(), - ), - ); - } - } - Ok(states - .into_values() - .map(|(mut keys, family, state)| { - keys.push(Value::Summary { family, state }); - keys - }) - .collect()) -} - -fn validate_readout( - family: &SummaryFamilyType, - statistic: crate::Statistic, - parameters: &std::collections::HashMap, -) -> Result<(), Error> { - use crate::Statistic as S; - use planner_types::post_asap::{ExactKind as E, SketchAlgorithm as A}; - let supported = match family { - SummaryFamilyType::ExactAggregate(kind, _) => matches!( - (kind, statistic), - (E::Sum, S::Sum) - | (E::Count, S::Count) - | (E::Min, S::Min) - | (E::Max, S::Max) - | (E::Rate, S::Rate) - | (E::Increase, S::Increase) - ), - SummaryFamilyType::Sketch(kind, _) => match kind.algorithm() { - A::Kll => statistic == S::Quantile, - A::DDSketch => matches!(statistic, S::Quantile | S::Count), - A::Hll => matches!(statistic, S::Cardinality | S::Count), - _ => false, - }, - _ => false, - }; - if !supported { - return Err(invalid( - "readout is not implemented for this summary family", - )); - } - if statistic == S::Quantile - && !parameters - .get("quantile") - .and_then(|s| s.parse::().ok()) - .is_some_and(|q| (0.0..=1.0).contains(&q)) - { - return Err(invalid("quantile readout requires quantile in [0,1]")); - } - Ok(()) -} diff --git a/crates/asap-physical-operators/src/dag/planner.rs b/crates/asap-physical-operators/src/dag/planner.rs deleted file mode 100644 index 1952f1790..000000000 --- a/crates/asap-physical-operators/src/dag/planner.rs +++ /dev/null @@ -1,479 +0,0 @@ -//! Bind a post-ASAP DAG to native operators. Sources are explicit execution -//! frontiers supplied by the deployment; unsupported computation is an error. -use super::{ - operators::{Expression, Operator, Reduction, SortKey}, - values::{Batch, Schema, Value}, - Error, NodeId, PhysicalDag, PhysicalOperator, -}; -use planner_types::{ - post_asap::{ - ExactOperation, ExecutableDag, ExecutableDagNode, ExecutableOperatorPayload as Payload, - SketchQuery, SummaryFamilyType, SummaryInputExpr, ValueOperation, - }, - pre_asap::{ - AggIntent, ColumnRef, CompareOpKind, DataType, GroupKeys, QueryExpr, - Reduction as PlannerReduction, ScalarValue, - }, -}; -use std::{ - collections::{BTreeMap, BTreeSet}, - sync::Arc, -}; -fn invalid(message: impl Into) -> Error { - Error::Invalid(message.into()) -} - -/// Source nodes cut the DAG at an installed storage/ingestion frontier. The -/// binding must have exactly the declared schema and no upstream dependencies. -/// A deployment must authorize these frontiers before calling this function. -pub type Source<'a> = Box + 'a>; - -pub fn bind<'a>( - dag: &ExecutableDag, - mut sources: BTreeMap>, - roots: &[NodeId], -) -> Result, Error> { - preflight_depth(dag)?; - dag.validate().map_err(|e| invalid(e.to_string()))?; - let nodes = dag - .nodes - .iter() - .map(|node| (u64::from(node.id.0), node)) - .collect::>(); - let mut dependencies = BTreeMap::>::new(); - for edge in &dag.edges { - dependencies - .entry(u64::from(edge.consumer.0)) - .or_default() - .push(u64::from(edge.producer.0)); - } - if sources.keys().any(|id| !nodes.contains_key(id)) { - return Err(invalid("source binding names an unknown node")); - } - let mut ordered = Vec::new(); - let mut seen = BTreeSet::new(); - let mut pending = roots.iter().map(|&id| (id, false)).collect::>(); - while let Some((id, expanded)) = pending.pop() { - if expanded { - ordered.push(id); - continue; - } - if !seen.insert(id) { - continue; - } - if !nodes.contains_key(&id) { - return Err(invalid(format!("missing root {id}"))); - } - pending.push((id, true)); - if !sources.contains_key(&id) { - for &input in dependencies.get(&id).into_iter().flatten() { - pending.push((input, false)); - } - } - } - let mut graph = PhysicalDag::default(); - let mut auxiliary = u64::MAX; - for id in ordered { - let node = nodes[&id]; - let output = Arc::new(node.output_schema.clone()); - super::values::validate_schema(&output)?; - let (operator, inputs) = if let Some(source) = sources.remove(&id) { - if !source.input_schemas().is_empty() || source.output_schema() != output { - return Err(invalid("frontier is not a source with the declared schema")); - } - ( - Box::new(CheckedSource { source, output }) as Source<'a>, - vec![], - ) - } else { - let mut inputs = dependencies.get(&id).cloned().unwrap_or_default(); - let mut schemas = inputs - .iter() - .map(|id| Arc::new(nodes[id].output_schema.clone())) - .collect::>(); - if matches!(node.payload, Payload::SummaryMerge { .. }) && inputs.len() > 1 { - if schemas.iter().any(|s| s != &schemas[0]) { - return Err(invalid("summary merge inputs have different schemas")); - } - graph.add( - auxiliary, - inputs, - Operator::union(schemas[0].clone(), schemas.len())?, - )?; - inputs = vec![auxiliary]; - auxiliary -= 1; - schemas.truncate(1); - } - let operator = bind_operation(node, &schemas) - .map_err(|error| invalid(format!("node {id}: {error}")))? - .with_output_schema(output)?; - (Box::new(operator) as Source<'a>, inputs) - }; - graph.add_boxed(id, inputs, operator)?; - } - graph.validate(roots)?; - Ok(graph) -} - -fn bind_operation(node: &ExecutableDagNode, inputs: &[Schema]) -> Result { - let [input] = inputs else { - return Err(invalid( - "native Planner binding currently requires a unary operation or an explicit source", - )); - }; - match &node.payload { - Payload::Value { operation, .. } => match operation { - ValueOperation::Project { cols, .. } => Operator::project( - input.clone(), - cols.iter() - .enumerate() - .map(|(i, col)| { - Ok(( - node.output_schema - .fields - .get(i) - .ok_or_else(|| invalid("projection width mismatch"))? - .name - .clone(), - expression(&col.expr)?, - )) - }) - .collect::>()?, - ), - ValueOperation::Filter { pred } => { - Operator::filter(input.clone(), expression(&pred.0)?) - } - ValueOperation::Sort { keys, partition_by } => Operator::sort( - input.clone(), - keys.iter() - .map(|key| { - let QueryExpr::Column(column) = key.expr else { - return Err(invalid( - "sort expression must be projected before sorting", - )); - }; - Ok(SortKey { - column, - descending: !key.ascending, - nulls_first: key.nulls_first, - }) - }) - .collect::>()?, - groups(input, partition_by)?, - ), - ValueOperation::Limit { n, offset } => { - Operator::limit(input.clone(), *n as u64, *offset as u64, vec![]) - } - ValueOperation::Exact(ExactOperation::Aggregate { - reduction, - measures, - output_names, - having: None, - }) => { - if measures.len() != output_names.len() { - return Err(invalid("aggregate output names differ from measures")); - } - let PlannerReduction::Reduce(keys) = reduction else { - return Err(invalid( - "per-entity aggregate requires an explicit entity binding", - )); - }; - let measures = measures - .iter() - .zip(output_names) - .map(|(m, name)| { - let column = |col: Option| { - col.map(Ok) - .unwrap_or_else(|| named_column(input, &ColumnRef::SampleValue)) - }; - let m = match m { - AggIntent::Count { .. } => Reduction::Count, - AggIntent::Sum { col } => Reduction::Sum(column(*col)?), - AggIntent::Avg { col } => Reduction::Avg(column(*col)?), - AggIntent::Min { col } => Reduction::Min(column(*col)?), - AggIntent::Max { col } => Reduction::Max(column(*col)?), - _ => { - return Err(invalid( - "aggregate intent has no native implementation", - )) - } - }; - Ok((name.clone(), m)) - }) - .collect::>()?; - Operator::aggregate(input.clone(), groups(input, keys)?, measures) - } - ValueOperation::FinalizeExactAccumulator => { - let state = summary_column(input)?; - use crate::Statistic as S; - use planner_types::post_asap::ExactKind as E; - let statistic = match &input.fields[state].dtype { - SummaryFamilyType::ExactAggregate(kind, _) => match kind { - E::Sum => S::Sum, - E::Count => S::Count, - E::Min => S::Min, - E::Max => S::Max, - E::Rate => S::Rate, - E::Increase => S::Increase, - _ => return Err(invalid("exact family readout is unsupported")), - }, - _ => return Err(invalid("exact finalization requires exact state")), - }; - Operator::readout(input.clone(), state, statistic, Default::default()) - } - _ => Err(invalid("value operation has no native implementation")), - }, - Payload::SummaryAgg { - family, - input: update, - reduction, - grouping, - } => { - if update.item.is_some() { - return Err(invalid("keyed summary update binding is not implemented")); - } - crate::capability::validate_summary_kernel(family, update, grouping) - .map_err(Error::Invalid)?; - let SummaryInputExpr::Column(column) = &update.weight else { - return Err(invalid( - "summary update expression must be projected to a column", - )); - }; - let PlannerReduction::Reduce(keys) = reduction else { - return Err(invalid( - "summary construction requires explicit grouping columns", - )); - }; - Operator::summary_build( - input.clone(), - family.clone(), - named_column(input, column)?, - input.time_index, - groups(input, keys)?, - ) - } - Payload::SummaryMerge { .. } => { - let state = summary_column(input)?; - Operator::summary_merge( - input.clone(), - state, - (0..input.fields.len()) - .filter(|&i| i != state && Some(i) != input.time_index) - .collect(), - ) - } - Payload::SummaryEstimate { query } => { - let mut params = std::collections::HashMap::new(); - let statistic = match query { - SketchQuery::Quantile { q } => { - params.insert("quantile".into(), q.to_string()); - crate::Statistic::Quantile - } - SketchQuery::Cardinality => crate::Statistic::Cardinality, - SketchQuery::PointCount { value: None, .. } => crate::Statistic::Count, - _ => return Err(invalid("summary readout is not implemented")), - }; - Operator::readout(input.clone(), summary_column(input)?, statistic, params) - } - _ => Err(invalid( - "physical operation has no native binding; no fallback is installed", - )), - } -} -fn summary_column(input: &Schema) -> Result { - let columns = input - .fields - .iter() - .enumerate() - .filter(|(_, f)| !matches!(f.dtype, SummaryFamilyType::Plain(_))) - .map(|(i, _)| i) - .collect::>(); - match columns.as_slice() { - [column] => Ok(*column), - _ => Err(invalid("one summary state column required")), - } -} -fn named_column(input: &Schema, column: &ColumnRef) -> Result { - let name = match column { - ColumnRef::Named(name) => name.as_str(), - ColumnRef::SampleValue => "value", - _ => { - return Err(invalid( - "summary update requires an unambiguous bound column", - )) - } - }; - let matches = input - .fields - .iter() - .enumerate() - .filter(|(_, field)| field.name == name) - .map(|(i, _)| i) - .collect::>(); - match matches.as_slice() { - [column] => Ok(*column), - _ => Err(invalid("summary update column missing or ambiguous")), - } -} -fn groups(input: &Schema, groups: &GroupKeys) -> Result, Error> { - if groups.is_without() { - return Err(invalid("grouping without requires resolved label columns")); - } - if groups.keys().iter().any(|&i| i >= input.fields.len()) { - return Err(invalid("grouping column out of range")); - } - Ok(groups.keys().to_vec()) -} -fn expression(expr: &QueryExpr) -> Result { - let bind = |e: &QueryExpr| expression(e).map(Box::new); - Ok(match expr { - QueryExpr::Column(i) => Expression::Column(*i), - QueryExpr::Literal(value) => { - let (value, dtype) = match value { - ScalarValue::Int64(v) => (Value::Int64(*v), DataType::Int64), - ScalarValue::Float64(v) => (Value::Float64(*v), DataType::Float64), - ScalarValue::Utf8(v) => (Value::Utf8(v.as_str().into()), DataType::Utf8), - ScalarValue::Boolean(v) => (Value::Bool(*v), DataType::Bool), - ScalarValue::Null => (Value::Null, DataType::Null), - ScalarValue::Interval { - months, - days, - nanos, - } => ( - Value::Interval { - months: *months, - days: *days, - nanos: *nanos, - }, - DataType::Interval, - ), - }; - Expression::Literal { value, dtype } - } - QueryExpr::Arithmetic { op, left, right } => Expression::Arithmetic { - op: op.clone(), - left: bind(left)?, - right: bind(right)?, - }, - QueryExpr::Compare { - left, - op: CompareOpKind::Eq, - right, - } => Expression::Equal(bind(left)?, bind(right)?), - QueryExpr::Compare { - left, - op: CompareOpKind::Lt, - right, - } => Expression::Less(bind(left)?, bind(right)?), - QueryExpr::Not(v) => Expression::Not(bind(v)?), - QueryExpr::IsNull(v) => Expression::IsNull(bind(v)?), - QueryExpr::IsNotNull(v) => Expression::Not(Box::new(Expression::IsNull(bind(v)?))), - QueryExpr::BoolAnd(items) | QueryExpr::BoolOr(items) => { - let and = matches!(expr, QueryExpr::BoolAnd(_)); - let mut result = Expression::Literal { - value: Value::Bool(and), - dtype: DataType::Bool, - }; - for item in items { - result = if and { - Expression::And(Box::new(result), bind(item)?) - } else { - Expression::Or(Box::new(result), bind(item)?) - }; - } - result - } - _ => return Err(invalid("expression has no native implementation")), - }) -} - -// Source adapters may perform I/O, but their actual batches must honor the -// schema accepted by the binder before a downstream expression sees a row. -struct CheckedSource<'a> { - source: Source<'a>, - output: Schema, -} -impl PhysicalOperator for CheckedSource<'_> { - fn name(&self) -> &str { - self.source.name() - } - fn input_schemas(&self) -> Vec { - vec![] - } - fn output_schema(&self) -> Schema { - self.output.clone() - } - fn output_bytes(&self, batch: &Batch) -> usize { - self.source.output_bytes(batch) - } - fn start<'a>( - &'a self, - inputs: Vec>, - context: super::RunContext, - ) -> Result, Error> { - use futures::StreamExt; - Ok(self - .source - .start(inputs, context)? - .map(|batch| { - let batch = batch?; - if batch.schema() != &self.output { - return Err(invalid("source batch differs from its bound schema")); - } - Ok(batch) - }) - .boxed_local()) - } -} - -// Bound recursion before invoking the upstream recursive provenance validator. -fn preflight_depth(dag: &ExecutableDag) -> Result<(), Error> { - let mut remaining = dag - .nodes - .iter() - .map(|node| (node.id, 0usize)) - .collect::>(); - if remaining.len() != dag.nodes.len() { - return Err(invalid("duplicate Planner node")); - } - let mut consumers = BTreeMap::<_, Vec<_>>::new(); - for edge in &dag.edges { - if !remaining.contains_key(&edge.producer) { - return Err(invalid("missing Planner edge producer")); - } - *remaining - .get_mut(&edge.consumer) - .ok_or_else(|| invalid("missing Planner edge consumer"))? += 1; - consumers - .entry(edge.producer) - .or_default() - .push(edge.consumer); - } - let mut ready = remaining - .iter() - .filter(|(_, n)| **n == 0) - .map(|(id, _)| *id) - .collect::>(); - let mut depths = BTreeMap::new(); - let mut visited = 0; - while let Some(id) = ready.pop_front() { - visited += 1; - let depth = *depths.get(&id).unwrap_or(&1usize); - if depth > 128 { - return Err(invalid("DAG exceeds the supported execution depth of 128")); - } - for &consumer in consumers.get(&id).into_iter().flatten() { - let next = depths.entry(consumer).or_insert(1); - *next = (*next).max(depth + 1); - let count = remaining.get_mut(&consumer).expect("validated endpoint"); - *count -= 1; - if *count == 0 { - ready.push_back(consumer); - } - } - } - if visited != dag.nodes.len() { - return Err(invalid("Planner DAG contains a cycle")); - } - Ok(()) -} diff --git a/crates/asap-physical-operators/src/dag/tests.rs b/crates/asap-physical-operators/src/dag/tests.rs deleted file mode 100644 index e051fa313..000000000 --- a/crates/asap-physical-operators/src/dag/tests.rs +++ /dev/null @@ -1,260 +0,0 @@ -use super::*; -use futures::{executor::block_on, stream, StreamExt}; - -struct Source { - starts: Rc>, - polls: Rc>, - fail: bool, - end: u64, -} -impl PhysicalOperator for Source { - fn name(&self) -> &str { - "CountingSource" - } - fn input_schemas(&self) -> Vec<()> { - vec![] - } - fn output_schema(&self) {} - fn output_bytes(&self, _: &u64) -> usize { - 8 - } - fn start<'a>( - &'a self, - _: Vec>, - _: RunContext, - ) -> Result, Error> { - self.starts.set(self.starts.get() + 1); - Ok(stream::iter(0..self.end) - .map(move |n| { - self.polls.set(self.polls.get() + 1); - if self.fail && n == 1 { - Err(Error::Operator("source failure".into())) - } else { - Ok(n) - } - }) - .boxed_local()) - } -} -struct Identity; -impl PhysicalOperator for Identity { - fn name(&self) -> &str { - "Identity" - } - fn input_schemas(&self) -> Vec<()> { - vec![()] - } - fn output_schema(&self) {} - fn output_bytes(&self, _: &u64) -> usize { - 8 - } - fn start<'a>( - &'a self, - mut inputs: Vec>, - _: RunContext, - ) -> Result, Error> { - Ok(inputs - .remove(0) - .map(|value| value.map(|v| *v)) - .boxed_local()) - } -} -fn context() -> RunContext { - RunContext::new( - Scope::Query { - evaluation_time_ms: 100, - revision: 1, - }, - Limits { - max_buffered_batches: 1, - max_bytes: 1024, - }, - ) - .unwrap() -} -fn source(fail: bool) -> (Source, Rc>, Rc>) { - let starts = Rc::new(Cell::new(0)); - let polls = Rc::new(Cell::new(0)); - ( - Source { - starts: starts.clone(), - polls: polls.clone(), - fail, - end: 4, - }, - starts, - polls, - ) -} - -// A shared producer runs once, and the slow reader bounds producer progress. -#[test] -fn shared_source_backpressure_and_reader_drop() { - let (source, starts, polls) = source(false); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], source).unwrap(); - let context = context(); - let mut readers = dag.execute(&[0, 0], context.clone()).unwrap(); - let mut slow = readers.pop().unwrap(); - let mut fast = readers.pop().unwrap(); - assert_eq!(starts.get(), 1); - let first = block_on(fast.next()).unwrap().unwrap(); - assert_eq!(*first, 0); - let mut cx = Context::from_waker(futures::task::noop_waker_ref()); - assert!(Pin::new(&mut fast).poll_next(&mut cx).is_pending()); - assert_eq!(polls.get(), 1); - let same = block_on(slow.next()).unwrap().unwrap(); - assert!(Arc::ptr_eq(&first.value, &same.value)); - drop(same); - drop(first); - assert_eq!(context.retained_bytes(), 0); - assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 1); - drop(slow); - assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 2); - assert_eq!(*block_on(fast.next()).unwrap().unwrap(), 3); - assert!(block_on(fast.next()).is_none()); - assert_eq!(polls.get(), 4); - drop(fast); - assert_eq!(context.retained_bytes(), 0); -} - -// Independent branches consume a common node concurrently without duplicate work. -#[test] -fn diamond_and_run_isolation() { - let (source, starts, polls) = source(false); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], source).unwrap(); - dag.add(1, vec![0], Identity).unwrap(); - dag.add(2, vec![0], Identity).unwrap(); - for _ in 0..2 { - let mut outputs = dag.execute(&[1, 2], context()).unwrap(); - let a = outputs.pop().unwrap(); - let b = outputs.pop().unwrap(); - let (a, b) = - block_on(async { futures::join!(a.collect::>(), b.collect::>()) }); - assert_eq!( - a.iter().map(|v| **v.as_ref().unwrap()).collect::>(), - vec![0, 1, 2, 3] - ); - assert_eq!( - b.iter().map(|v| **v.as_ref().unwrap()).collect::>(), - vec![0, 1, 2, 3] - ); - } - assert_eq!(starts.get(), 2); - assert_eq!(polls.get(), 8); -} - -// Failure reaches every subscriber; cancellation stops further producer work. -#[test] -fn broadcast_error_and_cancel() { - let (source, _, polls) = source(true); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], source).unwrap(); - let mut outputs = dag.execute(&[0, 0], context()).unwrap(); - let a = outputs.pop().unwrap(); - let b = outputs.pop().unwrap(); - let (a, b) = block_on(async { futures::join!(a.collect::>(), b.collect::>()) }); - for values in [a, b] { - assert_eq!(values.len(), 2); - assert!(matches!(values[1], Err(Error::AtNode { node: 0, .. }))); - } - assert_eq!(polls.get(), 2); - let run = context(); - let mut output = dag.execute(&[0], run.clone()).unwrap().remove(0); - run.cancel(); - assert!(matches!( - block_on(output.next()), - Some(Err(Error::Cancelled)) - )); - assert!(block_on(output.next()).is_none()); - assert_eq!(polls.get(), 2); -} - -// Retaining a consumer output retains its budget lease after queue eviction. -#[test] -fn retained_outputs_count_against_budget() { - let (source, _, _) = source(false); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], source).unwrap(); - let run = RunContext::new( - Scope::Query { - evaluation_time_ms: 0, - revision: 0, - }, - Limits { - max_buffered_batches: 1, - max_bytes: 8, - }, - ) - .unwrap(); - let mut input = dag.execute(&[0], run.clone()).unwrap().remove(0); - let held = block_on(input.next()).unwrap().unwrap(); - assert_eq!(run.retained_bytes(), 8); - assert!(matches!( - block_on(input.next()), - Some(Err(Error::MemoryLimit)) - )); - drop(input); - assert_eq!(run.retained_bytes(), 8); - drop(held); - assert_eq!(run.retained_bytes(), 0); -} - -// Invalid graphs fail before even starting a source. -#[test] -fn invalid_graphs_do_not_start_sources() { - let (source, starts, _) = source(false); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], source).unwrap(); - dag.add(1, vec![2], Identity).unwrap(); - dag.add(2, vec![1], Identity).unwrap(); - assert!(dag.execute(&[0, 1], context()).is_err()); - assert_eq!(starts.get(), 0); - let mut missing = PhysicalDag::default(); - missing.add(1, vec![9], Identity).unwrap(); - assert!(missing.validate(&[1]).is_err()); - let mut arity = PhysicalDag::default(); - arity.add(1, vec![], Identity).unwrap(); - assert!(arity.validate(&[1]).is_err()); -} - -// An always-ready source must yield so cancellation can be polled on this worker. -#[test] -fn ready_sources_cooperate_with_cancellation() { - let (mut source, _, polls) = source(false); - source.end = 10_000; - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], source).unwrap(); - let context = context(); - let mut input = dag.execute(&[0], context.clone()).unwrap().remove(0); - block_on(async { - let drain = async { - while let Some(result) = input.next().await { - if let Err(error) = result { - assert_eq!(error, Error::Cancelled); - return; - } - } - panic!("source completed without yielding"); - }; - let cancel = async { - context.cancel(); - }; - futures::join!(drain, cancel); - }); - assert_eq!(polls.get(), 32); - assert_eq!(context.retained_bytes(), 0); -} - -// Cached shorter paths must not hide an over-deep path through shared nodes. -#[test] -fn depth_limit_covers_shared_paths() { - let (source, _, _) = source(false); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], source).unwrap(); - for id in 1..129 { - dag.add(id, vec![id - 1], Identity).unwrap(); - } - assert!(dag.validate(&(0..129).collect::>()).is_err()); -} diff --git a/crates/asap-physical-operators/src/dag/values.rs b/crates/asap-physical-operators/src/dag/values.rs deleted file mode 100644 index 82b873143..000000000 --- a/crates/asap-physical-operators/src/dag/values.rs +++ /dev/null @@ -1,327 +0,0 @@ -//! Runtime values preserve Planner schemas; summary states are typed values too. -use super::Error; -use crate::AggregateCore; -use planner_types::{ - post_asap::{SummaryFamilyType, SummarySchema}, - pre_asap::DataType, -}; -use std::{cmp::Ordering, sync::Arc}; -pub type Schema = Arc; -#[derive(Clone)] -pub enum Value { - Null, - Bool(bool), - Int64(i64), - Float64(f64), - Utf8(Arc), - Timestamp(i64), - Date(i32), - Interval { - months: i32, - days: i32, - nanos: i64, - }, - List(Arc<[Value]>), - Struct(Arc<[Value]>), - Map(Arc<[(Value, Value)]>), - Summary { - family: SummaryFamilyType, - state: Arc, - }, -} -impl std::fmt::Debug for Value { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::Summary { family, .. } => f.debug_tuple("Summary").field(family).finish(), - _ => write!(f, "{:?}", self.key()), - } - } -} -impl Value { - pub fn bytes(&self) -> usize { - std::mem::size_of::() - + match self { - Self::Utf8(s) => s.len(), - Self::List(v) | Self::Struct(v) => v.iter().map(Self::bytes).sum(), - Self::Map(v) => v.iter().map(|(k, v)| k.bytes() + v.bytes()).sum(), - Self::Summary { state, .. } => state.approx_memory_bytes(), - _ => 0, - } - } - pub fn matches(&self, dtype: &DataType, nullable: bool) -> bool { - if matches!(self, Self::Null) { - return nullable || matches!(dtype, DataType::Null); - } - match (self, dtype) { - (Self::Bool(_), DataType::Bool) - | (Self::Int64(_), DataType::Int64) - | (Self::Float64(_), DataType::Float64) - | (Self::Utf8(_), DataType::Utf8) - | (Self::Timestamp(_), DataType::Timestamp) - | (Self::Date(_), DataType::Date) - | (Self::Interval { .. }, DataType::Interval) => true, - (Self::List(v), DataType::List { element }) => v - .iter() - .all(|v| v.matches(&element.dtype, element.nullable)), - (Self::Struct(v), DataType::Struct { fields }) => { - v.len() == fields.len() - && v.iter() - .zip(fields) - .all(|(v, f)| v.matches(&f.dtype, f.nullable)) - } - ( - Self::Map(v), - DataType::Map { - key, - value, - value_nullable, - }, - ) => v - .iter() - .all(|(k, v)| k.matches(key, false) && v.matches(value, *value_nullable)), - _ => false, - } - } - /// Stable typed equality key. Zero signs and NaN payloads form one group. - pub fn key(&self) -> Result, Error> { - let mut out = Vec::new(); - macro_rules! number { - ($tag:expr,$v:expr) => {{ - out.push($tag); - out.extend_from_slice(&$v.to_le_bytes()); - }}; - } - match self { - Self::Null => out.push(0), - Self::Bool(v) => out.extend([1, *v as u8]), - Self::Int64(v) => number!(2, v), - Self::Float64(v) => { - let bits = if *v == 0. { - 0 - } else if v.is_nan() { - f64::NAN.to_bits() - } else { - v.to_bits() - }; - number!(3, bits); - } - Self::Utf8(v) => { - out.push(4); - out.extend(v.as_bytes()); - } - Self::Timestamp(v) => number!(5, v), - Self::Date(v) => number!(6, v), - Self::Interval { - months, - days, - nanos, - } => { - number!(7, months); - number!(8, days); - number!(9, nanos); - } - Self::List(v) | Self::Struct(v) => { - out.push(if matches!(self, Self::List(_)) { - 10 - } else { - 11 - }); - for v in v.iter() { - let key = v.key()?; - out.extend((key.len() as u64).to_le_bytes()); - out.extend(key); - } - } - Self::Map(v) => { - out.push(12); - for (k, v) in v.iter() { - for value in [k, v] { - let key = value.key()?; - out.extend((key.len() as u64).to_le_bytes()); - out.extend(key); - } - } - } - Self::Summary { .. } => { - return Err(Error::Invalid( - "summary states cannot be grouping keys".into(), - )) - } - } - Ok(out) - } - pub fn compare(&self, other: &Self) -> Result { - Ok(match (self, other) { - (Self::Null, Self::Null) => Ordering::Equal, - (Self::Int64(a), Self::Int64(b)) | (Self::Timestamp(a), Self::Timestamp(b)) => a.cmp(b), - (Self::Float64(a), Self::Float64(b)) => { - if a == b { - Ordering::Equal - } else { - a.total_cmp(b) - } - } - (Self::Utf8(a), Self::Utf8(b)) => a.cmp(b), - (Self::Bool(a), Self::Bool(b)) => a.cmp(b), - (Self::Date(a), Self::Date(b)) => a.cmp(b), - _ => { - return Err(Error::Operator( - "values do not have a supported common ordering".into(), - )) - } - }) - } -} -#[derive(Clone, Debug)] -pub struct Batch { - schema: Schema, - rows: Vec>, -} -impl Batch { - pub fn try_new(schema: Schema, rows: Vec>) -> Result { - validate_schema(&schema)?; - for row in &rows { - if row.len() != schema.fields.len() { - return Err(Error::Invalid( - "row width differs from Planner schema".into(), - )); - } - for (value, field) in row.iter().zip(&schema.fields) { - let matches = match (&field.dtype, value) { - (SummaryFamilyType::Plain(dtype), value) => { - value.matches(dtype, field.nullable) - } - (expected, Value::Summary { family, state }) => { - expected == family && validate_state(family, state.as_ref()).is_ok() - } - _ => false, - }; - if !matches { - return Err(Error::Invalid(format!( - "value differs from type of {}", - field.name - ))); - } - } - } - Ok(Self { schema, rows }) - } - pub fn schema(&self) -> &Schema { - &self.schema - } - pub fn rows(&self) -> &[Vec] { - &self.rows - } - pub fn bytes(&self) -> usize { - std::mem::size_of::() - + self - .rows - .iter() - .flat_map(|r| r.iter()) - .map(Value::bytes) - .sum::() - } -} -pub(crate) fn group_key(row: &[Value], columns: &[usize]) -> Result>, Error> { - columns - .iter() - .map(|&i| { - row.get(i) - .ok_or_else(|| Error::Invalid("group column out of range".into()))? - .key() - }) - .collect() -} - -pub(crate) fn validate_family(family: &SummaryFamilyType) -> Result<(), Error> { - use planner_types::post_asap::SketchAlgorithm as A; - match family { - SummaryFamilyType::ExactAggregate(..) => {} - SummaryFamilyType::Sketch(kind, _) - if matches!(kind.algorithm(), A::Kll | A::DDSketch | A::Hll) => {} - _ => { - return Err(Error::Invalid( - "summary family has no native DAG state implementation".into(), - )) - } - } - crate::capability::validate_summary_kernel( - family, - &planner_types::post_asap::SummaryUpdate::column( - planner_types::pre_asap::ColumnRef::SampleValue, - ), - &Default::default(), - ) - .map_err(Error::Invalid) -} -fn validate_state(family: &SummaryFamilyType, state: &dyn AggregateCore) -> Result<(), Error> { - use crate::accumulators::{ - datasketches_kll_accumulator::DatasketchesKLLAccumulator, - dd_sketch_accumulator::DDSketchAccumulator, exact_accumulator::ExactAccumulator, - hll_sketch_accumulator::HllSketchAccumulator, - }; - use planner_types::post_asap::SketchParams; - validate_family(family)?; - let valid = match family { - SummaryFamilyType::ExactAggregate(..) => { - state - .as_any() - .downcast_ref::() - .is_some_and(|s| s.family() == family && !s.is_keyed()) - || (matches!( - family, - SummaryFamilyType::ExactAggregate( - planner_types::post_asap::ExactKind::Sum, - planner_types::post_asap::ExactParams::Sum - ) - ) && state.as_any().is::()) - } - SummaryFamilyType::Sketch(kind, _) => match kind.params() { - SketchParams::Kll { k } => state - .as_any() - .downcast_ref::() - .is_some_and(|s| u32::from(s.inner.k()) == *k), - SketchParams::DDSketch { alpha } => state - .as_any() - .downcast_ref::() - .is_some_and(|s| s.inner.alpha == *alpha && s.sample_p == 1.0), - SketchParams::Hll { precision } => state - .as_any() - .downcast_ref::() - .is_some_and(|s| s.inner.precision == u32::from(*precision) && s.sample_p == 1.0), - _ => false, - }, - _ => false, - }; - if valid { - Ok(()) - } else { - Err(Error::Invalid( - "state payload differs from declared family, parameters or population layout".into(), - )) - } -} - -pub(crate) fn validate_schema(schema: &Schema) -> Result<(), Error> { - if schema.time_index.is_some_and(|index| { - schema - .fields - .get(index) - .is_none_or(|field| field.dtype != SummaryFamilyType::Plain(DataType::Timestamp)) - }) { - return Err(Error::Invalid( - "time index must name a Timestamp column".into(), - )); - } - for field in &schema.fields { - if !matches!(field.dtype, SummaryFamilyType::Plain(_)) { - validate_family(&field.dtype)?; - if field.nullable { - return Err(Error::Invalid( - "nullable summary states are not supported".into(), - )); - } - } - } - Ok(()) -} diff --git a/crates/asap-physical-operators/src/factory.rs b/crates/asap-physical-operators/src/factory.rs deleted file mode 100644 index fe6bed87f..000000000 --- a/crates/asap-physical-operators/src/factory.rs +++ /dev/null @@ -1,2046 +0,0 @@ -use crate::accumulators::{ - CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, - CountSketchWithHeapAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, - HydraKllSketchAccumulator, IncreaseAccumulator, KeyedCounterState, KeyedMaxState, - KeyedMinState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, -}; -#[cfg(test)] -use crate::AggregationType; -use crate::{AggregateCore, KeyByLabelValues, Measurement}; -#[cfg(test)] -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 crate::accumulators::hll_sketch_accumulator::HllSketchAccumulator; -use crate::accumulators::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) - } - }; -} - -/// Shared update interface for query-time and maintenance-time accumulation. -/// -/// This provides a uniform interface over all accumulator types so that the -/// worker loop doesn't need to know which concrete type it's dealing with. -pub trait AccumulatorUpdater: Send { - /// Validate an immutable 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() { - use planner_types::post_asap::{ - EntityIdentity, NonNegativeWeightProof, SummaryInputExpr, SummaryUpdate, WeightDomain, - }; - let config = topk_config(AggregationType::CountMinSketchWithHeap, None); - let family = config.accumulator_spec().unwrap().family; - let input = SummaryUpdate { - item: Some(SummaryInputExpr::Column( - planner_types::pre_asap::ColumnRef::Named("host".into()), - )), - weight: SummaryInputExpr::ResetAwareCounterDelta { - value: planner_types::pre_asap::ColumnRef::SampleValue, - series: EntityIdentity::PromqlLabelSet { excluding: vec![] }, - }, - weight_domain: WeightDomain::NonNegative { - proof: NonNegativeWeightProof::ResetAwareCounterDerivative, - }, - }; - let mut updater = create_planner_accumulator(&family, &input, &Default::default()).unwrap(); - 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> { - crate::capability::validate_summary_kernel(family, input, grouping)?; - use planner_types::post_asap::GroupingStrategy; - if grouping != &GroupingStrategy::PerSubpopulationInstance { - return Err("shared summary grouping requires a supported Planner Hydra kernel".into()); - } - if matches!(family, SummaryFamilyType::ExactAggregate(..)) { - return Ok(Box::new(PlannerExactUpdater { - acc: crate::accumulators::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()); - } - // Heap counters use fixed-point storage for fractional counter deltas. - // This encodes the selected update; it does not choose another family. - let weight_scale = if matches!( - input.weight, - planner_types::post_asap::SummaryInputExpr::ResetAwareCounterDelta { .. } - ) { - 1_000_000.0 - } else { - 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: crate::accumulators::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 = crate::accumulators::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/crates/asap-physical-operators/src/key_by_label_values.rs b/crates/asap-physical-operators/src/key_by_label_values.rs deleted file mode 100644 index 34bc84899..000000000 --- a/crates/asap-physical-operators/src/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/crates/asap-physical-operators/src/lib.rs b/crates/asap-physical-operators/src/lib.rs deleted file mode 100644 index 266a516fa..000000000 --- a/crates/asap-physical-operators/src/lib.rs +++ /dev/null @@ -1,22 +0,0 @@ -#![doc = include_str!("../README.md")] - -pub mod accumulators; -pub mod key_by_label_values; -pub mod measurement; -pub mod traits; - -pub use asap_types::{AggregationType, Statistic}; -pub use key_by_label_values::KeyByLabelValues; -pub use measurement::Measurement; -pub use traits::*; - -pub mod arithmetic; -pub mod capability; -pub mod factory; - -/// The exact Planner contract used by these kernels. -pub use planner_types as planner; - -pub mod rows; - -pub mod dag; diff --git a/crates/asap-physical-operators/src/measurement.rs b/crates/asap-physical-operators/src/measurement.rs deleted file mode 100644 index 0fe1abc0d..000000000 --- a/crates/asap-physical-operators/src/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/crates/asap-physical-operators/src/rows.rs b/crates/asap-physical-operators/src/rows.rs deleted file mode 100644 index 869ac0c30..000000000 --- a/crates/asap-physical-operators/src/rows.rs +++ /dev/null @@ -1,88 +0,0 @@ -//! Composable row operators, independent of storage, query language and sketches. -use std::collections::{BTreeMap, BTreeSet}; - -/// A semijoin preserves value-row order and multiplicity; duplicate membership -/// keys never multiply rows. Missing membership keys are reported separately so -/// the deployment can enforce the pruning proof attached to its plan. -pub fn membership_filter( - members: impl IntoIterator, - values: Vec, - identity: impl Fn(&T) -> K, -) -> (Vec, BTreeSet) { - let members: BTreeSet = members.into_iter().collect(); - let mut missing = members.clone(); - let rows = values - .into_iter() - .filter(|row| { - let key = identity(row); - missing.remove(&key); - members.contains(&key) - }) - .collect(); - (rows, missing) -} - -/// Stable descending TopK per group. NaN sorts after numeric values; ties keep -/// input order. This operator does not know how its input was filtered or built. -pub fn grouped_topk( - values: Vec, - k: usize, - group_key: impl Fn(&T) -> K, - score: impl Fn(&T) -> f64, -) -> Vec { - let mut groups: BTreeMap> = BTreeMap::new(); - for row in values { - groups.entry(group_key(&row)).or_default().push(row); - } - groups - .into_values() - .flat_map(|mut rows| { - rows.sort_by(|a, b| { - let (a, b) = (score(a), score(b)); - match (a.is_nan(), b.is_nan()) { - (true, true) => std::cmp::Ordering::Equal, - (true, false) => std::cmp::Ordering::Greater, - (false, true) => std::cmp::Ordering::Less, - (false, false) => b.total_cmp(&a), - } - }); - rows.truncate(k); - rows - }) - .collect() -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn semijoin_preserves_values_order_and_duplicates_without_ranking() { - let (rows, missing) = membership_filter( - ["b", "c", "b", "missing"], - vec![("a", 100.), ("b", 2.), ("c", 9.), ("b", 3.)], - |row| row.0, - ); - assert_eq!(rows, vec![("b", 2.), ("c", 9.), ("b", 3.)]); - assert_eq!(missing, BTreeSet::from(["missing"])); - assert_eq!( - grouped_topk(rows, 2, |_| (), |r| r.1), - vec![("c", 9.), ("b", 3.)] - ); - } - - #[test] - fn grouped_ranking_preserves_ties_and_places_nan_last() { - let rows = vec![("x", 0, f64::NAN), ("x", 1, 2.), ("y", 2, 8.), ("x", 3, 2.)]; - let ranked = grouped_topk(rows, 2, |r| r.0, |r| r.2); - assert_eq!(ranked, vec![("x", 1, 2.), ("x", 3, 2.), ("y", 2, 8.)]); - assert!(grouped_topk(vec![1], 0, |_| (), |r| *r as f64).is_empty()); - } - - #[test] - fn empty_membership_removes_all_rows() { - let (rows, missing) = membership_filter([], vec![1, 2], |r| *r); - assert!(rows.is_empty()); - assert!(missing.is_empty()); - } -} diff --git a/crates/asap-physical-operators/src/traits.rs b/crates/asap-physical-operators/src/traits.rs deleted file mode 100644 index 2b99da789..000000000 --- a/crates/asap-physical-operators/src/traits.rs +++ /dev/null @@ -1,351 +0,0 @@ -use crate::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/crates/asap-physical-operators/tests/deployment.rs b/crates/asap-physical-operators/tests/deployment.rs deleted file mode 100644 index ffdd3aea5..000000000 --- a/crates/asap-physical-operators/tests/deployment.rs +++ /dev/null @@ -1,96 +0,0 @@ -//! Exercise the public library without a backend server, store, or scheduler. -use asap_physical_operators::planner::{ - post_asap::{ - GroupingStrategy, SketchAlgorithm, SketchKind, SketchParams, SummaryFamilyType, - SummaryUpdate, - }, - pre_asap::ColumnRef, -}; -use asap_physical_operators::{factory::create_planner_accumulator, AggregateCore, Statistic}; -use std::collections::HashMap; - -fn family(k: u32) -> SummaryFamilyType { - SummaryFamilyType::Sketch( - SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k }), - GroupingStrategy::PerSubpopulationInstance, - ) -} -fn build(values: &[f64]) -> Box { - let mut operator = create_planner_accumulator( - &family(512), - &SummaryUpdate::column(ColumnRef::SampleValue), - &Default::default(), - ) - .unwrap(); - for (at, value) in values.iter().enumerate() { - operator.validate_single_input(*value).unwrap(); - operator.update_single(*value, at as i64); - } - operator.into_accumulator() -} -fn read(state: &dyn AggregateCore) -> f64 { - state - .query_statistic( - Statistic::Quantile, - &None, - &HashMap::from([("quantile".into(), "0.5".into())]), - ) - .unwrap() -} - -// The same kernels work when every build is query-time, when only a prefix -// was precomputed, and when all state was precomputed before the readout. -#[test] -fn raw_partial_and_fully_precomputed_use_the_same_kernels() { - let raw: Vec = (0..128).map(f64::from).collect(); - let raw_only = build(&raw); - let stored_prefix = build(&raw[..64]); - let query_time_suffix = build(&raw[64..]); - let partial = stored_prefix.merge_with(&*query_time_suffix).unwrap(); - let stored_complete = build(&raw); - assert_eq!(read(&*raw_only), read(&*partial)); - assert_eq!(read(&*partial), read(&*stored_complete)); - assert!((read(&*raw_only) - 64.0).abs() <= 1.0); -} - -// A compiler must reject invalid physical parameters before starting execution. -#[test] -fn invalid_kll_parameters_are_rejected_at_binding() { - let result = create_planner_accumulator( - &family(0), - &SummaryUpdate::column(ColumnRef::SampleValue), - &Default::default(), - ); - assert!(result.is_err()); -} - -// Native CountSketch supports the confidence-sized depth used by the backend; -// a packed-wire column-bit budget must not be imposed on this constructor. -#[test] -fn native_count_sketch_dimensions_are_not_packed_wire_dimensions() { - use asap_physical_operators::planner::post_asap::SummaryInputExpr; - use asap_physical_operators::KeyByLabelValues; - let family = SummaryFamilyType::Sketch( - SketchKind::new( - SketchAlgorithm::CountSketchWithHeap, - SketchParams::CountSketchWithHeap { - width: 1200, - depth: 55, - heap_size: 3, - }, - ), - Default::default(), - ); - let mut update = SummaryUpdate::column(ColumnRef::SampleValue); - update.item = Some(SummaryInputExpr::Column(ColumnRef::Named("host".into()))); - let mut operator = create_planner_accumulator(&family, &update, &Default::default()).unwrap(); - let key = KeyByLabelValues::new_with_labels(vec!["a".into()]); - operator.update_keyed(&key, 7.0, 1000); - let state = operator.into_accumulator(); - assert_eq!( - state - .query_statistic(Statistic::Sum, &Some(key), &Default::default()) - .unwrap(), - 7.0 - ); -} diff --git a/crates/asap-physical-operators/tests/physical_dag.rs b/crates/asap-physical-operators/tests/physical_dag.rs deleted file mode 100644 index a46cd2d01..000000000 --- a/crates/asap-physical-operators/tests/physical_dag.rs +++ /dev/null @@ -1,663 +0,0 @@ -//! Acceptance tests use the library directly, without either backend engine. -use asap_physical_operators::{ - dag::{ - operators::{Expression, Operator, Reduction, SortKey}, - values::{Batch, Schema, Value}, - Limits, PhysicalDag, RunContext, Scope, - }, - Statistic, -}; -use futures::{executor::block_on, StreamExt}; -use planner_types::{ - post_asap::{ExactKind, ExactParams, SummaryFamilyType, SummaryField, SummarySchema}, - pre_asap::DataType, -}; -use std::sync::Arc; -fn schema(fields: &[(&str, DataType, bool)]) -> Schema { - Arc::new(SummarySchema { - fields: fields - .iter() - .map(|(name, dtype, nullable)| SummaryField { - name: (*name).into(), - dtype: SummaryFamilyType::Plain(dtype.clone()), - nullable: *nullable, - }) - .collect(), - time_index: None, - }) -} -fn run(dag: &PhysicalDag<'_, Batch, Schema>, root: u64, scope: Scope) -> Vec> { - let context = RunContext::new( - scope, - Limits { - max_buffered_batches: 1, - ..Limits::default() - }, - ) - .unwrap(); - block_on(async { - let mut stream = dag.execute(&[root], context.clone()).unwrap().remove(0); - let mut rows = vec![]; - while let Some(batch) = stream.next().await { - rows.extend(batch.unwrap().rows().iter().cloned()); - } - assert_eq!(context.retained_bytes(), 0); - rows - }) -} -fn query() -> Scope { - Scope::Query { - evaluation_time_ms: 1000, - revision: 2, - } -} -fn floats(rows: &[Vec], column: usize) -> Vec { - rows.iter() - .map(|r| { - if let Value::Float64(v) = r[column] { - v - } else { - panic!("not Float64") - } - }) - .collect() -} - -// Sort followed by partitioned Limit implements ranking independently per group. -#[test] -fn grouped_sort_limit_across_batches() { - let schema = schema(&[ - ("group", DataType::Int64, false), - ("score", DataType::Float64, false), - ]); - let batches = [ - vec![(1, 1.), (2, 4.), (1, 9.)], - vec![(2, 8.), (1, 5.), (2, 2.)], - ] - .into_iter() - .map(|rows| { - Batch::try_new( - schema.clone(), - rows.into_iter() - .map(|(g, v)| vec![Value::Int64(g), Value::Float64(v)]) - .collect(), - ) - .unwrap() - }) - .collect(); - let mut dag = PhysicalDag::default(); - dag.add( - 0, - vec![], - Operator::source(schema.clone(), batches).unwrap(), - ) - .unwrap(); - dag.add( - 1, - vec![0], - Operator::sort( - schema.clone(), - vec![SortKey { - column: 1, - descending: true, - nulls_first: false, - }], - vec![0], - ) - .unwrap(), - ) - .unwrap(); - dag.add(2, vec![1], Operator::limit(schema, 1, 1, vec![0]).unwrap()) - .unwrap(); - assert_eq!(floats(&run(&dag, 2, query()), 1), vec![5., 4.]); -} - -// The same computation runs in either engine scope with fresh per-run state. -#[test] -fn summary_construction_merge_and_readout_at_both_phases() { - let schema = schema(&[("v", DataType::Float64, false)]); - let batches = (1..=20) - .map(|v| Batch::try_new(schema.clone(), vec![vec![Value::Float64(v as f64)]]).unwrap()) - .collect(); - let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); - let build = Operator::summary_build(schema.clone(), family, 0, None, vec![]).unwrap(); - let state = build.schema(); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], Operator::source(schema, batches).unwrap()) - .unwrap(); - dag.add(1, vec![0], build).unwrap(); - dag.add(2, vec![1, 1], Operator::union(state.clone(), 2).unwrap()) - .unwrap(); - dag.add( - 3, - vec![2], - Operator::summary_merge(state.clone(), 0, vec![]).unwrap(), - ) - .unwrap(); - dag.add( - 4, - vec![3], - Operator::readout(state, 0, Statistic::Sum, Default::default()).unwrap(), - ) - .unwrap(); - for scope in [ - query(), - Scope::Ingestion { - window_start_ms: 0, - window_end_ms: 1000, - revision: 2, - }, - ] { - assert_eq!(floats(&run(&dag, 4, scope), 0), vec![420.]); - } -} - -// A semi-join can consume two branches of one producer with a one-batch buffer. -#[test] -fn diamond_semijoin_preserves_left_values_and_multiplicity() { - let schema = schema(&[("key", DataType::Int64, false)]); - let batches = [1, 2, 2, 3] - .into_iter() - .map(|v| Batch::try_new(schema.clone(), vec![vec![Value::Int64(v)]]).unwrap()) - .collect(); - let filter = Operator::filter( - schema.clone(), - Expression::Equal( - Box::new(Expression::Column(0)), - Box::new(Expression::Literal { - value: Value::Int64(2), - dtype: DataType::Int64, - }), - ), - ) - .unwrap(); - let mut dag = PhysicalDag::default(); - dag.add( - 0, - vec![], - Operator::source(schema.clone(), batches).unwrap(), - ) - .unwrap(); - dag.add(1, vec![0], filter).unwrap(); - dag.add( - 2, - vec![0, 1], - Operator::semi_join(schema.clone(), schema, vec![(0, 0)]).unwrap(), - ) - .unwrap(); - let rows = run(&dag, 2, query()); - assert_eq!(rows.len(), 2); - assert!(rows.iter().all(|r| matches!(r[0], Value::Int64(2)))); -} - -// Integer aggregation must not silently lose precision through Float64. -#[test] -fn exact_integer_and_empty_extrema() { - let schema = schema(&[("v", DataType::Int64, false)]); - let aggregate = Operator::aggregate( - schema.clone(), - vec![], - vec![("sum".into(), Reduction::Sum(0))], - ) - .unwrap(); - let mut dag = PhysicalDag::default(); - let value = 9_007_199_254_740_993; - dag.add( - 0, - vec![], - Operator::source( - schema.clone(), - vec![Batch::try_new( - schema.clone(), - vec![vec![Value::Int64(value)], vec![Value::Int64(2)]], - ) - .unwrap()], - ) - .unwrap(), - ) - .unwrap(); - dag.add(1, vec![0], aggregate).unwrap(); - assert!(matches!(run(&dag,1,query())[0][0],Value::Int64(v) if v==value+2)); - let mut empty = PhysicalDag::default(); - empty - .add(0, vec![], Operator::source(schema.clone(), vec![]).unwrap()) - .unwrap(); - empty - .add( - 1, - vec![0], - Operator::aggregate(schema, vec![], vec![("min".into(), Reduction::Min(0))]).unwrap(), - ) - .unwrap(); - assert!(matches!(run(&empty, 1, query())[0][0], Value::Null)); -} - -// Plain value operators are library implementations, including NaN comparison. -#[test] -fn scalar_negation_and_vector_conversion() { - let scalar = Operator::scalar(Value::Float64(7.), DataType::Float64).unwrap(); - let project = Operator::project( - scalar.schema(), - vec![( - "v".into(), - Expression::Negate(Box::new(Expression::Column(0))), - )], - ) - .unwrap(); - let convert = Operator::vector_to_scalar(project.schema(), 0).unwrap(); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], scalar).unwrap(); - dag.add(1, vec![0], project).unwrap(); - dag.add(2, vec![1], convert).unwrap(); - assert_eq!(floats(&run(&dag, 2, query()), 0), vec![-7.]); - let scalar = Operator::scalar(Value::Float64(f64::NAN), DataType::Float64).unwrap(); - let predicate = Expression::Equal( - Box::new(Expression::Column(0)), - Box::new(Expression::Column(0)), - ); - let filter = Operator::filter(scalar.schema(), predicate).unwrap(); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], scalar).unwrap(); - dag.add(1, vec![0], filter).unwrap(); - assert!(run(&dag, 1, query()).is_empty()); -} - -// Invalid operations fail at binding rather than becoming external fallbacks. -#[test] -fn binding_rejects_unsupported_operations() { - let schema = schema(&[("v", DataType::Float64, false)]); - assert!(Operator::summary_build( - schema.clone(), - SummaryFamilyType::ExactAggregate(ExactKind::Rate, ExactParams::Rate), - 0, - None, - vec![] - ) - .is_err()); - let sum = Operator::summary_build( - schema.clone(), - SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum), - 0, - None, - vec![], - ) - .unwrap(); - assert!(Operator::readout(sum.schema(), 0, Statistic::Quantile, Default::default()).is_err()); - assert!(Operator::filter(schema, Expression::Column(0)).is_err()); -} - -// KLL is one family example: precomputation changes input sources, not operators. -#[test] -fn kll_raw_partial_and_precomputed_are_native_dags() { - use planner_types::post_asap::{GroupingStrategy, SketchAlgorithm, SketchKind, SketchParams}; - let input = schema(&[("value", DataType::Float64, false)]); - let family = SummaryFamilyType::Sketch( - SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k: 512 }), - GroupingStrategy::PerSubpopulationInstance, - ); - let build = Operator::summary_build(input.clone(), family, 0, None, vec![]).unwrap(); - let state = build.schema(); - let build_range = |start: u32, end: u32| { - let mut dag = PhysicalDag::default(); - let batch = Batch::try_new( - input.clone(), - (start..end) - .map(|v| vec![Value::Float64(f64::from(v))]) - .collect(), - ) - .unwrap(); - dag.add( - 0, - vec![], - Operator::source(input.clone(), vec![batch]).unwrap(), - ) - .unwrap(); - dag.add(1, vec![0], build.clone()).unwrap(); - run( - &dag, - 1, - Scope::Ingestion { - window_start_ms: 0, - window_end_ms: 1000, - revision: 1, - }, - ) - }; - let prefix = build_range(0, 64); - let complete = build_range(0, 128); - let query_plan = |stored: Option>>, raw_start: Option| { - let mut dag = PhysicalDag::default(); - let mut states = vec![]; - if let Some(rows) = stored { - dag.add( - 0, - vec![], - Operator::source( - state.clone(), - vec![Batch::try_new(state.clone(), rows).unwrap()], - ) - .unwrap(), - ) - .unwrap(); - states.push(0); - } - if let Some(start) = raw_start { - dag.add( - 1, - vec![], - Operator::source( - input.clone(), - vec![Batch::try_new( - input.clone(), - (start..128) - .map(|v| vec![Value::Float64(f64::from(v))]) - .collect(), - ) - .unwrap()], - ) - .unwrap(), - ) - .unwrap(); - dag.add(2, vec![1], build.clone()).unwrap(); - states.push(2); - } - dag.add( - 3, - states.clone(), - Operator::union(state.clone(), states.len()).unwrap(), - ) - .unwrap(); - dag.add( - 4, - vec![3], - Operator::summary_merge(state.clone(), 0, vec![]).unwrap(), - ) - .unwrap(); - dag.add( - 5, - vec![4], - Operator::readout( - state.clone(), - 0, - Statistic::Quantile, - std::collections::HashMap::from([("quantile".into(), "0.5".into())]), - ) - .unwrap(), - ) - .unwrap(); - floats(&run(&dag, 5, query()), 0)[0] - }; - let raw = query_plan(None, Some(0)); - let partial = query_plan(Some(prefix), Some(64)); - let full = query_plan(Some(complete), None); - assert_eq!(raw, partial); - assert_eq!(partial, full); - assert!((raw - 64.).abs() <= 1.); -} - -// Restored state must retain its family; a mislabeled state is rejected. -#[test] -fn restored_exact_state_and_family_validation() { - use asap_physical_operators::{ - accumulators::exact_accumulator::ExactAccumulator, SerializableToSink, - }; - let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); - let mut acc = ExactAccumulator::new(family.clone(), false).unwrap(); - acc.update(None, 7., 0); - let acc = ExactAccumulator::deserialize_from_bytes(&acc.serialize_to_bytes()).unwrap(); - let schema = Arc::new(SummarySchema { - fields: vec![SummaryField { - name: "state".into(), - dtype: family.clone(), - nullable: false, - }], - time_index: None, - }); - let value = Value::Summary { - family: family.clone(), - state: Arc::new(acc), - }; - let mut dag = PhysicalDag::default(); - dag.add( - 0, - vec![], - Operator::source( - schema.clone(), - vec![Batch::try_new(schema.clone(), vec![vec![value]]).unwrap()], - ) - .unwrap(), - ) - .unwrap(); - dag.add( - 1, - vec![0], - Operator::readout(schema.clone(), 0, Statistic::Sum, Default::default()).unwrap(), - ) - .unwrap(); - assert_eq!(floats(&run(&dag, 1, query()), 0), vec![7.]); - let wrong = ExactAccumulator::new( - SummaryFamilyType::ExactAggregate(ExactKind::Max, ExactParams::Max), - false, - ) - .unwrap(); - assert!(Batch::try_new( - schema, - vec![vec![Value::Summary { - family, - state: Arc::new(wrong) - }]] - ) - .is_err()); -} - -// Planner binding rejects unknown computation instead of accepting a fallback. -#[test] -fn bind_post_asap_before_execution() { - use asap_physical_operators::dag::planner::bind; - use planner_types::{ - post_asap::{ - EdgeRole, ExecutableDag, ExecutableDagEdge, ExecutableDagNode, - ExecutableOperatorPayload, ExecutionDataState, ExecutionTiming, - GroupingEdgeCompatibility, PostAsapNodeId, ValueOperation, WindowEdgeCompatibility, - }, - pre_asap::{ArithmeticOpKind, ProjectItem, QueryExpr, ScalarValue}, - }; - use std::{collections::BTreeMap, rc::Rc}; - let schema = schema(&[("value", DataType::Float64, false)]); - let node = |id, payload| ExecutableDagNode { - id: PostAsapNodeId(id), - payload, - output_state: ExecutionDataState::QUERY_ROWS, - output_schema: (*schema).clone(), - guarantee: None, - }; - let mut dag = ExecutableDag { - nodes: vec![ - node( - 0, - ExecutableOperatorPayload::Fallback { - expression: QueryExpr::promql_scalar(1.), - }, - ), - node( - 1, - ExecutableOperatorPayload::Value { - timing: ExecutionTiming::QueryTime, - operation: ValueOperation::Project { - cols: vec![ProjectItem { - alias: None, - expr: QueryExpr::Arithmetic { - op: ArithmeticOpKind::Add, - left: Rc::new(QueryExpr::Column(0)), - right: Rc::new(QueryExpr::Literal(ScalarValue::Float64(2.))), - }, - }], - qualifier: None, - }, - }, - ), - ], - edges: vec![ExecutableDagEdge { - producer: PostAsapNodeId(0), - consumer: PostAsapNodeId(1), - role: EdgeRole::Input, - intermediate_schema: (*schema).clone(), - data_state: ExecutionDataState::QUERY_ROWS, - grouping: GroupingEdgeCompatibility::NotApplicable, - window: WindowEdgeCompatibility::NotApplicable, - }], - root: PostAsapNodeId(1), - }; - let sources = || -> BTreeMap> { - BTreeMap::from([( - 0, - Box::new( - Operator::source( - schema.clone(), - vec![Batch::try_new(schema.clone(), vec![vec![Value::Float64(1.)]]).unwrap()], - ) - .unwrap(), - ) as asap_physical_operators::dag::planner::Source<'static>, - )]) - }; - let native = bind(&dag, sources(), &[1]).unwrap(); - assert_eq!(floats(&run(&native, 1, query()), 0), vec![3.]); - assert!(bind(&dag, BTreeMap::new(), &[1]).is_err()); - dag.nodes[1].payload = ExecutableOperatorPayload::Value { - timing: ExecutionTiming::QueryTime, - operation: ValueOperation::Extension { - name: "unknown".into(), - }, - }; - assert!(bind(&dag, sources(), &[1]).is_err()); -} - -// A completed empty population has an exact zero count, with integer output. -#[test] -fn empty_exact_count_is_an_integer_state_readout() { - let input = schema(&[("value", DataType::Float64, false)]); - let build = Operator::summary_build( - input.clone(), - SummaryFamilyType::ExactAggregate(ExactKind::Count, ExactParams::Count), - 0, - None, - vec![], - ) - .unwrap(); - let read = Operator::readout(build.schema(), 0, Statistic::Count, Default::default()).unwrap(); - let mut dag = PhysicalDag::default(); - dag.add(0, vec![], Operator::source(input, vec![]).unwrap()) - .unwrap(); - dag.add(1, vec![0], build).unwrap(); - dag.add(2, vec![1], read).unwrap(); - assert!(matches!(run(&dag, 2, query())[0][0], Value::Int64(0))); -} - -// A deployment source cannot pass a different row shape to bound expressions. -#[test] -fn source_batches_must_match_the_bound_schema() { - use asap_physical_operators::dag::{self, PhysicalOperator}; - use planner_types::{ - post_asap::{ - ExecutableDag, ExecutableDagNode, ExecutableOperatorPayload, ExecutionDataState, - PostAsapNodeId, - }, - pre_asap::QueryExpr, - }; - use std::{cell::Cell, collections::BTreeMap, rc::Rc}; - struct WrongSource { - schema: Schema, - starts: Rc>, - } - impl PhysicalOperator for WrongSource { - fn name(&self) -> &str { - "ExternalSource" - } - fn input_schemas(&self) -> Vec { - vec![] - } - fn output_schema(&self) -> Schema { - self.schema.clone() - } - fn output_bytes(&self, value: &Batch) -> usize { - value.bytes() - } - fn start<'a>( - &'a self, - _: Vec>, - _: RunContext, - ) -> Result, dag::Error> { - self.starts.set(self.starts.get() + 1); - Ok( - futures::stream::once(async { Batch::try_new(schema(&[]), vec![vec![]]) }) - .boxed_local(), - ) - } - } - let expected = schema(&[("value", DataType::Float64, false)]); - let starts = Rc::new(Cell::new(0)); - let plan = ExecutableDag { - nodes: vec![ExecutableDagNode { - id: PostAsapNodeId(0), - payload: ExecutableOperatorPayload::Fallback { - expression: QueryExpr::promql_scalar(1.), - }, - output_state: ExecutionDataState::QUERY_ROWS, - output_schema: (*expected).clone(), - guarantee: None, - }], - edges: vec![], - root: PostAsapNodeId(0), - }; - let source = Box::new(WrongSource { - schema: expected, - starts: starts.clone(), - }) as dag::planner::Source<'static>; - let native = dag::planner::bind(&plan, BTreeMap::from([(0, source)]), &[0]).unwrap(); - assert_eq!(starts.get(), 0); - let context = RunContext::new(query(), Limits::default()).unwrap(); - let mut output = native.execute(&[0], context).unwrap().remove(0); - assert!(matches!( - block_on(output.next()), - Some(Err(dag::Error::AtNode { node: 0, .. })) - )); - assert_eq!(starts.get(), 1); -} - -// Float extrema have the same NaN behavior as the exact summary kernels. -#[test] -fn extrema_preserve_numeric_values_in_the_presence_of_nan() { - let input = schema(&[("v", DataType::Float64, false)]); - let mut dag = PhysicalDag::default(); - dag.add( - 0, - vec![], - Operator::source( - input.clone(), - vec![Batch::try_new( - input.clone(), - vec![vec![Value::Float64(-f64::NAN)], vec![Value::Float64(5.)]], - ) - .unwrap()], - ) - .unwrap(), - ) - .unwrap(); - dag.add( - 1, - vec![0], - Operator::aggregate( - input, - vec![], - vec![ - ("min".into(), Reduction::Min(0)), - ("max".into(), Reduction::Max(0)), - ], - ) - .unwrap(), - ) - .unwrap(); - let rows = run(&dag, 1, query()); - assert_eq!(floats(&rows, 0), vec![5.]); - assert_eq!(floats(&rows, 1), vec![5.]); -} 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/derived_input.rs b/crates/asap_types/src/derived_input.rs index e4596b7ca..759063303 100644 --- a/crates/asap_types/src/derived_input.rs +++ b/crates/asap_types/src/derived_input.rs @@ -224,7 +224,7 @@ mod tests { .into_iter() .map(|id| OwnedPostAsapNode { id: PostAsapNodeId(id), - payload: serde_json::json!({"kind":"summary_merge", "timing":"ingestion_time"}), + payload: serde_json::json!({"kind":"summary_merge"}), output_state: state, output_schema: serde_json::json!({"fields":[],"time_index":null}), guarantee: None, 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 7c1b947a5..1b087f3ae 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/precompute_plan.rs b/crates/asap_types/src/precompute_plan.rs index 12d9f421f..418f1b328 100644 --- a/crates/asap_types/src/precompute_plan.rs +++ b/crates/asap_types/src/precompute_plan.rs @@ -623,7 +623,7 @@ impl PrecomputePlan { .filter(|edge| edge.consumer == id) .collect(); use planner_types::post_asap::{ - ExecutableOperatorPayload as Payload, ExecutionTiming, ValueOperation, + ExecutableOperatorPayload as Payload, ValueOperation, }; if node.output_state != planner_types::post_asap::ExecutionDataState::INGESTION_ROWS @@ -633,32 +633,29 @@ impl PrecomputePlan { match &node.payload { Payload::Value { operation: ValueOperation::FinalizeExactAccumulator, - timing: ExecutionTiming::IngestionTime, } if children.len() == 1 && frontiers.contains_key(&children[0].producer) => {} - Payload::Binary { - operator, - timing: ExecutionTiming::IngestionTime, - } if children.len() == 2 - && children - .iter() - .filter(|edge| { - edge.role == planner_types::post_asap::EdgeRole::Left - }) - .count() - == 1 - && children - .iter() - .filter(|edge| { - edge.role == planner_types::post_asap::EdgeRole::Right - }) - .count() - == 1 - && operator.vector_match.is_none() - && matches!( - operator.kind, - planner_types::pre_asap::BinaryOpKind::Arithmetic(_) - ) => + Payload::Binary { operator } + if children.len() == 2 + && children + .iter() + .filter(|edge| { + edge.role == planner_types::post_asap::EdgeRole::Left + }) + .count() + == 1 + && children + .iter() + .filter(|edge| { + edge.role == planner_types::post_asap::EdgeRole::Right + }) + .count() + == 1 + && operator.vector_match.is_none() + && matches!( + operator.kind, + planner_types::pre_asap::BinaryOpKind::Arithmetic(_) + ) => { pending.extend(children.iter().map(|edge| edge.producer)); } @@ -1077,9 +1074,7 @@ mod source_window_cohort_tests { assert!(validate_maintenance_reduction(&config, &node).is_err()); config.partitioning = None; assert!(validate_maintenance_reduction(&config, &node).is_err()); - node.payload = ExecutableOperatorPayload::SummaryMerge { - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, - }; + node.payload = ExecutableOperatorPayload::SummaryMerge; assert!(validate_maintenance_reduction(&config, &node).is_err()); } diff --git a/crates/asap_types/src/query_plan.rs b/crates/asap_types/src/query_plan.rs index fc73f58c8..7282719b9 100644 --- a/crates/asap_types/src/query_plan.rs +++ b/crates/asap_types/src/query_plan.rs @@ -391,7 +391,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 } @@ -401,7 +411,7 @@ impl QueryPlanEntry { || guarantee.failure_probability.evaluate().is_none() ) { return Err(QueryPlanError::Invalid( - "invalid MembershipFilter completeness certificate".into(), + "invalid semi-join pruning certificate".into(), )); } } @@ -563,6 +573,7 @@ pub enum QueryPlanNode { 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, @@ -605,13 +616,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 { @@ -637,7 +641,6 @@ impl QueryPlanNode { 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 5e9c27b28..289105b96 100644 --- a/data_plane/Cargo.toml +++ b/data_plane/Cargo.toml @@ -7,7 +7,7 @@ edition.workspace = true # Internal crates (workspace) asap_types.workspace = true asap-physical-operators.workspace = true -asap_sketch_codec = { path = "../crates/asap_sketch_codec" } +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 @@ -65,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 = "026cd18c7b8c23ae6c46d4d683151ba562b8cd3a" } # 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 diff --git a/data_plane/src/precompute_engine/maintenance_runtime.rs b/data_plane/src/precompute_engine/maintenance_runtime.rs index f33317a79..569efc5ea 100644 --- a/data_plane/src/precompute_engine/maintenance_runtime.rs +++ b/data_plane/src/precompute_engine/maintenance_runtime.rs @@ -166,14 +166,12 @@ impl PrecomputeOperatorRegistry for OperatorAdapter<'_> { node: &ExecutableDagNode, inputs: &[Arc], ) -> Result { + if node.output_state.timing != planner_types::post_asap::ExecutionTiming::IngestionTime { + return Err("ingestion executor received a query-time node".into()); + } match &node.payload { - ExecutableOperatorPayload::SummaryMerge { - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, - } => merge_inputs(inputs), - ExecutableOperatorPayload::Binary { - operator, - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, - } => { + ExecutableOperatorPayload::SummaryMerge => merge_inputs(inputs), + ExecutableOperatorPayload::Binary { operator } => { if !self.inputs.frozen_inputs().is_some() || node.output_state != planner_types::post_asap::ExecutionDataState::INGESTION_ROWS @@ -185,7 +183,6 @@ impl PrecomputeOperatorRegistry for OperatorAdapter<'_> { ExecutableOperatorPayload::Value { operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, } => { if !self.inputs.frozen_inputs().is_some() { return Err( @@ -2251,9 +2248,7 @@ mod tests { fn node(id: u32) -> ExecutableDagNode { ExecutableDagNode { id: PostAsapNodeId(id), - payload: ExecutableOperatorPayload::SummaryMerge { - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, - }, + payload: ExecutableOperatorPayload::SummaryMerge, output_state: planner_types::post_asap::ExecutionDataState::INGESTION_SUMMARY, output_schema: SummarySchema { fields: vec![], @@ -2448,7 +2443,6 @@ mod tests { let mut read = node(2); read.payload = ExecutableOperatorPayload::Value { operation: planner_types::post_asap::ValueOperation::FinalizeExactAccumulator, - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, }; read.output_schema.fields = vec![SummaryField { name: "value".into(), @@ -2911,9 +2905,7 @@ mod tests { second_node.id = PostAsapNodeId(5); let mut merge = second_node.clone(); merge.id = PostAsapNodeId(6); - merge.payload = ExecutableOperatorPayload::SummaryMerge { - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, - }; + merge.payload = ExecutableOperatorPayload::SummaryMerge; dag.nodes.extend([second_node, merge]); let original = dag .edges @@ -3401,7 +3393,6 @@ mod tests { }; operation.payload = ExecutableOperatorPayload::Binary { operator: operator.clone(), - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, }; operation.output_state = planner_types::post_asap::ExecutionDataState::INGESTION_ROWS; assert!(frozen @@ -3420,11 +3411,12 @@ mod tests { .is_err()); operation.payload = ExecutableOperatorPayload::Binary { operator: operator.clone(), - timing: planner_types::post_asap::ExecutionTiming::QueryTime, }; + operation.output_state = planner_types::post_asap::ExecutionDataState::QUERY_ROWS; 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/subdag_scheduler.rs b/data_plane/src/precompute_engine/subdag_scheduler.rs index a9d10fec3..c00fab6fe 100644 --- a/data_plane/src/precompute_engine/subdag_scheduler.rs +++ b/data_plane/src/precompute_engine/subdag_scheduler.rs @@ -347,7 +347,6 @@ mod tests { } let mut binary = node(3); binary.payload = ExecutableOperatorPayload::Binary { - timing: planner_types::post_asap::ExecutionTiming::IngestionTime, operator: BinaryOperator { checked_relative_division: false, checked_finite_division: false, 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 30f39c8bf..108609100 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 @@ -456,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, @@ -666,7 +667,7 @@ mod tests { root, QueryPlanNode::Relational { input: sort, - operation: serde_json::to_value(ValueOperation::Limit { n: 1, offset: 0 }) + 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 c47b7b48e..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 @@ -484,7 +484,16 @@ impl ClickHouseRelationalAdapter { .rows .sort_by(|left, right| compare_sort_keys(left, right, keys, &schema)); } - ValueOperation::Limit { n, offset } => { + ValueOperation::Limit { + n, + offset, + partition_by, + } => { + if !partition_by.keys().is_empty() || partition_by.is_without() { + return Err(ClickHouseRelationalError::Unsupported( + "partitioned relation Limit is not bound".into(), + )); + } input.rows = input.rows.into_iter().skip(*offset).take(*n).collect(); } other => return Err(ClickHouseRelationalError::Unsupported(format!("{other:?}"))), @@ -1604,7 +1613,11 @@ mod tests { }], partition_by: GroupKeys::none(), }, - ValueOperation::Limit { n: 1, offset: 0 }, + ValueOperation::Limit { + n: 1, + offset: 0, + partition_by: planner_types::pre_asap::GroupKeys::none(), + }, ] { relation = adapter .apply_operation(&operation, &projected_schema, relation) 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 72cf22f6a..cb90c72b0 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 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 ba8356858..e82f54d10 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,9 @@ 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::RelationalJoin { inputs, .. } => { pending.extend(inputs.iter().map(|input| (*input, at))); } QueryPlanNode::ExternalExact { request, inputs } => { @@ -650,21 +650,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, 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 86b29c8f6..2c1e7a813 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, @@ -205,13 +206,34 @@ 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, + .. } => { - let candidates = vector(self.eval(inputs[0], at)?)?; - let values = vector(self.eval(inputs[1], at)?)?; - let (selected, warning) = membership_filter(candidates, values, &completeness)?; + 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"))?; + let keys = asap_physical_operators::dag::planner::equijoin_keys( + &predicate, + &left_schema, + &right_schema, + ) + .map_err(|error| miss(error.to_string()))? + .into_iter() + .map(|(left, right)| { + ( + left_schema.fields[left].name.clone(), + right_schema.fields[right].name.clone(), + ) + }) + .collect::>(); + let (selected, warning) = semi_join(candidates, values, &keys, pruning.as_ref())?; if let Some(warning) = warning { self.warnings.push(warning); } @@ -277,9 +299,15 @@ 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, + )?)) } ResidualQueryOperator::Binary { operation, @@ -350,22 +378,14 @@ 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, + )?)) } ResidualQueryOperator::HistogramQuantile => { let Value::Scalar(quantile) = self.eval(input(0)?, at)? else { @@ -420,27 +440,39 @@ impl Result> Evaluator<' } } -fn membership_filter( +fn semi_join( candidates: Vector, values: Vector, - completeness: &CandidateCompleteness, + keys: &[(String, String)], + completeness: Option<&CandidateCompleteness>, ) -> Result<(Vector, Option), EngineError> { - let identity = |labels: &Labels| { - let mut labels = labels.clone(); - labels.remove("__name__"); - labels + let left_key = |labels: &Labels| { + keys.iter() + .map(|(left, _)| labels.get(left).cloned().unwrap_or_default()) + .collect::>() }; - let (selected, missing) = asap_physical_operators::rows::membership_filter( - candidates.iter().map(|(labels, _)| identity(labels)), - values, - |(labels, _)| identity(labels), - ); - if !missing.is_empty() && matches!(completeness, CandidateCompleteness::Certified { .. }) { - return Err(miss("certified membership key has no authoritative value")); + let right_key = |labels: &Labels| { + keys.iter() + .map(|(_, right)| labels.get(right).cloned().unwrap_or_default()) + .collect::>() + }; + let available = values + .iter() + .map(|(labels, _)| left_key(labels)) + .collect::>(); + let missing = candidates + .iter() + .map(|(labels, _)| right_key(labels)) + .filter(|key| !available.contains(key)) + .collect::>(); + let selected = native_values::semi_join(values, &candidates, &left_key, &right_key)?; + if !missing.is_empty() && matches!(completeness, Some(CandidateCompleteness::Certified { .. })) + { + return Err(miss("certified pruning key has no authoritative value")); } let warning = match completeness { - CandidateCompleteness::Certified { .. } => None, - CandidateCompleteness::BestEffort { guarantee } => Some(match guarantee { + None | Some(CandidateCompleteness::Certified { .. }) => None, + Some(CandidateCompleteness::BestEffort { guarantee }) => Some(match guarantee { Some(guarantee) => format!( "ASAP membership pruning is approximate: {:?}", guarantee.metric @@ -506,13 +538,15 @@ 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 { - asap_physical_operators::rows::grouped_topk( - values, - usize::try_from(k).unwrap_or(usize::MAX), - |(labels, _)| grouping_key(labels, grouping), - |(_, value)| *value, + native_values::limit( + native_values::sort(values, grouping, true).unwrap(), + grouping, + k, + 0, ) + .unwrap() } fn binary( @@ -1070,8 +1104,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, @@ -1154,12 +1202,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( @@ -1205,20 +1254,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, @@ -1281,23 +1369,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..951fec4c2 --- /dev/null +++ b/data_plane/src/query_engines/asap_query_engine/logical_dag/native_values.rs @@ -0,0 +1,137 @@ +//! 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 context() -> Result { + dag::RunContext::new( + dag::Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + dag::Limits::default(), + ) + .map_err(|e| miss(e.to_string())) +} +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(|e| miss(e.to_string())) +} +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, +) -> 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(|e| miss(e.to_string()))?; + let result = batch_execution::evaluate_batch(batch, vec![op], context()?) + .map_err(|e| miss(e.to_string()))?; + output(values, result) +} +pub(super) fn limit( + values: Vector, + grouping: &Grouping, + n: u64, + offset: u64, +) -> Result { + let batch = ranked_batch(&values, grouping)?; + let op = Operator::limit(batch.schema().clone(), n, offset, vec![1]) + .map_err(|e| miss(e.to_string()))?; + let result = batch_execution::evaluate_batch(batch, vec![op], context()?) + .map_err(|e| miss(e.to_string()))?; + output(values, result) +} +pub(super) fn semi_join( + values: Vector, + candidates: &Vector, + left_key: &impl Fn(&Labels) -> Vec, + right_key: &impl Fn(&Labels) -> Vec, +) -> Result { + let schema = schema(&[("index", DataType::Int64), ("key", DataType::Utf8)]); + let batch = |rows: &Vector, identity: &dyn Fn(&Labels) -> Vec| { + Batch::try_new( + schema.clone(), + rows.iter() + .enumerate() + .map(|(i, (labels, _))| vec![Value::Int64(i as i64), key(&identity(labels))]) + .collect(), + ) + .map_err(|e| miss(e.to_string())) + }; + let op = Operator::semi_join(schema.clone(), schema.clone(), vec![(1, 1)]) + .map_err(|e| miss(e.to_string()))?; + let result = batch_execution::evaluate_inputs( + vec![batch(&values, left_key)?, batch(candidates, right_key)?], + op, + context()?, + ) + .map_err(|e| miss(e.to_string()))?; + output(values, result) +} 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 aeeceebc7..658f9e9c3 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 @@ -305,7 +305,6 @@ impl QueryNodeRuntime for PhysicalQueryRuntime<'_> { }) } QueryPlanNode::Logical { .. } - | QueryPlanNode::MembershipFilter { .. } | QueryPlanNode::Relational { .. } | QueryPlanNode::ExternalExact { .. } | QueryPlanNode::RelationalJoin { .. } => Err(PhysicalNodeError::Fallback( 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 02c2adac6..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,11 +223,10 @@ 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. - SummaryExpr::MembershipFilter { .. } => Err(ExecError::NotYetSupported("MembershipFilter")), SummaryExpr::BinaryOp { .. } => Err(ExecError::NotYetSupported("BinaryOp")), SummaryExpr::ValueOperation { .. } => Err(ExecError::NotYetSupported("ValueOperation")), SummaryExpr::SummarySubtract { .. } => Err(ExecError::NotYetSupported("SummarySubtract")), diff --git a/docs/design_docs/physical-operators.md b/docs/design_docs/physical-operators.md index 02d321a19..0ef951614 100644 --- a/docs/design_docs/physical-operators.md +++ b/docs/design_docs/physical-operators.md @@ -1,80 +1,34 @@ -# Shared physical operators and DAG execution +# Consuming Planner physical operators -## Decision +## Ownership and contract -ASAP owns an independent physical operator library and DAG runtime. Precompute -and query engines bind inputs and consume outputs from the same library. The -operator defines computation; the engine supplies ingestion time or query time, -window boundaries, storage access and publication. There is no second execution -algorithm selected by phase. +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-operators.md). -This foundation precedes the precompute integration in #763 and query integration -in #765. Its native operators can execute independently of either backend engine. -Engine integration must use these operators for computation, rather than merely -using the shared scheduler around a second implementation. +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. -## Execution contract +Both engines 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. -An immutable plan describes typed nodes and dependency edges. Each execution -creates its own operator state. One producer may have multiple consumers; the -producer executes once in that run and sends the same outputs to all consumers. -Separate runs, query evaluation times and ingestion windows do not share mutable -state. Request-local caching of intermediate results is scoped to execution. - -The runtime validates dependencies, schemas, arity and cycles before sources -start. Each consumer advances independently. Bounded queues apply backpressure; -dropping one consumer does not cancel other consumers. Whole-run cancellation -wakes readers and releases queued work as streams are polled or dropped. - -Execution runs on the caller's worker without an internal thread pool. Active -streams are worker-local. Deployments poll all consumers concurrently. The byte -budget accounts for retained outputs and native operator state, including outputs -held after queue eviction. It is not an RSS limit: source-owned data, temporary -allocation peaks and allocator overhead remain outside that estimate. Blocking -operators currently have no spill implementation. - -## Operator coverage - -Native operations include scalar sources, typed Project and Filter, arithmetic -and boolean expressions, exact grouped aggregation, semi-join, grouped Sort and -Limit, Union, vector-to-scalar conversion, and summary construction, merge and -readout. Grouped TopK composes Sort and Limit within each group; candidate -completeness is an earlier pruning obligation. - -Values retain Planner types and nullability. Native summary batches currently -support exact Sum/Count/Min/Max/Rate/Increase, KLL, DDSketch and HLL. Other available -low-level kernels do not imply native batch bindings. Unsupported expressions, -state families and parameters must be rejected during binding, without an -implicit external fallback. The installed engine adapters and remaining gaps -are tracked in the query DAG design's unified coverage table. - -Deployments provide explicit storage or ingestion source frontiers. A supplied -batch source is not a backend raw Scan implementation. Local backend raw Scan -is deferred; a raw-only library test does not establish that deployment capability. - -## DataFusion reuse vs independent implementation - -| Decision dimension | Reuse DataFusion | Independent ASAP implementation | -| --- | --- | --- | -| General computation | Reuse mature Arrow operators and expression execution | Implement and test the supported Planner vocabulary explicitly | -| Shared DAG producer | Shared plan references need an explicit execution-sharing and buffering policy | One producer and independent consumer cursors are part of the runtime contract | -| Summary lifecycle | Add custom summary state operators to the framework | Summary construction, merge and readout are native capabilities | -| Engine reuse | Adapt both engines to DataFusion's execution model | Both engines bind the same ASAP interfaces | -| Engineering cost | Less generic operator work; integration and semantic adaptation remain | More operator, typing, scheduling and resource-accounting responsibility | - -DataFusion is a design reference, not this library's execution dependency. This -choice does not claim that DataFusion cannot express shared dependencies. ASAP -chooses direct ownership of execution sharing and summary-state semantics across -both engines. Mathematical sketch kernels remain reusable implementation details. +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 -Independent tests must execute shared-producer diamonds without duplicated work -or deadlock, exercise slow and dropped consumers, propagate cancellation and -errors, retain memory accounting, and isolate separate executions. Operator tests -must cover types, nulls, grouped limits, state compatibility and unsupported -bindings. The same summary pipeline must run at ingestion time and query time. +Binding must reject an unsupported operation, expression, family, parameter or +schema before execution. 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. -#763 and #765 add deployment acceptance for source binding, window and revision -scope, durable publication and query output adaptation. External exact forwarding -does not count as evidence that a local operator was implemented. +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. From 16cc4d5770fd05c8a645bcfcbf31608f12549a96 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 22:54:45 +0000 Subject: [PATCH 10/15] build: pin Planner physical library with relational and window operators --- Cargo.lock | 14 +++++++------- Cargo.toml | 12 ++++++------ docs/design_docs/physical-operators.md | 15 ++++++++++++++- 3 files changed, 27 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3a765ca52..cfb94d286 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=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" 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=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" 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=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" dependencies = [ "asap-sql-function-catalog", "asap-types", @@ -396,7 +396,7 @@ dependencies = [ [[package]] name = "asap-physical-operators" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" dependencies = [ "asap-types", "asap_sketch_codec", @@ -416,12 +416,12 @@ dependencies = [ [[package]] name = "asap-sql-function-catalog" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" [[package]] name = "asap-types" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" dependencies = [ "serde", "serde_json", @@ -442,7 +442,7 @@ dependencies = [ [[package]] name = "asap_sketch_codec" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=14fceed73f0ea35b5d5b275e4920f35a34233a9c#14fceed73f0ea35b5d5b275e4920f35a34233a9c" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" dependencies = [ "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=026cd18c7b8c23ae6c46d4d683151ba562b8cd3a)", "prost", diff --git a/Cargo.toml b/Cargo.toml index a4fd489c4..8334ce379 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,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 = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } -asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } -asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } -asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } +planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } +asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } +asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } +asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } # Shared external deps (used by 2+ crates) serde = { version = "1.0", features = ["derive"] } @@ -37,8 +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 = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } -asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "14fceed73f0ea35b5d5b275e4920f35a34233a9c" } +asap-physical-operators = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } +asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } asap_types = { path = "crates/asap_types" } asap_otel_proto = { path = "crates/asap_otel_proto" } indexmap = { version = "2.0", features = ["serde"] } diff --git a/docs/design_docs/physical-operators.md b/docs/design_docs/physical-operators.md index 0ef951614..aea8fff82 100644 --- a/docs/design_docs/physical-operators.md +++ b/docs/design_docs/physical-operators.md @@ -11,7 +11,7 @@ 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. -Both engines use the independent ASAP DAG runtime. Deployment code binds input +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. @@ -32,3 +32,16 @@ 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 #770 consumes the pinned library; +#763 integrates ingestion DAG execution and #765 integrates query DAG execution. +The remaining backend stack builds on those integrations. #759 carries the +full-workload acceptance suite; its performance results must be reported +separately from library and process 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. From 72d751c61cc8b8173cbd29cc67e07aead3095d92 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 23:04:54 +0000 Subject: [PATCH 11/15] build: pin Planner exact-state finalization contract --- Cargo.lock | 14 +++++++------- Cargo.toml | 12 ++++++------ 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index cfb94d286..fb1c6ad93 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=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" 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=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" 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=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" dependencies = [ "asap-sql-function-catalog", "asap-types", @@ -396,7 +396,7 @@ dependencies = [ [[package]] name = "asap-physical-operators" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" dependencies = [ "asap-types", "asap_sketch_codec", @@ -416,12 +416,12 @@ dependencies = [ [[package]] name = "asap-sql-function-catalog" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" [[package]] name = "asap-types" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" dependencies = [ "serde", "serde_json", @@ -442,7 +442,7 @@ dependencies = [ [[package]] name = "asap_sketch_codec" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=46328f0e50c1450bc2385a4fab32c9c5631f4e0e#46328f0e50c1450bc2385a4fab32c9c5631f4e0e" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" dependencies = [ "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=026cd18c7b8c23ae6c46d4d683151ba562b8cd3a)", "prost", diff --git a/Cargo.toml b/Cargo.toml index 8334ce379..1eebd9cf9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,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 = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } -asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } -asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } -asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } +planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } +asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } +asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } +asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } # Shared external deps (used by 2+ crates) serde = { version = "1.0", features = ["derive"] } @@ -37,8 +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 = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } -asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "46328f0e50c1450bc2385a4fab32c9c5631f4e0e" } +asap-physical-operators = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } +asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } asap_types = { path = "crates/asap_types" } asap_otel_proto = { path = "crates/asap_otel_proto" } indexmap = { version = "2.0", features = ["serde"] } From 2020d3d889a203ca6449290d6f4ba7089f8026b6 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 23:05:51 +0000 Subject: [PATCH 12/15] test: keep independent operator tests in Planner workspace --- scripts/e2e.sh | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/scripts/e2e.sh b/scripts/e2e.sh index 19878aa6a..b495ade61 100755 --- a/scripts/e2e.sh +++ b/scripts/e2e.sh @@ -81,9 +81,8 @@ contracts() { say "contracts: shared policy and routing types" rust_test asap_types - CURRENT_STAGE="contracts/asap-physical-operators" - say "contracts: shared physical kernels and deployment-independent execution" - rust_test asap-physical-operators + # 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" From 3826a4aac39a06b6673fee7e0338b4312c0f1a43 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 23:17:14 +0000 Subject: [PATCH 13/15] test: assert grouped Sort and Limit in level-one plan acceptance --- control_plane/tests/issue754_level1.rs | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/control_plane/tests/issue754_level1.rs b/control_plane/tests/issue754_level1.rs index da5c9f8df..4b1f84851 100644 --- a/control_plane/tests/issue754_level1.rs +++ b/control_plane/tests/issue754_level1.rs @@ -92,7 +92,7 @@ fn expected_plan(name: &str) -> ExpectedPlan { family: Some(ExpectedFamily::Exact(ExactKind::Rate)), partitioning: "per_entity", readout: "rate", - root_operation: Some("top_k_selection"), + root_operation: Some("limit"), }, "quantile-ratio" => ExpectedPlan { family: None, @@ -209,7 +209,8 @@ fn assert_selected_plan(name: &str, plan: &CompiledPhysicalPlan) -> Option Option Date: Wed, 23 Sep 2026 23:28:59 +0000 Subject: [PATCH 14/15] test: wait for reactivated series publication before restart --- .../src/storage_engines/sketch_db/index/mod.rs | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) 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 fb9105cef..5dce3059b 100644 --- a/data_plane/src/storage_engines/sketch_db/index/mod.rs +++ b/data_plane/src/storage_engines/sketch_db/index/mod.rs @@ -5202,7 +5202,6 @@ mod tests { || !persistence.manifest.live_parts().is_empty(), Duration::from_secs(5) )); - let old_parts = persistence.manifest.live_parts().len(); store.remove_instance(old_sid).unwrap(); assert!(resolver .resolve_with_reactivation("metric", "group", "family", |sid| store @@ -5227,7 +5226,20 @@ mod tests { ); } assert!(wait_until( - || persistence.manifest.live_parts().len() > old_parts, + || { + // Old-series epochs may still publish after reactivation. Wait + // for this series, not an unrelated increase in part count. + persistence.manifest.live_parts().iter().any(|part| { + let path = + persistence::part::part_dir_path(&persistence.parts_root, part.part_id); + persistence::part::PartReader::open(&path).is_ok_and(|reader| { + reader + .index_records() + .iter() + .any(|row| row.agg_id == new_sid && row.start_ts < 90_000) + }) + }) + }, Duration::from_secs(5) )); persistence.shutdown(); From 240f12e5f537b3d09676f80a366fb7c3c767ee30 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 23 Sep 2026 23:38:28 +0000 Subject: [PATCH 15/15] fix: pin Planner temporal ranking composition --- Cargo.lock | 14 +++++++------- Cargo.toml | 12 ++++++------ 2 files changed, 13 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index fb1c6ad93..1e6eed93b 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=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=645cacb04a451e25147f0d23aaddd05da16b4746#645cacb04a451e25147f0d23aaddd05da16b4746" 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=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=645cacb04a451e25147f0d23aaddd05da16b4746#645cacb04a451e25147f0d23aaddd05da16b4746" 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=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=645cacb04a451e25147f0d23aaddd05da16b4746#645cacb04a451e25147f0d23aaddd05da16b4746" dependencies = [ "asap-sql-function-catalog", "asap-types", @@ -396,7 +396,7 @@ dependencies = [ [[package]] name = "asap-physical-operators" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=645cacb04a451e25147f0d23aaddd05da16b4746#645cacb04a451e25147f0d23aaddd05da16b4746" dependencies = [ "asap-types", "asap_sketch_codec", @@ -416,12 +416,12 @@ dependencies = [ [[package]] name = "asap-sql-function-catalog" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=645cacb04a451e25147f0d23aaddd05da16b4746#645cacb04a451e25147f0d23aaddd05da16b4746" [[package]] name = "asap-types" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=645cacb04a451e25147f0d23aaddd05da16b4746#645cacb04a451e25147f0d23aaddd05da16b4746" dependencies = [ "serde", "serde_json", @@ -442,7 +442,7 @@ dependencies = [ [[package]] name = "asap_sketch_codec" version = "0.1.0" -source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=93a0fa2b60219c624fe74a0218ff4a76cd97ea78#93a0fa2b60219c624fe74a0218ff4a76cd97ea78" +source = "git+https://github.com/ProjectASAP/ASAPPlanner?rev=645cacb04a451e25147f0d23aaddd05da16b4746#645cacb04a451e25147f0d23aaddd05da16b4746" dependencies = [ "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=026cd18c7b8c23ae6c46d4d683151ba562b8cd3a)", "prost", diff --git a/Cargo.toml b/Cargo.toml index 1eebd9cf9..f94bdb8c2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,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 = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } -asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } -asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } -asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } +planner-types = { package = "asap-types", git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "645cacb04a451e25147f0d23aaddd05da16b4746" } +asap-aware-mapping = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "645cacb04a451e25147f0d23aaddd05da16b4746" } +asap-frontend-promql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "645cacb04a451e25147f0d23aaddd05da16b4746" } +asap-frontend-sql = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "645cacb04a451e25147f0d23aaddd05da16b4746" } # Shared external deps (used by 2+ crates) serde = { version = "1.0", features = ["derive"] } @@ -37,8 +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 = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } -asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "93a0fa2b60219c624fe74a0218ff4a76cd97ea78" } +asap-physical-operators = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "645cacb04a451e25147f0d23aaddd05da16b4746" } +asap_sketch_codec = { git = "https://github.com/ProjectASAP/ASAPPlanner", rev = "645cacb04a451e25147f0d23aaddd05da16b4746" } asap_types = { path = "crates/asap_types" } asap_otel_proto = { path = "crates/asap_otel_proto" } indexmap = { version = "2.0", features = ["serde"] }