diff --git a/Cargo.lock b/Cargo.lock index d28e1bc5..f54e6055 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -162,7 +162,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64", + "base64 0.22.1", "chrono", "comfy-table", "half", @@ -309,7 +309,7 @@ version = "0.1.0" dependencies = [ "asap-frontend-promql", "asap-types", - "asap_sketchlib", + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib)", "serde", "serde_json", "thiserror 2.0.18", @@ -368,12 +368,36 @@ dependencies = [ "asap-aware-mapping", "asap-frontend-promql", "asap-frontend-sql", + "asap-physical-operators", "asap-types", - "asap_sketchlib", + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib)", + "futures", "serde_json", "tokio", ] +[[package]] +name = "asap-physical-operators" +version = "0.1.0" +dependencies = [ + "asap-aware-mapping", + "asap-frontend-promql", + "asap-types", + "asap_sketch_codec", + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", + "base64 0.21.7", + "bincode", + "futures", + "hex", + "prost", + "rmp-serde", + "serde", + "serde_json", + "thiserror 2.0.18", + "tracing", + "xxhash-rust", +] + [[package]] name = "asap-sql-function-catalog" version = "0.1.0" @@ -384,9 +408,36 @@ version = "0.1.0" dependencies = [ "serde", "serde_json", + "sha2", "thiserror 2.0.18", ] +[[package]] +name = "asap_sketch_codec" +version = "0.1.0" +dependencies = [ + "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", + "prost", +] + +[[package]] +name = "asap_sketchlib" +version = "0.3.0" +source = "git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369#5f03ccbd798ed5fec62bdd839bcb331123cab369" +dependencies = [ + "bincode", + "bytes", + "prost", + "rand 0.9.5", + "rmp-serde", + "serde", + "serde-big-array", + "serde_bytes", + "smallvec", + "twox-hash 2.1.2", + "xxhash-rust", +] + [[package]] name = "asap_sketchlib" version = "0.3.0" @@ -448,6 +499,12 @@ version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" +[[package]] +name = "base64" +version = "0.21.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d297deb1925b89f2ccc13d7635fa0714f12c87adce1c75356b39ca9b7178567" + [[package]] name = "base64" version = "0.22.1" @@ -942,7 +999,7 @@ checksum = "f52c4012648b34853e40a2c6bcaa8772f837831019b68aca384fb38436dba162" dependencies = [ "arrow", "arrow-buffer", - "base64", + "base64 0.22.1", "blake2", "blake3", "chrono", @@ -2168,7 +2225,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64", + "base64 0.22.1", "brotli", "bytes", "chrono", diff --git a/Cargo.toml b/Cargo.toml index 9d44070d..037a4d8b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,5 +1,7 @@ [workspace] members = [ + "crates/asap-physical-operators", + "crates/asap_sketch_codec", "crates/types", "crates/sql-function-catalog", "crates/asap-aware-mapping", diff --git a/crates/asap-aware-mapping/src/accuracy/composition.rs b/crates/asap-aware-mapping/src/accuracy/composition.rs index a284ab32..2fa1a11b 100644 --- a/crates/asap-aware-mapping/src/accuracy/composition.rs +++ b/crates/asap-aware-mapping/src/accuracy/composition.rs @@ -276,10 +276,10 @@ impl DefaultAccuracyModel { if inputs.len() != 2 || inputs .iter() - .any(|input| input.metric != ErrorMetric::RelativeValue) + .any(|input| input.metric != ErrorMetric::RelativeValue && !input.is_exact()) { return Err(unsupported( - "division needs exactly two RelativeValue guarantees".into(), + "division needs two relative-value or exact guarantees".into(), )); } let Some(numerator) = inputs[0].bound.evaluate() else { diff --git a/crates/asap-aware-mapping/src/cost_model.rs b/crates/asap-aware-mapping/src/cost_model.rs index 366f0289..b2c42953 100644 --- a/crates/asap-aware-mapping/src/cost_model.rs +++ b/crates/asap-aware-mapping/src/cost_model.rs @@ -876,6 +876,13 @@ pub trait CostModel { self.raw_query_recompute_cost(target) .map(|per_read| Cost(per_read.0 * expected_reads)) } + /// Physical feasibility evidence for a complete summary candidate. + /// `None` defers admission to physical/deployment compilation; `Some(false)` + /// excludes the candidate without changing its computation or parameters. + fn summary_support_evidence(&self, _summary: &SummaryNode) -> Option { + None + } + /// Which mixed exact/summary execution shapes the downstream runtime /// advertises (issue #171). Gates candidate *generation* in /// [`crate::exact_composition::ExactCompositionStrategy`]: a shape the diff --git a/crates/asap-aware-mapping/src/hll_confidence.rs b/crates/asap-aware-mapping/src/hll_confidence.rs new file mode 100644 index 00000000..d7abcd26 --- /dev/null +++ b/crates/asap-aware-mapping/src/hll_confidence.rs @@ -0,0 +1,208 @@ +//! Estimator-specific confidence for classic HLL's linear-counting branch. +//! +//! This is conditional on independent uniform bucket hashes and an enforced +//! upper bound on distinct items in the complete readout population (including +//! all merged panes). It is not an RSE-to-normal conversion or an ERP fit. + +use asap_types::post_asap::{ + BoundExpr, ErrorMetric, GuaranteeSource, ProbabilityExpr, ResultGuarantee, +}; + +/// A finite-population contract for `m * ln(m / zero_registers)` with the +/// classic HLL small-range switch. Hashing is assumed independent and uniform. +/// The deployment must establish the population bound; observations alone do +/// not establish it. Unsupported precisions/populations return no certificate. +#[derive(Debug, Clone, Copy)] +pub struct ClassicHllConfidence { + max_distinct: u32, + relative_error: f64, +} + +impl ClassicHllConfidence { + pub fn new(max_distinct: u32, relative_error: f64) -> Option { + (max_distinct > 0 + && max_distinct <= 4096 + && relative_error.is_finite() + && (1e-6..1.0).contains(&relative_error)) + .then_some(Self { + max_distinct, + relative_error, + }) + } + + pub fn guarantee(&self, precision: u8) -> Option { + let delta = self.failure_probability(precision)?; + Some(ResultGuarantee { + metric: ErrorMetric::Cardinality, + bound: BoundExpr::Constant { + value: self.relative_error, + }, + failure_probability: ProbabilityExpr::Constant { value: delta }, + provenance: vec![GuaranteeSource::SketchReadout { + algorithm: "Hll".into(), + contract: "classic_hll_linear_counting_collision_bound_v1".into(), + params: serde_json::json!({"precision": precision, + "max_distinct": self.max_distinct, "relative_error": self.relative_error, + "hash_assumption": "independent_uniform_buckets", + "population_scope": "complete_readout_including_merged_panes"}), + query: "Cardinality".into(), + }], + }) + } + + pub fn precision(&self, delta: f64) -> Option { + if !delta.is_finite() || !(0.0..1.0).contains(&delta) || delta == 0.0 { + return None; + } + (4..=18).find(|&p| self.failure_probability(p).is_some_and(|d| d <= delta)) + } + + /// Finite bound, not an asymptotic RSE fit. With N distinct hashes and K + /// occupied buckets, C=N-K collision arrivals satisfy + /// P(C>=t) <= lambda^t/t!, lambda=N(N-1)/(2m): each arrival's conditional + /// collision probability is at most (i-1)/m, and a union bound over t + /// arrivals is bounded by the t-th power of their sum divided by t!. + /// + /// N<=m/2 makes the classic raw estimate <=2*alpha_m*m<2.5m, + /// so the small-range switch always uses L=-m*ln(1-K/m). Then + /// K<=L<=N + N^2/(2(m-N)). The latter bounds overestimation + /// deterministically; underestimation implies C>epsilon*N. + /// We maximize the collision bound over EVERY integer N in the contract, + /// not just its upper endpoint (small-cardinality tails matter). + fn failure_probability(&self, precision: u8) -> Option { + if !(4..=18).contains(&precision) { + return None; + } + let m = f64::from(1u32 << precision); + let max_n = f64::from(self.max_distinct); + // Reserve numerical slack; do not certify sub-floating-point error. + let eps = self.relative_error - 1e-8; + if max_n > m / 2.0 || max_n / (2.0 * (m - max_n)) > eps { + return None; + } + let mut log_factorial = vec![0.0; self.max_distinct as usize + 1]; + for i in 1..log_factorial.len() { + log_factorial[i] = log_factorial[i - 1] + (i as f64).ln(); + } + let mut worst = 0.0_f64; + for n in 2..=self.max_distinct { + let nf = f64::from(n); + // Including a boundary collision event is conservative. + let t = ((eps * nf).floor() as usize + 1).min(n as usize); + let lambda = nf * (nf - 1.0) / (2.0 * m); + let log_tail = (t as f64) * lambda.ln() - log_factorial[t]; + worst = worst.max(log_tail.min(0.0).exp()); + } + // Never return a spurious zero from underflow or numeric cancellation. + Some((worst * (1.0 + 1e-10) + 1e-12).min(1.0)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A supported estimator contract supplies a probability, unlike generic HLL RSE. + #[test] + fn bounded_classic_hll_has_a_feasible_confidence_target() { + let model = ClassicHllConfidence::new(128, 0.05).unwrap(); + let precision = model.precision(0.01).expect("finite confidence-sized HLL"); + let guarantee = model.guarantee(precision).unwrap(); + assert!(!guarantee.has_unknown()); + assert!(guarantee.failure_probability.evaluate().unwrap() <= 0.01); + assert_eq!(guarantee.bound.evaluate(), Some(0.05)); + } + /// Tighter confidence must increase precision or explicitly become unavailable. + #[test] + fn sizing_and_domain_limits_are_consistent() { + let model = ClassicHllConfidence::new(128, 0.05).unwrap(); + assert!(model.precision(0.001).unwrap() > model.precision(0.01).unwrap()); + assert!(model.precision(1e-12).is_none()); + assert!(model.precision(0.0).is_none()); + assert!(model.precision(f64::NAN).is_none()); + assert!(model.guarantee(3).is_none()); + assert!(model.guarantee(19).is_none()); + assert!(model.guarantee(7).is_none()); + for (n, e) in [(0, 0.05), (4097, 0.05), (128, 0.0), (128, f64::NAN)] { + assert!(ClassicHllConfidence::new(n, e).is_none()); + } + } + + /// Exact occupancy probabilities independently check both tails for every N. + #[test] + fn probability_bound_dominates_exact_occupancy_distribution() { + for precision in 4..=10 { + let m = 1usize << precision; + let max_n = 64.min(m / 2); + for eps in [0.05, 0.2, 0.6] { + let model = ClassicHllConfidence::new(max_n as u32, eps).unwrap(); + let Some(bound) = model.failure_probability(precision) else { + continue; + }; + let mut occupancy = vec![0.0; max_n + 1]; + occupancy[0] = 1.0; + for n in 1..=max_n { + let mut next = vec![0.0; max_n + 1]; + for k in 0..n { + next[k] += occupancy[k] * k as f64 / m as f64; + next[k + 1] += occupancy[k] * (m - k) as f64 / m as f64; + } + occupancy = next; + let actual: f64 = occupancy + .iter() + .enumerate() + .filter_map(|(k, &prob)| { + let estimate = -(m as f64) * (-(k as f64) / (m as f64)).ln_1p(); + ((estimate - n as f64).abs() > eps * n as f64).then_some(prob) + }) + .sum(); + assert!( + actual <= bound + 1e-12, + "p={precision} n={n} eps={eps}: {actual}>{bound}" + ); + } + } + } + } + /// The model's readout formula matches the actual classic estimator after merge. + #[test] + fn native_classic_estimator_and_merged_registers_use_the_same_contract() { + use asap_sketchlib::sketches::hll::{Classic, HyperLogLogP16}; + let model = ClassicHllConfidence::new(128, 0.05).unwrap(); + assert!( + model + .guarantee(16) + .unwrap() + .failure_probability + .evaluate() + .unwrap() + < 0.01 + ); + let mut single = HyperLogLogP16::::new(); + let mut left = HyperLogLogP16::::new(); + let mut right = HyperLogLogP16::::new(); + for n in 0..128u64 { + // SplitMix64 supplies deterministic test hashes, not a proof of randomness. + let mut h = n.wrapping_add(0x9e3779b97f4a7c15); + h = (h ^ (h >> 30)).wrapping_mul(0xbf58476d1ce4e5b9); + h = (h ^ (h >> 27)).wrapping_mul(0x94d049bb133111eb); + h ^= h >> 31; + single.insert_with_hash(h); + if n % 2 == 0 { + left.insert_with_hash(h); + } else { + right.insert_with_hash(h); + } + } + left.merge(&right); + assert_eq!(single.registers_as_slice(), left.registers_as_slice()); + let zeroes = left + .registers_as_slice() + .iter() + .filter(|&&r| r == 0) + .count(); + let expected = (65536.0 * (65536.0 / zeroes as f64).ln()) as usize; + assert_eq!(left.estimate(), expected); + assert!((expected as f64 - 128.0).abs() / 128.0 <= 0.05); + } +} diff --git a/crates/asap-aware-mapping/src/lib.rs b/crates/asap-aware-mapping/src/lib.rs index 420d5a2f..b9356220 100644 --- a/crates/asap-aware-mapping/src/lib.rs +++ b/crates/asap-aware-mapping/src/lib.rs @@ -161,6 +161,7 @@ pub mod exact_composition; pub mod explanation; mod function_rules; pub mod grouping; +pub mod hll_confidence; pub mod pane_sharing; pub mod physical_handoff_cost; pub mod physical_operator_statistics; diff --git a/crates/asap-aware-mapping/src/maintained_population.rs b/crates/asap-aware-mapping/src/maintained_population.rs index 687a0f7f..94714602 100644 --- a/crates/asap-aware-mapping/src/maintained_population.rs +++ b/crates/asap-aware-mapping/src/maintained_population.rs @@ -50,6 +50,7 @@ fn recognize(root: &QueryExpr) -> Option<(MaintainedPopulation, PopulationReadou AggIntent::Quantile { q, col, .. } if q.is_finite() => { (*col, PopulationReadout::Quantile { q: *q }) } + AggIntent::TopK { k, .. } => (None, PopulationReadout::TopK { k: *k }), AggIntent::Sum { col } => (*col, PopulationReadout::Sum), AggIntent::Count { .. } => (None, PopulationReadout::Count), AggIntent::Avg { col } => (*col, PopulationReadout::Average), @@ -142,8 +143,11 @@ fn recognize(root: &QueryExpr) -> Option<(MaintainedPopulation, PopulationReadou if value_column.is_some_and(|c| schema.columns.get(c).is_none_or(|c| c.name != "value")) { return None; } - // Open time-series schemas distinguish instant PromQL populations from table rows. - if metric.is_empty() || schema.closed || schema.time_index.is_none() { + // PromQL can retain open labels or resolve them into a complete identity column. + if metric.is_empty() + || (schema.closed && !schema.has_promql_series_identity()) + || schema.time_index.is_none() + { return None; } let label = |col: usize| -> Option { diff --git a/crates/asap-aware-mapping/src/replacement.rs b/crates/asap-aware-mapping/src/replacement.rs index e2ce61d1..e40b0010 100644 --- a/crates/asap-aware-mapping/src/replacement.rs +++ b/crates/asap-aware-mapping/src/replacement.rs @@ -520,15 +520,15 @@ impl ReplacementSubDAG { ) } - /// Runtime support for this candidate. Summary implementations remain - /// unknown until backend binding; a pure logical rewrite needs no new - /// physical operator. `Some(false)` disproves mixed-operation support. + /// Physical feasibility evidence for this candidate. A pure logical + /// rewrite needs no new operator. Unknown support is checked during + /// physical/deployment compilation; explicit rejection prevents selection. pub fn runtime_support_evidence(&self, cost_model: &dyn CostModel) -> Option { match &self.replacement { Replacement::ExactComposition(composition) => { cost_model.value_operation_support_evidence(&composition.op, composition.placement) } - Replacement::Summary(_) => None, + Replacement::Summary(node) => cost_model.summary_support_evidence(node), Replacement::Rewrite(_) => Some(true), } } @@ -834,6 +834,11 @@ pub(crate) fn realizations_for_intent( )), ], AccuracyTarget::Exact => vec![exact_realization(intent)], + _ if matches!(intent, AggIntent::Count { .. }) => { + let mut candidates = sketch_realizations(intent, accuracy, cost_model); + candidates.push(exact_realization(intent)); + candidates + } _ => sketch_realizations(intent, accuracy, cost_model), }, @@ -1277,6 +1282,173 @@ impl<'a> SketchAlgorithmStrategy<'a> { } } + /// Preserve the canonical Sort/Limit representation while exploring heap + /// realizations of an instant-vector ranking under the caller's target. + /// The input must carry the complete dynamic series identity. This never + /// treats a range of historical samples as the instant vector. + pub fn current_series_topk_candidates( + &self, + root: &Rc, + accuracy: &AccuracyTarget, + ) -> Proposals { + let QueryExpr::Limit { + n, + offset: 0, + child, + } = root.as_ref() + else { + return Proposals::default(); + }; + let QueryExpr::Sort { + keys, + partition_by, + child, + } = child.as_ref() + else { + return Proposals::default(); + }; + let [key] = keys.as_slice() else { + return Proposals::default(); + }; + let QueryExpr::Column(value) = key.expr else { + return Proposals::default(); + }; + let Ok(schema) = child.output_schema() else { + return Proposals::default(); + }; + if key.ascending + || key.nulls_first + || partition_by.is_without() + || !schema.has_promql_series_identity() + || !schema + .columns + .get(value) + .is_some_and(|column| column.name == "value") + || !is_current_series_source(child) + { + return Proposals::default(); + } + let ranked = Rc::new(QueryExpr::Aggregate { + reduction: Reduction::Reduce(partition_by.clone()), + measures: vec![AggIntent::TopK { + k: *n, + accuracy: accuracy.clone(), + }], + output_names: vec![], + having: None, + child: Rc::clone(child), + }); + self.propose_with(&ranked, None) + } + + /// Fixed-window maintenance can finalize each series' counter state and + /// build a fresh heap or grouped Sum for that evaluation window. Deployment must provide + /// a complete, synchronized population and bind the matching window; this + /// candidate never incrementally adds one window's rates to another. + pub fn fixed_window_rate_candidates(&self, root: &Rc) -> Proposals { + fn place(node: &Rc) -> Option> { + let mut next = node.as_ref().clone(); + match &mut next.expr { + SummaryExpr::ValueOperation { + child, + operation: ValueOperation::FinalizeExactAccumulator, + timing, + } if matches!(&child.expr, SummaryExpr::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + reduction: Reduction::PerEntity, child: source, .. + } if matches!(&source.expr, SummaryExpr::KeepPreAsap(source) if matches!(source.as_ref(), QueryExpr::TimeRange { .. }))) => + { + *timing = ExecutionTiming::IngestionTime; + } + SummaryExpr::ValueOperation { child, .. } + | SummaryExpr::SummaryAgg { child, .. } => *child = place(child)?, + SummaryExpr::SummaryEstimate { summary_input, .. } => { + *summary_input = place(summary_input)? + } + _ => return None, + } + Some(Rc::new(next)) + } + let mut proposals = self.propose_with(root, None); + proposals.candidates.retain_mut(|candidate| { + let Replacement::Summary(node) = &candidate.replacement else { + return false; + }; + let Ok(dag) = asap_types::post_asap::compile_executable_dag(node) else { + return false; + }; + if !dag.nodes.iter().any(|node| match &node.payload { + asap_types::post_asap::ExecutableOperatorPayload::SummaryAgg { + family: SummaryFamilyType::Sketch(kind, _), + .. + } => matches!( + kind.algorithm(), + SketchAlgorithm::CmsWithHeap | SketchAlgorithm::CountSketchWithHeap + ), + asap_types::post_asap::ExecutableOperatorPayload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), + .. + } => true, + _ => false, + }) { + return false; + } + let Some(placed) = place(node) else { + return false; + }; + if asap_types::post_asap::compile_executable_dag(&placed).is_err() { + return false; + } + let Ok(placed) = finalize_query_candidate(placed, root) else { + return false; + }; + candidate.replacement = Replacement::Summary(placed); + candidate + .rationale + .push_str("; fixed-window precompute over complete per-series counter states"); + true + }); + proposals + } + + /// Retain grouped Sum after a per-series Rate readout as a query-time + /// candidate alongside its complete-window maintenance placement. + pub fn query_time_rate_aggregation_candidates(&self, root: &Rc) -> Proposals { + fn query_time(node: &Rc) -> Rc { + let mut next = node.as_ref().clone(); + match &mut next.expr { + SummaryExpr::ValueOperation { + child, + operation: ValueOperation::FinalizeExactAccumulator, + timing, + } if matches!( + &child.expr, + SummaryExpr::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + .. + } + ) => + { + *timing = ExecutionTiming::QueryTime; + } + SummaryExpr::ValueOperation { child, .. } + | SummaryExpr::SummaryAgg { child, .. } => *child = query_time(child), + _ => {} + } + Rc::new(next) + } + let mut proposals = self.fixed_window_rate_candidates(root); + proposals.candidates.retain_mut(|candidate| { + let Replacement::Summary(node) = &candidate.replacement else { return false }; + if !matches!(&node.expr, SummaryExpr::ValueOperation { child, operation: ValueOperation::FinalizeExactAccumulator, .. } + if matches!(&child.expr, SummaryExpr::SummaryAgg { family: SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), .. })) { return false; } + candidate.replacement = Replacement::Summary(query_time(node)); + candidate.rationale = "query-time grouped Sum over complete per-series Rate readouts".into(); + true + }); + proposals + } + pub(crate) fn from_planning_inputs(planning_inputs: CandidatePlanningInputs<'a>) -> Self { Self { planning_inputs } } @@ -1286,6 +1458,22 @@ impl<'a> SketchAlgorithmStrategy<'a> { /// differs — see [`realize_child_with`]). fn propose_with(&self, root: &Rc, intent_override: Option<&AggIntent>) -> Proposals { let mut proposals = Proposals::default(); + // A selected logical rewrite otherwise remains KeepPreAsap during DAG + // assembly. Also expose its concrete summary realization for selection. + if intent_override.is_none() { + if let Some(rewritten) = crate::rewrite::composed_aggregate_rewrite(root) { + if let Ok(node) = realize_child_with(&rewritten, self.planning_inputs, None) { + if !matches!(node.expr, SummaryExpr::KeepPreAsap(_)) { + proposals.candidates.push(ReplacementSubDAG { + replacement: Replacement::Summary(node), + strategy: "SketchAlgorithmStrategy", + provenance: ReplacementProvenance::SummaryRealization, + rationale: "realize a schema-preserving composition of temporal and grouped accumulators".into(), + }); + } + } + } + } if let Ok(Some(node)) = exact_topk_over_temporal_values(root, self.planning_inputs) { proposals.candidates.push(ReplacementSubDAG { replacement: Replacement::Summary(node), @@ -1525,7 +1713,7 @@ fn describe_realization(intent: &AggIntent, realization: &Realization) -> String match realization { Realization::Sketch(kind) => format!( "{} realizes as a {:?} sketch — one of summary_candidates' \ - alternatives for this intent (asap_aware_mapping::replacement::realizations_for_intent)", + candidates for this intent (asap_aware_mapping::replacement::realizations_for_intent)", describe_intent(intent), kind.algorithm() ), @@ -1638,11 +1826,7 @@ fn exact_topk_over_temporal_values( else { return Ok(None); }; - let [AggIntent::TopK { - k, - accuracy: AccuracyTarget::Exact, - }] = measures.as_slice() - else { + let [AggIntent::TopK { k, .. }] = measures.as_slice() else { return Ok(None); }; let QueryExpr::Aggregate { @@ -1665,7 +1849,7 @@ fn exact_topk_over_temporal_values( { return Ok(None); } - let values = finalize_exact_accumulator(values, child)?; + let values = finalize_query_candidate(values, child)?; let partition_by = reduction .group_keys() .ok_or(RealizationError::PhysicalRealization( @@ -1821,11 +2005,30 @@ fn realize_binary( return Ok(None); } let domains = lhs_domain.zip(rhs_domain).map(|(lhs, rhs)| [lhs, rhs]); + let has_mean = [lhs, rhs] + .iter() + .any(|expr| matches!(bindable_intent(expr), Some(AggIntent::Avg { .. }))); + if has_mean + && domains.as_ref().is_none_or(|domains| { + domains.iter().any(|domain| { + !(domain.lower.abs().max(domain.upper.abs()) * domain.max_samples as f64) + .is_finite() + }) + }) + { + return Ok(None); + } lhs_node = realize_ddsketch_quantile_operand(lhs, planning_inputs, &target)?; rhs_node = realize_ddsketch_quantile_operand(rhs, planning_inputs, &target)?; if let Some(domains) = domains.as_ref() { for (domain, node) in domains.iter().zip([&lhs_node, &rhs_node]) { if !ddsketch_quantile_alpha(node) + .or_else(|| { + node.guarantee + .as_ref() + .is_some_and(ResultGuarantee::is_exact) + .then_some(alpha) + }) .is_some_and(|alpha| domain.supports_ddsketch(alpha)) { return Ok(None); @@ -1900,16 +2103,20 @@ fn realize_binary( return Ok(None); } - lhs_node = finalize_exact_accumulator(lhs_node, lhs)?; - rhs_node = finalize_exact_accumulator(rhs_node, rhs)?; + lhs_node = finalize_query_candidate(lhs_node, lhs)?; + rhs_node = finalize_query_candidate(rhs_node, rhs)?; let has_ratio_domains = ratio_domains.is_some(); let guarantee = if matches!(op, BinaryOpKind::Arithmetic(ArithmeticOpKind::Div)) && direct_ddsketch_ratio && has_ratio_domains - && ddsketch_quantile_alpha(&lhs_node).is_some() - && ddsketch_quantile_alpha(&rhs_node).is_some() - { + && [&lhs_node, &rhs_node].iter().all(|node| { + ddsketch_quantile_alpha(node).is_some() + || node + .guarantee + .as_ref() + .is_some_and(ResultGuarantee::is_exact) + }) { [lhs_node.guarantee.clone(), rhs_node.guarantee.clone()] .into_iter() .collect::>>() @@ -1961,7 +2168,7 @@ fn realize_binary( /// Put an explicit read boundary between maintained exact state and a /// query-time value consumer. Approximate summaries must already carry a /// `SummaryEstimate`, so they deliberately do not pass this predicate. -fn finalize_exact_accumulator( +pub fn finalize_query_candidate( node: Rc, logical_output: &QueryExpr, ) -> Result, RealizationError> { @@ -2018,9 +2225,8 @@ fn is_promql_scalar(expr: &QueryExpr) -> bool { ) } -/// Both direct quantile operands inherit the workload target during PromQL -/// lowering. Reuse that one target for the ratio rather than interpreting it -/// as two independent error budgets. +/// Quantile operands inherit one workload target. A temporal mean is exact +/// on its checked finite domain and needs no approximation budget. fn shared_quantile_target(lhs: &QueryExpr, rhs: &QueryExpr) -> Option { let quantile_target = |expr: &QueryExpr| match bindable_intent(expr) { Some(AggIntent::Quantile { accuracy, q, .. }) @@ -2030,9 +2236,16 @@ fn shared_quantile_target(lhs: &QueryExpr, rhs: &QueryExpr) -> Option None, }; - let lhs = quantile_target(lhs)?; - let rhs = quantile_target(rhs)?; - (lhs == rhs).then_some(lhs) + match (quantile_target(lhs), quantile_target(rhs)) { + (Some(lhs), Some(rhs)) => (lhs == rhs).then_some(lhs), + (Some(target), None) if matches!(bindable_intent(rhs), Some(AggIntent::Avg { .. })) => { + Some(target) + } + (None, Some(target)) if matches!(bindable_intent(lhs), Some(AggIntent::Avg { .. })) => { + Some(target) + } + _ => None, + } } /// For `a / b`, two DDSketches with the same relative bound `alpha` produce @@ -2244,7 +2457,7 @@ pub(crate) fn construct_summary_with( child_target, allocation, )?; - if is_counter_weighted_topk(intent, child) { + if is_snapshot_weighted_topk(intent, child) { return finish_weighted_topk(candidate, expr, intent); } return Ok(candidate); @@ -2351,13 +2564,25 @@ fn finish_weighted_topk( Ok(result) } -fn is_counter_weighted_topk(intent: &AggIntent, child: &QueryExpr) -> bool { +fn is_current_series_source(child: &QueryExpr) -> bool { + let source = match child { + QueryExpr::TimeRange { child, .. } => child.as_ref(), + source => source, + }; + matches!(source, QueryExpr::Scan { + source: asap_types::pre_asap::Source::TimeSeries { .. }, schema, .. + } if schema.has_promql_series_identity()) +} + +fn is_snapshot_weighted_topk(intent: &AggIntent, child: &QueryExpr) -> bool { matches!(intent, AggIntent::TopK { .. }) - && matches!(child, + && (is_current_series_source(child) + || matches!(child, QueryExpr::Aggregate { measures, child, .. } - if matches!(measures.as_slice(), [AggIntent::Sum { .. }]) - && matches!(child.as_ref(), QueryExpr::Aggregate { measures, .. } - if matches!(measures.as_slice(), [AggIntent::Rate | AggIntent::Increase]))) + if matches!(measures.as_slice(), [AggIntent::Rate | AggIntent::Increase]) + || (matches!(measures.as_slice(), [AggIntent::Sum { .. }]) + && matches!(child.as_ref(), QueryExpr::Aggregate { measures, .. } + if matches!(measures.as_slice(), [AggIntent::Rate | AggIntent::Increase]))))) } /// Translate an [`Realization`] into the `(family, needs a @@ -2413,6 +2638,8 @@ type PhysicalSummaryInputRule = fn( /// `construct_summary_agg`. const PHYSICAL_SUMMARY_INPUT_RULES: &[PhysicalSummaryInputRule] = &[ realize_value_frequency_summary_input, + realize_counter_value_summary_input, + realize_current_series_summary_input, realize_keyed_additive_summary_input, ]; @@ -2568,10 +2795,10 @@ fn construct_summary_agg( SummaryFamilyType::Sketch(kind, _) if matches!(kind.algorithm(), SketchAlgorithm::CmsWithHeap | SketchAlgorithm::CountSketchWithHeap) ); - let rate_weighted = matches!(node, QueryExpr::Aggregate { child, .. } - if is_counter_weighted_topk(intent, child)); + let snapshot_weighted = matches!(node, QueryExpr::Aggregate { child, .. } + if is_snapshot_weighted_topk(intent, child)); let mut family = family; - let score_population = if rate_weighted { + let score_population = if snapshot_weighted { let bound = planning_inputs.evidence.topk_max_distinct_items(node); if bound.is_some_and(|n| n == 0 || n > (1u64 << 53)) { return Err(RealizationError::PhysicalRealization( @@ -2595,7 +2822,7 @@ fn construct_summary_agg( } else { None }; - let physical_reduction = if rate_weighted { + let physical_reduction = if snapshot_weighted { let QueryExpr::Aggregate { child, .. } = node else { unreachable!() }; @@ -2646,7 +2873,7 @@ fn construct_summary_agg( let state_idx = summary_col_index(&out_schema, &by, per_series); let readout_schema = if keyed_heap - && matches!(node, QueryExpr::Aggregate { child, .. } if is_counter_weighted_topk(intent, child)) + && matches!(node, QueryExpr::Aggregate { child, .. } if is_snapshot_weighted_topk(intent, child)) { keyed_heap_readout_schema(&input, node)? } else { @@ -2655,7 +2882,7 @@ fn construct_summary_agg( let summary_input = input.input; let query = estimate.then(|| { - if rate_weighted { + if snapshot_weighted { if let SummaryFamilyType::Sketch(kind, _) = &family { let capacity = match kind.params() { SketchParams::CmsWithHeap { heap_size, .. } @@ -2674,7 +2901,7 @@ fn construct_summary_agg( if keyed_heap { let mut state = state_schema.fields[state_idx].clone(); state.dtype = family.clone(); - let mut fields = if rate_weighted { + let mut fields = if snapshot_weighted { readout_schema.fields[..reduction.group_keys().map_or(0, |keys| keys.len())].to_vec() } else { Vec::new() @@ -2696,16 +2923,33 @@ fn construct_summary_agg( let bound_child = realize_child_with( &input.child, planning_inputs, - if rate_weighted { + if snapshot_weighted { Some(&AccuracyTarget::Exact) } else { child_target }, )?; - let bound_child = if rate_weighted { + let bound_child = if snapshot_weighted && is_current_series_source(&input.child) { + // Explicit snapshot selection prevents historical observations from + // becoming repeated weights in an instant-vector heap. + let root = Rc::new(node.clone()); + let population = crate::maintained_population::MaintainedPopulationStrategy::new( + std::slice::from_ref(&root), + ) + .candidate(&root) + .ok_or(RealizationError::PhysicalRealization( + "snapshot ranking requires a supported current-series population", + ))?; + let SummaryExpr::ValueOperation { child, .. } = &population.expr else { + return Err(RealizationError::PhysicalRealization( + "missing population readout", + )); + }; + Rc::clone(child) + } else if snapshot_weighted { // A fresh query-time summary consumes this evaluation's finalized rates. // Moving rate snapshots must never accumulate across evaluations. - finalize_exact_accumulator(bound_child, &input.child)? + finalize_query_candidate(bound_child, &input.child)? } else { let child = finalize_exact_accumulator_at( bound_child, @@ -2729,7 +2973,7 @@ fn construct_summary_agg( planning_inputs.evidence.estimator_contract(node), local_target, ); - let membership_query = if rate_weighted { + let membership_query = if snapshot_weighted { Some(readout(intent, &summary_input, planning_inputs.cost)) } else { query.clone() @@ -2744,7 +2988,7 @@ fn construct_summary_agg( allocation, )?; - if rate_weighted { + if snapshot_weighted { use asap_types::post_asap::{BoundExpr, ProbabilityExpr}; let target = accuracy_target(intent).expect("TopK target"); guarantee = if let Some(mut score) = @@ -2968,6 +3212,19 @@ fn ranking_score_index( logical: &QueryExpr, values: &SummarySchema, ) -> Result { + if is_current_series_source(logical) { + return values + .fields + .iter() + .position(|field| { + field.name == "value" + && field.dtype + == SummaryFamilyType::Plain(asap_types::pre_asap::DataType::Float64) + }) + .ok_or(RealizationError::PhysicalRealization( + "snapshot ranking requires the sample value column", + )); + } let QueryExpr::Aggregate { reduction, measures, @@ -3015,6 +3272,124 @@ fn ranking_score_index( Ok(index) } +/// Rebuild a heap from this evaluation's finalized per-series counter values. +/// The rate window is preserved; raw counter samples never become CMS weights. +fn realize_counter_value_summary_input( + intent: &AggIntent, + family: &SummaryFamilyType, + output_reduction: &Reduction, + child: &Rc, +) -> PhysicalSummaryInputRuleResult { + if !matches!(intent, AggIntent::TopK { .. }) + || !matches!(family, SummaryFamilyType::Sketch(kind, _) if matches!(kind.algorithm(), SketchAlgorithm::CmsWithHeap | SketchAlgorithm::CountSketchWithHeap)) + || !matches!(child.as_ref(), QueryExpr::Aggregate { reduction: Reduction::PerEntity, measures, .. } if matches!(measures.as_slice(), [AggIntent::Rate | AggIntent::Increase])) + { + return PhysicalSummaryInputRuleResult::NotApplicable; + } + let Ok(schema) = child.output_schema() else { + return PhysicalSummaryInputRuleResult::Unsupported( + "counter ranking needs a valid value schema", + ); + }; + if !schema.closed { + return PhysicalSummaryInputRuleResult::Unsupported( + "counter ranking needs the complete resolved series identity", + ); + } + let Reduction::Reduce(groups) = output_reduction else { + return PhysicalSummaryInputRuleResult::Unsupported( + "counter ranking requires explicit partitions", + ); + }; + if groups.is_without() { + return PhysicalSummaryInputRuleResult::Unsupported( + "counter ranking requires resolved partitions", + ); + } + // Retain the evaluation timestamp in each returned row. This sketch is a + // snapshot, not an additive history of successive rate evaluations. + let items = schema + .columns + .iter() + .enumerate() + .filter(|(index, column)| column.name != "value" && !groups.contains(index)) + .map(|(index, _)| schema_column_ref(child, index).map(SummaryInputExpr::Column)) + .collect::>>(); + let Some(items) = items.filter(|items| !items.is_empty()) else { + return PhysicalSummaryInputRuleResult::Unsupported("counter ranking has no item columns"); + }; + PhysicalSummaryInputRuleResult::Realized(PhysicalSummaryInput { + child: Rc::clone(child), + input: SummaryUpdate { + item: Some(SummaryInputExpr::Tuple(items)), + weight: SummaryInputExpr::Column(ColumnRef::SampleValue), + weight_domain: WeightDomain::NonNegative { + proof: NonNegativeWeightProof::ResetAwareCounterDerivative, + }, + }, + }) +} + +/// An instant-vector source has one current value per full series identity. +/// Rebuild the state for each evaluation; historical samples are not updates. +fn realize_current_series_summary_input( + intent: &AggIntent, + family: &SummaryFamilyType, + output_reduction: &Reduction, + child: &Rc, +) -> PhysicalSummaryInputRuleResult { + if !matches!(intent, AggIntent::TopK { .. }) || !is_current_series_source(child) { + return PhysicalSummaryInputRuleResult::NotApplicable; + } + let SummaryFamilyType::Sketch(kind, _) = family else { + return PhysicalSummaryInputRuleResult::NotApplicable; + }; + match kind.algorithm() { + SketchAlgorithm::CountSketchWithHeap => {} + SketchAlgorithm::CmsWithHeap => { + return PhysicalSummaryInputRuleResult::Unsupported( + "current sample values do not prove non-negative CMS weights", + ) + } + _ => return PhysicalSummaryInputRuleResult::NotApplicable, + } + let Reduction::Reduce(groups) = output_reduction else { + return PhysicalSummaryInputRuleResult::Unsupported( + "snapshot ranking requires explicit partitions", + ); + }; + if groups.is_without() { + return PhysicalSummaryInputRuleResult::Unsupported( + "snapshot ranking requires resolved partitions", + ); + } + let Ok(schema) = child.output_schema() else { + return PhysicalSummaryInputRuleResult::Unsupported( + "snapshot ranking requires a valid source schema", + ); + }; + let items = schema + .columns + .iter() + .enumerate() + .filter(|(index, column)| column.name != "value" && !groups.contains(index)) + .map(|(index, _)| schema_column_ref(child, index).map(SummaryInputExpr::Column)) + .collect::>>(); + let Some(items) = items.filter(|items| !items.is_empty()) else { + return PhysicalSummaryInputRuleResult::Unsupported( + "snapshot ranking has no item identity", + ); + }; + PhysicalSummaryInputRuleResult::Realized(PhysicalSummaryInput { + child: Rc::clone(child), + input: SummaryUpdate { + item: Some(SummaryInputExpr::Tuple(items)), + weight: SummaryInputExpr::Column(ColumnRef::SampleValue), + weight_domain: WeightDomain::UnknownOrSigned, + }, + }) +} + /// Realize the composite heavy-hitter realization for /// `TopK(Count GROUP BY key)`. The heap sketch consumes the raw keyed stream; /// it does not consume an independently materialized Count result. @@ -3665,6 +4040,256 @@ impl PlanSpace { } } +/// DAG candidates assembled from an unpriced search space. +/// This is an internal planning stage: callers must still validate lifecycle +/// requirements and compile supported physical operators before deployment. +/// The caller supplies a finite expansion budget; exceeding it is an error, +/// never a silently truncated inventory presented as exhaustive. +#[derive(Debug)] +pub struct CandidateDagInventory { + pub candidates: Vec)>>, + pub rejected_assemblies: Vec, +} + +type CandidateDagChoice<'a> = (Option<&'a ReplacementSubDAG>, Option>); + +impl PlanSpace { + pub fn enumerate_candidate_dags( + &self, + expansion_limit: usize, + ) -> Result, RealizationError> { + self.enumerate_candidate_roots(&self.roots, expansion_limit) + } + + /// Enumerate one workload root without expanding independent roots' choices. + /// Discovery and composition proofs still come from the shared workload + /// space. Deployment may price combinations lazily; this API does not rank + /// candidates or claim that independently cheapest roots minimize shared cost. + pub fn enumerate_candidate_dags_for_root( + &self, + id: &Id, + expansion_limit: usize, + ) -> Result, RealizationError> { + let roots = self + .roots + .iter() + .filter(|(candidate, _)| candidate == id) + .cloned() + .collect::>(); + if roots.len() != 1 { + return Err(RealizationError::PhysicalRealization( + "candidate enumeration requires one uniquely identified workload root", + )); + } + self.enumerate_candidate_roots(&roots, expansion_limit) + } + + fn enumerate_candidate_roots( + &self, + roots: &[(Id, Rc)], + expansion_limit: usize, + ) -> Result, RealizationError> { + let mut reachable = Vec::new(); + let mut nodes = HashMap::new(); + let mut counts = HashMap::new(); + for (_, root) in roots { + walk(root, &mut reachable, &mut nodes, &mut counts); + } + // Rewrites may introduce descendants absent from the original root. + let mut cursor = 0; + while cursor < reachable.len() { + let ptr = reachable[cursor]; + cursor += 1; + if let Some(group) = self.groups.get(&ptr) { + for candidate in &group.candidates { + if let Replacement::Rewrite(rewritten) = &candidate.replacement { + walk(rewritten, &mut reachable, &mut nodes, &mut counts); + } + } + } + } + let order = self + .order + .iter() + .copied() + .filter(|ptr| counts.contains_key(ptr)) + .collect::>(); + // Composition plans carry the proofs established during discovery. + // No cost ranking is consulted while expanding these choices. + let options: Vec>> = order + .iter() + .map(|ptr| { + let group = &self.groups[ptr]; + let mut choices = vec![(None, None)]; + for candidate in &group.candidates { + match &candidate.replacement { + Replacement::ExactComposition(operation) => { + for prepared in &self.composition_plans { + if prepared.target == *ptr + && prepared.operation.placement == operation.placement + && prepared.operation.op == operation.op + && Rc::ptr_eq( + &prepared.operation.child_target, + &operation.child_target, + ) + { + choices + .push((Some(candidate), Some(Rc::clone(&prepared.plan)))); + } + } + } + _ => choices.push((Some(candidate), None)), + } + } + choices + }) + .collect(); + let combinations = options + .iter() + .try_fold(1usize, |n, choices| n.checked_mul(choices.len())) + .filter(|n| *n <= expansion_limit) + .ok_or(RealizationError::PhysicalRealization( + "candidate expansion budget exceeded; no partial inventory returned", + ))?; + let mut inventory = CandidateDagInventory { + candidates: Vec::new(), + rejected_assemblies: Vec::new(), + }; + // Hash buckets avoid quadratic comparisons across a large workload + // inventory. Equality still decides deduplication, including collisions. + let mut seen = HashMap::>::new(); + for mut ordinal in 0..combinations { + let mut groups = HashMap::new(); + let mut assembled_nodes = HashMap::new(); + for (ptr, choices) in order.iter().zip(&options) { + let (chosen, prepared) = &choices[ordinal % choices.len()]; + ordinal /= choices.len(); + let group = &self.groups[ptr]; + if let Some(node) = prepared { + assembled_nodes.insert(*ptr, Rc::clone(node)); + } + groups.insert( + *ptr, + TargetSubDAGSelection { + target: &group.target, + consumer_count: group.consumer_count, + effective_consumer_count: group.consumer_count, + chosen: *chosen, + composition: None, + }, + ); + } + let assembly = GlobalSelection { + order: order.clone(), + groups, + assembled_nodes: RefCell::new(assembled_nodes), + }; + let roots = roots + .iter() + .map(|(id, root)| { + assembly + .assemble_target(root) + // Exposed query candidates return values. Internal assembly + // still retains accumulator states for sharing and storage. + .and_then(|node| finalize_query_candidate(node, root)) + .map(|node| (id.clone(), node)) + }) + .collect::, _>>(); + match roots { + Ok(roots) => { + let roots = asap_types::post_asap::share_common_summary_subtrees(roots); + use std::hash::{Hash, Hasher}; + let mut hash = std::collections::hash_map::DefaultHasher::new(); + let mut pending = roots + .iter() + .map(|(_, node)| node.as_ref()) + .collect::>(); + while let Some(node) = pending.pop() { + std::mem::discriminant(&node.expr).hash(&mut hash); + let raw = match &node.expr { + SummaryExpr::KeepPreAsap(raw) => Some(raw.as_ref()), + _ => None, + }; + let operation = match &node.expr { + SummaryExpr::ValueOperation { + timing, operation, .. + } => serde_json::json!((timing, operation)), + SummaryExpr::BinaryOp { + timing, operator, .. + } => serde_json::json!((timing, operator)), + SummaryExpr::SummaryMerge { timing, .. } => serde_json::json!(timing), + _ => serde_json::Value::Null, + }; + let mut value = + serde_json::to_value((&node.schema, &node.guarantee, raw, operation)) + .map_err(|_| { + RealizationError::PhysicalRealization( + "candidate identity serialization failed", + ) + })?; + fn normalize(value: &mut serde_json::Value) { + match value { + serde_json::Value::Number(number) + if number.as_f64() == Some(0.0) => + { + *value = serde_json::json!(0); + } + serde_json::Value::Array(values) => { + values.iter_mut().for_each(normalize) + } + serde_json::Value::Object(values) => { + values.values_mut().for_each(normalize) + } + _ => {} + } + } + normalize(&mut value); + value.sort_all_objects(); + value.to_string().hash(&mut hash); + match &node.expr { + SummaryExpr::KeepPreAsap(_) => {} + SummaryExpr::BinaryOp { lhs, rhs, .. } => { + pending.extend([lhs.as_ref(), rhs.as_ref()]) + } + SummaryExpr::RelationalJoin { left, right, .. } + | SummaryExpr::SummarySubtract { left, right } => { + pending.extend([left.as_ref(), right.as_ref()]) + } + SummaryExpr::ValueOperation { child, .. } + | SummaryExpr::SummaryAgg { child, .. } => pending.push(child.as_ref()), + SummaryExpr::SummaryJoin { outer, inner, .. } => { + pending.extend([outer.as_ref(), inner.as_ref()]) + } + SummaryExpr::SummaryDelete { summary_input, .. } + | SummaryExpr::SummaryEstimate { summary_input, .. } => { + pending.push(summary_input.as_ref()) + } + SummaryExpr::SummaryMerge { children, .. } => { + pending.extend(children.iter().map(|child| child.as_ref())) + } + } + } + let bucket = seen.entry(hash.finish()).or_default(); + if !bucket + .iter() + .any(|&index| inventory.candidates[index] == roots) + { + bucket.push(inventory.candidates.len()); + inventory.candidates.push(roots); + } + } + Err(error) => { + let reason = error.to_string(); + if !inventory.rejected_assemblies.contains(&reason) { + inventory.rejected_assemblies.push(reason); + } + } + } + } + Ok(inventory) + } +} + /// Lifecycle-aware whole-subplan costs keyed by target and candidate identity. #[derive(Default)] pub(crate) struct CandidateCostOverrides { @@ -4316,9 +4941,17 @@ fn rank_group<'a>( // purpose. `total_cmp` gives deterministic placement to a model's NaN // placeholders without dropping any candidate. ranked.sort_by(|a, b| { - cost_model - .estimate_cost(a, &target) - .total_cmp(&cost_model.estimate_cost(b, &target)) + match ( + cost_model.candidate_cost(a, &target), + cost_model.candidate_cost(b, &target), + ) { + (Some(a), Some(b)) => a.0.total_cmp(&b.0), + (Some(_), None) => std::cmp::Ordering::Less, + (None, Some(_)) => std::cmp::Ordering::Greater, + (None, None) => cost_model + .estimate_cost(a, &target) + .total_cmp(&cost_model.estimate_cost(b, &target)), + } }); ranked } @@ -4553,12 +5186,32 @@ impl<'a> GlobalSelection<'a> { self.assemble_target(target).map(Some) } + /// Assemble a complete query result, including an exact-state readout when + /// needed. `assemble_selected_dag` also serves internal state frontiers; + /// callers exposing query results must use this boundary instead. + pub fn assemble_selected_query( + &self, + target: &Rc, + ) -> Result>, RealizationError> { + self.assemble_selected_dag(target)? + .map(|node| finalize_query_candidate(node, target)) + .transpose() + } + fn assemble_target(&self, target: &Rc) -> Result, RealizationError> { let ptr = Rc::as_ptr(target); if let Some(node) = self.assembled_nodes.borrow().get(&ptr) { return Ok(Rc::clone(node)); } - let node = if query_time_nested_sum(target) { + let selected_composed_summary = self + .groups + .get(&ptr) + .and_then(|sel| sel.chosen) + .is_some_and(|candidate| matches!(&candidate.replacement, + Replacement::Summary(node) if matches!(&node.expr, + SummaryExpr::SummaryAgg { child, .. } + if matches!(&child.expr, SummaryExpr::KeepPreAsap(raw) if !contains_aggregate(raw))))); + let node = if query_time_nested_sum(target) && !selected_composed_summary { self.assemble_residual(target)? } else { match self @@ -4609,8 +5262,8 @@ impl<'a> GlobalSelection<'a> { let Some(pred) = normalized_pred else { return keep_pre_asap(target); }; - let left = finalize_exact_accumulator(self.assemble_target(left)?, left)?; - let right = finalize_exact_accumulator(self.assemble_target(right)?, right)?; + let left = finalize_query_candidate(self.assemble_target(left)?, left)?; + let right = finalize_query_candidate(self.assemble_target(right)?, right)?; let guarantee = relational_join_guarantee(left.guarantee.as_ref(), right.guarantee.as_ref()); let node = Rc::new(SummaryNode { @@ -4681,7 +5334,7 @@ impl<'a> GlobalSelection<'a> { ), _ => return keep_pre_asap(target), }; - let child = finalize_exact_accumulator(self.assemble_target(child_target)?, child_target)?; + let child = finalize_query_candidate(self.assemble_target(child_target)?, child_target)?; let guarantee = child.guarantee.clone(); let node = Rc::new(SummaryNode { expr: SummaryExpr::ValueOperation { @@ -4904,7 +5557,7 @@ fn composition_options<'a>( None => child_group.candidates.iter().collect(), }; for child_candidate in child_candidates { - if child_candidate.has_missing_accuracy_evidence() { + if !is_automatically_selectable(child_candidate, cost_model) { continue; } let Replacement::Summary(summary) = &child_candidate.replacement else { @@ -5081,7 +5734,7 @@ impl PlanSpace { .candidates .iter() .filter(|candidate| !is_composition_candidate(candidate)) - .filter(|candidate| is_automatically_selectable(candidate)) + .filter(|candidate| is_automatically_selectable(candidate, cost_model)) .filter_map(|candidate| { costs .get(&group.target, candidate) @@ -5107,7 +5760,7 @@ impl PlanSpace { .filter(|candidate| { !is_cse_candidate(candidate) && !is_composition_candidate(candidate) - && is_automatically_selectable(candidate) + && is_automatically_selectable(candidate, cost_model) }) .filter_map(|candidate| { cost_model @@ -5164,7 +5817,7 @@ impl PlanSpace { .filter(|candidate| { !is_cse_candidate(candidate) && !is_composition_candidate(candidate) - && is_automatically_selectable(candidate) + && is_automatically_selectable(candidate, cost_model) }) .filter_map(|candidate| { cost_model @@ -5209,7 +5862,7 @@ impl PlanSpace { // children (see `multiplier`'s `_ => effective` arm). None => rank_group(group, cost_model).into_iter().find(|candidate| { !is_composition_candidate(candidate) - && is_automatically_selectable(candidate) + && is_automatically_selectable(candidate, cost_model) && (cost_model .candidate_cost( candidate, @@ -5226,7 +5879,7 @@ impl PlanSpace { .find(|candidate| { !is_cse_candidate(candidate) && !is_composition_candidate(candidate) - && is_automatically_selectable(candidate) + && is_automatically_selectable(candidate, cost_model) && (cost_model .candidate_cost(candidate, &effective_target) .is_some() @@ -5312,8 +5965,9 @@ fn is_cse_candidate(candidate: &ReplacementSubDAG) -> bool { ) } -fn is_automatically_selectable(candidate: &ReplacementSubDAG) -> bool { +fn is_automatically_selectable(candidate: &ReplacementSubDAG, cost_model: &dyn CostModel) -> bool { !candidate.has_missing_accuracy_evidence() + && candidate.runtime_support_evidence(cost_model) != Some(false) } /// How much one direct reference to `parent_ptr` actually costs, once @@ -6232,6 +6886,150 @@ mod tests { use asap_types::types::AccuracyTarget; use std::collections::HashMap; + // Every exposed query result has a readout; internal accumulator frontiers stay states. + #[test] + fn query_candidate_roots_do_not_leak_exact_accumulator_state() { + for query in [ + "sum by(job)(rate(m[1m]))", + "sum by(job)(m)", + "sum_over_time(m[1m])", + ] { + let root = Rc::new(lower_promql(query, AccuracyTarget::Exact)); + let space = search_workload(vec![(0usize, root.clone())]); + let inventory = space.enumerate_candidate_dags(4096).unwrap(); + assert!(!inventory.candidates.is_empty()); + let strategy = SketchAlgorithmStrategy::new(&DefaultCostModel); + for candidate in strategy.propose(&TargetSubDAG::new(&root)).candidates { + if let Replacement::Summary(node) = candidate.replacement { + let output = finalize_query_candidate(node, &root).unwrap(); + assert!( + output + .schema + .fields + .iter() + .all(|field| matches!(field.dtype, SummaryFamilyType::Plain(_))), + "direct candidate {query} leaks state" + ); + } + } + let selected = space + .global_selection(&DefaultCostModel) + .assemble_selected_query(&space.roots[0].1) + .unwrap() + .unwrap(); + for node in inventory + .candidates + .iter() + .map(|forest| &forest[0].1) + .chain(std::iter::once(&selected)) + { + assert!( + node.schema + .fields + .iter() + .all(|field| matches!(field.dtype, SummaryFamilyType::Plain(_))), + "{query}: query root leaks state: {:?}", + node.schema + ); + } + } + } + + #[test] + fn unpriced_inventory_retains_quantile_families_and_raw_execution() { + let query = Rc::new(agg(vec![2], default_quantile(0.9), metric_scan(&["job"]))); + let space = search_workload(vec![(0usize, query)]); + let inventory = space.enumerate_candidate_dags(4096).unwrap(); + let roots = inventory + .candidates + .iter() + .map(|forest| format!("{:?}", forest[0].1)) + .collect::>(); + assert!(roots.iter().any(|root| root.contains("Kll"))); + assert!(roots.iter().any(|root| root.contains("DDSketch"))); + assert!(inventory + .candidates + .iter() + .any(|forest| matches!(forest[0].1.expr, SummaryExpr::KeepPreAsap(_)))); + } + + // Independent roots must not require materializing their Cartesian product. + #[test] + fn root_inventory_preserves_choices_without_workload_cartesian_expansion() { + let roots = (0..24usize) + .map(|id| { + ( + id, + Rc::new(agg( + vec![2], + default_quantile((id + 1) as f64 / 25.0), + metric_scan(&["job"]), + )), + ) + }) + .collect(); + let space = search_workload(roots); + assert!(space.enumerate_candidate_dags(4096).is_err()); + for id in 0..24 { + let inventory = space.enumerate_candidate_dags_for_root(&id, 4096).unwrap(); + assert!(inventory + .candidates + .iter() + .all(|forest| forest.len() == 1 && forest[0].0 == id)); + let descriptions = inventory + .candidates + .iter() + .map(|forest| format!("{:?}", forest[0].1)) + .collect::>(); + assert!(descriptions.iter().any(|node| node.contains("Kll"))); + assert!(descriptions.iter().any(|node| node.contains("DDSketch"))); + assert!(inventory + .candidates + .iter() + .any(|forest| matches!(forest[0].1.expr, SummaryExpr::KeepPreAsap(_)))); + } + assert!(space.enumerate_candidate_dags_for_root(&24, 4096).is_err()); + assert!(space.enumerate_candidate_dags_for_root(&0, 0).is_err()); + } + + // Factoring changes enumeration, not the set of root computations. + #[test] + fn root_inventory_matches_projection_of_exhaustive_workload_inventory() { + let roots = (0..2usize) + .map(|id| { + ( + id, + Rc::new(agg( + vec![2], + default_quantile(0.5 + id as f64 * 0.4), + metric_scan(&["job"]), + )), + ) + }) + .collect(); + let space = search_workload(roots); + let full = space.enumerate_candidate_dags(4096).unwrap(); + for id in 0..2 { + let inventory = space.enumerate_candidate_dags_for_root(&id, 4096).unwrap(); + for forest in &full.candidates { + let node = &forest.iter().find(|(root, _)| *root == id).unwrap().1; + assert!(inventory.candidates.iter().any(|one| &one[0].1 == node)); + } + for one in &inventory.candidates { + assert!(full.candidates.iter().any(|forest| forest + .iter() + .any(|(root, node)| *root == id && node == &one[0].1))); + } + } + } + + #[test] + fn inventory_budget_never_returns_a_silent_partial_search() { + let query = Rc::new(agg(vec![2], default_quantile(0.9), metric_scan(&["job"]))); + let space = search_workload(vec![(0usize, query)]); + assert!(space.enumerate_candidate_dags(0).is_err()); + } + fn equi_pred(left: ColumnId, right: ColumnId) -> Predicate { Predicate(Rc::new(QueryExpr::Compare { left: Rc::new(QueryExpr::Column(left)), @@ -6265,6 +7063,25 @@ mod tests { ); } + // Approximate requests also admit exact temporal ranking candidates. + #[test] + fn approximate_temporal_topk_admits_exact_maintained_values() { + let root = Rc::new(lower_promql( + "topk by(job)(1,count_over_time(a[5m]))", + AccuracyTarget::EpsilonDelta { + epsilon: 0.01, + delta: 0.01, + }, + )); + let planning_inputs = + CandidatePlanningInputs::with_default_accuracy(&crate::cost_model::DefaultCostModel); + let node = exact_topk_over_temporal_values(&root, planning_inputs) + .unwrap() + .expect("exact ranking is legal for an approximate request"); + assert!(node.guarantee.as_ref().unwrap().is_exact()); + asap_types::post_asap::compile_executable_dag(&node).unwrap(); + } + // Exact Top-K consumes the Planner's maintained temporal values. #[test] fn exact_temporal_topk_has_a_maintained_value_candidate() { @@ -6318,6 +7135,43 @@ mod tests { } } + // A bounded exact mean can share the relative division proof with a quantile. + #[test] + fn bounded_mean_quantile_ratio_is_certified() { + struct Domain; + impl AccuracyEvidenceProvider for Domain { + fn quantile_input_domain( + &self, + _: &QueryExpr, + ) -> Option { + Some(crate::accuracy::QuantileInputDomain { + lower: 1.0, + upper: 1000.0, + max_samples: 10000, + contract: "finite test population".into(), + }) + } + } + let target = AccuracyTarget::EpsilonDelta { + epsilon: 0.01, + delta: 0.01, + }; + let inputs = CandidatePlanningInputs { + evidence: &Domain, + ..CandidatePlanningInputs::with_default_accuracy(&DefaultCostModel) + }; + for query in [ + "avg_over_time(a[5m]) / quantile_over_time(0.5,a[5m])", + "quantile_over_time(0.5,a[5m]) / avg_over_time(a[5m])", + ] { + let root = Rc::new(lower_promql(query, target.clone())); + let node = realize_binary(&root, inputs, Some(&target)) + .unwrap() + .expect("bounded ratio candidate"); + assert!(DefaultAccuracyModel.satisfies(node.guarantee.as_ref().unwrap(), &target)); + } + } + // Missing domain proof permits an uncertified direct quantile ratio only. #[test] fn quantile_ratio_without_input_proof_has_no_root_guarantee() { @@ -6614,6 +7468,23 @@ mod tests { ); } + // Exact counting remains a legal candidate under an approximate target. + #[test] + fn approximate_count_includes_exact_accumulator_candidate() { + let intent = AggIntent::Count { + accuracy: eps(0.01), + }; + assert!(realizations_for_intent(&intent, &DefaultCostModel) + .iter() + .any(|candidate| matches!( + candidate, + Realization::ExactAggregate { + kind: ExactKind::Count, + .. + } + ))); + } + #[test] fn epsilon_delta_sizes_cms_depth() { let intent = AggIntent::Count { @@ -7399,7 +8270,7 @@ mod tests { assert_eq!(agg_group.consumer_count, 1); assert_eq!( agg_group.candidates.len(), - 5, + 6, "Hydra candidates with unknown evidence remain available: {:?}", agg_group.candidates ); @@ -8458,6 +9329,125 @@ mod tests { ); } + // A cheap but physically infeasible candidate must not be selected. + #[test] + fn explicit_summary_infeasibility_prevents_selection() { + struct Unsupported; + impl CostModel for Unsupported { + fn rank_candidates( + &self, + _: &AggIntent, + candidates: &[SketchAlgorithm], + ) -> Vec { + candidates.to_vec() + } + fn candidate_cost(&self, _: &ReplacementSubDAG, _: &TargetSubDAG<'_>) -> Option { + Some(Cost(1.0)) + } + fn summary_support_evidence(&self, _: &SummaryNode) -> Option { + Some(false) + } + } + let root = Rc::new(lower_promql("sum_over_time(a[1m])", AccuracyTarget::Exact)); + let space = search_workload(vec![("q", root)]); + let selected = space.global_selection(&Unsupported); + assert!(selected + .for_target(&space.roots[0].1) + .unwrap() + .chosen + .is_none()); + } + + // Composable temporal/grouped Sum must be executable as one producer. + #[test] + fn grouped_temporal_sum_has_one_summary_producer_candidate() { + let root = Rc::new(lower_promql( + "sum by(job)(sum_over_time(a[1m]))", + AccuracyTarget::Exact, + )); + let candidates = + SketchAlgorithmStrategy::default_cost_model().replacements(&TargetSubDAG::new(&root)); + assert!(candidates + .iter() + .any(|candidate| matches!(&candidate.replacement, + Replacement::Summary(node) if matches!(&node.expr, + SummaryExpr::SummaryAgg { reduction: Reduction::Reduce(_), child, .. } + if matches!(child.expr, SummaryExpr::KeepPreAsap(_)))))); + struct PreferComposed; + impl CostModel for PreferComposed { + fn rank_candidates( + &self, + _: &AggIntent, + candidates: &[SketchAlgorithm], + ) -> Vec { + candidates.to_vec() + } + fn candidate_cost( + &self, + candidate: &ReplacementSubDAG, + _: &TargetSubDAG<'_>, + ) -> Option { + Some(Cost( + if matches!(&candidate.replacement, + Replacement::Summary(node) if matches!(&node.expr, + SummaryExpr::SummaryAgg { reduction: Reduction::Reduce(_), child, .. } + if matches!(child.expr, SummaryExpr::KeepPreAsap(_)))) + { + 1.0 + } else { + 100.0 + }, + )) + } + } + let space = search_workload(vec![("q", root.clone())]); + let selected = space.global_selection(&PreferComposed); + let node = selected.assemble_target(&space.roots[0].1).unwrap(); + assert!(matches!(&node.expr, + SummaryExpr::SummaryAgg { reduction: Reduction::Reduce(_), child, .. } + if matches!(child.expr, SummaryExpr::KeepPreAsap(_)))); + } + + // Mixed candidate ranking must honor explicit costs, not legacy estimates. + #[test] + fn mixed_candidate_ranking_uses_explicit_candidate_costs() { + struct ExplicitCosts; + impl CostModel for ExplicitCosts { + fn rank_candidates( + &self, + _: &AggIntent, + candidates: &[SketchAlgorithm], + ) -> Vec { + candidates.to_vec() + } + fn candidate_cost( + &self, + candidate: &ReplacementSubDAG, + _: &TargetSubDAG<'_>, + ) -> Option { + Some(Cost( + if candidate.provenance == ReplacementProvenance::LogicalRewrite { + 1.0 + } else { + 100.0 + }, + )) + } + } + let root = Rc::new(lower_promql( + "sum by(job)(sum_over_time(a[1m]))", + AccuracyTarget::Exact, + )); + let space = search_workload(vec![("q", root)]); + let selection = space.global_selection(&ExplicitCosts); + let selected = selection + .for_target(&space.roots[0].1) + .unwrap() + .chosen + .unwrap(); + assert_eq!(selected.provenance, ReplacementProvenance::LogicalRewrite); + } + #[test] fn global_selection_compares_a_logical_rewrite_with_the_cse_choice() { struct PreferLogicalRewrite; diff --git a/crates/asap-aware-mapping/src/rewrite.rs b/crates/asap-aware-mapping/src/rewrite.rs index 6081105e..44b685f8 100644 --- a/crates/asap-aware-mapping/src/rewrite.rs +++ b/crates/asap-aware-mapping/src/rewrite.rs @@ -247,7 +247,7 @@ fn build_rewrite(root: &Rc) -> Option> { /// Compose adjacent per-entity and cross-entity accumulators when their /// algebra, rather than a query-language spelling, proves equivalence. -fn composed_aggregate_rewrite(root: &Rc) -> Option> { +pub(crate) fn composed_aggregate_rewrite(root: &Rc) -> Option> { let original_schema = root.output_schema().ok()?; let QueryExpr::Aggregate { reduction: outer_reduction @ Reduction::Reduce(_), diff --git a/crates/asap-physical-operators/Cargo.toml b/crates/asap-physical-operators/Cargo.toml new file mode 100644 index 00000000..e8ab710a --- /dev/null +++ b/crates/asap-physical-operators/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "asap-physical-operators" +version = "0.1.0" +edition = "2021" + +[dependencies] +futures = "0.3" +planner-types = { package = "asap-types", path = "../types" } +asap_sketch_codec = { path = "../asap_sketch_codec" } +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "5f03ccbd798ed5fec62bdd839bcb331123cab369" } +serde = { version = "1", features = ["derive", "rc"] } +serde_json = "1" +tracing = "0.1" +thiserror = "2" +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" +asap-aware-mapping = { path = "../asap-aware-mapping" } +asap-frontend-promql = { path = "../frontend-promql" } diff --git a/crates/asap-physical-operators/README.md b/crates/asap-physical-operators/README.md new file mode 100644 index 00000000..31bbc49a --- /dev/null +++ b/crates/asap-physical-operators/README.md @@ -0,0 +1,126 @@ +# 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. + +`plan::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. + +`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::{ + expressions::Expression, + operators::Operator, + values::Value, + plan::PhysicalDag, + runtime::{Limits, 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>(()) +``` + +`physical_planner::compile` accepts a post-ASAP DAG and typed input contracts. +The resulting candidate is instantiated with deployment readers after selection. 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 physical compiler 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, HLL and Float64 weighted CMS and CountSketch with candidate heaps. Binding checks family, +parameters and readout compatibility; source batches also validate state payloads. +Existing accumulator algorithms are reused as kernels behind these operators. + +This crate is owned by ASAPPlanner. Its `planner-types` dependency is the local +IR crate, so a contract change and its execution tests belong in the same PR. +Deployments supply storage/ingestion sources and adapt output protocols. The +library has no ASAPQuery-backend dependency. Backend raw Scan remains a separate +deployment capability. + +See [the design](../../docs/design_docs/physical-planning-and-deployment.md). + +## Module boundaries + +- `plan`: immutable graph, operator interface, schemas and execution properties. +- `runtime`: per-run streams, shared producers, memory reservations and cancellation. +- `expressions`: scalar evaluation; typed builders and the Planner expression adapter. +- `operators`: projection, filter, joins, aggregate/window, sort, limit and summary implementations. +- `sources`: raw-source interface, Scan and the memory connector. +- `physical_planner`: native operator lowering, typed input contracts and checked instantiation. +- `summary_kernels`: sketchlib state adapters, exact accumulators, update adapters and traits. +- `stored_state`: persisted-state decoding, delta reconstruction and readout. +- `capability`: explicit kernel and native-batch/readout validation. + +The old `dag`, `accumulators`, `factory`, `traits` and `arithmetic` paths remain re-exports for deployment +source compatibility. They contain no alternative execution implementations. + +A source must declare `Boundedness::Bounded` to feed a blocking operator. +The default for a custom raw source is `Unknown`; query or ingestion scope alone +does not promise that its cursor ends. `PhysicalDag::properties` validates these +requirements before any source starts and returns boundedness and emission mode +for every reachable node. The memory connector declares finite input. Custom +physical sources expose the same facts through `PhysicalOperator::properties`. + +Blocking operators reserve estimated workspace and yield cooperatively during +row processing and sort merges. Cancellation releases reservations when the +stream is polled or dropped. Individual scalar evaluations, bounded sort chunks +and sketch kernel calls are synchronous; this is not preemptive execution. +There is no spill or partitioned parallel execution in this implementation. + +## Physical compilation and deployment inputs + +`physical_planner::compile` accepts a Planner `ExecutableDag`, typed +`InputContract`s and output roots. It returns a reusable `CompiledPhysicalDag` +containing selected native operators and no live readers. Compilation validates +schemas, input ordering, sharing and boundedness before deployment source access. + +A deployment calls `CompiledPhysicalDag::instantiate` with exactly the declared +inputs. This checks source schemas and execution properties and constructs the +runnable graph without repeating logical lowering. The graph executes through +the shared runtime with independent per-run state. Window coverage, revision and +maintenance-policy admission remain deployment/planning contracts; this compiler +does not discover storage or silently change a selected maintenance strategy. + +`physical_planner::compile_temporal_pane_candidate` lowers a selected continuous +KLL lifecycle and Sliding/Tumbling framework into maintenance and query DAGs. +`TemporalPaneMaintenance` supplies pane geometry and a resolved complete entity +identity contract. The compiler inserts population guards, scan predicates, +pane construction, ordered state slots, a shared merge and quantile readouts. +Pane outputs have distinct physical identities from the logical whole-window +summary, and the returned candidate retains the maintenance contract for binding. +Each run checks phase, pane timestamps and duplicate entity states. The initial +realization uses complete bounded snapshots; partial edges, exponential +histograms and cross-run delta accumulation are unsupported. Storage identities, +revision selection, completeness/readiness evidence and scheduling stay with +deployment. diff --git a/crates/asap-physical-operators/src/aggregation_type.rs b/crates/asap-physical-operators/src/aggregation_type.rs new file mode 100644 index 00000000..647f7604 --- /dev/null +++ b/crates/asap-physical-operators/src/aggregation_type.rs @@ -0,0 +1,215 @@ +//! Shared aggregation vocabulary for configuration and accumulator dispatch. +//! The wire shape combines aggregation type, subtype, and parameters. +//! `AccumulatorSpec` provides a typed representation at conversion boundaries. + +use serde::{Deserialize, Serialize}; +use std::fmt; +use std::str::FromStr; + +/// Concrete aggregation/sketch type used in precompute configs and accumulator dispatch. +/// +/// `Display` outputs the canonical PascalCase name used in YAML/JSON configs. +/// `FromStr` accepts the canonical name plus legacy aliases (e.g. "KLL" → `DatasketchesKLL`). +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum AggregationType { + // ---------- single-population (non-keyed) ---------- + Sum, + Count, + Increase, + Rate, + Min, + Max, + DatasketchesKLL, + // ---------- multi-population (keyed) ---------- + HydraKLL, + CountMinSketch, + CountMinSketchWithHeap, + CountSketch, + CountSketchWithHeap, + // ---------- cardinality / set tracking ---------- + HLL, + UnivMon, + DDSketch, + // ---------- legacy config wrapper names ---------- + SingleSubpopulation, + MultipleSubpopulation, +} + +impl AggregationType { + /// Adapt a storage/processor tag to Planner's exact family. Keyed storage + /// changes the payload layout, not the semantic family. + pub fn planner_exact_family(self) -> Option { + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + let (kind, params) = match self { + Self::Sum => (ExactKind::Sum, ExactParams::Sum), + Self::Count => (ExactKind::Count, ExactParams::Count), + Self::Increase => (ExactKind::Increase, ExactParams::Increase), + Self::Rate => (ExactKind::Rate, ExactParams::Rate), + Self::Min => (ExactKind::Min, ExactParams::Min), + Self::Max => (ExactKind::Max, ExactParams::Max), + _ => return None, + }; + Some(SummaryFamilyType::ExactAggregate(kind, params)) + } + + pub fn as_str(self) -> &'static str { + match self { + AggregationType::Sum => "Sum", + AggregationType::Count => "Count", + AggregationType::Increase => "Increase", + AggregationType::Rate => "Rate", + AggregationType::Min => "Min", + AggregationType::Max => "Max", + AggregationType::DatasketchesKLL => "DatasketchesKLL", + AggregationType::HydraKLL => "HydraKLL", + AggregationType::CountMinSketch => "CountMinSketch", + AggregationType::CountMinSketchWithHeap => "CountMinSketchWithHeap", + AggregationType::CountSketch => "CountSketch", + AggregationType::CountSketchWithHeap => "CountSketchWithHeap", + AggregationType::HLL => "HLL", + AggregationType::UnivMon => "UnivMon", + AggregationType::DDSketch => "DDSketch", + AggregationType::SingleSubpopulation => "SingleSubpopulation", + AggregationType::MultipleSubpopulation => "MultipleSubpopulation", + } + } + + /// Returns `true` if this type produces keyed (multi-population) accumulators. + pub fn is_keyed(self) -> bool { + matches!( + self, + AggregationType::MultipleSubpopulation + | AggregationType::CountMinSketch + | AggregationType::CountMinSketchWithHeap + | AggregationType::CountSketch + | AggregationType::CountSketchWithHeap + | AggregationType::HydraKLL + ) + } +} + +impl fmt::Display for AggregationType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(self.as_str()) + } +} + +impl FromStr for AggregationType { + type Err = String; + + fn from_str(s: &str) -> Result { + match s { + // Canonical names + "Sum" => Ok(AggregationType::Sum), + "Count" => Ok(AggregationType::Count), + "Increase" => Ok(AggregationType::Increase), + "Rate" => Ok(AggregationType::Rate), + "Min" => Ok(AggregationType::Min), + "Max" => Ok(AggregationType::Max), + "DatasketchesKLL" => Ok(AggregationType::DatasketchesKLL), + "HydraKLL" => Ok(AggregationType::HydraKLL), + "CountMinSketch" => Ok(AggregationType::CountMinSketch), + "CountMinSketchWithHeap" => Ok(AggregationType::CountMinSketchWithHeap), + "CountSketch" => Ok(AggregationType::CountSketch), + "CountSketchWithHeap" => Ok(AggregationType::CountSketchWithHeap), + "HLL" | "HyperLogLog" => Ok(AggregationType::HLL), + "UnivMon" => Ok(AggregationType::UnivMon), + "DDSketch" | "DdSketch" => Ok(AggregationType::DDSketch), + "SingleSubpopulation" => Ok(AggregationType::SingleSubpopulation), + "MultipleSubpopulation" => Ok(AggregationType::MultipleSubpopulation), + // Legacy accumulator-suffixed aliases + "SumAccumulator" | "SumAggregator" | "sum" => Ok(AggregationType::Sum), + "IncreaseAccumulator" | "IncreaseAggregator" | "increase" => { + Ok(AggregationType::Increase) + } + "MinAccumulator" | "MinAggregator" | "min" => Ok(AggregationType::Min), + "MaxAccumulator" | "MaxAggregator" | "max" => Ok(AggregationType::Max), + "DatasketchesKLLAccumulator" | "KLL" | "kll" | "datasketches_kll" => { + Ok(AggregationType::DatasketchesKLL) + } + "HydraKllSketchAccumulator" | "hydra_kll" => Ok(AggregationType::HydraKLL), + "CountMinSketchAccumulator" | "CMS" | "cms" | "count_min_sketch" => { + Ok(AggregationType::CountMinSketch) + } + "CountMinSketchWithHeapAccumulator" => Ok(AggregationType::CountMinSketchWithHeap), + "CountSketchAccumulator" | "CS" | "cs" | "count_sketch" => { + Ok(AggregationType::CountSketch) + } + "CountSketchWithHeapAccumulator" => Ok(AggregationType::CountSketchWithHeap), + // Retired names. `MinMax` used to be one accumulator whose + // direction rode alongside in `aggregationSubType`; the two + // directions are separate types now, so there is no safe + // direction to guess here -- resolving a min workload as a + // max one is silently wrong, not merely imprecise. + "MinMax" + | "MinMaxAccumulator" + | "MinMaxAggregator" + | "min_max" + | "MultipleMinMax" + | "MultipleMinMaxAccumulator" + | "multiple_min_max" => Err(format!( + "Retired aggregation type: '{s}' -- min and max are separate types now, \ + use 'Min'/'Max'" + )), + _ => Err(format!("Unknown aggregation type: '{s}'")), + } + } +} + +impl Serialize for AggregationType { + fn serialize(&self, serializer: S) -> Result { + serializer.serialize_str(self.as_str()) + } +} + +impl<'de> Deserialize<'de> for AggregationType { + fn deserialize>(deserializer: D) -> Result { + let s = String::deserialize(deserializer)?; + s.parse().map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + + /// Removed layout tags cannot be installed as semantic families. + #[test] + fn rejects_keyed_family_aliases() { + for name in [ + "MultipleSum", + "MultipleIncrease", + "MultipleMin", + "MultipleMax", + ] { + assert!(name.parse::().is_err(), "{name}"); + } + } + + #[test] + fn storage_layout_tags_do_not_create_planner_families() { + for (storage, expected) in [ + (AggregationType::Sum, ExactKind::Sum), + (AggregationType::Count, ExactKind::Count), + (AggregationType::Increase, ExactKind::Increase), + (AggregationType::Rate, ExactKind::Rate), + ] { + let family = storage.planner_exact_family().unwrap(); + assert!( + matches!(family, SummaryFamilyType::ExactAggregate(kind, _) if kind == expected) + ); + } + assert_eq!( + AggregationType::Rate.planner_exact_family(), + Some(SummaryFamilyType::ExactAggregate( + ExactKind::Rate, + ExactParams::Rate + )) + ); + assert_ne!( + AggregationType::Rate.planner_exact_family(), + AggregationType::Increase.planner_exact_family() + ); + } +} diff --git a/crates/asap-physical-operators/src/capability.rs b/crates/asap-physical-operators/src/capability.rs new file mode 100644 index 00000000..12ecabcd --- /dev/null +++ b/crates/asap-physical-operators/src/capability.rs @@ -0,0 +1,219 @@ +//! Capability boundaries, checked without constructing accumulator state. +//! +//! `validate_summary_kernel` checks update kernels, including families without a +//! native batch representation. `validate_native_family` and +//! `validate_native_readout` check native state and scalar readout support. +//! Keyed weighted-frequency readouts are checked by `Operator::keyed_readout`. +//! A successful kernel check alone does not mean an executable DAG will bind. +//! +//! Persisted state uses `stored_state` decoding and readout contracts; support +//! there does not imply a native build/merge operator. Full plan acceptance is +//! owned by `binding`, which also validates schemas, expressions and inputs. +use crate::Error; +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() && !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() +} + +pub(crate) fn is_unit_sample_frequency(update: &planner_types::post_asap::SummaryUpdate) -> bool { + use planner_types::post_asap::{NonNegativeWeightProof, SummaryInputExpr, WeightDomain}; + matches!( + update.item, + Some(SummaryInputExpr::Column( + planner_types::pre_asap::ColumnRef::SampleValue + )) + ) && matches!(update.weight, SummaryInputExpr::Constant(1.0)) + && matches!( + update.weight_domain, + WeightDomain::NonNegative { + proof: NonNegativeWeightProof::UnitCount + } + ) +} + +pub fn validate_native_family(family: &SummaryFamilyType) -> Result<(), Error> { + use planner_types::post_asap::SketchAlgorithm as A; + if let SummaryFamilyType::Sketch(kind, grouping) = family { + if matches!(kind.algorithm(), A::CmsWithHeap | A::CountSketchWithHeap) { + let (_, width, depth, _) = + crate::summary_kernels::weighted_frequency::WeightedFrequency::configuration(kind)?; + return if valid_matrix(width as u32, depth as u32) && grouping == &Default::default() { + Ok(()) + } else { + Err(Error::Invalid( + "invalid weighted frequency dimensions or grouping strategy".into(), + )) + }; + } + } + 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) +} + +pub fn validate_native_readout( + family: &SummaryFamilyType, + statistic: crate::Statistic, + parameters: &std::collections::HashMap, +) -> Result<(), Error> { + validate_native_family(family)?; + 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(Error::Invalid( + "readout is not implemented for this summary family".into(), + )); + } + 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(Error::Invalid( + "quantile readout requires quantile in [0,1]".into(), + )); + } + Ok(()) +} 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 00000000..c7837672 --- /dev/null +++ b/crates/asap-physical-operators/src/dag/mod.rs @@ -0,0 +1,8 @@ +//! Compatibility imports. New code should use plan, runtime, operators, physical_planner and sources directly. +pub use crate::plan::{NodeId, PhysicalDag, PhysicalOperator}; +pub use crate::runtime::batch_execution; +pub use crate::runtime::{ + Input, Limits, OutputStream, Reservation, RunContext, Scope, SharedValue, +}; +pub use crate::Error; +pub use crate::{expressions, operators, physical_planner as planner, sources as scan, values}; diff --git a/crates/asap-physical-operators/src/error.rs b/crates/asap-physical-operators/src/error.rs new file mode 100644 index 00000000..afee33d8 --- /dev/null +++ b/crates/asap-physical-operators/src/error.rs @@ -0,0 +1,18 @@ +use crate::plan::NodeId; +#[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, +} diff --git a/crates/asap-physical-operators/src/expressions/arithmetic.rs b/crates/asap-physical-operators/src/expressions/arithmetic.rs new file mode 100644 index 00000000..30e277d4 --- /dev/null +++ b/crates/asap-physical-operators/src/expressions/arithmetic.rs @@ -0,0 +1,63 @@ +//! 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), + } +} + +/// Execute the Planner binary contract after a deployment has resolved matching rows. +pub fn evaluate_binary( + operator: &planner_types::post_asap::BinaryOperator, + left: f64, + right: f64, +) -> Result { + use crate::{values::Value, Error}; + use planner_types::pre_asap::{ArithmeticOpKind, BinaryOpKind, CompareOpKind}; + let invalid = + || Error::Invalid("unsupported binary operation or invalid checked-division domain".into()); + if operator.vector_match.is_some() { + return Err(invalid()); + } + if operator.checked_relative_division || operator.checked_finite_division { + if operator.kind != BinaryOpKind::Arithmetic(ArithmeticOpKind::Div) + || !left.is_finite() + || !right.is_finite() + || right == 0. + { + return Err(invalid()); + } + let value = left / right; + if !value.is_finite() || (operator.checked_relative_division && !value.is_normal()) { + return Err(invalid()); + } + return Ok(Value::Float64(value)); + } + Ok(match operator.kind { + BinaryOpKind::Arithmetic(ref op) => { + Value::Float64(evaluate_float64_arithmetic(op, left, right)) + } + BinaryOpKind::Compare(ref op) => Value::Bool(match op { + CompareOpKind::Eq => left == right, + CompareOpKind::Ne => left != right, + CompareOpKind::Lt => left < right, + CompareOpKind::Le => left <= right, + CompareOpKind::Gt => left > right, + CompareOpKind::Ge => left >= right, + _ => return Err(invalid()), + }), + _ => return Err(invalid()), + }) +} diff --git a/crates/asap-physical-operators/src/expressions/mod.rs b/crates/asap-physical-operators/src/expressions/mod.rs new file mode 100644 index 00000000..b9429220 --- /dev/null +++ b/crates/asap-physical-operators/src/expressions/mod.rs @@ -0,0 +1,255 @@ +//! Scalar semantics and typed expression binding. Planner expressions enter through CompiledExpression. +use crate::{ + values::{plain, Schema, Value}, + Error, +}; +use planner_types::pre_asap::{ArithmeticOpKind, DataType}; +pub mod arithmetic; +mod planner; +pub use planner::CompiledExpression; +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub enum Expression { + Binary { + operator: planner_types::post_asap::BinaryOperator, + left: Box, + right: Box, + }, + Planner(Box), + 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 { + pub fn planner(expression: crate::expressions::CompiledExpression) -> Self { + Self::Planner(Box::new(expression)) + } + pub(crate) fn dtype(&self, input: &Schema) -> Result<(DataType, bool), Error> { + use Expression::*; + match self { + Binary { + operator, + left, + right, + } => { + use planner_types::pre_asap::{BinaryOpKind, CompareOpKind}; + let (a, n) = left.dtype(input)?; + let (b, m) = right.dtype(input)?; + if a != DataType::Float64 || b != a || operator.vector_match.is_some() { + return Err(invalid( + "binary expression requires resolved Float64 operands", + )); + } + if (operator.checked_relative_division || operator.checked_finite_division) + && operator.kind != BinaryOpKind::Arithmetic(ArithmeticOpKind::Div) + { + return Err(invalid("checked division contract on non-division")); + } + let dtype = match operator.kind { + BinaryOpKind::Arithmetic(_) => DataType::Float64, + BinaryOpKind::Compare( + CompareOpKind::Eq + | CompareOpKind::Ne + | CompareOpKind::Lt + | CompareOpKind::Le + | CompareOpKind::Gt + | CompareOpKind::Ge, + ) => DataType::Bool, + _ => return Err(invalid("unsupported binary operation")), + }; + Ok((dtype, n || m)) + } + Planner(expression) => { + expression.validate_input(input)?; + Ok(expression.dtype()) + } + 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)) + } + } + } + pub(crate) fn evaluate(&self, row: &[Value]) -> Result { + use Expression::*; + Ok(match self { + Binary { + operator, + left, + right, + } => { + let (a, b) = (left.evaluate(row)?, right.evaluate(row)?); + if matches!(a, Value::Null) || matches!(b, Value::Null) { + Value::Null + } else { + let (Value::Float64(a), Value::Float64(b)) = (a, b) else { + return Err(invalid("binary value schema mismatch")); + }; + arithmetic::evaluate_binary(operator, a, b)? + } + } + Planner(expression) => expression.evaluate(row)?, + 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)), + }) + } +} +pub(crate) fn ordered(dtype: &DataType) -> bool { + if let DataType::Map { key, value, .. } = dtype { + return ordered(key) && ordered(value); + } + matches!( + dtype, + DataType::Null + | DataType::Int64 + | DataType::Float64 + | DataType::Utf8 + | DataType::Bool + | DataType::Timestamp + | DataType::Date + ) +} +pub(crate) 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(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")), + }) +} + +fn invalid(message: &str) -> Error { + Error::Invalid(message.into()) +} diff --git a/crates/asap-physical-operators/src/expressions/planner.rs b/crates/asap-physical-operators/src/expressions/planner.rs new file mode 100644 index 00000000..2130a2f7 --- /dev/null +++ b/crates/asap-physical-operators/src/expressions/planner.rs @@ -0,0 +1,549 @@ +//! Planner scalar expressions evaluated over native typed rows. +use crate::{ + values::{Schema, Value}, + Error, +}; +use planner_types::pre_asap::{ArithmeticOpKind, CompareOpKind, DataType, QueryExpr, ScalarValue}; +use std::{cmp::Ordering, sync::Arc}; + +pub(super) fn evaluate( + expr: &QueryExpr, + row: &[Value], + schema: &planner_types::pre_asap::Schema, +) -> Result { + match expr { + QueryExpr::Column(index) => row.get(*index).cloned().ok_or(Error::Invalid(format!( + "column {index} outside row width {}", + row.len() + ))), + QueryExpr::Literal(value) => Ok(match value { + ScalarValue::Interval { + months, + days, + nanos, + } => Value::Interval { + months: *months, + days: *days, + nanos: *nanos, + }, + ScalarValue::Int64(value) => Value::Int64(*value), + ScalarValue::Float64(value) => Value::Float64(*value), + ScalarValue::Utf8(value) => Value::Utf8(value.clone().into()), + ScalarValue::Boolean(value) => Value::Bool(*value), + ScalarValue::Null => Value::Null, + }), + QueryExpr::Compare { left, op, right } => { + let left = evaluate(left, row, schema)?; + let right = evaluate(right, row, schema)?; + compare(op, left, right) + } + QueryExpr::Arithmetic { op, left, right } => arithmetic( + op, + evaluate(left, row, schema)?, + evaluate(right, row, schema)?, + ), + QueryExpr::BoolAnd(parts) | QueryExpr::BoolOr(parts) => { + let and = matches!(expr, QueryExpr::BoolAnd(_)); + let mut null = false; + for part in parts { + match evaluate(part, row, schema)? { + Value::Bool(value) if value != and => return Ok(Value::Bool(value)), + Value::Bool(_) => {} + Value::Null => null = true, + _ => return Err(Error::Invalid("boolean predicate required".into())), + } + } + Ok(if null { Value::Null } else { Value::Bool(and) }) + } + QueryExpr::Not(value) => match evaluate(value, row, schema)? { + Value::Bool(value) => Ok(Value::Bool(!value)), + Value::Null => Ok(Value::Null), + _ => Err(Error::Invalid("boolean predicate required".into())), + }, + QueryExpr::IsNull(value) => Ok(Value::Bool(matches!( + evaluate(value, row, schema)?, + Value::Null + ))), + QueryExpr::IsNotNull(value) => Ok(Value::Bool(!matches!( + evaluate(value, row, schema)?, + Value::Null + ))), + QueryExpr::FunctionCall { name, args } => { + use planner_types::pre_asap::scalar_signature::MapScalarFunction; + if name.eq_ignore_ascii_case("asap_struct_field") { + expr.scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))?; + let DataType::Struct { fields } = args[0] + .scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))? + .0 + else { + unreachable!() + }; + let offset = match &args[1] { + QueryExpr::Literal(ScalarValue::Int64(index)) => { + usize::try_from(index - 1).ok() + } + QueryExpr::Literal(ScalarValue::Utf8(name)) => { + fields.iter().position(|field| &field.name == name) + } + _ => None, + } + .ok_or_else(|| Error::Invalid("struct field selector".into()))?; + let Value::Struct(values) = evaluate(&args[0], row, schema)? else { + return Err(Error::Invalid("struct field input".into())); + }; + return values + .get(offset) + .cloned() + .ok_or_else(|| Error::Invalid("struct field value".into())); + } + if name.eq_ignore_ascii_case("asap_element_access") { + let (output_type, _) = expr + .scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))?; + if let DataType::List { element } = args[0] + .scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))? + .0 + { + let Value::List(values) = evaluate(&args[0], row, schema)? else { + return Err(Error::Invalid("array access input".into())); + }; + let index = match evaluate(&args[1], row, schema)? { + Value::Null => return Ok(Value::Null), + Value::Int64(index) => index, + _ => return Err(Error::Invalid("array access index".into())), + }; + let offset = if index > 0 { + usize::try_from(index - 1).ok() + } else if index < 0 { + usize::try_from(index.unsigned_abs()) + .ok() + .and_then(|distance| values.len().checked_sub(distance)) + } else { + None + }; + return match offset.and_then(|offset| values.get(offset)) { + Some(value) => Ok(value.clone()), + None => default_collection_element(&output_type, element.nullable), + }; + } + } + let function = (if name.eq_ignore_ascii_case("asap_element_access") { + Some(MapScalarFunction::Access) + } else { + MapScalarFunction::from_name(name) + }) + .ok_or_else(|| Error::Invalid(format!("scalar function {name}")))?; + expr.scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))?; + let values = args + .iter() + .map(|arg| evaluate(arg, row, schema)) + .collect::, _>>()?; + match function { + MapScalarFunction::Construct => { + let mut values = values.into_iter(); + let mut entries = Vec::new(); + while let Some(key) = values.next() { + if !matches!(key, Value::Int64(_) | Value::Utf8(_) | Value::Bool(_)) { + return Err(Error::Invalid("map key value type".into())); + } + entries.push(( + key, + values + .next() + .ok_or_else(|| Error::Invalid("odd map argument count".into()))?, + )); + } + Ok(Value::Map(entries.into())) + } + MapScalarFunction::Concat => { + let mut entries = Vec::new(); + for value in values { + let Value::Map(next) = value else { + return Err(Error::Invalid("map concat argument".into())); + }; + entries.extend(next.iter().cloned()); + } + Ok(Value::Map(entries.into())) + } + MapScalarFunction::Access => { + let [Value::Map(entries), key] = values.as_slice() else { + return Err(Error::Invalid("map access arguments".into())); + }; + if matches!(key, Value::Null) { + return Ok(Value::Null); + } + if !matches!(key, Value::Int64(_) | Value::Utf8(_) | Value::Bool(_)) { + return Err(Error::Invalid("map lookup key type".into())); + } + if let Some((_, value)) = entries + .iter() + .find(|(candidate, _)| cell_cmp(candidate, key) == Some(Ordering::Equal)) + { + return Ok(value.clone()); + } + let ( + DataType::Map { + value, + value_nullable, + .. + }, + _, + ) = args[0] + .scalar_type(schema) + .map_err(|error| Error::Invalid(error.to_string()))? + else { + unreachable!() + }; + default_collection_element(&value, value_nullable) + } + } + } + other => Err(Error::Invalid(format!("scalar expression {other:?}"))), + } +} + +fn default_collection_element(dtype: &DataType, nullable: bool) -> Result { + if nullable { + return Ok(Value::Null); + } + Ok(match dtype { + DataType::Interval | DataType::Date => { + return Err(Error::Invalid("temporal value transport".into())) + } + DataType::Null => Value::Null, + DataType::Int64 => Value::Int64(0), + DataType::Float64 => Value::Float64(0.0), + DataType::Utf8 => Value::Utf8("".into()), + DataType::Bool => Value::Bool(false), + DataType::Map { .. } => Value::Map(Arc::from([])), + DataType::List { .. } => Value::List(Arc::from([])), + DataType::Struct { fields } => Value::Struct( + fields + .iter() + .map(|field| default_collection_element(&field.dtype, field.nullable)) + .collect::, _>>()? + .into(), + ), + _ => { + return Err(Error::Invalid( + "collection missing-element default type".into(), + )) + } + }) +} + +fn compare(op: &CompareOpKind, left: Value, right: Value) -> Result { + if matches!(left, Value::Null) || matches!(right, Value::Null) { + return Ok(Value::Null); + } + // NaN is unordered, not a type mismatch. Match the native scalar path. + if matches!(&left, Value::Float64(v) if v.is_nan()) + || matches!(&right, Value::Float64(v) if v.is_nan()) + { + return match op { + CompareOpKind::Ne => Ok(Value::Bool(true)), + CompareOpKind::Eq + | CompareOpKind::Lt + | CompareOpKind::Le + | CompareOpKind::Gt + | CompareOpKind::Ge => Ok(Value::Bool(false)), + _ => Err(Error::Invalid(format!("comparison {op:?}"))), + }; + } + let ordering = cell_cmp(&left, &right) + .ok_or_else(|| Error::Invalid("comparison of incompatible values".into()))?; + let value = match op { + CompareOpKind::Eq => ordering == Ordering::Equal, + CompareOpKind::Ne => ordering != Ordering::Equal, + CompareOpKind::Lt => ordering == Ordering::Less, + CompareOpKind::Le => ordering != Ordering::Greater, + CompareOpKind::Gt => ordering == Ordering::Greater, + CompareOpKind::Ge => ordering != Ordering::Less, + _ => return Err(Error::Invalid(format!("comparison {op:?}"))), + }; + Ok(Value::Bool(value)) +} + +fn arithmetic(op: &ArithmeticOpKind, left: Value, right: Value) -> Result { + let (left, right) = match (left, right) { + (Value::Int64(a), Value::Float64(b)) => (Value::Float64(a as f64), Value::Float64(b)), + (Value::Float64(a), Value::Int64(b)) => (Value::Float64(a), Value::Float64(b as f64)), + pair => pair, + }; + super::numeric(op, left, right) +} + +fn integer_float_cmp(integer: i64, float: f64) -> Option { + if float.is_nan() { + return None; + } + // These bounds are powers of two, exactly representable as Float64. + if float >= 9_223_372_036_854_775_808.0 { + return Some(Ordering::Less); + } + if float < -9_223_372_036_854_775_808.0 { + return Some(Ordering::Greater); + } + let integral = float as i64; + match integer.cmp(&integral) { + Ordering::Equal => 0.0_f64.partial_cmp(&float.fract()), + other => Some(other), + } +} + +fn cell_cmp(left: &Value, right: &Value) -> Option { + match (left, right) { + (Value::Int64(left), Value::Int64(right)) => Some(left.cmp(right)), + (Value::Float64(left), Value::Float64(right)) => left.partial_cmp(right), + (Value::Int64(left), Value::Float64(right)) => integer_float_cmp(*left, *right), + (Value::Float64(left), Value::Int64(right)) => { + integer_float_cmp(*right, *left).map(Ordering::reverse) + } + (Value::Utf8(left), Value::Utf8(right)) => Some(left.cmp(right)), + (Value::Bool(left), Value::Bool(right)) => Some(left.cmp(right)), + (Value::Timestamp(left), Value::Timestamp(right)) => Some(left.cmp(right)), + (Value::Map(left), Value::Map(right)) => { + for ((left_key, left_value), (right_key, right_value)) in left.iter().zip(right.iter()) + { + let order = cell_cmp(left_key, right_key)?; + if order != Ordering::Equal { + return Some(order); + } + let order = match (left_value, right_value) { + (Value::Null, Value::Null) => Ordering::Equal, + (Value::Null, _) => Ordering::Greater, + (_, Value::Null) => Ordering::Less, + _ => cell_cmp(left_value, right_value)?, + }; + if order != Ordering::Equal { + return Some(order); + } + } + Some(left.len().cmp(&right.len())) + } + _ => None, + } +} + +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub struct CompiledExpression { + expression: QueryExpr, + schema: planner_types::pre_asap::Schema, + output: (DataType, bool), +} +impl CompiledExpression { + pub(crate) fn expression(&self) -> &QueryExpr { + &self.expression + } + + pub fn compile(expression: &QueryExpr, input: &Schema) -> Result { + let schema = input + .fields + .iter() + .map(|field| { + let planner_types::post_asap::SummaryFamilyType::Plain(dtype) = &field.dtype else { + return Err(Error::Invalid( + "scalar expression cannot consume opaque summary state".into(), + )); + }; + Ok(planner_types::pre_asap::Column::new( + field.name.clone(), + dtype.clone(), + field.nullable, + )) + }) + .collect::, Error>>()?; + let schema = planner_types::pre_asap::Schema::new(schema); + validate(expression, &schema)?; + let output = expression + .scalar_type(&schema) + .map_err(|e| Error::Invalid(e.to_string()))?; + Ok(Self { + expression: expression.clone(), + schema, + output, + }) + } + pub(crate) fn dtype(&self) -> (DataType, bool) { + self.output.clone() + } + pub(crate) fn validate_input(&self, input: &Schema) -> Result<(), Error> { + let checked = Self::compile(&self.expression, input)?; + if checked.output != self.output { + return Err(Error::Invalid( + "persisted expression type differs from its semantics".into(), + )); + } + if input.fields.len() != self.schema.columns.len() + || input + .fields + .iter() + .zip(&self.schema.columns) + .any(|(field, column)| { + field.dtype + != planner_types::post_asap::SummaryFamilyType::Plain(column.dtype.clone()) + || field.nullable != column.nullable + }) + { + return Err(Error::Invalid( + "expression input differs from its bound schema".into(), + )); + } + Ok(()) + } + /// Evaluate a row under the same typed schema used when binding the expression. + pub fn evaluate(&self, row: &[Value]) -> Result { + if row.len() != self.schema.columns.len() + || row + .iter() + .zip(&self.schema.columns) + .any(|(value, column)| !value.matches(&column.dtype, column.nullable)) + { + return Err(Error::Invalid( + "expression input differs from its bound schema".into(), + )); + } + evaluate(&self.expression, row, &self.schema) + } +} +fn validate(expr: &QueryExpr, schema: &planner_types::pre_asap::Schema) -> Result<(), Error> { + let invalid = || Error::Invalid(format!("unsupported scalar expression: {expr:?}")); + expr.scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))?; + match expr { + QueryExpr::Column(_) | QueryExpr::Literal(_) => Ok(()), + QueryExpr::Arithmetic { left, right, .. } => { + for value in [left, right] { + validate(value, schema)?; + if !matches!( + value + .scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))? + .0, + DataType::Int64 | DataType::Float64 | DataType::Null + ) { + return Err(invalid()); + } + } + Ok(()) + } + QueryExpr::Compare { left, right, op } => { + if !matches!( + op, + CompareOpKind::Eq + | CompareOpKind::Ne + | CompareOpKind::Lt + | CompareOpKind::Le + | CompareOpKind::Gt + | CompareOpKind::Ge + ) { + return Err(invalid()); + } + validate(left, schema)?; + validate(right, schema)?; + let (a, _) = left + .scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))?; + let (b, _) = right + .scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))?; + fn comparable(dtype: &DataType) -> bool { + match dtype { + DataType::Null + | DataType::Int64 + | DataType::Float64 + | DataType::Utf8 + | DataType::Bool + | DataType::Timestamp => true, + DataType::Map { key, value, .. } => comparable(key) && comparable(value), + _ => false, + } + } + let numeric = |dtype: &DataType| matches!(dtype, DataType::Int64 | DataType::Float64); + if !comparable(&a) + || !comparable(&b) + || (a != b + && !matches!(a, DataType::Null) + && !matches!(b, DataType::Null) + && !(numeric(&a) && numeric(&b))) + { + return Err(invalid()); + } + Ok(()) + } + QueryExpr::FunctionCall { name, args } => { + if name != "asap_struct_field" + && name != "asap_element_access" + && planner_types::pre_asap::scalar_signature::MapScalarFunction::from_name(name) + .is_none() + { + return Err(invalid()); + } + for arg in args { + validate(arg, schema)?; + } + Ok(()) + } + QueryExpr::BoolAnd(parts) | QueryExpr::BoolOr(parts) => { + for part in parts { + validate(part, schema)?; + if !matches!( + part.scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))? + .0, + DataType::Bool | DataType::Null + ) { + return Err(invalid()); + } + } + Ok(()) + } + QueryExpr::Not(value) => { + validate(value, schema)?; + if !matches!( + value + .scalar_type(schema) + .map_err(|e| Error::Invalid(e.to_string()))? + .0, + DataType::Bool | DataType::Null + ) { + return Err(invalid()); + } + Ok(()) + } + QueryExpr::IsNull(value) | QueryExpr::IsNotNull(value) => validate(value, schema), + _ => Err(invalid()), + } +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn mixed_comparison_preserves_integer_precision_and_boundaries() { + assert_eq!( + integer_float_cmp(9_007_199_254_740_993, 9_007_199_254_740_992.0), + Some(Ordering::Greater) + ); + assert_eq!( + integer_float_cmp(i64::MAX, 9_223_372_036_854_775_808.0), + Some(Ordering::Less) + ); + assert_eq!( + integer_float_cmp(i64::MIN, -9_223_372_036_854_775_808.0), + Some(Ordering::Equal) + ); + assert_eq!(integer_float_cmp(-1, -1.5), Some(Ordering::Greater)); + assert_eq!(integer_float_cmp(1, 1.5), Some(Ordering::Less)); + assert_eq!(integer_float_cmp(0, f64::INFINITY), Some(Ordering::Less)); + assert_eq!( + integer_float_cmp(0, f64::NEG_INFINITY), + Some(Ordering::Greater) + ); + assert_eq!(integer_float_cmp(0, f64::NAN), None); + } +} 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 00000000..34bc8489 --- /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 00000000..5709a586 --- /dev/null +++ b/crates/asap-physical-operators/src/lib.rs @@ -0,0 +1,37 @@ +#![doc = include_str!("../README.md")] + +pub mod summary_kernels; +/// Compatibility alias for existing deployments. +pub use summary_kernels as accumulators; +pub mod key_by_label_values; +pub mod measurement; +pub use summary_kernels::traits; + +mod aggregation_type; +mod statistic; +pub use aggregation_type::AggregationType; +pub use key_by_label_values::KeyByLabelValues; +pub use measurement::Measurement; +pub use statistic::Statistic; +pub use traits::*; + +pub use expressions::arithmetic; +pub mod capability; +pub use summary_kernels::factory; + +/// The exact Planner contract used by these kernels. +pub use planner_types as planner; + +pub mod dag; + +pub mod stored_state; + +mod error; +pub use error::Error; +pub mod expressions; +pub mod operators; +pub mod physical_planner; +pub mod plan; +pub mod runtime; +pub mod sources; +pub mod values; diff --git a/crates/asap-physical-operators/src/measurement.rs b/crates/asap-physical-operators/src/measurement.rs new file mode 100644 index 00000000..0fe1abc0 --- /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/operators/aggregate/mod.rs b/crates/asap-physical-operators/src/operators/aggregate/mod.rs new file mode 100644 index 00000000..c57b4fce --- /dev/null +++ b/crates/asap-physical-operators/src/operators/aggregate/mod.rs @@ -0,0 +1,293 @@ +use super::*; +impl Operator { + 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, nullable) = plain(&input, *i)?; + if !ordered(t) { + return Err(invalid("ordered aggregate input required")); + } + (t.clone(), nullable || groups.is_empty()) + } + }; + 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 window( + input: Schema, + intent: planner_types::pre_asap::AggIntent, + coordinate: usize, + value: usize, + groups: Vec, + window: Option<(i64, i64)>, + ) -> Result { + use planner_types::pre_asap::AggIntent; + validate_groups(&input, &groups)?; + let histogram = matches!(intent, AggIntent::HistogramQuantile { .. }); + if !matches!( + intent, + AggIntent::Rate + | AggIntent::Increase + | AggIntent::Count { .. } + | AggIntent::Sum { col: None } + | AggIntent::Avg { col: None } + | AggIntent::Min { col: None } + | AggIntent::Max { col: None } + | AggIntent::HistogramQuantile { .. } + ) { + return Err(invalid( + "unsupported temporal intent or unresolved value column", + )); + } + let coordinate_type = if histogram { + DataType::Float64 + } else { + DataType::Timestamp + }; + if plain(&input, coordinate)? != (&coordinate_type, false) + || plain(&input, value)? != (&DataType::Float64, false) + { + return Err(invalid("window coordinate/value schema mismatch")); + } + if (!histogram && !matches!(window, Some((start, end)) if start < end)) + || (histogram && window.is_some()) + { + return Err(invalid("invalid temporal window")); + } + let mut fields = groups + .iter() + .map(|i| input.fields[*i].clone()) + .collect::>(); + fields.push(result_field( + "value", + if matches!(intent, AggIntent::Count { .. }) { + DataType::Int64 + } else { + DataType::Float64 + }, + false, + )); + Ok(Self { + kind: Kind::Window { + intent: Box::new(intent), + coordinate, + value, + groups, + window, + }, + inputs: vec![input], + output: schema(fields), + }) + } +} +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub enum Reduction { + Count, + Sum(usize), + Avg(usize), + Min(usize), + Max(usize), +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let result = match &operator.kind { + Kind::Window { + intent, + coordinate, + value, + groups, + window, + } => { + crate::operators::aggregate::temporal::reduce( + rows, + intent, + groups, + *coordinate, + *value, + *window, + &context, + ) + .await? + } + Kind::Aggregate { groups, measures } => { + reduce(rows, groups, measures, &operator.inputs[0], &context).await? + } + _ => unreachable!(), + }; + Batch::try_new(output, result) + }) + .boxed_local()) +} + +mod temporal; +async fn reduce( + rows: Vec>, + groups: &[usize], + measures: &[Reduction], + input: &Schema, + context: &RunContext, +) -> Result>, Error> { + let mut work = Cooperative::new(context); + let mut workspace = Workspace::new(context)?; + let mut grouped = BTreeMap::>, Vec>>::new(); + if rows.is_empty() && groups.is_empty() { + grouped.insert(vec![], vec![]); + } + for row in rows { + work.checkpoint().await?; + let key = group_key(&row, groups)?; + workspace.grow(std::mem::size_of::>())?; + if !grouped.contains_key(&key) { + workspace.grow(key_bytes(&key))?; + } + grouped.entry(key).or_default().push(row); + } + let mut output = Vec::new(); + for rows in grouped.into_values() { + work.checkpoint().await?; + let mut result = groups + .iter() + .map(|&i| rows[0][i].clone()) + .collect::>(); + for measure in measures { + result.push(reduce_one(&rows, measure, input, &mut work).await?); + } + workspace.grow(row_bytes(&result))?; + output.push(result); + } + Ok(output) +} + +async fn reduce_one( + rows: &[Vec], + measure: &Reduction, + input: &Schema, + work: &mut Cooperative, +) -> 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)); + 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 { + work.checkpoint().await?; + 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 { + work.checkpoint().await?; + 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 mut count = 0usize; + let dtype = plain(input, column)?.0; + if dtype == &DataType::Int64 { + let mut sum = 0i128; + for v in values { + work.checkpoint().await?; + let Value::Int64(v) = v else { + return Err(invalid("integer aggregate value required")); + }; + sum = sum + .checked_add(i128::from(*v)) + .ok_or_else(|| invalid("integer aggregate overflow"))?; + count += 1; + } + 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 mut sum = -0.0; + for v in values { + work.checkpoint().await?; + let Value::Float64(v) = v else { + return Err(invalid("floating aggregate value required")); + }; + sum += v; + count += 1; + } + Ok(Value::Float64(if matches!(measure, Reduction::Avg(_)) { + sum / count as f64 + } else { + sum + })) +} diff --git a/crates/asap-physical-operators/src/operators/aggregate/temporal.rs b/crates/asap-physical-operators/src/operators/aggregate/temporal.rs new file mode 100644 index 00000000..fb74965a --- /dev/null +++ b/crates/asap-physical-operators/src/operators/aggregate/temporal.rs @@ -0,0 +1,299 @@ +//! Windowed computations use Planner intents; deployments supply the input window. +use crate::{ + operators::{ + common::{key_bytes, row_bytes, Workspace}, + sort::cooperative_sort, + }, + runtime::{Cooperative, RunContext}, +}; +use crate::{ + values::{group_key, Value}, + Error, +}; +use planner_types::pre_asap::{AggIntent, ColumnRef}; +use std::collections::BTreeMap; + +pub(super) async fn reduce( + rows: Vec>, + intent: &AggIntent, + groups: &[usize], + coordinate: usize, + value: usize, + window: Option<(i64, i64)>, + context: &RunContext, +) -> Result>, Error> { + let mut work = Cooperative::new(context); + let mut workspace = Workspace::new(context)?; + let mut grouped = + BTreeMap::>, (Vec, Vec<(f64, f64)>, Vec<(i64, f64)>)>::new(); + for row in rows { + work.checkpoint().await?; + let key = group_key(&row, groups)?; + workspace.grow(32)?; + if !grouped.contains_key(&key) { + workspace.grow(key_bytes(&key) + row_bytes(&row))?; + } + let entry = grouped.entry(key).or_insert_with(|| { + ( + groups.iter().map(|i| row[*i].clone()).collect(), + vec![], + vec![], + ) + }); + let Value::Float64(v) = row[value] else { + return Err(Error::Invalid("window value must be Float64".into())); + }; + match row[coordinate] { + Value::Timestamp(t) => entry.2.push((t, v)), + Value::Float64(bound) => entry.1.push((bound, v)), + _ => return Err(Error::Invalid("invalid window coordinate".into())), + } + } + let mut output = Vec::new(); + for (_, (mut keys, buckets, points)) in grouped { + work.checkpoint().await?; + let result = if let AggIntent::HistogramQuantile { q } = intent { + Some(Value::Float64(bucket_quantile(*q, buckets, context).await?)) + } else { + let points = cooperative_sort(points, |a, b| a.0.cmp(&b.0), context).await?; + let (start, end) = + window.ok_or_else(|| Error::Invalid("missing temporal window".into()))?; + if points.iter().any(|p| p.0 < start || p.0 > end) + || points.windows(2).any(|p| p[0].0 == p[1].0) + { + return Err(Error::Invalid( + "duplicate or out-of-window timestamp".into(), + )); + } + match intent { + AggIntent::Rate => rate(&points, start, end).map(Value::Float64), + AggIntent::Increase => rate(&points, start, end) + .map(|v| Value::Float64(v * (end as f64 - start as f64) / 1000.)), + AggIntent::Count { .. } => Some(Value::Int64( + i64::try_from(points.len()) + .map_err(|_| Error::Invalid("count overflow".into()))?, + )), + AggIntent::Sum { .. } => Some(Value::Float64(points.iter().map(|p| p.1).sum())), + AggIntent::Avg { .. } => Some(Value::Float64( + points.iter().map(|p| p.1).sum::() / points.len() as f64, + )), + AggIntent::Min { .. } => { + Some(Value::Float64(points.iter().fold(f64::NAN, |a, p| { + if a.is_nan() || p.1 < a { + p.1 + } else { + a + } + }))) + } + AggIntent::Max { .. } => { + Some(Value::Float64(points.iter().fold(f64::NAN, |a, p| { + if a.is_nan() || p.1 > a { + p.1 + } else { + a + } + }))) + } + _ => return Err(Error::Invalid("unsupported temporal intent".into())), + } + }; + if let Some(result) = result { + keys.push(result); + output.push(keys); + } + } + Ok(output) +} + +fn rate(points: &[(i64, f64)], start: i64, end: i64) -> Option { + if points.len() < 2 { + return None; + } + let (first_t, first) = points[0]; + let (last_t, last) = *points.last()?; + let span = (last_t as f64 - first_t as f64) / 1000.; + if span <= 0. { + return None; + } + let mut delta = last - first; + for pair in points.windows(2) { + if pair[1].1 < pair[0].1 { + delta += pair[0].1; + } + } + let average = span / (points.len() - 1) as f64; + let mut to_start = (first_t as f64 - start as f64) / 1000.; + let mut to_end = (end as f64 - last_t as f64) / 1000.; + if to_start >= average * 1.1 { + to_start = average / 2.; + } + // Apply the zero bound after the sparse-window half-interval cap. + if delta > 0. && first >= 0. { + to_start = to_start.min(span * first / delta); + } + if to_end >= average * 1.1 { + to_end = average / 2.; + } + Some(delta * (span + to_start + to_end) / span / ((end as f64 - start as f64) / 1000.)) +} + +async fn bucket_quantile( + q: f64, + mut b: Vec<(f64, f64)>, + context: &RunContext, +) -> Result { + let mut work = Cooperative::new(context); + let _scratch = context.reserve(b.len().checked_mul(16).ok_or(Error::MemoryLimit)?)?; + if q.is_nan() { + return Ok(f64::NAN); + } + if q < 0. { + return Ok(f64::NEG_INFINITY); + } + if q > 1. { + return Ok(f64::INFINITY); + } + b.retain(|p| !p.0.is_nan()); + b = cooperative_sort(b, |a, b| a.0.total_cmp(&b.0), context).await?; + let mut buckets: Vec<(f64, f64)> = Vec::new(); + for p in b { + work.checkpoint().await?; + if let Some(last) = buckets.last_mut() { + if last.0 == p.0 { + last.1 += p.1; + continue; + } + } + buckets.push(p); + } + if buckets.len() < 2 || buckets.last().unwrap().0 != f64::INFINITY { + return Ok(f64::NAN); + } + let mut prev = buckets[0].1; + for p in buckets.iter_mut().skip(1) { + work.checkpoint().await?; + if p.1 < prev || (p.1 - prev).abs() <= 1e-12 * (p.1.abs() + prev.abs()) { + p.1 = prev; + } + prev = p.1; + } + let count = buckets.last().unwrap().1; + if count == 0. { + return Ok(f64::NAN); + } + let rank = q * count; + let idx = buckets[..buckets.len() - 1].partition_point(|p| p.1 < rank); + if idx == buckets.len() - 1 { + return Ok(buckets[idx - 1].0); + } + if idx == 0 && buckets[0].0 <= 0. { + return Ok(buckets[0].0); + } + let (start, base) = if idx == 0 { (0., 0.) } else { buckets[idx - 1] }; + let (end, upper) = buckets[idx]; + Ok(start + (end - start) * (rank - base) / (upper - base)) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::dag::{ + batch_execution::evaluate_batch, operators::Operator, values::Batch, Limits, RunContext, + Scope, + }; + use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, + types::AccuracyTarget, + }; + use std::sync::Arc; + + // The same window operator must give the same answer in either engine phase. + #[test] + fn temporal_windows_execute_in_both_phases_and_count_is_integer() { + let schema = Arc::new(SummarySchema { + fields: vec![ + SummaryField { + name: "time".into(), + dtype: SummaryFamilyType::Plain(DataType::Timestamp), + nullable: false, + }, + SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }, + ], + time_index: Some(0), + }); + for scope in [ + Scope::Query { + evaluation_time_ms: 2000, + revision: 1, + }, + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 2000, + revision: 1, + }, + ] { + for (intent, expected) in [ + (AggIntent::Rate, Value::Float64(2.)), + (AggIntent::Increase, Value::Float64(4.)), + ( + AggIntent::Count { + accuracy: AccuracyTarget::Exact, + }, + Value::Int64(3), + ), + ] { + let batch = Batch::try_new( + schema.clone(), + vec![ + vec![Value::Timestamp(0), Value::Float64(2.)], + vec![Value::Timestamp(1000), Value::Float64(4.)], + vec![Value::Timestamp(2000), Value::Float64(2.)], + ], + ) + .unwrap(); + let operator = + Operator::window(schema.clone(), intent, 0, 1, vec![], Some((0, 2000))) + .unwrap(); + let result = evaluate_batch( + batch, + vec![operator], + RunContext::new(scope.clone(), Limits::default()).unwrap(), + ) + .unwrap(); + assert_eq!( + format!("{:?}", result[0].rows()[0][0]), + format!("{expected:?}") + ); + } + } + assert!(Operator::window(schema, AggIntent::Rate, 0, 1, vec![], Some((1, 1))).is_err()); + } + + // Histogram interpolation requires an infinite terminal bucket and coalesces duplicates. + #[test] + fn histogram_boundaries_and_duplicate_buckets() { + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + let bucket_quantile = |q, buckets| { + futures::executor::block_on(super::bucket_quantile(q, buckets, &context)).unwrap() + }; + assert_eq!( + bucket_quantile(0.5, vec![(1., 1.), (1., 1.), (2., 4.), (f64::INFINITY, 4.)]), + 1. + ); + assert!(bucket_quantile(0.5, vec![(1., 2.), (2., 4.)]).is_nan()); + assert_eq!(bucket_quantile(-0.1, vec![]), f64::NEG_INFINITY); + } +} diff --git a/crates/asap-physical-operators/src/operators/common.rs b/crates/asap-physical-operators/src/operators/common.rs new file mode 100644 index 00000000..9f1709da --- /dev/null +++ b/crates/asap-physical-operators/src/operators/common.rs @@ -0,0 +1,75 @@ +use super::*; +pub(super) fn invalid(message: &str) -> Error { + Error::Invalid(message.into()) +} +pub(super) fn schema(fields: Vec) -> Schema { + Arc::new(SummarySchema { + fields, + time_index: None, + }) +} +pub(super) fn result_field(name: &str, dtype: DataType, nullable: bool) -> SummaryField { + SummaryField { + name: name.into(), + dtype: SummaryFamilyType::Plain(dtype), + nullable, + } +} + +pub(super) 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(()) +} +pub(super) async fn collect_rows( + mut input: Input<'_, Batch>, + context: &RunContext, +) -> Result<(Vec>, Vec), Error> { + let mut rows = Vec::new(); + let mut work = Cooperative::new(context); + let mut reservations = Vec::new(); + while let Some(batch) = input.next().await { + let batch = batch?; + reservations.push(context.reserve(batch.bytes())?); + for row in batch.rows() { + work.checkpoint().await?; + rows.push(row.clone()); + } + } + Ok((rows, reservations)) +} +/// Estimates retained workspace before growing collections. It is not an RSS limit. +pub(super) struct Workspace { + reservation: Reservation, + bytes: usize, +} +impl Workspace { + pub(super) fn new(context: &RunContext) -> Result { + Ok(Self { + reservation: context.reserve(0)?, + bytes: 0, + }) + } + pub(super) fn grow(&mut self, bytes: usize) -> Result<(), Error> { + self.bytes = self.bytes.checked_add(bytes).ok_or(Error::MemoryLimit)?; + self.reservation.resize(self.bytes) + } +} +pub(super) fn row_bytes(row: &[Value]) -> usize { + std::mem::size_of::>() + row.iter().map(Value::bytes).sum::() +} +pub(super) fn key_bytes(key: &[Vec]) -> usize { + 64 + key + .iter() + .map(|part| std::mem::size_of::>() + part.len()) + .sum::() +} diff --git a/crates/asap-physical-operators/src/operators/current_series.rs b/crates/asap-physical-operators/src/operators/current_series.rs new file mode 100644 index 00000000..937036c5 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/current_series.rs @@ -0,0 +1,135 @@ +//! A bounded instant-vector snapshot: select latest before removing stale markers. +use super::*; + +impl Operator { + pub fn current_series( + input: Schema, + identity: usize, + coordinate: usize, + value: usize, + lookback_ms: i64, + ) -> Result { + if lookback_ms <= 0 + || plain(&input, identity)? != (&DataType::Utf8, false) + || plain(&input, coordinate)? != (&DataType::Timestamp, false) + || plain(&input, value)? != (&DataType::Float64, false) + { + return Err(invalid("invalid current-series input contract")); + } + Ok(Self { + kind: Kind::CurrentSeries { + identity, + coordinate, + value, + lookback_ms, + }, + inputs: vec![input.clone()], + output: input, + }) + } +} + +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let Kind::CurrentSeries { + identity, + coordinate, + value, + lookback_ms, + } = operator.kind + else { + unreachable!() + }; + let input = inputs + .pop() + .ok_or_else(|| invalid("current-series input missing"))?; + let output = operator.output.clone(); + let (start, end) = window(lookback_ms, &context)?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let mut latest = BTreeMap::, usize>::new(); + let mut work = Cooperative::new(&context); + let mut workspace = Workspace::new(&context)?; + for (index, row) in rows.iter().enumerate() { + work.checkpoint().await?; + let Value::Timestamp(timestamp) = row[coordinate] else { + unreachable!() + }; + if timestamp <= start || timestamp > end { + continue; + } + let key = row[identity].key()?; + if let Some(&previous) = latest.get(&key) { + let Value::Timestamp(previous_time) = rows[previous][coordinate] else { + unreachable!() + }; + if timestamp < previous_time { + continue; + } + if timestamp == previous_time { + let (Value::Float64(a), Value::Float64(b)) = + (&row[value], &rows[previous][value]) + else { + unreachable!() + }; + if a.to_bits() != b.to_bits() { + return Err(invalid("conflicting samples for one series timestamp")); + } + continue; + } + } else { + workspace.grow(64 + key.len())?; + } + latest.insert(key, index); + } + let mut result = Vec::new(); + for index in latest.into_values() { + work.checkpoint().await?; + let Value::Float64(sample) = rows[index][value] else { + unreachable!() + }; + if sample.to_bits() == 0x7ff0_0000_0000_0002 { + continue; + } + workspace.grow(row_bytes(&rows[index]))?; + let mut row = rows[index].clone(); + row[coordinate] = Value::Timestamp(end); + result.push(row); + } + Batch::try_new(output, result) + }) + .boxed_local()) +} + +fn window(lookback_ms: i64, context: &RunContext) -> Result<(i64, i64), Error> { + let end = match context.scope { + crate::runtime::Scope::Query { + evaluation_time_ms, .. + } => evaluation_time_ms, + crate::runtime::Scope::Ingestion { window_end_ms, .. } => window_end_ms, + }; + let start = end + .checked_sub(lookback_ms) + .ok_or_else(|| invalid("current-series window overflows"))?; + if let crate::runtime::Scope::Ingestion { + window_start_ms, .. + } = context.scope + { + if window_start_ms != start { + return Err(invalid( + "current-series maintenance window differs from lookback", + )); + } + } + Ok((start, end)) +} + +pub(super) fn validate_context(operator: &Operator, context: &RunContext) -> Result<(), Error> { + if let Kind::CurrentSeries { lookback_ms, .. } = operator.kind { + window(lookback_ms, context)?; + } + Ok(()) +} diff --git a/crates/asap-physical-operators/src/operators/filter.rs b/crates/asap-physical-operators/src/operators/filter.rs new file mode 100644 index 00000000..8862aacc --- /dev/null +++ b/crates/asap-physical-operators/src/operators/filter.rs @@ -0,0 +1,39 @@ +use super::*; +impl Operator { + 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(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + Kind::Filter(predicate) => Ok(input + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + 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()), + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/operators/joins/mod.rs b/crates/asap-physical-operators/src/operators/joins/mod.rs new file mode 100644 index 00000000..1fece7c1 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/joins/mod.rs @@ -0,0 +1,187 @@ +use super::*; +impl Operator { + 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 relational_join( + left: Schema, + right: Schema, + kind: planner_types::pre_asap::JoinKind, + predicate: &planner_types::pre_asap::Predicate, + output: Schema, + ) -> Result { + use planner_types::pre_asap::JoinKind; + let mut joined = left.fields.clone(); + joined.extend(right.fields.clone()); + let predicate = + crate::expressions::CompiledExpression::compile(&predicate.0, &schema(joined.clone()))?; + if predicate.dtype().0 != DataType::Bool { + return Err(invalid("join predicate must be boolean")); + } + let fields = if matches!(kind, JoinKind::Semi | JoinKind::Anti) { + left.fields.clone() + } else { + for field in &mut joined[..left.fields.len()] { + if matches!(kind, JoinKind::Right | JoinKind::Full) { + field.nullable = true; + } + } + for field in &mut joined[left.fields.len()..] { + if matches!(kind, JoinKind::Left | JoinKind::Full) { + field.nullable = true; + } + } + joined + }; + Self { + kind: Kind::Join { + kind, + predicate: Box::new(predicate), + }, + inputs: vec![left, right], + output: schema(fields), + } + .with_output_schema(output) + } +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + if let Kind::Join { kind, predicate } = &operator.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 { + use planner_types::pre_asap::JoinKind; + let ((left, _left_memory), (right, _right_memory)) = + futures::try_join!(collect_rows(left, &context), collect_rows(right, &context))?; + let mut workspace = Workspace::new(&context)?; + let mut work = Cooperative::new(&context); + workspace.grow(right.len())?; + let mut result = Vec::new(); + let mut right_matched = vec![false; right.len()]; + for left_row in &left { + work.checkpoint().await?; + let mut matched = false; + for (i, right_row) in right.iter().enumerate() { + work.checkpoint().await?; + let mut joined = left_row.clone(); + joined.extend(right_row.iter().cloned()); + if *kind == JoinKind::Cross + || matches!(predicate.evaluate(&joined)?, Value::Bool(true)) + { + matched = true; + right_matched[i] = true; + match kind { + JoinKind::Semi => { + workspace.grow(row_bytes(left_row))?; + result.push(left_row.clone()); + break; + } + JoinKind::Anti => break, + _ => { + workspace.grow(row_bytes(&joined))?; + result.push(joined); + } + } + } + } + if !matched { + match kind { + JoinKind::Left | JoinKind::Full => { + let mut joined = left_row.clone(); + joined.resize( + joined.len() + operator.inputs[1].fields.len(), + Value::Null, + ); + workspace.grow(row_bytes(&joined))?; + result.push(joined); + } + JoinKind::Anti => { + workspace.grow(row_bytes(left_row))?; + result.push(left_row.clone()); + } + _ => {} + } + } + } + if matches!(kind, JoinKind::Right | JoinKind::Full) { + for (matched, row) in right_matched.into_iter().zip(right) { + work.checkpoint().await?; + if !matched { + let mut joined = vec![Value::Null; operator.inputs[0].fields.len()]; + joined.extend(row); + workspace.grow(row_bytes(&joined))?; + result.push(joined); + } + } + } + Batch::try_new(output, result) + }) + .boxed_local()); + } + if let Kind::SemiJoin { keys } = &operator.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 mut members = std::collections::BTreeSet::new(); + let mut workspace = Workspace::new(&context)?; + let mut work = Cooperative::new(&context); + for row in &right { + work.checkpoint().await?; + if right_cols.iter().all(|&i| matchable_key(&row[i])) { + let key = group_key(row, &right_cols)?; + if !members.contains(&key) { + workspace.grow(key_bytes(&key))?; + members.insert(key); + } + } + } + let mut rows = Vec::new(); + for row in left { + work.checkpoint().await?; + if left_cols.iter().all(|&i| matchable_key(&row[i])) + && members.contains(&group_key(&row, &left_cols)?) + { + workspace.grow(std::mem::size_of::>())?; + rows.push(row); + } + } + Batch::try_new(output, rows) + }) + .boxed_local()); + } + unreachable!() +} + +// Group keys canonicalize NaNs, but equality joins must not match them. +fn matchable_key(value: &Value) -> bool { + match value { + Value::Null => false, + Value::Float64(v) => !v.is_nan(), + _ => true, + } +} diff --git a/crates/asap-physical-operators/src/operators/limit.rs b/crates/asap-physical-operators/src/operators/limit.rs new file mode 100644 index 00000000..5f5e601f --- /dev/null +++ b/crates/asap-physical-operators/src/operators/limit.rs @@ -0,0 +1,69 @@ +use super::*; +impl Operator { + 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(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + 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()) + } + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/operators/mod.rs b/crates/asap-physical-operators/src/operators/mod.rs new file mode 100644 index 00000000..bef6b913 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/mod.rs @@ -0,0 +1,310 @@ +//! Native physical operators. Each module owns its constructors and execution. +use crate::plan::{Boundedness, Emission, PhysicalOperator, PlanProperties}; +use crate::{ + runtime::{Cooperative, Input, OutputStream, Reservation, RunContext}, + values::{field, group_key, plain, Batch, Schema, Value}, + Error, +}; +use futures::StreamExt; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema, SummaryUpdate}, + pre_asap::{ColumnRef, DataType}, +}; +use std::{collections::BTreeMap, sync::Arc}; +pub(crate) mod common; +use crate::expressions::ordered; +pub use crate::expressions::Expression; +use common::*; +mod aggregate; +mod current_series; +mod filter; +mod joins; +mod limit; +mod panes; +mod persisted; +mod projection; +mod sort; +mod source; +mod summary; +pub use aggregate::Reduction; +pub use sort::SortKey; +#[derive(Clone, serde::Serialize, serde::Deserialize)] +enum Kind { + #[serde(skip)] + Source(Vec), + PaneInput { + coordinate: usize, + layout: planner_types::post_asap::PaneLayout, + offset_ms: Option, + }, + ScopeTimestamp { + columns: Vec>, + }, + CurrentSeries { + identity: usize, + coordinate: usize, + value: usize, + lookback_ms: i64, + }, + Union, + VectorToScalar { + column: usize, + }, + Project(Vec), + Filter(Expression), + Limit { + n: u64, + offset: u64, + groups: Vec, + }, + Sort { + keys: Vec, + groups: Vec, + }, + Window { + intent: Box>, + coordinate: usize, + value: usize, + groups: Vec, + window: Option<(i64, i64)>, + }, + Aggregate { + groups: Vec, + measures: Vec, + }, + SemiJoin { + keys: Vec<(usize, usize)>, + }, + Join { + kind: planner_types::pre_asap::JoinKind, + predicate: Box, + }, + SummaryBuild { + family: SummaryFamilyType, + value: usize, + time: Option, + groups: Vec, + }, + KeyedSummaryBuild { + family: SummaryFamilyType, + value: usize, + items: Vec, + groups: Vec, + }, + KeyedReadout { + state: usize, + k: usize, + }, + 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, serde::Serialize, serde::Deserialize)] +#[serde(try_from = "persisted::StoredOperator")] +pub struct Operator { + kind: Kind, + inputs: Vec, + output: Schema, +} +impl Operator { + pub(crate) fn row_preserving_input(&self) -> Option { + match self.kind { + Kind::Filter(_) | Kind::Sort { .. } | Kind::Limit { .. } | Kind::SemiJoin { .. } => { + Some(0) + } + _ => None, + } + } + + pub(crate) fn is_counter_readout(&self) -> bool { + matches!( + self.kind, + Kind::Readout { + statistic: crate::Statistic::Rate | crate::Statistic::Increase, + .. + } + ) + } + pub(crate) fn with_counter_lookback(mut self, lookback: i64) -> Result { + if lookback <= 0 { + return Err(invalid("counter lookback must be positive")); + } + if let Kind::Readout { parameters, .. } = &mut self.kind { + parameters.insert("logical_lookback_ms".into(), lookback.to_string()); + } + Ok(self) + } + pub(super) fn readout_parameters( + &self, + context: &RunContext, + ) -> Result, Error> { + let Kind::Readout { parameters, .. } = &self.kind else { + return Ok(Default::default()); + }; + let mut parameters = parameters.clone(); + if let Some(lookback) = parameters.remove("logical_lookback_ms") { + let lookback: i64 = lookback + .parse() + .map_err(|_| invalid("invalid counter lookback"))?; + let end = match context.scope { + crate::runtime::Scope::Query { + evaluation_time_ms, .. + } => evaluation_time_ms, + crate::runtime::Scope::Ingestion { window_end_ms, .. } => window_end_ms, + }; + let start = end + .checked_sub(lookback) + .ok_or_else(|| invalid("counter window overflows Int64"))?; + if let crate::runtime::Scope::Ingestion { + window_start_ms, .. + } = context.scope + { + if window_start_ms != start { + return Err(invalid( + "maintenance window differs from logical counter window", + )); + } + } + parameters.insert("range_start_ms".into(), start.to_string()); + parameters.insert("range_end_ms".into(), end.to_string()); + } + Ok(parameters) + } + + 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() + } +} +impl PhysicalOperator for Operator { + fn requires_bounded_input(&self) -> bool { + matches!( + self.kind, + Kind::Sort { .. } + | Kind::CurrentSeries { .. } + | Kind::Aggregate { .. } + | Kind::Window { .. } + | Kind::Join { .. } + | Kind::SemiJoin { .. } + | Kind::SummaryBuild { .. } + | Kind::KeyedSummaryBuild { .. } + | Kind::SummaryMerge { .. } + | Kind::VectorToScalar { .. } + ) + } + fn properties(&self, inputs: &[PlanProperties]) -> PlanProperties { + let boundedness = match &self.kind { + Kind::Source(_) => Boundedness::Bounded, + Kind::Limit { groups, .. } if groups.is_empty() => Boundedness::Bounded, + _ => Boundedness::from_inputs(inputs), + }; + PlanProperties { + boundedness, + emission: if matches!( + self.kind, + Kind::PaneInput { .. } | Kind::ScopeTimestamp { .. } + ) { + inputs + .first() + .map_or(Emission::Unknown, |input| input.emission) + } else if self.requires_bounded_input() { + Emission::AfterInput + } else { + Emission::Incremental + }, + } + } + + fn name(&self) -> &str { + match self.kind { + Kind::Source(_) => "Source", + Kind::PaneInput { .. } => "PaneInput", + Kind::ScopeTimestamp { .. } => "ScopeTimestamp", + Kind::Union => "Union", + Kind::CurrentSeries { .. } => "CurrentSeries", + Kind::VectorToScalar { .. } => "VectorToScalar", + Kind::Project(_) => "Project", + Kind::Filter(_) => "Filter", + Kind::Limit { .. } => "Limit", + Kind::Sort { .. } => "Sort", + Kind::Aggregate { .. } => "Aggregate", + Kind::Window { .. } => "WindowAggregate", + Kind::SemiJoin { .. } => "SemiJoin", + Kind::Join { .. } => "RelationalJoin", + Kind::SummaryBuild { .. } | Kind::KeyedSummaryBuild { .. } => "SummaryAgg", + Kind::KeyedReadout { .. } => "SummaryEstimate", + Kind::SummaryMerge { .. } => "SummaryMerge", + Kind::Readout { .. } => "SummaryReadout", + } + } + fn validate_context(&self, context: &RunContext) -> Result<(), Error> { + panes::validate_context(self, context)?; + current_series::validate_context(self, context)?; + self.readout_parameters(context).map(|_| ()) + } + 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, + inputs: Vec>, + context: RunContext, + ) -> Result, Error> { + match self.kind { + Kind::Source(_) | Kind::Union | Kind::VectorToScalar { .. } => { + source::execute(self, inputs, context) + } + Kind::Project(_) => projection::execute(self, inputs, context), + Kind::CurrentSeries { .. } => current_series::execute(self, inputs, context), + Kind::PaneInput { .. } | Kind::ScopeTimestamp { .. } => { + panes::execute(self, inputs, context) + } + Kind::Filter(_) => filter::execute(self, inputs, context), + Kind::Limit { .. } => limit::execute(self, inputs, context), + Kind::Sort { .. } => sort::execute(self, inputs, context), + Kind::Window { .. } | Kind::Aggregate { .. } => { + aggregate::execute(self, inputs, context) + } + Kind::Join { .. } | Kind::SemiJoin { .. } => joins::execute(self, inputs, context), + Kind::SummaryMerge { .. } => summary::execute_merge(self, inputs, context), + Kind::SummaryBuild { .. } + | Kind::Readout { .. } + | Kind::KeyedSummaryBuild { .. } + | Kind::KeyedReadout { .. } => summary::execute(self, inputs, context), + } + } +} diff --git a/crates/asap-physical-operators/src/operators/panes.rs b/crates/asap-physical-operators/src/operators/panes.rs new file mode 100644 index 00000000..4748db09 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/panes.rs @@ -0,0 +1,228 @@ +//! Run-scoped pane population checks and timestamp restoration after reduction. +use super::*; +use crate::runtime::Scope; +use planner_types::post_asap::{validate_pane_coverage, PaneLayout, WindowEdgeCoverage}; + +impl Operator { + pub(crate) fn pane_input( + input: Schema, + coordinate: usize, + layout: PaneLayout, + offset_ms: Option, + ) -> Result { + if plain(&input, coordinate)? != (&DataType::Timestamp, false) { + return Err(invalid("pane input requires a non-null timestamp")); + } + if layout.pane_width_ms > i64::MAX as u64 { + return Err(invalid("pane width exceeds timestamp range")); + } + validate_pane_coverage( + &layout, + layout.pane_origin_ms, + &WindowEdgeCoverage::PaneAligned, + ) + .map_err(|error| Error::Invalid(format!("invalid pane layout: {error:?}")))?; + if offset_ms.is_some_and(|offset| offset < 0) { + return Err(invalid("negative pane offset")); + } + Ok(Self { + kind: Kind::PaneInput { + coordinate, + layout, + offset_ms, + }, + inputs: vec![input.clone()], + output: input, + }) + } + + pub(crate) fn scope_timestamp(input: Schema, output: Schema) -> Result { + crate::values::validate_schema(&output)?; + let coordinate = output + .time_index + .ok_or_else(|| invalid("temporal output requires a time index"))?; + if plain(&output, coordinate)? != (&DataType::Timestamp, false) { + return Err(invalid("temporal output requires a non-null timestamp")); + } + let mut columns = Vec::new(); + let mut used = std::collections::BTreeSet::new(); + for (index, field) in output.fields.iter().enumerate() { + if index == coordinate { + columns.push(None); + continue; + } + let matches: Vec<_> = input + .fields + .iter() + .enumerate() + .filter(|(_, candidate)| { + candidate.dtype == field.dtype + && candidate.nullable == field.nullable + && (candidate.name == field.name + || !matches!(field.dtype, SummaryFamilyType::Plain(_))) + }) + .map(|(index, _)| index) + .collect(); + let [column] = matches.as_slice() else { + return Err(invalid("temporal output column missing or ambiguous")); + }; + if !used.insert(*column) { + return Err(invalid("temporal output repeats an input column")); + } + columns.push(Some(*column)); + } + if used.len() != input.fields.len() { + return Err(invalid("temporal output drops an input column")); + } + Ok(Self { + kind: Kind::ScopeTimestamp { columns }, + inputs: vec![input], + output, + }) + } +} + +pub(super) fn validate_context(operator: &Operator, context: &RunContext) -> Result<(), Error> { + let Kind::PaneInput { + layout, offset_ms, .. + } = &operator.kind + else { + return Ok(()); + }; + let end = match (&context.scope, offset_ms) { + ( + Scope::Ingestion { + window_start_ms, + window_end_ms, + .. + }, + None, + ) => { + if window_end_ms.checked_sub(*window_start_ms) != Some(layout.pane_width_ms as i64) { + return Err(invalid("maintenance run must cover exactly one pane")); + } + *window_end_ms + } + ( + Scope::Query { + evaluation_time_ms, .. + }, + Some(offset), + ) => { + let end = evaluation_time_ms + .checked_sub(*offset) + .ok_or_else(|| invalid("query pane timestamp overflows"))?; + end.checked_sub(layout.pane_width_ms as i64) + .ok_or_else(|| invalid("query pane start overflows"))?; + end + } + _ => return Err(invalid("pane operator received the wrong execution scope")), + }; + validate_pane_coverage(layout, Some(end), &WindowEdgeCoverage::PaneAligned).map_err(|error| { + Error::Invalid(format!( + "query requires aligned panes or boundary residuals: {error:?}" + )) + }) +} + +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + validate_context(operator, &context)?; + let input = inputs.pop().ok_or_else(|| invalid("pane input missing"))?; + let output = operator.output.clone(); + let mut seen = std::collections::BTreeSet::new(); + let mut memory = context.reserve(0)?; + let mut key_bytes = 0; + Ok(input + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let batch = batch?; + match &operator.kind { + Kind::PaneInput { + coordinate, + offset_ms, + .. + } => { + let groups: Vec<_> = output + .fields + .iter() + .enumerate() + .filter(|(index, field)| { + *index != *coordinate + && matches!(field.dtype, SummaryFamilyType::Plain(_)) + }) + .map(|(index, _)| index) + .collect(); + for row in batch.rows() { + let Value::Timestamp(timestamp) = row[*coordinate] else { + return Err(invalid("pane timestamp type mismatch")); + }; + match (&context.scope, offset_ms) { + ( + Scope::Ingestion { + window_start_ms, + window_end_ms, + .. + }, + None, + ) if timestamp > *window_start_ms && timestamp <= *window_end_ms => {} + ( + Scope::Query { + evaluation_time_ms, .. + }, + Some(offset), + ) if timestamp + == evaluation_time_ms + .checked_sub(*offset) + .ok_or_else(|| invalid("pane timestamp overflows"))? => + { + let key = group_key(row, &groups)?; + if seen.contains(&key) { + return Err(invalid("duplicate entity state within a pane")); + } + key_bytes += key.iter().map(Vec::len).sum::() + + key.len() * std::mem::size_of::>() + + 64; + memory.resize(key_bytes)?; + seen.insert(key); + } + _ => { + return Err(invalid("input population differs from required pane")) + } + } + } + Ok(batch.value().clone()) + } + Kind::ScopeTimestamp { columns } => { + let timestamp = match context.scope { + Scope::Ingestion { window_end_ms, .. } => window_end_ms, + Scope::Query { + evaluation_time_ms, .. + } => evaluation_time_ms, + }; + let rows = batch + .rows() + .iter() + .map(|row| { + columns + .iter() + .map(|column| { + column.map_or(Value::Timestamp(timestamp), |column| { + row[column].clone() + }) + }) + .collect() + }) + .collect(); + Batch::try_new(output.clone(), rows) + } + _ => unreachable!(), + } + }) + .boxed_local()) +} diff --git a/crates/asap-physical-operators/src/operators/persisted.rs b/crates/asap-physical-operators/src/operators/persisted.rs new file mode 100644 index 00000000..85add57a --- /dev/null +++ b/crates/asap-physical-operators/src/operators/persisted.rs @@ -0,0 +1,124 @@ +//! Recovery validates operator contracts, without selecting or lowering a plan. +use super::*; + +#[derive(serde::Deserialize)] +#[serde(deny_unknown_fields)] +pub(super) struct StoredOperator { + kind: Kind, + inputs: Vec, + output: Schema, +} +impl TryFrom for Operator { + type Error = Error; + fn try_from(stored: StoredOperator) -> Result { + let expected_kind = + serde_json::to_value(&stored.kind).map_err(|error| invalid(&error.to_string()))?; + let StoredOperator { + kind, + inputs, + output, + } = stored; + for schema in inputs.iter().chain(std::iter::once(&output)) { + crate::values::validate_schema(schema)?; + } + let input = |index| { + inputs + .get(index) + .cloned() + .ok_or_else(|| invalid("missing persisted input")) + }; + let op = match kind { + Kind::Source(_) => return Err(invalid("physical plans cannot persist live sources")), + Kind::PaneInput { + coordinate, + layout, + offset_ms, + } => Operator::pane_input(input(0)?, coordinate, layout, offset_ms)?, + Kind::ScopeTimestamp { .. } => Operator::scope_timestamp(input(0)?, output.clone())?, + Kind::Union => Operator::union(input(0)?, inputs.len())?, + Kind::CurrentSeries { + identity, + coordinate, + value, + lookback_ms, + } => Operator::current_series(input(0)?, identity, coordinate, value, lookback_ms)?, + Kind::VectorToScalar { column } => Operator::vector_to_scalar(input(0)?, column)?, + Kind::Project(expressions) => { + if expressions.len() != output.fields.len() { + return Err(invalid("persisted projection width mismatch")); + } + Operator::project( + input(0)?, + output + .fields + .iter() + .zip(expressions) + .map(|(f, e)| (f.name.clone(), e)) + .collect(), + )? + } + Kind::Filter(expression) => Operator::filter(input(0)?, expression)?, + Kind::Limit { n, offset, groups } => Operator::limit(input(0)?, n, offset, groups)?, + Kind::Sort { keys, groups } => Operator::sort(input(0)?, keys, groups)?, + Kind::Window { + intent, + coordinate, + value, + groups, + window, + } => Operator::window(input(0)?, *intent, coordinate, value, groups, window)?, + Kind::Aggregate { groups, measures } => { + if groups.len() + measures.len() != output.fields.len() { + return Err(invalid("persisted aggregate width mismatch")); + } + let names = output.fields[groups.len()..].iter().map(|f| f.name.clone()); + Operator::aggregate(input(0)?, groups, names.zip(measures).collect())? + } + Kind::SemiJoin { keys } => Operator::semi_join(input(0)?, input(1)?, keys)?, + Kind::Join { kind, predicate } => Operator::relational_join( + input(0)?, + input(1)?, + kind, + &planner_types::pre_asap::Predicate(std::rc::Rc::new( + predicate.expression().clone(), + )), + output.clone(), + )?, + Kind::SummaryBuild { + family, + value, + time, + groups, + } => Operator::summary_build(input(0)?, family, value, time, groups)?, + Kind::KeyedSummaryBuild { + family, + value, + items, + groups, + } => Operator::keyed_summary_build(input(0)?, family, value, items, groups)?, + Kind::KeyedReadout { state, k } => { + Operator::keyed_readout(input(0)?, state, k, output.clone())? + } + Kind::SummaryMerge { state, groups } => { + Operator::summary_merge(input(0)?, state, groups)? + } + Kind::Readout { + state, + statistic, + parameters, + } => Operator::readout(input(0)?, state, statistic, parameters)?, + } + .with_output_schema(output)?; + if serde_json::to_value(&op.kind).map_err(|error| invalid(&error.to_string()))? + != expected_kind + { + return Err(invalid( + "persisted operator contains inconsistent compiled fields", + )); + } + if op.inputs != inputs { + return Err(invalid("persisted operator input contracts differ")); + } + Ok(op) + } +} diff --git a/crates/asap-physical-operators/src/operators/projection.rs b/crates/asap-physical-operators/src/operators/projection.rs new file mode 100644 index 00000000..9337f7ab --- /dev/null +++ b/crates/asap-physical-operators/src/operators/projection.rs @@ -0,0 +1,47 @@ +use super::*; +impl Operator { + 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(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + Kind::Project(expressions) => Ok(input + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + 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()), + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/operators/sort.rs b/crates/asap-physical-operators/src/operators/sort.rs new file mode 100644 index 00000000..71f81de0 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/sort.rs @@ -0,0 +1,172 @@ +use super::*; +impl Operator { + 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, + }) + } +} +#[derive(serde::Serialize, serde::Deserialize, Clone, Debug)] +pub struct SortKey { + pub column: usize, + pub descending: bool, + pub nulls_first: bool, +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let result = match &operator.kind { + Kind::Sort { keys, groups } => { + let mut grouped = BTreeMap::>, Vec>>::new(); + let mut work = Cooperative::new(&context); + let mut workspace = Workspace::new(&context)?; + for row in rows { + work.checkpoint().await?; + for key in keys { + if matches!(row[key.column], Value::Map(_)) && nested_nan(&row[key.column]) + { + return Err(invalid("NaN in collection sort key")); + } + } + let key = group_key(&row, groups)?; + workspace.grow(std::mem::size_of::>())?; + if !grouped.contains_key(&key) { + workspace.grow(key_bytes(&key))?; + } + grouped.entry(key).or_default().push(row); + } + let mut result = Vec::new(); + for rows in grouped.into_values() { + let rows = + cooperative_sort(rows, |a, b| compare_rows(a, b, keys), &context).await?; + result.extend(rows); + } + result + } + _ => unreachable!(), + }; + 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 nested_nan(value: &Value) -> bool { + match value { + Value::Float64(value) => value.is_nan(), + Value::Map(values) => values + .iter() + .any(|(key, value)| nested_nan(key) || nested_nan(value)), + Value::List(values) | Value::Struct(values) => values.iter().any(nested_nan), + _ => false, + } +} +/// Stable in-memory merge sort with bounded synchronous chunks. Scratch storage +/// is reserved before allocation; comparisons yield between merge steps. +pub(super) async fn cooperative_sort( + rows: Vec, + compare: impl Fn(&T, &T) -> std::cmp::Ordering, + context: &RunContext, +) -> Result, Error> { + use std::collections::VecDeque; + let bytes = rows + .len() + .checked_mul(std::mem::size_of::() + std::mem::size_of::>()) + .and_then(|n| n.checked_mul(3)) + .ok_or(Error::MemoryLimit)?; + let _scratch = context.reserve(bytes)?; + let mut work = Cooperative::new(context); + let mut rows = rows.into_iter(); + let mut runs = VecDeque::new(); + loop { + work.checkpoint().await?; + let mut chunk = rows.by_ref().take(256).collect::>(); + if chunk.is_empty() { + break; + } + chunk.sort_by(&compare); + runs.push_back(VecDeque::from(chunk)); + } + // Merge adjacent runs in rounds to preserve ties in original input order. + while runs.len() > 1 { + let mut next = VecDeque::new(); + while let Some(mut left) = runs.pop_front() { + let Some(mut right) = runs.pop_front() else { + next.push_back(left); + break; + }; + let mut merged = VecDeque::with_capacity(left.len() + right.len()); + while !left.is_empty() || !right.is_empty() { + work.checkpoint().await?; + let take_left = match (left.front(), right.front()) { + (Some(a), Some(b)) => !compare(a, b).is_gt(), + (Some(_), None) => true, + _ => false, + }; + merged.push_back(if take_left { + left.pop_front().unwrap() + } else { + right.pop_front().unwrap() + }); + } + next.push_back(merged); + } + runs = next; + } + Ok(runs.pop_front().unwrap_or_default().into()) +} diff --git a/crates/asap-physical-operators/src/operators/source.rs b/crates/asap-physical-operators/src/operators/source.rs new file mode 100644 index 00000000..060a5381 --- /dev/null +++ b/crates/asap-physical-operators/src/operators/source.rs @@ -0,0 +1,88 @@ +use super::*; +impl Operator { + pub fn source(output: Schema, batches: Vec) -> Result { + crate::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, + }) + } + 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]])?], + ) + } + 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 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(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + if let Kind::Source(batches) = &operator.kind { + return Ok(futures::stream::iter(batches.iter().cloned().map(Ok)).boxed_local()); + } + if matches!(operator.kind, Kind::Union) { + return Ok(futures::stream::select_all(inputs) + .map(|batch| batch.map(|batch| batch.value().clone())) + .boxed_local()); + } + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + Kind::VectorToScalar { column } => Ok(futures::stream::once(async move { + let mut input = input; + let mut work = Cooperative::new(&context); + let mut value = f64::NAN; + let mut count = 0usize; + while let Some(batch) = input.next().await { + for row in batch?.rows() { + work.checkpoint().await?; + 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()), + _ => unreachable!(), + } +} diff --git a/crates/asap-physical-operators/src/operators/summary/mod.rs b/crates/asap-physical-operators/src/operators/summary/mod.rs new file mode 100644 index 00000000..745b524f --- /dev/null +++ b/crates/asap-physical-operators/src/operators/summary/mod.rs @@ -0,0 +1,559 @@ +use super::*; +impl Operator { + pub fn keyed_summary_build( + input: Schema, + family: SummaryFamilyType, + value: usize, + items: Vec, + groups: Vec, + ) -> Result { + use crate::summary_kernels::weighted_frequency::WeightedFrequency; + crate::values::validate_family(&family)?; + let SummaryFamilyType::Sketch(kind, _) = &family else { + return Err(invalid("keyed sketch required")); + }; + WeightedFrequency::configuration(kind)?; + validate_groups(&input, &groups)?; + if items.is_empty() || plain(&input, value)? != (&DataType::Float64, false) { + return Err(invalid( + "keyed summary requires identities and non-null Float64 weights", + )); + } + for &item in &items { + if !matches!( + plain(&input, item)?.0, + DataType::Utf8 + | DataType::Timestamp + | DataType::Int64 + | DataType::Float64 + | DataType::Bool + | DataType::Null + ) { + return Err(invalid("unsupported keyed summary identity type")); + } + } + 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::KeyedSummaryBuild { + family, + value, + items, + groups, + }, + inputs: vec![input], + output: schema(fields), + }) + } + pub fn keyed_readout( + input: Schema, + state: usize, + k: usize, + output: Schema, + ) -> Result { + use crate::summary_kernels::weighted_frequency::WeightedFrequency; + crate::values::validate_family(&field(&input, state)?.dtype)?; + let SummaryFamilyType::Sketch(kind, _) = &field(&input, state)?.dtype else { + return Err(invalid("keyed readout requires summary state")); + }; + let (_, _, _, capacity) = WeightedFrequency::configuration(kind)?; + if k > capacity || output.fields.len() <= input.fields.len() { + return Err(invalid("invalid keyed readout shape or capacity")); + } + if state + 1 != input.fields.len() + || output.fields[..state] != input.fields[..state] + || output.fields.last().unwrap().dtype != SummaryFamilyType::Plain(DataType::Float64) + { + return Err(invalid( + "keyed readout must preserve partitions and return a Float64 score", + )); + } + crate::values::validate_schema(&output)?; + Ok(Self { + kind: Kind::KeyedReadout { state, k }, + inputs: vec![input], + output, + }) + } + pub fn summary_build( + input: Schema, + family: SummaryFamilyType, + value: usize, + time: Option, + groups: Vec, + ) -> Result { + crate::values::validate_family(&family)?; + validate_groups(&input, &groups)?; + if plain(&input, value)?.0 != &DataType::Float64 { + return Err(invalid("summary numeric update requires 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)?; + crate::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 { + crate::values::validate_family(&field(&input, state)?.dtype)?; + if matches!(field(&input, state)?.dtype, SummaryFamilyType::Plain(_)) { + return Err(invalid("summary state required")); + } + crate::capability::validate_native_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 + }; + // A state-only row represents the global population. Its extrema may + // be empty, just like an ordinary ungrouped MIN/MAX aggregate. + let nullable = fields.len() == 1 + && matches!( + fields[state].dtype, + SummaryFamilyType::ExactAggregate( + planner_types::post_asap::ExactKind::Min + | planner_types::post_asap::ExactKind::Max, + _ + ) + ); + fields[state] = result_field("value", result_type, nullable); + Ok(Self { + kind: Kind::Readout { + state, + statistic, + parameters, + }, + inputs: vec![input], + output: schema(fields), + }) + } +} +pub(super) fn execute<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let parameters = operator.readout_parameters(&context)?; + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + match &operator.kind { + 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::KeyedSummaryBuild { + family, + value, + items, + groups, + } => Ok(futures::stream::once(async move { + Batch::try_new( + output, + build_keyed_summary(input, family, *value, items, groups, &context).await?, + ) + }) + .boxed_local()), + Kind::KeyedReadout { state, k } => Ok(input + .map(move |batch| { + let batch = batch?; + let mut rows = Vec::new(); + for row in batch.rows() { + let Value::Summary { state: summary, .. } = &row[*state] else { + return Err(invalid("summary value required")); + }; + let summary = summary + .as_any() + .downcast_ref::() + .ok_or_else(|| invalid("weighted frequency typed state required"))?; + for items in summary.rows(*k) { + let mut values = row[..*state].to_vec(); + values.extend(items); + // The typed output schema restores epoch-millisecond + // timestamp keys from the kernel's Int64 representation. + for (value, field) in values.iter_mut().zip(&output.fields) { + if field.dtype == SummaryFamilyType::Plain(DataType::Timestamp) { + if let Value::Int64(time) = value { + *value = Value::Timestamp(*time); + } + } + } + rows.push(values); + } + } + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + Kind::Readout { + state, + statistic, + .. + } => Ok(input + .map(move |batch| { + let batch = batch?; + let mut rows = batch.rows().to_vec(); + rows.retain(|row| !matches!(&row[*state], Value::Summary { state: summary, .. } + if crate::stored_state::readout::insufficient_counter_samples(summary.as_ref(), *statistic))); + 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 if output.fields[*state].nullable + && matches!(statistic, crate::Statistic::Min | crate::Statistic::Max) + { + let stats = summary.aux_stats(); + let value = if *statistic == crate::Statistic::Min { stats.min } else { stats.max }; + value.map(Value::Float64).unwrap_or(Value::Null) + } else { + Value::Float64( + summary + .query_statistic(*statistic, &None, ¶meters) + .map_err(|e| Error::Operator(e.to_string()))?, + ) + }; + } + Batch::try_new(output.clone(), rows) + }) + .boxed_local()), + _ => unreachable!(), + } +} +pub(super) fn execute_merge<'a>( + operator: &'a Operator, + mut inputs: Vec>, + context: RunContext, +) -> Result, Error> { + let output = operator.output.clone(); + let input = inputs.pop().ok_or_else(|| invalid("input missing"))?; + Ok(futures::stream::once(async move { + let (rows, _memory) = collect_rows(input, &context).await?; + let result = match &operator.kind { + Kind::SummaryMerge { state, groups } => { + merge_summary(rows, *state, groups, &context).await? + } + _ => unreachable!(), + }; + Batch::try_new(output, result) + }) + .boxed_local()) +} + +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 work = Cooperative::new(context); + 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() { + work.checkpoint().await?; + 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"); + // SQL aggregates ignore NULL samples while retaining the group. + // A missing counter sample also contributes no observation. + let value = match row[value] { + Value::Float64(value) => value, + Value::Null => continue, + _ => 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()) +} +async fn merge_summary( + rows: Vec>, + state_column: usize, + groups: &[usize], + context: &RunContext, +) -> Result>, Error> { + type GroupState = (Vec, SummaryFamilyType, Arc); + let mut states: BTreeMap>, GroupState> = BTreeMap::new(); + let mut work = Cooperative::new(context); + let mut memory = context.reserve(0)?; + let mut retained = 0usize; + for row in rows { + work.checkpoint().await?; + 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")); + } + let old_bytes = existing.approx_memory_bytes(); + // Reserve an estimate for the replacement while both input states remain live. + memory.resize( + retained + .checked_add(old_bytes) + .and_then(|n| n.checked_add(state.approx_memory_bytes())) + .ok_or(Error::MemoryLimit)?, + )?; + *existing = Arc::from( + existing + .merge_with(state.as_ref()) + .map_err(|e| Error::Operator(e.to_string()))?, + ); + retained = retained + .checked_sub(old_bytes) + .and_then(|n| n.checked_add(existing.approx_memory_bytes())) + .ok_or(Error::MemoryLimit)?; + memory.resize(retained)?; + } else { + retained = retained + .checked_add(key_bytes(&key) + row_bytes(&row)) + .ok_or(Error::MemoryLimit)?; + memory.resize(retained)?; + 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()) +} + +async fn build_keyed_summary( + mut input: Input<'_, Batch>, + family: &SummaryFamilyType, + value: usize, + items: &[usize], + groups: &[usize], + context: &RunContext, +) -> Result>, Error> { + use crate::{summary_kernels::weighted_frequency::WeightedFrequency, AggregateCore}; + let SummaryFamilyType::Sketch(kind, _) = family else { + unreachable!() + }; + let (algorithm, width, depth, capacity) = WeightedFrequency::configuration(kind)?; + let mut work = Cooperative::new(context); + let mut states = + BTreeMap::>, (Vec, WeightedFrequency, Reservation, usize)>::new(); + while let Some(batch) = input.next().await { + let batch = batch?; + for row in batch.rows() { + work.checkpoint().await?; + let key = group_key(row, groups)?; + if !states.contains_key(&key) { + let labels = groups.iter().map(|&i| row[i].clone()).collect::>(); + let overhead = labels.iter().map(Value::bytes).sum::() + + key.iter().map(|v| v.len() + 24).sum::() + + 128; + let bytes = width + .checked_mul(depth) + .and_then(|n| n.checked_mul(8)) + .and_then(|n| n.checked_add(overhead)) + .ok_or_else(|| invalid("weighted frequency memory size overflow"))?; + let reservation = context.reserve(bytes)?; + states.insert( + key.clone(), + ( + labels, + WeightedFrequency::new(algorithm, width, depth, capacity)?, + reservation, + overhead, + ), + ); + } + let (_, summary, reservation, overhead) = states.get_mut(&key).unwrap(); + let Value::Float64(weight) = row[value] else { + return Err(invalid("weighted frequency weight type")); + }; + summary.update( + &items + .iter() + .map(|&i| match &row[i] { + Value::Timestamp(time) => Value::Int64(*time), + value => value.clone(), + }) + .collect::>(), + weight, + )?; + reservation.resize(summary.approx_memory_bytes() + *overhead)?; + } + } + Ok(states + .into_values() + .map(|(mut labels, summary, _, _)| { + labels.push(Value::Summary { + family: family.clone(), + state: Arc::new(summary), + }); + labels + }) + .collect()) +} diff --git a/crates/asap-physical-operators/src/physical_planner/candidates.rs b/crates/asap-physical-operators/src/physical_planner/candidates.rs new file mode 100644 index 00000000..148035d7 --- /dev/null +++ b/crates/asap-physical-operators/src/physical_planner/candidates.rs @@ -0,0 +1,295 @@ +//! Compile maintenance-selected frontiers without deployment-specific graph rewrites. +use super::*; + +/// One computation realization; lifecycle/window/revision requirements accompany +/// it during optimization and deployment. Stored outputs have no storage identity. +#[derive(Clone)] +pub struct PhysicalCandidate { + pub precompute: Option, + pub query: CompiledPhysicalDag, + pub materialized_outputs: BTreeMap, +} + +/// Compile an explicit materialization frontier selected by Planner maintenance +/// search. Operators upstream of that frontier run in precompute, including +/// readouts/reductions; query execution receives their typed output values. +/// Empty frontiers retain the full computation in the query DAG. +/// +/// Repeated windows must be instantiated with the same evaluation/population +/// contract used to build each output. This API never treats a result from a +/// different window or revision as interchangeable merely because types match. +pub fn compile_candidate( + dag: &ExecutableDag, + inputs: BTreeMap, + roots: &[NodeId], + frontier: &[NodeId], +) -> Result { + if frontier.is_empty() { + return Ok(PhysicalCandidate { + precompute: None, + query: compile(dag, inputs, roots)?, + materialized_outputs: BTreeMap::new(), + }); + } + let frontier_set: BTreeSet<_> = frontier.iter().copied().collect(); + if frontier_set.len() != frontier.len() || frontier.iter().any(|id| inputs.contains_key(id)) { + return Err(invalid("frontier must contain distinct computed outputs")); + } + let full = compile(dag, inputs.clone(), roots)?; + let precompute = compile(dag, inputs.clone(), frontier)?; + let mut materialized_outputs = BTreeMap::new(); + for &id in frontier { + // Also proves that the frontier is reachable from the requested roots. + full.output_contract(id)?; + let mut output = precompute.output_contract(id)?; + if output.properties.boundedness != Boundedness::Bounded { + return Err(invalid("materialized output requires bounded execution")); + } + // A stored reader may stream batches even when the producer blocked. + // Its timing is independent; the retained result still must be finite. + output.properties.emission = Emission::Unknown; + materialized_outputs.insert(id, output); + } + let mut query_inputs = inputs; + query_inputs.extend(materialized_outputs.clone()); + let query = compile(dag, query_inputs, roots)?; + let used: BTreeSet<_> = query.input_contracts().map(|(id, _)| id).collect(); + if !frontier.iter().all(|id| used.contains(id)) { + return Err(invalid( + "frontier contains an output shadowed by another boundary", + )); + } + Ok(PhysicalCandidate { + precompute: Some(precompute), + query, + materialized_outputs, + }) +} + +/// Enumerate bounded, reachable materialization frontiers above explicit inputs. +/// Each frontier is an antichain: storing an output and its ancestor together +/// would leave the ancestor unused by query execution. Lifecycle eligibility +/// and deployment feasibility are evaluated separately before cost selection. +/// Exceeding the search budget returns an error, never a partial inventory. +pub fn enumerate_frontiers( + dag: &ExecutableDag, + inputs: &BTreeMap, + roots: &[NodeId], + max_candidates: usize, +) -> Result>, Error> { + if max_candidates == 0 { + return Err(invalid( + "frontier search requires a positive candidate budget", + )); + } + let compiled = compile(dag, inputs.clone(), roots)?; + let mut ancestors = BTreeMap::>::new(); + let mut eligible = Vec::new(); + for node in &dag.nodes { + let id = u64::from(node.id.0); + if inputs.contains_key(&id) { + continue; + } + let Ok(contract) = compiled.output_contract(id) else { + continue; + }; + if contract.properties.boundedness != Boundedness::Bounded { + continue; + } + let mut seen = BTreeSet::new(); + let mut pending = vec![id]; + while let Some(current) = pending.pop() { + if !seen.insert(current) || inputs.contains_key(¤t) { + continue; + } + pending.extend( + dag.edges + .iter() + .filter(|edge| u64::from(edge.consumer.0) == current) + .map(|edge| u64::from(edge.producer.0)), + ); + } + ancestors.insert(id, seen); + eligible.push(id); + } + eligible.sort_unstable(); + let mut frontiers = vec![vec![]]; + for id in eligible { + let additions = frontiers + .iter() + .filter(|frontier| { + frontier.iter().all(|previous| { + !ancestors[&id].contains(previous) && !ancestors[previous].contains(&id) + }) + }) + .map(|frontier| { + let mut next = frontier.clone(); + next.push(id); + next + }) + .collect::>(); + if additions.len() > max_candidates.saturating_sub(frontiers.len()) { + return Err(invalid( + "materialization frontier search exceeds candidate budget", + )); + } + frontiers.extend(additions); + } + Ok(frontiers) +} + +/// Lower every maintenance candidate before feasibility/cost evaluation. Keep +/// individual failures visible; do not substitute another computation on error. +pub fn compile_candidates( + dag: &ExecutableDag, + inputs: BTreeMap, + roots: &[NodeId], + frontiers: &[Vec], +) -> Vec> { + frontiers + .iter() + .map(|frontier| compile_candidate(dag, inputs.clone(), roots, frontier)) + .collect() +} + +/// Complete workload cost supplied by scoped optimizer/deployment evidence. +/// The evaluator includes build/update work, retained state, shared producers +/// and recurrent reads over the same horizon; these are not per-query timings. +#[derive(Clone, Debug)] +pub struct CandidateCost { + pub workload_scope: String, + pub horizon_seconds: f64, + pub total_cost: f64, +} + +pub struct CandidateSelection { + pub candidate: T, + pub candidate_index: usize, + pub cost: CandidateCost, +} + +/// Select only compiled and deployment-feasible physical candidates. `None` +/// rejects an unbindable candidate before pricing. Comparable scoped costs are +/// required; deployment never rewrites the selected frontier after this step. +/// The payload is generic so deployments can retain binding/diagnostic metadata +/// alongside each compiled computation without duplicating winner selection. +pub fn select_candidate( + candidates: Vec>, + mut evaluate: impl FnMut(&T) -> Result, Error>, +) -> Result, Error> { + let mut scope: Option<(String, f64)> = None; + let mut selected: Option> = None; + for (candidate_index, candidate) in candidates.into_iter().enumerate() { + let Ok(candidate) = candidate else { continue }; + let Some(cost) = evaluate(&candidate)? else { + continue; + }; + if cost.workload_scope.is_empty() + || !cost.horizon_seconds.is_finite() + || cost.horizon_seconds <= 0. + || !cost.total_cost.is_finite() + || cost.total_cost < 0. + { + return Err(invalid( + "candidate cost lacks a valid workload scope/horizon", + )); + } + let current_scope = (cost.workload_scope.clone(), cost.horizon_seconds); + if scope.as_ref().is_some_and(|scope| scope != ¤t_scope) { + return Err(invalid( + "candidate costs describe different workloads or horizons", + )); + } + scope = Some(current_scope); + if selected + .as_ref() + .is_none_or(|selected| cost.total_cost < selected.cost.total_cost) + { + selected = Some(CandidateSelection { + candidate, + candidate_index, + cost, + }); + } + } + selected.ok_or_else(|| invalid("no feasible priced physical candidate")) +} + +#[derive(serde::Serialize, serde::Deserialize)] +#[serde(deny_unknown_fields)] +struct StoredCandidate { + version: u32, + precompute: Option, + query: serde_json::Value, + materialized_outputs: BTreeMap, +} + +impl PhysicalCandidate { + /// Validate the physical handoff, including the producer/reader boundary. + pub fn validate(&self) -> Result<(), Error> { + self.query.validate()?; + let Some(precompute) = &self.precompute else { + return if self.materialized_outputs.is_empty() { + Ok(()) + } else { + Err(invalid("materialized outputs have no producer DAG")) + }; + }; + precompute.validate()?; + let outputs: BTreeSet<_> = self.materialized_outputs.keys().copied().collect(); + if outputs.is_empty() || outputs != precompute.roots().iter().copied().collect() { + return Err(invalid("physical frontier differs from precompute outputs")); + } + let readers: BTreeMap<_, _> = self.query.input_contracts().collect(); + for (&id, contract) in &self.materialized_outputs { + let produced = precompute.output_contract(id)?; + // Direct frontiers retain their node IDs. Temporal candidates can + // read several window instances through distinct input slots; + // their deployment bindings must validate those slots separately. + let reader = readers.get(&id); + if contract.schema != produced.schema + || reader.is_some_and(|reader| contract.schema != reader.schema) + || produced.properties.boundedness != Boundedness::Bounded + || contract.properties.boundedness != Boundedness::Bounded + || reader + .is_some_and(|reader| reader.properties.boundedness != Boundedness::Bounded) + { + return Err(invalid("physical frontier schema or boundedness mismatch")); + } + } + Ok(()) + } + pub fn encode(&self) -> Result, Error> { + self.validate()?; + let graph = |dag: &CompiledPhysicalDag| -> Result { + serde_json::from_slice(&dag.encode()?).map_err(|error| invalid(error.to_string())) + }; + serde_json::to_vec(&StoredCandidate { + version: 1, + precompute: self.precompute.as_ref().map(graph).transpose()?, + query: graph(&self.query)?, + materialized_outputs: self.materialized_outputs.clone(), + }) + .map_err(|error| invalid(error.to_string())) + } + /// Recover the selected physical candidate; no logical IR is accepted here. + pub fn decode(bytes: &[u8]) -> Result { + let stored: StoredCandidate = + serde_json::from_slice(bytes).map_err(|error| invalid(error.to_string()))?; + if stored.version != 1 { + return Err(invalid("unsupported physical candidate format")); + } + let graph = |value| -> Result { + CompiledPhysicalDag::decode( + &serde_json::to_vec(&value).map_err(|error| invalid(error.to_string()))?, + ) + }; + let candidate = Self { + precompute: stored.precompute.map(graph).transpose()?, + query: graph(stored.query)?, + materialized_outputs: stored.materialized_outputs, + }; + candidate.validate()?; + Ok(candidate) + } +} diff --git a/crates/asap-physical-operators/src/physical_planner/compiled.rs b/crates/asap-physical-operators/src/physical_planner/compiled.rs new file mode 100644 index 00000000..da97def5 --- /dev/null +++ b/crates/asap-physical-operators/src/physical_planner/compiled.rs @@ -0,0 +1,246 @@ +//! Reader-independent physical computation and checked deployment instantiation. +use super::*; + +/// A typed execution boundary, without storage identity or a live reader. +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +pub struct InputContract { + pub schema: Schema, + pub properties: PlanProperties, +} +impl InputContract { + pub fn bounded(schema: Schema) -> Self { + Self { + schema, + properties: PlanProperties { + boundedness: Boundedness::Bounded, + emission: Emission::Unknown, + }, + } + } + pub fn from_source(source: &dyn PhysicalOperator) -> Self { + Self { + schema: source.output_schema(), + properties: source.properties(&[]), + } + } +} +#[derive(Clone, serde::Serialize, serde::Deserialize)] +enum Node { + Input(InputContract), + Operator { + inputs: Vec, + operator: Operator, + }, +} + +/// Selected native operators and input slots. Rebinding never repeats lowering. +#[derive(Clone)] +pub struct CompiledPhysicalDag { + nodes: BTreeMap, + roots: Vec, +} +#[derive(serde::Serialize, serde::Deserialize)] +#[serde(deny_unknown_fields)] +struct StoredDag { + version: u32, + nodes: BTreeMap, + roots: Vec, +} + +impl CompiledPhysicalDag { + /// Persist selected physical operators and input slots, never live readers + /// or mutable summary state. Recovery does not run logical plan lowering. + pub fn encode(&self) -> Result, Error> { + self.validate()?; + let bytes = serde_json::to_vec(&StoredDag { + version: 1, + nodes: self.nodes.clone(), + roots: self.roots.clone(), + }) + .map_err(|error| invalid(error.to_string()))?; + // JSON cannot preserve non-finite literal values. Fail at publication, + // rather than persisting a document that cannot be recovered. + Self::decode(&bytes)?; + Ok(bytes) + } + pub fn decode(bytes: &[u8]) -> Result { + let stored: StoredDag = + serde_json::from_slice(bytes).map_err(|error| invalid(error.to_string()))?; + if stored.version != 1 { + return Err(invalid("unsupported physical plan format")); + } + let result = Self { + nodes: stored.nodes, + roots: stored.roots, + }; + result.validate()?; + Ok(result) + } + + /// Assemble already-lowered operators and typed external inputs. This is + /// useful for engines that compose multiple compiled computation fragments. + pub fn from_operators( + inputs: BTreeMap, + operators: BTreeMap, Operator)>, + roots: Vec, + ) -> Result { + let mut result = Self::new(roots); + for (id, contract) in inputs { + result.add_input(id, contract)?; + } + for (id, (inputs, operator)) in operators { + result.add(id, inputs, operator)?; + } + result.validate()?; + Ok(result) + } + pub(super) fn new(roots: Vec) -> Self { + Self { + nodes: BTreeMap::new(), + roots, + } + } + pub(super) fn add_input(&mut self, id: NodeId, contract: InputContract) -> Result<(), Error> { + self.insert(id, Node::Input(contract)) + } + pub(super) fn add( + &mut self, + id: NodeId, + inputs: Vec, + operator: Operator, + ) -> Result<(), Error> { + self.insert(id, Node::Operator { inputs, operator }) + } + fn insert(&mut self, id: NodeId, node: Node) -> Result<(), Error> { + if self.nodes.insert(id, node).is_some() { + return Err(invalid(format!("duplicate physical node {id}"))); + } + Ok(()) + } + /// Identify the external input whose rows survive unchanged at this output. + /// Protocol adapters can retain labels that are outside a closed physical schema. + pub fn row_source(&self, id: NodeId) -> Option { + match self.nodes.get(&id)? { + Node::Input(_) => Some(id), + Node::Operator { inputs, operator } => { + let index = operator.row_preserving_input()?; + self.row_source(*inputs.get(index)?) + } + } + } + + /// Selected operator name, for plan inspection without decoding its wire format. + pub fn operator_name(&self, id: NodeId) -> Option<&str> { + match self.nodes.get(&id)? { + Node::Input(_) => Some("Input"), + Node::Operator { operator, .. } => Some(operator.name()), + } + } + + pub fn roots(&self) -> &[NodeId] { + &self.roots + } + pub fn input_contracts(&self) -> impl Iterator { + self.nodes.iter().filter_map(|(&id, node)| match node { + Node::Input(contract) => Some((id, contract)), + Node::Operator { .. } => None, + }) + } + /// Derive a reachable output contract without opening deployment readers. + pub fn output_contract(&self, id: NodeId) -> Result { + let sources = self + .input_contracts() + .map(|(id, contract)| (id, Box::new(contract.clone()) as Source<'_>)) + .collect(); + let graph = self.instantiate(sources)?; + let properties = graph.properties(&self.roots)?; + let properties = *properties + .get(&id) + .ok_or_else(|| invalid("output is not reachable"))?; + let schema = match self + .nodes + .get(&id) + .ok_or_else(|| invalid("missing output"))? + { + Node::Input(contract) => contract.schema.clone(), + Node::Operator { operator, .. } => operator.output_schema(), + }; + Ok(InputContract { schema, properties }) + } + /// Validate using contract-only sources. No deployment reader is available. + pub fn validate(&self) -> Result<(), Error> { + let sources = self + .input_contracts() + .map(|(id, c)| (id, Box::new(c.clone()) as Source<'_>)) + .collect(); + self.instantiate(sources).map(|_| ()) + } + /// Resolve exactly the declared inputs and validate before any source starts. + pub fn instantiate<'a>( + &self, + mut sources: BTreeMap>, + ) -> Result, Error> { + let mut graph = PhysicalDag::default(); + for (&id, node) in &self.nodes { + match node { + Node::Input(contract) => { + let source = sources + .remove(&id) + .ok_or_else(|| invalid(format!("missing physical input {id}")))?; + let actual = source.properties(&[]); + if !source.input_schemas().is_empty() + || source.output_schema() != contract.schema + || (contract.properties.boundedness != Boundedness::Unknown + && actual.boundedness != contract.properties.boundedness) + || (contract.properties.emission != Emission::Unknown + && actual.emission != contract.properties.emission) + { + return Err(invalid(format!( + "physical input {id} violates its compiled contract" + ))); + } + graph.add_boxed( + id, + vec![], + Box::new(CheckedSource { + source, + output: contract.schema.clone(), + }), + )?; + } + Node::Operator { inputs, operator } => { + graph.add(id, inputs.clone(), operator.clone())?; + } + } + } + if !sources.is_empty() { + return Err(invalid("unexpected physical input binding")); + } + graph.validate(&self.roots)?; + Ok(graph) + } +} +impl PhysicalOperator for InputContract { + fn name(&self) -> &str { + "UnresolvedInput" + } + fn input_schemas(&self) -> Vec { + vec![] + } + fn output_schema(&self) -> Schema { + self.schema.clone() + } + fn properties(&self, _: &[PlanProperties]) -> PlanProperties { + self.properties + } + fn output_bytes(&self, batch: &Batch) -> usize { + batch.bytes() + } + fn start<'a>( + &'a self, + _: Vec>, + _: crate::runtime::RunContext, + ) -> Result, Error> { + Err(invalid("physical input must be resolved before execution")) + } +} diff --git a/crates/asap-physical-operators/src/physical_planner/mod.rs b/crates/asap-physical-operators/src/physical_planner/mod.rs new file mode 100644 index 00000000..7064ce22 --- /dev/null +++ b/crates/asap-physical-operators/src/physical_planner/mod.rs @@ -0,0 +1,836 @@ +//! Compile logical computation to native operators with typed external inputs. +//! Compilation needs no readers; deployment resolves inputs after selection. +use crate::{ + operators::{Expression, Operator, Reduction, SortKey}, + plan::{Boundedness, Emission, NodeId, PhysicalDag, PhysicalOperator, PlanProperties}, + values::{Batch, Schema}, + Error, +}; +use planner_types::{ + post_asap::{ + ExactOperation, ExecutableDag, ExecutableDagNode, ExecutableOperatorPayload as Payload, + SketchQuery, SummaryFamilyType, SummaryInputExpr, ValueOperation, + }, + pre_asap::{ + AggIntent, ColumnRef, CompareOpKind, GroupKeys, QueryExpr, Reduction as PlannerReduction, + }, +}; +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 mod promql_rows; + +mod candidates; +pub use candidates::{ + compile_candidate, compile_candidates, enumerate_frontiers, select_candidate, CandidateCost, + CandidateSelection, PhysicalCandidate, +}; + +mod temporal_panes; +pub use temporal_panes::{ + compile_temporal_pane_candidate, TemporalEntityIdentity, TemporalPaneCandidate, + TemporalPaneMaintenance, +}; + +mod compiled; +pub use compiled::{CompiledPhysicalDag, InputContract}; + +/// Compile computation without opening or retaining deployment readers. +/// Input contracts identify explicit boundaries selected by maintenance planning. +pub fn compile( + dag: &ExecutableDag, + inputs: BTreeMap, + roots: &[NodeId], +) -> Result { + compile_internal(dag, inputs, roots) +} + +/// Convenience for callers that already resolved inputs. Lowering still uses +/// only their contracts, and instantiation checks those contracts again. +pub fn bind<'a>( + dag: &ExecutableDag, + sources: BTreeMap>, + roots: &[NodeId], +) -> Result, Error> { + let inputs = sources + .iter() + .map(|(&id, source)| (id, InputContract::from_source(source.as_ref()))) + .collect(); + compile(dag, inputs, roots)?.instantiate(sources) +} + +/// Resolve raw scan connectors before invoking the reader-independent compiler. +pub fn bind_with_data_sources<'a>( + dag: &ExecutableDag, + mut sources: BTreeMap>, + roots: &[NodeId], + data_sources: &crate::sources::DataSources, +) -> Result, Error> { + // Only resolve scans reachable below the selected input boundaries. + let mut pending = roots.to_vec(); + let mut seen = BTreeSet::new(); + while let Some(id) = pending.pop() { + if !seen.insert(id) || sources.contains_key(&id) { + continue; + } + let node = dag + .nodes + .iter() + .find(|n| u64::from(n.id.0) == id) + .ok_or_else(|| invalid(format!("missing node {id}")))?; + if let Payload::Fallback { + expression: expression @ QueryExpr::Scan { .. }, + } = &node.payload + { + sources.insert(id, Box::new(data_sources.bind(expression)?)); + } else { + pending.extend( + dag.edges + .iter() + .filter(|e| u64::from(e.consumer.0) == id) + .map(|e| u64::from(e.producer.0)), + ); + } + } + bind(dag, sources, roots) +} + +fn compile_internal( + dag: &ExecutableDag, + mut sources: BTreeMap, + roots: &[NodeId], +) -> Result { + 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(); + // Binary input order is semantic; serialized edge order is not. + let mut edges = dag.edges.iter().collect::>(); + edges.sort_by_key(|edge| { + ( + edge.consumer.0, + match edge.role { + planner_types::post_asap::EdgeRole::Left => 0, + planner_types::post_asap::EdgeRole::Input => 1, + planner_types::post_asap::EdgeRole::Right => 2, + }, + ) + }); + for edge in 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 = CompiledPhysicalDag::new(roots.to_vec()); + let mut auxiliary = u64::MAX; + for id in ordered { + let node = nodes[&id]; + let output = Arc::new(node.output_schema.clone()); + crate::values::validate_schema(&output)?; + if let Some(source) = sources.remove(&id) { + if source.schema != output { + return Err(invalid("frontier does not have the declared schema")); + } + graph.add_input(id, source)?; + } 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); + } + if let Payload::Value { + operation: ValueOperation::MaintainPopulation { population }, + } = &node.payload + { + use planner_types::post_asap::maintained_population::PopulationInput; + let PopulationInput::CurrentSeries(spec) = &population.input else { + return Err(invalid( + "native maintained population requires a current-series input", + )); + }; + let [input] = schemas.as_slice() else { + return Err(invalid("current-series population requires one input")); + }; + if spec.without { + return Err(invalid( + "dynamic without grouping requires label-set projection", + )); + } + let identity = named_column( + input, + &ColumnRef::Named(promql_rows::SERIES_IDENTITY_COLUMN.into()), + )?; + let coordinate = input + .time_index + .ok_or_else(|| invalid("current-series input lacks timestamp"))?; + let value = named_column(input, &ColumnRef::SampleValue)?; + let lookback = i64::try_from(spec.lookback_ms) + .map_err(|_| invalid("current-series lookback overflows"))?; + graph.add( + id, + inputs, + Operator::current_series(input.clone(), identity, coordinate, value, lookback)? + .with_output_schema(output)?, + )?; + continue; + } + if let Payload::Value { + operation: ValueOperation::ReadPopulation { readout }, + } = &node.payload + { + use planner_types::post_asap::maintained_population::{ + PopulationInput, PopulationReadout, + }; + let PopulationReadout::TopK { k } = readout else { + return Err(invalid( + "native population readout does not support this operation", + )); + }; + let [producer] = inputs.as_slice() else { + return Err(invalid("population readout requires one input")); + }; + let Payload::Value { + operation: ValueOperation::MaintainPopulation { population }, + } = &nodes[producer].payload + else { + return Err(invalid( + "population readout requires its declared population", + )); + }; + let PopulationInput::CurrentSeries(spec) = &population.input else { + return Err(invalid("current-series population required")); + }; + if spec.without { + return Err(invalid( + "dynamic without ranking requires label-set projection", + )); + } + let input = schemas[0].clone(); + let groups = spec + .grouping + .iter() + .map(|name| named_column(&input, &ColumnRef::Named(name.clone()))) + .collect::, _>>()?; + let value = named_column(&input, &ColumnRef::SampleValue)?; + graph.add( + auxiliary, + inputs, + Operator::sort( + input.clone(), + vec![SortKey { + column: value, + descending: true, + nulls_first: false, + }], + groups.clone(), + )?, + )?; + graph.add( + id, + vec![auxiliary], + Operator::limit(input, *k as u64, 0, groups)?.with_output_schema(output)?, + )?; + auxiliary -= 1; + continue; + } + // A closed row must include either all source labels or the explicit + // complete-label identity. Projected labels alone are insufficient. + if let Payload::SummaryAgg { + family, + input: update, + reduction: PlannerReduction::PerEntity, + grouping, + } = &node.payload + { + let [input_id] = inputs.as_slice() else { + return Err(invalid("per-entity summary requires one input")); + }; + let Payload::Fallback { + expression: QueryExpr::TimeRange { child, .. }, + } = &nodes[input_id].payload + else { + return Err(invalid( + "per-entity summary requires a resolved raw time range", + )); + }; + let QueryExpr::Scan { schema, .. } = child.as_ref() else { + return Err(invalid("per-entity summary requires a resolved source")); + }; + if !schema.closed || update.item.is_some() { + return Err(invalid( + "per-entity summary requires complete source identity", + )); + } + crate::capability::validate_summary_kernel(family, update, grouping) + .map_err(Error::Invalid)?; + let SummaryInputExpr::Column(value) = &update.weight else { + return Err(invalid( + "per-entity update requires a projected value column", + )); + }; + let input = schemas[0].clone(); + let value = named_column(&input, value)?; + let coordinate = input + .time_index + .ok_or_else(|| invalid("temporal input lacks time"))?; + let groups = (0..input.fields.len()) + .filter(|&column| column != value && column != coordinate) + .collect(); + let build = Operator::summary_build( + input, + family.clone(), + value, + Some(coordinate), + groups, + )?; + let compact = build.schema(); + graph.add(auxiliary, inputs, build)?; + graph.add( + id, + vec![auxiliary], + Operator::scope_timestamp(compact, output)?, + )?; + auxiliary -= 1; + continue; + } + let mut operator = compile_node(node, &schemas) + .map_err(|error| invalid(format!("node {id}: {error}")))?; + if operator.is_counter_readout() { + let mut pending = vec![id]; + let mut visited = BTreeSet::new(); + let mut ranges = BTreeSet::new(); + while let Some(ancestor) = pending.pop() { + if !visited.insert(ancestor) { + continue; + } + if let Payload::Fallback { + expression: QueryExpr::TimeRange { range, .. }, + } = &nodes[&ancestor].payload + { + ranges.insert( + i64::try_from(range.as_millis()) + .map_err(|_| invalid("counter lookback exceeds Int64"))?, + ); + continue; + } + pending.extend(dependencies.get(&ancestor).into_iter().flatten().copied()); + } + if ranges.len() > 1 { + return Err(invalid("counter readout has ambiguous logical windows")); + } + if let Some(lookback) = ranges.into_iter().next() { + operator = operator.with_counter_lookback(lookback)?; + } + } + graph.add(id, inputs, operator)?; + } + } + graph.validate()?; + Ok(graph) +} + +/// Bind a Planner node against the schemas supplied by its deployment edges. +/// This is the same checked path used by complete DAG binding. +pub fn compile_node(node: &ExecutableDagNode, inputs: &[Schema]) -> Result { + for schema in inputs { + crate::values::validate_schema(schema)?; + } + bind_operation(node, inputs)?.with_output_schema(Arc::new(node.output_schema.clone())) +} + +fn bind_operation(node: &ExecutableDagNode, inputs: &[Schema]) -> Result { + if let Payload::RelationalJoin { + join_kind, + pred, + pruning, + } = &node.payload + { + use planner_types::{post_asap::CandidateCompleteness, pre_asap::JoinKind}; + if pruning.is_some() && *join_kind != JoinKind::Semi { + return Err(invalid("pruning certificate requires a semi-join")); + } + if matches!(pruning,Some(CandidateCompleteness::Certified { guarantee }) if guarantee.has_unknown() || guarantee.metric != planner_types::post_asap::ErrorMetric::TopKMembership) + { + return Err(invalid("invalid pruning certificate")); + } + let [left, right] = inputs else { + return Err(invalid("join requires two inputs")); + }; + if *join_kind == JoinKind::Semi { + if let Ok(keys) = equijoin_keys(pred, left, right) { + return Operator::semi_join(left.clone(), right.clone(), keys); + } + } + return Operator::relational_join( + left.clone(), + right.clone(), + join_kind.clone(), + pred, + Arc::new(node.output_schema.clone()), + ); + } + 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, input)?, + )) + }) + .collect::>()?, + ), + ValueOperation::Filter { pred } => { + Operator::filter(input.clone(), expression(&pred.0, input)?) + } + 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, + partition_by, + } => Operator::limit( + input.clone(), + *n as u64, + *offset as u64, + groups(input, partition_by)?, + ), + 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 let Some(item) = &update.item { + let PlannerReduction::Reduce(keys) = reduction else { + return Err(invalid("keyed summary requires explicit partitions")); + }; + let SummaryInputExpr::Column(weight) = &update.weight else { + return Err(invalid( + "keyed summary weight must be a finalized value column", + )); + }; + if matches!(family, SummaryFamilyType::Sketch(kind, _) if kind.algorithm() == &planner_types::post_asap::SketchAlgorithm::CmsWithHeap) + && !matches!( + update.weight_domain, + planner_types::post_asap::WeightDomain::NonNegative { .. } + ) + { + return Err(invalid("CMS requires a nonnegative weight contract")); + } + fn columns( + expr: &SummaryInputExpr, + input: &Schema, + result: &mut Vec, + ) -> Result<(), Error> { + match expr { + SummaryInputExpr::Column(column) => { + result.push(named_column(input, column)?) + } + SummaryInputExpr::Tuple(items) => { + for item in items { + columns(item, input, result)?; + } + } + _ => return Err(invalid("keyed summary needs explicit item columns")), + } + Ok(()) + } + let mut items = Vec::new(); + columns(item, input, &mut items)?; + return Operator::keyed_summary_build( + input.clone(), + family.clone(), + named_column(input, weight)?, + items, + groups(input, keys)?, + ); + } + 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 } => { + if let SketchQuery::TopK { k } = query { + return Operator::keyed_readout( + input.clone(), + summary_column(input)?, + *k, + Arc::new(node.output_schema.clone()), + ); + } + 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 { + // Executable SummarySchema retains column names, not table qualifiers. + // Frontend binding has resolved the qualifier; still reject ambiguous + // names here rather than guessing a join side. + ColumnRef::Named(name) | ColumnRef::Qualified { 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, input: &Schema) -> Result { + Ok(Expression::planner( + crate::expressions::CompiledExpression::compile(expr, input)?, + )) +} + +struct CheckedSource<'a> { + source: Source<'a>, + output: Schema, +} +impl PhysicalOperator for CheckedSource<'_> { + fn properties(&self, inputs: &[crate::plan::PlanProperties]) -> crate::plan::PlanProperties { + self.source.properties(inputs) + } + + 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: crate::runtime::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(()) +} + +/// Join predicates address the concatenated left/right schema. +fn semi_join_keys( + expr: &QueryExpr, + left: usize, + right: usize, + keys: &mut Vec<(usize, usize)>, +) -> Result<(), Error> { + match expr { + QueryExpr::BoolAnd(parts) => { + for part in parts { + semi_join_keys(part, left, right, keys)?; + } + } + QueryExpr::Compare { + left: a, + op: CompareOpKind::Eq, + right: b, + } => { + let (QueryExpr::Column(a), QueryExpr::Column(b)) = (a.as_ref(), b.as_ref()) else { + return Err(invalid("semi-join requires column equality keys")); + }; + let (a, b) = if a < b { (*a, *b) } else { (*b, *a) }; + if a >= left || b < left || b >= left + right { + return Err(invalid("semi-join key must match left to right")); + } + keys.push((a, b - left)); + } + _ => return Err(invalid("unsupported semi-join predicate")), + } + Ok(()) +} + +/// Resolve equality keys against the Planner join's concatenated input schema. +/// Deployments may use these positions to bind their source columns. +pub fn equijoin_keys( + pred: &planner_types::pre_asap::Predicate, + left: &planner_types::post_asap::SummarySchema, + right: &planner_types::post_asap::SummarySchema, +) -> Result, Error> { + let mut keys = Vec::new(); + semi_join_keys(&pred.0, left.fields.len(), right.fields.len(), &mut keys)?; + if keys.is_empty() { + return Err(invalid("semi-join requires explicit matching keys")); + } + Ok(keys) +} diff --git a/crates/asap-physical-operators/src/physical_planner/promql_rows.rs b/crates/asap-physical-operators/src/physical_planner/promql_rows.rs new file mode 100644 index 00000000..83208f50 --- /dev/null +++ b/crates/asap-physical-operators/src/physical_planner/promql_rows.rs @@ -0,0 +1,389 @@ +//! A bounded PromQL source row carries the entire label set, not just labels +//! mentioned by the query. The source adapter owns this lossless encoding. +use super::*; +use planner_types::pre_asap::{Column, DataType, Source as LogicalSource}; +use std::rc::Rc; + +/// Not a legal PromQL label name, so it cannot shadow a user label. +pub use planner_types::pre_asap::schema::PROMQL_SERIES_IDENTITY as SERIES_IDENTITY_COLUMN; + +/// Canonical, reversible identity. JSON object encoding preserves label names, +/// empty values and escaping; sorting makes ingestion order irrelevant. +pub fn encode_series_identity(labels: &BTreeMap) -> Result { + serde_json::to_string(labels).map_err(|error| invalid(error.to_string())) +} + +pub fn decode_series_identity(encoded: &str) -> Result, Error> { + let labels: BTreeMap = + serde_json::from_str(encoded).map_err(|error| invalid(error.to_string()))?; + if encode_series_identity(&labels)? != encoded { + return Err(invalid("series identity is not canonically encoded")); + } + Ok(labels) +} + +/// Resolve the row representation before candidate search. `closed` describes +/// physical columns here: the final column contains every dynamic source label. +/// It does not assert that the query's projected labels are the full label set. +/// +/// This realization supports explicit `by` grouping and per-series computation. +/// Operators that rewrite or implicitly match dynamic label sets require their +/// own realization; they must not accidentally treat the opaque identity as a +/// user label or silently discard it. +pub fn with_series_identity(root: &QueryExpr) -> Result { + let mut root = root.clone(); + fn visit(node: &mut QueryExpr) -> Result<(), Error> { + use planner_types::pre_asap::Reduction; + match node { + QueryExpr::Scan { + source: LogicalSource::TimeSeries { .. }, + schema, + .. + } => { + if schema + .columns + .iter() + .any(|column| column.name == SERIES_IDENTITY_COLUMN) + { + return Err(invalid( + "source already contains a physical series identity", + )); + } + if schema.closed { + return Err(invalid( + "dynamic series identity requires an open PromQL source", + )); + } + schema + .columns + .push(Column::new(SERIES_IDENTITY_COLUMN, DataType::Utf8, false)); + schema.closed = true; + Ok(()) + } + QueryExpr::TimeRange { child, .. } | QueryExpr::Limit { child, .. } => { + visit(Rc::make_mut(child)) + } + QueryExpr::Aggregate { + child, reduction, .. + } => { + if matches!(reduction, Reduction::Reduce(keys) if keys.is_without()) { + return Err(invalid( + "dynamic without grouping requires label-set projection", + )); + } + visit(Rc::make_mut(child)) + } + QueryExpr::Sort { + child, + partition_by, + .. + } => { + if partition_by.is_without() { + return Err(invalid( + "dynamic without ranking requires label-set projection", + )); + } + visit(Rc::make_mut(child)) + } + _ => Err(invalid( + "operator has no dynamic series-identity realization", + )), + } + } + visit(&mut root)?; + root.output_schema() + .map_err(|error| invalid(error.to_string()))?; + Ok(root) +} + +/// Construct source rows only from full identities. The named label columns +/// are projections of that same identity and cannot independently redefine it. +pub fn series_row( + schema: &Schema, + labels: &BTreeMap, + timestamp: i64, + value: f64, +) -> Result, Error> { + use crate::values::Value; + let identity = encode_series_identity(labels)?; + let mut found = false; + let row = schema + .fields + .iter() + .enumerate() + .map(|(index, field)| { + if field.name == SERIES_IDENTITY_COLUMN { + if field.dtype != SummaryFamilyType::Plain(DataType::Utf8) + || field.nullable + || found + { + return Err(invalid("invalid series identity column")); + } + found = true; + Ok(Value::Utf8(identity.clone().into())) + } else if Some(index) == schema.time_index { + Ok(Value::Timestamp(timestamp)) + } else if field.name == "value" + && field.dtype == SummaryFamilyType::Plain(DataType::Float64) + { + Ok(Value::Float64(value)) + } else if field.dtype == SummaryFamilyType::Plain(DataType::Utf8) { + Ok(labels.get(&field.name).map_or_else( + || Value::Utf8("".into()), + |value| Value::Utf8(value.clone().into()), + )) + } else { + Err(invalid("unsupported PromQL source column")) + } + }) + .collect::, _>>()?; + if !found { + return Err(invalid("source lacks its full series identity")); + } + Ok(row) +} + +/// Compile the selected TopK computation above an existing maintained-population +/// source. The boundary supplies the complete eligible vector, not a truncated +/// TopK result; ranking remains a native physical operator. +pub fn compile_current_series_readout( + selected: &Rc, +) -> Result { + use planner_types::post_asap::{ + compile_executable_dag, maintained_population::PopulationReadout, SummaryField, + }; + let mut dag = compile_executable_dag(selected).map_err(|error| invalid(error.to_string()))?; + // Typed snapshot candidates already carry full identity throughout the DAG. + // Cut at the population output, preserving all selected heap/readout nodes. + let populations = dag.nodes.iter().filter(|node| matches!(&node.payload, + Payload::Value { operation: ValueOperation::MaintainPopulation { population } } + if matches!(population.input, planner_types::post_asap::maintained_population::PopulationInput::CurrentSeries(_)) + )).collect::>(); + if let [population] = populations.as_slice() { + if population + .output_schema + .fields + .iter() + .any(|field| field.name == SERIES_IDENTITY_COLUMN) + { + return compile( + &dag, + BTreeMap::from([( + u64::from(population.id.0), + InputContract::bounded(Arc::new(population.output_schema.clone())), + )]), + &[u64::from(dag.root.0)], + ); + } + } + if dag.nodes.len() != 3 + || !dag.nodes.iter().any(|node| { + node.id == dag.root + && matches!( + node.payload, + Payload::Value { + operation: ValueOperation::ReadPopulation { + readout: PopulationReadout::TopK { .. } + } + } + ) + }) + { + return Err(invalid( + "expected one selected current-series TopK computation", + )); + } + let mut frontier = None; + for node in &mut dag.nodes { + match &mut node.payload { + Payload::Fallback { expression } => { + *expression = with_series_identity(expression)?; + } + Payload::Value { + operation: ValueOperation::MaintainPopulation { .. }, + } => { + frontier = Some(u64::from(node.id.0)); + } + Payload::Value { + operation: + ValueOperation::ReadPopulation { + readout: PopulationReadout::TopK { .. }, + }, + } => {} + _ => return Err(invalid("unsupported current-series readout dependency")), + } + if node + .output_schema + .fields + .iter() + .any(|field| field.name == SERIES_IDENTITY_COLUMN) + { + return Err(invalid( + "current-series input already has a physical identity column", + )); + } + node.output_schema.fields.push(SummaryField { + name: SERIES_IDENTITY_COLUMN.into(), + dtype: SummaryFamilyType::Plain(DataType::Utf8), + nullable: false, + }); + } + for edge in &mut dag.edges { + edge.intermediate_schema = dag + .nodes + .iter() + .find(|node| node.id == edge.producer) + .unwrap() + .output_schema + .clone(); + } + let frontier = frontier.ok_or_else(|| invalid("missing current-series population"))?; + let schema = Arc::new( + dag.nodes + .iter() + .find(|node| u64::from(node.id.0) == frontier) + .unwrap() + .output_schema + .clone(), + ); + compile( + &dag, + BTreeMap::from([(frontier, InputContract::bounded(schema))]), + &[u64::from(dag.root.0)], + ) +} + +/// Compile selected ranking or aggregation above an exact per-series Rate +/// readout. Deployments bind complete window readouts at this boundary; +/// the heap is rebuilt independently for each evaluation. This does not move +/// that frontier to ingestion time or authorize combining finalized rates. +pub fn compile_rate_ranking( + selected: &Rc, +) -> Result< + ( + Rc, + CompiledPhysicalDag, + ), + Error, +> { + use planner_types::post_asap::{ + compile_executable_dag_with_node_ids, ExactKind, SummaryExpr, SummaryNode, + }; + fn frontier(node: &Rc) -> Option> { + match &node.expr { + SummaryExpr::ValueOperation { + child, + operation: ValueOperation::FinalizeExactAccumulator, + timing: planner_types::post_asap::ExecutionTiming::QueryTime, + } if matches!(&child.expr, SummaryExpr::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + reduction: planner_types::pre_asap::Reduction::PerEntity, + child: raw, .. + } if matches!(&raw.expr, SummaryExpr::KeepPreAsap(expr) if matches!(expr.as_ref(), QueryExpr::TimeRange { .. }))) => + { + Some(Rc::clone(node)) + } + SummaryExpr::ValueOperation { child, .. } | SummaryExpr::SummaryAgg { child, .. } => { + frontier(child) + } + SummaryExpr::SummaryEstimate { summary_input, .. } => frontier(summary_input), + _ => None, + } + } + let source = frontier(selected) + .ok_or_else(|| invalid("ranking requires one exact per-series Rate frontier"))?; + if !source + .schema + .fields + .iter() + .any(|field| field.name == SERIES_IDENTITY_COLUMN) + { + return Err(invalid("Rate ranking requires complete series identity")); + } + let compiled = compile_executable_dag_with_node_ids(selected) + .map_err(|error| invalid(error.to_string()))?; + let id = u64::from( + compiled + .node_ids + .node_id(&source) + .ok_or_else(|| invalid("missing Rate frontier"))? + .0, + ); + let program = compile( + &compiled.dag, + BTreeMap::from([(id, InputContract::bounded(Arc::new(source.schema.clone())))]), + &[u64::from(compiled.dag.root.0)], + )?; + Ok((source, program)) +} + +/// The selected logical placement requires fresh aggregate state per closed window. +/// Compile both physical graphs before deployment chooses storage or scheduling. +/// The input is the complete collection of per-series exact counter states. +pub fn compile_fixed_window_rate_aggregation( + selected: &Rc, +) -> Result { + use planner_types::post_asap::{ + compile_executable_dag, ExactKind, ExecutionTiming, SketchAlgorithm, + }; + let dag = compile_executable_dag(selected).map_err(|e| invalid(e.to_string()))?; + let sources = dag + .nodes + .iter() + .filter(|n| { + matches!( + &n.payload, + Payload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + reduction: planner_types::pre_asap::Reduction::PerEntity, + .. + } + ) + }) + .collect::>(); + let heaps = dag + .nodes + .iter() + .filter(|n| { + n.output_state.timing == ExecutionTiming::IngestionTime + && match &n.payload { + Payload::SummaryAgg { + family: SummaryFamilyType::Sketch(kind, _), + .. + } => matches!( + kind.algorithm(), + SketchAlgorithm::CmsWithHeap | SketchAlgorithm::CountSketchWithHeap + ), + Payload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), + .. + } => true, + _ => false, + } + }) + .collect::>(); + let ([source], [heap]) = (sources.as_slice(), heaps.as_slice()) else { + return Err(invalid( + "expected one selected fixed-window Rate aggregation", + )); + }; + if !source + .output_schema + .fields + .iter() + .any(|f| f.name == SERIES_IDENTITY_COLUMN) + { + return Err(invalid( + "fixed-window Rate aggregation requires complete series identity", + )); + } + compile_candidate( + &dag, + BTreeMap::from([( + u64::from(source.id.0), + InputContract::bounded(Arc::new(source.output_schema.clone())), + )]), + &[u64::from(dag.root.0)], + &[u64::from(heap.id.0)], + ) +} diff --git a/crates/asap-physical-operators/src/physical_planner/temporal_panes.rs b/crates/asap-physical-operators/src/physical_planner/temporal_panes.rs new file mode 100644 index 00000000..4e4d5722 --- /dev/null +++ b/crates/asap-physical-operators/src/physical_planner/temporal_panes.rs @@ -0,0 +1,332 @@ +//! Lower a selected temporal maintenance contract; deployment supplies readers. +use super::*; +use planner_types::post_asap::{ + EvaluationSchedule, OutputRepresentation, PaneLayout, SketchAlgorithm, + SummaryMaintenanceLifecycle, SummaryMaintenanceLifecycleGuarantee, SummaryMaintenanceMode, + SummaryWindowFramework, +}; + +/// Resolved source identity, supplied with physical capability evidence. +/// A schemaless PromQL projection cannot establish the complete label set. +#[derive(Clone, Debug)] +pub enum TemporalEntityIdentity { + /// The input resolver guarantees that the slot contains one entity. + SingleEntity, + /// All entity keys are represented by these columns; there are no hidden + /// labels distinguishing two rows with the same key. + Columns(Vec), +} + +/// Planner-selected lifecycle/window requirements for one temporal producer. +/// Pane geometry is semantic input, not a storage identity or scheduling policy. +#[derive(Clone, Debug)] +pub struct TemporalPaneMaintenance { + pub summary_node: NodeId, + pub lifecycle: SummaryMaintenanceLifecycleGuarantee, + pub framework: SummaryWindowFramework, + pub layout: PaneLayout, + pub entity_identity: TemporalEntityIdentity, +} + +/// Generated precompute and query computation. `pane_inputs` is ordered from +/// the oldest complete pane to the newest; each run checks actual timestamps. +#[derive(Clone)] +pub struct TemporalPaneCandidate { + pub physical: PhysicalCandidate, + pub maintenance: TemporalPaneMaintenance, + pub pane_inputs: Vec, + pub merged_state: NodeId, + pub window_width_ms: u64, +} + +/// Compile bounded pane construction and a shared pane merge for temporal KLL +/// quantile roots. The selected contract remains authoritative; unsupported +/// lifecycle/framework/operator shapes fail rather than being substituted. +/// This initial realization consumes complete pane populations and emits full +/// state snapshots. Cross-run delta accumulation belongs to other candidates. +pub fn compile_temporal_pane_candidate( + dag: &ExecutableDag, + inputs: BTreeMap, + roots: &[NodeId], + maintenance: &TemporalPaneMaintenance, +) -> Result { + dag.validate().map_err(|error| invalid(error.to_string()))?; + if maintenance.lifecycle.summary_maintenance_lifecycle + != SummaryMaintenanceLifecycle::ContinuouslyMaintained + || maintenance.lifecycle.summary_maintenance_mode != SummaryMaintenanceMode::Incremental + || maintenance.lifecycle.evaluation_schedule != EvaluationSchedule::PerUpdate + || maintenance.lifecycle.output_representation != OutputRepresentation::SummaryState + { + return Err(invalid( + "pane candidate requires continuous incremental summary maintenance", + )); + } + let build = dag + .nodes + .iter() + .find(|node| u64::from(node.id.0) == maintenance.summary_node) + .ok_or_else(|| invalid("unknown maintained producer"))?; + let Payload::SummaryAgg { + family, + input: update, + reduction: PlannerReduction::PerEntity, + grouping, + } = &build.payload + else { + return Err(invalid( + "pane candidate requires a temporal per-entity summary", + )); + }; + if !matches!(family, SummaryFamilyType::Sketch(kind, _) if kind.algorithm() == &SketchAlgorithm::Kll) + || update.item.is_some() + { + return Err(invalid("pane candidate supports unkeyed temporal KLL only")); + } + crate::capability::validate_summary_kernel(family, update, grouping).map_err(Error::Invalid)?; + let dependencies: Vec<_> = dag + .edges + .iter() + .filter(|edge| edge.consumer == build.id) + .map(|edge| edge.producer) + .collect(); + let [raw_id] = dependencies.as_slice() else { + return Err(invalid("temporal producer requires one raw input")); + }; + let raw = dag + .nodes + .iter() + .find(|node| node.id == *raw_id) + .ok_or_else(|| invalid("missing raw input"))?; + let Payload::Fallback { + expression: QueryExpr::TimeRange { range, child }, + } = &raw.payload + else { + return Err(invalid( + "temporal producer requires an explicit logical time range", + )); + }; + let QueryExpr::Scan { predicates, .. } = child.as_ref() else { + return Err(invalid("temporal pane source requires a raw scan")); + }; + let window_width_ms: u64 = range + .as_millis() + .try_into() + .map_err(|_| invalid("temporal window overflows"))?; + if window_width_ms == 0 + || window_width_ms > i64::MAX as u64 + || range.subsec_nanos() % 1_000_000 != 0 + { + return Err(invalid( + "temporal window requires positive integral milliseconds", + )); + } + let width = maintenance.layout.pane_width_ms; + if width == 0 || width > window_width_ms || !window_width_ms.is_multiple_of(width) { + return Err(invalid("temporal window must contain whole panes")); + } + match maintenance.framework { + SummaryWindowFramework::Sliding => {} + SummaryWindowFramework::Tumbling if width == window_width_ms => {} + _ => return Err(invalid("unsupported temporal window realization")), + } + let count = window_width_ms / width; + if count > 4096 { + return Err(invalid("temporal pane candidate exceeds input budget")); + } + let raw_id = u64::from(raw_id.0); + if inputs.len() != 1 { + return Err(invalid( + "pane candidate requires exactly its raw input contract", + )); + } + let contract = inputs + .get(&raw_id) + .ok_or_else(|| invalid("missing raw input contract"))?; + let raw_schema = Arc::new(raw.output_schema.clone()); + if contract.schema != raw_schema || contract.properties.boundedness != Boundedness::Bounded { + return Err(invalid("pane source requires its declared bounded schema")); + } + let coordinate = raw_schema + .time_index + .ok_or_else(|| invalid("temporal source requires a time index"))?; + let SummaryInputExpr::Column(value) = &update.weight else { + return Err(invalid("pane builder requires a value column")); + }; + let value = named_column(&raw_schema, value)?; + let groups: Vec<_> = (0..raw_schema.fields.len()) + .filter(|&index| index != coordinate && index != value) + .collect(); + match &maintenance.entity_identity { + TemporalEntityIdentity::SingleEntity if groups.is_empty() => {} + TemporalEntityIdentity::Columns(columns) + if !columns.is_empty() + && columns.len() == columns.iter().collect::>().len() + && columns.iter().copied().collect::>() + == groups.iter().copied().collect() => {} + _ => { + return Err(invalid( + "pane input requires its complete resolved entity identity", + )) + } + } + let mut next = dag + .nodes + .iter() + .map(|node| u64::from(node.id.0)) + .max() + .unwrap_or(0) + + 1; + let mut allocate = || { + let id = next; + next += 1; + id + }; + let mut operators = BTreeMap::new(); + let guard = allocate(); + operators.insert( + guard, + ( + vec![raw_id], + Operator::pane_input( + raw_schema.clone(), + coordinate, + maintenance.layout.clone(), + None, + )?, + ), + ); + let mut previous = guard; + for predicate in predicates { + let id = allocate(); + operators.insert( + id, + ( + vec![previous], + Operator::filter(raw_schema.clone(), expression(&predicate.0, &raw_schema)?)?, + ), + ); + previous = id; + } + let native = + Operator::summary_build(raw_schema, family.clone(), value, Some(coordinate), groups)?; + let compact_state = native.schema(); + let native_id = allocate(); + operators.insert(native_id, (vec![previous], native)); + let state_schema = Arc::new(build.output_schema.clone()); + let pane_output = allocate(); + operators.insert( + pane_output, + ( + vec![native_id], + Operator::scope_timestamp(compact_state, state_schema.clone())?, + ), + ); + let precompute = CompiledPhysicalDag::from_operators(inputs, operators, vec![pane_output])?; + let state_coordinate = state_schema + .time_index + .ok_or_else(|| invalid("pane state requires a time index"))?; + let state_column = summary_column(&state_schema)?; + let mut query_inputs = BTreeMap::new(); + let mut operators = BTreeMap::new(); + let mut pane_inputs = Vec::new(); + let mut guarded_inputs = Vec::new(); + for pane in 0..count { + let input = allocate(); + let guard = allocate(); + query_inputs.insert(input, InputContract::bounded(state_schema.clone())); + let offset = ((count - 1 - pane) * width) as i64; + operators.insert( + guard, + ( + vec![input], + Operator::pane_input( + state_schema.clone(), + state_coordinate, + maintenance.layout.clone(), + Some(offset), + )?, + ), + ); + pane_inputs.push(input); + guarded_inputs.push(guard); + } + let union = allocate(); + operators.insert( + union, + ( + guarded_inputs, + Operator::union(state_schema.clone(), count as usize)?, + ), + ); + let merge = Operator::summary_merge( + state_schema.clone(), + state_column, + (0..state_schema.fields.len()) + .filter(|&index| index != state_coordinate && index != state_column) + .collect(), + )?; + let merged_schema = merge.schema(); + let merged_state = allocate(); + operators.insert(merged_state, (vec![union], merge)); + if roots.is_empty() || roots.iter().copied().collect::>().len() != roots.len() { + return Err(invalid("temporal query requires distinct output roots")); + } + for &root in roots { + let node = dag + .nodes + .iter() + .find(|node| u64::from(node.id.0) == root) + .ok_or_else(|| invalid("unknown temporal output root"))?; + let Payload::SummaryEstimate { + query: SketchQuery::Quantile { q }, + } = node.payload + else { + return Err(invalid("temporal pane root must be a KLL quantile")); + }; + if !q.is_finite() || !(0. ..=1.).contains(&q) { + return Err(invalid("invalid temporal quantile")); + } + let dependencies: Vec<_> = dag + .edges + .iter() + .filter(|edge| edge.consumer == node.id) + .map(|edge| u64::from(edge.producer.0)) + .collect(); + if dependencies != [maintenance.summary_node] { + return Err(invalid( + "temporal readout must consume the maintained producer", + )); + } + let readout = Operator::readout( + merged_schema.clone(), + summary_column(&merged_schema)?, + crate::Statistic::Quantile, + std::collections::HashMap::from([("quantile".into(), q.to_string())]), + )?; + let readout_schema = readout.schema(); + let readout_id = allocate(); + operators.insert(readout_id, (vec![merged_state], readout)); + operators.insert( + root, + ( + vec![readout_id], + Operator::scope_timestamp(readout_schema, Arc::new(node.output_schema.clone()))?, + ), + ); + } + let query = CompiledPhysicalDag::from_operators(query_inputs, operators, roots.to_vec())?; + let mut output = precompute.output_contract(pane_output)?; + // Persisted readers have independent timing from the blocking builder. + output.properties.emission = Emission::Unknown; + Ok(TemporalPaneCandidate { + physical: PhysicalCandidate { + precompute: Some(precompute), + query, + materialized_outputs: BTreeMap::from([(pane_output, output)]), + }, + maintenance: maintenance.clone(), + pane_inputs, + merged_state, + window_width_ms, + }) +} diff --git a/crates/asap-physical-operators/src/plan/mod.rs b/crates/asap-physical-operators/src/plan/mod.rs new file mode 100644 index 00000000..0dce4dd8 --- /dev/null +++ b/crates/asap-physical-operators/src/plan/mod.rs @@ -0,0 +1,160 @@ +//! Immutable physical graph, operator contracts and pre-execution validation. +use crate::{ + runtime::{Input, OutputStream, RunContext}, + Error, +}; +use std::{ + collections::{BTreeMap, BTreeSet}, + fmt::Debug, +}; +pub type NodeId = u64; +mod properties; +pub use properties::{Boundedness, Emission, PlanProperties}; +/// 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; + /// Source implementations must explicitly declare finite input before feeding blocking operators. + fn properties(&self, inputs: &[PlanProperties]) -> PlanProperties { + PlanProperties { + boundedness: Boundedness::from_inputs(inputs), + emission: Emission::Unknown, + } + } + fn requires_bounded_input(&self) -> bool { + false + } + + /// Validate run-specific contracts before any source is opened. + fn validate_context(&self, _context: &RunContext) -> Result<(), Error> { + Ok(()) + } + 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; +} +pub(crate) struct Node<'a, V, S> { + pub(crate) inputs: Vec, + pub(crate) operator: Box + 'a>, +} +pub struct PhysicalDag<'a, V, S> { + pub(crate) 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> { + self.properties(roots).map(|_| ()) + } + /// Derive properties while checking topology and schemas, before starting sources. + pub fn properties(&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; + let mut input_properties = Vec::new(); + for (input, schema) in node.inputs.iter().zip(expected) { + depth = depth.max(1 + visit(dag, *input, active, done)?); + input_properties.push(done[input].1); + 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(), + )); + } + if node.operator.requires_bounded_input() + && input_properties + .iter() + .any(|p| p.boundedness != Boundedness::Bounded) + { + return Err(Error::Invalid(format!( + "node {id} ({}) requires bounded inputs", + node.operator.name() + ))); + } + let properties = node.operator.properties(&input_properties); + active.remove(&id); + done.insert(id, (depth, properties)); + 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(done + .into_iter() + .map(|(id, (_, properties))| (id, properties)) + .collect()) + } + pub fn execute<'r>( + &'r self, + roots: &[NodeId], + context: RunContext, + ) -> Result>, Error> + where + 'a: 'r, + { + crate::runtime::execute(self, roots, context) + } +} diff --git a/crates/asap-physical-operators/src/plan/properties.rs b/crates/asap-physical-operators/src/plan/properties.rs new file mode 100644 index 00000000..7ec770ea --- /dev/null +++ b/crates/asap-physical-operators/src/plan/properties.rs @@ -0,0 +1,32 @@ +//! Execution facts used to reject operators that cannot finish on their inputs. +#[derive(serde::Serialize, serde::Deserialize, Clone, Copy, Debug, PartialEq, Eq)] +pub enum Boundedness { + /// The source or operator promises a finite result for this run. + Bounded, + Unbounded, + /// No finite-input guarantee has been supplied. + Unknown, +} +impl Boundedness { + pub fn from_inputs(inputs: &[PlanProperties]) -> Self { + if inputs.iter().any(|p| p.boundedness == Self::Unbounded) { + Self::Unbounded + } else if inputs.is_empty() || inputs.iter().any(|p| p.boundedness == Self::Unknown) { + Self::Unknown + } else { + Self::Bounded + } + } +} +#[derive(serde::Serialize, serde::Deserialize, Clone, Copy, Debug, PartialEq, Eq)] +pub enum Emission { + Incremental, + /// Produces its result only after all inputs end, even if accumulation is incremental. + AfterInput, + Unknown, +} +#[derive(serde::Serialize, serde::Deserialize, Clone, Copy, Debug, PartialEq, Eq)] +pub struct PlanProperties { + pub boundedness: Boundedness, + pub emission: Emission, +} diff --git a/crates/asap-physical-operators/src/runtime/batch_execution.rs b/crates/asap-physical-operators/src/runtime/batch_execution.rs new file mode 100644 index 00000000..86809f2b --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/batch_execution.rs @@ -0,0 +1,206 @@ +//! 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 crate::{ + operators::Operator, + plan::PhysicalDag, + runtime::{RunContext, SharedValue}, + values::Batch, + Error, +}; +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) +} + +/// Bind the ordered in-memory inputs of a native multi-input operator. +pub fn evaluate_inputs( + inputs: Vec, + operator: Operator, + context: RunContext, +) -> Result>, Error> { + let mut graph = PhysicalDag::default(); + let root = inputs.len() as u64; + for (id, input) in inputs.into_iter().enumerate() { + graph.add( + id as u64, + vec![], + Operator::source(input.schema().clone(), vec![input])?, + )?; + } + graph.add(root, (0..root).collect(), operator)?; + 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, crate::values::Schema>, + root: crate::plan::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/runtime/context.rs b/crates/asap-physical-operators/src/runtime/context.rs new file mode 100644 index 00000000..c145a457 --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/context.rs @@ -0,0 +1,133 @@ +use crate::Error; +use std::{ + cell::{Cell, RefCell}, + rc::Rc, + task::Waker, +}; +/// 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, + } + } +} +pub(super) struct Control { + cancelled: Cell, + bytes: Cell, + peak: Cell, + pub(super) limits: Limits, + waiters: RefCell>, +} +#[derive(Clone)] +pub struct RunContext { + pub scope: Scope, + pub(super) 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), + }) + } + pub(super) 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, + pub(super) 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)); + } +} diff --git a/crates/asap-physical-operators/src/runtime/cooperative.rs b/crates/asap-physical-operators/src/runtime/cooperative.rs new file mode 100644 index 00000000..fae869d6 --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/cooperative.rs @@ -0,0 +1,40 @@ +//! Worker-local CPU loops yield so other consumers and cancellation can progress. +use super::RunContext; +use crate::Error; +use std::task::Poll; + +pub(crate) struct Cooperative { + context: RunContext, + remaining: usize, +} +impl Cooperative { + pub(crate) fn new(context: &RunContext) -> Self { + Self { + context: context.clone(), + remaining: 1024, + } + } + pub(crate) async fn checkpoint(&mut self) -> Result<(), Error> { + if self.context.is_cancelled() { + return Err(Error::Cancelled); + } + self.remaining -= 1; + if self.remaining == 0 { + self.remaining = 1024; + let mut yielded = false; + futures::future::poll_fn(|cx| { + if self.context.is_cancelled() { + return Poll::Ready(Err(Error::Cancelled)); + } + if yielded { + return Poll::Ready(Ok(())); + } + yielded = true; + cx.waker().wake_by_ref(); + Poll::Pending + }) + .await?; + } + Ok(()) + } +} diff --git a/crates/asap-physical-operators/src/runtime/mod.rs b/crates/asap-physical-operators/src/runtime/mod.rs new file mode 100644 index 00000000..f72a57ef --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/mod.rs @@ -0,0 +1,272 @@ +//! Per-run producer sharing, streams, backpressure and resource ownership. +use crate::{ + plan::{NodeId, PhysicalDag}, + Error, +}; +use futures::{stream::LocalBoxStream, Stream}; +use std::{ + cell::RefCell, + collections::{BTreeMap, VecDeque}, + fmt::Debug, + pin::Pin, + rc::Rc, + sync::Arc, + task::{Context, Poll, Waker}, +}; +mod context; +pub use context::{Limits, Reservation, RunContext, Scope}; +pub type OutputStream<'a, V> = LocalBoxStream<'a, Result>; +/// 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 + } +} + +pub(crate) fn execute<'r, V: 'r, S: Clone + PartialEq + Debug + 'r>( + dag: &'r PhysicalDag<'_, V, S>, + roots: &[NodeId], + context: RunContext, +) -> Result>, Error> { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + dag.validate(roots)?; + let mut pending = roots.to_vec(); + let mut visited = std::collections::BTreeSet::new(); + while let Some(id) = pending.pop() { + if visited.insert(id) { + let node = &dag.nodes[&id]; + node.operator.validate_context(&context)?; + pending.extend(node.inputs.iter().copied()); + } + } + 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(dag, 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 batch_execution; +#[cfg(test)] +mod tests; + +mod cooperative; +pub(crate) use cooperative::Cooperative; diff --git a/crates/asap-physical-operators/src/runtime/tests.rs b/crates/asap-physical-operators/src/runtime/tests.rs new file mode 100644 index 00000000..f683041c --- /dev/null +++ b/crates/asap-physical-operators/src/runtime/tests.rs @@ -0,0 +1,262 @@ +use super::*; +use crate::plan::PhysicalOperator; +use futures::{executor::block_on, stream, StreamExt}; +use std::cell::Cell; + +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/sources/memory.rs b/crates/asap-physical-operators/src/sources/memory.rs new file mode 100644 index 00000000..7f77c54f --- /dev/null +++ b/crates/asap-physical-operators/src/sources/memory.rs @@ -0,0 +1,44 @@ +use super::*; +/// Immutable in-memory raw data. The connector owns the resident input; each +/// cursor clones only the next requested batch, not the entire data set. +pub struct MemorySource { + schema: Schema, + batches: Vec, +} +impl MemorySource { + pub fn new(schema: Schema, batches: Vec) -> Result { + crate::values::validate_schema(&schema)?; + if schema + .fields + .iter() + .any(|f| !matches!(f.dtype, SummaryFamilyType::Plain(_))) + { + return Err(Error::Invalid( + "raw source cannot contain summary states".into(), + )); + } + if batches.iter().any(|batch| batch.schema() != &schema) { + return Err(Error::Invalid("memory source batch schema mismatch".into())); + } + Ok(Self { schema, batches }) + } +} +impl RawSource for MemorySource { + fn boundedness(&self) -> crate::plan::Boundedness { + crate::plan::Boundedness::Bounded + } + fn schema(&self) -> Schema { + self.schema.clone() + } + fn scan(&self, context: RunContext) -> Result, Error> { + Ok(stream::iter(self.batches.iter()) + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let _allocation = context.reserve(batch.bytes())?; + Ok(batch.clone()) + }) + .boxed_local()) + } +} diff --git a/crates/asap-physical-operators/src/sources/mod.rs b/crates/asap-physical-operators/src/sources/mod.rs new file mode 100644 index 00000000..4da778b0 --- /dev/null +++ b/crates/asap-physical-operators/src/sources/mod.rs @@ -0,0 +1,187 @@ +//! Raw data access. Connectors provide rows; Scan owns Planner predicate semantics. +use crate::{ + expressions::CompiledExpression, + plan::PhysicalOperator, + runtime::{Input, OutputStream, RunContext}, + values::{Batch, Schema, Value}, + Error, +}; +use futures::{stream, StreamExt}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::{DataType, QueryExpr, Source}, +}; +use std::sync::Arc; + +/// A bound data source. Metadata must be stable for the lifetime of the binding. +/// Each scan opens an independent cursor. Connectors return raw, unfiltered rows +/// and must honor cancellation and bound their own I/O buffers. Dropping a cursor +/// must release its resources. A connector error is never an empty successful scan. +pub trait RawSource { + fn schema(&self) -> Schema; + /// Declare a finite snapshot/window explicitly; execution scope alone does not bound a cursor. + fn boundedness(&self) -> crate::plan::Boundedness { + crate::plan::Boundedness::Unknown + } + fn scan(&self, context: RunContext) -> Result, Error>; +} + +/// Explicit source identities; no implicit network discovery or fallback. +#[derive(Default)] +pub struct DataSources { + sources: Vec<(Source, Arc)>, +} +impl DataSources { + pub fn register(&mut self, identity: Source, source: Arc) -> Result<(), Error> { + if self.sources.iter().any(|(key, _)| key == &identity) { + return Err(Error::Invalid("duplicate data source".into())); + } + crate::values::validate_schema(&source.schema())?; + self.sources.push((identity, source)); + Ok(()) + } + pub fn bind(&self, expression: &QueryExpr) -> Result { + let QueryExpr::Scan { + source, + predicates, + schema, + } = expression + else { + return Err(Error::Invalid( + "raw Scan requires a Planner Scan leaf".into(), + )); + }; + let output = Arc::new(SummarySchema { + fields: schema + .columns + .iter() + .map(|column| SummaryField { + name: column.name.clone(), + dtype: SummaryFamilyType::Plain(column.dtype.clone()), + nullable: column.nullable, + }) + .collect(), + time_index: schema.time_index, + }); + crate::values::validate_schema(&output)?; + let reader = self + .sources + .iter() + .find(|(key, _)| key == source) + .map(|(_, reader)| reader.clone()) + .ok_or_else(|| Error::Invalid(format!("unbound raw source: {source:?}")))?; + if reader.schema() != output { + return Err(Error::Invalid( + "raw source differs from Planner Scan schema".into(), + )); + } + let predicates = predicates + .iter() + .map(|predicate| { + let predicate = CompiledExpression::compile(&predicate.0, &output)?; + if predicate.dtype().0 != DataType::Bool { + return Err(Error::Invalid("Scan predicate must be boolean".into())); + } + Ok(predicate) + }) + .collect::, Error>>()?; + Ok(Scan { + reader, + output, + predicates, + }) + } +} + +pub struct Scan { + reader: Arc, + output: Schema, + predicates: Vec, +} +impl PhysicalOperator for Scan { + fn properties(&self, _: &[crate::plan::PlanProperties]) -> crate::plan::PlanProperties { + crate::plan::PlanProperties { + boundedness: self.reader.boundedness(), + emission: crate::plan::Emission::Incremental, + } + } + + fn name(&self) -> &str { + "Scan" + } + fn input_schemas(&self) -> Vec { + vec![] + } + fn output_schema(&self) -> Schema { + self.output.clone() + } + fn output_bytes(&self, batch: &Batch) -> usize { + batch.bytes() + } + fn start<'a>( + &'a self, + inputs: Vec>, + context: RunContext, + ) -> Result, Error> { + if !inputs.is_empty() { + return Err(Error::Invalid("Scan cannot have inputs".into())); + } + if context.is_cancelled() { + return Err(Error::Cancelled); + } + // Opening is lazy: validation and construction of a run perform no I/O. + let opening = context.clone(); + let stream = stream::once(async move { + if opening.is_cancelled() { + return Err(Error::Cancelled); + } + self.reader.scan(opening) + }); + use futures::TryStreamExt; + Ok(stream + .try_flatten() + .map(move |batch| { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let batch = batch?; + if batch.schema() != &self.output { + return Err(Error::Invalid( + "connector returned a different Scan schema".into(), + )); + } + if self.predicates.is_empty() { + return Ok(batch); + } + let _workspace = + context.reserve(batch.bytes().checked_mul(2).ok_or(Error::MemoryLimit)?)?; + let mut rows = Vec::new(); + for row in batch.rows() { + if context.is_cancelled() { + return Err(Error::Cancelled); + } + let mut keep = true; + for predicate in &self.predicates { + match predicate.evaluate(row)? { + Value::Bool(true) => {} + Value::Bool(false) | Value::Null => { + keep = false; + break; + } + _ => { + return Err(Error::Invalid("Scan predicate is not boolean".into())) + } + } + } + if keep { + rows.push(row.clone()); + } + } + Batch::try_new(self.output.clone(), rows) + }) + .boxed_local()) + } +} + +mod memory; +pub use memory::MemorySource; diff --git a/crates/asap-physical-operators/src/statistic.rs b/crates/asap-physical-operators/src/statistic.rs new file mode 100644 index 00000000..7053308d --- /dev/null +++ b/crates/asap-physical-operators/src/statistic.rs @@ -0,0 +1,67 @@ +use std::{fmt, str::FromStr}; +use tracing::debug; +#[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(()) + } +} diff --git a/crates/asap-physical-operators/src/stored_state/decoders.rs b/crates/asap-physical-operators/src/stored_state/decoders.rs new file mode 100644 index 00000000..37928f8f --- /dev/null +++ b/crates/asap-physical-operators/src/stored_state/decoders.rs @@ -0,0 +1,366 @@ +//! Shared sketch state reconstruction and decoding. +use asap_sketchlib::CountMinSketch; +use asap_sketchlib::CountMinSketchDelta; +use asap_sketchlib::CountMinSketchWithHeap; +use asap_sketchlib::CountSketch; +use asap_sketchlib::CountSketchDelta; +use asap_sketchlib::CountSketchWithHeap; +use asap_sketchlib::CsHeapItem; +use asap_sketchlib::MessagePackCodec; + +use crate::summary_kernels::count_min_sketch_with_heap::CountMinSketchWithHeapAccumulator; + +/// Decode a `CountMinSketch` from the modified-OTLP wire bytes. +/// MSGPACK path round-trips `CountMinSketch::deserialize_msgpack`; +/// PROTO path decodes a `SketchEnvelope{count_min: CountMinState}` +/// (or bare `CountMinState`) and re-projects to a flat matrix. Mirrors +/// `precompute_operators::count_min_sketch::from_sketchlib_proto_bytes`. +pub fn decode_cms_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::CountMin(st)) => st, + Some(_) => return Err("SketchEnvelope contains non-CountMin sketch".to_string()), + None => { + CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))? + } + }, + Err(_) => { + CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))? + } + }; + let rows = state.rows as usize; + let cols = state.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountMinState has zero dims (rows={rows}, cols={cols})" + )); + } + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type) + .map_err(|_| format!("CountMinState unknown counter_type {}", state.counter_type))?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountMinState counts_int has {} entries, expected {}", + state.counts_int.len(), + expected_len + )); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountMinState counts_float has {} entries, expected {}", + state.counts_float.len(), + expected_len + )); + } + state.counts_float.clone() + } + other => { + return Err(format!( + "CountMinState counter_type {other:?} not yet supported in reducer" + )); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(CountMinSketch::from_legacy_matrix(matrix, rows, cols)) +} + +/// Decode a `CountMinSketch` from msgpack bytes (sketch-core wire +/// format). Mirrors +/// `CountMinSketchAccumulator::from_msgpack_bytes`. +pub fn decode_cms_from_msgpack(buffer: &[u8]) -> Result { + CountMinSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketch msgpack: {e}")) +} + +/// Decode a `CountSketch` from the modified-OTLP proto wire bytes. +/// Mirrors +/// `precompute_operators::count_sketch::from_sketchlib_proto_bytes`. +pub fn decode_cs_from_proto(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountSketchState, CounterType, SketchEnvelope, + }; + use prost::Message; + + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::CountSketch(st)) => st, + Some(_) => return Err("SketchEnvelope contains non-CountSketch sketch".to_string()), + None => CountSketchState::decode(buffer) + .map_err(|e| format!("decode CountSketchState: {e}"))?, + }, + Err(_) => { + CountSketchState::decode(buffer).map_err(|e| format!("decode CountSketchState: {e}"))? + } + }; + let rows = state.rows as usize; + let cols = state.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountSketchState has zero dims (rows={rows}, cols={cols})" + )); + } + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { + format!( + "CountSketchState unknown counter_type {}", + state.counter_type + ) + })?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountSketchState counts_int has {} entries, expected {}", + state.counts_int.len(), + expected_len + )); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountSketchState counts_float has {} entries, expected {}", + state.counts_float.len(), + expected_len + )); + } + state.counts_float.clone() + } + other => { + return Err(format!( + "CountSketchState counter_type {other:?} not yet supported in reducer" + )); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(CountSketch::from_legacy_matrix(matrix, rows, cols)) +} + +/// Decode a `CountSketch` from msgpack bytes (sketch-core wire format). +pub fn decode_cs_from_msgpack(buffer: &[u8]) -> Result { + CountSketch::from_msgpack(buffer).map_err(|e| format!("deserialize CountSketch msgpack: {e}")) +} + +/// Decode a `CountMinSketchWithHeap` from msgpack bytes — the OTLP +/// `CountMinSketch` wire bytes when the gateway/precompute layer +/// marked the sid as CmsWithHeap (heap embedded in the +/// `CountMinSketchWithHeapSerialized` outer wrapper). Delegates to +/// `asap_sketchlib::CountMinSketchWithHeap::deserialize_msgpack`. +pub fn decode_cms_with_heap_from_msgpack(buffer: &[u8]) -> Result { + CountMinSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketchWithHeap msgpack: {e}")) +} + +/// Decode a `CountSketchWithHeap` (median-estimator, Count Sketch family) +/// from msgpack bytes. Distinct wire type from `CountMinSketchWithHeap` +/// (min-estimator, Count-Min family) even though both are heap-bearing +/// frequency sketches — see `asap_sketchlib::CountSketchWithHeap`. +/// Delegates to `asap_sketchlib::CountSketchWithHeap::from_msgpack`. +pub fn decode_cs_with_heap_from_msgpack(buffer: &[u8]) -> Result { + CountSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketchWithHeap msgpack: {e}")) +} + +// --------------------------------------------------------------------------- +// Delta decoders. Under the per-window-reset (PWR) contract +// (`asap-precompute-go/window.go`: a delta is that window's own state +// applied onto a freshly-reset per-series sketch), each stored *Delta +// frame reconstructs into the FULL window state when applied onto an +// EMPTY base of the frame's declared dimensions. The reducer's +// `FrequencyEstimate` / `FrequencyTopk` paths are per-window evaluations, +// so "empty + apply(this window's delta)" yields exactly the window's +// matrix/heap — no cross-window stitching needed (mirrors how the ingest +// accumulators reset_to_empty per window before applying). +// +// The proto path reuses the PUBLIC `asap_sketchlib::{CountSketch, +// CountMinSketch}::apply_delta`; the proto `*Delta` message is decoded via +// `asap_sketchlib::proto::sketchlib::{CountSketchDelta, CountMinDelta}`, +// exactly as `precompute_operators::{count_sketch, +// count_min_sketch}::apply_proto_delta_bytes` does. +// --------------------------------------------------------------------------- + +/// Decode a `CountMinSketch` PROTO_DELTA frame into a FULL sketch by +/// applying the sparse cell delta onto an empty base of the frame's +/// declared dimensions. Mirrors +/// `precompute_operators::count_min_sketch::apply_proto_delta_bytes`. +pub fn decode_cms_from_proto_delta(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountMinDelta packed-array length mismatch: cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + )); + } + let rows = pb.rows as usize; + let cols = pb.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountMinDelta has zero dims (rows={rows}, cols={cols})" + )); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + // hh_keys is parsed off the wire by the precompute accumulator but + // intentionally dropped (the vendored Go proto bindings don't yet + // populate it); match that to keep behavior identical. + let delta = CountMinSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l1: pb.l1, + l2: pb.l2, + hh_keys: Vec::new(), + }; + let mut cms = CountMinSketch::from_legacy_matrix(vec![vec![0.0; cols]; rows], rows, cols); + cms.apply_delta(&delta) + .map_err(|e| format!("apply CountMinDelta onto empty base: {e}"))?; + Ok(cms) +} + +/// Decode a `CountSketch` PROTO_DELTA frame into a FULL sketch by applying +/// the sparse cell delta onto an empty base of the frame's declared +/// dimensions. Mirrors +/// `precompute_operators::count_sketch::apply_proto_delta_bytes`. +pub fn decode_cs_from_proto_delta(buffer: &[u8]) -> Result { + use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountSketchDelta packed-array length mismatch: cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + )); + } + let rows = pb.rows as usize; + let cols = pb.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountSketchDelta has zero dims (rows={rows}, cols={cols})" + )); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + let delta = CountSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l2: pb.l2, + hh_keys: Vec::new(), + }; + let mut cs = CountSketch::from_legacy_matrix(vec![vec![0.0; cols]; rows], rows, cols); + cs.apply_delta(&delta) + .map_err(|e| format!("apply CountSketchDelta onto empty base: {e}"))?; + Ok(cs) +} + +/// Decode a heap-bearing CountSketch MSGPACK_DELTA frame into a FULL +/// `CountMinSketchWithHeap` by applying the sparse matrix delta + full +/// heap onto an empty base of the frame's declared dimensions. This +/// REUSES the ingest-side delta-heap apply logic +/// (`CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes` → +/// `apply_msgpack_heap_delta_bytes`), which decodes the frame generically +/// with `rmp_serde` — no `asap_sketchlib` delta API is added. +pub fn decode_cms_with_heap_from_msgpack_delta( + buffer: &[u8], +) -> Result { + let acc = CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes(buffer) + .map_err(|e| format!("reconstruct CountMinSketchWithHeap from delta: {e}"))?; + Ok(acc.inner) +} + +/// Decode a heap-bearing CountSketch (median-estimator) MSGPACK_DELTA frame +/// into a FULL `asap_sketchlib::CountSketchWithHeap` by applying the sparse +/// matrix delta + full heap onto an empty base of the frame's declared +/// dimensions. Same DELTA-HEAP wire shape as the CmsWithHeap delta frame +/// (see `HeapDeltaWire`/`MatrixDeltaWire` in +/// `count_min_sketch_with_heap.rs`), decoded here directly +/// with `rmp_serde` since there is no CountSketchWithHeap ingest +/// accumulator to delegate to. No `asap_sketchlib` delta API needed — the +/// public `from_legacy_matrix` rebuilds both the matrix and heap. +pub fn decode_cs_with_heap_from_msgpack_delta( + buffer: &[u8], +) -> Result { + #[derive(serde::Deserialize)] + struct HeapDeltaWire { + is_delta: bool, + matrix_delta: MatrixDeltaWire, + topk_heap: Vec<(String, f64)>, + heap_size: u64, + } + #[derive(serde::Deserialize)] + struct MatrixDeltaWire { + rows: u32, + cols: u32, + cells: Vec<(u32, u32, i64)>, + } + + let wire: HeapDeltaWire = rmp_serde::from_slice(buffer) + .map_err(|e| format!("decode CountSketchWithHeap delta msgpack: {e}"))?; + if !wire.is_delta { + return Err("CountSketchWithHeap delta frame has is_delta=false".to_string()); + } + let rows = wire.matrix_delta.rows as usize; + let cols = wire.matrix_delta.cols as usize; + if rows == 0 || cols == 0 { + return Err(format!( + "CountSketchWithHeap delta frame has zero dims (rows={rows}, cols={cols})" + )); + } + let mut matrix = vec![vec![0.0; cols]; rows]; + for (r, c, dc) in &wire.matrix_delta.cells { + let (r, c) = (*r as usize, *c as usize); + if r >= rows || c >= cols { + continue; + } + matrix[r][c] += *dc as f64; + } + let heap: Vec = wire + .topk_heap + .into_iter() + .map(|(key, value)| CsHeapItem { key, value }) + .collect(); + Ok(CountSketchWithHeap::from_legacy_matrix( + matrix, + heap, + rows, + cols, + wire.heap_size as usize, + )) +} diff --git a/crates/asap-physical-operators/src/stored_state/delta_apply.rs b/crates/asap-physical-operators/src/stored_state/delta_apply.rs new file mode 100644 index 00000000..3f638503 --- /dev/null +++ b/crates/asap-physical-operators/src/stored_state/delta_apply.rs @@ -0,0 +1,1276 @@ +//! Shared sketch state reconstruction and decoding. +use asap_sketchlib::CountMinSketch; +use asap_sketchlib::CountMinSketchWithHeap; +use asap_sketchlib::CountSketch; +use asap_sketchlib::CountSketchWithHeap; +use asap_sketchlib::DdSketch; +use asap_sketchlib::HllSketch; +use asap_sketchlib::HllVariant; +use asap_sketchlib::KllSketch; +use asap_sketchlib::MessagePackCodec; + +use super::decoders::{ + decode_cms_from_msgpack, decode_cms_from_proto, decode_cms_from_proto_delta, + decode_cms_with_heap_from_msgpack, decode_cms_with_heap_from_msgpack_delta, + decode_cs_from_msgpack, decode_cs_from_proto, decode_cs_from_proto_delta, + decode_cs_with_heap_from_msgpack, decode_cs_with_heap_from_msgpack_delta, +}; +use super::{SketchEncoding, SketchSampleState}; + +/// Which sketch family a candidate is, and the parameters needed to +/// *bootstrap an empty state* — required by the per-window-reset (PWR) +/// delta model where a window's FIRST frame is a delta-from-empty (no +/// carry-in Full). Most families' deltas embed their own params in the +/// wire fragment (decoded independently, then merged in — see +/// `SummaryState::apply_delta_bytes`); HLL register deltas and DD's +/// bucket-index deltas are applied onto a pre-sized structure instead, +/// so those two need the params known up front to allocate it. +#[derive(Debug, Clone, Copy)] +pub enum DeltaSketchKind { + UnivMon { + heap_size: u32, + sketch_rows: u32, + sketch_cols: u32, + layers: u8, + }, + DDSketch { + alpha: f64, + }, + Hll { + precision: u32, + }, + Kll { + k: u32, + }, + Cms { + rows: usize, + cols: usize, + }, + CountSketch { + rows: usize, + cols: usize, + }, + /// `CmsWithHeap` wraps `asap_sketchlib::CountMinSketchWithHeap` + /// (min-over-rows estimator) and `CountSketchWithHeap` wraps the + /// distinct `asap_sketchlib::CountSketchWithHeap` (median-of-signed-rows + /// estimator) -- different algorithms that happen to share a storage + /// shape. Kept as two variants (not one shared `Heap`) so + /// `merge_same_family` rejects merging one into the other the same + /// way it already rejects e.g. merging a `Cms` into a `Kll`; now the + /// type system enforces it too, since the two variants hold different + /// Rust types. + CmsWithHeap { + rows: usize, + cols: usize, + heap_size: usize, + }, + CountSketchWithHeap { + rows: usize, + cols: usize, + heap_size: usize, + }, +} + +impl DeltaSketchKind { + /// Construct an EMPTY state for this kind, used to seed a new window + /// when its first frame is a delta-from-empty (PWR). A delta applied + /// onto this empty base reconstructs exactly that window's state + /// (delta-from-empty ⊕ empty = window state). + fn bootstrap_empty(&self) -> SummaryState { + match self { + Self::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + } => SummaryState::UnivMon( + crate::summary_kernels::univmon::UnivMonAccumulator::new( + *heap_size as usize, + *sketch_rows as usize, + *sketch_cols as usize, + *layers as usize, + ) + .expect("validated UnivMon catalog dimensions"), + ), + DeltaSketchKind::DDSketch { alpha } => SummaryState::Dd(DdSketch::new(*alpha)), + DeltaSketchKind::Kll { k } => SummaryState::Kll(KllSketch::new(*k as u16)), + DeltaSketchKind::Hll { precision } => { + SummaryState::Hll(HllSketch::new(HllVariant::Regular, *precision)) + } + DeltaSketchKind::Cms { rows, cols } => { + SummaryState::Cms(CountMinSketch::new(*rows, *cols)) + } + DeltaSketchKind::CountSketch { rows, cols } => { + SummaryState::CountSketch(CountSketch::new(*rows, *cols)) + } + DeltaSketchKind::CmsWithHeap { + rows, + cols, + heap_size, + } => SummaryState::CmsWithHeap(CountMinSketchWithHeap::new(*rows, *cols, *heap_size)), + DeltaSketchKind::CountSketchWithHeap { + rows, + cols, + heap_size, + } => SummaryState::CountSketchWithHeap(CountSketchWithHeap::new( + *rows, *cols, *heap_size, + )), + } + } +} + +/// Try to decode a "full" sketch from the bytes (used by both +/// per-window and cumulative modes when the encoding is `*Full`). +fn decode_full( + kind: &DeltaSketchKind, + bytes: &[u8], + encoding: SketchEncoding, +) -> Result { + match (kind, encoding) { + ( + DeltaSketchKind::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + }, + SketchEncoding::MsgpackFull, + ) => { + let state = crate::summary_kernels::univmon::UnivMonAccumulator::from_bytes(bytes) + .map_err(|e| e.to_string())?; + if state.dimensions() + != ( + *heap_size as usize, + *sketch_rows as usize, + *sketch_cols as usize, + *layers as usize, + ) + { + return Err("UnivMon payload dimensions differ from installed catalog".into()); + } + Ok(SummaryState::UnivMon(state)) + } + (DeltaSketchKind::DDSketch { .. }, SketchEncoding::ProtoFull) => { + let sk = dd_from_proto(bytes)?; + Ok(SummaryState::Dd(sk)) + } + (DeltaSketchKind::DDSketch { .. }, SketchEncoding::MsgpackFull) => { + let sk = DdSketch::from_msgpack(bytes) + .map_err(|e| format!("deserialize DDSketch msgpack: {e}"))?; + Ok(SummaryState::Dd(sk)) + } + (DeltaSketchKind::Hll { .. }, SketchEncoding::ProtoFull) => { + let sk = hll_from_proto(bytes)?; + Ok(SummaryState::Hll(sk)) + } + (DeltaSketchKind::Hll { .. }, SketchEncoding::MsgpackFull) => { + let sk = HllSketch::from_msgpack(bytes) + .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?; + Ok(SummaryState::Hll(sk)) + } + (DeltaSketchKind::Kll { .. }, SketchEncoding::ProtoFull) => { + let sk = kll_from_proto(bytes)?; + Ok(SummaryState::Kll(sk)) + } + (DeltaSketchKind::Kll { .. }, SketchEncoding::MsgpackFull) => { + let sk = KllSketch::from_msgpack(bytes) + .map_err(|e| format!("deserialize KllSketch msgpack: {e}"))?; + Ok(SummaryState::Kll(sk)) + } + (DeltaSketchKind::Cms { .. }, SketchEncoding::ProtoFull) => { + Ok(SummaryState::Cms(decode_cms_from_proto(bytes)?)) + } + (DeltaSketchKind::Cms { .. }, SketchEncoding::MsgpackFull) => { + Ok(SummaryState::Cms(decode_cms_from_msgpack(bytes)?)) + } + (DeltaSketchKind::CountSketch { .. }, SketchEncoding::ProtoFull) => { + Ok(SummaryState::CountSketch(decode_cs_from_proto(bytes)?)) + } + (DeltaSketchKind::CountSketch { .. }, SketchEncoding::MsgpackFull) => { + Ok(SummaryState::CountSketch(decode_cs_from_msgpack(bytes)?)) + } + // The heap-bearing wire format is msgpack-only in this + // deployment; `decode_cms_with_heap_from_msgpack` is the same + // "Full" decoder the reducer's existing per-frame dispatch falls + // through to for any non-MsgpackDelta encoding. + ( + DeltaSketchKind::CmsWithHeap { .. }, + SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull, + ) => Ok(SummaryState::CmsWithHeap( + decode_cms_with_heap_from_msgpack(bytes)?, + )), + ( + DeltaSketchKind::CountSketchWithHeap { .. }, + SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull, + ) => Ok(SummaryState::CountSketchWithHeap( + decode_cs_with_heap_from_msgpack(bytes)?, + )), + (_, e) => Err(format!("decode_full called with non-Full encoding {e:?}")), + } +} + +/// The reconstructed state one candidate sid contributes — either +/// folded across a window (or several) via delta application, or merged +/// in from another sid's own reconstruction. +pub enum SummaryState { + UnivMon(crate::summary_kernels::univmon::UnivMonAccumulator), + Dd(DdSketch), + Hll(HllSketch), + Kll(KllSketch), + Cms(CountMinSketch), + CountSketch(CountSketch), + /// See `DeltaSketchKind::CmsWithHeap`/`CountSketchWithHeap` for why + /// these are two variants holding two different sketchlib types. + CmsWithHeap(CountMinSketchWithHeap), + CountSketchWithHeap(CountSketchWithHeap), +} + +impl SummaryState { + /// Apply a delta-encoded payload from a window sample. For DD / KLL, + /// the delta is interpreted as a "mergeable fragment" decoded + /// through the same full-state decoder and merged into the + /// rolling state. For HLL, the wire delta is a sparse register + /// update applied via the sketch's `apply_delta`. + /// + /// On encoding mismatch (e.g. trying to apply an HllDelta to a + /// DDSketch rolling state) returns Err. + pub fn apply_delta_bytes( + &mut self, + bytes: &[u8], + encoding: SketchEncoding, + ) -> Result<(), String> { + if !matches!( + encoding, + SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta + ) { + return Err(format!( + "apply_delta_bytes called with non-Delta encoding {encoding:?}" + )); + } + match self { + SummaryState::UnivMon(_) => Err("UnivMon requires full pane snapshots".into()), + SummaryState::Dd(sk) => { + match encoding { + // PROTO_DELTA: dispatch on the payload SHAPE, mirroring the + // supported DDSketch frame decoder, which tries the + // full-envelope decode first, then falls back to + // the bucket-delta proto. Two wire shapes can arrive on the + // ProtoDelta channel: + // + // 1. `SketchEnvelope{DdSketchState}` — a full-state + // fragment, mergeable via `DdSketch::merge`. (The edge + // sends this when `compute_delta_against` hits the + // empty-current / undecodable-prior fallback and ships + // a full snapshot tagged as a delta.) + // 2. `DDSketchDelta { buckets: [{index, d_count}] }` — a + // bucket-index delta proto, applied additively. This is + // the COMMON delta_transmission frame the edge emits + // under per-window-reset (`compute_delta(&empty)`). + // + // Before this fix the reducer decoded ONLY shape (1) via + // `decode_full`. A real shape-(2) frame failed with a wire- + // type mismatch on field 1 (delta field 1 = repeated + // submessage; state field 1 = `double alpha`) → the whole + // `quantile_over_time` returned `No result` for every + // delta_transmission DDSketch stream. We wrap the rolling + // `DdSketch` in a transient accumulator so the bucket-delta + // apply lands on `sk` in place. + SketchEncoding::ProtoDelta => { + // Shape (1): full envelope fragment → merge. Try this + // first (cheap decode attempt; a bucket-delta proto + // fails it on the field-1 wire-type mismatch). + if let Ok(SummaryState::Dd(other)) = decode_full( + &DeltaSketchKind::DDSketch { alpha: 0.0 }, + bytes, + SketchEncoding::ProtoFull, + ) { + sk.merge(&other) + .map_err(|e| format!("merge DDSketch delta envelope: {e}"))?; + return Ok(()); + } + // Shape (2): bucket-delta proto → additive apply via the + // SAME decoder the ingest delta path uses. + use crate::summary_kernels::dd_sketch::DDSketchAccumulator; + let mut acc = DDSketchAccumulator { + inner: std::mem::replace(sk, DdSketch::new(sk.alpha)), + sample_p: 1.0, + }; + let res = acc.apply_proto_delta_bytes(bytes); + *sk = acc.inner; + res.map_err(|e| format!("apply DDSketch proto bucket-delta: {e}"))?; + Ok(()) + } + // MSGPACK_DELTA: a serialized full-sketch fragment, mergeable + // via the full-state decoder. Kept for completeness — the + // edge wires PROTO_DELTA for DDSketch today. + SketchEncoding::MsgpackDelta => { + let other = match decode_full( + &DeltaSketchKind::DDSketch { alpha: 0.0 }, + bytes, + SketchEncoding::MsgpackFull, + ) { + Ok(SummaryState::Dd(s)) => s, + Ok(_) => { + return Err( + "decode_full(DDSketch) returned non-DDSketch state".to_string() + ) + } + Err(e) => return Err(e), + }; + sk.merge(&other) + .map_err(|e| format!("merge DDSketch delta: {e}"))?; + Ok(()) + } + _ => unreachable!(), + } + } + SummaryState::Hll(sk) => { + // HLL has a true sparse register delta in the proto + // wire format. Use the same path the precompute + // accumulator uses (`apply_proto_delta_bytes`-style). + if encoding == SketchEncoding::ProtoDelta { + apply_hll_proto_delta(sk, bytes) + } else { + // MsgpackDelta for HLL isn't a sparse encoding; + // it's a serialized HllSketch fragment, mergeable + // via `HllSketch::merge`. + let other = HllSketch::from_msgpack(bytes) + .map_err(|e| format!("deserialize HllSketch (delta-as-msgpack): {e}"))?; + sk.merge(&other) + .map_err(|e| format!("merge HLL delta: {e}"))?; + Ok(()) + } + } + SummaryState::Kll(sk) => { + let full_enc = match encoding { + SketchEncoding::ProtoDelta => SketchEncoding::ProtoFull, + SketchEncoding::MsgpackDelta => SketchEncoding::MsgpackFull, + _ => unreachable!(), + }; + let other = match decode_full(&DeltaSketchKind::Kll { k: 0 }, bytes, full_enc) { + Ok(SummaryState::Kll(s)) => s, + Ok(_) => return Err("decode_full(Kll) returned non-Kll state".to_string()), + Err(e) => return Err(e), + }; + sk.merge(&other) + .map_err(|e| format!("merge KLL delta: {e}"))?; + Ok(()) + } + // CMS/CountSketch/Heap have no true sparse in-place delta + // (unlike DD's bucket-index proto or HLL's register proto, + // above) — every delta frame already decodes into a + // complete, standalone state on its own (the PWR wire + // contract resets to empty at the source), so applying one + // is always "decode independently, then merge". + SummaryState::Cms(sk) => { + if encoding != SketchEncoding::ProtoDelta { + return Err( + "CountMin (heap-less) MSGPACK_DELTA is not a valid producer encoding \ + (msgpack-delta is the heap-bearing form)" + .to_string(), + ); + } + let other = decode_cms_from_proto_delta(bytes)?; + sk.merge(&other) + .map_err(|e| format!("merge CountMinSketch delta: {e}")) + } + SummaryState::CountSketch(sk) => { + if encoding != SketchEncoding::ProtoDelta { + return Err( + "CountSketch (heap-less) MSGPACK_DELTA is not a valid producer encoding \ + (msgpack-delta is the heap-bearing form)" + .to_string(), + ); + } + let other = decode_cs_from_proto_delta(bytes)?; + sk.merge(&other) + .map_err(|e| format!("merge CountSketch delta: {e}")) + } + SummaryState::CmsWithHeap(sk) => { + // Matches the existing per-frame reducer dispatch: only + // MsgpackDelta gets true delta treatment; ProtoDelta (not + // produced for this family in this deployment) falls + // through to the full-msgpack decoder, same as `decode_full`. + let other = if encoding == SketchEncoding::MsgpackDelta { + decode_cms_with_heap_from_msgpack_delta(bytes)? + } else { + decode_cms_with_heap_from_msgpack(bytes)? + }; + sk.merge(&other) + .map_err(|e| format!("merge CmsWithHeap delta: {e}")) + } + SummaryState::CountSketchWithHeap(sk) => { + let other = if encoding == SketchEncoding::MsgpackDelta { + decode_cs_with_heap_from_msgpack_delta(bytes)? + } else { + decode_cs_with_heap_from_msgpack(bytes)? + }; + sk.merge(&other) + .map_err(|e| format!("merge CountSketchWithHeap delta: {e}")) + } + } + } + + pub fn quantile(&self, q: f64) -> f64 { + match self { + SummaryState::Dd(sk) => sk.quantile(q).unwrap_or(0.0), + SummaryState::Kll(sk) => sk.quantile(q), + _ => 0.0, + } + } + + pub fn cardinality(&self) -> f64 { + match self { + SummaryState::Hll(sk) => sk.estimate(), + _ => 0.0, + } + } + + /// The bucket TOTAL — sum of row 0 of the underlying matrix. What a + /// bare `count_over_time`/`sum by (item) (rate(...))`-shaped query + /// (no specific item key) reads out. `0.0` for non-Frequency-family + /// states. + pub fn total(&self) -> f64 { + let matrix = match self { + SummaryState::Cms(c) => c.sketch(), + SummaryState::CountSketch(c) => c.sketch().clone(), + SummaryState::CmsWithHeap(h) => h.sketch_matrix(), + SummaryState::CountSketchWithHeap(h) => h.sketch_matrix(), + _ => return 0.0, + }; + matrix + .first() + .map(|row| row.iter().copied().sum::()) + .unwrap_or(0.0) + } + + /// Per-key point estimate — `count(metric{item="x"})`-shaped queries. + /// Unlike [`Self::topk_items`], no heap is needed: all four Frequency + /// variants (heap-bearing or not) already carry a keyed `estimate` + /// over their matrix. `None` for the quantile/cardinality states, + /// which have no item universe at all. + pub fn estimate(&self, key: &str) -> Option { + match self { + SummaryState::Cms(c) => Some(c.estimate(key)), + SummaryState::CountSketch(c) => Some(c.estimate(key)), + SummaryState::CmsWithHeap(h) => Some(h.estimate(key)), + SummaryState::CountSketchWithHeap(h) => Some(h.estimate(key)), + _ => None, + } + } + + /// Top-k `(key, value)` pairs from the heap, descending by value. + /// `None` for anything other than a heap-bearing state — the + /// heap-less Frequency states (`Cms`/`CountSketch`) carry no item + /// universe to enumerate, and the quantile/cardinality states have + /// no heap at all. + pub fn topk_items(&self) -> Option> { + match self { + SummaryState::CmsWithHeap(h) => Some( + h.topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(), + ), + SummaryState::CountSketchWithHeap(h) => Some( + h.topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(), + ), + _ => None, + } + } + + /// Merge `other` into `self` in place — both must be the same sketch + /// family. Used to combine several sids' reconstructed states + /// (`cumulative_summary_state`/`per_window_summary_states`) into one + /// cross-sid answer. `CmsWithHeap`/`CountSketchWithHeap` fall through + /// to the catch-all mismatch arm below like any other mixed pair — + /// and since the two variants now hold distinct sketchlib types + /// (`CountMinSketchWithHeap` vs `CountSketchWithHeap`), there is no + /// arm that could accidentally match them together — see their doc + /// on `DeltaSketchKind`. + pub fn merge_same_family(&mut self, other: &SummaryState) -> Result<(), String> { + match (self, other) { + (SummaryState::UnivMon(a), SummaryState::UnivMon(b)) => { + a.merge_in_place(b).map_err(|e| e.to_string()) + } + (SummaryState::Dd(a), SummaryState::Dd(b)) => { + a.merge(b).map_err(|e| format!("merge DDSketch: {e}")) + } + (SummaryState::Hll(a), SummaryState::Hll(b)) => { + a.merge(b).map_err(|e| format!("merge HLL: {e}")) + } + (SummaryState::Kll(a), SummaryState::Kll(b)) => { + a.merge(b).map_err(|e| format!("merge KLL: {e}")) + } + (SummaryState::Cms(a), SummaryState::Cms(b)) => { + a.merge(b).map_err(|e| format!("merge CountMinSketch: {e}")) + } + (SummaryState::CountSketch(a), SummaryState::CountSketch(b)) => { + a.merge(b).map_err(|e| format!("merge CountSketch: {e}")) + } + (SummaryState::CmsWithHeap(a), SummaryState::CmsWithHeap(b)) => { + a.merge(b).map_err(|e| format!("merge CmsWithHeap: {e}")) + } + (SummaryState::CountSketchWithHeap(a), SummaryState::CountSketchWithHeap(b)) => a + .merge(b) + .map_err(|e| format!("merge CountSketchWithHeap: {e}")), + (a, _) => Err(format!( + "SummaryState family mismatch in merge_same_family (self is {})", + a.family_name() + )), + } + } + + /// Diagnostic family name for error messages — not used for dispatch. + fn family_name(&self) -> &'static str { + match self { + SummaryState::UnivMon(_) => "UnivMon", + SummaryState::Dd(_) => "DDSketch", + SummaryState::Hll(_) => "Hll", + SummaryState::Kll(_) => "Kll", + SummaryState::Cms(_) => "Cms", + SummaryState::CountSketch(_) => "CountSketch", + SummaryState::CmsWithHeap(_) => "CmsWithHeap", + SummaryState::CountSketchWithHeap(_) => "CountSketchWithHeap", + } + } +} + +/// Fold every in-range window's frames for ONE series into a single +/// merged `SummaryState` (cumulative over `[t0, t1]`), returning `None` +/// if no Full frame ever landed (every sample was a leading delta). The +/// per-sid building block for a cross-sid answer: reconstruct each +/// candidate sid's state this way, then merge them (`merge_same_family`) +/// before reading out a quantile/cardinality over the combined data. +pub fn cumulative_summary_state( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, +) -> Result, String> { + let mut rolling: Option = None; + visit_window_summary_states(samples, kind, |_, state| { + if let Some(acc) = rolling.as_mut() { + acc.merge_same_family(&state)?; + } else { + rolling = Some(state); + } + Ok(()) + })?; + Ok(rolling) +} + +#[cfg(test)] +/// Walk a sorted-by-window-end slice of samples in time order and +/// produce ONE per-window scalar `(window_end_ms, scalar)`. +/// +/// ## Per-window-reset (PWR) delta model +/// +/// The edge emits frames grouped by window (all frames of one window +/// share the same `window_end` key; the key changes across windows). +/// The edge RESETS its snapshot base at each window boundary, so each +/// window's state is built *from empty*: +/// +/// * Within a window, frames accumulate to the window total. The first +/// frame may be a `Full` (window 1, or a periodic re-snapshot) or a +/// `Delta`-from-empty (windows 2+ under PWR); subsequent frames are +/// `Delta` INCREMENTS applied onto the window's running base. +/// * Across windows, the base MUST reset — a new `window_end` discards +/// the previous window's rolling state and starts from empty. Never +/// carry one window's state into the next (that would inflate via +/// cross-window accumulation). +/// +/// Concretely this fixes two bugs in the old "single rolling Option that +/// only ever resets on a Full" walk: +/// 1. A query range whose Full lives only in window 1 (or out of +/// range) left windows 2+ as deltas with `rolling=None`, all +/// skipped → empty result. +/// 2. A window 2+ delta applied onto window 1's leftover rolling state +/// → cross-window inflation. +/// +/// For a `Delta` that is the window's FIRST frame (the PWR delta-from- +/// empty case), we bootstrap an EMPTY rolling state of `kind` and apply +/// the delta onto it (delta-from-empty ⊕ empty = that window's state). +/// +/// The delta-OFF path (exactly one `Full` per window) still produces one +/// correct value per window: the window opens with a Full, has no +/// further frames, and emits that Full's scalar. +/// +/// `eval` reads a scalar from the rolling state (`quantile(q)` / +/// `cardinality()`). `skipped` counts frames that could not contribute +/// (a delta we genuinely couldn't bootstrap from — should be rare). +/// +/// Returns `Ok(per_window_samples, skipped)`. +pub fn per_window_evaluate( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, + eval: E, +) -> Result<(Vec<(i64, f64)>, usize), String> +where + E: Fn(&SummaryState) -> f64, +{ + let (states, skipped) = per_window_summary_states(samples, kind)?; + Ok(( + states.into_iter().map(|(w, rs)| (w, eval(&rs))).collect(), + skipped, + )) +} + +/// Walk a sorted-by-window-end slice of samples in time order and +/// reconstruct ONE sid's per-window `SummaryState` (same per-window-reset +/// walk as [`per_window_evaluate`], generalized to return the +/// reconstructed state itself instead of an already-evaluated scalar). +/// The per-sid building block for cross-sid per-window merging (unlike +/// [`cumulative_summary_state`], which folds a whole `[t0, t1]` range +/// into one answer, this keeps each window separate so a caller can +/// merge same-window states across several sids before evaluating -- +/// needed for a matrix/range-query answer, where each output point is +/// itself a cross-sid merge for that one window). +/// +/// Returns `Ok((per_window_states, skipped))`. +pub fn per_window_summary_states( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, +) -> Result<(Vec<(i64, SummaryState)>, usize), String> { + let mut out: Vec<(i64, SummaryState)> = Vec::new(); + let skipped = visit_window_summary_states(samples, kind, |end, state| { + out.push((end, state)); + Ok(()) + })?; + Ok((out, skipped)) +} + +// Both readout modes must reconstruct the same final pane population. The +// visitor lets cumulative merging stream panes without retaining every state. +fn visit_window_summary_states( + samples: &[(i64, &SketchSampleState)], + kind: DeltaSketchKind, + mut emit: impl FnMut(i64, SummaryState) -> Result<(), String>, +) -> Result { + let mut skipped = 0usize; + + // Rolling state for the CURRENT window only. Reset to None whenever + // `window_end` changes (a new window establishes its own base from + // empty). `cur_end` tracks which window `rolling` belongs to. + let mut rolling: Option = None; + let mut cur_end: Option = None; + + for (window_end, state) in samples { + // Window boundary: flush the previous window's final accumulated + // state, then reset the base so this window starts from empty. + if cur_end != Some(*window_end) { + if let (Some(prev_end), Some(rs)) = (cur_end, rolling.take()) { + emit(prev_end, rs)?; + } + cur_end = Some(*window_end); + } + + match state.encoding { + SketchEncoding::NativeBatchV1 => { + return Err("native physical outputs require the bound native batch decoder".into()) + } + SketchEncoding::ProtoFull | SketchEncoding::MsgpackFull => { + // A Full (re)sets this window's base. + rolling = Some(decode_full(&kind, &state.bytes, state.encoding)?); + } + SketchEncoding::ProtoDelta | SketchEncoding::MsgpackDelta => { + // Apply onto this window's running base. If this is the + // window's first frame (PWR delta-from-empty), bootstrap + // an empty base and apply onto it. + if rolling.is_none() { + rolling = Some(kind.bootstrap_empty()); + } + match rolling.as_mut() { + Some(rs) => rs.apply_delta_bytes(&state.bytes, state.encoding)?, + None => skipped += 1, + } + } + } + } + + // Flush the final window. + if let (Some(prev_end), Some(rs)) = (cur_end, rolling.take()) { + emit(prev_end, rs)?; + } + + Ok(skipped) +} + +// --------------------------------------------------------------------------- +// Proto-envelope decoders — P2-4: ONE decoder per family. +// +// These delegate to the precompute-side accumulators' +// `from_sketchlib_proto_bytes`, which are the single source of truth for +// the modified-OTLP proto wire format (envelope unwrapping, alpha/k/ +// precision validation, and — critically for HLL — SPARSE +// `registers_sparse` expansion). Folding the warm read path onto the +// same decoder the ingest path uses means the sparse-register fix (and +// any future format change) can never drift between the two copies again +// — the bug class P2-3 / P2-4 closed. We extract the accumulator's +// public `inner` sketch for the rolling-state merge. +// --------------------------------------------------------------------------- + +fn dd_from_proto(buffer: &[u8]) -> Result { + use crate::summary_kernels::dd_sketch::DDSketchAccumulator; + DDSketchAccumulator::from_sketchlib_proto_bytes(buffer) + .map(|acc| acc.inner) + .map_err(|e| e.to_string()) +} + +fn kll_from_proto(buffer: &[u8]) -> Result { + use crate::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(buffer) + .map(|acc| acc.inner) + .map_err(|e| e.to_string()) +} + +fn hll_from_proto(buffer: &[u8]) -> Result { + use crate::summary_kernels::hll_sketch::HllSketchAccumulator; + HllSketchAccumulator::from_sketchlib_proto_bytes(buffer) + .map(|acc| acc.inner) + .map_err(|e| e.to_string()) +} + +/// Apply a proto-encoded `HllDelta` frame onto the HLL register vector — the +/// delta is a varint-packed (index_delta, value) blob; decode + apply +/// (register-wise max) via the shared sketch library so the unpacking stays a +/// single source of truth. +fn apply_hll_proto_delta(sk: &mut HllSketch, buffer: &[u8]) -> Result<(), String> { + sk.apply_delta_bytes(buffer) + .map_err(|e| format!("apply HLLDelta: {e}"))?; + Ok(()) +} + +#[cfg(test)] +mod tests { + //! P2-3 / P2-4 regression tests for the consolidated single-decoder + //! path. These exercise the family proto decoders that now delegate + //! to the precompute accumulators (the single source of truth), so a + //! divergence between the warm read path and the ingest path — + //! notably the SPARSE-register HLL handling the deleted dead decoder + //! got wrong — fails the build. + use super::*; + use asap_sketchlib::HllVariant; + + fn encode_dd(sk: &DdSketch) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + let state = DdSketchState { + alpha: sk.alpha, + store_counts: sk.store_counts.clone(), + store_offset: sk.store_offset, + ..Default::default() + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), + ..Default::default() + } + .encode_to_vec() + } + + fn encode_kll(k: u16, items: &[f64]) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + let state = KllState { + k: k as u32, + items: items.to_vec(), + levels: vec![], + num_levels: 0, + ..Default::default() + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + } + .encode_to_vec() + } + + fn encode_hll_dense(sk: &HllSketch) -> Vec { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision: sk.precision, + registers: sk.registers.clone(), + hip_kxq0: sk.hip_kxq0, + hip_kxq1: sk.hip_kxq1, + hip_est: sk.hip_est, + registers_sparse: None, + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + } + .encode_to_vec() + } + + /// Build a SPARSE HLL proto frame: dense `registers` left empty, + /// `registers_sparse.packed` = varint (index_delta, value) pairs. + /// This is exactly the wire form a low-cardinality producer emits + /// (sketchlib-go below its dense/sparse crossover) — the frame the + /// DELETED `HllSketch_from_sketchlib_proto_bytes` hard-rejected with + /// "registers has 0 bytes". + fn encode_hll_sparse(precision: u32, nonzero: &[(u64, u8)]) -> Vec { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllSparseRegisters, HllVariant as ProtoVariant, HyperLogLogState, + SketchEnvelope, + }; + use prost::Message; + // Varint-pack (index_delta, value), ascending index order. + let mut packed: Vec = Vec::new(); + let mut prev: u64 = 0; + let put_uvarint = |buf: &mut Vec, mut v: u64| loop { + let b = (v & 0x7f) as u8; + v >>= 7; + if v != 0 { + buf.push(b | 0x80); + } else { + buf.push(b); + break; + } + }; + let mut sorted = nonzero.to_vec(); + sorted.sort_by_key(|(i, _)| *i); + for (idx, val) in &sorted { + put_uvarint(&mut packed, idx - prev); + put_uvarint(&mut packed, *val as u64); + prev = *idx; + } + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision, + registers: Vec::new(), // dense field empty → sparse path + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + // `num_registers` is informational — the decoder expands + // against `expected_len` from precision, not this field. + registers_sparse: Some(HllSparseRegisters { + num_registers: 1u32 << precision, + packed, + }), + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + } + .encode_to_vec() + } + + #[test] + fn hll_from_proto_accepts_sparse_frame() { + // The consolidated decoder must accept the sparse wire form (the + // deleted dead decoder rejected it). Build a sparse frame setting + // a handful of registers, decode it, and confirm those register + // slots came back set in the dense array. + let precision = 12u32; + let nonzero = [(3u64, 5u8), (100, 2), (4000, 7)]; + let bytes = encode_hll_sparse(precision, &nonzero); + let sk = hll_from_proto(&bytes).expect("sparse HLL frame must decode (P2-3 regression)"); + assert_eq!(sk.registers.len(), 1usize << precision); + for (idx, val) in nonzero { + assert_eq!( + sk.registers[idx as usize], val, + "sparse register {idx} expanded to wrong value" + ); + } + } + + #[test] + fn hll_from_proto_matches_accumulator_decoder() { + // P2-4: the warm read path and the ingest accumulator must decode + // the SAME bytes to the SAME sketch (one source of truth). + use crate::summary_kernels::hll_sketch::HllSketchAccumulator; + let mut sk = HllSketch::new(HllVariant::Regular, 12); + for i in 0..500u64 { + sk.update(format!("item-{i}").as_bytes()); + } + let bytes = encode_hll_dense(&sk); + let via_delta = hll_from_proto(&bytes).expect("delta_apply hll decode"); + let via_acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("accumulator hll decode") + .inner; + assert_eq!( + via_delta.registers, via_acc.registers, + "delta_apply and accumulator must produce identical HLL registers" + ); + assert!((via_delta.estimate() - via_acc.estimate()).abs() < 1e-9); + } + + #[test] + fn dd_from_proto_matches_accumulator_decoder() { + use crate::summary_kernels::dd_sketch::DDSketchAccumulator; + let mut sk = DdSketch::new(0.01); + for v in [1.0, 2.0, 5.0, 5.0, 9.0, 42.0] { + sk.update(v); + } + let bytes = encode_dd(&sk); + let via_delta = dd_from_proto(&bytes).expect("delta_apply dd decode"); + let via_acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("accumulator dd decode") + .inner; + // Same quantile answers from the same bytes through both paths. + assert_eq!(via_delta.quantile(0.5), via_acc.quantile(0.5)); + assert_eq!(via_delta.quantile(0.99), via_acc.quantile(0.99)); + } + + #[test] + fn kll_from_proto_matches_accumulator_decoder() { + use crate::summary_kernels::datasketches_kll::DatasketchesKLLAccumulator; + let items: Vec = (0..200).map(|i| i as f64).collect(); + let bytes = encode_kll(256, &items); + let via_delta = kll_from_proto(&bytes).expect("delta_apply kll decode"); + let via_acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("accumulator kll decode") + .inner; + assert_eq!(via_delta.quantile(0.5), via_acc.quantile(0.5)); + } + + // ----------------------------------------------------------------- + // Per-window-reset (PWR) delta-apply regression tests. + // + // The edge resets its snapshot base at every window boundary, so a + // window's first frame is either a Full (window 1 / re-snapshot) or + // a Delta-from-empty (windows 2+). The query-side walk must: + // * reset the rolling base when `window_end` changes, + // * bootstrap an empty base for a window's leading Delta, + // * emit ONE value per window (the window's final accumulated + // state), never per-frame and never cross-window-accumulated. + // ----------------------------------------------------------------- + + fn full(bytes: Vec) -> SketchSampleState { + SketchSampleState { + bytes, + encoding: SketchEncoding::ProtoFull, + } + } + fn delta(bytes: Vec) -> SketchSampleState { + SketchSampleState { + bytes, + encoding: SketchEncoding::ProtoDelta, + } + } + + fn dd_over(alpha: f64, vals: &[f64]) -> DdSketch { + let mut sk = DdSketch::new(alpha); + for &v in vals { + sk.update(v); + } + sk + } + + /// A full re-snapshot replaces its pane's earlier frames; cumulative + /// readout must merge the finalized panes without counting updates twice. + #[test] + fn cumulative_readout_counts_resnapshot_population_once() { + let first = full(encode_dd(&dd_over(0.01, &[1., 2.]))); + let updated = full(encode_dd(&dd_over(0.01, &[1., 2., 3.]))); + let next = delta(encode_dd(&dd_over(0.01, &[9.]))); + let samples = [(1000, &first), (1000, &updated), (2000, &next)]; + let state = cumulative_summary_state(&samples, DeltaSketchKind::DDSketch { alpha: 0.01 }) + .unwrap() + .unwrap(); + let SummaryState::Dd(state) = state else { + panic!("expected DDSketch state"); + }; + assert_eq!(state.store_counts.iter().sum::(), 4); + } + + /// PWR across 3 windows: window 1 is `[Full]`, windows 2 & 3 are + /// `[Delta-from-empty]` (NO Full carry-in). Each window must + /// reconstruct its OWN distribution's median — not empty (the old + /// "skip delta with no base" bug) and not cross-window-inflated. + #[test] + fn pwr_ddsketch_three_windows_delta_from_empty() { + let alpha = 0.01; + let w1 = dd_over(alpha, &[1.0, 2.0, 3.0, 4.0, 5.0]); + let w2 = dd_over(alpha, &[10.0, 20.0, 30.0, 40.0, 50.0]); + let w3 = dd_over(alpha, &[100.0, 200.0, 300.0, 400.0, 500.0]); + + // window 1 ships a Full; windows 2+ ship a delta-from-empty. + let s1 = full(encode_dd(&w1)); + let s2 = delta(encode_dd(&w2)); + let s3 = delta(encode_dd(&w3)); + let samples = vec![(1000_i64, &s1), (2000, &s2), (3000, &s3)]; + + let kind = DeltaSketchKind::DDSketch { alpha }; + let (out, skipped) = + per_window_evaluate(&samples, kind, |rs| rs.quantile(0.5)).expect("pwr eval"); + assert_eq!(skipped, 0, "PWR must not skip delta-from-empty frames"); + assert_eq!(out.len(), 3, "one value per window"); + + // Each window's median ≈ that window's own distribution median, + // independent of the others (no carry-in inflation). + let truth = [ + w1.quantile(0.5).unwrap(), + w2.quantile(0.5).unwrap(), + w3.quantile(0.5).unwrap(), + ]; + for (i, (w_end, est)) in out.iter().enumerate() { + assert_eq!(*w_end, (i as i64 + 1) * 1000); + let rel = (est - truth[i]).abs() / truth[i].max(1e-9); + assert!( + rel < 0.05, + "window {i}: est={est} truth={} rel={rel}", + truth[i] + ); + } + // Cross-window-inflation guard: window 2's median must NOT have + // absorbed window 1 (would pull it well below 30). + assert!( + out[1].1 > 20.0, + "window 2 median {} suggests cross-window accumulation", + out[1].1 + ); + } + + /// Sub-window producer: a SINGLE window carries multiple frames + /// `[Full, Delta, Delta]`, where each later delta is an increment + /// since the previous emit in that window. The walk must COLLAPSE + /// them to ONE value = the window's running total, not emit three. + #[test] + fn pwr_ddsketch_subwindow_frames_collapse_to_window_total() { + let alpha = 0.01; + // Three sub-window increments that together cover 1..=15. + let a = dd_over(alpha, &[1.0, 2.0, 3.0, 4.0, 5.0]); + let b = dd_over(alpha, &[6.0, 7.0, 8.0, 9.0, 10.0]); + let c = dd_over(alpha, &[11.0, 12.0, 13.0, 14.0, 15.0]); + let s_a = full(encode_dd(&a)); + let s_b = delta(encode_dd(&b)); + let s_c = delta(encode_dd(&c)); + // All three share the same window_end (one window, sub-window frames). + let samples = vec![(5000_i64, &s_a), (5000, &s_b), (5000, &s_c)]; + + let kind = DeltaSketchKind::DDSketch { alpha }; + let (out, skipped) = + per_window_evaluate(&samples, kind, |rs| rs.quantile(0.5)).expect("subwindow eval"); + assert_eq!(skipped, 0); + assert_eq!(out.len(), 1, "sub-window frames collapse to ONE value"); + assert_eq!(out[0].0, 5000); + + let truth = dd_over(alpha, &(1..=15).map(|v| v as f64).collect::>()) + .quantile(0.5) + .unwrap(); + let rel = (out[0].1 - truth).abs() / truth.max(1e-9); + assert!(rel < 0.05, "window total est={} truth={truth}", out[0].1); + } + + /// Same sub-window collapse, but the window's FIRST frame is a + /// Delta-from-empty (PWR window 2+ with sub-window frames): + /// `[Delta-from-empty, Delta, Delta]`. + #[test] + fn pwr_ddsketch_subwindow_first_frame_delta_from_empty() { + let alpha = 0.01; + let a = dd_over(alpha, &[1.0, 2.0, 3.0, 4.0, 5.0]); + let b = dd_over(alpha, &[6.0, 7.0, 8.0, 9.0, 10.0]); + let c = dd_over(alpha, &[11.0, 12.0, 13.0, 14.0, 15.0]); + let s_a = delta(encode_dd(&a)); // first frame is delta-from-empty + let s_b = delta(encode_dd(&b)); + let s_c = delta(encode_dd(&c)); + let samples = vec![(9000_i64, &s_a), (9000, &s_b), (9000, &s_c)]; + + let kind = DeltaSketchKind::DDSketch { alpha }; + let (out, skipped) = + per_window_evaluate(&samples, kind, |rs| rs.quantile(0.5)).expect("eval"); + assert_eq!(skipped, 0); + assert_eq!(out.len(), 1); + let truth = dd_over(alpha, &(1..=15).map(|v| v as f64).collect::>()) + .quantile(0.5) + .unwrap(); + let rel = (out[0].1 - truth).abs() / truth.max(1e-9); + assert!(rel < 0.05, "est={} truth={truth}", out[0].1); + } + + /// PWR for HLL across 3 windows, each a Delta-from-empty (sparse + /// register delta). Bootstrapping an EMPTY HLL of the right precision + /// is required (register deltas index into a pre-sized array). Each + /// window's cardinality must reflect its OWN item set. + #[test] + fn pwr_hll_three_windows_delta_from_empty() { + let precision = 12u32; + // Build per-window HLLs, then encode each as a register-delta + // against an EMPTY sketch (= that window's full register state, + // the PWR delta-from-empty wire form). + let empty = HllSketch::new(HllVariant::Regular, precision); + let mut frames = Vec::new(); + let truths = [200usize, 800, 1500]; + for (w, &n) in truths.iter().enumerate() { + let mut sk = HllSketch::new(HllVariant::Regular, precision); + let base = (w as u64) * 100_000; // disjoint item sets per window + for i in 0..n as u64 { + sk.update(format!("u-{}", base + i).as_bytes()); + } + let bytes = sk.compute_delta(&empty, 0); + frames.push((((w as u64) + 1) * 1000, delta(bytes))); + } + let samples: Vec<(i64, &SketchSampleState)> = + frames.iter().map(|(t, s)| (*t as i64, s)).collect(); + + let kind = DeltaSketchKind::Hll { precision }; + let (out, skipped) = + per_window_evaluate(&samples, kind, |rs| rs.cardinality()).expect("hll pwr eval"); + assert_eq!(skipped, 0, "HLL delta-from-empty must bootstrap, not skip"); + assert_eq!(out.len(), 3); + for (i, (_w_end, est)) in out.iter().enumerate() { + let n = truths[i] as f64; + let rel = (est - n).abs() / n; + assert!( + rel < 0.15, + "window {i}: HLL est={est} truth={n} rel={rel} (each window independent)" + ); + } + } + + /// `CmsWithHeap` (min-over-rows estimator, `CountMinSketchWithHeap`) + /// and `CountSketchWithHeap` (median-of-signed-rows estimator, the + /// distinct `CountSketchWithHeap` type) are different sketch + /// algorithms that merely happen to share a storage shape — merging + /// one into the other must be rejected as a family mismatch, the + /// same as merging a `Cms` into a `Kll` would be. Since the two + /// `SummaryState` variants now hold genuinely different Rust types, + /// this is also enforced at compile time — there is no arm in + /// `merge_same_family` that type-checks a mixed pair together. + #[test] + fn cms_with_heap_and_count_sketch_with_heap_are_not_the_same_family() { + use asap_sketchlib::{CountMinSketchWithHeap, CountSketchWithHeap, MessagePackCodec}; + + let mut cms_heap = CountMinSketchWithHeap::new(4, 256, 10); + cms_heap.update("a", 1.0); + let mut cs_heap = CountSketchWithHeap::new(4, 256, 10); + cs_heap.update("b", 1.0); + + let mut a = SummaryState::CmsWithHeap( + CountMinSketchWithHeap::from_msgpack(&cms_heap.to_msgpack().unwrap()).unwrap(), + ); + let b = SummaryState::CountSketchWithHeap( + CountSketchWithHeap::from_msgpack(&cs_heap.to_msgpack().unwrap()).unwrap(), + ); + + match a.merge_same_family(&b) { + Err(msg) => assert!( + msg.contains("family mismatch"), + "expected a family-mismatch error, got: {msg}" + ), + Ok(()) => panic!( + "CmsWithHeap must not merge with CountSketchWithHeap -- \ + different algorithms sharing only a storage shape" + ), + } + } + + fn encode_delta_heap( + rows: u32, + cols: u32, + cells: &[(u32, u32, i64)], + heap: &[(&str, f64)], + heap_size: u64, + ) -> Vec { + #[derive(serde::Serialize)] + struct W<'a>( + bool, + (u32, u32, &'a [(u32, u32, i64)]), + Vec<(String, f64)>, + u64, + ); + let heap_owned: Vec<(String, f64)> = + heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); + let w = W(true, (rows, cols, cells), heap_owned, heap_size); + rmp_serde::to_vec(&w).expect("encode delta-heap") + } + + /// `SummaryState::CountSketchWithHeap` must decode both FULL and + /// DELTA-HEAP msgpack frames through the genuine + /// `asap_sketchlib::CountSketchWithHeap` (median-of-signed-rows + /// estimator) rather than the CMS-family `CountMinSketchWithHeap` + /// (min-over-rows estimator) it used to alias — the bug this split + /// fixed. Built via real `update()` calls (not a hand-crafted matrix) + /// so the sign-hashed row semantics are genuinely exercised, then + /// checks both decode paths reproduce the same matrix and the same + /// `estimate()` as the in-memory sketch they were encoded from. + #[test] + fn count_sketch_with_heap_full_and_delta_decode_via_new_asap_sketchlib_type() { + use asap_sketchlib::{CountSketchWithHeap, MessagePackCodec}; + + let mut built = CountSketchWithHeap::new(4, 64, 10); + for _ in 0..50 { + built.update("k", 1.0); + } + let expected_matrix = built.sketch_matrix(); + let expected_estimate = built.estimate("k"); + + // FULL path. + let full_bytes = built.to_msgpack().expect("encode full CountSketchWithHeap"); + let full_state = decode_full( + &DeltaSketchKind::CountSketchWithHeap { + rows: 4, + cols: 64, + heap_size: 10, + }, + &full_bytes, + SketchEncoding::MsgpackFull, + ) + .expect("decode_full CountSketchWithHeap"); + match full_state { + SummaryState::CountSketchWithHeap(inner) => { + assert_eq!(inner.sketch_matrix(), expected_matrix); + assert_eq!(inner.estimate("k"), expected_estimate); + } + other => panic!( + "expected CountSketchWithHeap state, got {}", + other.family_name() + ), + } + + // DELTA-HEAP path: same cells + heap against an empty base (PWR + // contract), encoded the way the Go producer does. + let cells: Vec<(u32, u32, i64)> = expected_matrix + .iter() + .enumerate() + .flat_map(|(r, row)| { + row.iter().enumerate().filter_map(move |(c, v)| { + if *v != 0.0 { + Some((r as u32, c as u32, *v as i64)) + } else { + None + } + }) + }) + .collect(); + let heap_pairs: Vec<(String, f64)> = built + .topk_heap_items() + .into_iter() + .map(|item| (item.key, item.value)) + .collect(); + assert!(!heap_pairs.is_empty(), "expected \"k\" in the top-k heap"); + let heap_refs: Vec<(&str, f64)> = + heap_pairs.iter().map(|(k, v)| (k.as_str(), *v)).collect(); + let delta_bytes = encode_delta_heap(4, 64, &cells, &heap_refs, 10); + + let mut rolling = DeltaSketchKind::CountSketchWithHeap { + rows: 4, + cols: 64, + heap_size: 10, + } + .bootstrap_empty(); + rolling + .apply_delta_bytes(&delta_bytes, SketchEncoding::MsgpackDelta) + .expect("apply CountSketchWithHeap delta"); + match rolling { + SummaryState::CountSketchWithHeap(inner) => { + assert_eq!( + inner.sketch_matrix(), + expected_matrix, + "delta path must reconstruct the identical matrix" + ); + assert_eq!(inner.estimate("k"), expected_estimate); + } + other => panic!( + "expected CountSketchWithHeap state, got {}", + other.family_name() + ), + } + } +} diff --git a/crates/asap-physical-operators/src/stored_state/mod.rs b/crates/asap-physical-operators/src/stored_state/mod.rs new file mode 100644 index 00000000..89ab9578 --- /dev/null +++ b/crates/asap-physical-operators/src/stored_state/mod.rs @@ -0,0 +1,22 @@ +//! Portable stored-summary payloads and reconstruction, independent of storage engines. +pub mod decoders; +pub mod delta_apply; +pub mod native; +pub mod readout; + +#[derive(Debug, Clone)] +pub struct SketchSampleState { + pub bytes: Vec, + /// Wire-encoding hint from the OTLP DataPoint's `encoding` field. + pub encoding: SketchEncoding, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SketchEncoding { + ProtoFull, + ProtoDelta, + MsgpackFull, + MsgpackDelta, + /// Versioned typed physical output; never a legacy sketch frame. + NativeBatchV1, +} diff --git a/crates/asap-physical-operators/src/stored_state/native.rs b/crates/asap-physical-operators/src/stored_state/native.rs new file mode 100644 index 00000000..4cf0e41b --- /dev/null +++ b/crates/asap-physical-operators/src/stored_state/native.rs @@ -0,0 +1,368 @@ +//! Versioned physical output batches. Deployment identities and coverage remain +//! outside this payload and must be checked before decoding with the bound schema. +use crate::{ + summary_kernels::{ + datasketches_kll::DatasketchesKLLAccumulator, dd_sketch::DDSketchAccumulator, + exact::ExactAccumulator, hll_sketch::HllSketchAccumulator, + weighted_frequency::WeightedFrequency, SumAccumulator, + }, + values::{Batch, Schema, Value}, + AggregateCore, Error, +}; +use planner_types::post_asap::{SummaryFamilyType, SummarySchema}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; + +#[derive(Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +struct StoredBatch { + version: u32, + schema: SummarySchema, + rows: Vec>, +} + +#[derive(Serialize, Deserialize)] +enum Cell { + Plain(Value), + Summary { + family: SummaryFamilyType, + codec: StateCodec, + bytes: Vec, + }, +} + +// Codec identity is distinct from algorithm identity: old integer CMS/CS bytes +// must never be interpreted as Float64 weighted state with typed item tuples. +#[derive(Serialize, Deserialize)] +enum StateCodec { + WeightedFrequencyV1, + ExactAccumulatorV1, + SumAccumulatorV1, + KllMsgpackV1, + DdMsgpackV1, + HllMsgpackV1, +} +fn invalid(message: impl ToString) -> Error { + Error::Invalid(message.to_string()) +} +impl StateCodec { + fn for_state(state: &dyn AggregateCore) -> Result { + let state = state.as_any(); + if state.is::() { + Ok(Self::WeightedFrequencyV1) + } else if state.is::() { + Ok(Self::ExactAccumulatorV1) + } else if state.is::() { + Ok(Self::SumAccumulatorV1) + } else if state.is::() { + Ok(Self::KllMsgpackV1) + } else if state.is::() { + Ok(Self::DdMsgpackV1) + } else if state.is::() { + Ok(Self::HllMsgpackV1) + } else { + Err(invalid("physical summary has no persisted native codec")) + } + } + fn decode(&self, bytes: &[u8]) -> Result, Error> { + Ok(match self { + Self::WeightedFrequencyV1 => Arc::new(WeightedFrequency::from_bytes(bytes)?), + Self::ExactAccumulatorV1 => { + Arc::new(ExactAccumulator::deserialize_from_bytes(bytes).map_err(invalid)?) + } + Self::SumAccumulatorV1 => { + Arc::new(SumAccumulator::deserialize_from_bytes(bytes).map_err(invalid)?) + } + Self::KllMsgpackV1 => { + Arc::new(DatasketchesKLLAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + } + Self::DdMsgpackV1 => { + Arc::new(DDSketchAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + } + Self::HllMsgpackV1 => { + Arc::new(HllSketchAccumulator::from_msgpack_bytes(bytes).map_err(invalid)?) + } + }) + } +} + +/// Encode a validated physical output, preserving Float64 and typed identities. +/// This format is independent of the logical and physical plan wire formats. +pub fn encode_batch(batch: &Batch) -> Result, Error> { + let rows = batch + .rows() + .iter() + .map(|row| { + row.iter() + .map(|value| { + Ok(match value { + Value::Summary { family, state } => Cell::Summary { + family: family.clone(), + codec: StateCodec::for_state(state.as_ref())?, + bytes: state.serialize_to_bytes(), + }, + value => Cell::Plain(value.clone()), + }) + }) + .collect::, Error>>() + }) + .collect::, Error>>()?; + rmp_serde::to_vec_named(&StoredBatch { + version: 1, + schema: batch.schema().as_ref().clone(), + rows, + }) + .map_err(invalid) +} + +/// Decode only against the installed output contract. The caller supplies its +/// per-read payload limit; checking state parameters is part of Batch validation. +pub fn decode_batch(bytes: &[u8], expected: Schema, max_bytes: usize) -> Result { + if bytes.len() > max_bytes { + return Err(invalid("native output payload exceeds read budget")); + } + let stored: StoredBatch = rmp_serde::from_slice(bytes).map_err(invalid)?; + if stored.version != 1 { + return Err(invalid("unsupported native output format")); + } + if stored.schema != *expected { + return Err(invalid( + "native output schema differs from installed contract", + )); + } + let rows = stored + .rows + .into_iter() + .map(|row| { + row.into_iter() + .map(|cell| { + Ok(match cell { + Cell::Plain(value) => value, + Cell::Summary { + family, + codec, + bytes, + } => Value::Summary { + family, + state: codec.decode(&bytes)?, + }, + }) + }) + .collect::, Error>>() + }) + .collect::, Error>>()?; + Batch::try_new(expected, rows) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::summary_kernels::weighted_frequency::FrequencyAlgorithm; + use planner_types::{ + post_asap::{SketchAlgorithm, SketchKind, SketchParams, SummaryField}, + pre_asap::DataType, + }; + + fn weighted(algorithm: SketchAlgorithm) -> Batch { + let (native, params) = match algorithm { + SketchAlgorithm::CmsWithHeap => ( + FrequencyAlgorithm::Cms, + SketchParams::CmsWithHeap { + width: 64, + depth: 5, + heap_size: 8, + }, + ), + _ => ( + FrequencyAlgorithm::CountSketch, + SketchParams::CountSketchWithHeap { + width: 64, + depth: 5, + heap_size: 8, + }, + ), + }; + let family = + SummaryFamilyType::Sketch(SketchKind::new(algorithm, params), Default::default()); + let schema = Arc::new(SummarySchema { + fields: vec![ + SummaryField { + name: "group".into(), + dtype: SummaryFamilyType::Plain(DataType::Utf8), + nullable: false, + }, + SummaryField { + name: "state".into(), + dtype: family.clone(), + nullable: false, + }, + ], + time_index: None, + }); + let mut state = WeightedFrequency::new(native, 64, 5, 8).unwrap(); + state + .update(&[Value::Int64(7), Value::Utf8("service-a".into())], 0.125) + .unwrap(); + state + .update( + &[Value::Utf8("7".into()), Value::Utf8("service-b".into())], + 0.25, + ) + .unwrap(); + Batch::try_new( + schema, + vec![vec![ + Value::Utf8("job-a".into()), + Value::Summary { + family, + state: Arc::new(state), + }, + ]], + ) + .unwrap() + } + + // Fractional rates and distinct typed item tuples survive both heap codecs. + #[test] + fn weighted_outputs_roundtrip_without_integer_conversion() { + for algorithm in [ + SketchAlgorithm::CmsWithHeap, + SketchAlgorithm::CountSketchWithHeap, + ] { + let batch = weighted(algorithm); + let bytes = encode_batch(&batch).unwrap(); + let restored = decode_batch(&bytes, batch.schema().clone(), bytes.len()).unwrap(); + let scores = |b: &Batch| { + let Value::Summary { state, .. } = &b.rows()[0][1] else { + panic!() + }; + state + .as_any() + .downcast_ref::() + .unwrap() + .rows(8) + .iter() + .map(|row| row.iter().map(|v| v.key().unwrap()).collect::>()) + .collect::>() + }; + assert_eq!(scores(&batch), scores(&restored)); + assert!(decode_batch(&bytes, batch.schema().clone(), bytes.len() - 1).is_err()); + } + } + + // A storage tag cannot send weighted physical output through an integer heap decoder. + #[test] + fn legacy_sketch_reader_rejects_native_batch_frames() { + let sample = super::super::SketchSampleState { + bytes: encode_batch(&weighted(SketchAlgorithm::CmsWithHeap)).unwrap(), + encoding: super::super::SketchEncoding::NativeBatchV1, + }; + let result = super::super::delta_apply::per_window_summary_states( + &[(60_000, &sample)], + super::super::delta_apply::DeltaSketchKind::CmsWithHeap { + rows: 5, + cols: 64, + heap_size: 8, + }, + ); + assert!(result.is_err()); + } + + // Every admitted native summary codec survives the same typed boundary. + #[test] + fn native_summary_families_and_nonfinite_plain_values_roundtrip() { + use planner_types::post_asap::{ExactKind, ExactParams}; + let exact = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + let sketch = |algorithm, params| { + SummaryFamilyType::Sketch(SketchKind::new(algorithm, params), Default::default()) + }; + let cases: Vec<(SummaryFamilyType, Arc)> = vec![ + ( + exact.clone(), + Arc::new(ExactAccumulator::new(exact.clone(), false).unwrap()), + ), + (exact, Arc::new(SumAccumulator::new())), + ( + sketch(SketchAlgorithm::Kll, SketchParams::Kll { k: 200 }), + Arc::new(DatasketchesKLLAccumulator::new(200)), + ), + ( + sketch( + SketchAlgorithm::DDSketch, + SketchParams::DDSketch { alpha: 0.01 }, + ), + Arc::new(DDSketchAccumulator::new(0.01)), + ), + ( + sketch(SketchAlgorithm::Hll, SketchParams::Hll { precision: 12 }), + Arc::new(HllSketchAccumulator::new( + asap_sketchlib::HllVariant::Regular, + 12, + )), + ), + ]; + for (family, state) in cases { + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "state".into(), + dtype: family.clone(), + nullable: false, + }], + time_index: None, + }); + let batch = + Batch::try_new(schema.clone(), vec![vec![Value::Summary { family, state }]]) + .unwrap(); + let bytes = encode_batch(&batch).unwrap(); + let restored = decode_batch(&bytes, schema, bytes.len()).unwrap(); + assert_eq!(encode_batch(&restored).unwrap(), bytes); + } + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }], + time_index: None, + }); + let batch = Batch::try_new( + schema.clone(), + vec![ + vec![Value::Float64(f64::NAN)], + vec![Value::Float64(f64::INFINITY)], + ], + ) + .unwrap(); + let restored = decode_batch(&encode_batch(&batch).unwrap(), schema, usize::MAX).unwrap(); + assert!(matches!(restored.rows()[0][0], Value::Float64(v) if v.is_nan())); + assert!(matches!(restored.rows()[1][0], Value::Float64(v) if v == f64::INFINITY)); + } + + // Recovery validates the format, bound schema and actual sketch parameters. + #[test] + fn corrupt_or_relabelled_output_is_rejected() { + let batch = weighted(SketchAlgorithm::CmsWithHeap); + let bytes = encode_batch(&batch).unwrap(); + let wrong = weighted(SketchAlgorithm::CountSketchWithHeap); + assert!(decode_batch(&bytes, wrong.schema().clone(), usize::MAX).is_err()); + let mut stored: StoredBatch = rmp_serde::from_slice(&bytes).unwrap(); + stored.version = 2; + assert!(decode_batch( + &rmp_serde::to_vec_named(&stored).unwrap(), + batch.schema().clone(), + usize::MAX + ) + .is_err()); + stored.version = 1; + let Cell::Summary { bytes: payload, .. } = &mut stored.rows[0][1] else { + panic!() + }; + *payload = b"legacy integer heap".to_vec(); + assert!(decode_batch( + &rmp_serde::to_vec_named(&stored).unwrap(), + batch.schema().clone(), + usize::MAX + ) + .is_err()); + } +} diff --git a/crates/asap-physical-operators/src/stored_state/readout.rs b/crates/asap-physical-operators/src/stored_state/readout.rs new file mode 100644 index 00000000..df13f33c --- /dev/null +++ b/crates/asap-physical-operators/src/stored_state/readout.rs @@ -0,0 +1,303 @@ +//! Planner-declared readouts over reconstructed summary states. +use super::delta_apply::SummaryState; +use planner_types::{post_asap::SketchQuery, pre_asap::ColumnRef}; +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("{0}")] + Unsupported(&'static str), +} +pub fn sketch_query_value(rs: &SummaryState, query: &SketchQuery) -> Result { + if let SummaryState::UnivMon(state) = rs { + use crate::AggregateCore; + let statistic = match query { + SketchQuery::Cardinality => crate::Statistic::Cardinality, + SketchQuery::FrequencyL2 => crate::Statistic::FrequencyL2, + SketchQuery::FrequencyEntropy => crate::Statistic::FrequencyEntropy, + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + } => crate::Statistic::Count, + _ => return Err(Error::Unsupported("unsupported UnivMon readout")), + }; + return state + .query_statistic(statistic, &None, &Default::default()) + .map_err(|_| Error::Unsupported("UnivMon readout failed")); + } + match query { + SketchQuery::FrequencyL2 | SketchQuery::FrequencyEntropy => Err(Error::Unsupported( + "frequency moment readout requires UnivMon", + )), + SketchQuery::Quantile { q } => match rs { + // Typed PromQL/continuous-percentile readout uses interpolation; + // portable DDS `quantile` deliberately retains lower-rank parity. + SummaryState::Dd(sketch) => sketch.quantile_interpolated(*q).ok_or(Error::Unsupported( + "DDS interpolated quantile is unavailable", + )), + _ => Ok(rs.quantile(*q)), + }, + SketchQuery::Cardinality => Ok(rs.cardinality()), + // `key: ColumnRef::SampleValue, value: None` means "no specific + // item" -- the bare bucket total. `key: Named(_), value: Some(v)` + // is a per-item point lookup (e.g. `count(cms_metric{item="x"})`) + // -- `value` is where the filter's actual value lives (see + // `planner_types::post_asap::SketchQuery::PointCount`'s doc for why `readout` + // can't resolve it itself). Any other combination (e.g. a `Named` + // key with no value, or `SampleValue` with a value) is a shape + // this executor doesn't expect to see and reports rather than + // silently misreading. + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + } => Ok(rs.total()), + SketchQuery::PointCount { + key: ColumnRef::Named(_) | ColumnRef::Qualified { .. }, + value: Some(v), + } => rs.estimate(v).ok_or(Error::Unsupported( + "PointCount by key requires a Frequency-family sketch (Cms/CountSketch/..WithHeap)", + )), + SketchQuery::PointCount { .. } => Err(Error::Unsupported( + "unrecognized PointCount shape (key/value combination not expected)", + )), + // Both readout callers branch on `TopK` before ever calling this + // function (see `readout_cumulative`/`readout_per_window`), so + // this arm is unreachable in practice; kept for match + // exhaustiveness (`SketchQuery` has no `#[non_exhaustive]`) and to + // fail loudly rather than panic if that invariant is ever broken. + SketchQuery::TopK { .. } => Err(Error::Unsupported( + "TopK must be read out via topk_ranked, not sketch_query_value", + )), + } +} + +/// Rank a merged `SummaryState`'s top-k heap items descending by value and +/// cap at the requested `k`. The sort is load-bearing, not defensive +/// polish: `SummaryState::topk_items` reads back a bounded min-heap's +/// backing array as-is (`HHHeap::heap()`, asap_sketchlib) -- it does NOT +/// actually guarantee order despite its own doc wording. Errors for a +/// heap-less family (`Dd`/`Hll`/`Kll`/`Cms`/`CountSketch` -- no item +/// universe to rank), not for an empty heap (a heap-bearing family that +/// simply never received any updates yields `Ok(vec![])`, not an error). +pub fn topk_ranked(rs: &SummaryState, k: usize) -> Result, Error> { + let mut items = rs.topk_items().ok_or(Error::Unsupported( + "TopK requires a heap-bearing family (CmsWithHeap/CountSketchWithHeap) -- \ + this state's family carries no item universe to rank", + ))?; + items.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.0.cmp(&b.0)) // deterministic tie-break for equal counts + }); + items.truncate(k); + Ok(items) +} + +/// Merge already selected exact panes and finalize using the shared accumulator contract. +pub fn exact_readout( + states: impl IntoIterator>, + statistic: crate::Statistic, + key: &Option, + parameters: &std::collections::HashMap, +) -> Result { + let merged = merge_exact_states(states)?; + merged + .query_statistic(statistic, key, parameters) + .map_err(|e| e.to_string()) +} + +fn merge_exact_states( + states: impl IntoIterator>, +) -> Result, String> { + let mut states = states.into_iter(); + let first = states + .next() + .ok_or_else(|| "empty exact state input".to_string())?; + if let Some(first) = first + .as_any() + .downcast_ref::() + { + let mut merged = first.clone(); + for state in states { + let other = state + .as_any() + .downcast_ref::() + .ok_or_else(|| "merge requires Planner exact state".to_string())?; + merged + .merge_from(other) + .map_err(|error| error.to_string())?; + } + return Ok(Box::new(merged)); + } + let mut merged = first.clone_boxed_core(); + for state in states { + merged = merged + .merge_with(state.as_ref()) + .map_err(|e| e.to_string())?; + } + Ok(merged) +} + +/// PromQL counter readouts omit a series with fewer than two samples. Other +/// state/type/range failures remain errors rather than empty results. +pub fn insufficient_counter_samples( + state: &dyn crate::AggregateCore, + statistic: crate::Statistic, +) -> bool { + matches!( + statistic, + crate::Statistic::Rate | crate::Statistic::Increase + ) && (state + .as_any() + .downcast_ref::() + .is_some_and(|state| { + state.sample_count < 2 || state.last_seen_timestamp == state.starting_timestamp + }) + || state + .as_any() + .downcast_ref::() + .is_some_and(|state| state.insufficient_counter_samples(statistic, &None))) +} + +pub fn exact_readout_optional( + states: impl IntoIterator>, + statistic: crate::Statistic, + key: &Option, + parameters: &std::collections::HashMap, +) -> Result, String> { + let merged = merge_exact_states(states)?; + let counter = key.as_ref().and_then(|key| { + merged + .as_any() + .downcast_ref::() + .and_then(|state| state.increases.get(key)) + }); + let exact_insufficient = merged + .as_any() + .downcast_ref::() + .is_some_and(|state| state.insufficient_counter_samples(statistic, key)); + if exact_insufficient + || insufficient_counter_samples(merged.as_ref(), statistic) + || counter.is_some_and(|counter| insufficient_counter_samples(counter, statistic)) + { + return Ok(None); + } + merged + .query_statistic(statistic, key, parameters) + .map(Some) + .map_err(|error| error.to_string()) +} + +#[cfg(test)] +mod counter_tests { + use super::*; + use crate::{summary_kernels::IncreaseAccumulator, AggregateCore, Measurement, Statistic}; + use std::sync::Arc; + + #[test] + fn planner_counter_population_omits_insufficient_samples() { + use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; + for (kind, params, statistic) in [ + (ExactKind::Rate, ExactParams::Rate, Statistic::Rate), + ( + ExactKind::Increase, + ExactParams::Increase, + Statistic::Increase, + ), + ] { + for keyed in [false, true] { + let mut state = crate::summary_kernels::exact::ExactAccumulator::new( + SummaryFamilyType::ExactAggregate(kind.clone(), params.clone()), + keyed, + ) + .unwrap(); + let key = keyed + .then(|| crate::KeyByLabelValues::new_with_labels(vec!["checkout".into()])); + state.update(key.as_ref(), 10., 10_000); + assert_eq!( + exact_readout_optional( + [Arc::new(state) as Arc], + statistic, + &key, + &Default::default() + ) + .unwrap(), + None + ); + } + } + } + + #[test] + fn sparse_counter_is_absent_but_invalid_ranges_still_fail() { + let mut state = + IncreaseAccumulator::new(Measurement::new(10.), 10_000, Measurement::new(10.), 10_000); + let parameters = std::collections::HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + assert_eq!( + exact_readout_optional( + [Arc::new(state.clone()) as Arc], + Statistic::Rate, + &None, + ¶meters + ) + .unwrap(), + None + ); + let mut repeated = state.clone(); + repeated.update(Measurement::new(10.), 10_000); + assert_eq!( + exact_readout_optional( + [Arc::new(repeated) as Arc], + Statistic::Rate, + &None, + ¶meters + ) + .unwrap(), + None + ); + let mut keyed = crate::summary_kernels::KeyedCounterState::new(); + let label = crate::KeyByLabelValues::new_with_labels(vec!["checkout".into()]); + keyed.update(label.clone(), state.clone()); + assert_eq!( + exact_readout_optional( + [Arc::new(keyed.clone()) as Arc], + Statistic::Rate, + &Some(label), + ¶meters + ) + .unwrap(), + None + ); + assert!(exact_readout_optional( + [Arc::new(keyed) as Arc], + Statistic::Rate, + &Some(crate::KeyByLabelValues::new_with_labels(vec![ + "missing".into() + ])), + ¶meters + ) + .is_err()); + state.update(Measurement::new(20.), 20_000); + assert!(exact_readout_optional( + [Arc::new(state.clone()) as Arc], + Statistic::Rate, + &None, + ¶meters + ) + .unwrap() + .is_some()); + let invalid = std::collections::HashMap::from([ + ("range_start_ms".into(), "60000".into()), + ("range_end_ms".into(), "0".into()), + ]); + assert!(exact_readout_optional( + [Arc::new(state) as Arc], + Statistic::Rate, + &None, + &invalid + ) + .is_err()); + assert!(exact_readout_optional([], Statistic::Rate, &None, ¶meters).is_err()); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/count_min_sketch.rs b/crates/asap-physical-operators/src/summary_kernels/count_min_sketch.rs new file mode 100644 index 00000000..cffe90d2 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/count_min_sketch.rs @@ -0,0 +1,1323 @@ +use crate::summary_kernels::dd_sketch::normalize_sample_p; +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountMinSketch, CountMinSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Count-Min Sketch accumulator — wraps asap_sketchlib::CountMinSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. +/// This file retains QE-specific trait impls, legacy deserializers, and JSON output. +#[derive(Debug, Clone)] +pub struct CountMinSketchAccumulator { + pub inner: CountMinSketch, + /// Edge sampling probability `p ∈ (0,1]` carried on the producer's + /// `SketchEnvelope.sample_p`. The edge admits each insert with + /// probability `p`, so every stored cell count is ~`p`× the true count. + /// CMS is L1/additive and linear, so the unbiased rescale of BOTH a + /// point-frequency estimate (`query_key`) and the aggregate + /// total-event statistics (`Count`/`Sum`/`Increase`/`Rate`) is `×1/p`. + /// `1.0` (and the proto3 default `0.0`, dual-read as `1.0`) means no + /// sampling, so the rescale is a no-op and the behaviour is identical + /// to before. Mirrors `DDSketchAccumulator::sample_p`; set from the + /// envelope at the `from_sketchlib_proto_bytes` decode site and + /// preserved across `reset_to_empty` and `merge_with`. + pub sample_p: f64, +} + +impl CountMinSketchAccumulator { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + inner: CountMinSketch::new(row_num, col_num), + sample_p: 1.0, + } + } + + // Marked as _update and kept private; only called internally. + fn _update(&mut self, key: &KeyByLabelValues, value: f64) { + self.inner.update(&key.to_semicolon_str(), value); + } + + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + // The edge sampled inserts with probability `sample_p`, so the + // stored point-frequency estimate is ~`p`× the true frequency. + // CMS is linear/additive, so `×1/p` is the unbiased rescale. + // `sample_p == 1.0` (unsampled / legacy) makes this a no-op. + self.inner.estimate(&key.to_semicolon_str()) / self.sample_p + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let row_num = data["row_num"] + .as_f64() + .ok_or("Missing or invalid 'row_num' field")? as usize; + let col_num = data["col_num"] + .as_f64() + .ok_or("Missing or invalid 'col_num' field")? as usize; + + let sketch_data = data["sketch"] + .as_array() + .ok_or("Missing or invalid 'sketch' field")?; + + let mut sketch = Vec::new(); + for row in sketch_data { + let row_array = row.as_array().ok_or("Invalid row in sketch data")?; + let mut sketch_row = Vec::new(); + for cell in row_array { + let value = cell.as_f64().ok_or("Invalid cell value in sketch data")?; + sketch_row.push(value); + } + sketch.push(sketch_row); + } + + Ok(Self { + inner: CountMinSketch::from_legacy_matrix(sketch, row_num, col_num), + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `CountMinSketchDataPoint.sketch` bytes when + /// `encoding = COUNT_MIN_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `CountMinSketch` wire struct (same format the legacy Arroyo path + /// uses — this method is the modified-OTLP entrypoint for PR I). + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountMinSketch::from_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, + // The msgpack CountMinSketch struct carries no envelope/sample_p; + // the msgpack path is parity/test-only and is never edge-sampled. + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `CountMinSketchDataPoint.sketch` bytes — i.e. the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::CountMinState` message used by + /// DataCollector's `countminsketchprocessor` when emitting via + /// `Metric.data = CountMinSketch{…}` with + /// `encoding = COUNT_MIN_SKETCH_ENCODING_PROTO`. + /// + /// The resulting accumulator is constructed via + /// `CountMinSketch::from_legacy_matrix` after reshaping the flat + /// `counts_int` / `counts_float` field into a `Vec>`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + + // DataCollector's countminsketchprocessor wraps the state in a + // `SketchEnvelope{count_min: CountMinState}` via + // `SerializePortableFO` + `proto.Marshal`. Try decoding as envelope + // first, fall back to bare `CountMinState` for callers (e.g. unit + // tests) that encode the state directly. Capture the envelope's + // `sample_p` alongside the state so the point-frequency + // (`query_key`) and aggregate statistics rescale by `1/p`. Bare + // `CountMinState` bytes (no envelope) carry no sampling info → + // `sample_p` 1.0 (no rescale). Mirrors `DDSketchAccumulator`. + let (state, sample_p) = match SketchEnvelope::decode(buffer) { + Ok(env) => { + let sp = env.sample_p; + match env.sketch_state { + Some(sketch_envelope::SketchState::CountMin(st)) => (st, sp), + Some(other) => { + return Err(format!( + "SketchEnvelope contains non-CountMin sketch: {:?}", + std::mem::discriminant(&other) + ) + .into()); + } + // Envelope decoded but was empty (e.g. the buffer is a + // bare CountMinState that happened to parse as a default + // envelope). Fall through to bare decode. + None => ( + CountMinState::decode(buffer) + .map_err(|e| format!("decode CountMinState: {e}"))?, + 1.0, + ), + } + } + Err(_) => ( + CountMinState::decode(buffer).map_err(|e| format!("decode CountMinState: {e}"))?, + 1.0, + ), + }; + let rows = state.rows as usize; + let cols = state.cols as usize; + // Defensive dim validation BEFORE reconstructing the matrix: + // reject degenerate / narrow-hash-budget-violating / absurdly + // oversized dims so a malformed payload fails gracefully (the + // ingest caller skips the data point) instead of building a + // degenerate or huge matrix. + validate_sketch_dims("CountMinState", rows, cols)?; + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { + format!( + "CountMinState has unknown counter_type tag {}", + state.counter_type + ) + })?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountMinState counts_int has {} entries, expected rows*cols = {}", + state.counts_int.len(), + expected_len + ) + .into()); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountMinState counts_float has {} entries, expected rows*cols = {}", + state.counts_float.len(), + expected_len + ) + .into()); + } + state.counts_float.clone() + } + // INT128 stores (hi, lo) pairs and would have 2 * rows * cols + // entries in counts_int; defer to PR C if a producer ever uses it. + other => { + return Err(format!( + "CountMinState counter_type {other:?} not yet supported \ + (PR C will extend coverage)" + ) + .into()); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(Self { + inner: CountMinSketch::from_legacy_matrix(matrix, rows, cols), + sample_p: normalize_sample_p(sample_p), + }) + } + + /// Apply a proto-encoded `CountMinDelta` frame to this + /// accumulator's inner sketch — the decode path for + /// `COUNT_MIN_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountMinDelta: {e}"))?; + + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountMinDelta packed-array length mismatch: \ + cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + ) + .into()); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + let delta = CountMinSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l1: pb.l1, + l2: pb.l2, + // The Go-side CountMinDelta proto now carries an hh_keys field + // (heavy-hitter candidates), mirrored on asap_sketchlib's + // CountMinSketchDelta. The vendored Rust proto bindings here don't + // decode it yet, and CountMin has no TopK to rebuild, so pass an + // empty set — same handling as CountSketch's hh_keys. + hh_keys: Vec::new(), + }; + self.inner + .apply_delta(&delta) + .map_err(|e| format!("apply CountMinDelta: {e}"))?; + Ok(()) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 8 { + return Err("Buffer too short for row_num and col_num".into()); + } + + // TODO: this logic will need to be checked for i32 -> f64 + // Github Issue #11 + + let row_num = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; + let col_num = u32::from_le_bytes([buffer[4], buffer[5], buffer[6], buffer[7]]) as usize; + + let expected_size = 8 + (row_num * col_num * 4); + if buffer.len() < expected_size { + return Err("Buffer too short for sketch data".into()); + } + + let mut sketch = Vec::new(); + let mut offset = 8; + + for _ in 0..row_num { + let mut row = Vec::new(); + for _ in 0..col_num { + let value = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + row.push(value); + offset += 8; + } + sketch.push(row); + } + + Ok(Self { + inner: CountMinSketch::from_legacy_matrix(sketch, row_num, col_num), + sample_p: 1.0, + }) + } + + /// Merge multiple accumulators efficiently without cloning all of them. + pub fn merge_multiple( + accumulators: &[Box], + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut cms_accumulators = Vec::with_capacity(accumulators.len()); + for acc in accumulators { + if acc.get_accumulator_type() != AggregationType::CountMinSketch { + return Err(format!( + "Cannot merge CountMinSketchAccumulator with {:?}", + acc.get_accumulator_type() + ) + .into()); + } + let cms_acc = acc + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountMinSketchAccumulator")?; + cms_accumulators.push(cms_acc); + } + + // Check dimensions are consistent + let rows = cms_accumulators[0].inner.rows(); + let cols = cms_accumulators[0].inner.cols(); + for acc in &cms_accumulators { + if acc.inner.rows() != rows || acc.inner.cols() != cols { + return Err( + "Cannot merge CountMinSketch accumulators with different dimensions".into(), + ); + } + } + + let inner_refs: Vec<&CountMinSketch> = + cms_accumulators.iter().map(|acc| &acc.inner).collect(); + let merged_inner = CountMinSketch::merge_refs(&inner_refs)?; + // sample_p is a per-series config constant, so all operands carry the + // same value in practice. Mirror DDSketch's merge policy: prefer a + // sampled factor (< 1.0) over the no-sampling default so a merge with + // a freshly-reset (1.0) base keeps the series' sampling rate. + let sample_p = cms_accumulators + .iter() + .map(|acc| acc.sample_p) + .find(|&p| p < 1.0) + .unwrap_or(cms_accumulators[0].sample_p); + Ok(Self { + inner: merged_inner, + sample_p, + }) + } +} + +/// Defensive upper bound on the number of matrix cells (`rows * cols`) +/// we'll reconstruct from an inbound wire-declared CMS / CountSketch +/// dimension pair. A malformed / hostile payload could declare absurd +/// dims (e.g. `rows = cols = u32::MAX`) and trick the decoder into a +/// huge `Vec` allocation before the `counts_*.len() != rows*cols` +/// check ever runs. Realistic sketches are at most a few hundred rows +/// by tens-of-thousands of columns, so 8M cells (~64 MiB of f64) is a +/// generous ceiling that no legitimate producer reaches. +pub(crate) const MAX_SKETCH_CELLS: usize = 8 * 1024 * 1024; + +/// Validate an inbound, wire-declared `(rows, cols)` pair for a +/// matrix-backed frequency sketch (CMS / CountSketch) BEFORE any matrix +/// is reconstructed from it. Returns `Ok(())` for dimensions a +/// legitimate producer could have emitted, and an `Err` (never a panic) +/// for malformed / degenerate ones so the ingest path can skip the data +/// point and fall through to its existing decode-failure accounting. +/// +/// Rejections: +/// 1. `rows < 1` or `cols < 1` — a zero-dim matrix has no cells. +/// 2. Narrow-hash-budget violation. The cross-language wire hasher +/// (`sketchlib`'s `MatrixHashType::Packed64`) derives every row's +/// column index from disjoint bit-fields of a single 64-bit hash +/// word: row `r` reads `mask_bits = ceil(log2(cols))` bits at offset +/// `r * mask_bits`. Once `rows * mask_bits > 64` the per-row column +/// slices overflow / alias the 64-bit word and the matrix-cell +/// layout is no longer the one the producer hashed into — the sketch +/// is internally degenerate. This mirrors sketchlib's own +/// `MatrixFastHash::assert_compatible` budget (`rows * (mask_bits + 1) <= 64`); we check the column-index bits alone so realistic +/// configs (5x2048, 5x4096, 5x2000) — for which the sign bits share +/// the top of the word without affecting the cell layout — still +/// pass. +/// 3. Obviously-oversized dims: `rows * cols > MAX_SKETCH_CELLS`, +/// guarding against a huge allocation from a malformed payload. +/// +/// `what` names the wire struct for the error message (e.g. +/// `"CountMinState"`). +pub(crate) fn validate_sketch_dims(what: &str, rows: usize, cols: usize) -> Result<(), String> { + if rows < 1 || cols < 1 { + return Err(format!( + "{what} has degenerate dims (rows={rows}, cols={cols}); rejecting" + )); + } + // mask_bits = ceil(log2(cols)); cols >= 1 here. ilog2 is floor(log2). + let mask_bits = if cols.is_power_of_two() { + cols.ilog2() as usize + } else { + cols.ilog2() as usize + 1 + }; + if rows.saturating_mul(mask_bits) > 64 { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) exceed the 64-bit \ + packed-hash column budget (rows * ceil(log2(cols)) = {} > 64); \ + the sketch's matrix-cell layout is degenerate, rejecting", + rows.saturating_mul(mask_bits) + )); + } + if rows.saturating_mul(cols) > MAX_SKETCH_CELLS { + return Err(format!( + "{what} dims (rows={rows}, cols={cols}) declare {} cells, \ + exceeding the {MAX_SKETCH_CELLS}-cell ingest cap; rejecting to \ + avoid a huge allocation from a malformed payload", + rows.saturating_mul(cols) + )); + } + Ok(()) +} + +impl SerializableToSink for CountMinSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), + "sketch": self.inner.sketch() + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountMinSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountMinSketchAccumulator" + } + + /// Per-window base rotation: rebuild an empty counter matrix with + /// the same (rows, cols) so the next window's additive cell deltas + /// align to the identical hash geometry. `sample_p` is a per-series + /// config constant (not per-window data), so it is intentionally + /// preserved across the rotation — mirrors `DDSketchAccumulator`. + fn reset_to_empty(&mut self) { + self.inner = CountMinSketch::new(self.inner.rows(), self.inner.cols()); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountMinSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_cms = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountMinSketchAccumulator")?; + + let merged_inner = CountMinSketch::merge_refs(&[&self.inner, &other_cms.inner])?; + // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a + // per-series config constant, so both operands carry the same value + // in practice. Prefer a sampled factor over the no-sampling default + // so a merge with a freshly-reset (1.0) base keeps the series' + // sampling rate. + let sample_p = if self.sample_p < 1.0 { + self.sample_p + } else { + other_cms.sample_p + }; + Ok(Box::new(Self { + inner: merged_inner, + sample_p, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountMinSketch + } + + fn approx_memory_bytes(&self) -> usize { + // Conservative constant for the CountMinSketch counter matrix. + // Real per-instance sizing would require exposing rows/cols on + // the inner sketch; 16 KiB is a reasonable v1 default. + 16 * 1024 + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + use crate::Statistic; + + // Key-provided path: route to MultipleSubpopulationAggregate::query + // (the canonical "what's the count of this key?" lookup). + if let Some(key_val) = key.as_ref() { + return self.query(statistic, key_val, Some(query_kwargs)); + } + if let Some(k) = query_kwargs.get("key") { + let key_val = crate::KeyByLabelValues::new_with_labels(vec![k.clone()]); + return self.query(statistic, &key_val, Some(query_kwargs)); + } + + // No-key path: return total event volume. The min-row-sum is the + // canonical CMS estimator for "how many inserts were observed" — + // each insert increments exactly one cell per row, so every row + // sums to the true insert count (modulo collisions, which CMS + // never *underestimates*; min is the tightest upper bound). + // + // When the edge sampled this series (sample_p < 1.0), each insert + // was admitted w.p. `p`, so the stored min-row-sum is ~`p`× the + // true event count. CMS is L1/additive and linear, so rescale by + // `1/sample_p` for an unbiased estimate. `sample_p == 1.0` + // (unsampled / legacy) makes this a no-op. This rescales BOTH the + // Count/Sum/Increase statistics and (via the same closure) the + // Rate per-second readout. + let total_events = || -> f64 { + let matrix = self.inner.sketch(); + if matrix.is_empty() || matrix[0].is_empty() { + return 0.0; + } + let row_totals = matrix.iter().map(|r| r.iter().sum::()); + let min_total = row_totals.fold(f64::INFINITY, f64::min); + if min_total.is_finite() { + min_total / self.sample_p + } else { + 0.0 + } + }; + match statistic { + Statistic::Count | Statistic::Sum => Ok(total_events()), + // PR #111 honest-gap closure (in-the-bag for ASAP tier). + // CMS records insert counts but not timestamps, so per-second + // `rate(metric[range])` requires the engine to push the + // range duration via `query_kwargs["range_ms"]`. When + // present, divide the min-row-sum by `range_ms / 1000`. When + // absent (the engine has not been wired to inject range_ms + // for this query, e.g. instant `rate` calls outside the + // PromQL range-vector pattern), fall back to the raw event + // count so the answer is at least non-empty — the caller's + // caveat is that the units are events/window rather than + // events/second. Increase carries the same caveat. + Statistic::Rate => { + let total = total_events(); + let range_ms_str = query_kwargs.get("range_ms").map(String::as_str); + let Some(s) = range_ms_str else { + return Ok(total); + }; + let range_ms: f64 = s + .parse() + .map_err(|e| format!("CountMinSketchAccumulator: bad range_ms='{s}': {e}"))?; + if range_ms <= 0.0 { + return Err("CountMinSketchAccumulator: range_ms must be positive".into()); + } + Ok(total * 1000.0 / range_ms) + } + Statistic::Increase => Ok(total_events()), + other => Err(format!( + "CountMinSketchAccumulator: statistic {:?} not supported \ + without a key (only Count / Sum / Rate / Increase aggregate \ + over the whole sketch)", + other, + ) + .into()), + } + } +} + +impl MultipleSubpopulationAggregate for CountMinSketchAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountMinSketchAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_min_sketch_creation() { + let cms = CountMinSketchAccumulator::new(4, 1000); + assert_eq!(cms.inner.rows(), 4); + assert_eq!(cms.inner.cols(), 1000); + let sketch = cms.inner.sketch(); + assert_eq!(sketch.len(), 4); + assert_eq!(sketch[0].len(), 1000); + + for row in &sketch { + for &value in row { + assert_eq!(value, 0.0); + } + } + } + + #[test] + fn test_count_min_sketch_update() { + let mut cms = CountMinSketchAccumulator::new(2, 10); + let key = KeyByLabelValues::new(); + cms._update(&key, 1.0); + let result = cms.query_key(&key); + assert!(result >= 1.0); + } + + #[test] + fn test_count_min_sketch_query() { + let cms = CountMinSketchAccumulator::new(2, 10); + let key = KeyByLabelValues::new(); + assert_eq!(cms.query_key(&key), 0.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); + } + + #[test] + fn test_count_min_sketch_merge() { + // Build controlled state via from_legacy_matrix (works for both Legacy and Sketchlib backends). + let cms1 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![5.0, 0.0, 0.0], vec![0.0, 0.0, 10.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let cms2 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![3.0, 7.0, 0.0], vec![0.0, 0.0, 0.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + + let merged = CountMinSketchAccumulator::merge_accumulators(vec![cms1, cms2]).unwrap(); + + let merged_sketch = merged.inner.sketch(); + assert_eq!(merged_sketch[0][0], 8.0); + assert_eq!(merged_sketch[0][1], 7.0); + assert_eq!(merged_sketch[1][2], 10.0); + } + + #[test] + fn test_count_min_sketch_merge_dimension_mismatch() { + let cms1 = CountMinSketchAccumulator::new(2, 3); + let cms2 = CountMinSketchAccumulator::new(3, 3); + let result = CountMinSketchAccumulator::merge_accumulators(vec![cms1, cms2]); + assert!(result.is_err()); + } + + #[test] + fn test_count_min_sketch_as_aggregate_core() { + let cms = CountMinSketchAccumulator::new(2, 3); + assert_eq!(cms.type_name(), "CountMinSketchAccumulator"); + } + + #[test] + fn test_trait_object() { + let cms = CountMinSketchAccumulator::new(2, 3); + let trait_obj: Box = Box::new(cms); + assert_eq!(trait_obj.type_name(), "CountMinSketchAccumulator"); + } + + #[test] + fn test_count_min_sketch_key_query() { + let mut cms = CountMinSketchAccumulator::new(4, 100); + let key = KeyByLabelValues::new(); + assert_eq!(cms.query_key(&key), 0.0); + cms._update(&key, 5.0); + let result = cms.query_key(&key); + assert!(result >= 5.0); + } + + #[test] + fn test_update_and_query_use_same_key_encoding() { + // Regression test: _update and query_key must hash the same key string. + // Previously _update went through serialize_to_json (which returns a JSON + // array, so as_object() is always None) and always stored under key "". + // query_key correctly used key.labels.join(";"), so they never matched. + let mut cms = CountMinSketchAccumulator::new(4, 1000); + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string(), "prod".to_string()]); + cms._update(&key, 5.0); + let result = cms.query_key(&key); + assert!( + result >= 5.0, + "_update and query_key used different key encodings: got {result}" + ); + + // Also verify a different key does not interfere. + let other_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + // other_key was never updated; its estimate should be lower than key's. + let other_result = cms.query_key(&other_key); + // In a sketch this large there should be no collision, so other_result == 0. + assert_eq!( + other_result, 0.0, + "unrelated key returned non-zero: {other_result}" + ); + } + + #[test] + fn test_multiple_subpopulation_aggregate() { + let mut cms = CountMinSketchAccumulator::new(3, 50); + let key = KeyByLabelValues::new(); + cms._update(&key, 10.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert!(result >= 10.0); + + let keys = multi_trait.get_keys(); + assert!(keys.is_none()); + } + + #[test] + fn test_count_min_sketch_merge_multiple() { + // Build controlled state via from_legacy_matrix (works for both Legacy and Sketchlib backends). + let cms1 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![5.0, 0.0, 0.0], vec![0.0, 0.0, 10.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let cms2 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![3.0, 7.0, 0.0], vec![0.0, 0.0, 0.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let cms3 = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![2.0, 0.0, 0.0], vec![0.0, 0.0, 5.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + + let boxed_accs: Vec> = + vec![Box::new(cms1), Box::new(cms2), Box::new(cms3)]; + + let merged = CountMinSketchAccumulator::merge_multiple(&boxed_accs).unwrap(); + + let merged_sketch = merged.inner.sketch(); + assert_eq!(merged_sketch[0][0], 10.0); + assert_eq!(merged_sketch[0][1], 7.0); + assert_eq!(merged_sketch[1][2], 15.0); + } + + #[test] + fn test_count_min_sketch_merge_multiple_error_cases() { + let empty: Vec> = vec![]; + assert!(CountMinSketchAccumulator::merge_multiple(&empty).is_err()); + + let cms1 = CountMinSketchAccumulator::new(2, 3); + let cms2 = CountMinSketchAccumulator::new(3, 3); + let boxed_accs: Vec> = vec![Box::new(cms1), Box::new(cms2)]; + assert!(CountMinSketchAccumulator::merge_multiple(&boxed_accs).is_err()); + + use crate::summary_kernels::sum::SumAccumulator; + let cms = CountMinSketchAccumulator::new(2, 3); + let sum = SumAccumulator::new(); + let mixed_accs: Vec> = vec![Box::new(cms), Box::new(sum)]; + assert!(CountMinSketchAccumulator::merge_multiple(&mixed_accs).is_err()); + } + + #[test] + fn test_from_sketchlib_proto_bytes_int64() { + // Hand-build a CountMinState proto with INT64 counters and verify + // round-tripping through from_sketchlib_proto_bytes yields the same + // matrix that the modified-OTLP wire format would carry. + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + + let rows = 2u32; + let cols = 3u32; + // Row-major: row 0 = [1,2,3], row 1 = [4,5,6] + let counts_int: Vec = vec![1, 2, 3, 4, 5, 6]; + let state = CountMinState { + rows, + cols, + counter_type: CounterType::Int64 as i32, + counts_int: counts_int.clone(), + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let bytes = state.encode_to_vec(); + + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix.len(), rows as usize); + assert_eq!(matrix[0], vec![1.0, 2.0, 3.0]); + assert_eq!(matrix[1], vec![4.0, 5.0, 6.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's countminsketchprocessor emits: + // the state is wrapped in a `SketchEnvelope{count_min: ...}` + // via sketchlib-go's `SerializePortableFO` + `proto.Marshal`. + // Before the fix, the Rust decoder decoded the envelope bytes as + // a bare CountMinState, which produced "invalid wire type" + // errors on field `cols` and silently fell through to §5.2. + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + + let state = CountMinState { + rows: 2, + cols: 3, + counter_type: CounterType::Int64 as i32, + counts_int: vec![7, 8, 9, 10, 11, 12], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![7.0, 8.0, 9.0]); + assert_eq!(matrix[1], vec![10.0, 11.0, 12.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + // An envelope carrying a non-CountMin sketch should be rejected + // with a clear error rather than silently producing garbage. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + + let kll = KllState::default(); + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(kll)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_float64() { + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + + let state = CountMinState { + rows: 2, + cols: 2, + counter_type: CounterType::Float64 as i32, + counts_int: Vec::new(), + counts_float: vec![1.5, 2.5, 3.5, 4.5], + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let bytes = state.encode_to_vec(); + + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.5, 2.5]); + assert_eq!(matrix[1], vec![3.5, 4.5]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_dimension_mismatch() { + // counts_int has 5 entries but rows*cols = 6 → expect error + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + + let state = CountMinState { + rows: 2, + cols: 3, + counter_type: CounterType::Int64 as i32, + counts_int: vec![1, 2, 3, 4, 5], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let bytes = state.encode_to_vec(); + + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!( + result.unwrap_err().to_string().contains("counts_int"), + "error should mention counts_int dim mismatch" + ); + } + + #[test] + fn test_from_sketchlib_proto_bytes_zero_dims_rejected() { + use asap_sketchlib::proto::sketchlib::CountMinState; + use prost::Message; + + let state = CountMinState::default(); + let bytes = state.encode_to_vec(); + + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("degenerate dims")); + } + + #[test] + fn test_apply_proto_delta_bytes_round_trip() { + use asap_sketchlib::proto::sketchlib::CountMinDelta as PbDelta; + use prost::Message; + + let mut acc = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], + 2, + 3, + ), + sample_p: 1.0, + }; + let bytes = PbDelta { + rows: 2, + cols: 3, + cell_rows: vec![0, 1], + cell_cols: vec![0, 2], + d_counts: vec![10, 100], + l1: vec![], + l2: vec![], + ..Default::default() + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!( + acc.inner.sketch(), + vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 106.0]] + ); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = CountMinSketchAccumulator::new(2, 3); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ---------------------------------------------------------------- + // Statistic::Rate / Statistic::Increase — PR #111 honest-gap closure. + // CMS records insert counts but not timestamps. The Rate readout + // requires the engine to push `range_ms` via query_kwargs; without + // it the accumulator falls back to the raw event count (units of + // events/window) so the answer is at least non-empty. + // ---------------------------------------------------------------- + + #[test] + fn test_query_statistic_rate_with_range_ms() { + // Build a CMS whose min-row-sum is 100 events. With a 5-minute + // (300_000 ms) range, the per-second rate is 100 / 300 ≈ 0.333. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix( + vec![vec![100.0, 0.0], vec![100.0, 0.0]], + 2, + 2, + ), + sample_p: 1.0, + }; + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "300000".to_string()); + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect("Rate with range_ms is supported"); + assert!( + (v - (100.0 / 300.0)).abs() < 1e-9, + "expected 100/300 = {}, got {v}", + 100.0 / 300.0, + ); + } + + #[test] + fn test_query_statistic_rate_without_range_ms_falls_back_to_count() { + // Without `range_ms` in kwargs the accumulator returns the raw + // event volume (events/window units). Caller is responsible for + // surfacing that caveat to the user; this avoids `status=error` + // for instant rate-shape queries that bypass the matrix-selector + // code path. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix(vec![vec![42.0, 0.0], vec![42.0, 0.0]], 2, 2), + sample_p: 1.0, + }; + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Rate, &None, &HashMap::new()) + .expect("Rate without range_ms still answers (fallback)"); + assert_eq!(v, 42.0); + } + + #[test] + fn test_query_statistic_increase_returns_total_count() { + // Increase semantics on CMS: total events in the window — the + // same min-row-sum as Sum / Count. Differs from Rate only in + // that it never divides by range. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix(vec![vec![5.0, 7.0], vec![3.0, 9.0]], 2, 2), + sample_p: 1.0, + }; + let trait_obj: &dyn AggregateCore = &cms; + let v = trait_obj + .query_statistic(Statistic::Increase, &None, &HashMap::new()) + .expect("Increase is supported"); + // min-row-sum: row0 = 12, row1 = 12, min = 12. + assert_eq!(v, 12.0); + } + + // ---------------------------------------------------------------- + // Defensive inbound-dimension validation (harden/sketch-dim-validation). + // Malformed / degenerate / narrow-hash-budget-violating CMS dims must + // be rejected gracefully (Err, never a panic); valid configs the + // backend actually uses (5x2048, 5x4096, 5x2000) must still decode. + // ---------------------------------------------------------------- + + /// Build a bare `CountMinState` proto carrying the given dims and a + /// row-major INT64 counts vector sized to `rows*cols` so that, IF the + /// dims pass validation, the reshape also succeeds. Used to prove a + /// malformed-dim payload is rejected at the dim gate, not later. + fn cms_state_bytes(rows: u32, cols: u32) -> Vec { + use asap_sketchlib::proto::sketchlib::{CountMinState, CounterType}; + use prost::Message; + let n = (rows as usize).saturating_mul(cols as usize); + let state = CountMinState { + rows, + cols, + counter_type: CounterType::Int64 as i32, + counts_int: vec![0i64; n], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + state.encode_to_vec() + } + + #[test] + fn test_validate_sketch_dims_accepts_valid_configs() { + // The realistic configs the backend uses must pass unchanged. + for (r, c) in [(5usize, 2048usize), (5, 4096), (5, 2000), (4, 1000), (2, 3)] { + assert!( + validate_sketch_dims("CountMinState", r, c).is_ok(), + "valid config {r}x{c} was wrongly rejected" + ); + } + } + + #[test] + fn test_validate_sketch_dims_rejects_malformed() { + // Zero dims. + assert!(validate_sketch_dims("CountMinState", 0, 2048).is_err()); + assert!(validate_sketch_dims("CountMinState", 5, 0).is_err()); + // Narrow-hash-budget violation: 5 * ceil(log2(8192))=5*13=65 > 64. + let err = validate_sketch_dims("CountMinState", 5, 8192).unwrap_err(); + assert!(err.contains("budget"), "expected budget error, got: {err}"); + // Absurdly oversized: 1 x 16,777,216 = 16M cells > 8M cap. (1 row + // keeps the hash budget tiny — 1*24=24 — so the cap check, not the + // budget check, is what fires here.) + let err = validate_sketch_dims("CountMinState", 1, 16_777_216).unwrap_err(); + assert!(err.contains("cap"), "expected cell-cap error, got: {err}"); + // No panic on extreme dims (saturating_mul guards the products). + assert!(validate_sketch_dims("CountMinState", usize::MAX, usize::MAX).is_err()); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { + // A data point declaring narrow-hash-budget-violating dims must be + // skipped (Err returned, NOT a panic). The ingest caller turns + // this Err into a dropped data point + WARN log. + let bytes = cms_state_bytes(5, 8192); + let result = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "budget-violating dims should be rejected"); + assert!(result.unwrap_err().to_string().contains("rejecting")); + + // A valid neighbour (5x4096) on the same path still decodes fine. + let ok_bytes = cms_state_bytes(5, 4096); + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) + .expect("valid 5x4096 CMS should still decode"); + assert_eq!(acc.inner.rows(), 5); + assert_eq!(acc.inner.cols(), 4096); + } + + #[test] + fn test_query_statistic_rate_rejects_invalid_range_ms() { + let cms = CountMinSketchAccumulator::new(2, 2); + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "0".to_string()); + let trait_obj: &dyn AggregateCore = &cms; + let err = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect_err("range_ms=0 should error"); + assert!(err.to_string().contains("positive")); + + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "not-a-number".to_string()); + let err = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect_err("non-numeric range_ms should error"); + assert!(err.to_string().contains("bad range_ms")); + } + + // ---------------------------------------------------------------- + // sample_p rescale. The edge admits each insert with probability `p`, + // so every stored cell is ~p× the true count. CMS is L1/additive and + // linear, so BOTH the point-frequency (query_key) and the aggregate + // total-event statistics (Count/Sum/Increase/Rate) rescale by 1/p. + // ---------------------------------------------------------------- + + #[test] + fn test_query_key_rescaled_by_sample_p() { + // Same stored cell counts, two sample_p values: the p=0.25 sketch + // must report 4× the point-frequency of the unsampled one. + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let mut unsampled = CountMinSketchAccumulator::new(4, 1000); + unsampled._update(&key, 10.0); + let mut sampled = CountMinSketchAccumulator::new(4, 1000); + sampled._update(&key, 10.0); + sampled.sample_p = 0.25; + + let raw = unsampled.query_key(&key); + let rescaled = sampled.query_key(&key); + assert!( + raw >= 10.0, + "raw estimate should be >= inserted 10, got {raw}" + ); + assert!( + (rescaled - raw * 4.0).abs() < 1e-9, + "expected point-frequency rescaled ≈ 4×raw ({}), got {rescaled}", + raw * 4.0 + ); + } + + #[test] + fn test_aggregate_statistics_rescaled_by_sample_p() { + use crate::Statistic; + // Build a CMS with a known min-row-sum of 12 events, sampled at + // p=0.25 → every aggregate statistic should report 12 / 0.25 = 48. + let cms = CountMinSketchAccumulator { + inner: CountMinSketch::from_legacy_matrix(vec![vec![5.0, 7.0], vec![3.0, 9.0]], 2, 2), + sample_p: 0.25, + }; + let trait_obj: &dyn AggregateCore = &cms; + for stat in [Statistic::Count, Statistic::Sum, Statistic::Increase] { + let v = trait_obj + .query_statistic(stat, &None, &HashMap::new()) + .unwrap_or_else(|e| panic!("{stat:?} should be supported: {e}")); + // min-row-sum = 12, rescaled by 1/0.25 = 48. + assert!( + (v - 48.0).abs() < 1e-9, + "{stat:?}: expected rescaled 48, got {v}" + ); + } + // Rate also divides through the rescaled total: 48 events over a + // 6-second (6000 ms) range = 8 events/s. + let mut kwargs = HashMap::new(); + kwargs.insert("range_ms".to_string(), "6000".to_string()); + let r = trait_obj + .query_statistic(Statistic::Rate, &None, &kwargs) + .expect("rate ok"); + assert!((r - 8.0).abs() < 1e-9, "expected rate 8.0, got {r}"); + } + + #[test] + fn test_sample_p_unset_behaves_as_one() { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + // An envelope with no sample_p (proto3 default 0.0) must normalize + // to 1.0 (no rescale) — byte-compatible with legacy frames. + let state = CountMinState { + rows: 2, + cols: 2, + counter_type: CounterType::Int64 as i32, + counts_int: vec![1, 2, 3, 4], + counts_float: Vec::new(), + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let env = SketchEnvelope { + // sample_p left at proto3 default 0.0. + sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 1.0, "unset sample_p must normalize to 1.0"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { + use crate::Statistic; + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountMinState, CounterType, SketchEnvelope, + }; + use prost::Message; + // min-row-sum = 12 raw; sample_p 0.25 → Count = 48. + let state = CountMinState { + rows: 2, + cols: 2, + counter_type: CounterType::Float64 as i32, + counts_int: Vec::new(), + counts_float: vec![5.0, 7.0, 3.0, 9.0], + sum_counts: Vec::new(), + sum2_counts: Vec::new(), + l1: Vec::new(), + l2: Vec::new(), + }; + let env = SketchEnvelope { + sample_p: 0.25, + sketch_state: Some(sketch_envelope::SketchState::CountMin(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = CountMinSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 0.25); + let trait_obj: &dyn AggregateCore = &acc; + let v = trait_obj + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + assert!((v - 48.0).abs() < 1e-9, "expected rescaled 48, got {v}"); + } + + #[test] + fn test_reset_to_empty_preserves_sample_p() { + let mut acc = CountMinSketchAccumulator::new(2, 3); + acc.sample_p = 0.25; + acc.reset_to_empty(); + assert_eq!(acc.sample_p, 0.25, "window rotation must keep sample_p"); + } + + #[test] + fn test_merge_prefers_sampled_factor() { + let mut a = CountMinSketchAccumulator::new(2, 3); + a.sample_p = 0.25; + let b = CountMinSketchAccumulator::new(2, 3); // sample_p 1.0 + let merged = a.merge_with(&b).expect("merge ok"); + let merged = merged + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.sample_p, 0.25); + + // merge_multiple mirrors the same policy. + let mut c = CountMinSketchAccumulator::new(2, 3); + c.sample_p = 0.25; + let d = CountMinSketchAccumulator::new(2, 3); + let boxed: Vec> = vec![Box::new(d), Box::new(c)]; + let merged = CountMinSketchAccumulator::merge_multiple(&boxed).expect("merge ok"); + assert_eq!(merged.sample_p, 0.25); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/count_min_sketch_with_heap.rs b/crates/asap-physical-operators/src/summary_kernels/count_min_sketch_with_heap.rs new file mode 100644 index 00000000..f5d1369d --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/count_min_sketch_with_heap.rs @@ -0,0 +1,832 @@ +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CmsHeapItem, CountMinSketchWithHeap, MessagePackCodec}; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Local serde view of the DELTA-HEAP wire frame produced by sketchlib-go's +/// `CountSketch.SerializeMsgpackWithHeapDelta` (encoding `MSGPACK_DELTA`). +/// Decoded with `rmp_serde` directly in the backend so NO delta API needs to +/// be added to the public `asap_sketchlib`. +/// +/// rmp_serde compact layout — a 4-element positional array: +/// +/// [ +/// is_delta: bool (always true), +/// matrix_delta: ( rows:u32, cols:u32, cells: Vec<(u32,u32,i64)> ), +/// topk_heap: Vec<(String, f64)>, // FULL heap, [key, value] pairs +/// heap_size: u64, +/// ] +/// +/// Tuple structs deserialize from msgpack fixed arrays positionally, so this +/// matches the Go encoder's byte layout exactly (no field names on the wire). +#[derive(Debug, Deserialize)] +struct HeapDeltaWire { + is_delta: bool, + matrix_delta: MatrixDeltaWire, + topk_heap: Vec<(String, f64)>, + #[allow(dead_code)] + heap_size: u64, +} + +#[derive(Debug, Deserialize)] +struct MatrixDeltaWire { + rows: u32, + cols: u32, + cells: Vec<(u32, u32, i64)>, +} + +/// Validated/flattened view of a decoded DELTA-HEAP frame. +struct HeapDeltaFrame { + rows: u32, + cols: u32, + heap_size: u64, + cells: Vec<(u32, u32, i64)>, + heap: Vec<(String, f64)>, +} + +impl HeapDeltaFrame { + fn from_msgpack(buffer: &[u8]) -> Result> { + let wire: HeapDeltaWire = rmp_serde::from_slice(buffer) + .map_err(|e| format!("decode CountSketchWithHeap delta msgpack: {e}"))?; + if !wire.is_delta { + return Err("CountSketchWithHeap delta frame has is_delta=false".into()); + } + Ok(Self { + rows: wire.matrix_delta.rows, + cols: wire.matrix_delta.cols, + heap_size: wire.heap_size, + cells: wire.matrix_delta.cells, + heap: wire.topk_heap, + }) + } +} + +/// Count-Min Sketch with Heap accumulator — wraps `asap_sketchlib::CountMinSketchWithHeap`. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::message_pack_format::portable::countminsketch_topk`. +/// This file retains QE-specific trait impls, legacy deserializers, and JSON output. +#[derive(Debug, Clone)] +pub struct CountMinSketchWithHeapAccumulator { + pub inner: CountMinSketchWithHeap, +} + +// Re-export HeapItem so existing code using CountMinSketchWithHeapAccumulator::HeapItem still works. +pub use asap_sketchlib::CmsHeapItem as HeapItemReexport; + +impl CountMinSketchWithHeapAccumulator { + pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { + Self { + inner: CountMinSketchWithHeap::new(row_num, col_num, heap_size), + } + } + + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + let key_string = key.labels.join(";"); + self.inner.estimate(&key_string) + } + + /// Decode a heap-bearing CountSketch FULL msgpack frame + /// (`{sketch:[matrix,rows,cols], topk_heap, heap_size}`) into a heap + /// accumulator. This is the window-1 / full-frame base for the + /// DELTA-HEAP delta path: the backend caches THIS accumulator as the + /// per-series base so a later `MSGPACK_DELTA` frame applies its sparse + /// matrix delta onto a heap accumulator (not a plain CountSketch). + /// + /// Delegates to the PUBLIC `asap_sketchlib::CountMinSketchWithHeap:: + /// from_msgpack` (both heap-bearing frequency variants share the wire + /// shape; the CountSketch-with-heap promotion is decided by the ingest + /// router, not the bytes). + pub fn from_msgpack_with_heap_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountMinSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountMinSketchWithHeap msgpack: {e}"))?, + }) + } + + /// Apply a DELTA-HEAP msgpack frame (encoding `MSGPACK_DELTA`) onto this + /// accumulator IN PLACE, WITHOUT any change to the public + /// `asap_sketchlib`: the frame is decoded generically with `rmp_serde` + /// into local serde structs, the sparse signed cell deltas are added to + /// the stored matrix (read back via the public `sketch_matrix()`), and + /// the top-k heap is REPLACED with the frame's full heap. The rebuilt + /// inner is produced via the public `from_legacy_matrix`, which rounds + /// cells to the i64 storage and re-seeds the heap. + /// + /// Under the per-window-reset model (`docs/delta-baseline-contract.md` + /// §3) the ingest caller resets this accumulator to empty at a window + /// boundary before applying, so the delta — which is the window's own + /// matrix against an empty base — reconstructs the window's state. + pub fn apply_msgpack_heap_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + + let rows = self.inner.rows(); + let cols = self.inner.cols(); + let heap_size = self.inner.heap_size; + + // Read the current (post-reset, possibly empty) matrix and apply the + // sparse signed deltas additively. Cells outside the stored + // dimensions are skipped defensively (mirrors the plain-CountSketch + // delta apply). + let mut matrix = self.inner.sketch_matrix(); + for (r, c, dc) in &frame.cells { + let (r, c) = (*r as usize, *c as usize); + if r >= rows || c >= cols { + continue; + } + matrix[r][c] += *dc as f64; + } + + // Replace the heap with the frame's full heap. `from_legacy_matrix` + // re-seeds both the matrix and the heap from these inputs. + let heap: Vec = frame + .heap + .into_iter() + .map(|(key, value)| CmsHeapItem { key, value }) + .collect(); + + self.inner = + CountMinSketchWithHeap::from_legacy_matrix(matrix, heap, rows, cols, heap_size); + Ok(()) + } + + /// Reconstruct a heap accumulator STANDALONE from a single DELTA-HEAP + /// msgpack frame (encoding `MSGPACK_DELTA`), with NO cached per-series + /// base. Used by the read-side reducer's `FrequencyTopk` path, where — + /// unlike the ingest accumulator — there is no rolling base to apply + /// onto: under the per-window-reset contract + /// (`docs/delta-baseline-contract.md` §3) each window's delta encodes + /// that window's own state against an EMPTY base, so reconstruction is + /// "empty(dims) + apply(delta)". + /// + /// Reuses the exact ingest-side apply logic: read the (rows, cols, + /// heap_size) the frame declares, build an empty accumulator of those + /// dims (equivalent to `reset_to_empty` on a same-shape base), then + /// fold the frame in via `apply_msgpack_heap_delta_bytes`. No + /// `asap_sketchlib` change — the frame is decoded generically with + /// `rmp_serde`. + pub fn from_msgpack_heap_delta_bytes( + buffer: &[u8], + ) -> Result> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + if frame.rows == 0 || frame.cols == 0 { + return Err(format!( + "CountSketchWithHeap delta frame has zero dims (rows={}, cols={})", + frame.rows, frame.cols + ) + .into()); + } + let mut acc = Self::new( + frame.rows as usize, + frame.cols as usize, + frame.heap_size as usize, + ); + acc.apply_msgpack_heap_delta_bytes(buffer)?; + Ok(acc) + } + + /// This function seems will never be used anymore. Keep it for possible future use. + pub fn deserialize_from_json(data: &Value) -> Result> { + let row_num = data["row_num"] + .as_f64() + .ok_or("Missing or invalid 'row_num' field")? as usize; + let col_num = data["col_num"] + .as_f64() + .ok_or("Missing or invalid 'col_num' field")? as usize; + let heap_size = data["heap_size"] + .as_f64() + .ok_or("Missing or invalid 'heap_size' field")? as usize; + + let sketch_data = data["sketch"] + .as_array() + .ok_or("Missing or invalid 'sketch' field")?; + + let mut sketch = Vec::new(); + for row in sketch_data { + let row_array = row.as_array().ok_or("Invalid row in sketch data")?; + let mut sketch_row = Vec::new(); + for cell in row_array { + let value = cell.as_f64().ok_or("Invalid cell value in sketch data")?; + sketch_row.push(value); + } + sketch.push(sketch_row); + } + + let topk_heap_data = data["topk_heap"] + .as_array() + .ok_or("Missing or invalid 'topk_heap' field")?; + + let mut topk_heap = Vec::new(); + for item in topk_heap_data { + let key = item["key"] + .as_str() + .ok_or("Missing or invalid 'key' in heap item")? + .to_string(); + let value = item["value"] + .as_f64() + .ok_or("Missing or invalid 'value' in heap item")?; + topk_heap.push(CmsHeapItem { key, value }); + } + + Ok(Self { + inner: CountMinSketchWithHeap::from_legacy_matrix( + sketch, topk_heap, row_num, col_num, heap_size, + ), + }) + } + + pub fn deserialize_from_bytes(_buffer: &[u8]) -> Result> { + Err("deserialize_from_bytes for CountMinSketchWithHeapAccumulator not implemented".into()) + } + + /// VALUE-WEIGHTED heavy-hitter update (FIX: CountSketch/CMS topk + /// recall-0). The default ingest path inserts `+1` per occurrence keyed + /// by the raw `item`, so the heap ranks groups by OCCURRENCE COUNT — the + /// wrong answer for `topk(k, sum by (label) (metric))`, which asks for + /// the top groups by SUM OF VALUE. This update adds the sample `value` + /// (not `+1`) into both the CMS matrix and the top-k heap, keyed by the + /// GROUP LABEL (e.g. the `host` / `zone` value), so the heap's ranking is + /// by summed value. Repeated calls for the same `group_label` accumulate, + /// so after folding a window the heap holds Σvalue per group. + /// + /// Delegates to the library's value-weighted `CountMinSketchWithHeap:: + /// update(key, value)` (`sketchlib_cms_heap_update` → `insert_many(key, + /// round(value))`), which is the "separate update path" the evaluation + /// plan (Fig 3c) called for. + pub fn insert_value(&mut self, group_label: &str, value: f64) { + self.inner.update(group_label, value); + } + + /// Read the top-`k` GROUPS ranked by summed VALUE (descending), keyed by + /// the group label. Pairs with [`Self::insert_value`]: the heap built by + /// value-weighted updates ranks by Σvalue, so this returns the + /// value-weighted top-k (not the occurrence-count top-k the raw `item` + /// heap would give). Sorted descending by value; ties broken by key for + /// determinism; truncated to `k`. + pub fn topk_by_value(&self, k: usize) -> Vec<(String, f64)> { + let mut items: Vec<(String, f64)> = self + .inner + .topk_heap_items() + .into_iter() + .map(|it| (it.key, it.value)) + .collect(); + items.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.0.cmp(&b.0)) + }); + items.truncate(k); + items + } + + /// Get all keys from the top-k heap. + pub fn get_topk_keys(&self) -> Vec { + self.inner + .topk_heap_items() + .iter() + .map(|item| { + let labels: Vec = item.key.split(';').map(|s| s.to_string()).collect(); + KeyByLabelValues { labels } + }) + .collect() + } +} + +impl SerializableToSink for CountMinSketchWithHeapAccumulator { + fn serialize_to_json(&self) -> Value { + let heap_items: Vec = self + .inner + .topk_heap_items() + .iter() + .map(|item| { + serde_json::json!({ + "key": item.key, + "value": item.value + }) + }) + .collect(); + + serde_json::json!({ + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), + "heap_size": self.inner.heap_size, + "sketch": self.inner.sketch_matrix(), + "topk_heap": heap_items + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountMinSketchWithHeapAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountMinSketchWithHeapAccumulator" + } + + /// Per-window base rotation (`docs/delta-baseline-contract.md` §3): + /// rebuild an empty heap accumulator with the same (rows, cols, + /// heap_size) so the next window's DELTA-HEAP frame applies onto a clean, + /// same-shape base. Without this override the trait default is a no-op, + /// which would let the additive matrix delta accumulate across windows + /// (over-counting). Mirrors `CountSketchAccumulator::reset_to_empty`. + fn reset_to_empty(&mut self) { + self.inner = + CountMinSketchWithHeap::new(self.inner.rows(), self.inner.cols(), self.inner.heap_size); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountMinSketchWithHeapAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_cms = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountMinSketchWithHeapAccumulator")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_cms.clone()])?; + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountMinSketchWithHeap + } + + fn get_keys(&self) -> Option> { + Some(self.get_topk_keys()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for CountMinSketchWithHeapAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for CountMinSketchWithHeapAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountMinSketchWithHeapAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_min_sketch_with_heap_creation() { + let cms = CountMinSketchWithHeapAccumulator::new(4, 1000, 20); + assert_eq!(cms.inner.rows(), 4); + assert_eq!(cms.inner.cols(), 1000); + assert_eq!(cms.inner.heap_size, 20); + assert_eq!(cms.inner.topk_heap_items().len(), 0); + } + + #[test] + fn test_count_min_sketch_with_heap_query() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 10, 5); + let key = KeyByLabelValues::new(); + assert_eq!(cms.query_key(&key), 0.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); + } + + #[test] + fn test_count_min_sketch_with_heap_merge() { + // Build controlled state via from_legacy_matrix (works regardless of backend config). + let sketch1 = vec![ + vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap1 = vec![ + CmsHeapItem { + key: "key1".to_string(), + value: 100.0, + }, + CmsHeapItem { + key: "key2".to_string(), + value: 50.0, + }, + ]; + let sketch2 = vec![ + vec![5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap2 = vec![ + CmsHeapItem { + key: "key3".to_string(), + value: 75.0, + }, + CmsHeapItem { + key: "key1".to_string(), + value: 80.0, + }, + ]; + + let cms1 = CountMinSketchWithHeapAccumulator { + inner: CountMinSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), + }; + let cms2 = CountMinSketchWithHeapAccumulator { + inner: CountMinSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), + }; + + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.sketch_matrix()[0][0], 15.0); + assert_eq!(merged.inner.sketch_matrix()[1][1], 35.0); + assert_eq!(merged.inner.heap_size, 3); + assert!(merged.inner.topk_heap_items().len() <= 3); + } + + #[test] + fn test_count_min_sketch_with_heap_merge_single() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms.clone()]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.rows(), cms.inner.rows()); + assert_eq!(merged.inner.cols(), cms.inner.cols()); + assert_eq!(merged.inner.heap_size, cms.inner.heap_size); + } + + #[test] + fn test_count_min_sketch_with_heap_merge_dimension_mismatch() { + let cms1 = CountMinSketchWithHeapAccumulator::new(2, 10, 5); + let cms2 = CountMinSketchWithHeapAccumulator::new(3, 10, 5); + let result = CountMinSketchWithHeapAccumulator::merge_accumulators(vec![cms1, cms2]); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("dimension")); + } + + #[test] + fn test_count_min_sketch_with_heap_as_aggregate_core() { + let cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + assert_eq!(cms.type_name(), "CountMinSketchWithHeapAccumulator"); + } + + #[test] + fn test_get_topk_keys() { + let mut cms = CountMinSketchWithHeapAccumulator::new(2, 3, 5); + cms.inner.update("label1;label2", 100.0); + cms.inner.update("label3;label4", 50.0); + + let keys = cms.get_topk_keys(); + assert_eq!(keys.len(), 2); + // Top-k order can differ between Legacy and Sketchlib backends (heap ordering / estimates). + let label_sets: std::collections::HashSet<_> = + keys.iter().map(|k| k.labels.clone()).collect(); + assert!(label_sets.contains(&vec!["label1".to_string(), "label2".to_string()])); + assert!(label_sets.contains(&vec!["label3".to_string(), "label4".to_string()])); + } + + #[test] + fn test_multiple_subpopulation_aggregate() { + let cms = CountMinSketchWithHeapAccumulator::new(3, 50, 10); + let key = KeyByLabelValues::new(); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cms; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, 0.0); + + let keys = multi_trait.get_keys(); + assert!(keys.is_some()); + assert_eq!(keys.unwrap().len(), 0); + } + + // ---------------------------------------------------------------- + // DELTA-HEAP wire form (encoding MSGPACK_DELTA): apply a sparse matrix + // delta + replace the heap, decoded generically (rmp_serde) WITHOUT any + // asap_sketchlib delta API. The first test feeds a frame produced by the + // Go encoder (sketchlib-go `MarshalCountSketchWithHeapDelta`) to prove + // cross-language byte parity — mirrors how the full-heap parity is + // proven. The second proves PWR full -> delta -> delta reconstruction. + // ---------------------------------------------------------------- + + /// Cross-language byte-parity: this hex is the exact output of + /// sketchlib-go's `asapmsgpack.MarshalCountSketchWithHeapDelta(5, 1024, + /// cells=[(0,1,50),(1,3,-4),(4,1023,1_000_000)], + /// heap=[("/checkout",50),("/cart",20)], heap_size=20)` (captured via a + /// throw-away Go print test, identical methodology to the full-heap + /// golden in `sketchlib-go/.../count_sketch_with_heap_test.go`). If the + /// Go encoder or the rmp_serde layout ever shifts, this decode fails + /// loudly. + const GO_DELTA_HEAP_GOLDEN_HEX: &str = "94c39305cd04009393000132930103fc9304cd03ffce000f42409292a92f636865636b6f7574cb404900000000000092a52f63617274cb403400000000000014"; + + #[test] + fn test_apply_go_produced_delta_heap_frame_matrix_and_heap() { + let bytes = hex::decode(GO_DELTA_HEAP_GOLDEN_HEX).expect("hex"); + + // Base = empty heap accumulator with the frame's dims (what the + // ingest caller holds after the per-window base rotation). + let mut acc = CountMinSketchWithHeapAccumulator::new(5, 1024, 20); + acc.apply_msgpack_heap_delta_bytes(&bytes) + .expect("apply Go delta-heap frame"); + + // Matrix: the three sparse cells landed onto the empty base. + let m = acc.inner.sketch_matrix(); + assert_eq!(m.len(), 5); + assert_eq!(m[0].len(), 1024); + assert_eq!(m[0][1], 50.0, "cell (0,1)"); + assert_eq!(m[1][3], -4.0, "cell (1,3)"); + assert_eq!(m[4][1023], 1_000_000.0, "cell (4,1023)"); + // Everything else stays zero. + assert_eq!(m[2][2], 0.0); + assert_eq!(m[0][0], 0.0); + + // Heap: the frame's full heap, with /checkout ranked above /cart. + let mut items = acc.inner.topk_heap_items(); + items.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap()); + assert_eq!(items.len(), 2); + assert_eq!(items[0].key, "/checkout"); + assert_eq!(items[0].value, 50.0); + assert_eq!(items[1].key, "/cart"); + assert_eq!(items[1].value, 20.0); + } + + #[test] + fn test_pwr_full_then_delta_then_delta_reconstructs_per_window() { + use asap_sketchlib::MessagePackCodec; + + // Window 1 (full frame): build a heap-bearing CountSketch with mass + // and serialize the FULL `{sketch,topk_heap,heap_size}` frame, then + // decode it into a heap accumulator (the cached per-series base). + let w1 = CountMinSketchWithHeap::from_legacy_matrix( + vec![vec![300.0; 4]; 5], + vec![CmsHeapItem { + key: "k".into(), + value: 300.0, + }], + 5, + 4, + 20, + ); + let w1_bytes = w1.to_msgpack().expect("w1 full msgpack"); + let mut base = CountMinSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(&w1_bytes) + .expect("decode w1 full frame as heap accumulator"); + assert_eq!(base.inner.sketch_matrix()[0][0], 300.0); + + // Window 2 delta: this window's own state is matrix cells of value 50 + // against an EMPTY base + heap {k:50}. The DELTA-HEAP frame is encoded + // the same way the Go producer does (4-array, is_delta, sparse cells). + let w2_frame = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + // PWR: rotate base to empty at the window boundary, then apply. + base.reset_to_empty(); + assert_eq!( + base.inner.sketch_matrix()[0][0], + 0.0, + "reset_to_empty cleared matrix" + ); + base.apply_msgpack_heap_delta_bytes(&w2_frame) + .expect("apply w2 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 50.0, "window-2 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 50.0); + // No cross-window leakage from window 1's 300s. + assert_eq!(base.inner.sketch_matrix()[2][2], 0.0); + let h2: Vec<_> = base.inner.topk_heap_items(); + assert_eq!(h2.len(), 1); + assert_eq!(h2[0].key, "k"); + assert_eq!(h2[0].value, 50.0); + + // Window 3 delta: 80s against empty + heap {k:80}. + let w3_frame = encode_delta_heap(5, 4, &[(0, 0, 80)], &[("k", 80.0)], 20); + base.reset_to_empty(); + base.apply_msgpack_heap_delta_bytes(&w3_frame) + .expect("apply w3 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 80.0, "window-3 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 0.0, "no window-2 leakage"); + let h3 = base.inner.topk_heap_items(); + assert_eq!(h3.len(), 1); + assert_eq!(h3[0].value, 80.0); + } + + #[test] + fn test_rmp_serde_layout_is_byte_identical_to_go_encoder() { + // The rmp_serde positional encoding of the delta-heap frame must be + // BYTE-IDENTICAL to sketchlib-go's hand-rolled + // `MarshalCountSketchWithHeapDelta`. This hex is the Go encoder's + // output for (5, 4, cells=[(0,0,50),(1,1,50)], heap=[("k",50)], + // heap_size=20) — the same inputs `encode_delta_heap` uses below. + // Equality here proves both encode AND decode are cross-language + // byte-compatible (the decode path is exercised by the Go-golden + // test above). + const GO_PARITY_HEX: &str = "94c39305049293000032930101329192a16bcb404900000000000014"; + let rust_bytes = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + assert_eq!(hex::encode(&rust_bytes), GO_PARITY_HEX); + } + + #[test] + fn test_apply_delta_rejects_full_frame_and_garbage() { + use asap_sketchlib::MessagePackCodec; + let mut acc = CountMinSketchWithHeapAccumulator::new(2, 4, 5); + // A FULL frame (3-array, no is_delta marker) must NOT decode as a + // delta — the routing relies on the two shapes being distinct. + let full = CountMinSketchWithHeap::from_legacy_matrix( + vec![vec![1.0; 4]; 2], + vec![CmsHeapItem { + key: "a".into(), + value: 1.0, + }], + 2, + 4, + 5, + ) + .to_msgpack() + .unwrap(); + assert!(acc.apply_msgpack_heap_delta_bytes(&full).is_err()); + assert!(acc.apply_msgpack_heap_delta_bytes(b"not msgpack").is_err()); + } + + /// Encode a DELTA-HEAP frame the same way sketchlib-go's + /// `MarshalCountSketchWithHeapDelta` does (rmp_serde positional layout), + /// so the test exercises the real decode path. Tuple structs serialize + /// as msgpack fixed arrays — byte-identical to the Go hand-rolled writer. + fn encode_delta_heap( + rows: u32, + cols: u32, + cells: &[(u32, u32, i64)], + heap: &[(&str, f64)], + heap_size: u64, + ) -> Vec { + #[derive(serde::Serialize)] + struct W<'a>( + bool, + (u32, u32, &'a [(u32, u32, i64)]), + Vec<(String, f64)>, + u64, + ); + let heap_owned: Vec<(String, f64)> = + heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); + let w = W(true, (rows, cols, cells), heap_owned, heap_size); + rmp_serde::to_vec(&w).expect("encode delta-heap") + } + + // ---------------------------------------------------------------- + // FIX 1 — VALUE-WEIGHTED top-k (recall 0 → correct). + // + // `topk(k, sum by (host) (cpu_load))` asks for the top-k hosts by + // SUM OF VALUE. The heavy-hitter heap built by the default `+1`-per- + // occurrence update ranks by COUNT keyed by `item`, so its recall + // against the value-weighted ground truth is 0 when the busiest host + // (most samples) is NOT the heaviest host (largest Σvalue). + // `insert_value(group_label, value)` adds the sample VALUE keyed by the + // GROUP LABEL, so `topk_by_value` ranks by Σvalue — correct recall. + // ---------------------------------------------------------------- + + /// Crafted adversarial dataset: the host with the MOST samples + /// (`h_chatty`, 100 tiny samples) is NOT the host with the largest + /// value-sum (`h_heavy`, a handful of huge samples). A COUNT-ranked + /// heap would surface `h_chatty`; the value-weighted top-k must surface + /// the true heavy hitters by Σvalue, giving recall 1.0 against the + /// ground-truth top-k-by-value-sum. + #[test] + fn value_weighted_topk_has_full_recall_vs_count_topk() { + // (host, per-sample value, sample count) → true Σvalue: + // h_heavy : 1000 × 3 = 3000 (few samples, huge value) + // h_mid : 200 × 5 = 1000 + // h_small : 50 × 6 = 300 + // h_chatty: 1 × 100 = 100 (MOST samples, tiny value) + let data: &[(&str, f64, usize)] = &[ + ("h_heavy", 1000.0, 3), + ("h_mid", 200.0, 5), + ("h_small", 50.0, 6), + ("h_chatty", 1.0, 100), + ]; + + // Wide CMS + heap large enough to hold every group exactly (4 groups) + // so the estimate equals the true Σvalue with no hash collisions. + let mut acc = CountMinSketchWithHeapAccumulator::new(5, 4096, 16); + let mut truth: std::collections::HashMap<&str, f64> = std::collections::HashMap::new(); + for (host, value, count) in data { + for _ in 0..*count { + acc.insert_value(host, *value); + } + *truth.entry(*host).or_insert(0.0) += value * (*count as f64); + } + + // Ground-truth top-2 by value-sum: h_heavy (3000), h_mid (1000). + let mut truth_ranked: Vec<(&str, f64)> = truth.into_iter().collect(); + truth_ranked.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap()); + let truth_top2: std::collections::HashSet<&str> = + truth_ranked.iter().take(2).map(|(k, _)| *k).collect(); + assert!( + truth_top2.contains("h_heavy") && truth_top2.contains("h_mid"), + "ground-truth top-2 by value-sum should be h_heavy + h_mid" + ); + + // Value-weighted top-2 from the heap. + let got = acc.topk_by_value(2); + assert_eq!(got.len(), 2, "k=2 → two groups: {got:?}"); + let got_keys: std::collections::HashSet<&str> = + got.iter().map(|(k, _)| k.as_str()).collect(); + + // RECALL = |got ∩ truth| / |truth| must be 1.0. + let hits = got_keys.intersection(&truth_top2).count(); + let recall = hits as f64 / truth_top2.len() as f64; + assert_eq!( + recall, 1.0, + "value-weighted top-k recall must be 1.0 (count-ranked heap would \ + surface h_chatty and miss h_heavy → recall < 1): got={got:?}" + ); + + // The busiest-by-count host (h_chatty) must NOT be in the top-2, + // proving we rank by value-sum, not occurrence count. + assert!( + !got_keys.contains("h_chatty"), + "h_chatty (most samples, smallest value-sum) must be excluded: {got:?}" + ); + + // Estimates are exact here (no collisions, heap holds all groups): + // top-1 must be h_heavy with Σvalue 3000. + assert_eq!(got[0].0, "h_heavy"); + assert!( + (got[0].1 - 3000.0).abs() < 1e-6, + "h_heavy value-sum estimate ≈ 3000, got {}", + got[0].1 + ); + assert_eq!(got[1].0, "h_mid"); + assert!( + (got[1].1 - 1000.0).abs() < 1e-6, + "h_mid value-sum estimate ≈ 1000, got {}", + got[1].1 + ); + } + + /// A single value-weighted insert must put the full value (not +1) into + /// the heap, and repeated inserts for the same group must accumulate. + #[test] + fn insert_value_accumulates_summed_value_in_heap() { + let mut acc = CountMinSketchWithHeapAccumulator::new(4, 1024, 8); + acc.insert_value("g", 10.0); + acc.insert_value("g", 25.0); + let top = acc.topk_by_value(1); + assert_eq!(top.len(), 1); + assert_eq!(top[0].0, "g"); + assert!( + (top[0].1 - 35.0).abs() < 1e-6, + "summed value should be 35 (10+25), got {}", + top[0].1 + ); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/count_sketch.rs b/crates/asap-physical-operators/src/summary_kernels/count_sketch.rs new file mode 100644 index 00000000..78c1f7d7 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/count_sketch.rs @@ -0,0 +1,678 @@ +//! CountSketch accumulator backed by `asap_sketchlib::CountSketch`. +//! +//! Supports worker merge, persistence serialization, and modified-OTLP proto +//! decoding. Per-key queries delegate to sketchlib's median-of-signed-rows +//! estimator so query and ingest use the same hash specification. Top-k +//! requires the separate heap-bearing accumulator. + +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountSketch, CountSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Count Sketch accumulator — inner matrix of signed counts. +#[derive(Debug, Clone)] +pub struct CountSketchAccumulator { + pub inner: CountSketch, +} + +impl CountSketchAccumulator { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + inner: CountSketch::new(row_num, col_num), + } + } + + /// Median-of-signed-rows point estimate for `key`, via the real + /// `asap_sketchlib::CountSketch::estimate` — the canonical, hash-spec- + /// compatible estimator (see `AggregateCore::query_statistic`'s doc for + /// why this replaced a hand-rolled, non-compatible hash). + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + self.inner.estimate(&key.to_semicolon_str()) + } + + /// Decode from the modified OTLP wire format's + /// `CountSketchDataPoint.sketch` bytes when + /// `encoding = COUNT_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `CountSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketch msgpack: {e}"))?, + }) + } + + /// Decode from the modified OTLP wire format's + /// `CountSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::CountSketchState` message + /// that DataCollector's `countsketchprocessor` emits when + /// `encoding = COUNT_SKETCH_ENCODING_PROTO`. + /// + /// Mirrors `CountMinSketchAccumulator::from_sketchlib_proto_bytes` + /// but on the signed-counter `CountSketchState`. The resulting + /// accumulator is constructed via + /// `CountSketch::from_legacy_matrix` after reshaping the flat + /// `counts_int` / `counts_float` field into a `Vec>`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountSketchState, CounterType, SketchEnvelope, + }; + use prost::Message; + + // DataCollector's countsketchprocessor wraps the state in a + // `SketchEnvelope{count_sketch: CountSketchState}` via + // sketchlib-go's `SerializePortableFO` + `proto.Marshal`. Try + // decoding as envelope first, fall back to bare + // `CountSketchState` for callers (e.g. unit tests) that + // encode the state directly. Mirrors the PR #14 fix on + // `CountMinSketchAccumulator::from_sketchlib_proto_bytes`. + let state = match SketchEnvelope::decode(buffer) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::CountSketch(st)) => st, + Some(other) => { + return Err(format!( + "SketchEnvelope contains non-CountSketch sketch: {:?}", + std::mem::discriminant(&other) + ) + .into()); + } + None => CountSketchState::decode(buffer) + .map_err(|e| format!("decode CountSketchState: {e}"))?, + }, + Err(_) => CountSketchState::decode(buffer) + .map_err(|e| format!("decode CountSketchState: {e}"))?, + }; + let rows = state.rows as usize; + let cols = state.cols as usize; + // Defensive dim validation BEFORE reconstructing the matrix: + // reject degenerate / narrow-hash-budget-violating / absurdly + // oversized dims so a malformed payload fails gracefully (the + // ingest caller skips the data point) instead of building a + // degenerate or huge matrix. Shares the CMS validator since the + // CountSketch matrix uses the same packed-hash column layout. + crate::summary_kernels::count_min_sketch::validate_sketch_dims( + "CountSketchState", + rows, + cols, + )?; + let expected_len = rows * cols; + let counter_type = CounterType::try_from(state.counter_type).map_err(|_| { + format!( + "CountSketchState has unknown counter_type tag {}", + state.counter_type + ) + })?; + let flat: Vec = match counter_type { + CounterType::Int32 | CounterType::Int64 => { + if state.counts_int.len() != expected_len { + return Err(format!( + "CountSketchState counts_int has {} entries, expected rows*cols = {}", + state.counts_int.len(), + expected_len + ) + .into()); + } + state.counts_int.iter().map(|&v| v as f64).collect() + } + CounterType::Float64 => { + if state.counts_float.len() != expected_len { + return Err(format!( + "CountSketchState counts_float has {} entries, expected rows*cols = {}", + state.counts_float.len(), + expected_len + ) + .into()); + } + state.counts_float.clone() + } + other => { + return Err(format!( + "CountSketchState counter_type {other:?} not yet supported \ + (INT128 stores interleaved hi/lo pairs; will be added when needed)" + ) + .into()); + } + }; + let mut matrix = Vec::with_capacity(rows); + for r in 0..rows { + let start = r * cols; + matrix.push(flat[start..start + cols].to_vec()); + } + Ok(Self { + inner: CountSketch::from_legacy_matrix(matrix, rows, cols), + }) + } + + /// Apply a proto-encoded `CountSketchDelta` frame to this + /// accumulator's inner sketch — the decode path for + /// `COUNT_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + /// + /// Cells apply additively: `matrix[cell_rows[i]][cell_cols[i]] + /// += d_counts[i]`. Per-row L2 is parsed off the wire but + /// ignored at application time — it's a downstream error- + /// accounting signal, not a merge input. + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode CountSketchDelta: {e}"))?; + + if pb.cell_rows.len() != pb.cell_cols.len() || pb.cell_rows.len() != pb.d_counts.len() { + return Err(format!( + "CountSketchDelta packed-array length mismatch: \ + cell_rows={}, cell_cols={}, d_counts={}", + pb.cell_rows.len(), + pb.cell_cols.len(), + pb.d_counts.len() + ) + .into()); + } + let cells = pb + .cell_rows + .iter() + .zip(pb.cell_cols.iter()) + .zip(pb.d_counts.iter()) + .map(|((r, c), dc)| (*r, *c, *dc)) + .collect(); + // This is the heap-less matrix kernel; ranked membership is handled + // by the explicit heap-bearing operator, not inferred from delta keys. + let delta = CountSketchDelta { + rows: pb.rows, + cols: pb.cols, + cells, + l2: pb.l2, + hh_keys: Vec::new(), + }; + self.inner + .apply_delta(&delta) + .map_err(|e| format!("apply CountSketchDelta: {e}"))?; + Ok(()) + } +} + +impl SerializableToSink for CountSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "row_num": self.inner.rows, + "col_num": self.inner.cols, + "sketch": self.inner.sketch(), + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountSketchAccumulator" + } + + /// Per-window base rotation: rebuild an empty signed-counter matrix + /// with the same (rows, cols) so the next window's additive cell + /// deltas align to the identical hash geometry. + fn reset_to_empty(&mut self) { + self.inner = CountSketch::new(self.inner.rows, self.inner.cols); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_cs = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountSketchAccumulator")?; + + let merged_inner = CountSketch::merge_refs(&[&self.inner, &other_cs.inner])?; + Ok(Box::new(Self { + inner: merged_inner, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountSketch + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &HashMap, + ) -> Result> { + use crate::Statistic; + // Key-provided path: route to MultipleSubpopulationAggregate::query + // (the canonical "what's the count of this key?" lookup), same + // pattern as CountMinSketchAccumulator. Fixed from a hand-rolled + // `DefaultHasher`-based estimator that did NOT use the sketchlib + // hash spec (its own doc admitted this — "not the sketchlib hash + // spec... the canonical compatibility path requires plumbing the + // sketchlib seeds through") — `asap_sketchlib::CountSketch::estimate` + // already hashes against the correct portable spec, so this is a + // genuine correctness fix, not just a refactor. + if let Some(key_val) = key.as_ref() { + return self.query(statistic, key_val, Some(query_kwargs)); + } + if let Some(k) = query_kwargs.get("key") { + let key_val = KeyByLabelValues::new_with_labels(vec![k.clone()]); + return self.query(statistic, &key_val, Some(query_kwargs)); + } + // No-key path: unchanged from before this fix -- CountSketch's + // signed rows have no CMS-style "min-row-sum = true total" + // property, so these are documented approximations, not a + // heavy-hitter answer. Not touched by this fix (only the + // key-provided path above had the hash-compatibility bug). + match statistic { + Statistic::Topk | Statistic::Count => { + let matrix = self.inner.sketch(); + let total: f64 = matrix.iter().flatten().map(|v| v.abs()).sum(); + let rows = matrix.len() as f64; + Ok(if rows > 0.0 { total / rows } else { 0.0 }) + } + Statistic::Sum => { + let matrix = self.inner.sketch(); + let total: f64 = matrix.iter().flatten().sum(); + let rows = matrix.len() as f64; + Ok(if rows > 0.0 { total / rows } else { 0.0 }) + } + other => Err(format!( + "CountSketchAccumulator: statistic {:?} not supported (only Topk / Count / Sum, with optional `key` in query_kwargs)", + other, + ) + .into()), + } + } +} + +impl MultipleSubpopulationAggregate for CountSketchAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountSketchAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_query_key_uses_real_sketchlib_estimator() { + // `query_key` must match sketchlib's estimator and hash specification. + let mut cs = CountSketchAccumulator::new(4, 1000); + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + cs.inner.update(&key.to_semicolon_str(), 10.0); + assert_eq!( + cs.query_key(&key), + cs.inner.estimate(&key.to_semicolon_str()) + ); + } + + #[test] + fn test_multiple_subpopulation_aggregate_query() { + let mut cs = CountSketchAccumulator::new(4, 1000); + let key = KeyByLabelValues::new_with_labels(vec!["checkout".to_string()]); + cs.inner.update(&key.to_semicolon_str(), 25.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, cs.query_key(&key)); + + // query_statistic (the AggregateCore entry point) must route a + // provided key through the same path. + let core: &dyn AggregateCore = &cs; + let via_core = core + .query_statistic(Statistic::Sum, &Some(key.clone()), &HashMap::new()) + .unwrap(); + assert_eq!(via_core, cs.query_key(&key)); + } + + #[test] + fn test_mergeable_accumulator_merge_accumulators() { + let cs1 = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![1.0, -2.0], vec![3.0, -4.0]], 2, 2), + }; + let cs2 = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![-1.0, 2.0], vec![-3.0, 4.0]], 2, 2), + }; + let merged = CountSketchAccumulator::merge_accumulators(vec![cs1, cs2]).unwrap(); + assert_eq!(merged.inner.sketch(), &vec![vec![0.0, 0.0], vec![0.0, 0.0]]); + } + + #[test] + fn test_mergeable_accumulator_rejects_empty() { + let result = CountSketchAccumulator::merge_accumulators(vec![]); + assert!(result.is_err()); + } + + fn encode_state( + rows: u32, + cols: u32, + counter_type: i32, + counts_int: Vec, + counts_float: Vec, + ) -> Vec { + use asap_sketchlib::proto::sketchlib::CountSketchState; + use prost::Message; + let state = CountSketchState { + rows, + cols, + counter_type, + counts_int, + counts_float, + l2: Vec::new(), + topk: None, + }; + state.encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_int64() { + use asap_sketchlib::proto::sketchlib::CounterType; + // Signed 2x3 matrix: row 0 = [1,-2,3], row 1 = [-4,5,-6] + let bytes = encode_state( + 2, + 3, + CounterType::Int64 as i32, + vec![1, -2, 3, -4, 5, -6], + Vec::new(), + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.0, -2.0, 3.0]); + assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's countsketchprocessor emits: + // the state wrapped in a `SketchEnvelope{count_sketch: ...}` + // via sketchlib-go's `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, CountSketchState, CounterType, SketchEnvelope, + }; + use prost::Message; + + let state = CountSketchState { + rows: 2, + cols: 3, + counter_type: CounterType::Int64 as i32, + counts_int: vec![1, -2, 3, -4, 5, -6], + counts_float: Vec::new(), + ..Default::default() + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::CountSketch(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.0, -2.0, 3.0]); + assert_eq!(matrix[1], vec![-4.0, 5.0, -6.0]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + // An envelope carrying a non-CountSketch sketch should be + // rejected with a clear error rather than silently producing + // garbage. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_float64() { + use asap_sketchlib::proto::sketchlib::CounterType; + let bytes = encode_state( + 2, + 2, + CounterType::Float64 as i32, + Vec::new(), + vec![1.5, -2.5, 3.5, -4.5], + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + let matrix = acc.inner.sketch(); + assert_eq!(matrix[0], vec![1.5, -2.5]); + assert_eq!(matrix[1], vec![3.5, -4.5]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_dimension_mismatch() { + use asap_sketchlib::proto::sketchlib::CounterType; + // 2x3 declared but only 5 int entries + let bytes = encode_state( + 2, + 3, + CounterType::Int64 as i32, + vec![1, 2, 3, 4, 5], + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!( + result.unwrap_err().to_string().contains("counts_int"), + "error should mention counts_int dim mismatch" + ); + } + + #[test] + fn test_from_sketchlib_proto_bytes_zero_dims_rejected() { + use asap_sketchlib::proto::sketchlib::CountSketchState; + use prost::Message; + let state = CountSketchState::default(); + let bytes = state.encode_to_vec(); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("degenerate dims")); + } + + #[test] + fn test_aggregate_core_merge_matches_matrix_add() { + let a = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![1.0, -2.0], vec![3.0, -4.0]], 2, 2), + }; + let b = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix(vec![vec![-1.0, 2.0], vec![-3.0, 4.0]], 2, 2), + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + let m = merged.inner.sketch(); + assert_eq!(m[0], vec![0.0, 0.0]); + assert_eq!(m[1], vec![0.0, 0.0]); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::summary_kernels::count_min_sketch::CountMinSketchAccumulator; + let cs = CountSketchAccumulator::new(2, 3); + let cms = CountMinSketchAccumulator::new(2, 3); + let result = cs.merge_with(&cms); + assert!(result.is_err()); + } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = CountSketch::from_legacy_matrix( + vec![vec![1.0, -2.0, 3.0], vec![-4.0, 5.0, -6.0]], + 2, + 3, + ); + let bytes = original.to_msgpack().unwrap(); + let acc = CountSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.rows, 2); + assert_eq!(acc.inner.cols, 3); + assert_eq!(acc.inner.sketch(), original.sketch()); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = CountSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } + + #[test] + fn test_apply_proto_delta_bytes_round_trip() { + use asap_sketchlib::proto::sketchlib::CountSketchDelta as PbDelta; + use prost::Message; + + let mut acc = CountSketchAccumulator { + inner: CountSketch::from_legacy_matrix( + vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]], + 2, + 3, + ), + }; + let bytes = PbDelta { + rows: 2, + cols: 3, + cell_rows: vec![0, 1], + cell_cols: vec![0, 2], + d_counts: vec![10, -6], + l2: vec![], + ..Default::default() + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!( + acc.inner.sketch(), + &vec![vec![11.0, 2.0, 3.0], vec![4.0, 5.0, 0.0]] + ); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = CountSketchAccumulator::new(2, 3); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ---------------------------------------------------------------- + // Defensive inbound-dimension validation (harden/sketch-dim-validation). + // Malformed / narrow-hash-budget-violating CountSketch dims must be + // rejected gracefully (Err, never a panic); valid configs the backend + // actually uses (5x2048, 5x4096, 5x2000) must still decode. + // ---------------------------------------------------------------- + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_bad_dims_no_panic() { + use asap_sketchlib::proto::sketchlib::CounterType; + // 5 * ceil(log2(8192))=5*13=65 > 64 — narrow-hash-budget violation. + // counts sized to rows*cols so rejection is on dims, not length. + let n = 5usize * 8192usize; + let bytes = encode_state( + 5, + 8192, + CounterType::Int64 as i32, + vec![0i64; n], + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "budget-violating dims should be rejected"); + assert!(result.unwrap_err().to_string().contains("rejecting")); + + // A valid neighbour (5x4096) on the same path still decodes fine. + let n_ok = 5usize * 4096usize; + let ok_bytes = encode_state( + 5, + 4096, + CounterType::Int64 as i32, + vec![0i64; n_ok], + Vec::new(), + ); + let acc = CountSketchAccumulator::from_sketchlib_proto_bytes(&ok_bytes) + .expect("valid 5x4096 CountSketch should still decode"); + assert_eq!(acc.inner.rows, 5); + assert_eq!(acc.inner.cols, 4096); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_oversized_dims() { + use asap_sketchlib::proto::sketchlib::CounterType; + // Declare 1 x 16,777,216 = 16M cells (> 8M cap) but send an empty + // counts vector: validation must reject on the dim cap BEFORE the + // decoder tries to allocate/reshape a 16M-entry matrix. (1 row keeps + // the hash budget tiny so the cap check, not the budget check, fires.) + let bytes = encode_state( + 1, + 16_777_216, + CounterType::Int64 as i32, + Vec::new(), + Vec::new(), + ); + let result = CountSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "oversized dims should be rejected"); + let msg = result.unwrap_err().to_string(); + assert!(msg.contains("cap"), "expected cell-cap error, got: {msg}"); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/count_sketch_with_heap.rs b/crates/asap-physical-operators/src/summary_kernels/count_sketch_with_heap.rs new file mode 100644 index 00000000..1f614bcb --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/count_sketch_with_heap.rs @@ -0,0 +1,575 @@ +//! Count Sketch with Heap accumulator — wraps +//! `asap_sketchlib::CountSketchWithHeap`. +//! +//! Port of `count_min_sketch_with_heap.rs` for the distinct +//! `CountSketchWithHeap` (median-of-signed-rows estimator) rather than +//! `CountMinSketchWithHeap` (min-over-rows estimator). The two are +//! different sketch algorithms that happen to share a storage shape and +//! wire layout -- see `asap_sketchlib::CountSketchWithHeap`'s own doc and +//! this session's `delta_apply.rs`/`decoders.rs` fix on the read side. +//! Before this file existed, `accumulator_factory.rs`'s raw-metric +//! ingest dispatch built a `CountMinSketchWithHeapAccumulator` (CMS math) +//! for `SketchAlgorithm::CountSketchWithHeap` sids -- the same conflation bug +//! already fixed on the read side, now closed on the write side too. + +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{CountSketchWithHeap, CsHeapItem, MessagePackCodec}; +use serde::Deserialize; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Local serde view of the DELTA-HEAP wire frame (encoding `MSGPACK_DELTA`). +/// Identical shape to `count_min_sketch_with_heap.rs`'s +/// `HeapDeltaWire`/`MatrixDeltaWire` -- the wire frame is generic (sparse +/// cell deltas + a full heap), not CMS-specific. See that file's doc for +/// the exact rmp_serde positional layout. +#[derive(Debug, Deserialize)] +struct HeapDeltaWire { + is_delta: bool, + matrix_delta: MatrixDeltaWire, + topk_heap: Vec<(String, f64)>, + #[allow(dead_code)] + heap_size: u64, +} + +#[derive(Debug, Deserialize)] +struct MatrixDeltaWire { + rows: u32, + cols: u32, + cells: Vec<(u32, u32, i64)>, +} + +/// Validated/flattened view of a decoded DELTA-HEAP frame. +struct HeapDeltaFrame { + rows: u32, + cols: u32, + heap_size: u64, + cells: Vec<(u32, u32, i64)>, + heap: Vec<(String, f64)>, +} + +impl HeapDeltaFrame { + fn from_msgpack(buffer: &[u8]) -> Result> { + let wire: HeapDeltaWire = rmp_serde::from_slice(buffer) + .map_err(|e| format!("decode CountSketchWithHeap delta msgpack: {e}"))?; + if !wire.is_delta { + return Err("CountSketchWithHeap delta frame has is_delta=false".into()); + } + Ok(Self { + rows: wire.matrix_delta.rows, + cols: wire.matrix_delta.cols, + heap_size: wire.heap_size, + cells: wire.matrix_delta.cells, + heap: wire.topk_heap, + }) + } +} + +/// Count Sketch with Heap accumulator — wraps `asap_sketchlib::CountSketchWithHeap`. +/// Core struct, update/merge/serde logic live in +/// `asap_sketchlib::message_pack_format::portable::countsketch_topk`. This +/// file retains QE-specific trait impls, legacy deserializers, and JSON +/// output -- same split as `CountMinSketchWithHeapAccumulator`. +#[derive(Debug, Clone)] +pub struct CountSketchWithHeapAccumulator { + pub inner: CountSketchWithHeap, +} + +impl CountSketchWithHeapAccumulator { + pub fn new(row_num: usize, col_num: usize, heap_size: usize) -> Self { + Self { + inner: CountSketchWithHeap::new(row_num, col_num, heap_size), + } + } + + pub fn query_key(&self, key: &KeyByLabelValues) -> f64 { + let key_string = key.labels.join(";"); + self.inner.estimate(&key_string) + } + + /// Decode a heap-bearing CountSketch FULL msgpack frame into a heap + /// accumulator -- the window-1 / full-frame base for the DELTA-HEAP + /// delta path. Mirrors `CountMinSketchWithHeapAccumulator::from_msgpack_with_heap_bytes`. + pub fn from_msgpack_with_heap_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: CountSketchWithHeap::from_msgpack(buffer) + .map_err(|e| format!("deserialize CountSketchWithHeap msgpack: {e}"))?, + }) + } + + /// Apply a DELTA-HEAP msgpack frame (encoding `MSGPACK_DELTA`) onto this + /// accumulator IN PLACE. Mirrors + /// `CountMinSketchWithHeapAccumulator::apply_msgpack_heap_delta_bytes` + /// exactly -- the frame decode/apply logic is generic, not tied to + /// which estimator the rebuilt sketch uses. + pub fn apply_msgpack_heap_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + + let rows = self.inner.rows(); + let cols = self.inner.cols(); + let heap_size = self.inner.heap_size; + + let mut matrix = self.inner.sketch_matrix(); + for (r, c, dc) in &frame.cells { + let (r, c) = (*r as usize, *c as usize); + if r >= rows || c >= cols { + continue; + } + matrix[r][c] += *dc as f64; + } + + let heap: Vec = frame + .heap + .into_iter() + .map(|(key, value)| CsHeapItem { key, value }) + .collect(); + + self.inner = CountSketchWithHeap::from_legacy_matrix(matrix, heap, rows, cols, heap_size); + Ok(()) + } + + /// Reconstruct a heap accumulator STANDALONE from a single DELTA-HEAP + /// msgpack frame, with no cached per-series base. Mirrors + /// `CountMinSketchWithHeapAccumulator::from_msgpack_heap_delta_bytes`. + pub fn from_msgpack_heap_delta_bytes( + buffer: &[u8], + ) -> Result> { + let frame = HeapDeltaFrame::from_msgpack(buffer)?; + if frame.rows == 0 || frame.cols == 0 { + return Err(format!( + "CountSketchWithHeap delta frame has zero dims (rows={}, cols={})", + frame.rows, frame.cols + ) + .into()); + } + let mut acc = Self::new( + frame.rows as usize, + frame.cols as usize, + frame.heap_size as usize, + ); + acc.apply_msgpack_heap_delta_bytes(buffer)?; + Ok(acc) + } + + /// Value-weighted heavy-hitter update -- see + /// `CountMinSketchWithHeapAccumulator::insert_value`'s doc for why + /// this (not a `+1`-per-occurrence update) is the correct semantics + /// for `topk(k, sum by (label) (metric))`-shaped queries. + pub fn insert_value(&mut self, group_label: &str, value: f64) { + self.inner.update(group_label, value); + } + + /// Read the top-`k` groups ranked by summed value (descending, tie-broken + /// by key for determinism). Mirrors `CountMinSketchWithHeapAccumulator::topk_by_value`. + pub fn topk_by_value(&self, k: usize) -> Vec<(String, f64)> { + let mut items: Vec<(String, f64)> = self + .inner + .topk_heap_items() + .into_iter() + .map(|it| (it.key, it.value)) + .collect(); + items.sort_by(|a, b| { + b.1.partial_cmp(&a.1) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| a.0.cmp(&b.0)) + }); + items.truncate(k); + items + } + + /// Get all keys from the top-k heap. + pub fn get_topk_keys(&self) -> Vec { + self.inner + .topk_heap_items() + .iter() + .map(|item| { + let labels: Vec = item.key.split(';').map(|s| s.to_string()).collect(); + KeyByLabelValues { labels } + }) + .collect() + } +} + +impl SerializableToSink for CountSketchWithHeapAccumulator { + fn serialize_to_json(&self) -> Value { + let heap_items: Vec = self + .inner + .topk_heap_items() + .iter() + .map(|item| { + serde_json::json!({ + "key": item.key, + "value": item.value + }) + }) + .collect(); + + serde_json::json!({ + "row_num": self.inner.rows(), + "col_num": self.inner.cols(), + "heap_size": self.inner.heap_size, + "sketch": self.inner.sketch_matrix(), + "topk_heap": heap_items + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for CountSketchWithHeapAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "CountSketchWithHeapAccumulator" + } + + /// Per-window base rotation -- mirrors + /// `CountMinSketchWithHeapAccumulator::reset_to_empty`. + fn reset_to_empty(&mut self) { + self.inner = + CountSketchWithHeap::new(self.inner.rows(), self.inner.cols(), self.inner.heap_size); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge CountSketchWithHeapAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_cs = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to CountSketchWithHeapAccumulator")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_cs.clone()])?; + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::CountSketchWithHeap + } + + fn get_keys(&self) -> Option> { + Some(self.get_topk_keys()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for CountSketchWithHeapAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for CountSketchWithHeapAccumulator { + fn query( + &self, + _statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Ok(self.query_key(key)) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for CountSketchWithHeapAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_count_sketch_with_heap_creation() { + let cs = CountSketchWithHeapAccumulator::new(4, 1000, 20); + assert_eq!(cs.inner.rows(), 4); + assert_eq!(cs.inner.cols(), 1000); + assert_eq!(cs.inner.heap_size, 20); + assert_eq!(cs.inner.topk_heap_items().len(), 0); + } + + #[test] + fn test_count_sketch_with_heap_query() { + let cs = CountSketchWithHeapAccumulator::new(2, 10, 5); + let key = KeyByLabelValues::new(); + assert_eq!(cs.query_key(&key), 0.0); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + assert_eq!(multi_trait.query(Statistic::Sum, &key, None).unwrap(), 0.0); + } + + #[test] + fn test_count_sketch_with_heap_merge() { + let sketch1 = vec![ + vec![10.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 20.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap1 = vec![ + CsHeapItem { + key: "key1".to_string(), + value: 100.0, + }, + CsHeapItem { + key: "key2".to_string(), + value: 50.0, + }, + ]; + let sketch2 = vec![ + vec![5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + vec![0.0, 15.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ]; + let heap2 = vec![ + CsHeapItem { + key: "key3".to_string(), + value: 75.0, + }, + CsHeapItem { + key: "key1".to_string(), + value: 80.0, + }, + ]; + + let cs1 = CountSketchWithHeapAccumulator { + inner: CountSketchWithHeap::from_legacy_matrix(sketch1, heap1, 2, 10, 5), + }; + let cs2 = CountSketchWithHeapAccumulator { + inner: CountSketchWithHeap::from_legacy_matrix(sketch2, heap2, 2, 10, 3), + }; + + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs1, cs2]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.sketch_matrix()[0][0], 15.0); + assert_eq!(merged.inner.sketch_matrix()[1][1], 35.0); + assert_eq!(merged.inner.heap_size, 3); + assert!(merged.inner.topk_heap_items().len() <= 3); + } + + #[test] + fn test_count_sketch_with_heap_merge_single() { + let cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs.clone()]); + assert!(result.is_ok()); + let merged = result.unwrap(); + assert_eq!(merged.inner.rows(), cs.inner.rows()); + assert_eq!(merged.inner.cols(), cs.inner.cols()); + assert_eq!(merged.inner.heap_size, cs.inner.heap_size); + } + + #[test] + fn test_count_sketch_with_heap_merge_dimension_mismatch() { + let cs1 = CountSketchWithHeapAccumulator::new(2, 10, 5); + let cs2 = CountSketchWithHeapAccumulator::new(3, 10, 5); + let result = CountSketchWithHeapAccumulator::merge_accumulators(vec![cs1, cs2]); + assert!(result.is_err()); + } + + #[test] + fn test_count_sketch_with_heap_as_aggregate_core() { + let cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + assert_eq!(cs.type_name(), "CountSketchWithHeapAccumulator"); + } + + #[test] + fn test_get_topk_keys() { + let mut cs = CountSketchWithHeapAccumulator::new(2, 3, 5); + cs.inner.update("label1;label2", 100.0); + cs.inner.update("label3;label4", 50.0); + + let keys = cs.get_topk_keys(); + assert_eq!(keys.len(), 2); + let label_sets: std::collections::HashSet<_> = + keys.iter().map(|k| k.labels.clone()).collect(); + assert!(label_sets.contains(&vec!["label1".to_string(), "label2".to_string()])); + assert!(label_sets.contains(&vec!["label3".to_string(), "label4".to_string()])); + } + + #[test] + fn test_multiple_subpopulation_aggregate() { + let cs = CountSketchWithHeapAccumulator::new(3, 50, 10); + let key = KeyByLabelValues::new(); + + let multi_trait: &dyn MultipleSubpopulationAggregate = &cs; + let result = multi_trait.query(Statistic::Sum, &key, None).unwrap(); + assert_eq!(result, 0.0); + + let keys = multi_trait.get_keys(); + assert!(keys.is_some()); + assert_eq!(keys.unwrap().len(), 0); + } + + #[test] + fn test_pwr_full_then_delta_then_delta_reconstructs_per_window() { + use asap_sketchlib::MessagePackCodec; + + let w1 = CountSketchWithHeap::from_legacy_matrix( + vec![vec![300.0; 4]; 5], + vec![CsHeapItem { + key: "k".into(), + value: 300.0, + }], + 5, + 4, + 20, + ); + let w1_bytes = w1.to_msgpack().expect("w1 full msgpack"); + let mut base = CountSketchWithHeapAccumulator::from_msgpack_with_heap_bytes(&w1_bytes) + .expect("decode w1 full frame as heap accumulator"); + assert_eq!(base.inner.sketch_matrix()[0][0], 300.0); + + let w2_frame = encode_delta_heap(5, 4, &[(0, 0, 50), (1, 1, 50)], &[("k", 50.0)], 20); + base.reset_to_empty(); + assert_eq!( + base.inner.sketch_matrix()[0][0], + 0.0, + "reset_to_empty cleared matrix" + ); + base.apply_msgpack_heap_delta_bytes(&w2_frame) + .expect("apply w2 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 50.0, "window-2 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 50.0); + assert_eq!(base.inner.sketch_matrix()[2][2], 0.0); + let h2: Vec<_> = base.inner.topk_heap_items(); + assert_eq!(h2.len(), 1); + assert_eq!(h2[0].key, "k"); + assert_eq!(h2[0].value, 50.0); + + let w3_frame = encode_delta_heap(5, 4, &[(0, 0, 80)], &[("k", 80.0)], 20); + base.reset_to_empty(); + base.apply_msgpack_heap_delta_bytes(&w3_frame) + .expect("apply w3 delta"); + assert_eq!(base.inner.sketch_matrix()[0][0], 80.0, "window-3 cell"); + assert_eq!(base.inner.sketch_matrix()[1][1], 0.0, "no window-2 leakage"); + let h3 = base.inner.topk_heap_items(); + assert_eq!(h3.len(), 1); + assert_eq!(h3[0].value, 80.0); + } + + #[test] + fn test_apply_delta_rejects_full_frame_and_garbage() { + use asap_sketchlib::MessagePackCodec; + let mut acc = CountSketchWithHeapAccumulator::new(2, 4, 5); + let full = CountSketchWithHeap::from_legacy_matrix( + vec![vec![1.0; 4]; 2], + vec![CsHeapItem { + key: "a".into(), + value: 1.0, + }], + 2, + 4, + 5, + ) + .to_msgpack() + .unwrap(); + assert!(acc.apply_msgpack_heap_delta_bytes(&full).is_err()); + assert!(acc.apply_msgpack_heap_delta_bytes(b"not msgpack").is_err()); + } + + fn encode_delta_heap( + rows: u32, + cols: u32, + cells: &[(u32, u32, i64)], + heap: &[(&str, f64)], + heap_size: u64, + ) -> Vec { + #[derive(serde::Serialize)] + struct W<'a>( + bool, + (u32, u32, &'a [(u32, u32, i64)]), + Vec<(String, f64)>, + u64, + ); + let heap_owned: Vec<(String, f64)> = + heap.iter().map(|(k, v)| (k.to_string(), *v)).collect(); + let w = W(true, (rows, cols, cells), heap_owned, heap_size); + rmp_serde::to_vec(&w).expect("encode delta-heap") + } + + #[test] + fn insert_value_accumulates_summed_value_in_heap() { + let mut acc = CountSketchWithHeapAccumulator::new(4, 1024, 8); + acc.insert_value("g", 10.0); + acc.insert_value("g", 25.0); + let top = acc.topk_by_value(1); + assert_eq!(top.len(), 1); + assert_eq!(top[0].0, "g"); + assert!( + (top[0].1 - 35.0).abs() < 1e-6, + "summed value should be 35 (10+25), got {}", + top[0].1 + ); + } + + /// The core proof this file exists at all: `CountSketchWithHeapAccumulator` + /// wraps the real, distinct `asap_sketchlib::CountSketchWithHeap` -- + /// not the CMS-family `CountMinSketchWithHeap` a collapsed dispatch + /// used to substitute (the exact bug this file fixes on the ingest + /// side, mirroring the already-fixed read side). Two different Rust + /// types means `merge_with` rejects mixing them at the type-check + /// level, same as any other mismatched-family merge attempt -- + /// verified directly rather than via a numeric estimate comparison + /// (asap_sketchlib's own test suite already proves the median vs + /// min-over-rows divergence at the sketch-math level). + #[test] + fn test_rejects_merge_with_cms_family_accumulator() { + use crate::summary_kernels::count_min_sketch_with_heap::CountMinSketchWithHeapAccumulator; + + let cs = CountSketchWithHeapAccumulator::new(4, 64, 10); + let cms = CountMinSketchWithHeapAccumulator::new(4, 64, 10); + let result = cs.merge_with(&cms); + assert!( + result.is_err(), + "CountSketchWithHeapAccumulator must not merge with CountMinSketchWithHeapAccumulator \ + -- different algorithms sharing only a storage shape" + ); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/datasketches_kll.rs b/crates/asap-physical-operators/src/summary_kernels/datasketches_kll.rs new file mode 100644 index 00000000..1ab4df79 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/datasketches_kll.rs @@ -0,0 +1,727 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use asap_sketchlib::{KllSketch, MessagePackCodec}; +use base64::{engine::general_purpose, Engine as _}; +use serde_json::Value; +use std::collections::HashMap; +#[cfg(feature = "extra_debugging")] +use std::time::Instant; +use tracing::debug; + +use crate::Statistic; + +/// KLL sketch accumulator — wraps asap_sketchlib::KllSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. +/// This file retains QE-specific trait impls and JSON output. +pub struct DatasketchesKLLAccumulator { + pub inner: KllSketch, +} + +impl DatasketchesKLLAccumulator { + pub fn new(k: u16) -> Self { + Self { + inner: KllSketch::new(k), + } + } + + pub fn update(&mut self, value: f64) { + self.inner.update(value); + } + + pub fn get_quantile(&self, quantile: f64) -> f64 { + self.inner.quantile(quantile) + } + + /// Decode from the modified OTLP wire format's + /// `KLLSketchDataPoint.sketch` bytes when + /// `encoding = KLL_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `KllSketch` struct — PR I parity entrypoint. Unlike the + /// `_ENCODING_PROTO` path (which does lossy statistical + /// reconstruction via `update()` replay), the msgpack path is a + /// bit-identical round-trip because sketch-core's `KllSketch` + /// serializes its full internal state to msgpack. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: KllSketch::from_msgpack(buffer) + .map_err(|e| -> Box { e.to_string().into() })?, + }) + } + + /// Decode from the modified OTLP wire format's + /// `KLLSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::KllState` message that + /// DataCollector's `kllprocessor` emits when + /// `encoding = KLL_SKETCH_ENCODING_PROTO`. + /// + /// The neutral codec decodes the sketchlib envelope. + /// The level-aware constructor below preserves the supplied retained + /// sample layout without replaying updates. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + let state = asap_sketch_codec::kll_state(buffer)?; + if state.k < 8 { + return Err(format!("KllState.k must be >= 8 (got {})", state.k).into()); + } + if state.k > u16::MAX as u32 { + return Err(format!( + "KllState.k does not fit in u16 (got {}, max {})", + state.k, + u16::MAX + ) + .into()); + } + // Validate the levels[] boundary array if it is populated. The + // proto contract says `levels[0] == 0` and + // `levels[num_levels] == items.len()`. If the producer left + // levels empty (common when num_levels is zero), skip. + if !state.levels.is_empty() { + if state.levels.len() as u32 != state.num_levels + 1 { + return Err(format!( + "KllState levels length = {}, expected num_levels+1 = {}", + state.levels.len(), + state.num_levels + 1 + ) + .into()); + } + if state.levels[0] != 0 { + return Err(format!("KllState.levels[0] = {}, expected 0", state.levels[0]).into()); + } + if *state.levels.last().unwrap() as usize != state.items.len() { + return Err(format!( + "KllState.levels[{}] = {}, expected items.len() = {}", + state.num_levels, + state.levels.last().unwrap(), + state.items.len() + ) + .into()); + } + } + let k = state.k as u16; + // Direct, bit-exact reconstruction from the portable state (no per-item + // `update()` replay) whenever the producer supplied the `levels[]` + // boundary array — which it does for any non-empty sketch. Falls back to + // the statistical replay only when `levels` is absent (empty sketch). + if !state.levels.is_empty() { + // KllState is highest-level first; the in-memory constructor + // expects L0 first. Replaying or copying the wire order changes + // retained-item weights after the first compaction. + let mut items = Vec::with_capacity(state.items.len()); + let mut levels = vec![0]; + if state + .levels + .windows(2) + .any(|bounds| bounds[0] > bounds[1] || bounds[1] as usize > state.items.len()) + { + return Err("KllState levels must be monotonic and within items".into()); + } + for bounds in state.levels.windows(2).rev() { + items.extend_from_slice(&state.items[bounds[0] as usize..bounds[1] as usize]); + levels.push(items.len()); + } + return Ok(Self { + inner: KllSketch::from_portable_state( + k, + &items, + &levels, + state.num_levels as usize, + ) + .map_err(|e| -> Box { e.into() })?, + }); + } + let mut acc = Self::new(k); + for item in &state.items { + acc.update(*item); + } + Ok(acc) + } + + /// Merge multiple accumulators efficiently without cloning all of them. + pub fn merge_multiple( + accumulators: &[Box], + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut kll_accumulators = Vec::with_capacity(accumulators.len()); + for acc in accumulators { + if acc.get_accumulator_type() != AggregationType::DatasketchesKLL { + return Err(format!( + "Cannot merge DatasketchesKLLAccumulator with {:?}", + acc.get_accumulator_type() + ) + .into()); + } + let kll_acc = acc + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to DatasketchesKLLAccumulator")?; + kll_accumulators.push(kll_acc); + } + + let inner_refs: Vec<&KllSketch> = kll_accumulators.iter().map(|acc| &acc.inner).collect(); + let merged_inner = KllSketch::merge_refs(&inner_refs)?; + Ok(Self { + inner: merged_inner, + }) + } +} + +// Manual trait implementations since the C++ library doesn't provide them +impl Clone for DatasketchesKLLAccumulator { + fn clone(&self) -> Self { + Self { + inner: self.inner.clone(), + } + } +} + +impl std::fmt::Debug for DatasketchesKLLAccumulator { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DatasketchesKLLAccumulator") + .field("k", &self.inner.k) + .field("sketch_n", &self.inner.count()) + .finish() + } +} + +// TODO: verify this +// Thread safety: The C++ library is not thread-safe by default, but since we're using it +// in a single-threaded context per accumulator instance and only sharing read-only operations, +// this should be safe. +unsafe impl Send for DatasketchesKLLAccumulator {} +unsafe impl Sync for DatasketchesKLLAccumulator {} + +impl SerializableToSink for DatasketchesKLLAccumulator { + fn serialize_to_json(&self) -> Value { + // Mirror Python implementation: {"sketch": base64_encoded_string} + let sketch_bytes = self.inner.sketch_bytes(); + let sketch_b64 = general_purpose::STANDARD.encode(&sketch_bytes); + serde_json::json!({ "sketch": sketch_b64 }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for DatasketchesKLLAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "DatasketchesKLLAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + #[cfg(feature = "extra_debugging")] + let merge_with_start = Instant::now(); + #[cfg(feature = "extra_debugging")] + debug!( + "[PERF] DatasketchesKLLAccumulator::merge_with() started - self.k={}, self.n={}", + self.inner.k, + self.inner.count() + ); + + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge DatasketchesKLLAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_kll = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to DatasketchesKLLAccumulator")?; + + let merged_inner = KllSketch::merge_refs(&[&self.inner, &other_kll.inner])?; + let merged = Self { + inner: merged_inner, + }; + + #[cfg(feature = "extra_debugging")] + debug!( + "[PERF] DatasketchesKLLAccumulator::merge_with() TOTAL TIME: {:?}", + merge_with_start.elapsed() + ); + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::DatasketchesKLL + } + + fn approx_memory_bytes(&self) -> usize { + // KLL with default k=200 holds ~2*k items (~3 KiB). Round up + // for overhead. + 4 * 1024 + } + + fn aux_stats(&self) -> AuxStats { + // KLL natively tracks `count` (n, samples observed). min/max + // are available from the underlying sketch but only via a + // O(k) quantile extraction at quantile=0/1, which is not + // a cheap trait-method call. sum is not retained by KLL. + // + // Surface only count here; follow-up PR may add min/max via a + // dedicated accessor on sketch-core. `sum_over_time` queries + // on KLL fall back to query_statistic as they do today. + AuxStats { + count: Some(self.inner.count()), + ..AuxStats::empty() + } + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query(statistic, Some(query_kwargs)) + } +} + +impl SingleSubpopulationAggregate for DatasketchesKLLAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Quantile => { + debug!( + "Querying DatasketchesKLLAccumulator for quantile with kwargs: {:?}", + query_kwargs + ); + let quantile = query_kwargs + .and_then(|kwargs| kwargs.get("quantile")) + .ok_or("Missing quantile parameter for quantile query")? + .parse::() + .map_err(|_| "Invalid quantile parameter format")?; + + if !(0.0..=1.0).contains(&quantile) { + return Err("Quantile must be between 0.0 and 1.0".into()); + } + + Ok(self.get_quantile(quantile)) + } + _ => Err( + format!("Unsupported statistic in DatasketchesKLLAccumulator: {statistic:?}") + .into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for DatasketchesKLLAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use prost::Message; + + fn encode_state(state: asap_sketchlib::proto::sketchlib::KllState) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + } + .encode_to_vec() + } + + #[test] + fn test_datasketches_kll_creation() { + let kll = DatasketchesKLLAccumulator::new(200); + assert!(kll.inner.count() == 0); + assert_eq!(kll.inner.k, 200); + } + + #[test] + fn test_datasketches_kll_update() { + let mut kll = DatasketchesKLLAccumulator::new(200); + kll.update(10.0); + kll.update(20.0); + kll.update(15.0); + assert_eq!(kll.inner.count(), 3); + } + + #[test] + fn test_datasketches_kll_quantile() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + assert_eq!(kll.get_quantile(0.0), 1.0); + assert_eq!(kll.get_quantile(1.0), 10.0); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + let q50 = kll.get_quantile(0.5); + assert!((q50 - 6.0).abs() <= 1.0, "expected median ~6, got {q50}"); + } + + #[test] + fn test_datasketches_kll_query() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + + let mut query_kwargs = HashMap::new(); + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + assert!( + (result - 6.0).abs() <= 1.0, + "expected median ~6, got {result}" + ); + + assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); + } + + #[test] + fn test_datasketches_kll_merge() { + let mut kll1 = DatasketchesKLLAccumulator::new(200); + let mut kll2 = DatasketchesKLLAccumulator::new(200); + + for i in 1..=5 { + kll1.update(i as f64); + } + for i in 6..=10 { + kll2.update(i as f64); + } + + let merged = DatasketchesKLLAccumulator::merge_accumulators(vec![kll1, kll2]).unwrap(); + assert_eq!(merged.inner.count(), 10); + assert_eq!(merged.get_quantile(0.0), 1.0); + assert_eq!(merged.get_quantile(1.0), 10.0); + } + + #[test] + fn test_datasketches_kll_get_keys() { + let kll = DatasketchesKLLAccumulator::new(200); + assert_eq!(kll.type_name(), "DatasketchesKLLAccumulator"); + } + + #[test] + fn test_trait_object() { + let mut kll = DatasketchesKLLAccumulator::new(200); + kll.update(5.0); + let trait_obj: Box = Box::new(kll); + assert_eq!(trait_obj.type_name(), "DatasketchesKLLAccumulator"); + } + + #[test] + fn test_datasketches_kll_query_with_kwargs() { + let mut kll = DatasketchesKLLAccumulator::new(200); + for i in 1..=10 { + kll.update(i as f64); + } + + let mut query_kwargs = HashMap::new(); + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.5 quantile of 1..10 may be 5, 6, or 7. + assert!( + (result - 6.0).abs() <= 1.0, + "expected median ~6, got {result}" + ); + + query_kwargs.insert("quantile".to_string(), "0.9".to_string()); + let result = kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(); + // Sketchlib KLL is approximate; 0.9 quantile of 1..10 may be 9 or 10. + assert!( + (9.0..=10.0).contains(&result), + "expected 0.9 quantile in [9,10], got {result}" + ); + + query_kwargs.insert("quantile".to_string(), "0.0".to_string()); + assert_eq!( + kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), + 1.0 + ); + + query_kwargs.insert("quantile".to_string(), "1.0".to_string()); + assert_eq!( + kll.query(Statistic::Quantile, Some(&query_kwargs)).unwrap(), + 10.0 + ); + + assert!(kll.query(Statistic::Quantile, None).is_err()); + + query_kwargs.insert("quantile".to_string(), "invalid".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "1.5".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "-0.1".to_string()); + assert!(kll.query(Statistic::Quantile, Some(&query_kwargs)).is_err()); + + query_kwargs.insert("quantile".to_string(), "0.5".to_string()); + assert!(kll.query(Statistic::Sum, Some(&query_kwargs)).is_err()); + } + + #[test] + fn test_datasketches_kll_merge_multiple() { + let mut kll1 = DatasketchesKLLAccumulator::new(200); + let mut kll2 = DatasketchesKLLAccumulator::new(200); + let mut kll3 = DatasketchesKLLAccumulator::new(200); + + for i in 1..=5 { + kll1.update(i as f64); + } + for i in 6..=10 { + kll2.update(i as f64); + } + for i in 11..=15 { + kll3.update(i as f64); + } + + let boxed_accs: Vec> = + vec![Box::new(kll1), Box::new(kll2), Box::new(kll3)]; + + let merged = DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).unwrap(); + assert_eq!(merged.inner.count(), 15); + assert_eq!(merged.get_quantile(0.0), 1.0); + assert_eq!(merged.get_quantile(1.0), 15.0); + assert_eq!(merged.get_quantile(0.5), 8.0); + } + + #[test] + fn test_datasketches_kll_merge_multiple_error_cases() { + let empty: Vec> = vec![]; + assert!(DatasketchesKLLAccumulator::merge_multiple(&empty).is_err()); + + let kll1 = DatasketchesKLLAccumulator::new(200); + let kll2 = DatasketchesKLLAccumulator::new(100); + let boxed_accs: Vec> = vec![Box::new(kll1), Box::new(kll2)]; + assert!(DatasketchesKLLAccumulator::merge_multiple(&boxed_accs).is_err()); + + use crate::summary_kernels::sum::SumAccumulator; + let kll = DatasketchesKLLAccumulator::new(200); + let sum = SumAccumulator::new(); + let mixed_accs: Vec> = vec![Box::new(kll), Box::new(sum)]; + assert!(DatasketchesKLLAccumulator::merge_multiple(&mixed_accs).is_err()); + } + + #[test] + fn test_from_sketchlib_proto_bytes_reconstructs_quantiles() { + // Build a KllState with 64 items in level order; the decoder + // replays every item through `update()` so the reconstructed + // sketch is statistically equivalent — quantile estimates + // match the ground truth (sorted items) within KLL's own + // rank-error bound for k=200. + use asap_sketchlib::proto::sketchlib::KllState; + + let items: Vec = (0..64).map(|i| i as f64).collect(); + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 64], + items: items.clone(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + + let acc = + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.count(), 64); + // For 64 values 0..63, the true median is 31.5 and quantile + // error is ~1% × range = 0.63. KLL's own point query can + // legally be off by up to ε × N ~= 0.01 × 64 = 0.64. Allow a + // generous tolerance since the important invariant is "the + // decoded sketch is queryable and returns a sensible value". + let median = acc.get_quantile(0.5); + assert!( + (median - 31.5).abs() <= 10.0, + "reconstructed median {median} is outside tolerance of true median 31.5" + ); + let q01 = acc.get_quantile(0.01); + let q99 = acc.get_quantile(0.99); + assert!( + q01 <= q99, + "quantile monotonicity violated: q01={q01}, q99={q99}" + ); + } + + // Compacted portable state is highest-level first, unlike the runtime buffer. + #[test] + fn compacted_wire_state_preserves_count_and_quantiles() { + use asap_sketchlib::{proto::sketchlib::KllState, sketches::KLL}; + let mut source = KLL::::init_kll_with_seed(32, 123); + for i in 0..1000 { + source.update(&(((i * 7919 + 17) % 1009) as f64 / 1009.0)); + } + assert!(source.wire_num_levels() > 1); + let state = KllState { + k: 32, + m: source.wire_m(), + num_levels: source.wire_num_levels(), + levels: source.wire_levels(), + items: source.wire_items(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: vec![], + }; + let decoded = + DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&encode_state(state)).unwrap(); + assert_eq!(decoded.inner.count(), source.count() as u64); + for q in [0.0, 0.1, 0.5, 0.9, 1.0] { + assert_eq!(decoded.inner.quantile(q), source.quantile(q), "q={q}"); + } + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's kllprocessor emits: the state + // wrapped in a `SketchEnvelope{kll: ...}` via sketchlib-go's + // `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + + let items: Vec = (0..64).map(|i| i as f64).collect(); + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 64], + items, + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + assert_eq!(acc.inner.count(), 64); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, CountMinState, SketchEnvelope}; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::CountMin( + CountMinState::default(), + )), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_small_k() { + use asap_sketchlib::proto::sketchlib::KllState; + let state = KllState { + k: 4, // < minimum of 8 + m: 2, + num_levels: 0, + levels: Vec::new(), + items: Vec::new(), + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("k must be >= 8")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_inconsistent_levels() { + use asap_sketchlib::proto::sketchlib::KllState; + // num_levels=1 but levels array has 3 entries instead of 2 + let state = KllState { + k: 200, + m: 8, + num_levels: 1, + levels: vec![0, 5, 10], + items: vec![1.0, 2.0, 3.0, 4.0, 5.0], + coin: None, + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + }; + let bytes = encode_state(state); + let result = DatasketchesKLLAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("levels length")); + } + + #[test] + fn aux_stats_exposes_count_via_kll_n() { + let mut acc = DatasketchesKLLAccumulator::new(200); + for i in 0..50 { + acc.update(i as f64); + } + let aux = acc.aux_stats(); + assert_eq!(aux.count, Some(50)); + // KLL doesn't natively expose min/max cheaply and doesn't + // track sum at all — those fields must be None so callers + // fall through to query_statistic. + assert_eq!(aux.sum, None); + assert_eq!(aux.min, None); + assert_eq!(aux.max, None); + } + + #[test] + fn aux_stats_empty_kll_has_zero_count() { + let acc = DatasketchesKLLAccumulator::new(200); + assert_eq!(acc.aux_stats().count, Some(0)); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/dd_sketch.rs b/crates/asap-physical-operators/src/summary_kernels/dd_sketch.rs new file mode 100644 index 00000000..1c25d940 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/dd_sketch.rs @@ -0,0 +1,726 @@ +//! DDSketch accumulator — wraps `asap_sketchlib::DdSketch`. +//! +//! Concrete accumulator reached from the modified-OTLP +//! `Metric.data = DDSketch{…}` hot path (PR C-CountSketch follow-up). +//! Merge via bucket-index alignment on the inner sketch, serialize as +//! MessagePack for the sink, and decode from the sketchlib +//! `DDSketchState` proto. +//! +//! Query semantics follow the STRICT policy after the DataPoint-level +//! METRIC scalars were dropped from the wire format +//! (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57): the sketch serves +//! Quantile (log-bucket estimation) and Count (sum of bucket counts). +//! Sum/Min/Max are no longer derivable from the wire bytes and are +//! served by controller-provisioned exact aggregations — `query_statistic` +//! returns the unavailable-statistic error for them. + +use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::{DdSketch, DdSketchDelta, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +/// DDSketch accumulator — inner log-bucketed sketch. +#[derive(Debug, Clone)] +pub struct DDSketchAccumulator { + pub inner: DdSketch, + /// Edge sampling probability `p ∈ (0,1]` carried on the producer's + /// `SketchEnvelope.sample_p`. The edge admits each value with probability + /// `p` (NitroSketch geometric skip), so `inner.total_count()` is ~`p`× the + /// true count and a `Count` query must rescale by `1/p`. Quantiles are + /// rank-preserving and need NO rescale. `1.0` (and the proto3 default `0.0`, + /// dual-read as `1.0`) means no sampling, so the rescale is a no-op and the + /// behaviour is identical to before. The factor is a per-series config + /// constant: it is set from the first (always-full, otel.rs ingest + /// contract) frame and preserved across delta applies, window-boundary + /// `reset_to_empty`, and `merge_with`. + pub sample_p: f64, +} + +/// Normalize a wire `sample_p` to a usable rescale denominator. `0.0` (proto3 +/// default), `>= 1.0`, and non-finite all collapse to `1.0` (no sampling), so a +/// `Count` rescale by `1/p` is a no-op on unsampled / legacy frames. +pub(crate) fn normalize_sample_p(p: f64) -> f64 { + if p.is_finite() && p > 0.0 && p < 1.0 { + p + } else { + 1.0 + } +} + +impl DDSketchAccumulator { + pub fn new(alpha: f64) -> Self { + Self { + inner: DdSketch::new(alpha), + sample_p: 1.0, + } + } + + /// Read the normalized edge sampling probability from a full-frame + /// `SketchEnvelope`'s `sample_p`. Returns `1.0` (no sampling) for bare + /// `DdSketchState` bytes or any decode failure — the primary production + /// decode path (`reconstruct_via_runtime`) discards the envelope's + /// `sample_p`, so the ingest call site re-reads it from the same bytes. + pub fn sample_p_from_envelope_bytes(buffer: &[u8]) -> f64 { + use asap_sketchlib::proto::sketchlib::SketchEnvelope; + use prost::Message; + SketchEnvelope::decode(buffer) + .map(|env| normalize_sample_p(env.sample_p)) + .unwrap_or(1.0) + } + + /// Decode from the modified OTLP wire format's + /// `DDSketchDataPoint.sketch` bytes when + /// `encoding = DDSKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `DdSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: DdSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize DdSketch msgpack: {e}"))?, + // The msgpack DdSketch struct carries no envelope/sample_p; the + // msgpack path is parity/test-only and is never edge-sampled. + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `DDSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::DDSketchState` message that + /// DataCollector's `ddsketchprocessor` emits when + /// `encoding = DD_SKETCH_ENCODING_PROTO`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + let (state, sample_p) = asap_sketch_codec::ddsketch_state(buffer)?; + if !(state.alpha > 0.0 && state.alpha < 1.0) { + return Err(format!( + "DDSketchState alpha {} out of range (expected 0 < alpha < 1)", + state.alpha + ) + .into()); + } + // Preserve positive, negative and zero stores from the sketchlib wire state. + let inner = DdSketch::from_proto(state); + Ok(Self { + inner, + sample_p: normalize_sample_p(sample_p), + }) + } + + /// Apply a proto-encoded `DDSketchDelta` frame to this + /// accumulator's inner sketch — the decode path for + /// `DD_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + /// + /// Called against an accumulator that already carries the base + /// sketch state; the caller is the per-series snapshot cache in + /// the ingest path. Bytes are the + /// `asap_sketchlib::proto::sketchlib::DdSketchDelta` message. + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + use asap_sketchlib::proto::sketchlib::DdSketchDelta as PbDelta; + use prost::Message; + + let pb = PbDelta::decode(buffer).map_err(|e| format!("decode DDSketchDelta: {e}"))?; + + // The delta no longer carries d_count/d_sum/min/max + // (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). Apply the + // bucket deltas only; `DdSketch` recomputes its total count from + // the merged bucket counts (`total_count()`). + let buckets = pb + .buckets + .into_iter() + .map(|b| (b.index, b.d_count)) + .collect(); + let delta = DdSketchDelta { + buckets, + negative_buckets: pb + .negative_buckets + .into_iter() + .map(|b| (b.index, b.d_count)) + .collect(), + zero_count: pb.zero_count, + ..Default::default() + }; + self.inner + .apply_delta(&delta) + .map_err(|error| format!("apply DDSketchDelta: {error}"))?; + Ok(()) + } +} + +impl SerializableToSink for DDSketchAccumulator { + fn serialize_to_json(&self) -> Value { + // The DataPoint-level scalars (sum/min/max) are no longer carried + // by `DdSketch` (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). + // `count` is the bucket-derived total via `total_count()`. + serde_json::json!({ + "alpha": self.inner.alpha, + "store_offset": self.inner.store_offset, + "bucket_count": self.inner.store_counts.len(), + // Raw bucket-derived count (admitted samples). `sample_p` is the + // scale factor a consumer applies (count / sample_p) to estimate + // the true count; `query_statistic(Count)` already does this. + "count": self.inner.total_count(), + "sample_p": self.sample_p, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for DDSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "DDSketchAccumulator" + } + + /// Per-window base rotation: drop all bucket counts but keep the + /// relative-accuracy parameter so the next window's bucket deltas + /// index into the same log-bucket layout. `sample_p` is a per-series + /// config constant (not per-window data), so it is intentionally + /// preserved across the rotation — the next window's deltas are sampled + /// at the same rate and must rescale identically. + fn reset_to_empty(&mut self) { + self.inner = DdSketch::new(self.inner.alpha); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge DDSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_dd = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to DDSketchAccumulator")?; + let merged_inner = DdSketch::merge_refs(&[&self.inner, &other_dd.inner])?; + // sample_p is a per-series config constant, so both operands carry the + // same value in practice. Prefer a sampled factor over the no-sampling + // default so a merge with a freshly-reset (1.0) base keeps the series' + // sampling rate. + let sample_p = if self.sample_p < 1.0 { + self.sample_p + } else { + other_dd.sample_p + }; + Ok(Box::new(Self { + inner: merged_inner, + sample_p, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::DDSketch + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + query_kwargs: &HashMap, + ) -> Result> { + use crate::Statistic; + + match statistic { + Statistic::Quantile => { + // PromQL `histogram_quantile(q, …)` and + // `quantile_over_time(q, …)` both land here with + // `q` in `query_kwargs["quantile"]`. Default to + // 0.99 when the caller didn't provide one + // (defensive — pattern-matched queries in + // `inference_config.yaml` always populate it). + let q: f64 = query_kwargs + .get("quantile") + .and_then(|s| s.parse().ok()) + .unwrap_or(0.99); + if !(0.0..=1.0).contains(&q) { + return Err(format!("DDSketchAccumulator: quantile {q} out of [0,1]").into()); + } + self.inner.quantile(q).ok_or_else(|| { + "DDSketchAccumulator: quantile() returned None (sketch empty?)".into() + }) + } + // Count is derived by summing the bucket store counts — the only + // DataPoint-level scalar that survives the wire-format trim + // (ProjectASAP/sketchlib-go#243 / asap_sketchlib#57). When the edge + // sampled this series (sample_p < 1.0), the stored count is ~p× the + // true count, so rescale by 1/sample_p to recover an unbiased + // estimate. sample_p == 1.0 (unsampled / legacy) makes this a no-op. + Statistic::Count => Ok(self.inner.total_count() as f64 / self.sample_p), + // STRICT policy: the Sum/Min/Max scalars were removed from + // the DDSketch wire format. They are now served by the + // controller-provisioned exact aggregations (an exact `Sum` + // and an exact `MinMax`), NOT estimated from the buckets. + // Surface the unavailable-statistic error so the query path + // routes to those aggregations instead of returning a wrong + // (0 / panicked) value. + Statistic::Sum => Err( + "DDSketchAccumulator: Sum not available from DDSketch wire format \ + (ProjectASAP/sketchlib-go#243); use an exact Sum aggregation" + .into(), + ), + Statistic::Min | Statistic::Max => Err(format!( + "DDSketchAccumulator: {statistic:?} not available from DDSketch wire format \ + (ProjectASAP/sketchlib-go#243); use an exact MinMax aggregation", + ) + .into()), + other => Err(format!( + "DDSketchAccumulator: statistic {other:?} not supported (only Quantile / Count)", + ) + .into()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // The DataPoint-level METRIC scalars (count/sum/min/max) were dropped + // from `DdSketchState` (ProjectASAP/sketchlib-go#243 / + // asap_sketchlib#57); the proto now carries only + // `alpha`/`store_counts`/`store_offset`. + fn encode_state(alpha: f64, store_counts: Vec, store_offset: i32) -> Vec { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + let state = DdSketchState { + alpha, + store_counts, + store_offset, + ..Default::default() + }; + SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), + ..Default::default() + } + .encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_round_trip() { + let bytes = encode_state(0.01, vec![1, 2, 3, 4], -2); + let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.alpha, 0.01); + assert_eq!(acc.inner.store_counts, vec![1, 2, 3, 4]); + assert_eq!(acc.inner.store_offset, -2); + // `count` is recovered by summing the bucket store counts. + assert_eq!(acc.inner.total_count(), 10); + } + + #[test] + fn test_from_sketchlib_proto_bytes_rejects_invalid_alpha() { + let bytes = encode_state(0.0, vec![1], 0); + let result = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("alpha")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's ddsketchprocessor emits: the + // state wrapped in a `SketchEnvelope{ddsketch: ...}` via + // sketchlib-go's `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + + let state = DdSketchState { + alpha: 0.01, + store_counts: vec![1, 2, 3, 4], + store_offset: -2, + ..Default::default() + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + assert_eq!(acc.inner.alpha, 0.01); + assert_eq!(acc.inner.total_count(), 10); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_aggregate_core_merge_aligns_buckets() { + let a = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 1, 1], -1), + sample_p: 1.0, + }; + let b = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![10, 10, 10], 0), + sample_p: 1.0, + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.inner.store_counts, vec![1, 11, 11, 10]); + assert_eq!(merged.inner.store_offset, -1); + assert_eq!(merged.inner.total_count(), 33); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::summary_kernels::count_sketch::CountSketchAccumulator; + let dd = DDSketchAccumulator::new(0.01); + let cs = CountSketchAccumulator::new(2, 3); + assert!(dd.merge_with(&cs).is_err()); + } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = DdSketch::from_raw(0.01, vec![5, 10, 15, 20], -2); + let bytes = original.to_msgpack().unwrap(); + let acc = DDSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.alpha, 0.01); + assert_eq!(acc.inner.store_counts, vec![5, 10, 15, 20]); + assert_eq!(acc.inner.store_offset, -2); + // `count` is recovered by summing the bucket store counts. + assert_eq!(acc.inner.total_count(), 50); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = DDSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } + + #[test] + fn test_apply_proto_delta_bytes_round_trip() { + use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use prost::Message; + + let mut acc = DDSketchAccumulator::new(0.01); + acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); + + // The wire delta now carries only bucket deltas (tags 2-7 + // reserved); `DdSketchBucketDelta` has just `index` + `d_count`. + let bytes = PbDelta { + buckets: vec![ + DdSketchBucketDelta { + index: 0, + d_count: 10, + }, + DdSketchBucketDelta { + index: 2, + d_count: 20, + }, + ], + ..Default::default() + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&bytes).expect("apply ok"); + assert_eq!(acc.inner.store_counts, vec![11, 2, 23]); + // `count` recomputed from the merged buckets: 11 + 2 + 23 = 36. + assert_eq!(acc.inner.total_count(), 36); + } + + /// A valid protobuf with an inadmissible span must not acknowledge a dropped update. + #[test] + fn test_apply_proto_delta_rejects_span_without_mutating_state() { + use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use prost::Message; + let mut acc = DDSketchAccumulator::new(0.01); + acc.inner = DdSketch::from_raw(0.01, vec![1, 2, 3], 0); + let bytes = PbDelta { + buckets: vec![DdSketchBucketDelta { + index: i32::MAX, + d_count: 1, + }], + ..Default::default() + } + .encode_to_vec(); + assert!(acc.apply_proto_delta_bytes(&bytes).is_err()); + assert_eq!(acc.inner.store_counts, vec![1, 2, 3]); + assert_eq!(acc.inner.store_offset, 0); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = DDSketchAccumulator::new(0.01); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ----- query_statistic STRICT policy ----- + // + // After the DataPoint-level METRIC scalars were dropped from the + // DDSketch wire format (ProjectASAP/sketchlib-go#243 / + // asap_sketchlib#57), DDSketch serves only quantiles and Count. + // Sum/Min/Max move to controller-provisioned exact aggregations and + // MUST surface the unavailable-statistic error (never a panic / 0). + + fn sample_accumulator() -> DDSketchAccumulator { + // Build the in-memory sketch from bucket counts only — no scalars. + DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), + sample_p: 1.0, + } + } + + #[test] + fn test_query_statistic_quantile_is_sketch_derived() { + use crate::Statistic; + let acc = sample_accumulator(); + let mut kwargs = HashMap::new(); + kwargs.insert("quantile".to_string(), "0.5".to_string()); + let v = acc + .query_statistic(Statistic::Quantile, &None, &kwargs) + .expect("quantile should be served from the sketch buckets"); + assert!( + v.is_finite() && v > 0.0, + "quantile estimate should be positive finite, got {v}" + ); + } + + #[test] + fn test_query_statistic_count_is_bucket_derived() { + use crate::Statistic; + let acc = sample_accumulator(); + let v = acc + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count should be derivable from the bucket store"); + // 1 + 2 + 3 + 4 = 10. + assert_eq!(v, 10.0); + } + + #[test] + fn test_query_statistic_sum_min_max_return_unavailable_error() { + use crate::Statistic; + let acc = sample_accumulator(); + for stat in [Statistic::Sum, Statistic::Min, Statistic::Max] { + let result = acc.query_statistic(stat, &None, &HashMap::new()); + assert!( + result.is_err(), + "{stat:?} must return the unavailable-statistic error (not a panic / 0)" + ); + let msg = result.unwrap_err().to_string(); + assert!( + msg.contains("not available"), + "{stat:?} error should explain the statistic is unavailable, got: {msg}" + ); + } + } + + // ----- sample_p count rescale ----- + // + // When the edge sampled a DDSketch (sample_p < 1.0), the stored count is + // ~p× the true count, so Count rescales by 1/p. Quantiles are + // rank-preserving and must NOT be rescaled. + + #[test] + fn test_count_is_rescaled_by_sample_p() { + use crate::Statistic; + let acc = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), + sample_p: 0.1, + }; + let c = acc + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + // Raw bucket sum 10, rescaled by 1/0.1 = 100. + assert!((c - 100.0).abs() < 1e-9, "expected rescaled 100, got {c}"); + } + + #[test] + fn test_quantile_ignores_sample_p() { + use crate::Statistic; + let mut kwargs = HashMap::new(); + kwargs.insert("quantile".to_string(), "0.5".to_string()); + let unsampled = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), + sample_p: 1.0, + }; + let sampled = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3, 4], -2), + sample_p: 0.1, + }; + let qu = unsampled + .query_statistic(Statistic::Quantile, &None, &kwargs) + .expect("q ok"); + let qs = sampled + .query_statistic(Statistic::Quantile, &None, &kwargs) + .expect("q ok"); + assert_eq!(qu, qs, "quantile must be sample_p-invariant"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { + use crate::Statistic; + use asap_sketchlib::proto::sketchlib::{sketch_envelope, DdSketchState, SketchEnvelope}; + use prost::Message; + + let env = SketchEnvelope { + sample_p: 0.25, + sketch_state: Some(sketch_envelope::SketchState::Ddsketch(DdSketchState { + alpha: 0.01, + store_counts: vec![2, 4, 6, 8], + store_offset: -2, + ..Default::default() + })), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 0.25); + // Raw 20, rescaled 20 / 0.25 = 80. + let c = acc + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + assert!((c - 80.0).abs() < 1e-9, "expected rescaled 80, got {c}"); + } + + #[test] + fn test_sample_p_normalization() { + // proto3 default (0.0), >=1.0, and non-finite all mean no sampling. + assert_eq!(normalize_sample_p(0.0), 1.0); + assert_eq!(normalize_sample_p(1.0), 1.0); + assert_eq!(normalize_sample_p(1.5), 1.0); + assert_eq!(normalize_sample_p(f64::NAN), 1.0); + assert_eq!(normalize_sample_p(-0.1), 1.0); + assert_eq!(normalize_sample_p(0.5), 0.5); + } + + #[test] + fn test_sample_p_from_envelope_bytes_defaults_to_one() { + use asap_sketchlib::proto::sketchlib::DdSketchState; + use prost::Message; + // Bare DdSketchState bytes (no envelope) → no sampling info → 1.0. + let bare = DdSketchState { + alpha: 0.01, + store_counts: vec![1, 2, 3], + store_offset: 0, + ..Default::default() + } + .encode_to_vec(); + assert_eq!( + DDSketchAccumulator::sample_p_from_envelope_bytes(&bare), + 1.0 + ); + } + + #[test] + fn test_reset_to_empty_preserves_sample_p() { + let mut acc = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 2, 3], 0), + sample_p: 0.2, + }; + acc.reset_to_empty(); + assert_eq!(acc.sample_p, 0.2, "window rotation must keep sample_p"); + assert_eq!(acc.inner.total_count(), 0, "buckets cleared"); + } + + #[test] + fn test_merge_prefers_sampled_factor() { + // A sampled base merged with a freshly-reset (1.0) operand keeps the + // series' sampling rate. + let a = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 1, 1], 0), + sample_p: 0.1, + }; + let b = DDSketchAccumulator { + inner: DdSketch::from_raw(0.01, vec![1, 1, 1], 0), + sample_p: 1.0, + }; + let merged = a.merge_with(&b).expect("merge ok"); + let merged = merged + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.sample_p, 0.1); + } +} + +#[cfg(test)] +mod dependency_upgrade_tests { + use super::*; + // The upgraded sketchlib state must retain negative and zero stores through both adapters. + #[test] + fn signed_state_survives_codec_and_accumulator_roundtrip() { + let mut inner = DdSketch::new(0.01); + for value in [-4.0, 0.0, 8.0] { + inner.update(value); + } + let bytes = asap_sketch_codec::encode_ddsketch(&inner); + let (wire, _) = asap_sketch_codec::ddsketch_state(&bytes).unwrap(); + assert_eq!(wire.zero_count, 1); + assert_eq!(wire.negative_store_counts.iter().sum::(), 1); + let restored = DDSketchAccumulator::from_sketchlib_proto_bytes(&bytes).unwrap(); + assert_eq!(restored.inner.total_count(), 3); + assert_eq!(restored.inner.alpha, inner.wire_alpha()); + assert_eq!(restored.inner.store_counts, inner.store_counts); + assert_eq!(restored.inner.store_offset, inner.store_offset); + assert_eq!( + restored.inner.negative_store_counts, + inner.negative_store_counts + ); + assert_eq!( + restored.inner.negative_store_offset, + inner.negative_store_offset + ); + assert_eq!(restored.inner.zero_count, inner.zero_count); + } + // Negative and zero delta fields added by sketchlib must not be discarded by the adapter. + #[test] + fn signed_delta_survives_adapter() { + use asap_sketchlib::proto::sketchlib::{DdSketchBucketDelta, DdSketchDelta as PbDelta}; + use prost::Message; + let mut accumulator = DDSketchAccumulator::new(0.01); + let bytes = PbDelta { + negative_buckets: vec![DdSketchBucketDelta { + index: 0, + d_count: 2, + }], + zero_count: 3, + ..Default::default() + } + .encode_to_vec(); + accumulator.apply_proto_delta_bytes(&bytes).unwrap(); + assert_eq!(accumulator.inner.total_count(), 5); + assert_eq!(accumulator.inner.zero_count, 3); + assert_eq!( + accumulator.inner.negative_store_counts.iter().sum::(), + 2 + ); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/exact.rs b/crates/asap-physical-operators/src/summary_kernels/exact.rs new file mode 100644 index 00000000..8652bec4 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/exact.rs @@ -0,0 +1,350 @@ +//! Exact summary state identified by Planner family, independent of keyed layout. +use super::increase::IncreaseAccumulator; +use crate::Statistic; +use crate::{ + AggregateCore, AggregationType, AuxStats, KeyByLabelValues, Measurement, SerializableToSink, +}; +use planner_types::post_asap::{ExactKind, ExactParams, SummaryFamilyType}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +type Error = Box; + +#[derive(Debug, Clone, Serialize, Deserialize)] +enum ScalarState { + Sum(f64), + Count(u64), + Min(Option), + Max(Option), + Counter(Option), +} + +/// Both the family and population layout survive persistence. Sharing counter +/// arithmetic never authorizes a Rate state to answer an Increase readout. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ExactAccumulator { + family: SummaryFamilyType, + scalar: ScalarState, + keyed: Option>, +} + +impl ExactAccumulator { + /// Accumulate into run-local scratch state. Persistent input states remain + /// immutable; a failed merge discards this scratch state. + pub(crate) fn merge_from(&mut self, other: &Self) -> Result<(), Error> { + if self.family != other.family || self.is_keyed() != other.is_keyed() { + return Err("cannot merge different Planner families or layouts".into()); + } + if let (Some(target), Some(source)) = (&mut self.keyed, &other.keyed) { + for (key, state) in source { + let combined = match target.get(key) { + Some(old) => merge_scalar(old, state)?, + None => state.clone(), + }; + target.insert(key.clone(), combined); + } + } else { + self.scalar = merge_scalar(&self.scalar, &other.scalar)?; + } + Ok(()) + } + + pub fn new(family: SummaryFamilyType, keyed: bool) -> Result { + use ExactKind as K; + use ExactParams as P; + let scalar = match &family { + SummaryFamilyType::ExactAggregate(K::Sum, P::Sum) => ScalarState::Sum(0.0), + SummaryFamilyType::ExactAggregate(K::Count, P::Count) => ScalarState::Count(0), + SummaryFamilyType::ExactAggregate(K::Min, P::Min) => ScalarState::Min(None), + SummaryFamilyType::ExactAggregate(K::Max, P::Max) => ScalarState::Max(None), + SummaryFamilyType::ExactAggregate(K::Rate, P::Rate) + | SummaryFamilyType::ExactAggregate(K::Increase, P::Increase) => { + ScalarState::Counter(None) + } + _ => return Err(format!("unsupported exact Planner family: {family:?}")), + }; + Ok(Self { + family, + scalar, + keyed: keyed.then(HashMap::new), + }) + } + + pub fn family(&self) -> &SummaryFamilyType { + &self.family + } + pub(crate) fn insufficient_counter_samples( + &self, + statistic: Statistic, + key: &Option, + ) -> bool { + if statistic != self.statistic() { + return false; + } + let state = match (&self.keyed, key) { + (Some(states), Some(key)) => states.get(key), + (None, None) => Some(&self.scalar), + _ => None, + }; + match state { + Some(ScalarState::Counter(None)) => true, + Some(ScalarState::Counter(Some(counter))) => { + counter.sample_count < 2 + || counter.last_seen_timestamp == counter.starting_timestamp + } + _ => false, + } + } + pub fn is_keyed(&self) -> bool { + self.keyed.is_some() + } + + pub fn update(&mut self, key: Option<&KeyByLabelValues>, value: f64, timestamp: i64) { + let state = match (&mut self.keyed, key) { + (Some(states), Some(key)) => states + .entry(key.clone()) + .or_insert_with(|| self.scalar.clone()), + (None, None) => &mut self.scalar, + _ => panic!("exact update population layout differs from installed DAG"), + }; + match state { + ScalarState::Sum(sum) => *sum += value, + ScalarState::Count(count) => { + *count = count.checked_add(1).expect("exact count overflow") + } + ScalarState::Min(current) => { + *current = Some(current.map_or(value, |old| old.min(value))) + } + ScalarState::Max(current) => { + *current = Some(current.map_or(value, |old| old.max(value))) + } + ScalarState::Counter(current) => match current { + Some(counter) => counter.update(Measurement::new(value), timestamp), + None => { + *current = Some(IncreaseAccumulator::new( + Measurement::new(value), + timestamp, + Measurement::new(value), + timestamp, + )) + } + }, + } + } + + pub fn deserialize_from_bytes(bytes: &[u8]) -> Result { + let state: Self = rmp_serde::from_slice(bytes)?; + let expected = Self::new(state.family.clone(), state.is_keyed())?; + let same_variant = |value: &ScalarState| { + std::mem::discriminant(value) == std::mem::discriminant(&expected.scalar) + }; + if !same_variant(&state.scalar) + || state + .keyed + .as_ref() + .is_some_and(|states| states.values().any(|s| !same_variant(s))) + { + return Err("exact payload differs from declared Planner family".into()); + } + Ok(state) + } + + fn statistic(&self) -> Statistic { + match self.family { + SummaryFamilyType::ExactAggregate(ExactKind::Sum, _) => Statistic::Sum, + SummaryFamilyType::ExactAggregate(ExactKind::Count, _) => Statistic::Count, + SummaryFamilyType::ExactAggregate(ExactKind::Min, _) => Statistic::Min, + SummaryFamilyType::ExactAggregate(ExactKind::Max, _) => Statistic::Max, + SummaryFamilyType::ExactAggregate(ExactKind::Rate, _) => Statistic::Rate, + SummaryFamilyType::ExactAggregate(ExactKind::Increase, _) => Statistic::Increase, + _ => unreachable!("validated exact family"), + } + } +} + +fn merge_scalar(left: &ScalarState, right: &ScalarState) -> Result { + Ok(match (left, right) { + (ScalarState::Sum(a), ScalarState::Sum(b)) => ScalarState::Sum(a + b), + (ScalarState::Count(a), ScalarState::Count(b)) => { + ScalarState::Count(a.checked_add(*b).ok_or("exact count overflow")?) + } + (ScalarState::Min(a), ScalarState::Min(b)) => { + ScalarState::Min(a.iter().chain(b).copied().reduce(f64::min)) + } + (ScalarState::Max(a), ScalarState::Max(b)) => { + ScalarState::Max(a.iter().chain(b).copied().reduce(f64::max)) + } + (ScalarState::Counter(a), ScalarState::Counter(b)) => ScalarState::Counter(match (a, b) { + (Some(a), Some(b)) => Some(IncreaseAccumulator::merge_pair(a, b)), + (a, b) => a.clone().or_else(|| b.clone()), + }), + _ => return Err("exact scalar state families differ".into()), + }) +} + +impl SerializableToSink for ExactAccumulator { + fn serialize_to_json(&self) -> serde_json::Value { + serde_json::json!({"family": self.family, "scalar": self.scalar, "keyed": self.keyed.as_ref().map(|m|m.iter().collect::>())}) + } + fn serialize_to_bytes(&self) -> Vec { + rmp_serde::to_vec_named(self).expect("exact state encoding") + } +} + +impl AggregateCore for ExactAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + fn type_name(&self) -> &'static str { + "PlannerExactAccumulatorV1" + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + fn merge_with(&self, other: &dyn AggregateCore) -> Result, Error> { + let other = other + .as_any() + .downcast_ref::() + .ok_or("merge requires Planner exact state")?; + let mut merged = self.clone(); + merged.merge_from(other)?; + Ok(Box::new(merged)) + } + fn get_accumulator_type(&self) -> AggregationType { + match self.statistic() { + Statistic::Sum => AggregationType::Sum, + Statistic::Count => AggregationType::Count, + Statistic::Min => AggregationType::Min, + Statistic::Max => AggregationType::Max, + Statistic::Rate => AggregationType::Rate, + Statistic::Increase => AggregationType::Increase, + _ => unreachable!(), + } + } + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::() + + self.keyed.as_ref().map_or(0, |m| { + m.keys() + .map(|k| { + std::mem::size_of::() + + k.labels.iter().map(String::len).sum::() + }) + .sum::() + }) + } + fn aux_stats(&self) -> AuxStats { + if self.is_keyed() { + return AuxStats::empty(); + } + match self.scalar { + ScalarState::Sum(value) => AuxStats { + sum: Some(value), + ..AuxStats::empty() + }, + ScalarState::Count(value) => AuxStats { + count: Some(value), + ..AuxStats::empty() + }, + ScalarState::Min(value) => AuxStats { + min: value, + ..AuxStats::empty() + }, + ScalarState::Max(value) => AuxStats { + max: value, + ..AuxStats::empty() + }, + ScalarState::Counter(_) => AuxStats::empty(), + } + } + fn get_keys(&self) -> Option> { + self.keyed.as_ref().map(|m| m.keys().cloned().collect()) + } + fn query_statistic( + &self, + statistic: Statistic, + key: &Option, + kwargs: &HashMap, + ) -> Result { + if statistic != self.statistic() { + return Err("readout differs from Planner exact family".into()); + } + let state = match (&self.keyed, key) { + (Some(states), Some(key)) => states.get(key).ok_or("unknown exact population")?, + (None, None) => &self.scalar, + _ => return Err("readout population differs from installed layout".into()), + }; + match state { + ScalarState::Sum(sum) => Ok(*sum), + ScalarState::Count(count) => Ok(*count as f64), + ScalarState::Min(value) | ScalarState::Max(value) => { + value.ok_or_else(|| "empty exact population".into()) + } + ScalarState::Counter(Some(counter)) => { + counter.query_statistic(statistic, &None, kwargs) + } + ScalarState::Counter(None) => Err("empty counter population".into()), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Identity, population isolation, and readout survive the persisted format. + #[test] + fn exact_families_roundtrip_and_reject_cross_family_operations() { + let families = [ + (ExactKind::Sum, ExactParams::Sum, Statistic::Sum, 16.0), + (ExactKind::Count, ExactParams::Count, Statistic::Count, 3.0), + (ExactKind::Min, ExactParams::Min, Statistic::Min, 2.0), + (ExactKind::Max, ExactParams::Max, Statistic::Max, 8.0), + (ExactKind::Rate, ExactParams::Rate, Statistic::Rate, 3.0), + ( + ExactKind::Increase, + ExactParams::Increase, + Statistic::Increase, + 6.0, + ), + ]; + for keyed in [false, true] { + let key = keyed.then(|| KeyByLabelValues::new_with_labels(vec!["a".into()])); + let mut states = Vec::new(); + for (kind, params, stat, value) in &families { + let mut state = ExactAccumulator::new( + SummaryFamilyType::ExactAggregate(kind.clone(), params.clone()), + keyed, + ) + .unwrap(); + for (ts, v) in [(1000, 8.0), (2000, 2.0), (3000, 6.0)] { + state.update(key.as_ref(), v, ts); + } + let restored = + ExactAccumulator::deserialize_from_bytes(&state.serialize_to_bytes()).unwrap(); + assert_eq!(restored.family(), state.family()); + assert_eq!( + restored + .query_statistic(*stat, &key, &HashMap::new()) + .unwrap(), + *value + ); + for (_, _, wrong, _) in &families { + if wrong != stat { + assert!(restored + .query_statistic(*wrong, &key, &HashMap::new()) + .is_err()); + } + } + states.push(restored); + } + for (i, a) in states.iter().enumerate() { + for (j, b) in states.iter().enumerate() { + assert_eq!(a.merge_with(b).is_ok(), i == j); + } + } + } + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/factory.rs b/crates/asap-physical-operators/src/summary_kernels/factory.rs new file mode 100644 index 00000000..cd57454a --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/factory.rs @@ -0,0 +1,1161 @@ +use crate::summary_kernels::{ + CountMinSketchAccumulator, CountMinSketchWithHeapAccumulator, CountSketchAccumulator, + CountSketchWithHeapAccumulator, DDSketchAccumulator, DatasketchesKLLAccumulator, + HydraKllSketchAccumulator, IncreaseAccumulator, KeyedCounterState, KeyedMaxState, + KeyedMinState, KeyedSumCountAccumulator, MaxAccumulator, MinAccumulator, SumAccumulator, +}; +use crate::{AggregateCore, KeyByLabelValues, Measurement}; +// Production dispatch consumes Planner SummaryAgg payloads directly. The +// config adapter below is compiled only for isolated historical kernel tests. +use crate::summary_kernels::hll_sketch::HllSketchAccumulator; +use crate::summary_kernels::univmon::UnivMonAccumulator; +use planner_types::post_asap::{ExactKind, SketchAlgorithm, SketchParams, SummaryFamilyType}; + +/// Generate the two boilerplate clone-based `AccumulatorUpdater` methods +/// for updaters whose inner `acc` field implements `Clone + AggregateCore`. +/// Not applicable to `IncreaseAccumulatorUpdater` (its `acc` is `Option<_>` +/// with non-trivial `None` handling). +macro_rules! impl_clone_accumulator_methods { + ($acc_field:ident) => { + fn take_accumulator(&mut self) -> Box { + let result = Box::new(self.$acc_field.clone()); + self.reset(); + result + } + + fn snapshot_accumulator(&self) -> Box { + Box::new(self.$acc_field.clone()) + } + + fn into_accumulator(self: Box) -> Box { + // Consume the updater and MOVE the accumulator out — no clone. + // Avoids the expensive `Clone` (a full msgpack serialize/deserialize + // round-trip for sketch accumulators) when a pane is evicted at + // window close. + let this = *self; + Box::new(this.$acc_field) + } + }; +} + +/// Shared update interface for query-time and precompute-time accumulation. +/// +/// This provides a uniform interface over all accumulator types so that the +/// worker loop doesn't need to know which concrete type it's dealing with. +pub trait AccumulatorUpdater: Send { + /// Validate an immutable precompute input before an updater can silently + /// discard a value outside its representable domain. + fn validate_single_input(&self, value: f64) -> Result<(), String> { + if value.is_finite() { + Ok(()) + } else { + Err("accumulator input must be finite".into()) + } + } + + /// Feed a single (value, timestamp_ms) pair — for SingleSubpopulation types. + fn update_single(&mut self, value: f64, timestamp_ms: i64); + + /// Feed a keyed (key, value, timestamp_ms) triple — for MultipleSubpopulation types. + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp_ms: i64); + + /// Extract the final accumulator as a boxed `AggregateCore`. + fn take_accumulator(&mut self) -> Box; + + /// Non-destructive read of the current accumulator state (clone without reset). + /// Used by pane-based sliding windows to read shared panes. + fn snapshot_accumulator(&self) -> Box; + + /// Consume the updater and return its accumulator BY MOVE, avoiding the + /// `Clone` that `take_accumulator`/`snapshot_accumulator` pay (for sketch + /// accumulators that clone is a full msgpack serialize/deserialize + /// round-trip). Used by `merge_panes_for_window` when a pane is evicted at + /// window close. Default falls back to a clone for updaters that can't + /// cheaply move their inner accumulator out. + fn into_accumulator(self: Box) -> Box { + self.snapshot_accumulator() + } + + /// Reset internal state for reuse (avoids re-allocation). + fn reset(&mut self); + + /// Whether this updater is keyed (MultipleSubpopulation). + fn is_keyed(&self) -> bool; + + /// Estimated memory usage in bytes. + fn memory_usage_bytes(&self) -> usize; +} + +// --------------------------------------------------------------------------- +// SumAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct SumAccumulatorUpdater { + acc: SumAccumulator, +} + +impl SumAccumulatorUpdater { + pub fn new() -> Self { + Self { + acc: SumAccumulator::new(), + } + } +} + +impl Default for SumAccumulatorUpdater { + fn default() -> Self { + Self::new() + } +} + +impl AccumulatorUpdater for SumAccumulatorUpdater { + fn update_single(&mut self, value: f64, _timestamp_ms: i64) { + self.acc.update(value); + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = SumAccumulator::new(); + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + } +} + +// --------------------------------------------------------------------------- +// MinAccumulatorUpdater / MaxAccumulatorUpdater +// --------------------------------------------------------------------------- + +macro_rules! extremum_updater { + ($updater:ident, $acc:ty) => { + #[derive(Default)] + pub struct $updater { + acc: $acc, + } + + impl $updater { + pub fn new() -> Self { + Self::default() + } + } + + impl AccumulatorUpdater for $updater { + fn update_single(&mut self, value: f64, _timestamp_ms: i64) { + self.acc.update(value); + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = <$acc>::new(); + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::<$acc>() + } + } + }; +} + +extremum_updater!(MinAccumulatorUpdater, MinAccumulator); +extremum_updater!(MaxAccumulatorUpdater, MaxAccumulator); + +// --------------------------------------------------------------------------- +// IncreaseAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct IncreaseAccumulatorUpdater { + acc: Option, +} + +impl IncreaseAccumulatorUpdater { + pub fn new() -> Self { + Self { acc: None } + } +} + +impl Default for IncreaseAccumulatorUpdater { + fn default() -> Self { + Self::new() + } +} + +impl AccumulatorUpdater for IncreaseAccumulatorUpdater { + fn update_single(&mut self, value: f64, timestamp_ms: i64) { + let measurement = Measurement::new(value); + match &mut self.acc { + Some(acc) => acc.update(measurement, timestamp_ms), + None => { + self.acc = Some(IncreaseAccumulator::new( + measurement.clone(), + timestamp_ms, + measurement, + timestamp_ms, + )); + } + } + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + // Hand-written: acc is Option<_> with non-trivial None handling. + fn take_accumulator(&mut self) -> Box { + let acc = self.acc.take().unwrap_or_else(|| { + IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(0.0), 0) + }); + let result = Box::new(acc); + self.reset(); + result + } + + fn snapshot_accumulator(&self) -> Box { + match &self.acc { + Some(acc) => Box::new(acc.clone()), + None => Box::new(IncreaseAccumulator::new( + Measurement::new(0.0), + 0, + Measurement::new(0.0), + 0, + )), + } + } + + fn reset(&mut self) { + self.acc = None; + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::>() + } +} + +// --------------------------------------------------------------------------- +// KllAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct KllAccumulatorUpdater { + acc: DatasketchesKLLAccumulator, + k: u16, +} + +impl KllAccumulatorUpdater { + pub fn new(k: u16) -> Self { + Self { + acc: DatasketchesKLLAccumulator::new(k), + k, + } + } +} + +impl AccumulatorUpdater for KllAccumulatorUpdater { + fn update_single(&mut self, value: f64, _timestamp_ms: i64) { + self.acc.update(value); + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = DatasketchesKLLAccumulator::new(self.k); + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + // KLL sketch size is hard to estimate precisely; use a rough estimate + std::mem::size_of::() + 4096 + } +} + +// --------------------------------------------------------------------------- +// DDSketchAccumulatorUpdater — pendant to KllAccumulatorUpdater +// --------------------------------------------------------------------------- +// +// Drives the agent-aggregated DDSketch path: the worker either +// (a) merges an inbound `DDSketchAccumulator` from the +// modified-OTLP `Data::Ddsketch` ingest (via the worker's +// `merge_with`), or (b) consumes raw values via `update_single` +// when an OTLP scalar datapoint matches an aggregation typed as +// DDSketch. (b) is the less common path but it lets the same +// aggregation slot serve both pre-aggregated agent sketches and +// raw OTLP gauges. +pub struct DDSketchAccumulatorUpdater { + acc: DDSketchAccumulator, + alpha: f64, +} + +impl DDSketchAccumulatorUpdater { + pub fn new(alpha: f64) -> Self { + Self { + acc: DDSketchAccumulator::new(alpha), + alpha, + } + } +} + +impl AccumulatorUpdater for DDSketchAccumulatorUpdater { + fn validate_single_input(&self, value: f64) -> Result<(), String> { + let (minimum, maximum) = + asap_sketchlib::sketches::ddsketch::ddsketch_indexable_bounds(self.alpha); + if value.is_finite() && value > 0.0 && value >= minimum && value <= maximum { + Ok(()) + } else { + Err("DDS maintenance input is outside its positive representable domain".into()) + } + } + + fn update_single(&mut self, value: f64, _timestamp_ms: i64) { + // sketch-core's DdSketch (the inner of DDSketchAccumulator) + // exposes `update(f64)` for single-value ingestion. The + // worker calls this when a raw OTLP datapoint matches an + // aggregation typed as DDSketch — the sketch-merge path + // uses `merge_with` directly. + self.acc.inner.update(value); + } + + fn update_keyed(&mut self, _key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = DDSketchAccumulator::new(self.alpha); + } + + fn is_keyed(&self) -> bool { + false + } + + fn memory_usage_bytes(&self) -> usize { + // Bucket store is variable; rough estimate matches KLL. + std::mem::size_of::() + 4096 + } +} + +// --------------------------------------------------------------------------- +// KeyedSumCountAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct KeyedSumCountAccumulatorUpdater { + acc: KeyedSumCountAccumulator, +} + +impl KeyedSumCountAccumulatorUpdater { + pub fn new() -> Self { + Self::for_family(ExactKind::Sum) + } + + pub fn for_family(family: ExactKind) -> Self { + Self { + acc: KeyedSumCountAccumulator::for_family(family), + } + } +} + +impl Default for KeyedSumCountAccumulatorUpdater { + fn default() -> Self { + Self::new() + } +} + +impl AccumulatorUpdater for KeyedSumCountAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.update(key.clone(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = KeyedSumCountAccumulator::for_family(self.acc.family.clone()); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.acc.sums.len() * (std::mem::size_of::() + 16) + } +} + +// --------------------------------------------------------------------------- +// KeyedMinStateUpdater / KeyedMaxStateUpdater +// --------------------------------------------------------------------------- + +macro_rules! multiple_extremum_updater { + ($updater:ident, $acc:ty) => { + #[derive(Default)] + pub struct $updater { + acc: $acc, + } + + impl $updater { + pub fn new() -> Self { + Self::default() + } + } + + impl AccumulatorUpdater for $updater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.update(key.clone(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = <$acc>::new(); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::<$acc>() + + self.acc.values.len() * (std::mem::size_of::() + 8) + } + } + }; +} + +multiple_extremum_updater!(KeyedMinStateUpdater, KeyedMinState); +multiple_extremum_updater!(KeyedMaxStateUpdater, KeyedMaxState); + +// --------------------------------------------------------------------------- +// KeyedCounterStateUpdater +// --------------------------------------------------------------------------- + +pub struct KeyedCounterStateUpdater { + acc: KeyedCounterState, +} + +impl KeyedCounterStateUpdater { + pub fn new() -> Self { + Self { + acc: KeyedCounterState::new(), + } + } +} + +impl Default for KeyedCounterStateUpdater { + fn default() -> Self { + Self::new() + } +} + +impl AccumulatorUpdater for KeyedCounterStateUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + let measurement = Measurement::new(value); + match self.acc.increases.entry(key.clone()) { + std::collections::hash_map::Entry::Occupied(mut e) => { + e.get_mut().update(measurement, timestamp_ms); + } + std::collections::hash_map::Entry::Vacant(e) => { + e.insert(IncreaseAccumulator::new( + measurement.clone(), + timestamp_ms, + measurement, + timestamp_ms, + )); + } + } + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = KeyedCounterState::new(); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.acc.increases.len() + * (std::mem::size_of::() + + std::mem::size_of::()) + } +} + +// --------------------------------------------------------------------------- +// CmsAccumulatorUpdater (CountMinSketch) +// --------------------------------------------------------------------------- + +/// Keyed weighted-frequency updater. +/// +/// A raw Prometheus sample represents the observed metric value, so a bare CMS +/// adds `value` for its key. Counting each received sample as one is a distinct +/// event-count operation and requires an explicit typed plan contract; it must +/// not be inferred from the sketch algorithm alone. +pub struct CmsAccumulatorUpdater { + acc: CountMinSketchAccumulator, + row_num: usize, + col_num: usize, +} + +impl CmsAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + acc: CountMinSketchAccumulator::new(row_num, col_num), + row_num, + col_num, + } + } +} + +impl AccumulatorUpdater for CmsAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.inner.update(&key.to_semicolon_str(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = CountMinSketchAccumulator::new(self.row_num, self.col_num); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + } +} + +// --------------------------------------------------------------------------- +// CmsHeapAccumulatorUpdater — value-weighted / count-weighted top-k +// --------------------------------------------------------------------------- + +/// What quantity the top-k heap ranks keys by. +/// +/// These are DIFFERENT query semantics and must be chosen explicitly: +/// +/// * [`TopkWeight::Value`] — accumulate **Σ of the datapoint value** per key. +/// This answers "top-k by total " (e.g. "top-k hosts by +/// total CPU"). The heap value is the summed metric value, so the read-side +/// reducer's "sort heap descending by value" yields the correct ranking. +/// +/// * [`TopkWeight::Count`] — accumulate **+1 per event** per key (occurrence +/// frequency), the textbook heavy-hitter / frequency-top-k semantics +/// ("which keys appear most often"). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TopkWeight { + /// Σ datapoint value per key (value-weighted top-k). + Value, + /// +1 per event per key (count-weighted / frequency top-k). + Count, +} + +/// Keyed top-k updater backed by a real `CountMinSketchWithHeap` (a CMS +/// matrix PLUS a size-`heap_size` top-k heap). Unlike the heap-LESS +/// `CmsAccumulatorUpdater`, this enumerates top-k keys at read time +/// (`get_topk_keys` / `topk_heap_items`), which is what `topk(...)` queries +/// need. +/// +/// The key is the configured group-by (`aggregated_labels`) value vector — +/// e.g. `host` — formed by `extract_aggregated_key_from_series` in the worker, +/// NOT the hardcoded metric label `item`. The accumulated quantity is selected +/// by [`TopkWeight`]: +/// * `Value` → `inner.update(key, value)` adds the datapoint value (Σ value). +/// * `Count` → `inner.update(key, 1.0)` adds one per event (Σ count). +/// +/// Both `CountMinSketchWithHeap` and `CountSketchWithHeap` raw-input policies +/// route here; the heap is the shared distinguishing payload. +pub struct CmsHeapAccumulatorUpdater { + acc: CountMinSketchWithHeapAccumulator, + row_num: usize, + col_num: usize, + heap_size: usize, + weight: TopkWeight, +} + +impl CmsHeapAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize, heap_size: usize, weight: TopkWeight) -> Self { + Self { + acc: CountMinSketchWithHeapAccumulator::new(row_num, col_num, heap_size), + row_num, + col_num, + heap_size, + weight, + } + } +} + +impl AccumulatorUpdater for CmsHeapAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + // Heap key = the group-by label-value vector (e.g. `host`), joined the + // same way the read-side `get_topk_keys` splits it back apart (`;`). + let weighted = match self.weight { + // Σ value: feed the datapoint value. sketchlib's CMS-heap + // `update(key, w)` adds `w.round()` occurrences of `key`, so the + // heap value accumulates the (rounded) summed metric value. + TopkWeight::Value => value, + // Σ count: one occurrence per event, regardless of value. + TopkWeight::Count => 1.0, + }; + self.acc.inner.update(&key.to_semicolon_str(), weighted); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = + CountMinSketchWithHeapAccumulator::new(self.row_num, self.col_num, self.heap_size); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + + self.heap_size * (std::mem::size_of::() + 32) + } +} + +// --------------------------------------------------------------------------- +// CountSketchAccumulatorUpdater (real median-of-signed-rows CountSketch) +// --------------------------------------------------------------------------- + +/// Keyed point-frequency updater backed by a real `asap_sketchlib::CountSketch` +/// (signed rows, median-of-rows estimator) — distinct math from +/// `CmsAccumulatorUpdater`'s CMS (min-of-rows). Closes, on the raw-metric +/// ingest path, the conflation bug where `SketchAlgorithm::CountSketch` silently +/// shared `CmsAccumulatorUpdater` with bare CMS. +/// +/// As with bare CMS, each raw Prometheus sample contributes its `value`. +/// Unit event counting must be selected explicitly by a future typed plan +/// contract rather than being implied by `SketchAlgorithm::CountSketch`. +pub struct CountSketchAccumulatorUpdater { + acc: CountSketchAccumulator, + row_num: usize, + col_num: usize, +} + +impl CountSketchAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize) -> Self { + Self { + acc: CountSketchAccumulator::new(row_num, col_num), + row_num, + col_num, + } + } +} + +impl AccumulatorUpdater for CountSketchAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.inner.update(&key.to_semicolon_str(), value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = CountSketchAccumulator::new(self.row_num, self.col_num); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + } +} + +// --------------------------------------------------------------------------- +// CountSketchWithHeapAccumulatorUpdater (real CountSketch + top-k heap) +// --------------------------------------------------------------------------- + +/// Keyed top-k updater backed by a real `CountSketchWithHeap` (signed-row +/// CountSketch matrix PLUS a size-`heap_size` top-k heap). Distinct math from +/// `CmsHeapAccumulatorUpdater`'s CMS-with-heap (min-of-rows); shares the same +/// [`TopkWeight`] semantics and heap payload shape. +pub struct CountSketchWithHeapAccumulatorUpdater { + acc: CountSketchWithHeapAccumulator, + row_num: usize, + col_num: usize, + heap_size: usize, + weight: TopkWeight, +} + +impl CountSketchWithHeapAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize, heap_size: usize, weight: TopkWeight) -> Self { + Self { + acc: CountSketchWithHeapAccumulator::new(row_num, col_num, heap_size), + row_num, + col_num, + heap_size, + weight, + } + } +} + +impl AccumulatorUpdater for CountSketchWithHeapAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + let weighted = match self.weight { + TopkWeight::Value => value, + TopkWeight::Count => 1.0, + }; + self.acc.inner.update(&key.to_semicolon_str(), weighted); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = CountSketchWithHeapAccumulator::new(self.row_num, self.col_num, self.heap_size); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + std::mem::size_of::() + + self.row_num * self.col_num * std::mem::size_of::() + + self.heap_size * (std::mem::size_of::() + 32) + } +} + +// --------------------------------------------------------------------------- +// HydraKllAccumulatorUpdater +// --------------------------------------------------------------------------- + +pub struct HydraKllAccumulatorUpdater { + acc: HydraKllSketchAccumulator, + row_num: usize, + col_num: usize, + k: u16, +} + +impl HydraKllAccumulatorUpdater { + pub fn new(row_num: usize, col_num: usize, k: u16) -> Self { + Self { + acc: HydraKllSketchAccumulator::new(row_num, col_num, k), + row_num, + col_num, + k, + } + } +} + +impl AccumulatorUpdater for HydraKllAccumulatorUpdater { + fn update_single(&mut self, _value: f64, _timestamp_ms: i64) { + debug_assert!( + false, + "update_single called on keyed updater; use update_keyed" + ); + } + + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, _timestamp_ms: i64) { + self.acc.update(key, value); + } + + impl_clone_accumulator_methods!(acc); + + fn reset(&mut self) { + self.acc = HydraKllSketchAccumulator::new(self.row_num, self.col_num, self.k); + } + + fn is_keyed(&self) -> bool { + true + } + + fn memory_usage_bytes(&self) -> usize { + // Rough estimate: each cell is a KLL sketch + std::mem::size_of::() + self.row_num * self.col_num * 4096 + } +} + +// --------------------------------------------------------------------------- +// Config helpers +// --------------------------------------------------------------------------- + +fn cms_dims(params: &SketchParams) -> (usize, usize) { + match params { + SketchParams::Cms { width, depth } | SketchParams::CountSketch { width, depth } => { + (*depth as usize, *width as usize) + } + other => unreachable!( + "accumulator_spec() paired SketchAlgorithm::Cms/CountSketch with unexpected params: {other:?}" + ), + } +} + +/// Read `(rows = depth, columns = width, heap_size)` out of `SketchParams::CmsWithHeap` +/// or `::CountSketchWithHeap`. +fn cms_heap_dims(params: &SketchParams) -> (usize, usize, usize) { + match params { + SketchParams::CmsWithHeap { + width, + depth, + heap_size, + } + | SketchParams::CountSketchWithHeap { + width, + depth, + heap_size, + } => (*depth as usize, *width as usize, *heap_size as usize), + other => unreachable!( + "accumulator_spec() paired a WithHeap SketchAlgorithm with unexpected params: {other:?}" + ), + } +} + +/// Construct the kernel declared by a Planner SummaryAgg. No backend config +/// tags participate in this dispatch and unsupported payloads are errors. +pub fn create_planner_accumulator( + family: &SummaryFamilyType, + input: &planner_types::post_asap::SummaryUpdate, + grouping: &planner_types::post_asap::GroupingStrategy, +) -> Result, String> { + if input.item.is_some() + && matches!( + input.weight_domain, + planner_types::post_asap::WeightDomain::NonNegative { + proof: + planner_types::post_asap::NonNegativeWeightProof::ResetAwareCounterDerivative + } + ) + { + return Err("window-weighted summaries require typed DAG binding; integer heap updaters cannot consume rates".into()); + } + + 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::summary_kernels::exact::ExactAccumulator::new( + family.clone(), + input.item.is_some(), + )?, + })); + } + let SummaryFamilyType::Sketch(kind, family_grouping) = family else { + return Err(format!("unsupported Planner summary family {family:?}")); + }; + if family_grouping != grouping { + return Err("Planner family and operator grouping disagree".into()); + } + let updater: Box = match (kind.algorithm(), kind.params()) { + (SketchAlgorithm::Kll, SketchParams::Kll { k }) => Box::new(KllAccumulatorUpdater::new( + u16::try_from(*k).map_err(|_| "KLL k exceeds runtime bound")?, + )), + (SketchAlgorithm::DDSketch, SketchParams::DDSketch { alpha }) => { + Box::new(DDSketchAccumulatorUpdater::new(*alpha)) + } + (SketchAlgorithm::Cms, params @ SketchParams::Cms { .. }) => { + let (r, c) = cms_dims(params); + Box::new(CmsAccumulatorUpdater::new(r, c)) + } + (SketchAlgorithm::CountSketch, params @ SketchParams::CountSketch { .. }) => { + let (r, c) = cms_dims(params); + Box::new(CountSketchAccumulatorUpdater::new(r, c)) + } + (SketchAlgorithm::CmsWithHeap, params @ SketchParams::CmsWithHeap { .. }) => { + let (r, c, h) = cms_heap_dims(params); + Box::new(CmsHeapAccumulatorUpdater::new(r, c, h, TopkWeight::Value)) + } + ( + SketchAlgorithm::CountSketchWithHeap, + params @ SketchParams::CountSketchWithHeap { .. }, + ) => { + let (r, c, h) = cms_heap_dims(params); + Box::new(CountSketchWithHeapAccumulatorUpdater::new( + r, + c, + h, + TopkWeight::Value, + )) + } + (SketchAlgorithm::Hll, SketchParams::Hll { precision }) => Box::new(HllUpdater { + acc: HllSketchAccumulator::new( + asap_sketchlib::HllVariant::Regular, + u32::from(*precision), + ), + }), + ( + SketchAlgorithm::UnivMon, + SketchParams::UnivMon { + heap_size, + sketch_rows, + sketch_cols, + layers, + }, + ) => Box::new(UnivMonUpdater { + acc: UnivMonAccumulator::new( + *heap_size as usize, + *sketch_rows as usize, + *sketch_cols as usize, + *layers as usize, + ) + .map_err(|e| e.to_string())?, + }), + _ => { + return Err(format!( + "unsupported Planner algorithm/parameters: {kind:?}" + )) + } + }; + if updater.is_keyed() != input.item.is_some() + && !crate::capability::is_unit_sample_frequency(input) + { + return Err("Planner item expression does not match the selected kernel layout".into()); + } + Ok(updater) +} + +struct PlannerExactUpdater { + acc: crate::summary_kernels::exact::ExactAccumulator, +} +impl AccumulatorUpdater for PlannerExactUpdater { + fn update_single(&mut self, value: f64, timestamp: i64) { + self.acc.update(None, value, timestamp); + } + fn update_keyed(&mut self, key: &KeyByLabelValues, value: f64, timestamp: i64) { + self.acc.update(Some(key), value, timestamp); + } + impl_clone_accumulator_methods!(acc); + fn reset(&mut self) { + self.acc = crate::summary_kernels::exact::ExactAccumulator::new( + self.acc.family().clone(), + self.acc.is_keyed(), + ) + .expect("installed exact family"); + } + fn is_keyed(&self) -> bool { + self.acc.is_keyed() + } + fn memory_usage_bytes(&self) -> usize { + self.acc.approx_memory_bytes() + } +} + +struct UnivMonUpdater { + acc: UnivMonAccumulator, +} + +struct HllUpdater { + acc: HllSketchAccumulator, +} + +impl AccumulatorUpdater for HllUpdater { + fn is_keyed(&self) -> bool { + false + } + fn memory_usage_bytes(&self) -> usize { + self.acc.approx_memory_bytes() + } + fn update_single(&mut self, value: f64, _: i64) { + if !value.is_nan() { + let bits = if value == 0.0 { 0 } else { value.to_bits() }; + self.acc.inner.update(&bits.to_le_bytes()); + } + } + fn update_keyed(&mut self, _: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + impl_clone_accumulator_methods!(acc); + fn reset(&mut self) { + self.acc.reset_to_empty(); + } +} + +impl AccumulatorUpdater for UnivMonUpdater { + fn is_keyed(&self) -> bool { + false + } + fn memory_usage_bytes(&self) -> usize { + self.acc.approx_memory_bytes() + } + fn update_single(&mut self, value: f64, _: i64) { + self.acc + .insert_sample(value) + .expect("UnivMon sample counter overflow"); + } + fn update_keyed(&mut self, _: &KeyByLabelValues, value: f64, timestamp_ms: i64) { + self.update_single(value, timestamp_ms); + } + impl_clone_accumulator_methods!(acc); + fn reset(&mut self) { + self.acc.reset_to_empty(); + } +} + +#[cfg(test)] +mod planner_parameter_regression { + use super::*; + use planner_types::post_asap::{SketchKind, SummaryInputExpr, SummaryUpdate}; + + // Planner width is the bucket count; depth is the independent hash-row count. + #[test] + fn planner_sketch_dimensions_are_not_transposed() { + for (algorithm, params) in [ + ( + SketchAlgorithm::Cms, + SketchParams::Cms { + width: 128, + depth: 3, + }, + ), + ( + SketchAlgorithm::CountSketch, + SketchParams::CountSketch { + width: 128, + depth: 3, + }, + ), + ( + SketchAlgorithm::CmsWithHeap, + SketchParams::CmsWithHeap { + width: 128, + depth: 3, + heap_size: 8, + }, + ), + ( + SketchAlgorithm::CountSketchWithHeap, + SketchParams::CountSketchWithHeap { + width: 128, + depth: 3, + heap_size: 8, + }, + ), + ] { + let family = SummaryFamilyType::Sketch( + SketchKind::new(algorithm.clone(), params), + Default::default(), + ); + let update = SummaryUpdate { + item: Some(SummaryInputExpr::Column( + planner_types::pre_asap::ColumnRef::Named("host".into()), + )), + weight: SummaryInputExpr::Constant(1.0), + weight_domain: Default::default(), + }; + let state = create_planner_accumulator(&family, &update, &Default::default()) + .unwrap() + .snapshot_accumulator(); + let dims = match algorithm { + SketchAlgorithm::Cms => { + let s = state + .as_any() + .downcast_ref::() + .unwrap(); + (s.inner.rows(), s.inner.cols()) + } + SketchAlgorithm::CountSketch => { + let s = state + .as_any() + .downcast_ref::() + .unwrap(); + (s.inner.rows, s.inner.cols) + } + SketchAlgorithm::CmsWithHeap => { + let s = state + .as_any() + .downcast_ref::() + .unwrap(); + (s.inner.rows(), s.inner.cols()) + } + SketchAlgorithm::CountSketchWithHeap => { + let s = state + .as_any() + .downcast_ref::() + .unwrap(); + (s.inner.rows(), s.inner.cols()) + } + _ => unreachable!(), + }; + assert_eq!(dims, (3, 128), "{algorithm:?}"); + } + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/hll_sketch.rs b/crates/asap-physical-operators/src/summary_kernels/hll_sketch.rs new file mode 100644 index 00000000..4c16156d --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/hll_sketch.rs @@ -0,0 +1,788 @@ +//! HLL accumulator — wraps `asap_sketchlib::HllSketch`. +//! +//! Concrete accumulator reached from the modified-OTLP +//! `Metric.data = HLLSketch{…}` hot path (PR C-CountSketch follow-up). +//! Mirrors the CountSketch accumulator's shape: merge via register-wise +//! max on the inner sketch, serialize as MessagePack for the sink, and +//! decode from the sketchlib `HyperLogLogState` proto. +//! +//! Query semantics (cardinality estimation via the three HLL variants' +//! estimators) are intentionally deferred — the wire format carries the +//! registers + variant + HIP accumulators losslessly, so the merge + +//! store round-trip works end-to-end without that richer query surface. + +use crate::summary_kernels::dd_sketch::normalize_sample_p; +use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::{HllSketch, HllVariant, MessagePackCodec}; +use serde_json::Value; +use std::collections::HashMap; + +/// Decode one protobuf base-128 varint (LEB128) from the front of `buf`. +/// Returns `(value, bytes_consumed)`, or `None` if the buffer is truncated +/// or the varint overflows u64. +pub(crate) fn read_uvarint(buf: &[u8]) -> Option<(u64, usize)> { + let mut result: u64 = 0; + let mut shift: u32 = 0; + for (i, &b) in buf.iter().enumerate() { + if shift >= 64 { + return None; + } + result |= u64::from(b & 0x7f) << shift; + if b & 0x80 == 0 { + return Some((result, i + 1)); + } + shift += 7; + } + None +} + +/// Expand sketchlib-go's sparse HLL register encoding +/// (`HLLSparseRegisters.packed`) into the dense `num_registers`-byte array. +/// +/// Layout (sketchlib-go `proto/hll/hll.proto`): varint-packed +/// `(index_delta, value)` pairs in ascending index order; `prev_index` +/// starts at 0, so each register's absolute index is the running sum of the +/// deltas. Mirrors the Go encoder in `sketches/HLL/sparse.go` +/// (`encodeSparseRegisters`). The reconstructed array is byte-identical to +/// the dense `registers` field a high-cardinality producer would have sent. +pub(crate) fn expand_sparse_hll_registers( + packed: &[u8], + num_registers: usize, +) -> Result, Box> { + let mut regs = vec![0u8; num_registers]; + let mut prev: u64 = 0; + let mut pos = 0usize; + while pos < packed.len() { + let (delta, n1) = read_uvarint(&packed[pos..]) + .ok_or("HLLSparseRegisters.packed: truncated index_delta varint")?; + pos += n1; + let (value, n2) = read_uvarint(&packed[pos..]) + .ok_or("HLLSparseRegisters.packed: truncated value varint")?; + pos += n2; + let idx = prev + delta; + let i = usize::try_from(idx) + .map_err(|_| format!("HLLSparseRegisters: index {idx} overflows usize"))?; + if i >= num_registers { + return Err(format!( + "HLLSparseRegisters: register index {i} >= num_registers {num_registers}" + ) + .into()); + } + regs[i] = u8::try_from(value) + .map_err(|_| format!("HLLSparseRegisters: register value {value} > 255"))?; + prev = idx; + } + Ok(regs) +} + +/// HLL accumulator — inner register array + variant metadata. +#[derive(Debug, Clone)] +pub struct HllSketchAccumulator { + pub inner: HllSketch, + /// Edge sampling probability `p ∈ (0,1]` carried on the producer's + /// `SketchEnvelope.sample_p`. HLL uses HASH-THRESHOLD sampling — each + /// DISTINCT key is admitted into the sketch with probability `p`, so the + /// register-derived distinct-count estimate is ~`p`× the true + /// cardinality and a `Cardinality`/`Count` query must rescale by `1/p`. + /// `1.0` (and the proto3 default `0.0`, dual-read as `1.0`) means no + /// sampling, so the rescale is a no-op and the behaviour is identical to + /// before. Mirrors `DDSketchAccumulator::sample_p`; set from the envelope + /// at the `from_sketchlib_proto_bytes` decode site and preserved across + /// `reset_to_empty` and `merge_with`. + /// + /// NOTE: HLL edge sampling is currently force-disabled in the edge + /// (`warm_sketch.go` HLL case always emits `sample_p = 1.0`), so in + /// practice `p = 1.0` today and this is a latent-correctness fix that + /// activates if HLL sampling is ever enabled. + pub sample_p: f64, +} + +impl HllSketchAccumulator { + pub fn new(variant: HllVariant, precision: u32) -> Self { + Self { + inner: HllSketch::new(variant, precision), + sample_p: 1.0, + } + } + + /// Decode from the modified OTLP wire format's + /// `HLLSketchDataPoint.sketch` bytes when + /// `encoding = HLL_SKETCH_ENCODING_MSGPACK`. The bytes are the + /// MessagePack serialization of the cross-language sketch-core + /// `HllSketch` struct — PR I parity entrypoint. + pub fn from_msgpack_bytes(buffer: &[u8]) -> Result> { + Ok(Self { + inner: HllSketch::from_msgpack(buffer) + .map_err(|e| format!("deserialize HllSketch msgpack: {e}"))?, + // The msgpack HllSketch struct carries no envelope/sample_p; the + // msgpack path is parity/test-only and is never edge-sampled. + sample_p: 1.0, + }) + } + + /// Decode from the modified OTLP wire format's + /// `HLLSketchDataPoint.sketch` bytes — the protobuf-encoded + /// `asap_sketchlib::proto::sketchlib::HyperLogLogState` message + /// that DataCollector's `hllprocessor` emits when + /// `encoding = HLL_SKETCH_ENCODING_PROTO`. + pub fn from_sketchlib_proto_bytes(buffer: &[u8]) -> Result> { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + + // DataCollector's hllprocessor wraps the state in a + // `SketchEnvelope{hll: HyperLogLogState}` via sketchlib-go's + // `SerializePortableFO` + `proto.Marshal`. Try envelope first, + // fall back to bare `HyperLogLogState` for callers (e.g. unit + // tests) that encode the state directly. Mirrors the PR #14 + // fix on `CountMinSketchAccumulator::from_sketchlib_proto_bytes`. + // Capture the envelope's `sample_p` alongside the state so a + // Cardinality query can rescale the distinct-count estimate by + // `1/p`. Bare `HyperLogLogState` bytes (no envelope) carry no + // sampling info → `sample_p` 1.0 (no rescale). Mirrors + // `DDSketchAccumulator`. + let (state, sample_p) = match SketchEnvelope::decode(buffer) { + Ok(env) => { + let sp = env.sample_p; + match env.sketch_state { + Some(sketch_envelope::SketchState::Hll(st)) => (st, sp), + Some(other) => { + return Err(format!( + "SketchEnvelope contains non-HLL sketch: {:?}", + std::mem::discriminant(&other) + ) + .into()); + } + None => ( + HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?, + 1.0, + ), + } + } + Err(_) => ( + HyperLogLogState::decode(buffer) + .map_err(|e| format!("decode HyperLogLogState: {e}"))?, + 1.0, + ), + }; + if state.precision == 0 || state.precision > 20 { + return Err(format!( + "HyperLogLogState precision {} out of range (expected 1..=20)", + state.precision + ) + .into()); + } + let expected_len = 1usize << state.precision; + // Register resolution. sketchlib-go emits the SPARSE + // `registers_sparse` (proto tag 7) form below its dense/sparse + // crossover (~6000 non-zero registers — see + // sketchlib-go/sketches/HLL/sparse.go); low-cardinality producers + // (the common case) therefore leave the dense `registers` (tag 3) + // field empty. The proto contract (hll.proto) is: read whichever of + // `registers` / `registers_sparse` is present; if both are empty the + // sketch is all-zero. Reconstruct the dense 2^precision array in all + // three cases so the inner `HllSketch` always gets a full register + // vector. + let dense_registers: Vec = if state.registers.len() == expected_len { + state.registers.clone() + } else if !state.registers.is_empty() { + // A non-empty dense field of the wrong length is a malformed frame. + return Err(format!( + "HyperLogLogState registers has {} bytes, expected 2^precision = {}", + state.registers.len(), + expected_len + ) + .into()); + } else if let Some(sparse) = state.registers_sparse.as_ref() { + expand_sparse_hll_registers(&sparse.packed, expected_len)? + } else { + // Neither representation populated → all-zero register array. + vec![0u8; expected_len] + }; + let proto_variant = ProtoVariant::try_from(state.variant) + .map_err(|_| format!("HyperLogLogState has unknown variant tag {}", state.variant))?; + let variant = match proto_variant { + ProtoVariant::Unspecified => HllVariant::Unspecified, + ProtoVariant::Regular => HllVariant::Regular, + ProtoVariant::ErtlMle => HllVariant::Datafusion, + ProtoVariant::Hip => HllVariant::Hip, + }; + let inner = HllSketch::from_raw( + variant, + state.precision, + dense_registers, + state.hip_kxq0, + state.hip_kxq1, + state.hip_est, + ); + Ok(Self { + inner, + sample_p: normalize_sample_p(sample_p), + }) + } + + /// Apply a proto-encoded `HLLDelta` frame to this accumulator's + /// inner sketch — the decode path for + /// `HLL_SKETCH_ENCODING_PROTO_DELTA` (paper §6.2 B3 / B4). + /// + /// Called against an accumulator that already carries the base + /// sketch state; the caller is the per-series snapshot cache in + /// the ingest path. Bytes are the + /// `asap_sketchlib::proto::sketchlib::HllDelta` message. + pub fn apply_proto_delta_bytes( + &mut self, + buffer: &[u8], + ) -> Result<(), Box> { + // The HLLDelta wire format is a varint-packed (index_delta, value) blob; + // decode + apply (register-wise max) via the shared sketch library so + // the unpacking stays a single source of truth. + self.inner + .apply_delta_bytes(buffer) + .map_err(|e| format!("apply HLLDelta: {e}"))?; + Ok(()) + } +} + +impl SerializableToSink for HllSketchAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "variant": format!("{:?}", self.inner.variant), + "precision": self.inner.precision, + "register_bytes": self.inner.registers.len(), + "hip_kxq0": self.inner.hip_kxq0, + "hip_kxq1": self.inner.hip_kxq1, + "hip_est": self.inner.hip_est, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl AggregateCore for HllSketchAccumulator { + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::().saturating_add(self.inner.registers.capacity()) + } + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "HllSketchAccumulator" + } + + /// Per-window base rotation: zero the registers but keep the variant + /// and precision. Critical for HLL — its register-wise `max` merge + /// has no inverse, so a never-reset base accumulates the all-time-max + /// across windows (`docs/delta-baseline-contract.md` §1.5); rotating + /// to an empty register array makes per-window cardinality correct. + /// `sample_p` is a per-series config constant (not per-window data), so + /// it is intentionally preserved across the rotation — mirrors + /// `DDSketchAccumulator`. + fn reset_to_empty(&mut self) { + self.inner = HllSketch::new(self.inner.variant, self.inner.precision); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge HllSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_hll = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to HllSketchAccumulator")?; + let merged_inner = HllSketch::merge_refs(&[&self.inner, &other_hll.inner])?; + // Mirror DDSketchAccumulator's merge policy exactly: sample_p is a + // per-series config constant, so both operands carry the same value + // in practice. Prefer a sampled factor over the no-sampling default + // so a merge with a freshly-reset (1.0) base keeps the series' + // sampling rate. + let sample_p = if self.sample_p < 1.0 { + self.sample_p + } else { + other_hll.sample_p + }; + Ok(Box::new(Self { + inner: merged_inner, + sample_p, + })) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::HLL + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + _query_kwargs: &HashMap, + ) -> Result> { + use crate::Statistic; + match statistic { + // HLL's natural answer is unique-cardinality. PromQL's + // `count_over_time(...)` and `count(...)` both surface + // as `Statistic::Count` after pattern matching but + // semantically they mean "how many distinct values + // were observed in this window" when the underlying + // aggregator is HLL — that's the cardinality estimate, + // not a sample-count. Accept both. + Statistic::Cardinality | Statistic::Count => { + // HLL uses hash-threshold sampling — each distinct key is + // admitted with probability `sample_p`, so the register- + // derived distinct-count estimate is ~`p`× the true + // cardinality. Rescale by `1/sample_p` for an unbiased + // estimate. `sample_p == 1.0` (unsampled / legacy / edge + // HLL sampling currently force-disabled) makes this a no-op. + Ok(hll_cardinality_estimate(&self.inner.registers) / self.sample_p) + } + other => Err(format!( + "HllSketchAccumulator: statistic {:?} not supported (only Cardinality / Count)", + other, + ) + .into()), + } + } +} + +/// Standard HyperLogLog cardinality estimate with the canonical +/// `α_m × m² / Σ 2^(-register[i])` formula plus the small-range +/// (linear-counting) and large-range (32-bit space) corrections +/// from the original Flajolet et al. paper. +/// +/// Inlined here rather than added as a method on `asap_sketchlib::HllSketch` +/// because the existing `asap_sketchlib::asap` types only expose merge / +/// serialize today; adding a query method there would force a +/// cross-crate change. +fn hll_cardinality_estimate(registers: &[u8]) -> f64 { + let m = registers.len() as f64; + if m == 0.0 { + return 0.0; + } + let alpha = match registers.len() { + 16 => 0.673, + 32 => 0.697, + 64 => 0.709, + _ => 0.7213 / (1.0 + 1.079 / m), + }; + + let mut sum = 0.0f64; + let mut zero_registers = 0usize; + for &r in registers { + sum += 2f64.powi(-(r as i32)); + if r == 0 { + zero_registers += 1; + } + } + let raw = alpha * m * m / sum; + + // Small-range (linear-counting) correction. + if raw <= 2.5 * m && zero_registers > 0 { + return m * (m / zero_registers as f64).ln(); + } + + // Large-range correction (only meaningful with 32-bit register + // spaces; sketch-core uses up to 64-bit hashes so this branch + // rarely fires in practice — kept for completeness). + let two_pow_32 = 4_294_967_296f64; + if raw > two_pow_32 / 30.0 { + return -two_pow_32 * (1.0 - raw / two_pow_32).ln(); + } + raw +} + +#[cfg(test)] +mod tests { + use super::*; + + fn encode_state( + variant: i32, + precision: u32, + registers: Vec, + hip_kxq0: f64, + hip_kxq1: f64, + hip_est: f64, + ) -> Vec { + use asap_sketchlib::proto::sketchlib::HyperLogLogState; + use prost::Message; + let state = HyperLogLogState { + variant, + precision, + registers, + hip_kxq0, + hip_kxq1, + hip_est, + registers_sparse: None, + }; + state.encode_to_vec() + } + + #[test] + fn test_from_sketchlib_proto_bytes_regular() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + let bytes = encode_state( + ProtoVariant::Regular as i32, + 2, + vec![1, 2, 3, 4], + 0.0, + 0.0, + 0.0, + ); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Regular); + assert_eq!(acc.inner.precision, 2); + assert_eq!(acc.inner.registers, vec![1, 2, 3, 4]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_hip_preserves_accumulators() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + let bytes = encode_state( + ProtoVariant::Hip as i32, + 2, + vec![0, 0, 0, 0], + 1.5, + 2.5, + 42.0, + ); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Hip); + assert_eq!(acc.inner.hip_kxq0, 1.5); + assert_eq!(acc.inner.hip_kxq1, 2.5); + assert_eq!(acc.inner.hip_est, 42.0); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrapped() { + // Mirrors what DataCollector's hllprocessor emits: the state + // wrapped in a `SketchEnvelope{hll: ...}` via sketchlib-go's + // `SerializePortableFO` + `proto.Marshal`. + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision: 2, + registers: vec![1, 2, 3, 4], + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + registers_sparse: None, + }; + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes) + .expect("envelope-wrapped decode should succeed"); + assert_eq!(acc.inner.variant, HllVariant::Regular); + assert_eq!(acc.inner.registers, vec![1, 2, 3, 4]); + } + + #[test] + fn test_from_sketchlib_proto_bytes_envelope_wrong_sketch_type() { + use asap_sketchlib::proto::sketchlib::{sketch_envelope, KllState, SketchEnvelope}; + use prost::Message; + + let env = SketchEnvelope { + sketch_state: Some(sketch_envelope::SketchState::Kll(KllState::default())), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + + let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err(), "wrong-sketch envelope should error"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_register_length_mismatch() { + use asap_sketchlib::proto::sketchlib::HllVariant as ProtoVariant; + // precision=2 → expected 4 registers; supply only 3 + let bytes = encode_state( + ProtoVariant::Regular as i32, + 2, + vec![1, 2, 3], + 0.0, + 0.0, + 0.0, + ); + let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + assert!(result.unwrap_err().to_string().contains("registers")); + } + + #[test] + fn test_from_sketchlib_proto_bytes_zero_precision_rejected() { + use asap_sketchlib::proto::sketchlib::HyperLogLogState; + use prost::Message; + let state = HyperLogLogState::default(); + let bytes = state.encode_to_vec(); + let result = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes); + assert!(result.is_err()); + } + + #[test] + fn test_aggregate_core_merge_matches_register_max() { + let a = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 5, 3, 7], 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let b = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![4, 2, 6, 0], 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let merged_box = a.merge_with(&b).expect("merge ok"); + let merged = merged_box + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.inner.registers, vec![4, 5, 6, 7]); + } + + #[test] + fn test_aggregate_core_merge_wrong_type_rejects() { + use crate::summary_kernels::count_sketch::CountSketchAccumulator; + let hll = HllSketchAccumulator::new(HllVariant::Regular, 2); + let cs = CountSketchAccumulator::new(2, 3); + assert!(hll.merge_with(&cs).is_err()); + } + + #[test] + fn test_from_msgpack_bytes_round_trip() { + let original = HllSketch::from_raw( + HllVariant::Hip, + 3, + vec![0, 1, 2, 3, 4, 5, 6, 7], + 1.5, + 2.5, + 42.0, + ); + let bytes = original.to_msgpack().unwrap(); + let acc = HllSketchAccumulator::from_msgpack_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.inner.variant, HllVariant::Hip); + assert_eq!(acc.inner.precision, 3); + assert_eq!(acc.inner.registers, vec![0, 1, 2, 3, 4, 5, 6, 7]); + assert_eq!(acc.inner.hip_kxq0, 1.5); + } + + #[test] + fn test_from_msgpack_bytes_rejects_garbage() { + let result = HllSketchAccumulator::from_msgpack_bytes(b"not valid msgpack"); + assert!(result.is_err()); + } + + #[test] + fn test_apply_proto_delta_bytes_round_trip() { + use asap_sketchlib::proto::sketchlib::HllDelta as PbDelta; + use prost::Message; + + let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); + acc.inner.registers = vec![1, 5, 3, 7]; + + // Packed (index_delta, value) blob for updates {0:4, 2:6}: + // varint(0),varint(4),varint(2),varint(6). + let delta_bytes = PbDelta { + packed_updates: vec![0, 4, 2, 6], + } + .encode_to_vec(); + + acc.apply_proto_delta_bytes(&delta_bytes).expect("apply ok"); + // Max semantics: reg[0]=max(1,4)=4, reg[2]=max(3,6)=6; others unchanged. + assert_eq!(acc.inner.registers, vec![4, 5, 6, 7]); + } + + #[test] + fn test_apply_proto_delta_bytes_rejects_garbage() { + let mut acc = HllSketchAccumulator::new(HllVariant::Regular, 2); + assert!(acc.apply_proto_delta_bytes(b"not valid proto").is_err()); + } + + // ----- sample_p cardinality rescale ----- + // + // HLL uses hash-threshold sampling: each distinct key is admitted into + // the sketch with probability `p`, so the register-derived cardinality + // estimate is ~p× the true distinct count and must be rescaled by 1/p. + + #[test] + fn test_cardinality_is_rescaled_by_sample_p() { + use crate::Statistic; + // Build two accumulators with identical registers but different + // sample_p. The sampled one (p=0.25) must report ~4× the unsampled + // estimate. Use precision 8 (256 registers) with a spread of + // register values so the estimate is a non-trivial positive number. + let mut registers = vec![0u8; 256]; + for (i, r) in registers.iter_mut().enumerate() { + *r = ((i % 7) + 1) as u8; + } + let unsampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 8, registers.clone(), 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let sampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 8, registers, 0.0, 0.0, 0.0), + sample_p: 0.25, + }; + let raw = unsampled + .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) + .expect("cardinality ok"); + let rescaled = sampled + .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) + .expect("cardinality ok"); + assert!(raw > 0.0, "raw estimate should be positive, got {raw}"); + // Exact algebraic relationship: rescaled == raw / 0.25 == raw * 4. + assert!( + (rescaled - raw * 4.0).abs() < 1e-9, + "expected rescaled ≈ 4×raw ({}), got {rescaled}", + raw * 4.0 + ); + } + + #[test] + fn test_count_statistic_also_rescaled_by_sample_p() { + use crate::Statistic; + // Count maps to the same cardinality estimate for HLL, so it must + // rescale identically. + let registers = vec![3u8; 16]; + let unsampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 4, registers.clone(), 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let sampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 4, registers, 0.0, 0.0, 0.0), + sample_p: 0.25, + }; + let raw = unsampled + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + let rescaled = sampled + .query_statistic(Statistic::Count, &None, &HashMap::new()) + .expect("count ok"); + assert!((rescaled - raw * 4.0).abs() < 1e-9); + } + + #[test] + fn test_sample_p_unset_behaves_as_one() { + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + // An envelope with no sample_p set (proto3 default 0.0) must + // normalize to 1.0 (no rescale) — byte-compatible with legacy frames. + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision: 4, + registers: vec![2u8; 16], + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + registers_sparse: None, + }; + let env = SketchEnvelope { + // sample_p left at proto3 default 0.0. + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 1.0, "unset sample_p must normalize to 1.0"); + } + + #[test] + fn test_from_sketchlib_proto_bytes_reads_envelope_sample_p() { + use crate::Statistic; + use asap_sketchlib::proto::sketchlib::{ + sketch_envelope, HllVariant as ProtoVariant, HyperLogLogState, SketchEnvelope, + }; + use prost::Message; + + let registers = vec![3u8; 16]; + let state = HyperLogLogState { + variant: ProtoVariant::Regular as i32, + precision: 4, + registers: registers.clone(), + hip_kxq0: 0.0, + hip_kxq1: 0.0, + hip_est: 0.0, + registers_sparse: None, + }; + let env = SketchEnvelope { + sample_p: 0.25, + sketch_state: Some(sketch_envelope::SketchState::Hll(state)), + ..Default::default() + }; + let bytes = env.encode_to_vec(); + let acc = HllSketchAccumulator::from_sketchlib_proto_bytes(&bytes).expect("decode ok"); + assert_eq!(acc.sample_p, 0.25); + + // Compare against the unsampled estimate over the same registers. + let unsampled = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 4, registers, 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let raw = unsampled + .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) + .expect("cardinality ok"); + let rescaled = acc + .query_statistic(Statistic::Cardinality, &None, &HashMap::new()) + .expect("cardinality ok"); + assert!( + (rescaled - raw * 4.0).abs() < 1e-9, + "expected 4×raw rescale" + ); + } + + #[test] + fn test_reset_to_empty_preserves_sample_p() { + let mut acc = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 4, vec![3u8; 16], 0.0, 0.0, 0.0), + sample_p: 0.25, + }; + acc.reset_to_empty(); + assert_eq!(acc.sample_p, 0.25, "window rotation must keep sample_p"); + assert_eq!(acc.inner.registers, vec![0u8; 16], "registers cleared"); + } + + #[test] + fn test_merge_prefers_sampled_factor() { + let a = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 1, 1, 1], 0.0, 0.0, 0.0), + sample_p: 0.25, + }; + let b = HllSketchAccumulator { + inner: HllSketch::from_raw(HllVariant::Regular, 2, vec![1, 1, 1, 1], 0.0, 0.0, 0.0), + sample_p: 1.0, + }; + let merged = a.merge_with(&b).expect("merge ok"); + let merged = merged + .as_any() + .downcast_ref::() + .expect("downcast ok"); + assert_eq!(merged.sample_p, 0.25); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/hydra_kll.rs b/crates/asap-physical-operators/src/summary_kernels/hydra_kll.rs new file mode 100644 index 00000000..167bcae3 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/hydra_kll.rs @@ -0,0 +1,165 @@ +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, +}; +use asap_sketchlib::{HydraKllSketch, MessagePackCodec}; +use base64::{engine::general_purpose, Engine as _}; +use std::collections::HashMap; + +use crate::Statistic; + +/// HydraKLL sketch accumulator — wraps asap_sketchlib::HydraKllSketch. +/// Core struct, update/merge/serde logic live in `asap_sketchlib::sketches`. +/// This file retains QE-specific trait impls and JSON output. +#[derive(Debug, Clone)] +pub struct HydraKllSketchAccumulator { + pub inner: HydraKllSketch, +} + +impl HydraKllSketchAccumulator { + pub fn new(row_num: usize, col_num: usize, k: u16) -> Self { + Self { + inner: HydraKllSketch::new(row_num, col_num, k), + } + } + + pub fn update(&mut self, key: &KeyByLabelValues, value: f64) { + self.inner.update(&key.to_semicolon_str(), value); + } + + pub fn deserialize_from_bytes(_buffer: &[u8]) -> Result> { + Err("deserialize_from_bytes for HydraKllSketchAccumulator not implemented".into()) + } + + pub fn query_key(&self, key: &KeyByLabelValues, quantile: f64) -> f64 { + self.inner.quantile(&key.to_semicolon_str(), quantile) + } +} + +impl SerializableToSink for HydraKllSketchAccumulator { + fn serialize_to_json(&self) -> serde_json::Value { + // Mirror Python implementation: {"sketch": base64_encoded_string} + let sketch_bytes = self.inner.to_msgpack().unwrap_or_default(); + let sketch_b64 = general_purpose::STANDARD.encode(&sketch_bytes); + serde_json::json!({ "sketch": sketch_b64 }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_msgpack().unwrap_or_default() + } +} + +impl MergeableAccumulator for HydraKllSketchAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut iter = accumulators.into_iter(); + let mut merged = iter.next().unwrap(); + for acc in iter { + merged.inner.merge(&acc.inner)?; + } + Ok(merged) + } +} + +impl AggregateCore for HydraKllSketchAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "HydraKllSketchAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge HydraKllSketchAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + let hk = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to HydraKllSketchAccumulator")?; + + let merged = Self::merge_accumulators(vec![self.clone(), hk.clone()])?; + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::HydraKLL + } + + fn approx_memory_bytes(&self) -> usize { + // HydraKLL is a row*col grid of KLL sketches; typical instances + // are on the order of tens of KiB. 32 KiB is a conservative + // per-instance default. + 32 * 1024 + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for HydraKllSketchAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for HydraKllSketchAccumulator { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Quantile => { + let quantile = query_kwargs + .and_then(|kwargs| kwargs.get("quantile")) + .ok_or("Missing quantile parameter for quantile query")? + .parse::() + .map_err(|_| "Invalid quantile parameter format")?; + + if !(0.0..=1.0).contains(&quantile) { + return Err("Quantile must be between 0.0 and 1.0".into()); + } + + Ok(self.query_key(key, quantile)) + } + _ => Err( + format!("Unsupported statistic in HydraKllSketchAccumulator: {statistic:?}").into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/increase.rs b/crates/asap-physical-operators/src/summary_kernels/increase.rs new file mode 100644 index 00000000..a6cdb19c --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/increase.rs @@ -0,0 +1,736 @@ +use crate::{ + AggregateCore, AggregationType, Measurement, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +const RESET_AWARE_WIRE_MAGIC: &[u8; 8] = b"ASAPINC2"; +const RESET_AWARE_WIRE_EXTENSION_LEN: usize = 8 + 8 + 8; + +/// Accumulator for tracking increases in counter metrics +/// Stores the starting and last seen measurements with timestamps +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct IncreaseAccumulator { + pub starting_measurement: Measurement, + pub starting_timestamp: i64, + pub last_seen_measurement: Measurement, + pub last_seen_timestamp: i64, + /// Sum of monotonic deltas, adding the post-reset value whenever the + /// counter decreases. This is the reset correction Prometheus applies. + #[serde(default)] + pub total_increase: f64, + #[serde(default)] + pub sample_count: u64, +} + +impl IncreaseAccumulator { + /// Merge two counter intervals without a temporary collection. Ties retain + /// the left input, matching the stable ordering of multi-pane merges. + pub(crate) fn merge_pair(left: &Self, right: &Self) -> Self { + let (first, second) = if left.starting_timestamp <= right.starting_timestamp { + (left, right) + } else { + (right, left) + }; + let mut merged = first.clone(); + if second.starting_timestamp > merged.last_seen_timestamp { + merged.total_increase += + if second.starting_measurement.value >= merged.last_seen_measurement.value { + second.starting_measurement.value - merged.last_seen_measurement.value + } else { + second.starting_measurement.value + }; + } + merged.total_increase += second.total_increase; + merged.sample_count = merged.sample_count.saturating_add(second.sample_count); + if second.last_seen_timestamp > merged.last_seen_timestamp { + merged.last_seen_measurement = second.last_seen_measurement.clone(); + merged.last_seen_timestamp = second.last_seen_timestamp; + } + + merged + } + + /// Return the number of bytes occupied by one accumulator at the start of + /// `buffer`. Old persisted values end after `last_seen_timestamp`; reset- + /// aware values carry a magic-prefixed extension. The magic makes this + /// safe when the buffer also contains the next keyed entry. + pub(crate) fn serialized_len_from_prefix( + buffer: &[u8], + ) -> Result> { + if buffer.len() < 4 { + return Err("Buffer too short for starting measurement length".into()); + } + let starting_len = u32::from_le_bytes(buffer[0..4].try_into()?) as usize; + let last_len_offset = 4usize + .checked_add(starting_len) + .and_then(|offset| offset.checked_add(8)) + .ok_or("IncreaseAccumulator length overflow")?; + if buffer.len() < last_len_offset + 4 { + return Err("Buffer too short for last seen measurement length".into()); + } + let last_len = + u32::from_le_bytes(buffer[last_len_offset..last_len_offset + 4].try_into()?) as usize; + let legacy_len = last_len_offset + .checked_add(4) + .and_then(|offset| offset.checked_add(last_len)) + .and_then(|offset| offset.checked_add(8)) + .ok_or("IncreaseAccumulator length overflow")?; + if buffer.len() < legacy_len { + return Err("Buffer too short for last seen timestamp".into()); + } + let has_extension = buffer.len() >= legacy_len + RESET_AWARE_WIRE_EXTENSION_LEN + && &buffer[legacy_len..legacy_len + RESET_AWARE_WIRE_MAGIC.len()] + == RESET_AWARE_WIRE_MAGIC; + Ok(legacy_len + + if has_extension { + RESET_AWARE_WIRE_EXTENSION_LEN + } else { + 0 + }) + } + + pub fn new( + starting_measurement: Measurement, + starting_timestamp: i64, + last_seen_measurement: Measurement, + last_seen_timestamp: i64, + ) -> Self { + let total_increase = if last_seen_timestamp <= starting_timestamp { + 0.0 + } else if last_seen_measurement.value >= starting_measurement.value { + last_seen_measurement.value - starting_measurement.value + } else { + last_seen_measurement.value + }; + let sample_count = if last_seen_timestamp > starting_timestamp { + 2 + } else { + 1 + }; + Self { + starting_measurement, + starting_timestamp, + last_seen_measurement, + last_seen_timestamp, + total_increase, + sample_count, + } + } + + pub fn update(&mut self, measurement: Measurement, timestamp: i64) { + if timestamp < self.last_seen_timestamp { + return; + } + if timestamp == self.last_seen_timestamp { + return; + } + if measurement.value >= self.last_seen_measurement.value { + self.total_increase += measurement.value - self.last_seen_measurement.value; + } else { + self.total_increase += measurement.value; + } + self.last_seen_measurement = measurement; + self.last_seen_timestamp = timestamp; + self.sample_count = self.sample_count.saturating_add(1); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let starting_measurement = + Measurement::deserialize_from_json(&data["starting_measurement"])?; + let starting_timestamp = data["starting_timestamp"] + .as_i64() + .ok_or("Missing or invalid 'starting_timestamp' field")?; + let last_seen_measurement = + Measurement::deserialize_from_json(&data["last_seen_measurement"])?; + let last_seen_timestamp = data["last_seen_timestamp"] + .as_i64() + .ok_or("Missing or invalid 'last_seen_timestamp' field")?; + + let mut accumulator = Self::new( + starting_measurement, + starting_timestamp, + last_seen_measurement, + last_seen_timestamp, + ); + accumulator.total_increase = data["total_increase"] + .as_f64() + .unwrap_or(accumulator.total_increase); + accumulator.sample_count = data["sample_count"] + .as_u64() + .unwrap_or(accumulator.sample_count); + Ok(accumulator) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut offset = 0; + + // Read starting measurement length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for starting measurement length".into()); + } + let starting_measurement_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + starting_measurement_length { + return Err("Buffer too short for starting measurement".into()); + } + let starting_measurement = Measurement::deserialize_from_bytes( + &buffer[offset..offset + starting_measurement_length], + )?; + offset += starting_measurement_length; + + // Read starting timestamp + if buffer.len() < offset + 8 { + return Err("Buffer too short for starting timestamp".into()); + } + let starting_timestamp = i64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + + // Read last seen measurement length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for last seen measurement length".into()); + } + let last_seen_measurement_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + last_seen_measurement_length { + return Err("Buffer too short for last seen measurement".into()); + } + let last_seen_measurement = Measurement::deserialize_from_bytes( + &buffer[offset..offset + last_seen_measurement_length], + )?; + offset += last_seen_measurement_length; + + // Read last seen timestamp + if buffer.len() < offset + 8 { + return Err("Buffer too short for last seen timestamp".into()); + } + let last_seen_timestamp = i64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + + let mut accumulator = Self::new( + starting_measurement, + starting_timestamp, + last_seen_measurement, + last_seen_timestamp, + ); + offset += 8; + if buffer.len() >= offset + RESET_AWARE_WIRE_EXTENSION_LEN + && &buffer[offset..offset + RESET_AWARE_WIRE_MAGIC.len()] == RESET_AWARE_WIRE_MAGIC + { + offset += RESET_AWARE_WIRE_MAGIC.len(); + accumulator.total_increase = f64::from_le_bytes( + buffer[offset..offset + 8] + .try_into() + .expect("checked total-increase bytes"), + ); + offset += 8; + accumulator.sample_count = u64::from_le_bytes( + buffer[offset..offset + 8] + .try_into() + .expect("checked sample-count bytes"), + ); + } + Ok(accumulator) + } +} + +impl SerializableToSink for IncreaseAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "starting_measurement": self.starting_measurement.serialize_to_json(), + "starting_timestamp": self.starting_timestamp, + "last_seen_measurement": self.last_seen_measurement.serialize_to_json(), + "last_seen_timestamp": self.last_seen_timestamp, + "total_increase": self.total_increase, + "sample_count": self.sample_count, + }) + } + + fn serialize_to_bytes(&self) -> Vec { + let starting_measurement_bytes = self.starting_measurement.serialize_to_bytes(); + let last_seen_measurement_bytes = self.last_seen_measurement.serialize_to_bytes(); + + let mut buffer = Vec::new(); + + // Starting measurement length and data + buffer.extend_from_slice(&(starting_measurement_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&starting_measurement_bytes); + + // Starting timestamp + buffer.extend_from_slice(&self.starting_timestamp.to_le_bytes()); + + // Last seen measurement length and data + buffer.extend_from_slice(&(last_seen_measurement_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&last_seen_measurement_bytes); + + // Last seen timestamp + buffer.extend_from_slice(&self.last_seen_timestamp.to_le_bytes()); + buffer.extend_from_slice(RESET_AWARE_WIRE_MAGIC); + buffer.extend_from_slice(&self.total_increase.to_le_bytes()); + buffer.extend_from_slice(&self.sample_count.to_le_bytes()); + + buffer + } +} + +impl MergeableAccumulator for IncreaseAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut accumulators = accumulators; + accumulators.sort_by_key(|accumulator| accumulator.starting_timestamp); + let mut result = accumulators[0].clone(); + + for acc in &accumulators[1..] { + result = Self::merge_pair(&result, acc); + } + + Ok(result) + } +} + +impl AggregateCore for IncreaseAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "IncreaseAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + // Check if other is also an IncreaseAccumulator + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge IncreaseAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + // Downcast to IncreaseAccumulator + let other_increase = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to IncreaseAccumulator")?; + + let merged = Self::merge_pair(self, other_increase); + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Increase + } + + fn approx_memory_bytes(&self) -> usize { + // Two Measurements + two i64s. Measurements are a few f64 fields. + std::mem::size_of::() + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query( + statistic, + (!query_kwargs.is_empty()).then_some(query_kwargs), + ) + } +} + +impl SingleSubpopulationAggregate for IncreaseAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Increase => Ok(self.extrapolated_value(query_kwargs, false)?), + Statistic::Rate => Ok(self.extrapolated_value(query_kwargs, true)?), + // For instant `sum [by (...)] (counter_metric)` Prometheus + // sums the latest cumulative value of each matching series. + // The IncreaseAccumulator already tracks that latest value + // in `last_seen_measurement`, so per-series Sum is just + // that scalar; the engine's outer aggregation groups by the + // `by` labels and adds the per-series totals across keys. + // + // See PR #108 audit conclusion (commit 4359e10) and issue + // ProjectASAP/ASAPCollector#46: pre-fix the ASAP tier ingested + // counters as IncreaseAccumulator and bare `sum by (...) ()` + // capability-missed because this trait did not answer Sum. + Statistic::Sum => Ok(self.last_seen_measurement.value), + _ => Err(format!("Unsupported statistic in IncreaseAccumulator: {statistic:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl IncreaseAccumulator { + fn extrapolated_value( + &self, + query_kwargs: Option<&HashMap>, + is_rate: bool, + ) -> Result> { + if self.sample_count < 2 || self.last_seen_timestamp <= self.starting_timestamp { + return Err("at least two ordered counter samples are required".into()); + } + let sampled_interval = (self.last_seen_timestamp - self.starting_timestamp) as f64 / 1000.0; + let Some(kwargs) = query_kwargs else { + return Ok(if is_rate { + self.total_increase / sampled_interval + } else { + self.total_increase + }); + }; + let range_start = kwargs + .get("range_start_ms") + .ok_or("missing range_start_ms")? + .parse::()?; + let range_end = kwargs + .get("range_end_ms") + .ok_or("missing range_end_ms")? + .parse::()?; + if range_end <= range_start { + return Err("invalid counter evaluation range".into()); + } + + let mut duration_to_start = + (self.starting_timestamp.saturating_sub(range_start)) as f64 / 1000.0; + let duration_to_end = (range_end.saturating_sub(self.last_seen_timestamp)) as f64 / 1000.0; + let average_sample_interval = sampled_interval / (self.sample_count - 1) as f64; + let extrapolation_threshold = average_sample_interval * 1.1; + + if self.total_increase > 0.0 && self.starting_measurement.value >= 0.0 { + let duration_to_zero = + sampled_interval * (self.starting_measurement.value / self.total_increase); + duration_to_start = duration_to_start.min(duration_to_zero); + } + let mut extrapolate_to = sampled_interval; + extrapolate_to += if duration_to_start < extrapolation_threshold { + duration_to_start.max(0.0) + } else { + average_sample_interval / 2.0 + }; + extrapolate_to += if duration_to_end < extrapolation_threshold { + duration_to_end.max(0.0) + } else { + average_sample_interval / 2.0 + }; + let mut factor = extrapolate_to / sampled_interval; + if is_rate { + factor /= (range_end - range_start) as f64 / 1000.0; + } + Ok(self.total_increase * factor) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_increase_accumulator_creation() { + let starting_measurement = Measurement::new(10.0); + let last_seen_measurement = Measurement::new(25.0); + let acc = IncreaseAccumulator::new( + starting_measurement.clone(), + 1000, + last_seen_measurement.clone(), + 2000, + ); + + assert_eq!(acc.starting_measurement.value, 10.0); + assert_eq!(acc.starting_timestamp, 1000); + assert_eq!(acc.last_seen_measurement.value, 25.0); + assert_eq!(acc.last_seen_timestamp, 2000); + } + + #[test] + fn test_increase_accumulator_update() { + let starting_measurement = Measurement::new(10.0); + let mut acc = IncreaseAccumulator::new( + starting_measurement.clone(), + 1000, + starting_measurement.clone(), + 1000, + ); + + let new_measurement = Measurement::new(25.0); + acc.update(new_measurement.clone(), 2000); + + assert_eq!(acc.last_seen_measurement.value, 25.0); + assert_eq!(acc.last_seen_timestamp, 2000); + assert_eq!(acc.starting_measurement.value, 10.0); // Should remain unchanged + } + + #[test] + fn test_increase_accumulator_query() { + let starting_measurement = Measurement::new(10.0); + let last_seen_measurement = Measurement::new(25.0); + let acc = IncreaseAccumulator::new( + starting_measurement, + 1000, + last_seen_measurement, + 3000, // 2 second difference + ); + + // Test increase calculation + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, None).unwrap(), + 15.0 + ); + + // Test rate calculation (per second) + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, None).unwrap(), + 7.5 + ); // 15.0 / 2.0 + + // Statistic::Sum returns the latest cumulative counter value, + // matching Prometheus semantics for instant `sum()`. + // (Issue ProjectASAP/ASAPCollector#46, PR #108 diagnosis.) + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), + 25.0 + ); + + // Unsupported statistics still error. + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); + } + + #[test] + fn prometheus_counter_reset_and_boundary_extrapolation() { + let mut acc = IncreaseAccumulator::new( + Measurement::new(10.0), + 10_000, + Measurement::new(10.0), + 10_000, + ); + acc.update(Measurement::new(20.0), 20_000); + acc.update(Measurement::new(3.0), 30_000); + acc.update(Measurement::new(13.0), 50_000); + assert_eq!(acc.total_increase, 23.0); + assert_eq!(acc.sample_count, 4); + + let kwargs = HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + let increase = + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Increase, Some(&kwargs)) + .unwrap(); + let rate = crate::SingleSubpopulationAggregate::query(&acc, Statistic::Rate, Some(&kwargs)) + .unwrap(); + assert!((increase - 34.5).abs() < 1e-12); + assert!((rate - 0.575).abs() < 1e-12); + } + + #[test] + fn pane_merge_preserves_resets_and_prometheus_extrapolation() { + let mut left = IncreaseAccumulator::new( + Measurement::new(10.0), + 10_000, + Measurement::new(10.0), + 10_000, + ); + left.update(Measurement::new(20.0), 20_000); + let mut right = + IncreaseAccumulator::new(Measurement::new(3.0), 30_000, Measurement::new(3.0), 30_000); + right.update(Measurement::new(13.0), 50_000); + let merged = IncreaseAccumulator::merge_accumulators(vec![right, left]).unwrap(); + assert_eq!(merged.total_increase, 23.0); + assert_eq!(merged.sample_count, 4); + let kwargs = HashMap::from([ + ("range_start_ms".into(), "0".into()), + ("range_end_ms".into(), "60000".into()), + ]); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&merged, Statistic::Increase, Some(&kwargs)) + .unwrap(), + 34.5 + ); + } + + #[test] + fn counter_sds_state_is_constant_size_per_pane() { + let mut acc = IncreaseAccumulator::new(Measurement::new(0.0), 0, Measurement::new(0.0), 0); + let initial = acc.serialize_to_bytes().len(); + for second in 1..=86_400 { + acc.update(Measurement::new(second as f64), second * 1_000); + } + assert_eq!(acc.serialize_to_bytes().len(), initial); + assert_eq!(acc.sample_count, 86_401); + assert_eq!( + acc.approx_memory_bytes(), + std::mem::size_of::() + ); + } + + #[test] + fn test_increase_accumulator_sum_is_latest_cumulative_value() { + // Instant `sum ()` semantics: the per-series summand is + // the latest cumulative counter value. Two series with latest + // values 100 and 50 (started at 10 and 5 respectively) should + // each report Sum = 100 and Sum = 50 — the engine's `sum by` + // outer aggregation does the cross-series total. + let acc_a = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000); + let acc_b = + IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc_a, Statistic::Sum, None).unwrap(), + 100.0 + ); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc_b, Statistic::Sum, None).unwrap(), + 50.0 + ); + } + + #[test] + fn test_increase_accumulator_merge() { + let acc1 = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(20.0), 2000); + let acc2 = IncreaseAccumulator::new( + Measurement::new(5.0), + 500, // Earlier start + Measurement::new(15.0), + 1500, + ); + let acc3 = IncreaseAccumulator::new( + Measurement::new(20.0), + 2000, + Measurement::new(30.0), + 3000, // Later end + ); + + let merged = + >::merge_accumulators( + vec![acc1, acc2, acc3], + ) + .unwrap(); + + // Should use earliest start and latest end + assert_eq!(merged.starting_measurement.value, 5.0); + assert_eq!(merged.starting_timestamp, 500); + assert_eq!(merged.last_seen_measurement.value, 30.0); + assert_eq!(merged.last_seen_timestamp, 3000); + } + + #[test] + fn test_increase_accumulator_serialization() { + let acc = + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(25.0), 2000); + + // Test JSON serialization + let json = acc.serialize_to_json(); + let deserialized = IncreaseAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!( + acc.starting_measurement.value, + deserialized.starting_measurement.value + ); + assert_eq!(acc.starting_timestamp, deserialized.starting_timestamp); + assert_eq!( + acc.last_seen_measurement.value, + deserialized.last_seen_measurement.value + ); + assert_eq!(acc.last_seen_timestamp, deserialized.last_seen_timestamp); + + // Test byte serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = IncreaseAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + acc.starting_measurement.value, + deserialized_bytes.starting_measurement.value + ); + assert_eq!( + acc.starting_timestamp, + deserialized_bytes.starting_timestamp + ); + assert_eq!( + acc.last_seen_measurement.value, + deserialized_bytes.last_seen_measurement.value + ); + assert_eq!( + acc.last_seen_timestamp, + deserialized_bytes.last_seen_timestamp + ); + assert_eq!(acc.total_increase, deserialized_bytes.total_increase); + assert_eq!(acc.sample_count, deserialized_bytes.sample_count); + + let legacy = &bytes[..bytes.len() - RESET_AWARE_WIRE_EXTENSION_LEN]; + let legacy_value = IncreaseAccumulator::deserialize_from_bytes(legacy).unwrap(); + assert_eq!(legacy_value.total_increase, 15.0); + assert_eq!(legacy_value.sample_count, 2); + } + + #[test] + fn test_trait_object() { + let acc: Box = Box::new(IncreaseAccumulator::new( + Measurement::new(10.0), + 1000, + Measurement::new(25.0), + 2000, + )); + + assert_eq!(acc.type_name(), "IncreaseAccumulator"); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/keyed_counter_state.rs b/crates/asap-physical-operators/src/summary_kernels/keyed_counter_state.rs new file mode 100644 index 00000000..8db0ee7e --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/keyed_counter_state.rs @@ -0,0 +1,529 @@ +use crate::summary_kernels::IncreaseAccumulator; +use crate::{ + AggregateCore, AggregationType, KeyByLabelValues, MergeableAccumulator, + MultipleSubpopulationAggregate, SerializableToSink, SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Accumulator that maintains separate increase accumulators for multiple keys +/// Allows tracking rate/increase for different label combinations +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KeyedCounterState { + pub increases: HashMap, +} + +impl KeyedCounterState { + pub fn new() -> Self { + Self { + increases: HashMap::new(), + } + } + + pub fn update(&mut self, key: KeyByLabelValues, accumulator: IncreaseAccumulator) { + self.increases.insert(key, accumulator); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let mut accumulator = Self::new(); + + if let Some(entries) = data["entries"].as_array() { + for entry in entries { + let key = KeyByLabelValues::deserialize_from_json(&entry["key"])?; + let increase_data = + IncreaseAccumulator::deserialize_from_json(&entry["increase_data"])?; + accumulator.increases.insert(key, increase_data); + } + } + + Ok(accumulator) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut accumulator = Self::new(); + let mut offset = 0; + + // Read number of entries + if buffer.len() < 4 { + return Err("Buffer too short for entry count".into()); + } + let num_entries = u32::from_le_bytes([buffer[0], buffer[1], buffer[2], buffer[3]]) as usize; + offset += 4; + + for _ in 0..num_entries { + // Read key length and key + if offset + 4 > buffer.len() { + return Err("Buffer too short for key length".into()); + } + let key_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if offset + key_length > buffer.len() { + return Err("Buffer too short for key data".into()); + } + let key = + KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; + offset += key_length; + + // Read IncreaseAccumulator data + if offset >= buffer.len() { + return Err("Buffer too short for increase accumulator data".into()); + } + let consumed_bytes = + IncreaseAccumulator::serialized_len_from_prefix(&buffer[offset..])?; + let increase_data = IncreaseAccumulator::deserialize_from_bytes( + &buffer[offset..offset + consumed_bytes], + )?; + offset += consumed_bytes; + + accumulator.increases.insert(key, increase_data); + } + + Ok(accumulator) + } +} + +impl Default for KeyedCounterState { + fn default() -> Self { + Self::new() + } +} + +impl SerializableToSink for KeyedCounterState { + fn serialize_to_json(&self) -> Value { + let entries: Vec = self + .increases + .iter() + .map(|(key, data)| { + serde_json::json!({ + "key": key.serialize_to_json(), + "increase_data": data.serialize_to_json() + }) + }) + .collect(); + + serde_json::json!({ + "entries": entries + }) + } + + fn serialize_to_bytes(&self) -> Vec { + let mut buffer = Vec::new(); + + // Write number of entries + buffer.extend_from_slice(&(self.increases.len() as u32).to_le_bytes()); + + // Write each key-value pair + for (key, data) in &self.increases { + let key_bytes = key.serialize_to_bytes(); + buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&key_bytes); + + let data_bytes = data.serialize_to_bytes(); + buffer.extend_from_slice(&data_bytes); + } + + buffer + } +} + +impl AggregateCore for KeyedCounterState { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "KeyedCounterState" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + // Check if other is also a KeyedCounterState + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge KeyedCounterState with {}", + other.get_accumulator_type() + ) + .into()); + } + + // Downcast to KeyedCounterState + let other_multiple_increase = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to KeyedCounterState")?; + + // Clone self once, then merge each matching counter with the same + // reset-aware, boundary-aware implementation used by the unkeyed path. + let mut merged = self.clone(); + for (key, data) in &other_multiple_increase.increases { + if let Some(existing_data) = merged.increases.get_mut(key) { + *existing_data = IncreaseAccumulator::merge_accumulators(vec![ + existing_data.clone(), + data.clone(), + ])?; + } else { + merged.increases.insert(key.clone(), data.clone()); + } + } + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Increase + } + + fn approx_memory_bytes(&self) -> usize { + // HashMap. IncreaseAccumulator is ~64 B, + // per-entry key/overhead is ~96 B. + const BYTES_PER_ENTRY: usize = 160; + std::mem::size_of::() + self.increases.len() * BYTES_PER_ENTRY + } + + fn get_keys(&self) -> Option> { + Some(self.increases.keys().cloned().collect()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedCounterState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedCounterState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + query_kwargs: Option<&HashMap>, + ) -> Result> { + let data = self + .increases + .get(key) + .ok_or_else(|| format!("Key {key} not found in KeyedCounterState"))?; + + data.query(statistic, query_kwargs) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedCounterState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedCounterState::new(); + + for accumulator in accumulators { + for (key, data) in accumulator.increases { + if let Some(existing_data) = result.increases.get_mut(&key) { + *existing_data = + IncreaseAccumulator::merge_accumulators(vec![existing_data.clone(), data])?; + } else { + result.increases.insert(key, data); + } + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::Measurement; + + fn create_test_increase_accumulator(start_val: f64, end_val: f64) -> IncreaseAccumulator { + IncreaseAccumulator::new( + Measurement::new(start_val), + 1000, + Measurement::new(end_val), + 2000, + ) + } + + fn create_test_increase_accumulator_with_time( + start_val: f64, + start_time: i64, + end_val: f64, + end_time: i64, + ) -> IncreaseAccumulator { + IncreaseAccumulator::new( + Measurement::new(start_val), + start_time, + Measurement::new(end_val), + end_time, + ) + } + + #[test] + fn test_keyed_counter_state_creation() { + let acc = KeyedCounterState::new(); + assert!(acc.increases.is_empty()); + } + + #[test] + fn test_keyed_counter_state_update() { + let mut acc = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + let increase1 = create_test_increase_accumulator(10.0, 25.0); + let increase2 = create_test_increase_accumulator(5.0, 15.0); + + acc.update(key1.clone(), increase1); + acc.update(key2.clone(), increase2); + + assert_eq!(acc.increases.len(), 2); + assert!(acc.increases.contains_key(&key1)); + assert!(acc.increases.contains_key(&key2)); + } + + #[test] + fn test_keyed_counter_state_query() { + let mut acc = KeyedCounterState::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let increase_acc = create_test_increase_accumulator(10.0, 25.0); + acc.update(key.clone(), increase_acc); + + // Test increase query + assert_eq!(acc.query(Statistic::Increase, &key, None).unwrap(), 15.0); + + // Test rate query (15.0 increase over 1 second = 15.0 per second) + assert_eq!(acc.query(Statistic::Rate, &key, None).unwrap(), 15.0); + + // Sum returns the latest cumulative counter value for the + // queried key (per-series Prometheus `sum()` semantics; + // see issue ProjectASAP/ASAPCollector#46 and PR #108 diagnosis). + // The series here was created with last_seen=25.0. + assert_eq!(acc.query(Statistic::Sum, &key, None).unwrap(), 25.0); + + // Unsupported statistic still errors. + assert!(acc.query(Statistic::Min, &key, None).is_err()); + + let unknown_key = KeyByLabelValues::new(); + assert!(acc.query(Statistic::Increase, &unknown_key, None).is_err()); + } + + #[test] + fn test_keyed_counter_state_sum_per_key() { + // `sum by (zone) (counter)` reaches KeyedCounterState + // only when the ASAP-tier ingest groups multiple series under + // a single accumulator (the `Multiple*` variant). In that case + // each per-key Sum should be the series' latest cumulative + // value; the engine's outer `by` aggregation does the cross-key + // grouping. (Issue ProjectASAP/ASAPCollector#46.) + let mut acc = KeyedCounterState::new(); + let east = KeyByLabelValues::new_with_labels(vec!["us-east-1".to_string()]); + let west = KeyByLabelValues::new_with_labels(vec!["us-west-2".to_string()]); + + acc.update( + east.clone(), + IncreaseAccumulator::new(Measurement::new(10.0), 1000, Measurement::new(100.0), 2000), + ); + acc.update( + west.clone(), + IncreaseAccumulator::new(Measurement::new(5.0), 1000, Measurement::new(50.0), 2000), + ); + + assert_eq!(acc.query(Statistic::Sum, &east, None).unwrap(), 100.0); + assert_eq!(acc.query(Statistic::Sum, &west, None).unwrap(), 50.0); + } + + #[test] + fn test_keyed_counter_state_merge() { + let mut acc1 = KeyedCounterState::new(); + let mut acc2 = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + // Add different keys to each accumulator + acc1.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); + acc2.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); + + // Also add overlapping key with different time ranges (later timestamps) + acc2.update( + key1.clone(), + create_test_increase_accumulator_with_time(15.0, 2000, 30.0, 3000), + ); // Later time range + + let merged = KeyedCounterState::merge_accumulators(vec![acc1, acc2]).unwrap(); + + assert_eq!(merged.increases.len(), 2); + assert!(merged.increases.contains_key(&key1)); + assert!(merged.increases.contains_key(&key2)); + + // The merged key1 should have the full range (earliest start to latest end) + let merged_key1 = merged.increases.get(&key1).unwrap(); + assert_eq!(merged_key1.starting_measurement.value, 10.0); // Earlier start + assert_eq!(merged_key1.last_seen_measurement.value, 30.0); // Later end + } + + #[test] + fn test_keyed_counter_state_serialization() { + let mut acc = KeyedCounterState::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let second_key = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + let mut reset_aware = create_test_increase_accumulator(10.0, 25.0); + reset_aware.update(Measurement::new(3.0), 3000); + acc.update(key.clone(), reset_aware); + acc.update( + second_key.clone(), + create_test_increase_accumulator(4.0, 9.0), + ); + + // Test JSON serialization + let json_value = acc.serialize_to_json(); + let deserialized = KeyedCounterState::deserialize_from_json(&json_value).unwrap(); + + assert_eq!(deserialized.increases.len(), 2); + let deserialized_acc = deserialized.increases.get(&key).unwrap(); + assert_eq!(deserialized_acc.starting_measurement.value, 10.0); + assert_eq!(deserialized_acc.last_seen_measurement.value, 3.0); + assert_eq!(deserialized_acc.total_increase, 18.0); + + // Test binary serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = KeyedCounterState::deserialize_from_bytes(&bytes).unwrap(); + + assert_eq!(deserialized_bytes.increases.len(), 2); + let deserialized_acc_bytes = deserialized_bytes.increases.get(&key).unwrap(); + assert_eq!(deserialized_acc_bytes.starting_measurement.value, 10.0); + assert_eq!(deserialized_acc_bytes.last_seen_measurement.value, 3.0); + assert_eq!(deserialized_acc_bytes.total_increase, 18.0); + assert_eq!( + deserialized_bytes + .increases + .get(&second_key) + .unwrap() + .last_seen_measurement + .value, + 9.0 + ); + } + + #[test] + fn test_keyed_counter_state_get_keys() { + let mut acc = KeyedCounterState::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.update(key1.clone(), create_test_increase_accumulator(10.0, 20.0)); + acc.update(key2.clone(), create_test_increase_accumulator(5.0, 15.0)); + + let keys = acc.get_keys().unwrap(); + assert_eq!(keys.len(), 2); + assert!(keys.contains(&key1)); + assert!(keys.contains(&key2)); + } + + #[test] + fn test_trait_object() { + let mut acc = KeyedCounterState::new(); + let key = KeyByLabelValues::new(); + acc.update(key.clone(), create_test_increase_accumulator(10.0, 25.0)); + + let trait_obj: Box = Box::new(acc); + assert_eq!( + trait_obj.query(Statistic::Increase, &key, None).unwrap(), + 15.0 + ); + + let keys = trait_obj.get_keys().unwrap(); + assert_eq!(keys.len(), 1); + } + + // #[test] + // fn test_keyed_counter_state_arroyo_deserialization() { + // // Create test data in Arroyo MessagePack format + // // Format: {key: [starting_value, starting_timestamp, last_seen_value, last_seen_timestamp]} + // let mut test_data = std::collections::HashMap::new(); + // test_data.insert("web;service".to_string(), vec![10.0, 1000.0, 25.0, 2000.0]); + // test_data.insert("api;service".to_string(), vec![5.0, 1500.0, 15.0, 2500.0]); + + // // Serialize to MessagePack + // let arroyo_buffer = rmp_serde::to_vec(&test_data).unwrap(); + + // // Test Arroyo deserialization + // let deserialized_acc = + // KeyedCounterState::deserialize_from_bytes_arroyo(&arroyo_buffer).unwrap(); + + // // Verify the deserialized accumulator has the correct data + // assert_eq!(deserialized_acc.increases.len(), 2); + + // // Check first key (web;service) + // let keys: Vec<_> = deserialized_acc.increases.keys().collect(); + // let key1 = keys + // .iter() + // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "web")) + // .unwrap(); + + // let increase1 = deserialized_acc.increases.get(key1).unwrap(); + // assert_eq!(increase1.starting_measurement.value, 10.0); + // assert_eq!(increase1.starting_timestamp, 1000); + // assert_eq!(increase1.last_seen_measurement.value, 25.0); + // assert_eq!(increase1.last_seen_timestamp, 2000); + + // // Check second key (api;service) + // let key2 = keys + // .iter() + // .find(|k| k.labels.get("label_0").is_some_and(|v| v == "api")) + // .unwrap(); + + // let increase2 = deserialized_acc.increases.get(key2).unwrap(); + // assert_eq!(increase2.starting_measurement.value, 5.0); + // assert_eq!(increase2.starting_timestamp, 1500); + // assert_eq!(increase2.last_seen_measurement.value, 15.0); + // assert_eq!(increase2.last_seen_timestamp, 2500); + + // // Test querying + // assert_eq!( + // deserialized_acc.query(Statistic::Increase, key1).unwrap(), + // 15.0 + // ); // 25.0 - 10.0 + // assert_eq!( + // deserialized_acc.query(Statistic::Increase, key2).unwrap(), + // 10.0 + // ); // 15.0 - 5.0 + // } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/keyed_max_state.rs b/crates/asap-physical-operators/src/summary_kernels/keyed_max_state.rs new file mode 100644 index 00000000..304b3a37 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/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 crate::Statistic; + +/// Exact per-key maximum over many populations, mergeable by comparison. +/// +/// The minimum direction is +/// [`KeyedMinState`](super::keyed_min_state::KeyedMinState), +/// a separate type: these used to be one `MultipleMinMaxAccumulator` whose +/// direction lived in a `sub_type` string that every layer above had to carry +/// alongside the family. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct KeyedMaxState { + pub values: HashMap, +} + +impl KeyedMaxState { + pub fn new() -> Self { + Self::default() + } + + pub fn new_with_values(values: HashMap) -> Self { + Self { values } + } + + pub fn update(&mut self, key: KeyByLabelValues, value: f64) { + let current = self.values.entry(key).or_insert(f64::NEG_INFINITY); + if value > *current { + *current = value; + } + } + + pub fn add_value(&mut self, key: KeyByLabelValues, value: f64) { + self.values.insert(key, value); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let values_data = data["values"] + .as_object() + .ok_or("Missing or invalid 'values' field")?; + + let mut values = HashMap::new(); + for (key_str, value) in values_data { + let key_json: Value = serde_json::from_str(key_str)?; + let key = KeyByLabelValues::deserialize_from_json(&key_json)?; + let val = value.as_f64().ok_or("Invalid value")?; + values.insert(key, val); + } + + Ok(Self { values }) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut offset = 0; + + // Read number of entries + if buffer.len() < 4 { + return Err("Buffer too short for entry count".into()); + } + let num_entries = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + let mut values = HashMap::new(); + + for _ in 0..num_entries { + // Read key length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for key length".into()); + } + let key_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + key_length { + return Err("Buffer too short for key data".into()); + } + let key = + KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; + offset += key_length; + + // Read value + if buffer.len() < offset + 8 { + return Err("Buffer too short for value".into()); + } + let value = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + + values.insert(key, value); + } + + Ok(Self { values }) + } +} + +impl SerializableToSink for KeyedMaxState { + fn serialize_to_json(&self) -> Value { + let mut values_obj = serde_json::Map::new(); + for (key, value) in &self.values { + let key_json = key.serialize_to_json(); + let key_str = serde_json::to_string(&key_json).unwrap(); + values_obj.insert( + key_str, + Value::Number(serde_json::Number::from_f64(*value).unwrap()), + ); + } + + serde_json::json!({ "values": values_obj }) + } + + fn serialize_to_bytes(&self) -> Vec { + let mut buffer = Vec::new(); + + // Write number of entries + buffer.extend_from_slice(&(self.values.len() as u32).to_le_bytes()); + + // Write each key-value pair + for (key, value) in &self.values { + let key_bytes = key.serialize_to_bytes(); + + // Write key length and data + buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&key_bytes); + + // Write value + buffer.extend_from_slice(&value.to_le_bytes()); + } + + buffer + } +} + +impl AggregateCore for KeyedMaxState { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "KeyedMaxState" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge KeyedMaxState with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_multiple = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to KeyedMaxState")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_multiple.clone()])?; + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Max + } + + fn approx_memory_bytes(&self) -> usize { + const BYTES_PER_ENTRY: usize = 96; + std::mem::size_of::() + self.values.len() * BYTES_PER_ENTRY + } + + fn get_keys(&self) -> Option> { + Some(self.values.keys().cloned().collect()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedMaxState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedMaxState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Max => self + .values + .get(key) + .copied() + .ok_or_else(|| format!("Key {key} not found in KeyedMaxState").into()), + other => Err(format!("Unsupported statistic in KeyedMaxState: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedMaxState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedMaxState::new(); + + for acc in accumulators { + for (key, value) in acc.values { + result.update(key, value); + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(value: &str) -> KeyByLabelValues { + KeyByLabelValues::new_with_labels(vec![value.to_string()]) + } + + #[test] + fn keeps_the_largest_per_key() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 10.0); + acc.update(key("a"), 5.0); + acc.update(key("a"), 15.0); + acc.update(key("b"), 7.0); + + assert_eq!(acc.query(Statistic::Max, &key("a"), None).unwrap(), 15.0); + assert_eq!(acc.query(Statistic::Max, &key("b"), None).unwrap(), 7.0); + } + + #[test] + fn refuses_the_opposite_statistic_and_unknown_keys() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 1.0); + assert!(acc.query(Statistic::Min, &key("a"), None).is_err()); + assert!(acc.query(Statistic::Max, &key("missing"), None).is_err()); + } + + #[test] + fn merges_per_key() { + let mut left = KeyedMaxState::new(); + left.update(key("a"), 10.0); + let mut right = KeyedMaxState::new(); + right.update(key("a"), 5.0); + right.update(key("b"), 3.0); + + let merged = + >::merge_accumulators(vec![ + left, right, + ]) + .unwrap(); + + assert_eq!(merged.query(Statistic::Max, &key("a"), None).unwrap(), 10.0); + assert_eq!(merged.query(Statistic::Max, &key("b"), None).unwrap(), 3.0); + } + + #[test] + fn refuses_to_merge_with_the_opposite_direction() { + use super::super::keyed_min_state::KeyedMinState; + let mine = KeyedMaxState::new(); + let theirs = KeyedMinState::new(); + assert!(mine.merge_with(&theirs).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let mut acc = KeyedMaxState::new(); + acc.update(key("a"), 4.0); + + let json = acc.serialize_to_json(); + let from_json = KeyedMaxState::deserialize_from_json(&json).unwrap(); + assert_eq!( + from_json.query(Statistic::Max, &key("a"), None).unwrap(), + 4.0 + ); + + let bytes = acc.serialize_to_bytes(); + let from_bytes = KeyedMaxState::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + from_bytes.query(Statistic::Max, &key("a"), None).unwrap(), + 4.0 + ); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/keyed_min_state.rs b/crates/asap-physical-operators/src/summary_kernels/keyed_min_state.rs new file mode 100644 index 00000000..5c40da24 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/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 crate::Statistic; + +/// Exact per-key minimum over many populations, mergeable by comparison. +/// +/// The maximum direction is +/// [`KeyedMaxState`](super::keyed_max_state::KeyedMaxState), +/// a separate type: these used to be one `MultipleMinMaxAccumulator` whose +/// direction lived in a `sub_type` string that every layer above had to carry +/// alongside the family. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct KeyedMinState { + pub values: HashMap, +} + +impl KeyedMinState { + pub fn new() -> Self { + Self::default() + } + + pub fn new_with_values(values: HashMap) -> Self { + Self { values } + } + + pub fn update(&mut self, key: KeyByLabelValues, value: f64) { + let current = self.values.entry(key).or_insert(f64::INFINITY); + if value < *current { + *current = value; + } + } + + pub fn add_value(&mut self, key: KeyByLabelValues, value: f64) { + self.values.insert(key, value); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let values_data = data["values"] + .as_object() + .ok_or("Missing or invalid 'values' field")?; + + let mut values = HashMap::new(); + for (key_str, value) in values_data { + let key_json: Value = serde_json::from_str(key_str)?; + let key = KeyByLabelValues::deserialize_from_json(&key_json)?; + let val = value.as_f64().ok_or("Invalid value")?; + values.insert(key, val); + } + + Ok(Self { values }) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut offset = 0; + + // Read number of entries + if buffer.len() < 4 { + return Err("Buffer too short for entry count".into()); + } + let num_entries = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + let mut values = HashMap::new(); + + for _ in 0..num_entries { + // Read key length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for key length".into()); + } + let key_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + key_length { + return Err("Buffer too short for key data".into()); + } + let key = + KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; + offset += key_length; + + // Read value + if buffer.len() < offset + 8 { + return Err("Buffer too short for value".into()); + } + let value = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + + values.insert(key, value); + } + + Ok(Self { values }) + } +} + +impl SerializableToSink for KeyedMinState { + fn serialize_to_json(&self) -> Value { + let mut values_obj = serde_json::Map::new(); + for (key, value) in &self.values { + let key_json = key.serialize_to_json(); + let key_str = serde_json::to_string(&key_json).unwrap(); + values_obj.insert( + key_str, + Value::Number(serde_json::Number::from_f64(*value).unwrap()), + ); + } + + serde_json::json!({ "values": values_obj }) + } + + fn serialize_to_bytes(&self) -> Vec { + let mut buffer = Vec::new(); + + // Write number of entries + buffer.extend_from_slice(&(self.values.len() as u32).to_le_bytes()); + + // Write each key-value pair + for (key, value) in &self.values { + let key_bytes = key.serialize_to_bytes(); + + // Write key length and data + buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&key_bytes); + + // Write value + buffer.extend_from_slice(&value.to_le_bytes()); + } + + buffer + } +} + +impl AggregateCore for KeyedMinState { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "KeyedMinState" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge KeyedMinState with {}", + other.get_accumulator_type() + ) + .into()); + } + + let other_multiple = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to KeyedMinState")?; + + let merged = Self::merge_accumulators(vec![self.clone(), other_multiple.clone()])?; + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Min + } + + fn approx_memory_bytes(&self) -> usize { + const BYTES_PER_ENTRY: usize = 96; + std::mem::size_of::() + self.values.len() * BYTES_PER_ENTRY + } + + fn get_keys(&self) -> Option> { + Some(self.values.keys().cloned().collect()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key.as_ref().ok_or("Key required for KeyedMinState")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedMinState { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match statistic { + Statistic::Min => self + .values + .get(key) + .copied() + .ok_or_else(|| format!("Key {key} not found in KeyedMinState").into()), + other => Err(format!("Unsupported statistic in KeyedMinState: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedMinState { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let mut result = KeyedMinState::new(); + + for acc in accumulators { + for (key, value) in acc.values { + result.update(key, value); + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(value: &str) -> KeyByLabelValues { + KeyByLabelValues::new_with_labels(vec![value.to_string()]) + } + + #[test] + fn keeps_the_smallest_per_key() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 10.0); + acc.update(key("a"), 5.0); + acc.update(key("a"), 15.0); + acc.update(key("b"), 7.0); + + assert_eq!(acc.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); + assert_eq!(acc.query(Statistic::Min, &key("b"), None).unwrap(), 7.0); + } + + #[test] + fn refuses_the_opposite_statistic_and_unknown_keys() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 1.0); + assert!(acc.query(Statistic::Max, &key("a"), None).is_err()); + assert!(acc.query(Statistic::Min, &key("missing"), None).is_err()); + } + + #[test] + fn merges_per_key() { + let mut left = KeyedMinState::new(); + left.update(key("a"), 10.0); + let mut right = KeyedMinState::new(); + right.update(key("a"), 5.0); + right.update(key("b"), 3.0); + + let merged = + >::merge_accumulators(vec![ + left, right, + ]) + .unwrap(); + + assert_eq!(merged.query(Statistic::Min, &key("a"), None).unwrap(), 5.0); + assert_eq!(merged.query(Statistic::Min, &key("b"), None).unwrap(), 3.0); + } + + #[test] + fn refuses_to_merge_with_the_opposite_direction() { + use super::super::keyed_max_state::KeyedMaxState; + let mine = KeyedMinState::new(); + let theirs = KeyedMaxState::new(); + assert!(mine.merge_with(&theirs).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let mut acc = KeyedMinState::new(); + acc.update(key("a"), 4.0); + + let json = acc.serialize_to_json(); + let from_json = KeyedMinState::deserialize_from_json(&json).unwrap(); + assert_eq!( + from_json.query(Statistic::Min, &key("a"), None).unwrap(), + 4.0 + ); + + let bytes = acc.serialize_to_bytes(); + let from_bytes = KeyedMinState::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!( + from_bytes.query(Statistic::Min, &key("a"), None).unwrap(), + 4.0 + ); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/keyed_sum_count.rs b/crates/asap-physical-operators/src/summary_kernels/keyed_sum_count.rs new file mode 100644 index 00000000..7bf2b5ce --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/keyed_sum_count.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 crate::Statistic; +use planner_types::post_asap::ExactKind; + +fn sum_family() -> ExactKind { + ExactKind::Sum +} + +/// Accumulator that maintains separate sum values for multiple keys +/// Allows querying sums for specific label combinations +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KeyedSumCountAccumulator { + #[serde(default = "sum_family")] + pub family: ExactKind, + pub sums: HashMap, + #[serde(default)] + pub counts: HashMap, +} + +impl KeyedSumCountAccumulator { + pub fn new() -> Self { + Self::for_family(ExactKind::Sum) + } + + pub fn for_family(family: ExactKind) -> Self { + assert!(matches!(family, ExactKind::Sum | ExactKind::Count)); + Self { + family, + sums: HashMap::new(), + counts: HashMap::new(), + } + } + + pub fn update(&mut self, key: KeyByLabelValues, value: f64) { + let is_new = !self.sums.contains_key(&key); + *self.sums.entry(key.clone()).or_insert(0.0) += value; + if let Some(count) = self.counts.get(&key).copied() { + if let Some(next) = count.checked_add(1).filter(|next| *next != u64::MAX) { + self.counts.insert(key, next); + } else { + self.counts.remove(&key); + } + } else if is_new { + self.counts.insert(key, 1); + } + } + + pub fn add_sum(&mut self, key: KeyByLabelValues, sum: f64) { + self.counts.remove(&key); + self.sums.insert(key, sum); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let sums_data = data["sums"] + .as_object() + .ok_or("Missing or invalid 'sums' field")?; + + let mut sums = HashMap::new(); + for (key_str, value) in sums_data { + let key_json: Value = serde_json::from_str(key_str)?; + let key = KeyByLabelValues::deserialize_from_json(&key_json)?; + let sum = value.as_f64().ok_or("Invalid sum value")?; + sums.insert(key, sum); + } + + let mut counts = HashMap::new(); + if let Some(counts_data) = data.get("counts").and_then(Value::as_object) { + for (key_str, value) in counts_data { + let key_json: Value = serde_json::from_str(key_str)?; + let key = KeyByLabelValues::deserialize_from_json(&key_json)?; + let count = value.as_u64().ok_or("Invalid count value")?; + if !sums.contains_key(&key) { + return Err("Count key missing from sums".into()); + } + counts.insert(key, count); + } + } + let family = match data.get("family").and_then(Value::as_str) { + None | Some("Sum") => ExactKind::Sum, + Some("Count") => ExactKind::Count, + _ => return Err("Invalid keyed additive family".into()), + }; + Ok(Self { + family, + sums, + counts, + }) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + let mut offset = 0; + + // Read number of entries + if buffer.len() < 4 { + return Err("Buffer too short for entry count".into()); + } + let num_entries = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + let mut sums = HashMap::new(); + let mut keys = Vec::new(); + + for _ in 0..num_entries { + // Read key length and data + if buffer.len() < offset + 4 { + return Err("Buffer too short for key length".into()); + } + let key_length = u32::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + ]) as usize; + offset += 4; + + if buffer.len() < offset + key_length { + return Err("Buffer too short for key data".into()); + } + let key = + KeyByLabelValues::deserialize_from_bytes(&buffer[offset..offset + key_length])?; + offset += key_length; + + // Read sum value + if buffer.len() < offset + 8 { + return Err("Buffer too short for sum value".into()); + } + let sum = f64::from_le_bytes([ + buffer[offset], + buffer[offset + 1], + buffer[offset + 2], + buffer[offset + 3], + buffer[offset + 4], + buffer[offset + 5], + buffer[offset + 6], + buffer[offset + 7], + ]); + offset += 8; + + keys.push(key.clone()); + sums.insert(key, sum); + } + let remaining = buffer.len() - offset; + let count_bytes = num_entries + .checked_mul(8) + .ok_or("Count section too large")?; + if remaining != 0 && remaining != count_bytes && remaining != count_bytes + 1 { + return Err("Invalid count section length".into()); + } + let mut counts = HashMap::new(); + if count_bytes != 0 && remaining >= count_bytes { + for key in keys { + let count = u64::from_le_bytes(buffer[offset..offset + 8].try_into()?); + offset += 8; + if count != u64::MAX { + counts.insert(key, count); + } + } + } + let family = if remaining == count_bytes + 1 { + match buffer[offset] { + 0 => ExactKind::Sum, + 1 => ExactKind::Count, + _ => return Err("Invalid keyed additive family tag".into()), + } + } else { + ExactKind::Sum + }; + Ok(Self { + family, + sums, + counts, + }) + } +} + +impl Default for KeyedSumCountAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl SerializableToSink for KeyedSumCountAccumulator { + fn serialize_to_json(&self) -> Value { + let mut sums_obj = serde_json::Map::new(); + for (key, sum) in &self.sums { + let key_json = key.serialize_to_json(); + let key_str = serde_json::to_string(&key_json).unwrap(); + sums_obj.insert( + key_str, + Value::Number(serde_json::Number::from_f64(*sum).unwrap()), + ); + } + + let mut counts_obj = serde_json::Map::new(); + for (key, count) in &self.counts { + let key_str = serde_json::to_string(&key.serialize_to_json()).unwrap(); + counts_obj.insert(key_str, Value::from(*count)); + } + + serde_json::json!({ + "family": if self.family == ExactKind::Count { "Count" } else { "Sum" }, + "sums": sums_obj, + "counts": counts_obj + }) + } + + fn serialize_to_bytes(&self) -> Vec { + let mut buffer = Vec::new(); + + // Write number of entries + buffer.extend_from_slice(&(self.sums.len() as u32).to_le_bytes()); + + // Write each key-value pair + let mut ordered_keys = Vec::with_capacity(self.sums.len()); + for (key, sum) in &self.sums { + ordered_keys.push(key); + let key_bytes = key.serialize_to_bytes(); + + // Write key length and data + buffer.extend_from_slice(&(key_bytes.len() as u32).to_le_bytes()); + buffer.extend_from_slice(&key_bytes); + + // Write sum value + buffer.extend_from_slice(&sum.to_le_bytes()); + } + + for key in ordered_keys { + buffer.extend_from_slice( + &self + .counts + .get(key) + .copied() + .unwrap_or(u64::MAX) + .to_le_bytes(), + ); + } + + buffer.push(if self.family == ExactKind::Count { + 1 + } else { + 0 + }); + + buffer + } +} + +impl AggregateCore for KeyedSumCountAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "KeyedSumCountAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + // Check if other is also a KeyedSumCountAccumulator + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge KeyedSumCountAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + // Downcast to KeyedSumCountAccumulator + let other_multiple_sum = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to KeyedSumCountAccumulator")?; + + // Use the existing merge_accumulators method + let merged = Self::merge_accumulators(vec![self.clone(), other_multiple_sum.clone()])?; + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + if self.family == ExactKind::Count { + AggregationType::Count + } else { + AggregationType::Sum + } + } + + fn approx_memory_bytes(&self) -> usize { + // HashMap. Label strings dominate; use a + // conservative per-entry estimate plus HashMap overhead. + const BYTES_PER_ENTRY: usize = 112; + std::mem::size_of::() + self.sums.len() * BYTES_PER_ENTRY + } + + fn get_keys(&self) -> Option> { + Some(self.sums.keys().cloned().collect()) + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + key: &Option, + query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::MultipleSubpopulationAggregate; + let key_val = key + .as_ref() + .ok_or("Key required for KeyedSumCountAccumulator")?; + self.query(statistic, key_val, Some(query_kwargs)) + } +} + +impl MultipleSubpopulationAggregate for KeyedSumCountAccumulator { + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + match (&self.family, statistic) { + (ExactKind::Sum, Statistic::Sum) => self.sums.get(key).copied().ok_or_else(|| { + "Key not found in KeyedSumCountAccumulator" + .to_string() + .into() + }), + (ExactKind::Count, Statistic::Count) => self + .counts + .get(key) + .map(|count| *count as f64) + .ok_or_else(|| { + "Sample count unavailable in KeyedSumCountAccumulator" + .to_string() + .into() + }), + _ => Err( + format!("Unsupported statistic in KeyedSumCountAccumulator: {statistic:?}").into(), + ), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for KeyedSumCountAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + + let family = accumulators[0].family.clone(); + if accumulators.iter().any(|acc| acc.family != family) { + return Err("Cannot merge different keyed additive families".into()); + } + let mut result = KeyedSumCountAccumulator::for_family(family); + + for acc in accumulators { + for key in acc.sums.keys() { + match ( + result.counts.get(key).copied(), + acc.counts.get(key).copied(), + ) { + (None, Some(count)) if !result.sums.contains_key(key) => { + result.counts.insert(key.clone(), count); + } + (Some(existing), Some(count)) => { + if let Some(total) = existing.checked_add(count) { + result.counts.insert(key.clone(), total); + } else { + result.counts.remove(key); + } + } + _ => { + result.counts.remove(key); + } + } + } + for (key, sum) in acc.sums { + *result.sums.entry(key).or_insert(0.0) += sum; + } + } + + Ok(result) + } +} + +#[cfg(test)] +mod tests { + use std::vec; + + use super::*; + + #[test] + fn test_keyed_sum_count_accumulator_creation() { + let acc = KeyedSumCountAccumulator::new(); + assert!(acc.sums.is_empty()); + } + + #[test] + fn test_keyed_sum_count_accumulator_update() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.update(key1.clone(), 10.0); + acc.update(key2.clone(), 20.0); + acc.update(key1.clone(), 5.0); // Should add to existing + + assert_eq!(acc.sums.get(&key1), Some(&15.0)); + assert_eq!(acc.sums.get(&key2), Some(&20.0)); + } + + #[test] + fn grouped_count_reads_sample_count_and_survives_merge_and_round_trip() { + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + let mut first = KeyedSumCountAccumulator::for_family(ExactKind::Count); + first.update(key.clone(), 10.0); + first.update(key.clone(), 20.0); + let mut second = KeyedSumCountAccumulator::for_family(ExactKind::Count); + second.update(key.clone(), 7.0); + let merged = KeyedSumCountAccumulator::merge_accumulators(vec![first, second]).unwrap(); + for acc in [ + merged.clone(), + KeyedSumCountAccumulator::deserialize_from_json(&merged.serialize_to_json()).unwrap(), + KeyedSumCountAccumulator::deserialize_from_bytes(&merged.serialize_to_bytes()).unwrap(), + ] { + assert_eq!(acc.family, ExactKind::Count); + assert!(acc.query(Statistic::Sum, &key, None).is_err()); + assert_eq!(acc.query(Statistic::Count, &key, None).unwrap(), 3.0); + } + } + + #[test] + fn keyed_additive_merge_rejects_different_planner_families() { + assert!(KeyedSumCountAccumulator::merge_accumulators(vec![ + KeyedSumCountAccumulator::for_family(ExactKind::Sum), + KeyedSumCountAccumulator::for_family(ExactKind::Count), + ]) + .is_err()); + } + + #[test] + fn test_keyed_sum_count_accumulator_query() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); + + acc.add_sum(key.clone(), 42.0); + + // Test total queries (querying with the specific key) + assert_eq!( + crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Sum, &key, None).unwrap(), + 42.0 + ); + + // Test error cases + assert!( + crate::MultipleSubpopulationAggregate::query(&acc, Statistic::Min, &key, None).is_err() + ); + } + + #[test] + fn test_keyed_sum_count_accumulator_get_keys() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc.add_sum(key1.clone(), 10.0); + acc.add_sum(key2.clone(), 20.0); + + let keys = crate::AggregateCore::get_keys(&acc).unwrap(); + assert_eq!(keys.len(), 2); + assert!(keys.contains(&key1)); + assert!(keys.contains(&key2)); + } + + #[test] + fn test_keyed_sum_count_accumulator_merge() { + let mut acc1 = KeyedSumCountAccumulator::new(); + let mut acc2 = KeyedSumCountAccumulator::new(); + + let key1 = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + let key2 = KeyByLabelValues::new_with_labels(vec!["api".to_string()]); + + acc1.add_sum(key1.clone(), 10.0); + acc1.add_sum(key2.clone(), 20.0); + + acc2.add_sum(key1.clone(), 5.0); // Same key, different accumulator + + let merged = >::merge_accumulators(vec![acc1, acc2]).unwrap(); + + assert_eq!(merged.sums.get(&key1), Some(&15.0)); // Should be merged + assert_eq!(merged.sums.get(&key2), Some(&20.0)); // Should be preserved + } + + #[test] + fn test_keyed_sum_count_accumulator_serialization() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["service".to_string()]); + + acc.add_sum(key.clone(), 42.5); + + // Test JSON serialization + let json = acc.serialize_to_json(); + let deserialized = KeyedSumCountAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!(deserialized.sums.get(&key), Some(&42.5)); + + // Test byte serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = KeyedSumCountAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(deserialized_bytes.sums.get(&key), Some(&42.5)); + } + + #[test] + fn test_trait_object() { + let mut acc = KeyedSumCountAccumulator::new(); + + let key = KeyByLabelValues::new_with_labels(vec!["web".to_string()]); + + acc.add_sum(key.clone(), 42.0); + + let trait_obj: Box = Box::new(acc); + + // Test type name through trait object + assert_eq!(trait_obj.type_name(), "KeyedSumCountAccumulator"); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/max.rs b/crates/asap-physical-operators/src/summary_kernels/max.rs new file mode 100644 index 00000000..0c954708 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/max.rs @@ -0,0 +1,248 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Exact maximum over one population, mergeable by comparison. +/// +/// See [`MinAccumulator`](super::min::MinAccumulator) for why the +/// two directions are separate types rather than one accumulator carrying a +/// `sub_type` string. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MaxAccumulator { + pub value: f64, +} + +impl Default for MaxAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl MaxAccumulator { + pub fn new() -> Self { + Self { + value: f64::NEG_INFINITY, + } + } + + pub fn with_value(value: f64) -> Self { + Self { value } + } + + pub fn update(&mut self, value: f64) { + if value > self.value { + self.value = value; + } + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let value = data["value"] + .as_f64() + .ok_or("Missing or invalid 'value' field")?; + Ok(Self::with_value(value)) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 8 { + return Err("Buffer too short".into()); + } + let value = f64::from_le_bytes([ + buffer[0], buffer[1], buffer[2], buffer[3], buffer[4], buffer[5], buffer[6], buffer[7], + ]); + Ok(Self::with_value(value)) + } +} + +impl SerializableToSink for MaxAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ "value": self.value }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.value.to_le_bytes().to_vec() + } +} + +impl MergeableAccumulator for MaxAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut result = MaxAccumulator::new(); + for acc in accumulators { + result.update(acc.value); + } + Ok(result) + } +} + +impl AggregateCore for MaxAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "MaxAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge MaxAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_max = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to MaxAccumulator")?; + let mut merged = self.clone(); + merged.update(other_max.value); + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Max + } + + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::() + } + + fn aux_stats(&self) -> AuxStats { + // The sentinel `f64::NEG_INFINITY` from `new()` is surfaced as-is; the + // query engine already treats it as "no data yet", the same way it + // does for `query_statistic`. + AuxStats { + max: Some(self.value), + ..AuxStats::empty() + } + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + _query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query(statistic, None) + } +} + +impl SingleSubpopulationAggregate for MaxAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + if query_kwargs.is_some() { + return Err("MaxAccumulator does not support query parameters".into()); + } + match statistic { + Statistic::Max => Ok(self.value), + other => Err(format!("Unsupported statistic in MaxAccumulator: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keeps_the_largest_update() { + let mut acc = MaxAccumulator::new(); + acc.update(10.0); + acc.update(5.0); + acc.update(15.0); + + assert_eq!(acc.value, 15.0); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Max, None).unwrap(), + 15.0 + ); + } + + #[test] + fn refuses_to_answer_a_minimum_query() { + let acc = MaxAccumulator::with_value(15.0); + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); + } + + #[test] + fn merges_by_taking_the_largest() { + let merged = + >::merge_accumulators(vec![ + MaxAccumulator::with_value(10.0), + MaxAccumulator::with_value(5.0), + MaxAccumulator::with_value(15.0), + ]) + .unwrap(); + assert_eq!(merged.value, 15.0); + } + + #[test] + fn refuses_to_merge_with_a_minimum() { + use super::super::min::MinAccumulator; + let max = MaxAccumulator::with_value(15.0); + let min = MinAccumulator::with_value(5.0); + assert!(max.merge_with(&min).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let acc = MaxAccumulator::with_value(42.5); + + let json = acc.serialize_to_json(); + assert_eq!( + MaxAccumulator::deserialize_from_json(&json).unwrap().value, + 42.5 + ); + + let bytes = acc.serialize_to_bytes(); + assert_eq!( + MaxAccumulator::deserialize_from_bytes(&bytes) + .unwrap() + .value, + 42.5 + ); + } + + #[test] + fn aux_stats_expose_max_only() { + let aux = MaxAccumulator::with_value(99.0).aux_stats(); + assert_eq!(aux.max, Some(99.0)); + assert_eq!(aux.min, None); + assert_eq!(aux.try_answer(Statistic::Max), Some(99.0)); + assert_eq!(aux.try_answer(Statistic::Min), None); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/min.rs b/crates/asap-physical-operators/src/summary_kernels/min.rs new file mode 100644 index 00000000..ff2ad548 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/min.rs @@ -0,0 +1,253 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +/// Exact minimum over one population, mergeable by comparison. +/// +/// The sibling [`MaxAccumulator`](super::max::MaxAccumulator) is a +/// separate type on purpose: these two used to be one `MinMaxAccumulator` +/// whose direction lived in a `sub_type: String`, which meant every layer +/// above -- the wire `aggregationSubType`, the accumulator factory, the +/// summary catalog -- had to carry the direction alongside the family and +/// could silently answer a `min_over_time` read from maximum state. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MinAccumulator { + pub value: f64, +} + +impl Default for MinAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl MinAccumulator { + pub fn new() -> Self { + Self { + value: f64::INFINITY, + } + } + + pub fn with_value(value: f64) -> Self { + Self { value } + } + + pub fn update(&mut self, value: f64) { + if value < self.value { + self.value = value; + } + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let value = data["value"] + .as_f64() + .ok_or("Missing or invalid 'value' field")?; + Ok(Self::with_value(value)) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 8 { + return Err("Buffer too short".into()); + } + let value = f64::from_le_bytes([ + buffer[0], buffer[1], buffer[2], buffer[3], buffer[4], buffer[5], buffer[6], buffer[7], + ]); + Ok(Self::with_value(value)) + } +} + +impl SerializableToSink for MinAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ "value": self.value }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.value.to_le_bytes().to_vec() + } +} + +impl MergeableAccumulator for MinAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + if accumulators.is_empty() { + return Err("No accumulators to merge".into()); + } + let mut result = MinAccumulator::new(); + for acc in accumulators { + result.update(acc.value); + } + Ok(result) + } +} + +impl AggregateCore for MinAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "MinAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge MinAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + let other_min = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to MinAccumulator")?; + let mut merged = self.clone(); + merged.update(other_min.value); + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Min + } + + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::() + } + + fn aux_stats(&self) -> AuxStats { + // The sentinel `f64::INFINITY` from `new()` is surfaced as-is; the + // query engine already treats it as "no data yet", the same way it + // does for `query_statistic`. + AuxStats { + min: Some(self.value), + ..AuxStats::empty() + } + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + _query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query(statistic, None) + } +} + +impl SingleSubpopulationAggregate for MinAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + if query_kwargs.is_some() { + return Err("MinAccumulator does not support query parameters".into()); + } + match statistic { + Statistic::Min => Ok(self.value), + other => Err(format!("Unsupported statistic in MinAccumulator: {other:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keeps_the_smallest_update() { + let mut acc = MinAccumulator::new(); + acc.update(10.0); + acc.update(5.0); + acc.update(15.0); + + assert_eq!(acc.value, 5.0); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).unwrap(), + 5.0 + ); + } + + #[test] + fn refuses_to_answer_a_maximum_query() { + let acc = MinAccumulator::with_value(5.0); + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Max, None).is_err()); + } + + #[test] + fn merges_by_taking_the_smallest() { + let merged = + >::merge_accumulators(vec![ + MinAccumulator::with_value(10.0), + MinAccumulator::with_value(5.0), + MinAccumulator::with_value(15.0), + ]) + .unwrap(); + assert_eq!(merged.value, 5.0); + } + + #[test] + fn refuses_to_merge_with_a_maximum() { + use super::super::max::MaxAccumulator; + let min = MinAccumulator::with_value(5.0); + let max = MaxAccumulator::with_value(15.0); + assert!(min.merge_with(&max).is_err()); + } + + #[test] + fn round_trips_through_both_serializations() { + let acc = MinAccumulator::with_value(42.5); + + let json = acc.serialize_to_json(); + assert_eq!( + MinAccumulator::deserialize_from_json(&json).unwrap().value, + 42.5 + ); + + let bytes = acc.serialize_to_bytes(); + assert_eq!( + MinAccumulator::deserialize_from_bytes(&bytes) + .unwrap() + .value, + 42.5 + ); + } + + #[test] + fn aux_stats_expose_min_only() { + let aux = MinAccumulator::with_value(3.5).aux_stats(); + assert_eq!(aux.min, Some(3.5)); + assert_eq!(aux.max, None); + assert_eq!(aux.count, None); + assert_eq!(aux.sum, None); + assert_eq!(aux.try_answer(Statistic::Min), Some(3.5)); + assert_eq!(aux.try_answer(Statistic::Max), None); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/mod.rs b/crates/asap-physical-operators/src/summary_kernels/mod.rs new file mode 100644 index 00000000..8ebc5d55 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/mod.rs @@ -0,0 +1,42 @@ +//! ASAP state adapters and exact accumulators; sketch algorithms live in `asap_sketchlib`. +pub mod count_min_sketch; +pub mod count_min_sketch_with_heap; +pub mod count_sketch; +pub mod count_sketch_with_heap; +pub mod datasketches_kll; +pub mod dd_sketch; +pub mod exact; +pub mod hll_sketch; +pub mod hydra_kll; +pub mod increase; +pub mod keyed_counter_state; +pub mod keyed_max_state; +pub mod keyed_min_state; +pub mod keyed_sum_count; +pub mod max; +pub mod min; +pub mod sketch_envelope; +pub mod sum; +pub mod univmon; + +pub use count_min_sketch::*; +pub use count_min_sketch_with_heap::*; +pub use count_sketch::*; +pub use count_sketch_with_heap::*; +pub use datasketches_kll::*; +pub use dd_sketch::*; +pub use hll_sketch::*; +pub use hydra_kll::*; +pub use increase::*; +pub use keyed_counter_state::*; +pub use keyed_max_state::*; +pub use keyed_min_state::*; +pub use keyed_sum_count::*; +pub use max::*; +pub use min::*; +pub use sketch_envelope::*; +pub use sum::*; + +pub mod factory; +pub mod traits; +pub mod weighted_frequency; diff --git a/crates/asap-physical-operators/src/summary_kernels/sketch_envelope.rs b/crates/asap-physical-operators/src/summary_kernels/sketch_envelope.rs new file mode 100644 index 00000000..475d7f51 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/sketch_envelope.rs @@ -0,0 +1,154 @@ +//! SketchEnvelopeAccumulator — wraps a raw SketchEnvelope protobuf payload +//! received via OTLP ingest so it can be stored through the `Store` trait. +//! +//! The accumulator preserves the opaque proto bytes and decodes them lazily +//! (via `SketchEnvelope::decode`) only when merge or query operations need +//! the inner sketch type. + +use crate::{AggregateCore, KeyByLabelValues, SerializableToSink}; +use asap_sketchlib::proto::sketchlib::{sketch_envelope, SketchEnvelope}; +use prost::Message; +use serde_json::Value; +use std::collections::HashMap; + +use crate::AggregationType; +use crate::Statistic; + +/// Accumulator that stores a serialized `SketchEnvelope` protobuf. +/// +/// This is the simplest viable path for OTLP sketch ingest: the OTel Collector +/// has already computed the sketch, so the backend just stores the bytes and +/// serves them back at query time. +#[derive(Debug, Clone)] +pub struct SketchEnvelopeAccumulator { + /// Raw protobuf-encoded `SketchEnvelope`. + pub payload: Vec, + /// Sketch type string cached from decoding (e.g. "CountMin", "KLL"). + pub sketch_type: String, +} + +impl SketchEnvelopeAccumulator { + /// Create from raw protobuf bytes. Decodes the envelope once to cache + /// the sketch type; the full payload is kept for later use. + pub fn from_proto_bytes( + payload: Vec, + ) -> Result> { + let sketch_type = match SketchEnvelope::decode(payload.as_slice()) { + Ok(env) => match env.sketch_state { + Some(sketch_envelope::SketchState::CountMin(_)) => "CountMin".to_string(), + Some(sketch_envelope::SketchState::CountSketch(_)) => "CountSketch".to_string(), + Some(sketch_envelope::SketchState::Kll(_)) => "KLL".to_string(), + Some(sketch_envelope::SketchState::Hll(_)) => "HLL".to_string(), + Some(sketch_envelope::SketchState::Ddsketch(_)) => "DDSketch".to_string(), + Some(sketch_envelope::SketchState::Univmon(_)) => "UnivMon".to_string(), + Some(sketch_envelope::SketchState::Hydra(_)) => "Hydra".to_string(), + Some(sketch_envelope::SketchState::Coco(_)) => "CocoSketch".to_string(), + Some(sketch_envelope::SketchState::Elastic(_)) => "Elastic".to_string(), + None => "Unknown".to_string(), + }, + Err(e) => { + return Err(format!("Failed to decode SketchEnvelope: {}", e).into()); + } + }; + + Ok(Self { + payload, + sketch_type, + }) + } +} + +// --------------------------------------------------------------------------- +// Trait implementations +// --------------------------------------------------------------------------- + +impl SerializableToSink for SketchEnvelopeAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "type": "SketchEnvelopeAccumulator", + "sketch_type": self.sketch_type, + "payload_bytes": self.payload.len(), + }) + } + + fn serialize_to_bytes(&self) -> Vec { + self.payload.clone() + } +} + +impl AggregateCore for SketchEnvelopeAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "SketchEnvelopeAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge SketchEnvelopeAccumulator with {:?}", + other.get_accumulator_type() + ) + .into()); + } + + // For now, merging opaque envelopes is not supported — each window is + // a self-contained sketch produced by the OTel Collector. Return self + // as-is so the store can still call merge_with without panicking. + Ok(Box::new(self.clone())) + } + + fn get_accumulator_type(&self) -> AggregationType { + // Opaque wrapper — report as the generic multi-subpopulation bucket. + // Direct dispatch is not supported; native sketch query path must + // decode the envelope and delegate to the correct accumulator. + AggregationType::MultipleSubpopulation + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + _statistic: Statistic, + _key: &Option, + _query_kwargs: &HashMap, + ) -> Result> { + Err( + "SketchEnvelopeAccumulator: query_statistic not supported; decode envelope first" + .into(), + ) + } +} + +impl crate::MultipleSubpopulationAggregate for SketchEnvelopeAccumulator { + fn query( + &self, + _statistic: Statistic, + _key: &KeyByLabelValues, + _query_kwargs: Option<&HashMap>, + ) -> Result> { + Err( + "SketchEnvelopeAccumulator: direct query not supported; use native sketch query path" + .into(), + ) + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/sum.rs b/crates/asap-physical-operators/src/summary_kernels/sum.rs new file mode 100644 index 00000000..4e74a45a --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/sum.rs @@ -0,0 +1,413 @@ +use crate::{ + AggregateCore, AggregationType, AuxStats, MergeableAccumulator, SerializableToSink, + SingleSubpopulationAggregate, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +use crate::Statistic; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SumAccumulator { + pub sum: f64, + /// None for scalar-only payloads; a sum does not establish a sample count. + #[serde(default)] + pub observation_count: Option, +} + +impl SumAccumulator { + pub fn new() -> Self { + Self { + sum: 0.0, + observation_count: Some(0), + } + } + + pub fn with_sum(sum: f64) -> Self { + Self { + sum, + observation_count: None, + } + } + + pub fn update(&mut self, value: f64) { + self.sum += value; + self.observation_count = self + .observation_count + .and_then(|count| count.checked_add(1)); + } + + pub fn deserialize_from_json(data: &Value) -> Result> { + let sum = data["sum"] + .as_f64() + .ok_or("Missing or invalid 'sum' field")?; + Ok(Self { + sum, + observation_count: data.get("observation_count").and_then(Value::as_u64), + }) + } + + pub fn deserialize_from_bytes(buffer: &[u8]) -> Result> { + match buffer.len() { + // Legacy Python scalar sums carry no sample-count evidence. + 4 => Ok(Self::with_sum(f32::from_le_bytes(buffer.try_into()?) as f64)), + // Counted sums use the same fixed layout as the Collector Sum payload. + 16 => Self::from_sum_bytes(buffer), + len => { + Err(format!("Invalid persisted Sum payload length: {len} (want 4 or 16)").into()) + } + } + } + + /// Decode the fixed Sum payload produced by the first-class Sum + /// AggregationType path (asap-precompute-go's SumWrapper): float64 sum + /// (little-endian) followed by uint64 count (little-endian), 16 bytes. + /// + /// Sum is an aggregation, NOT a sketch, so this deliberately does NOT + /// depend on the sketchlib sketch-envelope proto — the payload is a small + /// self-contained fixed layout. It decodes into the SAME + /// `AggregationType::Sum` accumulator as a plain-OTLP Sum, so the SumAgg + /// envelope and a plain Sum land on one identity (`exact_agg:Sum`) with no + /// new SketchAlgorithm. The supplied observation count is retained for + /// exact sample-count readouts; scalar-only legacy payloads leave it unknown. + pub fn from_sum_bytes(buffer: &[u8]) -> Result> { + if buffer.len() < 16 { + return Err(format!("Sum payload too short: {} bytes (want 16)", buffer.len()).into()); + } + let sum = f64::from_le_bytes(buffer[0..8].try_into().unwrap()); + let count = u64::from_le_bytes(buffer[8..16].try_into().unwrap()); + Ok(Self { + sum, + observation_count: Some(count), + }) + } +} + +impl Default for SumAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl SerializableToSink for SumAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({ + "sum": self.sum, + "observation_count": self.observation_count + }) + } + + fn serialize_to_bytes(&self) -> Vec { + match self.observation_count { + Some(count) => { + let mut bytes = Vec::with_capacity(16); + bytes.extend_from_slice(&self.sum.to_le_bytes()); + bytes.extend_from_slice(&count.to_le_bytes()); + bytes + } + None => (self.sum as f32).to_le_bytes().to_vec(), + } + } +} + +impl AggregateCore for SumAccumulator { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + + fn type_name(&self) -> &'static str { + "SumAccumulator" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + // Check if other is also a SumAccumulator + if other.get_accumulator_type() != self.get_accumulator_type() { + return Err(format!( + "Cannot merge SumAccumulator with {}", + other.get_accumulator_type() + ) + .into()); + } + + // Downcast to SumAccumulator + let other_sum = other + .as_any() + .downcast_ref::() + .ok_or("Failed to downcast to SumAccumulator")?; + + // Use the existing merge_accumulators method + let merged = Self::merge_accumulators(vec![self.clone(), other_sum.clone()])?; + + Ok(Box::new(merged)) + } + + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::Sum + } + + fn approx_memory_bytes(&self) -> usize { + // Single f64 + struct overhead. + std::mem::size_of::() + } + + fn aux_stats(&self) -> AuxStats { + AuxStats { + sum: Some(self.sum), + count: self.observation_count, + ..AuxStats::empty() + } + } + + fn get_keys(&self) -> Option> { + None + } + + fn query_statistic( + &self, + statistic: crate::Statistic, + _key: &Option, + _query_kwargs: &std::collections::HashMap, + ) -> Result> { + use crate::SingleSubpopulationAggregate; + self.query(statistic, None) + } +} + +impl SingleSubpopulationAggregate for SumAccumulator { + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result> { + // SumAccumulator doesn't use query_kwargs, assert it's None + if query_kwargs.is_some() { + return Err("SumAccumulator does not support query parameters".into()); + } + + match statistic { + Statistic::Sum => Ok(self.sum), + Statistic::Count => self + .observation_count + .map(|count| count as f64) + .ok_or_else(|| "sample count is unavailable for this Sum payload".into()), + _ => Err(format!("Unsupported statistic in SumAccumulator: {statistic:?}").into()), + } + } + + fn clone_boxed(&self) -> Box { + Box::new(self.clone()) + } +} + +impl MergeableAccumulator for SumAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> { + let total_sum = accumulators.iter().map(|acc| acc.sum).sum(); + let observation_count = accumulators + .iter() + .try_fold(0u64, |total, acc| total.checked_add(acc.observation_count?)); + Ok(SumAccumulator { + sum: total_sum, + observation_count, + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // Sample counts must survive updates and merges independently of the sum. + #[test] + fn observation_count_survives_merge() { + let mut first = SumAccumulator::new(); + first.update(10.0); + first.update(20.0); + let mut second = SumAccumulator::new(); + second.update(100.0); + let merged = SumAccumulator::merge_accumulators(vec![first, second]).unwrap(); + assert_eq!(merged.sum, 130.0); + assert_eq!(merged.aux_stats().count, Some(3)); + } + + // A legacy scalar sum has no evidence of how many observations produced it. + #[test] + fn legacy_sum_does_not_invent_observation_count() { + let mut raw = SumAccumulator::new(); + raw.update(10.0); + let merged = + SumAccumulator::merge_accumulators(vec![raw, SumAccumulator::with_sum(20.0)]).unwrap(); + assert_eq!(merged.aux_stats().count, None); + } + + // Persistence retains known counts, including zero and the full u64 range. + #[test] + fn counted_sum_binary_round_trip() { + for count in [0, 3, u64::MAX] { + let acc = SumAccumulator { + sum: 1.0000000000001, + observation_count: Some(count), + }; + let bytes = acc.serialize_to_bytes(); + assert_eq!(bytes.len(), 16); + let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(restored.sum, acc.sum); + assert_eq!(restored.observation_count, Some(count)); + } + } + + // Existing scalar-only files remain readable without inventing counts. + #[test] + fn legacy_binary_sum_has_unknown_count() { + let bytes = 42.5f32.to_le_bytes(); + let restored = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(restored.sum, 42.5); + assert_eq!(restored.observation_count, None); + assert_eq!(restored.serialize_to_bytes(), bytes); + } + + // Truncated counted payloads must not silently decode as scalar sums. + #[test] + fn persisted_sum_rejects_invalid_lengths() { + for len in [0, 3, 5, 8, 15, 17] { + assert!(SumAccumulator::deserialize_from_bytes(&vec![0; len]).is_err()); + } + } + + #[test] + fn test_sum_accumulator_creation() { + let acc = SumAccumulator::new(); + assert_eq!(acc.sum, 0.0); + + let acc2 = SumAccumulator::with_sum(42.5); + assert_eq!(acc2.sum, 42.5); + } + + #[test] + fn test_sum_accumulator_update() { + let mut acc = SumAccumulator::new(); + acc.update(10.0); + acc.update(20.0); + assert_eq!(acc.sum, 30.0); + } + + #[test] + fn test_sum_accumulator_query() { + let acc = SumAccumulator::with_sum(42.0); + + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), + 42.0 + ); + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Count, None).is_err()); + + assert!(crate::SingleSubpopulationAggregate::query(&acc, Statistic::Min, None).is_err()); + // SumAccumulator is a single subpopulation accumulator, doesn't need key-based queries + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Sum, None).unwrap(), + 42.0 + ); + } + + #[test] + fn count_readout_uses_observation_count_not_sum() { + let mut acc = SumAccumulator::new(); + acc.update(10.0); + acc.update(20.0); + assert_eq!( + crate::SingleSubpopulationAggregate::query(&acc, Statistic::Count, None).unwrap(), + 2.0 + ); + } + + #[test] + fn test_sum_accumulator_merge() { + let acc1 = SumAccumulator::with_sum(10.0); + let acc2 = SumAccumulator::with_sum(20.0); + let acc3 = SumAccumulator::with_sum(30.0); + + let merged = + >::merge_accumulators(vec![ + acc1, acc2, acc3, + ]) + .unwrap(); + assert_eq!(merged.sum, 60.0); + } + + #[test] + fn test_sum_accumulator_serialization() { + let acc = SumAccumulator::with_sum(42.5); + + // Test JSON serialization + let json = acc.serialize_to_json(); + let deserialized = SumAccumulator::deserialize_from_json(&json).unwrap(); + assert_eq!(acc.sum, deserialized.sum); + + // Test byte serialization + let bytes = acc.serialize_to_bytes(); + let deserialized_bytes = SumAccumulator::deserialize_from_bytes(&bytes).unwrap(); + assert_eq!(acc.sum, deserialized_bytes.sum); + } + + #[test] + fn test_trait_object() { + let acc: Box = Box::new(SumAccumulator::with_sum(42.0)); + + assert_eq!(acc.type_name(), "SumAccumulator"); + } + + #[test] + fn from_sum_bytes_decodes_go_sum_payload() { + // GOLDEN: the 16-byte payload asap-precompute-go's + // SumWrapper{10,20,30,40}.Snapshot() emits — float64 sum (LE) followed + // by uint64 count (LE), sum=100, count=4. Proves the Rust backend + // decodes the first-class Sum payload the Go agent produces + // (cross-language wire parity, no sketchlib proto dependency). + let go_bytes: &[u8] = &[ + 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x59, 0x40, // 100.0 f64 LE + 0x04, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // 4 u64 LE + ]; + let acc = SumAccumulator::from_sum_bytes(go_bytes).expect("decode Go Sum payload"); + assert_eq!(acc.sum, 100.0, "decoded Go SumWrapper payload sum"); + } + + #[test] + fn from_sum_bytes_rejects_short_payload() { + // A short buffer is rejected (the ingest path then skips the point). + assert!(SumAccumulator::from_sum_bytes(&[]).is_err()); + assert!(SumAccumulator::from_sum_bytes(&[0u8; 8]).is_err()); + } + + #[test] + fn aux_stats_exposes_sum_only() { + let acc = SumAccumulator::with_sum(123.5); + let aux = acc.aux_stats(); + assert_eq!(aux.sum, Some(123.5)); + assert_eq!(aux.count, None); + assert_eq!(aux.min, None); + assert_eq!(aux.max, None); + } + + #[test] + fn aux_stats_try_answer_on_sum_statistic() { + use crate::Statistic; + let acc = SumAccumulator::with_sum(42.0); + // Sum statistic is covered by aux without deserialising. + assert_eq!(acc.aux_stats().try_answer(Statistic::Sum), Some(42.0)); + // Count is not tracked by SumAccumulator. + assert_eq!(acc.aux_stats().try_answer(Statistic::Count), None); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/traits.rs b/crates/asap-physical-operators/src/summary_kernels/traits.rs new file mode 100644 index 00000000..ae05ca06 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/traits.rs @@ -0,0 +1,357 @@ +use crate::KeyByLabelValues; +use std::collections::HashMap; + +use crate::AggregationType; +use crate::Statistic; + +use serde_json::Value; + +/// Trait for objects that can be serialized to different formats +pub trait SerializableToSink { + fn serialize_to_json(&self) -> Value; + fn serialize_to_bytes(&self) -> Vec; +} + +/// Core trait for all aggregates containing shared functionality +/// This trait provides common operations like serialization, cloning, and type identification +pub trait AggregateCore: SerializableToSink + Send + Sync { + /// Clone this accumulator into a boxed trait object + fn clone_boxed_core(&self) -> Box; + + /// Get the type name of this accumulator + fn type_name(&self) -> &'static str; + + /// Downcast to Any for type checking + fn as_any(&self) -> &dyn std::any::Any; + + /// Mutable downcast to Any. Used by ingest paths that need to + /// mutate a boxed accumulator in place — e.g. the PROTO_DELTA + /// delta-merge applier in `drivers::ingest::otel::apply_modified_otlp_delta_bytes`. + fn as_any_mut(&mut self) -> &mut dyn std::any::Any; + + /// Merge this accumulator with another accumulator of the same type + /// Returns a new merged accumulator, leaving the original unchanged + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box>; + + /// Get the accumulator type identifier for merge compatibility checking + fn get_accumulator_type(&self) -> AggregationType; + + /// Get all keys stored in this accumulator + fn get_keys(&self) -> Option>; + + /// Dispatch a statistic query without downcasting. + /// + /// Replaces the 12-arm `match get_accumulator_type()` in the engine. + /// Single-subpopulation types ignore `key`; multiple-subpopulation types + /// require it and return `Err` when it is `None`. + /// Special cases (DeltaSetAggregator, SetAggregator) fall back to a + /// cardinality value when `key` is `None`. + fn query_statistic( + &self, + statistic: Statistic, + key: &Option, + query_kwargs: &HashMap, + ) -> Result>; + + /// Approximate in-memory byte footprint of this accumulator. + /// + /// Used by the `SketchStore` persistence layer to drive its + /// memory-pressure trigger. Not required to be exact — the flusher + /// only needs rough proportionality. The default is a conservative + /// 4 KiB constant; concrete types should override it with a + /// type-aware estimate (e.g. KLL: `k * 8` plus overhead). + /// + /// Implementors must not call `serialize_to_bytes` here — this is + /// on the insert hot path. + fn approx_memory_bytes(&self) -> usize { + 4096 + } + + /// Typed auxiliary statistics — `count`, `sum`, `min`, `max` — + /// exposed as first-class scalars alongside the sketch payload. + /// + /// The overwhelming majority of production queries + /// (`count_over_time`, `sum_over_time`, `min_over_time`, + /// `max_over_time`, and the additive aggregations built on + /// them) only need these scalars. Returning them directly here + /// lets callers avoid deserialising the full sketch bytes. + /// + /// Returning fields as `None` means the accumulator doesn't + /// track that statistic exactly (e.g. a pure HLL doesn't carry + /// sum/min/max). Callers then fall back to the sketch's + /// `query_statistic` method. + /// + /// This is the phase-1 piece of the sketch DB design + /// (docs/design_docs/summary-storage.md). + fn aux_stats(&self) -> AuxStats { + AuxStats::empty() + } + + /// Reset the sketch state to empty **in place**, preserving its + /// shape / configuration (dimensions, relative accuracy, register + /// width, …) so a subsequent delta-apply lands on a clean, + /// same-shape base. + /// + /// Used by the OTLP ingest path's per-window base rotation: when a + /// delta frame opens a new tumbling window for a series, the cached + /// base is reset here before the new window's delta is applied, so + /// the reconstructed state reflects that window only rather than an + /// all-time accumulation across windows (see + /// `docs/delta-baseline-contract.md` §3). + /// + /// The default is a no-op: only the delta-capable, additive families + /// (DDSketch, CMS, CountSketch, HLL) ever reach the rotation path and + /// override this. KLL never deltas, and the non-sketch accumulators + /// are never cached as a delta base. + fn reset_to_empty(&mut self) {} +} + +/// Four typed auxiliary scalars tracked alongside every sketch entry: +/// `count`, `sum`, `min`, `max`. Exposed so the query engine can +/// serve Count / Sum / Min / Max statistics without touching sketch +/// bytes. +/// +/// Each field is `Option<…>` because not every accumulator tracks +/// every stat (e.g. HLL has cardinality but no meaningful +/// sum / min / max; DeltaSetAggregator tracks set transitions, not +/// numeric aggregates). +#[derive(Debug, Default, Clone, Copy, PartialEq)] +pub struct AuxStats { + pub count: Option, + pub sum: Option, + pub min: Option, + pub max: Option, +} + +impl AuxStats { + pub const fn empty() -> Self { + Self { + count: None, + sum: None, + min: None, + max: None, + } + } + + /// Attempt to fulfil a `Statistic` purely from the typed aux + /// columns, without needing to deserialise the sketch. Returns + /// `None` if the requested statistic isn't covered by aux + /// (e.g. Quantile, Cardinality, TopK) or if the corresponding + /// aux field is `None`. + pub fn try_answer(&self, statistic: Statistic) -> Option { + match statistic { + Statistic::Count => self.count.map(|c| c as f64), + Statistic::Sum => self.sum, + Statistic::Min => self.min, + Statistic::Max => self.max, + // Increase / Rate need two samples; aux columns carry + // window totals, so one entry's aux is insufficient. + // Cardinality / Quantile / Topk are sketch-native and + // must go through query_statistic. + _ => None, + } + } + + /// Merge two aux stats the way the corresponding sketch merge + /// would. Count / sum add, min / max take the extremum. When + /// either side is `None` the result is the other side (so a + /// window that only has partial aux still contributes). + pub fn merge(self, other: Self) -> Self { + fn add_opt_u(a: Option, b: Option) -> Option { + match (a, b) { + (Some(x), Some(y)) => Some(x.saturating_add(y)), + (x, None) => x, + (None, y) => y, + } + } + fn add_opt_f(a: Option, b: Option) -> Option { + match (a, b) { + (Some(x), Some(y)) => Some(x + y), + (x, None) => x, + (None, y) => y, + } + } + fn min_opt(a: Option, b: Option) -> Option { + match (a, b) { + (Some(x), Some(y)) => Some(x.min(y)), + (x, None) => x, + (None, y) => y, + } + } + fn max_opt(a: Option, b: Option) -> Option { + match (a, b) { + (Some(x), Some(y)) => Some(x.max(y)), + (x, None) => x, + (None, y) => y, + } + } + Self { + count: add_opt_u(self.count, other.count), + sum: add_opt_f(self.sum, other.sum), + min: min_opt(self.min, other.min), + max: max_opt(self.max, other.max), + } + } +} + +/// Trait for accumulators that support a single subpopulation +/// These accumulators store a single aggregate value (e.g., Sum, Increase) +pub trait SingleSubpopulationAggregate: AggregateCore { + /// Query the accumulator for a specific statistic + fn query( + &self, + statistic: Statistic, + query_kwargs: Option<&HashMap>, + ) -> Result>; + + /// Clone this accumulator into a boxed trait object + fn clone_boxed(&self) -> Box; +} + +/// Trait for accumulators that support multiple subpopulations identified by keys +/// These accumulators store separate values for different label combinations +pub trait MultipleSubpopulationAggregate: AggregateCore { + /// Query the accumulator for a specific statistic and key + fn query( + &self, + statistic: Statistic, + key: &KeyByLabelValues, + query_kwargs: Option<&HashMap>, + ) -> Result>; + + /// Clone this accumulator into a boxed trait object + fn clone_boxed(&self) -> Box; +} + +/// Trait for merging multiple accumulators of the same type +pub trait MergeableAccumulator { + fn merge_accumulators( + accumulators: Vec, + ) -> Result> + where + T: Sized; +} + +// Implement Clone for the new trait objects +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_boxed_core() + } +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_boxed() + } +} + +impl Clone for Box { + fn clone(&self) -> Self { + self.clone_boxed() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn aux_stats_empty_answers_nothing() { + let e = AuxStats::empty(); + assert_eq!(e.try_answer(Statistic::Count), None); + assert_eq!(e.try_answer(Statistic::Sum), None); + assert_eq!(e.try_answer(Statistic::Min), None); + assert_eq!(e.try_answer(Statistic::Max), None); + } + + #[test] + fn aux_stats_try_answer_covers_typed_stats() { + let a = AuxStats { + count: Some(7), + sum: Some(42.0), + min: Some(1.5), + max: Some(9.25), + }; + assert_eq!(a.try_answer(Statistic::Count), Some(7.0)); + assert_eq!(a.try_answer(Statistic::Sum), Some(42.0)); + assert_eq!(a.try_answer(Statistic::Min), Some(1.5)); + assert_eq!(a.try_answer(Statistic::Max), Some(9.25)); + } + + #[test] + fn aux_stats_try_answer_skips_sketch_native_stats() { + let a = AuxStats { + count: Some(100), + sum: Some(500.0), + min: Some(1.0), + max: Some(10.0), + }; + assert_eq!(a.try_answer(Statistic::Quantile), None); + assert_eq!(a.try_answer(Statistic::Cardinality), None); + assert_eq!(a.try_answer(Statistic::Topk), None); + assert_eq!(a.try_answer(Statistic::Increase), None); + assert_eq!(a.try_answer(Statistic::Rate), None); + } + + #[test] + fn aux_stats_merge_adds_count_and_sum_takes_extrema() { + let a = AuxStats { + count: Some(10), + sum: Some(50.0), + min: Some(1.0), + max: Some(9.0), + }; + let b = AuxStats { + count: Some(5), + sum: Some(20.0), + min: Some(0.5), + max: Some(12.0), + }; + let merged = a.merge(b); + assert_eq!(merged.count, Some(15)); + assert_eq!(merged.sum, Some(70.0)); + assert_eq!(merged.min, Some(0.5)); + assert_eq!(merged.max, Some(12.0)); + } + + #[test] + fn aux_stats_merge_handles_partial_sides() { + // HLL-like (count only) merged with Sum-only side. + let hll_like = AuxStats { + count: Some(100), + ..AuxStats::empty() + }; + let sum_like = AuxStats { + sum: Some(500.0), + ..AuxStats::empty() + }; + let merged = hll_like.merge(sum_like); + assert_eq!(merged.count, Some(100)); + assert_eq!(merged.sum, Some(500.0)); + assert_eq!(merged.min, None); + assert_eq!(merged.max, None); + } + + #[test] + fn aux_stats_merge_is_empty_plus_empty() { + let merged = AuxStats::empty().merge(AuxStats::empty()); + assert_eq!(merged, AuxStats::empty()); + } + + #[test] + fn aux_stats_count_saturates_on_overflow() { + let a = AuxStats { + count: Some(u64::MAX - 1), + ..AuxStats::empty() + }; + let b = AuxStats { + count: Some(100), + ..AuxStats::empty() + }; + let merged = a.merge(b); + assert_eq!(merged.count, Some(u64::MAX)); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/univmon.rs b/crates/asap-physical-operators/src/summary_kernels/univmon.rs new file mode 100644 index 00000000..4a52bf1d --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/univmon.rs @@ -0,0 +1,234 @@ +//! One frequency state shared by count, distinct, L2 and entropy readouts. + +use crate::{AggregateCore, AuxStats, KeyByLabelValues, SerializableToSink}; +use crate::{AggregationType, Statistic}; +use asap_sketchlib::{DataInput, UnivMon}; +use serde_json::Value; +use std::collections::HashMap; + +type Error = Box; + +#[derive(Debug, Clone)] +pub struct UnivMonAccumulator { + inner: UnivMon, +} + +impl UnivMonAccumulator { + pub fn new(heap_size: usize, rows: usize, cols: usize, layers: usize) -> Result { + if heap_size == 0 || cols == 0 || !(1..=20).contains(&rows) || !(1..=64).contains(&layers) { + return Err("invalid UnivMon dimensions".into()); + } + rows.checked_mul(cols) + .and_then(|n| n.checked_mul(layers)) + .ok_or("UnivMon dimensions overflow")?; + Ok(Self { + inner: UnivMon::init_univmon(heap_size, rows, cols, layers), + }) + } + + /// Each non-NaN sample is one occurrence. Signed zero has one identity. + pub fn insert_sample(&mut self, value: f64) -> Result<(), Error> { + if value.is_nan() { + return Ok(()); + } + self.inner + .bucket_size + .checked_add(1) + .ok_or("UnivMon count overflow")?; + let bits = if value == 0.0 { 0 } else { value.to_bits() }; + self.inner.insert(&DataInput::U64(bits), 1); + Ok(()) + } + + pub fn from_bytes(bytes: &[u8]) -> Result { + let inner = UnivMon::deserialize_from_bytes(bytes) + .map_err(|e| format!("invalid UnivMon state: {e}"))?; + if !inner.accepts_standard_updates() { + return Err( + "terminal-mode UnivMon state cannot enter the standard-update accumulator".into(), + ); + } + Ok(Self { inner }) + } + + fn compatible(&self, other: &Self) -> bool { + ( + self.inner.heap_size, + self.inner.sketch_row, + self.inner.sketch_col, + self.inner.layer_size, + ) == ( + other.inner.heap_size, + other.inner.sketch_row, + other.inner.sketch_col, + other.inner.layer_size, + ) + } + + pub fn dimensions(&self) -> (usize, usize, usize, usize) { + ( + self.inner.heap_size, + self.inner.sketch_row, + self.inner.sketch_col, + self.inner.layer_size, + ) + } + + pub fn merge_in_place(&mut self, other: &Self) -> Result<(), Error> { + if !self.compatible(other) { + return Err("incompatible UnivMon dimensions".into()); + } + self.inner + .bucket_size + .checked_add(other.inner.bucket_size) + .ok_or("UnivMon count overflow")?; + self.inner.merge(&other.inner); + Ok(()) + } +} + +impl SerializableToSink for UnivMonAccumulator { + fn serialize_to_json(&self) -> Value { + serde_json::json!({"count": self.inner.bucket_size}) + } + + fn serialize_to_bytes(&self) -> Vec { + self.inner + .serialize_to_bytes() + .expect("validated unit-frequency UnivMon state") + } +} + +impl AggregateCore for UnivMonAccumulator { + fn approx_memory_bytes(&self) -> usize { + std::mem::size_of::().saturating_add( + self.inner.layer_size.saturating_mul( + self.inner + .sketch_row + .saturating_mul(self.inner.sketch_col) + .saturating_mul(16) + .saturating_add(self.inner.heap_size.saturating_mul(256)), + ), + ) + } + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + fn type_name(&self) -> &'static str { + "UnivMonAccumulator" + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + fn get_accumulator_type(&self) -> AggregationType { + AggregationType::UnivMon + } + fn get_keys(&self) -> Option> { + None + } + fn reset_to_empty(&mut self) { + self.inner.free(); + } + + fn merge_with(&self, other: &dyn AggregateCore) -> Result, Error> { + let other = other + .as_any() + .downcast_ref::() + .ok_or("expected UnivMon state")?; + let mut merged = self.clone(); + merged.merge_in_place(other)?; + Ok(Box::new(merged)) + } + + fn query_statistic( + &self, + statistic: Statistic, + key: &Option, + _: &HashMap, + ) -> Result { + if key.is_some() { + return Err("UnivMon population is selected by the catalog binding".into()); + } + match statistic { + Statistic::Count => Ok(self.inner.calc_l1()), + Statistic::Cardinality => Ok(self.inner.calc_card()), + Statistic::FrequencyL2 => Ok(self.inner.calc_l2()), + Statistic::FrequencyEntropy => Ok(self.inner.calc_entropy()), + _ => Err("unsupported UnivMon readout".into()), + } + } + + fn aux_stats(&self) -> AuxStats { + AuxStats { + count: Some(self.inner.bucket_size as u64), + ..AuxStats::empty() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn read(state: &dyn AggregateCore, stat: Statistic) -> f64 { + state.query_statistic(stat, &None, &HashMap::new()).unwrap() + } + + /// Duplicate samples affect frequency but not cardinality, including signed zero. + #[test] + fn shared_readouts_survive_serialization() { + let mut state = UnivMonAccumulator::new(32, 5, 1024, 4).unwrap(); + for value in [0.0, -0.0, 2.0, 2.0, f64::NAN] { + state.insert_sample(value).unwrap(); + } + let restored = UnivMonAccumulator::from_bytes(&state.serialize_to_bytes()).unwrap(); + for stat in [ + Statistic::Count, + Statistic::Cardinality, + Statistic::FrequencyL2, + Statistic::FrequencyEntropy, + ] { + assert_eq!(read(&state, stat), read(&restored, stat)); + } + assert_eq!(read(&restored, Statistic::Count), 4.0); + assert!((read(&restored, Statistic::Cardinality) - 2.0).abs() < 0.01); + assert!((read(&restored, Statistic::FrequencyL2) - 8.0f64.sqrt()).abs() < 0.01); + assert!((read(&restored, Statistic::FrequencyEntropy) - 1.0).abs() < 0.01); + } + + /// Terminal-mode serialization is valid sketchlib state but not this accumulator's update domain. + #[test] + fn terminal_state_is_rejected_before_ingestion_or_merge() { + let mut state = UnivMon::init_univmon(4, 3, 16, 2); + state.fast_insert(&DataInput::U64(1), 1); + let bytes = state.serialize_to_bytes().unwrap(); + assert!(UnivMonAccumulator::from_bytes(&bytes).is_err()); + state.free(); + assert!(UnivMonAccumulator::from_bytes(&state.serialize_to_bytes().unwrap()).is_ok()); + } + + /// Pane merge preserves overlapping keys and reset removes the previous window. + #[test] + fn merge_and_reset_preserve_frequency_semantics() { + let mut left = UnivMonAccumulator::new(32, 5, 1024, 4).unwrap(); + let mut right = left.clone(); + for value in [1.0, 2.0] { + left.insert_sample(value).unwrap(); + } + for value in [2.0, 3.0] { + right.insert_sample(value).unwrap(); + } + let merged = left.merge_with(&right).unwrap(); + assert_eq!(read(merged.as_ref(), Statistic::Count), 4.0); + assert!((read(merged.as_ref(), Statistic::Cardinality) - 3.0).abs() < 0.01); + left.reset_to_empty(); + assert_eq!(read(&left, Statistic::Count), 0.0); + assert_eq!(read(&left, Statistic::FrequencyEntropy), 0.0); + assert!(left + .merge_with(&UnivMonAccumulator::new(16, 5, 1024, 4).unwrap()) + .is_err()); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/weighted_frequency.rs b/crates/asap-physical-operators/src/summary_kernels/weighted_frequency.rs new file mode 100644 index 00000000..042dce65 --- /dev/null +++ b/crates/asap-physical-operators/src/summary_kernels/weighted_frequency.rs @@ -0,0 +1,236 @@ +//! ASAP type and trait adapter for sketchlib's Float64 weighted frequency kernel. +use crate::{values::Value, Error}; +use crate::{AggregateCore, AggregationType, KeyByLabelValues, SerializableToSink, Statistic}; +pub use asap_sketchlib::FrequencyAlgorithm; +use asap_sketchlib::{FrequencyIdentity, WeightedFrequency as Kernel, WeightedFrequencyError}; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +fn adapt_error(error: WeightedFrequencyError) -> Error { + match error { + WeightedFrequencyError::Invalid(message) => Error::Invalid(message), + WeightedFrequencyError::Update(message) => Error::Operator(message), + } +} +fn identity(value: &Value) -> Result { + Ok(match value { + Value::Null => FrequencyIdentity::Null, + Value::Bool(v) => FrequencyIdentity::Bool(*v), + Value::Int64(v) => FrequencyIdentity::Int64(*v), + Value::Float64(v) => FrequencyIdentity::Float64(*v), + Value::Utf8(v) => FrequencyIdentity::Utf8(v.to_string()), + _ => { + return Err(Error::Invalid( + "unsupported weighted frequency identity".into(), + )) + } + }) +} +fn value(identity: FrequencyIdentity) -> Value { + match identity { + FrequencyIdentity::Null => Value::Null, + FrequencyIdentity::Bool(v) => Value::Bool(v), + FrequencyIdentity::Int64(v) => Value::Int64(v), + FrequencyIdentity::Float64(v) => Value::Float64(v), + FrequencyIdentity::Utf8(v) => Value::Utf8(v.into()), + } +} +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(transparent)] +pub struct WeightedFrequency { + inner: Kernel, +} +impl WeightedFrequency { + pub(crate) fn configuration( + kind: &planner_types::post_asap::SketchKind, + ) -> Result<(FrequencyAlgorithm, usize, usize, usize), Error> { + use planner_types::post_asap::{SketchAlgorithm as A, SketchParams as P}; + let (algorithm, width, depth, capacity) = match (kind.algorithm(), kind.params()) { + ( + A::CmsWithHeap, + P::CmsWithHeap { + width, + depth, + heap_size, + }, + ) => (FrequencyAlgorithm::Cms, *width, *depth, *heap_size), + ( + A::CountSketchWithHeap, + P::CountSketchWithHeap { + width, + depth, + heap_size, + }, + ) if depth % 2 == 1 => (FrequencyAlgorithm::CountSketch, *width, *depth, *heap_size), + _ => { + return Err(Error::Invalid( + "unsupported weighted frequency family or depth".into(), + )) + } + }; + if width == 0 || depth == 0 || capacity == 0 { + return Err(Error::Invalid( + "invalid weighted frequency dimensions".into(), + )); + } + Ok((algorithm, width as usize, depth as usize, capacity as usize)) + } + + pub(crate) fn algorithm(&self) -> FrequencyAlgorithm { + self.inner.algorithm() + } + pub(crate) fn shape(&self) -> (usize, usize, usize) { + self.inner.shape() + } + pub fn new( + algorithm: FrequencyAlgorithm, + width: usize, + depth: usize, + capacity: usize, + ) -> Result { + Kernel::new(algorithm, width, depth, capacity) + .map(|inner| Self { inner }) + .map_err(adapt_error) + } + pub fn from_bytes(bytes: &[u8]) -> Result { + Kernel::from_bytes(bytes) + .map(|inner| Self { inner }) + .map_err(adapt_error) + } + pub fn update(&mut self, values: &[Value], weight: f64) -> Result<(), Error> { + let values = values.iter().map(identity).collect::, _>>()?; + self.inner.update(&values, weight).map_err(adapt_error) + } + pub fn rows(&self, n: usize) -> Vec> { + self.inner + .topk(n) + .into_iter() + .map(|(items, score)| { + let mut row = items.into_iter().map(value).collect::>(); + row.push(Value::Float64(score)); + row + }) + .collect() + } +} +impl SerializableToSink for WeightedFrequency { + fn serialize_to_json(&self) -> serde_json::Value { + serde_json::to_value(self).expect("finite validated frequency state") + } + fn serialize_to_bytes(&self) -> Vec { + self.inner.to_bytes() + } +} +impl AggregateCore for WeightedFrequency { + fn clone_boxed_core(&self) -> Box { + Box::new(self.clone()) + } + fn type_name(&self) -> &'static str { + "WeightedFrequency" + } + fn as_any(&self) -> &dyn std::any::Any { + self + } + fn as_any_mut(&mut self) -> &mut dyn std::any::Any { + self + } + fn merge_with( + &self, + other: &dyn AggregateCore, + ) -> Result, Box> { + let other = other + .as_any() + .downcast_ref::() + .ok_or("weighted frequency state type mismatch")?; + Ok(Box::new(Self { + inner: self.inner.merge(&other.inner)?, + })) + } + fn get_accumulator_type(&self) -> AggregationType { + match self.inner.algorithm() { + FrequencyAlgorithm::Cms => AggregationType::CountMinSketchWithHeap, + FrequencyAlgorithm::CountSketch => AggregationType::CountSketchWithHeap, + } + } + fn get_keys(&self) -> Option> { + None + } + fn query_statistic( + &self, + _: Statistic, + _: &Option, + _: &HashMap, + ) -> Result> { + Err("weighted frequency uses typed row readout".into()) + } + fn approx_memory_bytes(&self) -> usize { + self.inner.approx_memory_bytes() + } +} +#[cfg(test)] +mod tests { + use super::*; + // Signed fractional updates and merges retain numeric ranking, not magnitude ranking. + #[test] + fn count_sketch_signed_updates_roundtrip_and_merge() { + let mut left = WeightedFrequency::new(FrequencyAlgorithm::CountSketch, 4096, 5, 8).unwrap(); + left.update(&[Value::Int64(1)], -10.5).unwrap(); + left.update(&[Value::Null], 0.125).unwrap(); + left.update(&[Value::Null], -0.0625).unwrap(); + let mut right = + WeightedFrequency::new(FrequencyAlgorithm::CountSketch, 4096, 5, 8).unwrap(); + right.update(&[Value::Null], 0.25).unwrap(); + let merged = left.merge_with(&right).unwrap(); + let merged = merged.as_any().downcast_ref::().unwrap(); + let decoded = WeightedFrequency::from_bytes(&merged.serialize_to_bytes()).unwrap(); + let rows = decoded.rows(2); + assert!(matches!(rows[0][0], Value::Null)); + assert!(matches!(rows[0][1], Value::Float64(0.3125))); + assert!(matches!(rows[1][1], Value::Float64(-10.5))); + assert!(left + .merge_with(&WeightedFrequency::new(FrequencyAlgorithm::Cms, 4096, 5, 8).unwrap()) + .is_err()); + let before = left.serialize_to_bytes(); + for weight in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + assert!(left.update(&[Value::Null], weight).is_err()); + assert_eq!(left.serialize_to_bytes(), before); + } + } + + // Invalid rates must not mutate state; typed keys cannot collide by formatting. + #[test] + fn fractional_updates_typed_identities_and_invalid_weights() { + let mut state = WeightedFrequency::new(FrequencyAlgorithm::Cms, 4096, 5, 8).unwrap(); + state.update(&[Value::Int64(1)], 0.125).unwrap(); + state.update(&[Value::Int64(1)], 0.125).unwrap(); + state.update(&[Value::Utf8("1".into())], 0.5).unwrap(); + state.update(&[Value::Null], 0.75).unwrap(); + let before = state.serialize_to_bytes(); + let decoded = WeightedFrequency::from_bytes(&before).unwrap(); + assert_eq!(decoded.rows(8).len(), 3); + assert!(WeightedFrequency::from_bytes(b"old integer state").is_err()); + for weight in [-1.0, f64::INFINITY, f64::NAN] { + assert!(state.update(&[Value::Null], weight).is_err()); + assert_eq!(state.serialize_to_bytes(), before); + } + let rows = state.rows(8); + assert_eq!(rows.len(), 3); + assert!(matches!(rows[0][0], Value::Null)); + assert!(matches!(rows[1][0], Value::Utf8(_))); + assert!(matches!(rows[2][1], Value::Float64(0.25))); + } + // Merge uses the same Float64 state representation and rejects other shapes. + #[test] + fn compatible_merge_preserves_fractional_weights() { + let mut left = WeightedFrequency::new(FrequencyAlgorithm::Cms, 4096, 5, 8).unwrap(); + let mut right = left.clone(); + left.update(&[Value::Int64(7)], 0.125).unwrap(); + right.update(&[Value::Int64(7)], 0.25).unwrap(); + let merged = left.merge_with(&right).unwrap(); + let merged = merged.as_any().downcast_ref::().unwrap(); + assert!(matches!(merged.rows(1)[0][1], Value::Float64(0.375))); + assert!(left + .merge_with(&WeightedFrequency::new(FrequencyAlgorithm::Cms, 32, 5, 8).unwrap()) + .is_err()); + } +} diff --git a/crates/asap-physical-operators/src/values.rs b/crates/asap-physical-operators/src/values.rs new file mode 100644 index 00000000..60dc1c70 --- /dev/null +++ b/crates/asap-physical-operators/src/values.rs @@ -0,0 +1,416 @@ +//! Runtime values preserve Planner schemas; summary states are typed values too. +use crate::AggregateCore; +use crate::Error; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, +}; +use std::{cmp::Ordering, sync::Arc}; +pub type Schema = Arc; +#[derive(Clone, serde::Serialize, serde::Deserialize)] +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)]>), + #[serde(skip)] + 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), + (Self::Map(left), Self::Map(right)) => { + let mut result = Ordering::Equal; + for ((lk, lv), (rk, rv)) in left.iter().zip(right.iter()) { + result = lk.compare(rk)?; + if result != Ordering::Equal { + break; + } + result = match (lv, rv) { + (Self::Null, Self::Null) => Ordering::Equal, + (Self::Null, _) => Ordering::Greater, + (_, Self::Null) => Ordering::Less, + _ => lv.compare(rv)?, + }; + if result != Ordering::Equal { + break; + } + } + if result == Ordering::Equal { + left.len().cmp(&right.len()) + } else { + result + } + } + _ => { + 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.capacity() * 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) use crate::capability::validate_native_family as validate_family; + +fn validate_state(family: &SummaryFamilyType, state: &dyn AggregateCore) -> Result<(), Error> { + use crate::summary_kernels::{ + datasketches_kll::DatasketchesKLLAccumulator, dd_sketch::DDSketchAccumulator, + exact::ExactAccumulator, hll_sketch::HllSketchAccumulator, + }; + use planner_types::post_asap::SketchParams; + validate_family(family)?; + let valid = match family { + SummaryFamilyType::Sketch(kind, _) + if matches!( + kind.params(), + SketchParams::CmsWithHeap { .. } | SketchParams::CountSketchWithHeap { .. } + ) => + { + use crate::summary_kernels::weighted_frequency::WeightedFrequency; + let (algorithm, width, depth, capacity) = WeightedFrequency::configuration(kind)?; + state + .as_any() + .downcast_ref::() + .is_some_and(|state| { + state.algorithm() == algorithm && state.shape() == (width, depth, capacity) + }) + } + + 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(()) +} + +pub(crate) fn field(schema: &Schema, column: usize) -> Result<&SummaryField, Error> { + schema + .fields + .get(column) + .ok_or_else(|| Error::Invalid("column out of range".into())) +} +pub(crate) fn plain(schema: &Schema, column: usize) -> Result<(&DataType, bool), Error> { + let f = field(schema, column)?; + let SummaryFamilyType::Plain(dtype) = &f.dtype else { + return Err(Error::Invalid("plain value required".into())); + }; + Ok((dtype, f.nullable)) +} + +#[cfg(test)] +mod weighted_state_tests { + use super::*; + use crate::summary_kernels::weighted_frequency::{FrequencyAlgorithm, WeightedFrequency}; + use planner_types::post_asap::{SketchAlgorithm, SketchKind, SketchParams}; + + // A state cannot acquire a different family or shape merely by relabeling its batch. + #[test] + fn weighted_state_family_and_shape_must_match() { + let cms = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::CmsWithHeap, + SketchParams::CmsWithHeap { + width: 32, + depth: 5, + heap_size: 8, + }, + ), + Default::default(), + ); + let cs = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::CountSketchWithHeap, + SketchParams::CountSketchWithHeap { + width: 32, + depth: 5, + heap_size: 8, + }, + ), + Default::default(), + ); + let state = WeightedFrequency::new(FrequencyAlgorithm::CountSketch, 32, 5, 8).unwrap(); + assert!(validate_state(&cs, &state).is_ok()); + assert!(validate_state(&cms, &state).is_err()); + let wrong_shape = + WeightedFrequency::new(FrequencyAlgorithm::CountSketch, 64, 5, 8).unwrap(); + assert!(validate_state(&cs, &wrong_shape).is_err()); + let even_depth = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::CountSketchWithHeap, + SketchParams::CountSketchWithHeap { + width: 32, + depth: 4, + heap_size: 8, + }, + ), + Default::default(), + ); + assert!(validate_family(&even_depth).is_err()); + } +} diff --git a/crates/asap-physical-operators/tests/blocking_resources.rs b/crates/asap-physical-operators/tests/blocking_resources.rs new file mode 100644 index 00000000..19381175 --- /dev/null +++ b/crates/asap-physical-operators/tests/blocking_resources.rs @@ -0,0 +1,240 @@ +//! Blocking operators enforce resources before returning their first batch. +use asap_physical_operators::dag::{ + operators::Operator, + values::{Batch, Schema, Value}, + Error, Limits, PhysicalDag, PhysicalOperator, RunContext, Scope, +}; +use futures::{executor::block_on, FutureExt, StreamExt}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::{DataType, JoinKind, Predicate, QueryExpr, ScalarValue}, +}; +use std::sync::Arc; + +fn schema(width: usize) -> Schema { + Arc::new(SummarySchema { + fields: (0..width) + .map(|i| SummaryField { + name: format!("v{i}"), + dtype: SummaryFamilyType::Plain(DataType::Int64), + nullable: false, + }) + .collect(), + time_index: None, + }) +} +fn context(max_bytes: usize) -> RunContext { + RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits { + max_bytes, + ..Limits::default() + }, + ) + .unwrap() +} +fn source(n: usize) -> PhysicalDag<'static, Batch, Schema> { + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source( + schema(1), + vec![Batch::try_new(schema(1), vec![vec![Value::Int64(1)]; n]).unwrap()], + ) + .unwrap(), + ) + .unwrap(); + dag +} +fn cross_join() -> Operator { + Operator::relational_join( + schema(1), + schema(1), + JoinKind::Cross, + &Predicate(std::rc::Rc::new(QueryExpr::Literal(ScalarValue::Boolean( + true, + )))), + schema(2), + ) + .unwrap() +} + +// Even callers starting an operator directly cannot bypass its workspace budget. +#[test] +fn join_reserves_result_growth_before_returning_output() { + let sources = source(64); + let run = context(32 * 1024); + let inputs = sources.execute(&[0, 0], run.clone()).unwrap(); + let join = cross_join(); + let mut output = join.start(inputs, run.clone()).unwrap(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::MemoryLimit)) + )); + drop(output); + assert_eq!(run.retained_bytes(), 0); +} + +// A single large input batch must not monopolize the worker during a join. +#[test] +fn join_yields_during_computation_and_observes_cancellation() { + let sources = source(64); + let run = context(16 * 1024 * 1024); + let inputs = sources.execute(&[0, 0], run.clone()).unwrap(); + let join = cross_join(); + let mut output = join.start(inputs, run.clone()).unwrap(); + assert!( + output.next().now_or_never().is_none(), + "join should yield before producing all 4096 rows" + ); + run.cancel(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::Cancelled)) + )); + drop(output); + assert_eq!(run.retained_bytes(), 0); +} + +// Sorting and grouping yield even for one large batch. +#[test] +fn blocking_reductions_yield_and_release_memory_on_cancellation() { + use asap_physical_operators::{ + operators::{Reduction, SortKey}, + plan::PhysicalOperator, + }; + let operators = vec![ + Operator::sort( + schema(1), + vec![SortKey { + column: 0, + descending: false, + nulls_first: false, + }], + vec![], + ) + .unwrap(), + Operator::aggregate(schema(1), vec![], vec![("sum".into(), Reduction::Sum(0))]).unwrap(), + ]; + for operator in operators { + let sources = source(768); + let run = context(16 * 1024 * 1024); + let inputs = sources.execute(&[0], run.clone()).unwrap(); + let mut output = operator.start(inputs, run.clone()).unwrap(); + assert!(output.next().now_or_never().is_none()); + run.cancel(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::Cancelled)) + )); + drop(output); + assert_eq!(run.retained_bytes(), 0); + } +} + +// Merge-sort rounds preserve input order for tied keys across chunk boundaries. +#[test] +fn cooperative_sort_preserves_ties_across_chunks() { + use asap_physical_operators::operators::SortKey; + let batch = Batch::try_new( + schema(2), + (0..1025) + .rev() + .map(|i| vec![Value::Int64(i % 3), Value::Int64(i)]) + .collect(), + ) + .unwrap(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], Operator::source(schema(2), vec![batch]).unwrap()) + .unwrap(); + dag.add( + 1, + vec![0], + Operator::sort( + schema(2), + vec![SortKey { + column: 0, + descending: false, + nulls_first: false, + }], + vec![], + ) + .unwrap(), + ) + .unwrap(); + let mut output = dag + .execute(&[1], context(16 * 1024 * 1024)) + .unwrap() + .remove(0); + let batch = block_on(output.next()).unwrap().unwrap(); + let expected = (0..3) + .flat_map(|key| (0..1025).rev().filter(move |i| i % 3 == key)) + .collect::>(); + for (row, expected) in batch.rows().iter().zip(expected) { + assert!(matches!(row[1], Value::Int64(i) if i == expected)); + } + assert_eq!(batch.rows().len(), 1025); +} + +// The integrated weighted-summary path obeys the same cooperative cancellation contract. +#[test] +fn weighted_summary_build_yields_within_a_batch() { + use planner_types::post_asap::{SketchAlgorithm, SketchKind, SketchParams}; + let input = Arc::new(SummarySchema { + fields: vec![ + SummaryField { + name: "item".into(), + dtype: SummaryFamilyType::Plain(DataType::Int64), + nullable: false, + }, + SummaryField { + name: "weight".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }, + ], + time_index: None, + }); + let mut sources = PhysicalDag::default(); + let batch = Batch::try_new( + input.clone(), + (0..1500) + .map(|i| vec![Value::Int64(i % 8), Value::Float64(0.25)]) + .collect(), + ) + .unwrap(); + sources + .add( + 0, + vec![], + Operator::source(input.clone(), vec![batch]).unwrap(), + ) + .unwrap(); + let family = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::CmsWithHeap, + SketchParams::CmsWithHeap { + width: 64, + depth: 3, + heap_size: 8, + }, + ), + Default::default(), + ); + let operator = Operator::keyed_summary_build(input, family, 1, vec![0], vec![]).unwrap(); + let run = context(16 * 1024 * 1024); + let inputs = sources.execute(&[0], run.clone()).unwrap(); + let mut output = operator.start(inputs, run.clone()).unwrap(); + assert!(output.next().now_or_never().is_none()); + run.cancel(); + assert!(matches!( + block_on(output.next()), + Some(Err(Error::Cancelled)) + )); + drop(output); + assert_eq!(run.retained_bytes(), 0); +} diff --git a/crates/asap-physical-operators/tests/current_series_heap.rs b/crates/asap-physical-operators/tests/current_series_heap.rs new file mode 100644 index 00000000..7abbeb04 --- /dev/null +++ b/crates/asap-physical-operators/tests/current_series_heap.rs @@ -0,0 +1,399 @@ +//! Spatial heap weights come from a fresh instant vector, never sample history. +use asap_physical_operators::{ + operators::Operator, + physical_planner::{ + promql_rows::{decode_series_identity, series_row, SERIES_IDENTITY_COLUMN}, + CompiledPhysicalDag, InputContract, Source, + }, + runtime::{Limits, RunContext, Scope}, + values::{Batch, Value}, +}; +use futures::{executor::block_on, StreamExt}; +use planner_types::{post_asap::*, pre_asap::DataType}; +use std::{collections::BTreeMap, sync::Arc}; + +fn schema() -> Arc { + Arc::new(SummarySchema { + fields: [ + ("ts", DataType::Timestamp), + ("value", DataType::Float64), + ("job", DataType::Utf8), + (SERIES_IDENTITY_COLUMN, DataType::Utf8), + ] + .into_iter() + .map(|(name, dtype)| SummaryField { + name: name.into(), + dtype: SummaryFamilyType::Plain(dtype), + nullable: false, + }) + .collect(), + time_index: Some(0), + }) +} +fn run(program: &CompiledPhysicalDag, data: Batch, end: i64) -> Result, String> { + let recovered = CompiledPhysicalDag::decode(&program.encode().map_err(|e| e.to_string())?) + .map_err(|e| e.to_string())?; + let input_id = recovered.input_contracts().next().unwrap().0; + let graph = recovered + .instantiate(BTreeMap::from([( + input_id, + Box::new(Operator::source(data.schema().clone(), vec![data]).unwrap()) as Source<'_>, + )])) + .unwrap(); + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: end, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + block_on(async { + let mut stream = graph.execute(recovered.roots(), context).unwrap().remove(0); + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push((*batch.map_err(|e| e.to_string())?).clone()); + } + Ok(batches) + }) +} +fn input(samples: &[(&str, i64, f64)]) -> Batch { + let schema = schema(); + let rows = samples + .iter() + .map(|(instance, time, value)| { + series_row( + &schema, + &BTreeMap::from([ + ("job".into(), "api".into()), + ("hidden_instance".into(), (*instance).into()), + ]), + *time, + *value, + ) + .unwrap() + }) + .collect(); + Batch::try_new(schema, rows).unwrap() +} +fn snapshot_plan() -> CompiledPhysicalDag { + CompiledPhysicalDag::from_operators( + BTreeMap::from([(0, InputContract::bounded(schema()))]), + BTreeMap::from([( + 1, + ( + vec![0], + Operator::current_series(schema(), 3, 0, 1, 60_000).unwrap(), + ), + )]), + vec![1], + ) + .unwrap() +} + +// Replacement, expiry and stale markers act before sketch updates. Hidden labels +// survive even when every series has the same projected `job` value. +#[test] +fn latest_snapshot_replaces_decreases_expires_and_retains_full_identity() { + let plan = snapshot_plan(); + let batches = run( + &plan, + input(&[ + ("decrease", 10_000, 100.), + ("decrease", 50_000, 1.), + ("steady", 40_000, 20.), + ("expired", 0, 1_000.), + ("stale", 20_000, 500.), + ("stale", 55_000, f64::from_bits(0x7ff0_0000_0000_0002)), + ("future", 60_001, 2_000.), + ]), + 60_000, + ) + .unwrap(); + let values = batches + .iter() + .flat_map(|batch| batch.rows()) + .map(|row| { + let Value::Utf8(identity) = &row[3] else { + panic!() + }; + let Value::Float64(value) = row[1] else { + panic!() + }; + assert!(matches!(row[0], Value::Timestamp(60_000))); + ( + decode_series_identity(identity).unwrap()["hidden_instance"].clone(), + value, + ) + }) + .collect::>(); + assert_eq!( + values, + BTreeMap::from([("decrease".into(), 1.), ("steady".into(), 20.)]) + ); + assert!(run(&plan, input(&[("steady", 40_000, 20.)]), 100_000) + .unwrap() + .iter() + .all(|batch| batch.rows().is_empty())); + assert!(run( + &plan, + input(&[("conflict", 50_000, 1.), ("conflict", 50_000, 2.)]), + 60_000 + ) + .is_err()); +} + +#[test] +fn spatial_heap_ranks_latest_values_in_independent_runs() { + for algorithm in [ + SketchAlgorithm::CmsWithHeap, + SketchAlgorithm::CountSketchWithHeap, + ] { + let params = match algorithm { + SketchAlgorithm::CmsWithHeap => SketchParams::CmsWithHeap { + width: 2048, + depth: 5, + heap_size: 100, + }, + _ => SketchParams::CountSketchWithHeap { + width: 2048, + depth: 5, + heap_size: 100, + }, + }; + let family = + SummaryFamilyType::Sketch(SketchKind::new(algorithm, params), Default::default()); + let build = Operator::keyed_summary_build(schema(), family, 1, vec![3], vec![2]).unwrap(); + let output = Arc::new(SummarySchema { + fields: vec![ + schema().fields[2].clone(), + schema().fields[3].clone(), + schema().fields[1].clone(), + ], + time_index: None, + }); + let read = Operator::keyed_readout(build.schema(), 1, 1, output).unwrap(); + let plan = CompiledPhysicalDag::from_operators( + BTreeMap::from([(0, InputContract::bounded(schema()))]), + BTreeMap::from([ + ( + 1, + ( + vec![0], + Operator::current_series(schema(), 3, 0, 1, 60_000).unwrap(), + ), + ), + (2, (vec![1], build)), + (3, (vec![2], read)), + ]), + vec![3], + ) + .unwrap(); + for (samples, end, winner, score) in [ + ( + vec![("a", 10_000, 100.), ("a", 50_000, 1.), ("b", 50_000, 20.)], + 60_000, + "b", + 20., + ), + ( + vec![("a", 110_000, 3.), ("b", 50_000, 20.)], + 120_000, + "a", + 3., + ), + ] { + let batches = run(&plan, input(&samples), end).unwrap(); + let rows = batches + .iter() + .flat_map(|batch| batch.rows()) + .collect::>(); + assert_eq!(rows.len(), 1); + let Value::Utf8(encoded) = &rows[0][1] else { + panic!() + }; + assert_eq!( + decode_series_identity(encoded).unwrap()["hidden_instance"], + winner + ); + assert!(matches!(rows[0][2], Value::Float64(actual) if actual == score)); + } + } +} + +// Blocking membership selection shares the run's cancellation and byte budget. +#[test] +fn current_series_observes_resource_limits() { + use asap_physical_operators::Error; + let plan = snapshot_plan(); + for cancelled in [false, true] { + let data = input(&[("one", 50_000, 1.)]); + let graph = plan + .instantiate(BTreeMap::from([( + 0, + Box::new(Operator::source(data.schema().clone(), vec![data]).unwrap()) + as Source<'_>, + )])) + .unwrap(); + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 60_000, + revision: 0, + }, + Limits { + max_bytes: if cancelled { 1 << 20 } else { 1 }, + ..Limits::default() + }, + ) + .unwrap(); + if cancelled { + context.cancel(); + } + let result = match graph.execute(&[1], context.clone()) { + Err(error) => Err(error), + Ok(mut streams) => block_on(streams.remove(0).next()).unwrap().map(|_| ()), + }; + assert!(matches!( + (cancelled, result), + (true, Err(Error::Cancelled)) | (false, Err(Error::MemoryLimit)) + )); + assert_eq!(context.retained_bytes(), 0); + } +} + +#[test] +fn identity_encoding_is_lossless_and_rejects_noncanonical_inputs() { + use asap_physical_operators::physical_planner::promql_rows::encode_series_identity; + let labels = BTreeMap::from([ + ("a".into(), "quote\"slash\\".into()), + ("other".into(), "".into()), + ]); + assert_eq!( + decode_series_identity(&encode_series_identity(&labels).unwrap()).unwrap(), + labels + ); + for invalid in [ + "[]", + "{\"a\":1}", + "{\"a\":\"x\",\"a\":\"x\"}", + "{ \"a\":\"x\"}", + ] { + assert!(decode_series_identity(invalid).is_err(), "{invalid}"); + } +} + +// The actual Planner population candidate lowers to native operators; this +// test does not manually assemble the computation or its dependency edges. +#[test] +fn planner_current_series_candidate_compiles_with_dynamic_identity() { + use asap_physical_operators::physical_planner::{compile, promql_rows::with_series_identity}; + use planner_types::{types::AccuracyTarget, workload::*}; + use std::rc::Rc; + let workload = PlanningWorkload { + query_workload: QueryWorkload { + language: QueryLanguage::PromQL, + query_batch: Some(vec![BatchEntry { + query: Query("topk by(job)(1, m)".into()), + requirements: QueryRequirements { + accuracy: AccuracyRequirement::Explicit(AccuracyTarget::Exact), + ..Default::default() + }, + predictability: Predictability::Unknown, + invocations: 1, + execute_at: None, + time_selection: TimeSelection::default(), + }]), + repeating_queries: None, + }, + data_workload: Some(DataWorkload { + data_ingestion_interval: Evidence { + value: Some(DurationMs(60_000)), + ..Default::default() + }, + ..Default::default() + }), + }; + let original = asap_frontend_promql::lower_promql_workload(&workload, 0) + .unwrap() + .remove(0); + let open_root = Rc::new(original.clone()); + let open_selected = + asap_aware_mapping::maintained_population::MaintainedPopulationStrategy::new( + std::slice::from_ref(&open_root), + ) + .candidate(&open_root) + .unwrap(); + let snapshot_program = + asap_physical_operators::physical_planner::promql_rows::compile_current_series_readout( + &open_selected, + ) + .unwrap(); + let encoded = String::from_utf8(snapshot_program.encode().unwrap()).unwrap(); + assert!( + !encoded.contains("CurrentSeries"), + "maintained input must not be rebuilt" + ); + assert!(encoded.contains("Sort") && encoded.contains("Limit")); + assert_eq!(snapshot_program.input_contracts().count(), 1); + let root = Rc::new(with_series_identity(&original).unwrap()); + let selected = asap_aware_mapping::maintained_population::MaintainedPopulationStrategy::new( + std::slice::from_ref(&root), + ) + .candidate(&root) + .unwrap(); + let logical = compile_executable_dag(&selected).unwrap(); + let raw = logical + .nodes + .iter() + .find(|node| matches!(node.payload, ExecutableOperatorPayload::Fallback { .. })) + .unwrap(); + let raw_schema = Arc::new(raw.output_schema.clone()); + let physical = compile( + &logical, + BTreeMap::from([( + u64::from(raw.id.0), + InputContract::bounded(raw_schema.clone()), + )]), + &[u64::from(logical.root.0)], + ) + .unwrap(); + let bytes = String::from_utf8(physical.encode().unwrap()).unwrap(); + assert!(bytes.contains("CurrentSeries")); + assert!(bytes.contains("Sort")); + assert!(bytes.contains("Limit")); + let rows = [("a", 10_000, 100.), ("a", 50_000, 1.), ("b", 50_000, 20.)] + .into_iter() + .map(|(member, at, value)| { + series_row( + &raw_schema, + &BTreeMap::from([ + ("job".into(), "api".into()), + ("unreferenced".into(), member.into()), + ]), + at, + value, + ) + .unwrap() + }) + .collect(); + let batches = run(&physical, Batch::try_new(raw_schema, rows).unwrap(), 60_000).unwrap(); + let rows = batches + .iter() + .flat_map(|batch| batch.rows()) + .collect::>(); + assert_eq!(rows.len(), 1); + assert!(matches!(rows[0][1], Value::Float64(20.))); + let id = batches[0] + .schema() + .fields + .iter() + .position(|field| field.name == SERIES_IDENTITY_COLUMN) + .unwrap(); + let Value::Utf8(encoded) = &rows[0][id] else { + panic!() + }; + assert_eq!( + decode_series_identity(encoded).unwrap()["unreferenced"], + "b" + ); +} diff --git a/crates/asap-physical-operators/tests/deployment.rs b/crates/asap-physical-operators/tests/deployment.rs new file mode 100644 index 00000000..ffdd3aea --- /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 00000000..e19cbed2 --- /dev/null +++ b/crates/asap-physical-operators/tests/physical_dag.rs @@ -0,0 +1,1191 @@ +//! 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::{summary_kernels::exact::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, 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 { + 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 { + 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.]); +} + +// Planner wire nodes, including grouping and edge roles, are executable at either phase. +#[test] +fn planner_semijoin_sort_limit_contract_at_both_phases() { + use asap_physical_operators::dag::planner::{bind, Source}; + use planner_types::{ + post_asap::*, + pre_asap::{CompareOpKind, GroupKeys, JoinKind, Predicate, QueryExpr, SortKey}, + }; + use std::{collections::BTreeMap, rc::Rc}; + let rows_schema = schema(&[ + ("group", DataType::Utf8, false), + ("key", DataType::Utf8, false), + ("score", DataType::Float64, false), + ]); + let keys_schema = schema(&[("key", DataType::Utf8, false)]); + let node = |id, payload, schema: &Schema| ExecutableDagNode { + id: PostAsapNodeId(id), + payload, + output_schema: (**schema).clone(), + output_state: ExecutionDataState::QUERY_ROWS, + guarantee: None, + }; + let edge = |producer, consumer, role, schema: &Schema| ExecutableDagEdge { + producer: PostAsapNodeId(producer), + consumer: PostAsapNodeId(consumer), + role, + intermediate_schema: (**schema).clone(), + data_state: ExecutionDataState::QUERY_ROWS, + grouping: GroupingEdgeCompatibility::NotApplicable, + window: WindowEdgeCompatibility::NotApplicable, + }; + let groups = GroupKeys::by(vec![0]); + let dag = ExecutableDag { + nodes: vec![ + node( + 0, + ExecutableOperatorPayload::Fallback { + expression: QueryExpr::promql_scalar(0.), + }, + &rows_schema, + ), + node( + 1, + ExecutableOperatorPayload::Fallback { + expression: QueryExpr::promql_scalar(0.), + }, + &keys_schema, + ), + node( + 2, + ExecutableOperatorPayload::RelationalJoin { + join_kind: JoinKind::Semi, + pruning: None, + pred: Predicate(Rc::new(QueryExpr::Compare { + left: Rc::new(QueryExpr::Column(1)), + op: CompareOpKind::Eq, + right: Rc::new(QueryExpr::Column(3)), + })), + }, + &rows_schema, + ), + node( + 3, + ExecutableOperatorPayload::Value { + operation: ValueOperation::Sort { + keys: vec![SortKey { + expr: QueryExpr::Column(2), + ascending: false, + nulls_first: false, + }], + partition_by: groups.clone(), + }, + }, + &rows_schema, + ), + node( + 4, + ExecutableOperatorPayload::Value { + operation: ValueOperation::Limit { + n: 1, + offset: 0, + partition_by: groups, + }, + }, + &rows_schema, + ), + ], + // Deliberately put Right before Left: list order must not swap inputs. + edges: vec![ + edge(1, 2, EdgeRole::Right, &keys_schema), + edge(0, 2, EdgeRole::Left, &rows_schema), + edge(2, 3, EdgeRole::Input, &rows_schema), + edge(3, 4, EdgeRole::Input, &rows_schema), + ], + root: PostAsapNodeId(4), + }; + let text = |v: &str| Value::Utf8(v.into()); + for (phase, scope) in [ + (ExecutionTiming::QueryTime, query()), + ( + ExecutionTiming::IngestionTime, + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 1000, + revision: 2, + }, + ), + ] { + let dag = dag + .with_execution_phases(&dag.nodes.iter().map(|node| (node.id, phase)).collect()) + .unwrap(); + let sources: BTreeMap> = BTreeMap::from([ + ( + 0, + Box::new( + Operator::source( + rows_schema.clone(), + vec![Batch::try_new( + rows_schema.clone(), + vec![ + vec![text("a"), text("x"), Value::Float64(8.)], + vec![text("a"), text("y"), Value::Float64(9.)], + vec![text("b"), text("x"), Value::Float64(2.)], + vec![text("b"), text("z"), Value::Float64(99.)], + ], + ) + .unwrap()], + ) + .unwrap(), + ) as Source<'static>, + ), + ( + 1, + Box::new( + Operator::source( + keys_schema.clone(), + vec![Batch::try_new( + keys_schema.clone(), + vec![vec![text("x")], vec![text("y")]], + ) + .unwrap()], + ) + .unwrap(), + ) as Source<'static>, + ), + ]); + let native = bind(&dag, sources, &[4]).unwrap(); + let mut scores = floats(&run(&native, 4, scope), 2); + scores.sort_by(f64::total_cmp); + assert_eq!(scores, vec![2., 9.]); + } +} + +// Planner scalar signatures, collection access and null predicates share native execution. +#[test] +fn planner_expressions_preserve_collection_and_nullable_types() { + use asap_physical_operators::dag::expressions::CompiledExpression; + use planner_types::pre_asap::{CompareOpKind, QueryExpr, ScalarValue}; + use std::rc::Rc; + let input_schema = schema(&[( + "items", + DataType::Map { + key: Box::new(DataType::Utf8), + value: Box::new(DataType::Int64), + value_nullable: false, + }, + false, + )]); + let access = QueryExpr::FunctionCall { + name: "asap_element_access".into(), + args: vec![ + QueryExpr::Column(0), + QueryExpr::Literal(ScalarValue::Utf8("count".into())), + ], + }; + let project = Operator::project( + input_schema.clone(), + vec![( + "count".into(), + Expression::planner(CompiledExpression::compile(&access, &input_schema).unwrap()), + )], + ) + .unwrap(); + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source( + input_schema.clone(), + vec![Batch::try_new( + input_schema.clone(), + vec![ + vec![Value::Map( + vec![(Value::Utf8("count".into()), Value::Int64(7))].into(), + )], + vec![Value::Map(Arc::from([]))], + ], + ) + .unwrap()], + ) + .unwrap(), + ) + .unwrap(); + let projected = project.schema(); + dag.add(1, vec![0], project).unwrap(); + let predicate = QueryExpr::Compare { + left: Rc::new(QueryExpr::Column(0)), + op: CompareOpKind::Ge, + right: Rc::new(QueryExpr::Literal(ScalarValue::Int64(1))), + }; + dag.add( + 2, + vec![1], + Operator::filter( + projected.clone(), + Expression::planner(CompiledExpression::compile(&predicate, &projected).unwrap()), + ) + .unwrap(), + ) + .unwrap(); + let rows = run(&dag, 2, query()); + assert!(matches!(rows.as_slice(),[row] if matches!(row.as_slice(),[Value::Int64(7)]))); + let unknown = QueryExpr::FunctionCall { + name: "unregistered_function".into(), + args: vec![QueryExpr::Column(0)], + }; + assert!(CompiledExpression::compile(&unknown, &input_schema).is_err()); +} + +// Outer, semi and anti joins share Planner predicates and preserve SQL null behavior. +#[test] +fn native_relational_join_kinds_preserve_unmatched_rows() { + use planner_types::pre_asap::{CompareOpKind, JoinKind, Predicate, QueryExpr}; + use std::rc::Rc; + let input = schema(&[("key", DataType::Int64, true)]); + let predicate = Predicate(Rc::new(QueryExpr::Compare { + left: Rc::new(QueryExpr::Column(0)), + op: CompareOpKind::Eq, + right: Rc::new(QueryExpr::Column(1)), + })); + for (kind, count) in [ + (JoinKind::Inner, 1), + (JoinKind::Left, 3), + (JoinKind::Right, 3), + (JoinKind::Full, 5), + (JoinKind::Semi, 1), + (JoinKind::Anti, 2), + (JoinKind::Cross, 9), + ] { + let output = if matches!(kind, JoinKind::Semi | JoinKind::Anti) { + input.clone() + } else { + schema(&[ + ("left", DataType::Int64, true), + ("right", DataType::Int64, true), + ]) + }; + let mut dag = PhysicalDag::default(); + for (id, rows) in [ + ( + 0, + vec![ + vec![Value::Int64(1)], + vec![Value::Int64(2)], + vec![Value::Null], + ], + ), + ( + 1, + vec![ + vec![Value::Int64(2)], + vec![Value::Int64(3)], + vec![Value::Null], + ], + ), + ] { + dag.add( + id, + vec![], + Operator::source( + input.clone(), + vec![Batch::try_new(input.clone(), rows).unwrap()], + ) + .unwrap(), + ) + .unwrap(); + } + dag.add( + 2, + vec![0, 1], + Operator::relational_join( + input.clone(), + input.clone(), + kind.clone(), + &predicate, + output, + ) + .unwrap(), + ) + .unwrap(); + assert_eq!(run(&dag, 2, query()).len(), count, "{kind:?}"); + } +} + +// Per-series fractional rates feed either weighted frequency family per job, in either scope. +#[test] +fn weighted_rate_topk_preserves_partitions_fractional_scores_and_evaluation_scope() { + for count_sketch in [false, true] { + assert_weighted_rate_topk(count_sketch); + } +} +fn assert_weighted_rate_topk(count_sketch: bool) { + use planner_types::post_asap::{SketchAlgorithm, SketchKind, SketchParams}; + let raw = schema(&[ + ("service", DataType::Utf8, false), + ("job", DataType::Utf8, false), + ("instance", DataType::Int64, false), + ("t", DataType::Timestamp, false), + ("value", DataType::Float64, false), + ]); + let mut rows = Vec::new(); + // Multiple instances of auth accumulate. Batch has a very different scale. + for (service, job, instance, rate) in [ + ("auth", "api", 1, 0.125), + ("auth", "api", 2, 0.25), + ("checkout", "api", 1, 0.3125), + ("search", "api", 1, 0.0625), + ("ingest", "batch", 1, 100.0), + ("export", "batch", 1, 80.0), + ("cleanup", "batch", 1, 20.0), + ] { + for (t, value) in [(0, 0.0), (30_000, rate * 30.0), (60_000, rate * 60.0)] { + rows.push(vec![ + Value::Utf8(service.into()), + Value::Utf8(job.into()), + Value::Int64(instance), + Value::Timestamp(t), + Value::Float64(value), + ]); + } + } + let rates = Operator::window( + raw.clone(), + planner_types::pre_asap::AggIntent::Rate, + 3, + 4, + vec![0, 1, 2], + Some((0, 60_000)), + ) + .unwrap(); + let family = SummaryFamilyType::Sketch( + SketchKind::new( + if count_sketch { + SketchAlgorithm::CountSketchWithHeap + } else { + SketchAlgorithm::CmsWithHeap + }, + if count_sketch { + SketchParams::CountSketchWithHeap { + width: 4096, + depth: 5, + heap_size: 8, + } + } else { + SketchParams::CmsWithHeap { + width: 4096, + depth: 5, + heap_size: 8, + } + }, + ), + Default::default(), + ); + let build = Operator::keyed_summary_build(rates.schema(), family, 3, vec![0], vec![1]).unwrap(); + let output = schema(&[ + ("job", DataType::Utf8, false), + ("service", DataType::Utf8, false), + ("score", DataType::Float64, false), + ]); + let readout = Operator::keyed_readout(build.schema(), 1, 8, output.clone()).unwrap(); + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source(raw.clone(), vec![Batch::try_new(raw, rows).unwrap()]).unwrap(), + ) + .unwrap(); + dag.add(1, vec![0], rates).unwrap(); + dag.add(2, vec![1], build).unwrap(); + dag.add(3, vec![2], readout).unwrap(); + dag.add( + 4, + vec![3], + Operator::sort( + output.clone(), + vec![SortKey { + column: 2, + descending: true, + nulls_first: false, + }], + vec![0], + ) + .unwrap(), + ) + .unwrap(); + dag.add(5, vec![4], Operator::limit(output, 2, 0, vec![0]).unwrap()) + .unwrap(); + for scope in [ + query(), + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 60_000, + revision: 2, + }, + query(), + ] { + let result = run(&dag, 5, scope); + assert_eq!(result.len(), 4); + assert_eq!(floats(&result, 2), vec![0.375, 0.3125, 100.0, 80.0]); + let services = result + .iter() + .map(|row| match &row[1] { + Value::Utf8(v) => v.as_ref(), + _ => panic!("service"), + }) + .collect::>(); + assert_eq!(services, vec!["auth", "checkout", "ingest", "export"]); + } +} + +// The grouped temporal reducer's sample schema must survive physical Sort/Limit binding. +#[test] +fn grouped_temporal_schema_compiles_and_executes_topk() { + use asap_physical_operators::physical_planner::{ + compile_node, CompiledPhysicalDag, InputContract, Source, + }; + use planner_types::post_asap::{ + ExecutableDagNode, ExecutableOperatorPayload, ExecutionDataState, PostAsapNodeId, + ValueOperation, + }; + use planner_types::pre_asap::{ + aggregate_output_schema, AggIntent, Column, GroupKeys, QueryExpr, Reduction as IrReduction, + Schema as IrSchema, + }; + let grouped = IrSchema::new(vec![ + Column::new("job", DataType::Utf8, false), + Column::new("sum", DataType::Float64, false), + ]); + let output = aggregate_output_schema( + &grouped, + &IrReduction::PerEntity, + &[AggIntent::Avg { col: None }], + &[], + ) + .unwrap(); + let input = schema( + &output + .columns + .iter() + .map(|c| (c.name.as_str(), c.dtype.clone(), c.nullable)) + .collect::>(), + ); + let node = |id, operation| ExecutableDagNode { + id: PostAsapNodeId(id), + payload: ExecutableOperatorPayload::Value { operation }, + output_state: ExecutionDataState::QUERY_ROWS, + output_schema: (*input).clone(), + guarantee: None, + }; + let sort = compile_node( + &node( + 1, + ValueOperation::Sort { + keys: vec![planner_types::pre_asap::SortKey { + expr: QueryExpr::Column(1), + ascending: false, + nulls_first: false, + }], + partition_by: GroupKeys::none(), + }, + ), + std::slice::from_ref(&input), + ) + .unwrap(); + let limit = compile_node( + &node( + 2, + ValueOperation::Limit { + n: 1, + offset: 0, + partition_by: GroupKeys::none(), + }, + ), + std::slice::from_ref(&input), + ) + .unwrap(); + let compiled = CompiledPhysicalDag::from_operators( + [(0, InputContract::bounded(input.clone()))].into(), + [(1, (vec![0], sort)), (2, (vec![1], limit))].into(), + vec![2], + ) + .unwrap(); + let recovered = CompiledPhysicalDag::decode(&compiled.encode().unwrap()).unwrap(); + assert_eq!(recovered.row_source(2), Some(0)); + assert_eq!(recovered.operator_name(2), Some("Limit")); + let expected = vec![Value::Utf8("api".into()), Value::Float64(9.)]; + let batch = Batch::try_new( + input.clone(), + vec![ + vec![Value::Utf8("worker".into()), Value::Float64(2.)], + expected.clone(), + ], + ) + .unwrap(); + let source = Box::new(Operator::source(input, vec![batch]).unwrap()) as Source<'_>; + let physical = recovered.instantiate([(0, source)].into()).unwrap(); + let mut stream = physical + .execute(&[2], RunContext::new(query(), Limits::default()).unwrap()) + .unwrap() + .remove(0); + let rows = block_on(async { + let mut rows = vec![]; + while let Some(batch) = stream.next().await { + rows.extend_from_slice(batch.unwrap().rows()); + } + rows + }); + assert_eq!(rows.len(), 1); + assert!(matches!(&rows[0][0], Value::Utf8(label) if label.as_ref() == "api")); + assert!(matches!(rows[0][1], Value::Float64(9.))); +} diff --git a/crates/asap-physical-operators/tests/physical_plan_recovery.rs b/crates/asap-physical-operators/tests/physical_plan_recovery.rs new file mode 100644 index 00000000..15a00ba5 --- /dev/null +++ b/crates/asap-physical-operators/tests/physical_plan_recovery.rs @@ -0,0 +1,102 @@ +//! Persisted physical plans recover selected operators without logical lowering. +use asap_physical_operators::{ + operators::{Operator, SortKey}, + physical_planner::{CompiledPhysicalDag, InputContract}, +}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::DataType, +}; +use std::{collections::BTreeMap, sync::Arc}; + +fn sorted() -> CompiledPhysicalDag { + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }], + time_index: None, + }); + CompiledPhysicalDag::from_operators( + BTreeMap::from([(0, InputContract::bounded(schema.clone()))]), + BTreeMap::from([( + 1, + ( + vec![0], + Operator::sort( + schema, + vec![SortKey { + column: 0, + descending: true, + nulls_first: false, + }], + vec![], + ) + .unwrap(), + ), + )]), + vec![1], + ) + .unwrap() +} + +#[test] +fn recovery_retains_selected_operator_and_rejects_invalid_contracts() { + let bytes = sorted().encode().unwrap(); + let recovered = CompiledPhysicalDag::decode(&bytes).unwrap(); + assert_eq!(recovered.encode().unwrap(), bytes); + for mutation in ["version", "column", "edge", "output"] { + let mut wire: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + match mutation { + "version" => wire["version"] = 999.into(), + "column" => { + wire["nodes"]["1"]["Operator"]["operator"]["kind"]["Sort"]["keys"][0]["column"] = + 7.into() + } + "edge" => wire["nodes"]["1"]["Operator"]["inputs"][0] = 999.into(), + "output" => { + wire["nodes"]["1"]["Operator"]["operator"]["output"]["fields"][0]["dtype"] = + serde_json::json!({"Plain":"utf8"}) + } + _ => unreachable!(), + } + assert!( + CompiledPhysicalDag::decode(&serde_json::to_vec(&wire).unwrap()).is_err(), + "accepted {mutation}" + ); + } +} + +#[test] +fn candidate_recovery_preserves_materialization_boundary() { + use asap_physical_operators::physical_planner::PhysicalCandidate; + let precompute = sorted(); + let output = InputContract::bounded(precompute.output_contract(1).unwrap().schema); + let query = CompiledPhysicalDag::from_operators( + BTreeMap::from([(1, output.clone())]), + BTreeMap::from([( + 2, + ( + vec![1], + Operator::limit(output.schema.clone(), 3, 0, vec![]).unwrap(), + ), + )]), + vec![2], + ) + .unwrap(); + let candidate = PhysicalCandidate { + precompute: Some(precompute), + query, + materialized_outputs: BTreeMap::from([(1, output)]), + }; + let bytes = candidate.encode().unwrap(); + let restored = PhysicalCandidate::decode(&bytes).unwrap(); + assert_eq!(restored.precompute.as_ref().unwrap().roots(), &[1]); + assert_eq!(restored.query.roots(), &[2]); + assert_eq!(restored.encode().unwrap(), bytes); + let mut wire: serde_json::Value = serde_json::from_slice(&bytes).unwrap(); + wire["materialized_outputs"]["1"]["schema"]["fields"][0]["dtype"] = + serde_json::json!({"Plain":"utf8"}); + assert!(PhysicalCandidate::decode(&serde_json::to_vec(&wire).unwrap()).is_err()); +} diff --git a/crates/asap-physical-operators/tests/physical_semantics.rs b/crates/asap-physical-operators/tests/physical_semantics.rs new file mode 100644 index 00000000..7270bd1e --- /dev/null +++ b/crates/asap-physical-operators/tests/physical_semantics.rs @@ -0,0 +1,685 @@ +//! Contract tests inspired by DataFusion's limit, sort and join test matrices. +//! Expectations follow ASAP's IR (notably row-count and IEEE NaN equality). +//! Reference: apache/datafusion e2ca7f3, physical-plan/src/{limit.rs,sorts/sort.rs}. +use asap_physical_operators::{ + expressions::CompiledExpression, + operators::{Expression, Operator, Reduction, SortKey}, + plan::PhysicalDag, + runtime::{Limits, RunContext, Scope}, + values::{Batch, Schema, Value}, +}; +use futures::{executor::block_on, StreamExt}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::{CompareOpKind, DataType, JoinKind, Predicate, QueryExpr}, +}; +use std::{rc::Rc, 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 context() -> RunContext { + RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 1, + }, + Limits { + max_buffered_batches: 1, + ..Limits::default() + }, + ) + .unwrap() +} +fn collect(dag: &PhysicalDag<'_, Batch, Schema>, root: u64) -> Vec> { + let run = context(); + let rows = block_on(async { + let mut stream = dag.execute(&[root], run.clone()).unwrap().remove(0); + let mut rows = vec![]; + while let Some(batch) = stream.next().await { + rows.extend_from_slice(batch.unwrap().rows()); + } + rows + }); + assert_eq!(run.retained_bytes(), 0); + rows +} +fn unary(input: Schema, batches: Vec>>, op: Operator) -> Vec> { + let mut dag = PhysicalDag::default(); + let batches = batches + .into_iter() + .map(|rows| Batch::try_new(input.clone(), rows).unwrap()) + .collect(); + dag.add(0, vec![], Operator::source(input, batches).unwrap()) + .unwrap(); + dag.add(1, vec![0], op).unwrap(); + collect(&dag, 1) +} +fn keys(rows: &[Vec]) -> Vec>> { + rows.iter() + .map(|r| r.iter().map(|v| v.key().unwrap()).collect()) + .collect() +} +fn eq_predicate() -> Predicate { + Predicate(Rc::new(QueryExpr::Compare { + left: Rc::new(QueryExpr::Column(0)), + op: CompareOpKind::Eq, + right: Rc::new(QueryExpr::Column(1)), + })) +} +fn join(left: Vec, right: Vec, kind: JoinKind, keyed: bool) -> Vec> { + let input = schema(&[("key", DataType::Float64, true)]); + let output = if matches!(kind, JoinKind::Semi | JoinKind::Anti) { + input.clone() + } else { + schema(&[ + ("left", DataType::Float64, true), + ("right", DataType::Float64, true), + ]) + }; + let op = if keyed { + Operator::semi_join(input.clone(), input.clone(), vec![(0, 0)]).unwrap() + } else { + Operator::relational_join(input.clone(), input.clone(), kind, &eq_predicate(), output) + .unwrap() + }; + let mut dag = PhysicalDag::default(); + for (id, values) in [(0, left), (1, right)] { + let batches = values + .into_iter() + .map(|v| Batch::try_new(input.clone(), vec![vec![v]]).unwrap()) + .collect(); + dag.add( + id, + vec![], + Operator::source(input.clone(), batches).unwrap(), + ) + .unwrap(); + } + dag.add(2, vec![0, 1], op).unwrap(); + collect(&dag, 2) +} + +// OFFSET/FETCH must be invariant to empty batches and input batch boundaries. +#[test] +fn limit_offset_fetch_matrix() { + let input = schema(&[("v", DataType::Int64, false)]); + for chunk in [1, 2, 5, 12] { + let values = (0..9).map(|n| vec![Value::Int64(n)]).collect::>(); + let mut batches = vec![vec![]]; + for rows in values.chunks(chunk) { + batches.push(rows.to_vec()); + batches.push(vec![]); + } + for offset in [0, 1, 8, 9, 10, u64::MAX] { + for n in [0, 1, 3, 12, u64::MAX] { + let rows = unary( + input.clone(), + batches.clone(), + Operator::limit(input.clone(), n, offset, vec![]).unwrap(), + ); + let expected = values + .iter() + .skip(offset.min(9) as usize) + .take(n.min(9) as usize) + .cloned() + .collect::>(); + assert_eq!( + keys(&rows), + keys(&expected), + "chunk={chunk}, offset={offset}, n={n}" + ); + } + } + } +} + +// Zero-column batches still have rows: LIMIT must not infer cardinality from columns. +#[test] +fn limit_preserves_zero_column_row_count() { + let input = schema(&[]); + let rows = unary( + input.clone(), + vec![vec![vec![]; 5], vec![vec![]; 5]], + Operator::limit(input, 4, 3, vec![]).unwrap(), + ); + assert_eq!(rows.len(), 4); +} + +// NULL placement is independent of sort direction; ties retain original row order. +#[test] +fn sort_direction_null_placement_and_ties() { + let input = schema(&[("v", DataType::Int64, true), ("id", DataType::Int64, false)]); + let values = [Some(2), None, Some(1), Some(2), None]; + let rows = values + .iter() + .enumerate() + .map(|(i, v)| { + vec![ + v.map(Value::Int64).unwrap_or(Value::Null), + Value::Int64(i as i64), + ] + }) + .collect::>(); + for (descending, nulls_first, expected) in [ + (false, false, vec![2, 0, 3, 1, 4]), + (false, true, vec![1, 4, 2, 0, 3]), + (true, false, vec![0, 3, 2, 1, 4]), + (true, true, vec![1, 4, 0, 3, 2]), + ] { + let op = Operator::sort( + input.clone(), + vec![SortKey { + column: 0, + descending, + nulls_first, + }], + vec![], + ) + .unwrap(); + let result = unary( + input.clone(), + vec![rows[..2].to_vec(), vec![], rows[2..].to_vec()], + op, + ); + let ids = result + .iter() + .map(|r| match r[1] { + Value::Int64(n) => n, + _ => unreachable!(), + }) + .collect::>(); + assert_eq!(ids, expected); + } +} + +// Outer joins preserve unmatched NULLs, while semi/anti joins preserve left multiplicity. +#[test] +fn joins_nulls_duplicates_and_empty_sides() { + for (kind, expected_len) in [ + (JoinKind::Inner, 4), + (JoinKind::Left, 6), + (JoinKind::Right, 6), + (JoinKind::Full, 8), + (JoinKind::Semi, 2), + (JoinKind::Anti, 2), + ] { + let left = vec![ + Value::Float64(1.), + Value::Float64(1.), + Value::Float64(2.), + Value::Null, + ]; + let right = vec![ + Value::Float64(1.), + Value::Float64(1.), + Value::Float64(3.), + Value::Null, + ]; + let result = join(left, right, kind.clone(), false); + assert_eq!(result.len(), expected_len, "{kind:?}"); + } + for (kind, expected_len) in [ + (JoinKind::Inner, 0), + (JoinKind::Left, 1), + (JoinKind::Right, 0), + (JoinKind::Full, 1), + (JoinKind::Semi, 0), + (JoinKind::Anti, 1), + ] { + assert_eq!( + join(vec![Value::Float64(7.)], vec![], kind.clone(), false).len(), + expected_len, + "{kind:?}" + ); + } + let result = join(vec![Value::Float64(7.)], vec![], JoinKind::Left, false); + assert!(matches!( + result[0].as_slice(), + [Value::Float64(7.), Value::Null] + )); +} + +// Changing the semi-join algorithm must not turn IEEE NaN != NaN into a match. +#[test] +fn keyed_semijoin_obeys_ieee_equality_for_nan_and_zero() { + let left = vec![ + Value::Float64(f64::NAN), + Value::Float64(-0.), + Value::Float64(0.), + Value::Null, + ]; + let right = vec![Value::Float64(f64::NAN), Value::Float64(0.), Value::Null]; + let keyed = join(left, right, JoinKind::Semi, true); + let expected = vec![vec![Value::Float64(-0.)], vec![Value::Float64(0.)]]; + assert_eq!(keys(&keyed), keys(&expected)); +} + +// Group equality intentionally differs from predicate equality: NULL and NaNs group together. +#[test] +fn grouping_canonicalizes_null_nan_and_signed_zero() { + let input = schema(&[("v", DataType::Float64, true)]); + let op = Operator::aggregate( + input.clone(), + vec![0], + vec![("count".into(), Reduction::Count)], + ) + .unwrap(); + let values = vec![ + Value::Null, + Value::Null, + Value::Float64(0.), + Value::Float64(-0.), + Value::Float64(f64::NAN), + Value::Float64(f64::from_bits(0x7ff8000000000001)), + ]; + let result = unary( + input, + values.into_iter().map(|v| vec![vec![v]]).collect(), + op, + ); + assert_eq!(result.len(), 3); + assert!(result.iter().all(|r| matches!(r[1], Value::Int64(2)))); +} + +// Global empty input yields one aggregate row; grouped empty input yields none. +#[test] +fn aggregate_empty_and_all_null_follow_asap_contract() { + let input = schema(&[("v", DataType::Int64, true)]); + for batches in [ + vec![], + vec![vec![]], + vec![vec![vec![Value::Null], vec![Value::Null]]], + ] { + let n = batches.iter().map(Vec::len).sum::(); + let op = Operator::aggregate( + input.clone(), + vec![], + vec![ + ("count".into(), Reduction::Count), + ("min".into(), Reduction::Min(0)), + ("max".into(), Reduction::Max(0)), + ], + ) + .unwrap(); + let result = unary(input.clone(), batches, op); + assert_eq!(result.len(), 1); + assert!(matches!(result[0][0], Value::Int64(v) if v == n as i64)); + assert!(matches!(result[0][1], Value::Null)); + assert!(matches!(result[0][2], Value::Null)); + } + let op = Operator::aggregate( + input.clone(), + vec![0], + vec![("count".into(), Reduction::Count)], + ) + .unwrap(); + assert!(unary(input, vec![], op).is_empty()); +} + +// A precompiled expression with a different input contract must fail during binding. +#[test] +fn projection_rejects_expression_bound_to_another_schema() { + let original = schema(&[("a", DataType::Int64, false), ("b", DataType::Int64, false)]); + let current = schema(&[("a", DataType::Int64, false)]); + let expr = CompiledExpression::compile(&QueryExpr::Column(1), &original).unwrap(); + assert!(Operator::project(current, vec![("b".into(), Expression::planner(expr))]).is_err()); +} + +// A valid Planner MIN/MAX schema must bind even for a non-null input column. +#[test] +fn global_extrema_bind_with_planner_derived_schema() { + use asap_physical_operators::physical_planner::compile_node; + use planner_types::{ + post_asap::*, + pre_asap::{AggIntent, Column, GroupKeys, Reduction as PlanReduction}, + }; + let input = schema(&[("v", DataType::Int64, false)]); + for measure in [ + AggIntent::Min { col: Some(0) }, + AggIntent::Max { col: Some(0) }, + ] { + let planner_input = + planner_types::pre_asap::Schema::new(vec![Column::new("v", DataType::Int64, false)]); + let derived = planner_types::pre_asap::query_expr::aggregate_output_schema( + &planner_input, + &PlanReduction::Reduce(GroupKeys::by(vec![])), + std::slice::from_ref(&measure), + &[], + ) + .unwrap(); + let result = derived.columns[0].clone(); + let output = schema(&[(&result.name, result.dtype, result.nullable)]); + let node = ExecutableDagNode { + id: PostAsapNodeId(1), + payload: ExecutableOperatorPayload::Value { + operation: ValueOperation::Exact(ExactOperation::Aggregate { + reduction: PlanReduction::Reduce(GroupKeys::by(vec![])), + measures: vec![measure], + output_names: vec![result.name], + having: None, + }), + }, + output_state: ExecutionDataState::QUERY_ROWS, + output_schema: (*output).clone(), + guarantee: None, + }; + let operator = compile_node(&node, std::slice::from_ref(&input)) + .expect("global extremum should bind to its Planner schema"); + assert!(operator.schema().fields[0].nullable); + let empty = unary(input.clone(), vec![], operator.clone()); + assert!(matches!(empty[0][0], Value::Null)); + let nonempty = unary(input.clone(), vec![vec![vec![Value::Int64(7)]]], operator); + assert!(matches!(nonempty[0][0], Value::Int64(7))); + } +} + +// NaN is a valid numeric input, not a schema error; all six comparisons obey IEEE rules. +#[test] +fn planner_comparisons_handle_nan_without_execution_errors() { + let input = schema(&[ + ("a", DataType::Float64, false), + ("b", DataType::Float64, false), + ]); + for op in [ + CompareOpKind::Eq, + CompareOpKind::Ne, + CompareOpKind::Lt, + CompareOpKind::Le, + CompareOpKind::Gt, + CompareOpKind::Ge, + ] { + let expression = QueryExpr::Compare { + left: Rc::new(QueryExpr::Column(0)), + op: op.clone(), + right: Rc::new(QueryExpr::Column(1)), + }; + let compiled = CompiledExpression::compile(&expression, &input).unwrap(); + for row in [ + [Value::Float64(f64::NAN), Value::Float64(1.)], + [Value::Float64(1.), Value::Float64(f64::NAN)], + [Value::Float64(f64::NAN), Value::Float64(f64::NAN)], + ] { + let actual = compiled.evaluate(&row).unwrap(); + assert!(matches!(actual,Value::Bool(value) if value == (op == CompareOpKind::Ne))); + } + } +} + +// A bounded LIMIT branch must unsubscribe so another branch can drain the producer. +#[test] +fn limit_branch_finishes_without_blocking_shared_sibling() { + let input = schema(&[("v", DataType::Int64, false)]); + let mut dag = PhysicalDag::default(); + let batches = (0..100) + .map(|v| Batch::try_new(input.clone(), vec![vec![Value::Int64(v)]]).unwrap()) + .collect(); + dag.add(0, vec![], Operator::source(input.clone(), batches).unwrap()) + .unwrap(); + dag.add( + 1, + vec![0], + Operator::limit(input.clone(), 1, 0, vec![]).unwrap(), + ) + .unwrap(); + dag.add(2, vec![0, 1], Operator::union(input, 2).unwrap()) + .unwrap(); + // Bound polls as well as rows so a backpressure regression cannot hang the suite. + use futures::{task::noop_waker_ref, Stream}; + use std::{ + pin::Pin, + task::{Context, Poll}, + }; + let run = context(); + let mut stream = dag.execute(&[2], run.clone()).unwrap().remove(0); + let mut cx = Context::from_waker(noop_waker_ref()); + let mut count = 0; + for _ in 0..2000 { + match Pin::new(&mut stream).poll_next(&mut cx) { + Poll::Ready(Some(batch)) => count += batch.unwrap().rows().len(), + Poll::Ready(None) => { + assert_eq!(count, 101); + drop(stream); + assert_eq!(run.retained_bytes(), 0); + return; + } + Poll::Pending => {} + } + } + panic!("shared LIMIT/Union failed to make progress"); +} + +// Mixed numeric comparisons must not round Int64 values through f64 before comparing. +#[test] +fn mixed_numeric_comparisons_preserve_large_integer_precision() { + let input = schema(&[ + ("a", DataType::Int64, false), + ("b", DataType::Float64, false), + ]); + let expr = QueryExpr::Compare { + left: Rc::new(QueryExpr::Column(0)), + op: CompareOpKind::Gt, + right: Rc::new(QueryExpr::Column(1)), + }; + let compiled = CompiledExpression::compile(&expr, &input).unwrap(); + for (a, b, expected) in [ + (9_007_199_254_740_993, 9_007_199_254_740_992.0, true), + (i64::MAX, 9_223_372_036_854_775_808.0, false), + (i64::MIN, f64::NEG_INFINITY, true), + ] { + assert!( + matches!(compiled.evaluate(&[Value::Int64(a),Value::Float64(b)]).unwrap(), Value::Bool(v) if v == expected) + ); + } +} + +// Both expression paths must implement all nine combinations of three-valued booleans. +#[test] +fn boolean_truth_tables_agree_between_expression_paths() { + let input = schema(&[("a", DataType::Bool, true), ("b", DataType::Bool, true)]); + for and in [true, false] { + for a in [None, Some(false), Some(true)] { + for b in [None, Some(false), Some(true)] { + let parts = vec![QueryExpr::Column(0), QueryExpr::Column(1)]; + let planner = if and { + QueryExpr::BoolAnd(parts) + } else { + QueryExpr::BoolOr(parts) + }; + let native = if and { + Expression::And( + Box::new(Expression::Column(0)), + Box::new(Expression::Column(1)), + ) + } else { + Expression::Or( + Box::new(Expression::Column(0)), + Box::new(Expression::Column(1)), + ) + }; + let expected = match (a, b, and) { + (Some(false), _, true) | (_, Some(false), true) => Some(false), + (Some(true), _, false) | (_, Some(true), false) => Some(true), + (None, _, _) | (_, None, _) => None, + (Some(a), Some(b), true) => Some(a && b), + (Some(a), Some(b), false) => Some(a || b), + } + .map(Value::Bool) + .unwrap_or(Value::Null); + let row = vec![ + a.map(Value::Bool).unwrap_or(Value::Null), + b.map(Value::Bool).unwrap_or(Value::Null), + ]; + let compiled = CompiledExpression::compile(&planner, &input).unwrap(); + assert_eq!( + compiled.evaluate(&row).unwrap().key().unwrap(), + expected.key().unwrap() + ); + let op = Operator::project(input.clone(), vec![("result".into(), native)]).unwrap(); + let result = unary(input.clone(), vec![vec![row]], op); + assert_eq!(result[0][0].key().unwrap(), expected.key().unwrap()); + } + } + } +} + +// Partial/final execution must agree with one build for an uncompacted KLL population. +#[test] +fn kll_partial_merge_and_multiple_readouts_preserve_population() { + use asap_physical_operators::Statistic; + use planner_types::post_asap::{SketchAlgorithm, SketchKind, SketchParams}; + let input = schema(&[("v", DataType::Float64, false)]); + let family = SummaryFamilyType::Sketch( + SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k: 512 }), + Default::default(), + ); + let mut dag = PhysicalDag::default(); + for (id, range) in [(0, 0..64), (1, 64..128), (2, 0..128)] { + let rows = range.map(|n| vec![Value::Float64(n as f64)]).collect(); + dag.add( + id, + vec![], + Operator::source( + input.clone(), + vec![Batch::try_new(input.clone(), rows).unwrap()], + ) + .unwrap(), + ) + .unwrap(); + dag.add( + id + 3, + vec![id], + Operator::summary_build(input.clone(), family.clone(), 0, None, vec![]).unwrap(), + ) + .unwrap(); + } + let state = Operator::summary_build(input, family, 0, None, vec![]) + .unwrap() + .schema(); + dag.add(6, vec![3, 4], Operator::union(state.clone(), 2).unwrap()) + .unwrap(); + dag.add( + 7, + vec![6], + Operator::summary_merge(state.clone(), 0, vec![]).unwrap(), + ) + .unwrap(); + let mut roots = vec![]; + for (i, q) in [0.0, 0.5, 1.0].into_iter().enumerate() { + for (j, build) in [5, 7].into_iter().enumerate() { + let id = 8 + (i * 2 + j) as u64; + dag.add( + id, + vec![build], + Operator::readout( + state.clone(), + 0, + Statistic::Quantile, + std::collections::HashMap::from([("quantile".into(), q.to_string())]), + ) + .unwrap(), + ) + .unwrap(); + roots.push(id); + } + } + for _ in 0..2 { + let run = context(); + let outputs = block_on(futures::future::join_all( + dag.execute(&roots, run.clone()) + .unwrap() + .into_iter() + .map(|s| s.collect::>()), + )); + for (pair, expected) in outputs.chunks(2).zip([0., 64., 127.]) { + let value = |batches: &[Result< + asap_physical_operators::runtime::SharedValue, + asap_physical_operators::Error, + >]| { + assert_eq!(batches.len(), 1); + match batches[0].as_ref().unwrap().rows()[0][0] { + Value::Float64(v) => v, + _ => panic!("quantile must be Float64"), + } + }; + assert_eq!(value(&pair[0]), value(&pair[1])); + assert!((value(&pair[0]) - expected).abs() <= 1.); + } + drop(outputs); + assert_eq!(run.retained_bytes(), 0); + } +} + +// Retained zero-column rows still own Vec headers and must consume the output budget. +#[test] +fn zero_column_output_obeys_memory_limit() { + use asap_physical_operators::Error; + let input = schema(&[]); + let batch = Batch::try_new(input.clone(), vec![vec![]; 200]).unwrap(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], Operator::source(input, vec![batch]).unwrap()) + .unwrap(); + let run = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits { + max_bytes: 1024, + ..Limits::default() + }, + ) + .unwrap(); + let mut stream = dag.execute(&[0], run.clone()).unwrap().remove(0); + assert!(matches!( + block_on(stream.next()), + Some(Err(Error::MemoryLimit)) + )); + drop(stream); + assert_eq!(run.retained_bytes(), 0); +} + +// Empty exact-state finalization must preserve ordinary global MIN/MAX null semantics. +#[test] +fn empty_exact_summary_extrema_agree_with_ordinary_aggregation() { + use asap_physical_operators::Statistic; + use planner_types::post_asap::{ExactKind, ExactParams}; + let input = schema(&[("v", DataType::Float64, false)]); + for (kind, params, statistic) in [ + (ExactKind::Min, ExactParams::Min, Statistic::Min), + (ExactKind::Max, ExactParams::Max, Statistic::Max), + ] { + let build = Operator::summary_build( + input.clone(), + SummaryFamilyType::ExactAggregate(kind, params), + 0, + None, + vec![], + ) + .unwrap(); + let state = build.schema(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], Operator::source(input.clone(), vec![]).unwrap()) + .unwrap(); + dag.add(1, vec![0], build).unwrap(); + dag.add( + 2, + vec![1], + Operator::readout(state, 0, statistic, Default::default()).unwrap(), + ) + .unwrap(); + let rows = collect(&dag, 2); + assert_eq!(rows.len(), 1); + assert!(matches!(rows[0][0], Value::Null)); + } +} diff --git a/crates/asap-physical-operators/tests/plan_properties.rs b/crates/asap-physical-operators/tests/plan_properties.rs new file mode 100644 index 00000000..f43ab23e --- /dev/null +++ b/crates/asap-physical-operators/tests/plan_properties.rs @@ -0,0 +1,155 @@ +//! Finite-input contracts are validated before source execution. +use asap_physical_operators::{ + operators::{Operator, SortKey}, + plan::{Boundedness, Emission, PhysicalDag}, + runtime::{Limits, OutputStream, RunContext, Scope}, + sources::{DataSources, RawSource}, + values::{Batch, Schema}, + Error, +}; +use planner_types::{ + post_asap::{SummaryFamilyType, SummaryField, SummarySchema}, + pre_asap::{Column, DataType, QueryExpr, Schema as LogicalSchema, Source}, +}; +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; +struct DeclaredSource { + schema: Schema, + boundedness: Boundedness, + opens: Arc, +} +impl RawSource for DeclaredSource { + fn schema(&self) -> Schema { + self.schema.clone() + } + fn boundedness(&self) -> Boundedness { + self.boundedness + } + fn scan(&self, _: RunContext) -> Result, Error> { + self.opens.fetch_add(1, Ordering::SeqCst); + Ok(Box::pin(futures::stream::empty())) + } +} +// A blocking parent must reject unknown and unbounded Scan inputs without opening a reader. +#[test] +fn blocking_inputs_require_an_explicit_finite_source() { + let schema = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "v".into(), + dtype: SummaryFamilyType::Plain(DataType::Int64), + nullable: false, + }], + time_index: None, + }); + for boundedness in [ + Boundedness::Unknown, + Boundedness::Unbounded, + Boundedness::Bounded, + ] { + let opens = Arc::new(AtomicUsize::new(0)); + let mut registry = DataSources::default(); + let identity = Source::Table { + table_ref: "t".into(), + }; + registry + .register( + identity.clone(), + Arc::new(DeclaredSource { + schema: schema.clone(), + boundedness, + opens: opens.clone(), + }), + ) + .unwrap(); + let scan = registry + .bind(&QueryExpr::Scan { + source: identity, + schema: LogicalSchema::new(vec![Column::new("v", DataType::Int64, false)]), + predicates: vec![], + }) + .unwrap(); + let mut dag = PhysicalDag::default(); + dag.add(0, vec![], scan).unwrap(); + dag.add( + 1, + vec![0], + Operator::sort( + schema.clone(), + vec![SortKey { + column: 0, + descending: false, + nulls_first: false, + }], + vec![], + ) + .unwrap(), + ) + .unwrap(); + let run = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + if boundedness == Boundedness::Bounded { + let properties = dag.properties(&[1]).unwrap(); + assert_eq!(properties[&1].emission, Emission::AfterInput); + assert_eq!(properties[&1].boundedness, Boundedness::Bounded); + assert!(dag.execute(&[1], run).is_ok()); + } else { + assert!( + matches!(dag.execute(&[1], run), Err(Error::Invalid(message)) if message.contains("requires bounded inputs")) + ); + } + assert_eq!(opens.load(Ordering::SeqCst), 0); + } +} + +// Kernel support must not be mistaken for executable native state/readout support. +#[test] +fn summary_capability_levels_are_distinct() { + use asap_physical_operators::{ + capability::{validate_native_family, validate_native_readout, validate_summary_kernel}, + Statistic, + }; + use planner_types::{ + post_asap::{GroupingStrategy, SketchAlgorithm, SketchKind, SketchParams, SummaryUpdate}, + pre_asap::ColumnRef, + }; + let grouping = GroupingStrategy::default(); + let cms = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::Cms, + SketchParams::Cms { + width: 64, + depth: 4, + }, + ), + grouping.clone(), + ); + let update = SummaryUpdate { + item: Some(planner_types::post_asap::SummaryInputExpr::Column( + ColumnRef::Named("host".into()), + )), + weight: planner_types::post_asap::SummaryInputExpr::Constant(1.0), + weight_domain: Default::default(), + }; + assert!(validate_summary_kernel(&cms, &update, &grouping).is_ok()); + assert!(validate_native_family(&cms).is_err()); + let kll = SummaryFamilyType::Sketch( + SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k: 128 }), + grouping, + ); + assert!(validate_native_family(&kll).is_ok()); + assert!(validate_native_readout(&kll, Statistic::Quantile, &Default::default()).is_err()); + assert!(validate_native_readout( + &kll, + Statistic::Quantile, + &[("quantile".into(), "0.5".into())].into() + ) + .is_ok()); +} diff --git a/crates/asap-physical-operators/tests/precompute_candidates.rs b/crates/asap-physical-operators/tests/precompute_candidates.rs new file mode 100644 index 00000000..2acb39a0 --- /dev/null +++ b/crates/asap-physical-operators/tests/precompute_candidates.rs @@ -0,0 +1,560 @@ +//! Materialized frontiers are compiled by Planner, never rewritten by deployment. +use asap_aware_mapping::{cost_model::DefaultCostModel, search_workload}; +use asap_physical_operators::{ + factory::create_planner_accumulator, + operators::Operator, + physical_planner::{ + compile_candidates, select_candidate, CandidateCost, CompiledPhysicalDag, InputContract, + Source, + }, + runtime::{Limits, RunContext, Scope}, + values::{Batch, Value}, +}; +use futures::{executor::block_on, StreamExt}; +use planner_types::{post_asap::*, pre_asap::DataType, types::AccuracyTarget, workload::*}; +use std::{collections::BTreeMap, rc::Rc, sync::Arc}; + +fn grouped_rate_space() -> asap_aware_mapping::PlanSpace<&'static str> { + let workload = PlanningWorkload { + query_workload: QueryWorkload { + language: QueryLanguage::PromQL, + query_batch: Some(vec![BatchEntry { + query: Query("sum by(job)(rate(m[1m]))".into()), + requirements: QueryRequirements { + accuracy: AccuracyRequirement::Explicit(AccuracyTarget::Exact), + ..Default::default() + }, + predictability: Predictability::Unknown, + invocations: 1, + execute_at: None, + time_selection: TimeSelection::default(), + }]), + repeating_queries: None, + }, + data_workload: Some(DataWorkload { + data_ingestion_interval: Evidence { + value: Some(DurationMs(1000)), + ..Default::default() + }, + ..Default::default() + }), + }; + let root = Rc::new( + asap_frontend_promql::lower_promql_workload(&workload, 0) + .unwrap() + .remove(0), + ); + let root = Rc::new( + asap_physical_operators::physical_planner::promql_rows::with_series_identity(&root) + .unwrap(), + ); + search_workload(vec![("grouped-rate", root)]) +} + +fn grouped_rate() -> ExecutableDag { + let space = grouped_rate_space(); + let selected = space + .global_selection(&DefaultCostModel) + .assemble_selected_dag(&space.roots[0].1) + .unwrap() + .unwrap(); + compile_executable_dag(&selected).unwrap() +} +fn run(plan: &CompiledPhysicalDag, inputs: BTreeMap, scope: Scope) -> Vec { + let sources = inputs + .into_iter() + .map(|(id, batch)| { + let source = Operator::source(batch.schema().clone(), vec![batch]).unwrap(); + (id, Box::new(source) as Source<'_>) + }) + .collect(); + let dag = plan.instantiate(sources).unwrap(); + block_on(async { + let context = RunContext::new(scope, Limits::default()).unwrap(); + let mut output = dag.execute(plan.roots(), context).unwrap().remove(0); + let mut batches = vec![]; + while let Some(batch) = output.next().await { + batches.push((*batch.unwrap()).clone()); + } + batches + }) +} + +/// Rate readouts and grouped Sum can run together during bounded precompute; +/// storing per-series rates instead leaves the same Sum in the query DAG. +#[test] +fn grouped_rate_can_be_materialized_before_or_after_grouped_sum() { + let dag = grouped_rate(); + let state = dag + .nodes + .iter() + .find(|node| { + matches!( + node.payload, + ExecutableOperatorPayload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + .. + } + ) + }) + .unwrap(); + let readout = dag + .nodes + .iter() + .find(|node| { + matches!( + node.payload, + ExecutableOperatorPayload::Value { + operation: ValueOperation::FinalizeExactAccumulator + } + ) + }) + .unwrap(); + let input_schema = Arc::new(state.output_schema.clone()); + let (family, update, grouping) = match &state.payload { + ExecutableOperatorPayload::SummaryAgg { + family, + input, + grouping, + .. + } => (family, input, grouping), + _ => unreachable!(), + }; + let range_parameters = std::collections::HashMap::from([ + ("range_start_ms".into(), "-58000".into()), + ("range_end_ms".into(), "2000".into()), + ]); + let mut expected_rate_sum = 0.; + let rows = [[100., 0., 100.], [100., 200., 0.]] + .into_iter() + .enumerate() + .map(|(index, values)| { + let mut accumulator = create_planner_accumulator(family, update, grouping).unwrap(); + for (i, value) in values.into_iter().enumerate() { + accumulator.update_single(value, i as i64 * 1000); + } + let state = accumulator.into_accumulator(); + expected_rate_sum += state + .query_statistic( + asap_physical_operators::Statistic::Rate, + &None, + &range_parameters, + ) + .unwrap(); + let summary = Value::Summary { + family: family.clone(), + state: Arc::from(state), + }; + input_schema + .fields + .iter() + .map(|field| match &field.dtype { + SummaryFamilyType::ExactAggregate(..) => summary.clone(), + SummaryFamilyType::Plain(DataType::Timestamp) => Value::Timestamp(2000), + SummaryFamilyType::Plain(DataType::Utf8) => { + Value::Utf8(if field.name == "job" { + "api".into() + } else { + format!("series-{index}").into() + }) + } + _ => panic!("unexpected input field {field:?}"), + }) + .collect() + }) + .collect(); + let batch = Batch::try_new(input_schema.clone(), rows).unwrap(); + let root = u64::from(dag.root.0); + let state_id = u64::from(state.id.0); + let rate_id = u64::from(readout.id.0); + let frontiers = asap_physical_operators::physical_planner::enumerate_frontiers( + &dag, + &BTreeMap::from([(state_id, InputContract::bounded(input_schema.clone()))]), + &[root], + 128, + ) + .unwrap(); + assert!(frontiers.contains(&vec![])); + assert!(frontiers.contains(&vec![rate_id])); + assert!(frontiers.contains(&vec![root])); + assert!(!frontiers.contains(&vec![root, rate_id])); + assert!( + asap_physical_operators::physical_planner::enumerate_frontiers( + &dag, + &BTreeMap::from([(state_id, InputContract::bounded(input_schema.clone()))]), + &[root], + 1, + ) + .is_err() + ); + let candidates = compile_candidates( + &dag, + BTreeMap::from([(state_id, InputContract::bounded(input_schema))]), + &[root], + &[vec![], vec![rate_id], vec![root]], + ); + // Scoped cost fixtures select either precompute boundary. No readers are + // opened during candidate construction or selection. + for prefer_grouped in [false, true] { + let inventory = compile_candidates( + &dag, + BTreeMap::from([( + state_id, + InputContract::bounded(Arc::new(state.output_schema.clone())), + )]), + &[root], + &[vec![999], vec![rate_id], vec![root]], + ); + assert!(inventory[0].is_err()); + let mut evaluated = 0; + let selected = select_candidate(inventory, |candidate| { + evaluated += 1; + let grouped = candidate.materialized_outputs.contains_key(&root); + Ok(Some(CandidateCost { + workload_scope: "reset-counter-workload".into(), + horizon_seconds: 300., + total_cost: if grouped == prefer_grouped { 1. } else { 100. }, + })) + }) + .unwrap(); + assert_eq!( + selected.candidate.materialized_outputs.contains_key(&root), + prefer_grouped + ); + assert_eq!(selected.cost.total_cost, 1.); + assert_eq!(evaluated, 2, "uncompilable candidates must never be priced"); + let candidate = selected.candidate; + let precompute = candidate.precompute.as_ref().unwrap(); + let stored = run( + precompute, + BTreeMap::from([(state_id, batch.clone())]), + Scope::Ingestion { + window_start_ms: -58_000, + window_end_ms: 2000, + revision: 1, + }, + ); + let output = run( + &candidate.query, + BTreeMap::from([(precompute.roots()[0], stored[0].clone())]), + Scope::Query { + evaluation_time_ms: 2000, + revision: 1, + }, + ); + assert!( + matches!(output[0].rows()[0][1], Value::Float64(value) if value == expected_rate_sum) + ); + } + let contracts = BTreeMap::from([( + state_id, + InputContract::bounded(Arc::new(state.output_schema.clone())), + )]); + for frontier in [vec![rate_id, rate_id], vec![root, rate_id], vec![999]] { + assert!( + asap_physical_operators::physical_planner::compile_candidate( + &dag, + contracts.clone(), + &[root], + &frontier + ) + .is_err() + ); + } + let inventory = compile_candidates( + &dag, + contracts.clone(), + &[root], + &[vec![rate_id], vec![root]], + ); + let selected = select_candidate(inventory, |candidate| { + if candidate.materialized_outputs.contains_key(&root) { + return Ok(None); + } + Ok(Some(CandidateCost { + workload_scope: "same-workload".into(), + horizon_seconds: 300., + total_cost: 100., + })) + }) + .unwrap(); + assert!(selected + .candidate + .materialized_outputs + .contains_key(&rate_id)); + let inventory = compile_candidates(&dag, contracts, &[root], &[vec![rate_id], vec![root]]); + assert!( + select_candidate(inventory, |candidate| Ok(Some(CandidateCost { + workload_scope: "same-workload".into(), + horizon_seconds: if candidate.materialized_outputs.contains_key(&root) { + 60. + } else { + 300. + }, + total_cost: 1., + }))) + .is_err() + ); + let query_scope = Scope::Query { + evaluation_time_ms: 2000, + revision: 1, + }; + let maintenance_scope = Scope::Ingestion { + window_start_ms: -58_000, + window_end_ms: 2000, + revision: 1, + }; + let mut results = vec![]; + for candidate in candidates { + let candidate = candidate.unwrap(); + let inputs = if let Some(precompute) = &candidate.precompute { + let source = Operator::source(batch.schema().clone(), vec![batch.clone()]).unwrap(); + let invalid = precompute + .instantiate(BTreeMap::from([(state_id, Box::new(source) as Source<'_>)])) + .unwrap(); + let context = RunContext::new( + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 2000, + revision: 1, + }, + Limits::default(), + ) + .unwrap(); + assert!(invalid.execute(precompute.roots(), context).is_err()); + let stored = run( + precompute, + BTreeMap::from([(state_id, batch.clone())]), + maintenance_scope.clone(), + ); + assert_eq!(stored.len(), 1); + let boundary = precompute.roots()[0]; + assert_eq!( + candidate.materialized_outputs[&boundary].schema, + *stored[0].schema() + ); + BTreeMap::from([(boundary, stored[0].clone())]) + } else { + BTreeMap::from([(state_id, batch.clone())]) + }; + let output = run(&candidate.query, inputs, query_scope.clone()); + assert_eq!(output.len(), 1); + assert_eq!(output[0].rows().len(), 1); + assert!(matches!(&output[0].rows()[0][0], Value::Utf8(job) if job.as_ref() == "api")); + assert!( + matches!(output[0].rows()[0][1], Value::Float64(value) if value == expected_rate_sum) + ); + results.push( + output[0].rows()[0] + .iter() + .map(|value| value.key().unwrap()) + .collect::>(), + ); + } + assert_eq!(results[0], results[1]); + assert_eq!(results[1], results[2]); + let mut wrong_order = create_planner_accumulator(family, update, grouping).unwrap(); + for (i, value) in [200., 200., 100.].into_iter().enumerate() { + wrong_order.update_single(value, i as i64 * 1000); + } + let rate_of_sum = wrong_order + .into_accumulator() + .query_statistic( + asap_physical_operators::Statistic::Rate, + &None, + &range_parameters, + ) + .unwrap(); + assert_ne!( + expected_rate_sum, rate_of_sum, + "counter resets prohibit moving Sum before Rate" + ); +} + +/// Enumerated frontiers include both grouped-result and per-series readout +/// persistence; an explicit Rate-state input retains its original semantics. +#[test] +fn bounded_inventory_exposes_grouped_rate_physical_frontiers() { + use asap_physical_operators::physical_planner::enumerate_frontiers; + let dag = grouped_rate(); + let state = dag + .nodes + .iter() + .find(|node| { + matches!( + &node.payload, + ExecutableOperatorPayload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + .. + } + ) + }) + .unwrap(); + let inputs = BTreeMap::from([( + u64::from(state.id.0), + InputContract::bounded(Arc::new(state.output_schema.clone())), + )]); + let roots = [u64::from(dag.root.0)]; + let frontiers = enumerate_frontiers(&dag, &inputs, &roots, 4096).unwrap(); + let candidates = compile_candidates(&dag, inputs.clone(), &roots, &frontiers) + .into_iter() + .collect::, _>>() + .unwrap(); + assert!(candidates.iter().any(|c| c.precompute.is_none())); + assert!(candidates + .iter() + .any(|c| c.materialized_outputs.contains_key(&roots[0]))); + assert!(candidates + .iter() + .any(|c| !c.materialized_outputs.is_empty() + && !c.materialized_outputs.contains_key(&roots[0]))); + assert!(enumerate_frontiers(&dag, &inputs, &roots, 1).is_err()); +} + +#[test] +fn enumerated_grouped_rate_candidates_execute_numeric_query_outputs() { + let inventory = grouped_rate_space().enumerate_candidate_dags(4096).unwrap(); + let mut executed = 0; + for forest in inventory.candidates { + let root = &forest[0].1; + let dag = compile_executable_dag(root).unwrap(); + let Some(state) = dag.nodes.iter().find(|node| { + matches!( + node.payload, + ExecutableOperatorPayload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + .. + } + ) + }) else { + continue; + }; + let boundary = dag + .nodes + .iter() + .find(|node| { + matches!( + node.payload, + ExecutableOperatorPayload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Sum, _), + .. + } + ) + }) + .map(|node| u64::from(node.id.0)) + .unwrap_or(u64::from(dag.root.0)); + let physical_candidates = compile_candidates( + &dag, + BTreeMap::from([( + u64::from(state.id.0), + InputContract::bounded(Arc::new(state.output_schema.clone())), + )]), + &[u64::from(dag.root.0)], + &[vec![], vec![boundary]], + ); + let (family, input, grouping) = match &state.payload { + ExecutableOperatorPayload::SummaryAgg { + family, + input, + grouping, + .. + } => (family, input, grouping), + _ => unreachable!(), + }; + let schema = Arc::new(state.output_schema.clone()); + let rows = ["a", "b"] + .into_iter() + .map(|instance| { + let mut accumulator = create_planner_accumulator(family, input, grouping).unwrap(); + for (timestamp, value) in [(1_000, 1.), (31_000, 31.), (59_000, 59.)] { + accumulator.update_single(value, timestamp); + } + let summary = Value::Summary { + family: family.clone(), + state: Arc::from(accumulator.into_accumulator()), + }; + schema + .fields + .iter() + .map(|field| match &field.dtype { + SummaryFamilyType::ExactAggregate(..) => summary.clone(), + SummaryFamilyType::Plain(DataType::Timestamp) => Value::Timestamp(60_000), + SummaryFamilyType::Plain(DataType::Utf8) + if field.name == "$promql_series_identity" => + { + Value::Utf8( + serde_json::to_string(&BTreeMap::from([ + ("job", "api"), + ("instance", instance), + ])) + .unwrap() + .into(), + ) + } + SummaryFamilyType::Plain(DataType::Utf8) => Value::Utf8("api".into()), + _ => panic!("unexpected input field {field:?}"), + }) + .collect() + }) + .collect(); + let batch = Batch::try_new(schema, rows).unwrap(); + for physical in physical_candidates { + let physical = physical.unwrap(); + let inputs = if let Some(precompute) = &physical.precompute { + let source_id = precompute.input_contracts().next().unwrap().0; + let stored = run( + precompute, + BTreeMap::from([(source_id, batch.clone())]), + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 60_000, + revision: 1, + }, + ); + assert_eq!(stored.len(), 1); + // Persist/recover the actual materialization boundary before reading it. + let bytes = asap_physical_operators::stored_state::native::encode_batch(&stored[0]) + .unwrap(); + let recovered = asap_physical_operators::stored_state::native::decode_batch( + &bytes, + stored[0].schema().clone(), + 1 << 20, + ) + .unwrap(); + BTreeMap::from([(precompute.roots()[0], recovered)]) + } else { + BTreeMap::from([( + physical.query.input_contracts().next().unwrap().0, + batch.clone(), + )]) + }; + let output = run( + &physical.query, + inputs, + Scope::Query { + evaluation_time_ms: 60_000, + revision: 1, + }, + ); + assert_eq!(output.len(), 1); + assert_eq!(output[0].rows().len(), 1); + assert!(output[0] + .schema() + .fields + .iter() + .all(|field| matches!(field.dtype, SummaryFamilyType::Plain(_)))); + assert!( + output[0].rows()[0] + .iter() + .any(|value| matches!(value, Value::Float64(x) if (*x - 2.).abs() < 1e-12)), + "{:?}", + output[0].rows() + ); + executed += 1; + } + } + assert!( + executed >= 2, + "must execute both stored and query-time grouped Rate candidates: {executed}" + ); +} diff --git a/crates/asap-physical-operators/tests/raw_scan.rs b/crates/asap-physical-operators/tests/raw_scan.rs new file mode 100644 index 00000000..25851c4f --- /dev/null +++ b/crates/asap-physical-operators/tests/raw_scan.rs @@ -0,0 +1,386 @@ +//! Scan acceptance uses the public connector contract and Planner physical DAGs. +use asap_physical_operators::dag::{ + planner::bind_with_data_sources, + scan::{DataSources, MemorySource, RawSource}, + values::{Batch, Schema, Value}, + Error, Limits, OutputStream, RunContext, Scope, +}; +use futures::{executor::block_on, stream, StreamExt}; +use planner_types::{ + post_asap::*, + pre_asap::{Column, DataType, GroupKeys, Predicate, QueryExpr, Source}, +}; +use std::{ + collections::BTreeMap, + rc::Rc, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, +}; + +fn fixture() -> (QueryExpr, Schema, Vec) { + let schema = + planner_types::pre_asap::Schema::new(vec![Column::new("value", DataType::Int64, true)]); + let output = Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Int64), + nullable: true, + }], + time_index: None, + }); + let scan = QueryExpr::Scan { + source: Source::Table { + table_ref: "numbers".into(), + }, + predicates: vec![Predicate(Rc::new(QueryExpr::IsNotNull(Rc::new( + QueryExpr::Column(0), + ))))], + schema, + }; + let batches = vec![ + Batch::try_new( + output.clone(), + vec![vec![Value::Int64(3)], vec![Value::Null]], + ) + .unwrap(), + Batch::try_new( + output.clone(), + vec![vec![Value::Int64(9)], vec![Value::Int64(2)]], + ) + .unwrap(), + ]; + (scan, output, batches) +} +fn plan(scan: QueryExpr, schema: &Schema, state: ExecutionDataState) -> ExecutableDag { + let node = |id, payload| ExecutableDagNode { + id: PostAsapNodeId(id), + payload, + output_state: state, + output_schema: (**schema).clone(), + guarantee: None, + }; + let edge = |producer, consumer| ExecutableDagEdge { + producer: PostAsapNodeId(producer), + consumer: PostAsapNodeId(consumer), + role: EdgeRole::Input, + intermediate_schema: (**schema).clone(), + data_state: state, + grouping: GroupingEdgeCompatibility::NotApplicable, + window: WindowEdgeCompatibility::NotApplicable, + }; + ExecutableDag { + nodes: vec![ + node(0, ExecutableOperatorPayload::Fallback { expression: scan }), + node( + 1, + ExecutableOperatorPayload::Value { + operation: ValueOperation::Sort { + keys: vec![planner_types::pre_asap::SortKey { + expr: QueryExpr::Column(0), + ascending: false, + nulls_first: false, + }], + partition_by: GroupKeys::by(vec![]), + }, + }, + ), + node( + 2, + ExecutableOperatorPayload::Value { + operation: ValueOperation::Limit { + n: 2, + offset: 0, + partition_by: GroupKeys::by(vec![]), + }, + }, + ), + ], + edges: vec![edge(0, 1), edge(1, 2)], + root: PostAsapNodeId(2), + } +} +fn registry(source: Arc) -> DataSources { + let mut r = DataSources::default(); + r.register( + Source::Table { + table_ref: "numbers".into(), + }, + source, + ) + .unwrap(); + r +} +fn context() -> RunContext { + RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 1, + }, + Limits::default(), + ) + .unwrap() +} + +// Raw-only execution filters nulls and ranks across batches at either phase. +#[test] +fn raw_scan_to_sort_limit_at_both_phases() { + let (scan, schema, batches) = fixture(); + let sources = registry(Arc::new( + MemorySource::new(schema.clone(), batches).unwrap(), + )); + for state in [ + ExecutionDataState::QUERY_ROWS, + ExecutionDataState::INGESTION_ROWS, + ] { + let dag = plan(scan.clone(), &schema, state); + let bound = bind_with_data_sources(&dag, BTreeMap::new(), &[2], &sources).unwrap(); + let ctx = if state == ExecutionDataState::QUERY_ROWS { + context() + } else { + RunContext::new( + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 1, + revision: 1, + }, + Limits::default(), + ) + .unwrap() + }; + let rows = block_on(async { + let mut output = bound.execute(&[2], ctx.clone()).unwrap().remove(0); + let mut rows = vec![]; + while let Some(batch) = output.next().await { + rows.extend(batch.unwrap().rows().iter().cloned()); + } + rows + }); + assert!( + matches!(rows.as_slice(), [a,b] if matches!(a.as_slice(), [Value::Int64(9)]) && matches!(b.as_slice(), [Value::Int64(3)])) + ); + assert_eq!(ctx.retained_bytes(), 0); + } +} +struct CountingSource { + schema: Schema, + opened: Arc, + fail: bool, +} +impl RawSource for CountingSource { + fn boundedness(&self) -> asap_physical_operators::plan::Boundedness { + asap_physical_operators::plan::Boundedness::Bounded + } + fn schema(&self) -> Schema { + self.schema.clone() + } + fn scan(&self, _: RunContext) -> Result, Error> { + self.opened.fetch_add(1, Ordering::SeqCst); + if self.fail { + return Err(Error::Operator("reader failed".into())); + } + Ok(stream::iter(vec![Batch::try_new( + self.schema.clone(), + vec![vec![Value::Int64(7)]], + )]) + .boxed_local()) + } +} +// Binding and cancellation do not perform I/O; fan-out opens one cursor per run. +#[test] +fn lazy_open_shared_producer_and_cancellation() { + let (scan, schema, _) = fixture(); + let opened = Arc::new(AtomicUsize::new(0)); + let sources = registry(Arc::new(CountingSource { + schema: schema.clone(), + opened: opened.clone(), + fail: false, + })); + let plan = plan(scan, &schema, ExecutionDataState::QUERY_ROWS); + let bound = bind_with_data_sources(&plan, BTreeMap::new(), &[0, 2], &sources).unwrap(); + let ctx = context(); + let streams = bound.execute(&[0, 2], ctx.clone()).unwrap(); + assert_eq!(opened.load(Ordering::SeqCst), 0); + ctx.cancel(); + drop(streams); + assert_eq!(opened.load(Ordering::SeqCst), 0); + for _ in 0..2 { + block_on(async { + let streams = bound.execute(&[0, 2], context()).unwrap(); + let all = + futures::future::join_all(streams.into_iter().map(|s| s.collect::>())).await; + assert!(all.iter().flatten().all(Result::is_ok)); + }); + } + assert_eq!(opened.load(Ordering::SeqCst), 2); +} +// Unavailable sources and unsupported predicates fail before opening any cursor. +#[test] +fn binding_errors_and_reader_errors_are_not_empty_results() { + let (mut scan, schema, _) = fixture(); + assert!(DataSources::default().bind(&scan).is_err()); + let opened = Arc::new(AtomicUsize::new(0)); + let sources = registry(Arc::new(CountingSource { + schema: schema.clone(), + opened: opened.clone(), + fail: true, + })); + if let QueryExpr::Scan { predicates, .. } = &mut scan { + predicates.push(Predicate(Rc::new(QueryExpr::Column(0)))); + } + assert!(sources.bind(&scan).is_err()); + assert_eq!(opened.load(Ordering::SeqCst), 0); + let (scan, _, _) = fixture(); + let plan = plan(scan, &schema, ExecutionDataState::QUERY_ROWS); + let bound = bind_with_data_sources(&plan, BTreeMap::new(), &[2], &sources).unwrap(); + block_on(async { + let mut stream = bound.execute(&[2], context()).unwrap().remove(0); + assert!(stream.next().await.unwrap().is_err()); + }); +} + +// Schema drift cannot enter the DAG, and connector batches obey execution limits. +#[test] +fn schema_drift_and_memory_limits_fail_the_scan() { + struct Drift { + expected: Schema, + batch: Batch, + } + impl RawSource for Drift { + fn schema(&self) -> Schema { + self.expected.clone() + } + fn scan(&self, _: RunContext) -> Result, Error> { + Ok(stream::once(async { Ok(self.batch.clone()) }).boxed_local()) + } + } + let (scan, schema, batches) = fixture(); + let mut different = (*schema).clone(); + different.fields[0].name = "wrong".into(); + let bad = Batch::try_new(Arc::new(different), vec![vec![Value::Int64(1)]]).unwrap(); + let sources = registry(Arc::new(Drift { + expected: schema.clone(), + batch: bad, + })); + let plan = plan(scan, &schema, ExecutionDataState::QUERY_ROWS); + let graph = bind_with_data_sources(&plan, BTreeMap::new(), &[0], &sources).unwrap(); + block_on(async { + let mut s = graph.execute(&[0], context()).unwrap().remove(0); + assert!(s.next().await.unwrap().is_err()); + }); + let sources = registry(Arc::new(MemorySource::new(schema, batches).unwrap())); + let graph = bind_with_data_sources(&plan, BTreeMap::new(), &[0], &sources).unwrap(); + let ctx = RunContext::new( + Scope::Query { + evaluation_time_ms: 0, + revision: 1, + }, + Limits { + max_bytes: 1, + max_buffered_batches: 1, + }, + ) + .unwrap(); + block_on(async { + let mut s = graph.execute(&[0], ctx.clone()).unwrap().remove(0); + assert!(s.next().await.unwrap().is_err()); + }); + assert_eq!(ctx.retained_bytes(), 0); +} + +// An empty table is a valid empty scan; nullable comparisons retain only TRUE. +#[test] +fn empty_sources_and_three_valued_predicates() { + use planner_types::pre_asap::{CompareOpKind, ScalarValue}; + let (mut scan, schema, batches) = fixture(); + if let QueryExpr::Scan { + predicates, source, .. + } = &mut scan + { + *source = Source::TimeSeries { + metric: "samples".into(), + }; + *predicates = vec![Predicate(Rc::new(QueryExpr::Compare { + left: Rc::new(QueryExpr::Column(0)), + op: CompareOpKind::Gt, + right: Rc::new(QueryExpr::Literal(ScalarValue::Int64(2))), + }))]; + } + for (batches, expected) in [(vec![], 0), (batches, 2)] { + let mut sources = DataSources::default(); + sources + .register( + Source::TimeSeries { + metric: "samples".into(), + }, + Arc::new(MemorySource::new(schema.clone(), batches).unwrap()), + ) + .unwrap(); + let plan = plan(scan.clone(), &schema, ExecutionDataState::QUERY_ROWS); + let graph = bind_with_data_sources(&plan, BTreeMap::new(), &[0], &sources).unwrap(); + block_on(async { + let mut s = graph.execute(&[0], context()).unwrap().remove(0); + let mut count = 0; + while let Some(b) = s.next().await { + count += b.unwrap().rows().len(); + } + assert_eq!(count, expected); + }); + } +} + +// A physical candidate can be compiled once without readers and rebound per run. +#[test] +fn compile_without_readers_and_rebind_inputs() { + use asap_physical_operators::{ + operators::Operator, + physical_planner::{compile, InputContract, Source}, + }; + let (scan, schema, batches) = fixture(); + let dag = plan(scan, &schema, ExecutionDataState::QUERY_ROWS); + let compiled = compile( + &dag, + BTreeMap::from([(0, InputContract::bounded(schema.clone()))]), + &[2], + ) + .unwrap(); + assert_eq!(compiled.input_contracts().count(), 1); + for _ in 0..2 { + let sources = BTreeMap::from([( + 0, + Box::new(Operator::source(schema.clone(), batches.clone()).unwrap()) as Source<'_>, + )]); + let graph = compiled.instantiate(sources).unwrap(); + let mut outputs = graph.execute(compiled.roots(), context()).unwrap(); + let result = block_on(outputs.remove(0).collect::>()); + assert!(result.iter().all(Result::is_ok)); + assert_eq!( + result + .iter() + .map(|b| b.as_ref().unwrap().rows().len()) + .sum::(), + 2 + ); + } + assert!(compiled.instantiate(BTreeMap::new()).is_err()); +} + +// Input boundedness must be proved during compilation, before readers exist. +#[test] +fn compilation_rejects_unknown_boundedness_for_sort() { + use asap_physical_operators::{ + physical_planner::{compile, InputContract}, + plan::{Boundedness, Emission, PlanProperties}, + }; + let (scan, schema, _) = fixture(); + let dag = plan(scan, &schema, ExecutionDataState::QUERY_ROWS); + let input = InputContract { + schema, + properties: PlanProperties { + boundedness: Boundedness::Unknown, + emission: Emission::Unknown, + }, + }; + assert!(compile(&dag, BTreeMap::from([(0, input)]), &[2]).is_err()); +} diff --git a/crates/asap-physical-operators/tests/weighted_topk_binding.rs b/crates/asap-physical-operators/tests/weighted_topk_binding.rs new file mode 100644 index 00000000..e4be66c1 --- /dev/null +++ b/crates/asap-physical-operators/tests/weighted_topk_binding.rs @@ -0,0 +1,1016 @@ +//! Planner output binds directly to the shared runtime at a declared rate-value frontier. +use asap_aware_mapping::{ + accuracy::{ + AccuracyEvidenceProvider, DefaultAccuracyModel, EqualSplitAllocator, PropagationStats, + }, + cost_model::DefaultCostModel, + Replacement, ReplacementStrategy, SketchAlgorithmStrategy, TargetSubDAG, +}; +use asap_physical_operators::dag::{ + operators::Operator, + planner::{compile, InputContract, Source}, + values::{Batch, Value}, + Limits, RunContext, Scope, +}; +use futures::{executor::block_on, StreamExt}; +use planner_types::{ + post_asap::*, + pre_asap::{DataType, QueryExpr}, + types::AccuracyTarget, +}; +use std::{collections::BTreeMap, rc::Rc, sync::Arc}; +struct Evidence; +impl AccuracyEvidenceProvider for Evidence { + fn topk_max_distinct_items(&self, _: &QueryExpr) -> Option { + Some(1000) + } + fn propagation_stats( + &self, + op: &CompositionOperator, + _: &SummaryFamilyType, + _: Option<&SketchQuery>, + ) -> PropagationStats { + if matches!(op, CompositionOperator::TopKSelection) { + PropagationStats { + topk_selected_lower_bound: Some(101.), + topk_excluded_upper_bound: Some(100.), + topk_interval_failure_probability: Some(0.001), + ..Default::default() + } + } else { + Default::default() + } + } +} +// The evidence here exercises binding; it is not inferred from the sample data. +#[test] +fn planner_weighted_topk_binds_at_either_deployment_phase() { + assert_weighted_binding(&Evidence, SketchAlgorithm::CmsWithHeap); + assert_weighted_binding(&Evidence, SketchAlgorithm::CountSketchWithHeap); +} + +// Binding validates representation, while deployment owns evidence acceptance. +#[test] +fn physical_binding_does_not_impose_an_accuracy_acceptance_policy() { + assert_weighted_binding( + &asap_aware_mapping::accuracy::NoAccuracyEvidence, + SketchAlgorithm::CmsWithHeap, + ); + assert_weighted_binding( + &asap_aware_mapping::accuracy::NoAccuracyEvidence, + SketchAlgorithm::CountSketchWithHeap, + ); +} + +fn assert_weighted_binding(evidence: &dyn AccuracyEvidenceProvider, algorithm: SketchAlgorithm) { + let root = Rc::new( + lower_promql( + "topk by(job)(2, sum by(service, job)(rate(m[1m])))", + AccuracyTarget::Epsilon(0.1), + ) + .unwrap(), + ); + let strategy = SketchAlgorithmStrategy::new_with_planning_inputs_and_evidence( + &DefaultCostModel, + &DefaultAccuracyModel, + &EqualSplitAllocator, + evidence, + ); + let plan = strategy + .replacements(&TargetSubDAG::new(&root)) + .into_iter() + .find_map(|candidate| match candidate.replacement { + Replacement::Summary(node) + if candidate.rationale.contains(&format!("{algorithm:?}")) => + { + Some(node) + } + _ => None, + }) + .unwrap(); + let dag = compile_executable_dag(&plan).unwrap(); + let build=dag.nodes.iter().find(|node|matches!(&node.payload,ExecutableOperatorPayload::SummaryAgg{family:SummaryFamilyType::Sketch(kind,_),..}if kind.algorithm()==&algorithm)).unwrap(); + let rate_id = dag + .edges + .iter() + .find(|edge| edge.consumer == build.id) + .unwrap() + .producer; + let rates = Arc::new( + dag.nodes + .iter() + .find(|node| node.id == rate_id) + .unwrap() + .output_schema + .clone(), + ); + let rows = [ + ("auth", "api", 0.125), + ("auth", "api", 0.25), + ("checkout", "api", 0.3125), + ("search", "api", 0.0625), + ("ingest", "batch", 100.), + ("export", "batch", 80.), + ("cleanup", "batch", 20.), + ] + .into_iter() + .map(|(service, job, value)| { + rates + .fields + .iter() + .map(|field| match field.name.as_str() { + "service" => Value::Utf8(service.into()), + "job" => Value::Utf8(job.into()), + "value" => Value::Float64(value), + _ => match field.dtype { + SummaryFamilyType::Plain(DataType::Timestamp) => Value::Timestamp(60_000), + _ => panic!("unexpected rate column {field:?}"), + }, + }) + .collect() + }) + .collect(); + let batch = Batch::try_new(rates.clone(), rows).unwrap(); + for (phase, scope) in [ + ( + ExecutionTiming::IngestionTime, + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 60_000, + revision: 1, + }, + ), + ( + ExecutionTiming::QueryTime, + Scope::Query { + evaluation_time_ms: 60_000, + revision: 1, + }, + ), + ] { + let placed = dag + .with_execution_phases(&dag.nodes.iter().map(|node| (node.id, phase)).collect()) + .unwrap(); + let source = Box::new(Operator::source(rates.clone(), vec![batch.clone()]).unwrap()) + as Source<'static>; + let compiled = compile( + &placed, + BTreeMap::from([(rate_id.0 as u64, InputContract::bounded(rates.clone()))]), + &[dag.root.0 as u64], + ) + .unwrap(); + let graph = compiled + .instantiate(BTreeMap::from([(rate_id.0 as u64, source)])) + .unwrap(); + let context = RunContext::new(scope, Limits::default()).unwrap(); + let output = block_on(async { + let mut output = Vec::new(); + let mut stream = graph + .execute(&[dag.root.0 as u64], context) + .unwrap() + .remove(0); + while let Some(batch) = stream.next().await { + output.extend(batch.unwrap().rows().iter().cloned()); + } + output + }); + assert_eq!(output.len(), 4); + let mut scores = output + .iter() + .map(|row| { + row.iter() + .find_map(|v| { + if let Value::Float64(v) = v { + Some(*v) + } else { + None + } + }) + .unwrap() + }) + .collect::>(); + scores.sort_by(f64::total_cmp); + assert_eq!(scores, vec![0.3125, 0.375, 80., 100.]); + } +} + +use asap_frontend_promql::lower_promql_workload; +use planner_types::workload::{ + AccuracyRequirement, BatchEntry, DataWorkload, DurationMs, Evidence as WorkloadEvidence, + PlanningWorkload, Predictability, Query, QueryLanguage, QueryRequirements, QueryWorkload, + TimeSelection, +}; +pub fn lower_promql( + query: &str, + accuracy: AccuracyTarget, +) -> Result { + let workload = PlanningWorkload { + query_workload: QueryWorkload { + language: QueryLanguage::PromQL, + query_batch: Some(vec![BatchEntry { + query: Query(query.into()), + requirements: QueryRequirements { + accuracy: AccuracyRequirement::Explicit(accuracy), + ..Default::default() + }, + predictability: Predictability::Unknown, + invocations: 1, + execute_at: None, + time_selection: TimeSelection::default(), + }]), + repeating_queries: None, + }, + data_workload: Some(DataWorkload { + data_ingestion_interval: WorkloadEvidence { + value: Some(DurationMs(1_000)), + ..Default::default() + }, + ..Default::default() + }), + }; + let mut lowered = lower_promql_workload(&workload, 0)?; + Ok(lowered.remove(0)) +} + +// The old untyped heap updater must not silently round a Planner rate update. +#[test] +fn rate_updates_cannot_enter_integer_heap_factory() { + let family = SummaryFamilyType::Sketch( + SketchKind::new( + SketchAlgorithm::CmsWithHeap, + SketchParams::CmsWithHeap { + width: 272, + depth: 5, + heap_size: 100, + }, + ), + Default::default(), + ); + let input = SummaryUpdate { + item: Some(SummaryInputExpr::Column( + planner_types::pre_asap::ColumnRef::Named("service".into()), + )), + weight: SummaryInputExpr::Column(planner_types::pre_asap::ColumnRef::SampleValue), + weight_domain: WeightDomain::NonNegative { + proof: NonNegativeWeightProof::ResetAwareCounterDerivative, + }, + }; + assert!( + asap_physical_operators::factory::create_planner_accumulator( + &family, + &input, + &Default::default() + ) + .is_err() + ); +} + +/// A catalog-resolved per-series rate can feed a heap sketch directly, without +/// requiring an otherwise unnecessary grouped Sum between Rate and TopK. +#[test] +fn direct_rate_topk_exposes_heap_candidates_with_complete_series_identity() { + check_direct_rate_topk(false); +} + +// Unreferenced labels still distinguish series throughout Rate and heap readout. +#[test] +fn direct_rate_topk_preserves_dynamic_unreferenced_labels() { + check_direct_rate_topk(true); +} + +fn check_direct_rate_topk(dynamic: bool) { + use asap_physical_operators::physical_planner::promql_rows::{ + decode_series_identity, series_row, with_series_identity, SERIES_IDENTITY_COLUMN, + }; + let mut logical = + lower_promql("topk by(job)(2, rate(m[1m]))", AccuracyTarget::Epsilon(0.1)).unwrap(); + fn resolve_catalog(node: &mut QueryExpr) { + match node { + QueryExpr::Aggregate { child, .. } | QueryExpr::TimeRange { child, .. } => { + resolve_catalog(Rc::make_mut(child)) + } + QueryExpr::Scan { schema, .. } => { + schema.closed = true; + schema + .columns + .push(planner_types::pre_asap::schema::Column::new( + "service", + DataType::Utf8, + false, + )); + } + _ => panic!("unexpected input shape: {node:?}"), + } + } + if dynamic { + logical = with_series_identity(&logical).unwrap(); + } else { + resolve_catalog(&mut logical); + } + let root = Rc::new(logical); + let strategy = SketchAlgorithmStrategy::new_with_planning_inputs_and_evidence( + &DefaultCostModel, + &DefaultAccuracyModel, + &EqualSplitAllocator, + &Evidence, + ); + let candidates = strategy.replacements(&TargetSubDAG::new(&root)); + for algorithm in [ + SketchAlgorithm::CmsWithHeap, + SketchAlgorithm::CountSketchWithHeap, + ] { + let candidate = candidates + .iter() + .find_map(|candidate| match &candidate.replacement { + Replacement::Summary(node) + if candidate.rationale.contains(&format!("{algorithm:?}")) => + { + Some(node) + } + _ => None, + }) + .unwrap_or_else(|| panic!("missing {algorithm:?} over direct Rate")); + if dynamic { + let (source, ranked) = + asap_physical_operators::physical_planner::promql_rows::compile_rate_ranking( + candidate, + ) + .unwrap(); + assert!(matches!( + source.expr, + SummaryExpr::ValueOperation { + operation: ValueOperation::FinalizeExactAccumulator, + .. + } + )); + assert_eq!(ranked.input_contracts().count(), 1); + let encoded = String::from_utf8(ranked.encode().unwrap()).unwrap(); + assert!(encoded.contains("KeyedSummaryBuild")); + assert!(encoded.contains("KeyedReadout")); + assert!( + !encoded.contains("\"Rate\""), + "Rate must be supplied by its exact stored-state readout" + ); + } + let dag = compile_executable_dag(candidate).unwrap(); + assert!(dag.nodes.iter().any(|node| matches!(&node.payload, + ExecutableOperatorPayload::SummaryAgg { family: SummaryFamilyType::Sketch(kind, _), .. } if kind.algorithm() == &algorithm))); + let build = dag.nodes.iter().find(|node| matches!(&node.payload, + ExecutableOperatorPayload::SummaryAgg { family: SummaryFamilyType::Sketch(kind, _), .. } if kind.algorithm() == &algorithm)).unwrap(); + let input_id = dag + .edges + .iter() + .find(|edge| edge.consumer == build.id) + .unwrap() + .producer; + let schema = Arc::new( + dag.nodes + .iter() + .find(|node| node.id == input_id) + .unwrap() + .output_schema + .clone(), + ); + let raw = dag + .nodes + .iter() + .find(|node| { + matches!( + &node.payload, + ExecutableOperatorPayload::Fallback { + expression: QueryExpr::TimeRange { .. } + } + ) + }) + .unwrap_or_else(|| panic!("no raw counter source: {dag:?}")); + let raw_schema = Arc::new(raw.output_schema.clone()); + let raw_compiled = compile( + &dag, + BTreeMap::from([( + u64::from(raw.id.0), + InputContract::bounded(raw_schema.clone()), + )]), + &[u64::from(dag.root.0)], + ) + .unwrap(); + let bytes = raw_compiled.encode().unwrap(); + let raw_compiled = + asap_physical_operators::physical_planner::CompiledPhysicalDag::decode(&bytes).unwrap(); + // Each evaluation receives a complete raw window. A reset, a stopped + // series and an expired leader must not retain last run's heap weights. + for (end, series, expected) in [ + ( + 60_000, + vec![ + ("auth", vec![10., 30., 50.]), + ("checkout", vec![10., 50., 90.]), + ("search", vec![10., 70., 130.]), + ], + vec![11. / 6., 8. / 3.], + ), + ( + 120_000, + vec![ + ("auth", vec![100., 10., 50.]), + ("checkout", vec![100., 100., 100.]), + ], + vec![0., 1.25], + ), + ] { + let mut raw_rows = Vec::new(); + for (service, samples) in series { + for (offset, value) in [10_000, 30_000, 50_000].into_iter().zip(samples) { + if dynamic { + raw_rows.push( + series_row( + &raw_schema, + &BTreeMap::from([ + ("job".into(), "api".into()), + ("service".into(), service.into()), + ("unreferenced".into(), format!("{service}-extra")), + ]), + end - 60_000 + offset, + value, + ) + .unwrap(), + ); + continue; + } + raw_rows.push( + raw_schema + .fields + .iter() + .map(|field| match field.name.as_str() { + "service" => Value::Utf8(service.into()), + "job" => Value::Utf8("api".into()), + "value" => Value::Float64(value), + "ts" => Value::Timestamp(end - 60_000 + offset), + _ => panic!("unexpected raw field"), + }) + .collect(), + ); + } + } + let raw_batch = Batch::try_new(raw_schema.clone(), raw_rows).unwrap(); + for scope in [ + Scope::Ingestion { + window_start_ms: end - 60_000, + window_end_ms: end, + revision: 1, + }, + Scope::Query { + evaluation_time_ms: end, + revision: 1, + }, + ] { + let source = Box::new( + Operator::source(raw_schema.clone(), vec![raw_batch.clone()]).unwrap(), + ) as Source<'static>; + let graph = raw_compiled + .instantiate(BTreeMap::from([(u64::from(raw.id.0), source)])) + .unwrap(); + let context = RunContext::new(scope, Limits::default()).unwrap(); + let mut raw_scores = block_on(async { + let mut scores = Vec::new(); + let mut stream = graph + .execute(&[u64::from(dag.root.0)], context) + .unwrap() + .remove(0); + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + for row in batch.rows() { + if dynamic { + let column = batch + .schema() + .fields + .iter() + .position(|field| field.name == SERIES_IDENTITY_COLUMN) + .unwrap(); + let Value::Utf8(encoded) = &row[column] else { + panic!("identity lost"); + }; + let labels = decode_series_identity(encoded).unwrap(); + assert_eq!(labels["job"], "api"); + assert_eq!( + labels["unreferenced"], + format!("{}-extra", labels["service"]) + ); + } + assert!(row.iter().any( + |value| matches!(value, Value::Timestamp(time) if *time == end) + )); + scores.extend(row.iter().filter_map(|value| match value { + Value::Float64(value) => Some(*value), + _ => None, + })); + } + } + scores + }); + raw_scores.sort_by(f64::total_cmp); + assert_eq!(raw_scores.len(), expected.len()); + for (actual, expected) in raw_scores.iter().zip(&expected) { + assert!( + (actual - expected).abs() < 1e-12, + "raw counter semantics must precede heap ranking: {raw_scores:?}" + ); + } + } + } + let compiled = compile( + &dag, + BTreeMap::from([( + u64::from(input_id.0), + InputContract::bounded(schema.clone()), + )]), + &[u64::from(dag.root.0)], + ) + .unwrap(); + for (time, values, expected) in [ + ( + 60_000, + vec![("auth", 3.), ("checkout", 2.), ("search", 1.)], + vec![2., 3.], + ), + ( + 61_000, + vec![("auth", 0.), ("checkout", 2.), ("search", 4.)], + vec![2., 4.], + ), + (62_000, vec![("auth", 0.), ("checkout", 2.)], vec![0., 2.]), + ] { + let rows = values + .into_iter() + .map(|(service, value)| { + if dynamic { + return series_row( + &schema, + &BTreeMap::from([ + ("job".into(), "api".into()), + ("service".into(), service.into()), + ]), + time, + value, + ) + .unwrap(); + } + schema + .fields + .iter() + .map(|field| match field.name.as_str() { + "service" => Value::Utf8(service.into()), + "job" => Value::Utf8("api".into()), + "value" => Value::Float64(value), + "ts" => Value::Timestamp(time), + _ => panic!("unexpected rate field {field:?}"), + }) + .collect() + }) + .collect(); + let batch = Batch::try_new(schema.clone(), rows).unwrap(); + for scope in [ + Scope::Query { + evaluation_time_ms: time, + revision: 1, + }, + Scope::Ingestion { + window_start_ms: time - 60_000, + window_end_ms: time, + revision: 1, + }, + ] { + let source = + Box::new(Operator::source(schema.clone(), vec![batch.clone()]).unwrap()) + as Source<'static>; + let graph = compiled + .instantiate(BTreeMap::from([(u64::from(input_id.0), source)])) + .unwrap(); + let context = RunContext::new(scope, Limits::default()).unwrap(); + let mut scores = block_on(async { + let mut scores = vec![]; + let mut stream = graph + .execute(&[u64::from(dag.root.0)], context) + .unwrap() + .remove(0); + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + for row in batch.rows() { + assert!(row.iter().any( + |value| matches!(value, Value::Timestamp(actual) if *actual == time) + )); + scores.push( + row.iter() + .find_map(|value| { + if let Value::Float64(value) = value { + Some(*value) + } else { + None + } + }) + .unwrap(), + ); + } + } + scores + }); + scores.sort_by(f64::total_cmp); + assert_eq!( + scores, expected, + "heap snapshots must not accumulate across evaluations" + ); + } + } + } +} + +// Spatial ranking consumes one eligible instant vector. Signed values require +// CountSketch; a raw metric does not establish the non-negative CMS contract. +#[test] +fn spatial_topk_exposes_signed_heap_candidate_over_complete_snapshot() { + use asap_physical_operators::physical_planner::promql_rows::{ + decode_series_identity, series_row, with_series_identity, SERIES_IDENTITY_COLUMN, + }; + let logical = lower_promql("topk by(job)(1, m)", AccuracyTarget::Epsilon(0.1)).unwrap(); + let root = Rc::new(with_series_identity(&logical).unwrap()); + let strategy = SketchAlgorithmStrategy::new_with_planning_inputs_and_evidence( + &DefaultCostModel, + &DefaultAccuracyModel, + &EqualSplitAllocator, + &Evidence, + ); + let candidates = strategy + .current_series_topk_candidates(&root, &AccuracyTarget::Epsilon(0.1)) + .candidates; + assert!(!candidates + .iter() + .any(|c| c.rationale.contains("CmsWithHeap"))); + let selected = candidates + .iter() + .find_map(|candidate| match &candidate.replacement { + Replacement::Summary(node) if candidate.rationale.contains("CountSketchWithHeap") => { + Some(node) + } + _ => None, + }) + .expect("signed spatial TopK must expose CountSketch with heap"); + let dag = compile_executable_dag(selected).unwrap(); + let raw = dag + .nodes + .iter() + .find(|node| { + matches!( + &node.payload, + ExecutableOperatorPayload::Fallback { + expression: QueryExpr::TimeRange { .. } + } + ) + }) + .unwrap(); + let schema = Arc::new(raw.output_schema.clone()); + let program = compile( + &dag, + BTreeMap::from([(u64::from(raw.id.0), InputContract::bounded(schema.clone()))]), + &[u64::from(dag.root.0)], + ) + .unwrap(); + let snapshot_program = + asap_physical_operators::physical_planner::promql_rows::compile_current_series_readout( + selected, + ) + .unwrap(); + let encoded: serde_json::Value = + serde_json::from_slice(&snapshot_program.encode().unwrap()).unwrap(); + assert!(!encoded.to_string().contains("CurrentSeries")); + assert!(encoded.to_string().contains("KeyedSummaryBuild")); + assert!(encoded.to_string().contains("KeyedReadout")); + for (values, expected, score) in [ + ([100., 20.], "a", 100.), + ([1., 20.], "b", 20.), + ([-10., -2.], "b", -2.), + ] { + let rows = ["a", "b"] + .into_iter() + .zip(values) + .map(|(instance, value)| { + series_row( + &schema, + &BTreeMap::from([ + ("job".into(), "api".into()), + ("unreferenced".into(), instance.into()), + ]), + 60_000, + value, + ) + .unwrap() + }) + .collect(); + let batch = Batch::try_new(schema.clone(), rows).unwrap(); + let graph = program + .instantiate(BTreeMap::from([( + u64::from(raw.id.0), + Box::new(Operator::source(schema.clone(), vec![batch]).unwrap()) as Source<'_>, + )])) + .unwrap(); + block_on(async { + let context = RunContext::new( + Scope::Query { + evaluation_time_ms: 60_000, + revision: 0, + }, + Limits::default(), + ) + .unwrap(); + let mut stream = graph.execute(program.roots(), context).unwrap().remove(0); + let mut result = Vec::new(); + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + let identity = batch + .schema() + .fields + .iter() + .position(|f| f.name == SERIES_IDENTITY_COLUMN) + .unwrap(); + let value = batch + .schema() + .fields + .iter() + .position(|f| f.name == "value") + .unwrap(); + for row in batch.rows() { + let Value::Utf8(labels) = &row[identity] else { + panic!() + }; + let Value::Float64(v) = row[value] else { + panic!() + }; + result.push(( + decode_series_identity(labels).unwrap()["unreferenced"].clone(), + v, + )); + } + } + assert_eq!(result, vec![(expected.into(), score)]); + }); + } +} + +// Placement changes execution ownership only. Every fixed-window candidate +// contains Rate finalization before a fresh heap, with query readout downstream. +#[test] +fn planner_exposes_fixed_window_rate_heap_precompute_candidates() { + use asap_physical_operators::physical_planner::{ + compile_candidate, promql_rows::with_series_identity, + }; + let root = Rc::new( + with_series_identity( + &lower_promql("topk by(job)(2, rate(m[1m]))", AccuracyTarget::Epsilon(0.1)).unwrap(), + ) + .unwrap(), + ); + let strategy = SketchAlgorithmStrategy::new_with_planning_inputs_and_evidence( + &DefaultCostModel, + &DefaultAccuracyModel, + &EqualSplitAllocator, + &Evidence, + ); + let candidates = strategy.fixed_window_rate_candidates(&root).candidates; + assert_eq!(candidates.len(), 2); + for candidate in candidates { + let Replacement::Summary(root) = candidate.replacement else { + panic!() + }; + let dag = compile_executable_dag(&root).unwrap(); + let state = dag + .nodes + .iter() + .find(|node| { + matches!( + &node.payload, + ExecutableOperatorPayload::SummaryAgg { + family: SummaryFamilyType::ExactAggregate(ExactKind::Rate, _), + .. + } + ) + }) + .unwrap(); + let heap = dag + .nodes + .iter() + .find(|node| { + matches!( + &node.payload, + ExecutableOperatorPayload::SummaryAgg { + family: SummaryFamilyType::Sketch(..), + .. + } + ) + }) + .unwrap(); + assert_eq!(heap.output_state.timing, ExecutionTiming::IngestionTime); + let physical = compile_candidate( + &dag, + BTreeMap::from([( + u64::from(state.id.0), + InputContract::bounded(Arc::new(state.output_schema.clone())), + )]), + &[u64::from(dag.root.0)], + &[u64::from(heap.id.0)], + ) + .unwrap(); + let exported = asap_physical_operators::physical_planner::promql_rows::compile_fixed_window_rate_aggregation(&root).unwrap(); + assert_eq!(exported.encode().unwrap(), physical.encode().unwrap()); + assert!( + asap_physical_operators::physical_planner::promql_rows::compile_rate_ranking(&root) + .is_err(), + "query binding must not move the selected precompute frontier" + ); + // Execute the selected split across a state serialization boundary. + // Each run builds fresh weights from that window's counters. + let execute = |plan: &asap_physical_operators::physical_planner::CompiledPhysicalDag, + input: Batch, + scope: Scope| { + let id = plan.input_contracts().next().unwrap().0; + let source = Box::new(Operator::source(input.schema().clone(), vec![input]).unwrap()) + as Source<'static>; + let graph = plan.instantiate(BTreeMap::from([(id, source)])).unwrap(); + block_on(async { + let mut stream = graph + .execute( + plan.roots(), + RunContext::new(scope, Limits::default()).unwrap(), + ) + .unwrap() + .remove(0); + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push((*batch.unwrap()).clone()); + } + assert_eq!(batches.len(), 1); + batches.remove(0) + }) + }; + let (family, input, grouping) = match &state.payload { + ExecutableOperatorPayload::SummaryAgg { + family, + input, + grouping, + .. + } => (family, input, grouping), + _ => unreachable!(), + }; + for (end, samples, leader) in [ + ( + 60_000, + [[0., 100., 200.], [0., 10., 20.], [0., 1., 2.]], + "a", + ), + ( + 120_000, + [[200., 200., 200.], [100., 0., 300.], [2., 3., 4.]], + "b", + ), + ] { + let schema = Arc::new(state.output_schema.clone()); + let rows = samples + .into_iter() + .zip(["a", "b", "c"]) + .map(|(samples, label)| { + let mut accumulator = + asap_physical_operators::factory::create_planner_accumulator( + family, input, grouping, + ) + .unwrap(); + for (offset, value) in [10_000, 30_000, 50_000].into_iter().zip(samples) { + accumulator.update_single(value, end - 60_000 + offset); + } + let summary = Value::Summary { + family: family.clone(), + state: Arc::from(accumulator.into_accumulator()), + }; + schema + .fields + .iter() + .map(|field| match &field.dtype { + SummaryFamilyType::ExactAggregate(..) => summary.clone(), + SummaryFamilyType::Plain(DataType::Timestamp) => Value::Timestamp(end), + SummaryFamilyType::Plain(DataType::Utf8) + if field.name == "$promql_series_identity" => + { + Value::Utf8( + serde_json::to_string(&BTreeMap::from([ + ("job", "api"), + ("instance", label), + ])) + .unwrap() + .into(), + ) + } + SummaryFamilyType::Plain(DataType::Utf8) => Value::Utf8("api".into()), + _ => panic!("unexpected state field {field:?}"), + }) + .collect() + }) + .collect(); + let batch = Batch::try_new(schema, rows).unwrap(); + let precompute = physical.precompute.as_ref().unwrap(); + let heap = execute( + precompute, + batch, + Scope::Ingestion { + window_start_ms: end - 60_000, + window_end_ms: end, + revision: 1, + }, + ); + let bytes = asap_physical_operators::stored_state::native::encode_batch(&heap).unwrap(); + let restored = asap_physical_operators::stored_state::native::decode_batch( + &bytes, + heap.schema().clone(), + 1 << 24, + ) + .unwrap(); + let result = execute( + &physical.query, + restored, + Scope::Query { + evaluation_time_ms: end, + revision: 1, + }, + ); + let identity = result + .schema() + .fields + .iter() + .position(|f| f.name == "$promql_series_identity") + .unwrap(); + let Value::Utf8(encoded) = &result.rows()[0][identity] else { + panic!() + }; + let labels: BTreeMap = serde_json::from_str(encoded).unwrap(); + assert_eq!(labels["instance"], leader); + assert_eq!(result.rows().len(), 2); + } + let precompute = String::from_utf8(physical.precompute.unwrap().encode().unwrap()).unwrap(); + assert!(precompute.contains("KeyedSummaryBuild")); + assert!(precompute.contains("Rate")); + assert!(!String::from_utf8(physical.query.encode().unwrap()) + .unwrap() + .contains("KeyedSummaryBuild")); + } +} + +// Grouped Rate has a legal stored Sum candidate as well as query-time reduction. +#[test] +fn grouped_rate_exposes_precomputed_sum_with_query_readout() { + let root = Rc::new( + asap_physical_operators::physical_planner::promql_rows::with_series_identity( + &lower_promql("sum by(job)(rate(m[1m]))", AccuracyTarget::Exact).unwrap(), + ) + .unwrap(), + ); + let strategy = SketchAlgorithmStrategy::new_with_planning_inputs_and_evidence( + &DefaultCostModel, + &DefaultAccuracyModel, + &EqualSplitAllocator, + &Evidence, + ); + let direct = strategy.query_time_rate_aggregation_candidates(&root); + assert!( + direct.candidates.iter().any(|candidate| { + let Replacement::Summary(root) = &candidate.replacement else { + return false; + }; + let Ok((_, program)) = + asap_physical_operators::physical_planner::promql_rows::compile_rate_ranking(root) + else { + return false; + }; + let output = program.output_contract(program.roots()[0]).unwrap(); + output + .schema + .fields + .iter() + .all(|field| matches!(field.dtype, SummaryFamilyType::Plain(_))) + }), + "query-time grouped Rate must finalize Sum inside the physical graph" + ); + let candidates = strategy.fixed_window_rate_candidates(&root).candidates; + assert!( + !candidates.is_empty(), + "Planner must expose Rate -> grouped Sum at ingestion" + ); + for candidate in candidates { + let Replacement::Summary(root) = candidate.replacement else { + panic!() + }; + let physical = asap_physical_operators::physical_planner::promql_rows::compile_fixed_window_rate_aggregation(&root).unwrap(); + let precompute = String::from_utf8(physical.precompute.unwrap().encode().unwrap()).unwrap(); + assert!( + precompute.contains("SummaryBuild") + && precompute.contains("Rate") + && precompute.contains("Sum") + ); + let query = String::from_utf8(physical.query.encode().unwrap()).unwrap(); + assert!(query.contains("Readout") && !query.contains("SummaryBuild")); + } +} diff --git a/crates/asap_sketch_codec/Cargo.toml b/crates/asap_sketch_codec/Cargo.toml new file mode 100644 index 00000000..2a1634b8 --- /dev/null +++ b/crates/asap_sketch_codec/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "asap_sketch_codec" +version = "0.1.0" +edition = "2021" + +[dependencies] +asap_sketchlib = { git = "https://github.com/ProjectASAP/asap_sketchlib", rev = "5f03ccbd798ed5fec62bdd839bcb331123cab369" } +prost = "0.13" diff --git a/crates/asap_sketch_codec/src/lib.rs b/crates/asap_sketch_codec/src/lib.rs new file mode 100644 index 00000000..00efb284 --- /dev/null +++ b/crates/asap_sketch_codec/src/lib.rs @@ -0,0 +1,80 @@ +//! Runtime-independent decoding of the sketchlib protobuf envelope. + +use asap_sketchlib::proto::sketchlib::{ + sketch_envelope::SketchState, DdSketchState, KllState, SketchEnvelope, +}; +use asap_sketchlib::DdSketch; +use prost::Message; + +pub fn envelope_state(bytes: &[u8]) -> Result, String> { + SketchEnvelope::decode(bytes) + .map(|envelope| envelope.sketch_state) + .map_err(|error| format!("decode SketchEnvelope: {error}")) +} + +pub fn ddsketch_state(bytes: &[u8]) -> Result<(DdSketchState, f64), String> { + let envelope = + SketchEnvelope::decode(bytes).map_err(|error| format!("decode SketchEnvelope: {error}"))?; + match envelope.sketch_state { + Some(SketchState::Ddsketch(state)) => Ok((state, envelope.sample_p)), + _ => Err("SketchEnvelope contains no DDSketch state".into()), + } +} + +pub fn reconstruct_ddsketch(bytes: &[u8]) -> Result<(DdSketch, f64), String> { + let (state, sample_p) = ddsketch_state(bytes)?; + if !state.alpha.is_finite() || !(0.0..1.0).contains(&state.alpha) || state.alpha == 0.0 { + return Err("DDSketch alpha must be finite and between zero and one".into()); + } + Ok(( + DdSketch::from_raw(state.alpha, state.store_counts, state.store_offset), + sample_p, + )) +} + +pub fn kll_state(bytes: &[u8]) -> Result { + let envelope = + SketchEnvelope::decode(bytes).map_err(|error| format!("decode SketchEnvelope: {error}"))?; + match envelope.sketch_state { + Some(SketchState::Kll(state)) => Ok(state), + _ => Err("SketchEnvelope contains no KLL state".into()), + } +} + +pub fn encode_ddsketch(sketch: &DdSketch) -> Vec { + let envelope = SketchEnvelope { + format_version: 1, + producer: None, + hash_spec: None, + sample_p: 0.0, + sketch_state: Some(SketchState::Ddsketch(sketch.to_proto())), + }; + envelope.encode_to_vec() +} + +pub fn encode_kll(sketch: &asap_sketchlib::sketches::kll::KLL) -> Vec { + use asap_sketchlib::proto::sketchlib::CoinState; + let (state, bit_cache, remaining_bits) = sketch.wire_coin(); + SketchEnvelope { + format_version: 1, + producer: None, + hash_spec: None, + sample_p: 0.0, + sketch_state: Some(SketchState::Kll(KllState { + k: sketch.wire_k(), + m: sketch.wire_m(), + num_levels: sketch.wire_num_levels(), + levels: sketch.wire_levels(), + items: sketch.wire_items(), + coin: Some(CoinState { + state, + bit_cache, + remaining_bits, + }), + offset: 0.0, + value_scale: 0, + residuals: Vec::new(), + })), + } + .encode_to_vec() +} diff --git a/crates/frontend-promql/src/promql.rs b/crates/frontend-promql/src/promql.rs index f0cf907f..c3df6c67 100644 --- a/crates/frontend-promql/src/promql.rs +++ b/crates/frontend-promql/src/promql.rs @@ -625,7 +625,10 @@ fn build_over_subtree(outer: Outer, keys: Vec, child: Unresolved) -> && matches!(sum_child.as_ref(), Unresolved::Aggregate { measures, .. } if matches!(measures.as_slice(), [AggIntent::Rate | AggIntent::Increase])) ); - if descending && weighted_counter_ranking { + let direct_counter_ranking = matches!(&child, Unresolved::Aggregate { + measures, reduction: Reduction::PerEntity, .. + } if matches!(measures.as_slice(), [AggIntent::Rate | AggIntent::Increase])); + if descending && (weighted_counter_ranking || direct_counter_ranking) { return Ok(outer_aggregate( keys, AggIntent::TopK { @@ -1487,10 +1490,27 @@ fn build(inner: Inner, keys: Vec, outer: Outer) -> Result }) } Outer::TopK { k, descending } => { + // Preserve the counter-value ranking intent. Physical candidates + // may rebuild a heap over finalized rates or use exact Sort/Limit; + // neither is allowed to sum raw counter samples as ranking weights. + if descending && matches!(inner.func, Some(InnerFunc::Rate | InnerFunc::Increase)) { + let intent = inner_intent(inner.func.as_ref().expect("counter function")); + let ranked = windowed_aggregate(inner, vec![], intent); + return Ok(Unresolved::Aggregate { + reduction: Reduction::Reduce(keys.into()), + measures: vec![AggIntent::TopK { + k: k as usize, + accuracy: current_accuracy(), + }], + output_names: vec![], + having: None, + child: Rc::new(ranked), + }); + } // Heavy-hitter only when ranking by an additive measure (`count` // or `sum`): that is a // first-class aggregate intent → `TopK`. Any other ranking (topk - // over avg/quantile/rate, a bare selector's raw value, all bottomk) + // over avg/quantile, a bare selector's raw value, all bottomk) // is a generic order-by-value + limit and stays as the `Sort + Limit` // operator pair. The descending-plus-measure rule is shared with the // canonicalize-pass promotion so the two cannot drift (issue #38). diff --git a/crates/integration-tests/Cargo.toml b/crates/integration-tests/Cargo.toml index 7b0c0d3e..5a5de990 100644 --- a/crates/integration-tests/Cargo.toml +++ b/crates/integration-tests/Cargo.toml @@ -13,3 +13,6 @@ asap-aware-mapping = { path = "../asap-aware-mapping" } asap_sketchlib = { workspace = true } serde_json = "1" tokio = { version = "1", features = ["rt", "macros", "rt-multi-thread"] } + +asap-physical-operators = { path = "../asap-physical-operators" } +futures = "0.3" diff --git a/crates/integration-tests/tests/kll_pane_execution.rs b/crates/integration-tests/tests/kll_pane_execution.rs new file mode 100644 index 00000000..2273f965 --- /dev/null +++ b/crates/integration-tests/tests/kll_pane_execution.rs @@ -0,0 +1,294 @@ +//! Maintenance -> wire state -> independently bound query execution. +mod physical_common; +use asap_physical_operators::{ + operators::Operator, + physical_planner::{CompiledPhysicalDag, InputContract, Source}, + plan::{PhysicalDag, PhysicalOperator, PlanProperties}, + runtime::{Input, Limits, OutputStream, RunContext, Scope}, + summary_kernels::datasketches_kll::DatasketchesKLLAccumulator, + values::{Batch, Schema, Value}, + Error, Statistic, +}; +use asap_types::{ + post_asap::{ + SketchAlgorithm, SketchKind, SketchParams, SummaryFamilyType, SummaryField, SummarySchema, + }, + pre_asap::DataType, +}; +use futures::{executor::block_on, StreamExt}; +use std::{ + collections::{BTreeMap, HashMap}, + sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, + }, +}; + +fn family(k: u32) -> SummaryFamilyType { + SummaryFamilyType::Sketch( + SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k }), + Default::default(), + ) +} +fn raw_schema() -> Schema { + Arc::new(SummarySchema { + fields: vec![SummaryField { + name: "value".into(), + dtype: SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }], + time_index: None, + }) +} +fn query_scope() -> Scope { + Scope::Query { + evaluation_time_ms: 300_000, + revision: 1, + } +} +fn fixture() -> (CompiledPhysicalDag, Operator, Schema) { + let raw = raw_schema(); + let build = Operator::summary_build(raw.clone(), family(200), 0, None, vec![]).unwrap(); + let state = build.schema(); + let maintenance = CompiledPhysicalDag::from_operators( + BTreeMap::from([(0, InputContract::bounded(raw))]), + BTreeMap::from([(1, (vec![0], build))]), + vec![1], + ) + .unwrap(); + let merge = Operator::summary_merge(state.clone(), 0, vec![]).unwrap(); + (maintenance, merge, state) +} +fn pane_bytes(maintenance: &CompiledPhysicalDag, pane: i64) -> Vec { + // Twenty samples in each (start,end] one-minute pane; k=200 avoids + // compaction so quantiles and sample counts have deterministic oracles. + let raw = raw_schema(); + let rows = (0..20) + .map(|i| vec![Value::Float64((pane * 20 + i) as f64)]) + .collect(); + let state = physical_common::execute( + maintenance, + BTreeMap::from([(0, Batch::try_new(raw, rows).unwrap())]), + Scope::Ingestion { + window_start_ms: pane * 60_000, + window_end_ms: (pane + 1) * 60_000, + revision: 1, + }, + ); + let Value::Summary { state, .. } = &state[0][0].rows()[0][0] else { + panic!("missing KLL") + }; + state.serialize_to_bytes() +} +fn restore(schema: Schema, bytes: &[Vec]) -> Batch { + Batch::try_new( + schema, + bytes + .iter() + .map(|bytes| { + vec![Value::Summary { + family: family(200), + state: Arc::new(DatasketchesKLLAccumulator::from_msgpack_bytes(bytes).unwrap()), + }] + }) + .collect(), + ) + .unwrap() +} +fn readout(schema: Schema, q: f64) -> Operator { + Operator::readout( + schema, + 0, + Statistic::Quantile, + HashMap::from([("quantile".into(), q.to_string())]), + ) + .unwrap() +} +struct CountStarts { + operator: Operator, + starts: Arc, +} +impl PhysicalOperator for CountStarts { + fn name(&self) -> &str { + self.operator.name() + } + fn properties(&self, inputs: &[PlanProperties]) -> PlanProperties { + self.operator.properties(inputs) + } + fn requires_bounded_input(&self) -> bool { + self.operator.requires_bounded_input() + } + fn input_schemas(&self) -> Vec { + self.operator.input_schemas() + } + fn output_schema(&self) -> Schema { + self.operator.output_schema() + } + fn output_bytes(&self, batch: &Batch) -> usize { + self.operator.output_bytes(batch) + } + fn start<'a>( + &'a self, + inputs: Vec>, + context: RunContext, + ) -> Result, Error> { + self.starts.fetch_add(1, Ordering::SeqCst); + self.operator.start(inputs, context) + } +} + +/// Actual codec bytes survive destruction of maintenance state; one shared +/// native merge supplies p50, p99 and the population-count oracle per run. +#[test] +fn five_panes_roundtrip_and_shared_merge_runs_once() { + let (maintenance, merge, schema) = fixture(); + let bytes: Vec<_> = (0..6).map(|pane| pane_bytes(&maintenance, pane)).collect(); + drop(maintenance); + let compiled = CompiledPhysicalDag::from_operators( + (0..5) + .map(|id| (id, InputContract::bounded(schema.clone()))) + .collect(), + BTreeMap::from([ + ( + 5, + ( + vec![0, 1, 2, 3, 4], + Operator::union(schema.clone(), 5).unwrap(), + ), + ), + (6, (vec![5], merge.clone())), + (7, (vec![6], readout(schema.clone(), 0.5))), + (8, (vec![6], readout(schema.clone(), 0.99))), + ]), + vec![6, 7, 8], + ) + .unwrap(); + let layout = asap_types::post_asap::PaneLayout { + pane_width_ms: 60_000, + pane_origin_ms: Some(0), + }; + assert!( + asap_types::post_asap::validate_pane_coverage( + &layout, + Some(330_000), + &asap_types::post_asap::WindowEdgeCoverage::PaneAligned + ) + .is_err(), + "moving window edges require residual computation" + ); + for offset in [0, 1] { + let restored = restore(schema.clone(), &bytes[offset..offset + 5]); + let evaluation_time_ms = (5 + offset as i64) * 60_000; + asap_types::post_asap::validate_pane_coverage( + &layout, + Some(evaluation_time_ms), + &asap_types::post_asap::WindowEdgeCoverage::PaneAligned, + ) + .unwrap(); + let inputs: BTreeMap<_, _> = (0..5) + .map(|id| { + ( + id as u64, + restore(schema.clone(), &bytes[offset + id..offset + id + 1]), + ) + }) + .collect(); + // A five-pane deployment cannot bind only four state slots. + let incomplete: BTreeMap<_, _> = inputs + .iter() + .take(4) + .map(|(&id, batch)| { + ( + id, + Box::new(Operator::source(schema.clone(), vec![batch.clone()]).unwrap()) + as Source<'_>, + ) + }) + .collect(); + assert!(compiled.instantiate(incomplete).is_err()); + let result = physical_common::execute( + &compiled, + inputs, + Scope::Query { + evaluation_time_ms, + revision: 1, + }, + ); + let Value::Summary { state, .. } = &result[0][0].rows()[0][0] else { + panic!("missing merged state") + }; + let kll = state + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(kll.inner.count(), 100); + let value = |index: usize| match result[index][0].rows()[0][0] { + Value::Float64(value) => value, + _ => panic!("missing quantile"), + }; + assert!((value(1) - (50 + offset * 20) as f64).abs() <= 1.); + assert!((value(2) - (99 + offset * 20) as f64).abs() <= 1.); + let starts = Arc::new(AtomicUsize::new(0)); + let mut dag = PhysicalDag::default(); + dag.add( + 0, + vec![], + Operator::source(schema.clone(), vec![restored]).unwrap(), + ) + .unwrap(); + dag.add( + 1, + vec![0], + CountStarts { + operator: merge.clone(), + starts: starts.clone(), + }, + ) + .unwrap(); + dag.add(2, vec![1], readout(schema.clone(), 0.5)).unwrap(); + dag.add(3, vec![1], readout(schema.clone(), 0.99)).unwrap(); + let outputs = block_on(futures::future::join_all( + dag.execute( + &[2, 3], + RunContext::new(query_scope(), Limits::default()).unwrap(), + ) + .unwrap() + .into_iter() + .map(|stream| stream.collect::>()), + )); + assert!(outputs + .iter() + .all(|output| output.len() == 1 && output[0].is_ok())); + assert_eq!(starts.load(Ordering::SeqCst), 1); + } +} + +/// Corrupt wire data, relabelled KLL parameters and missing bindings fail +/// explicitly; decoding success is not a license to change the state contract. +#[test] +fn restored_panes_reject_corruption_parameters_schema_and_missing_binding() { + let (maintenance, merge, schema) = fixture(); + let bytes = pane_bytes(&maintenance, 0); + assert!(DatasketchesKLLAccumulator::from_msgpack_bytes(&bytes[..bytes.len() / 2]).is_err()); + let wrong = Value::Summary { + family: family(200), + state: Arc::new(DatasketchesKLLAccumulator::new(128)), + }; + assert!(Batch::try_new(schema.clone(), vec![vec![wrong]]).is_err()); + let compiled = CompiledPhysicalDag::from_operators( + BTreeMap::from([(0, InputContract::bounded(schema))]), + BTreeMap::from([(1, (vec![0], merge))]), + vec![1], + ) + .unwrap(); + assert!(compiled.instantiate(BTreeMap::new()).is_err()); + let raw = raw_schema(); + let source = Operator::source( + raw.clone(), + vec![Batch::try_new(raw, vec![vec![Value::Float64(1.)]]).unwrap()], + ) + .unwrap(); + assert!(compiled + .instantiate(BTreeMap::from([(0, Box::new(source) as Source<'_>)])) + .is_err()); +} diff --git a/crates/integration-tests/tests/physical_common/mod.rs b/crates/integration-tests/tests/physical_common/mod.rs new file mode 100644 index 00000000..aca4d4ef --- /dev/null +++ b/crates/integration-tests/tests/physical_common/mod.rs @@ -0,0 +1,42 @@ +use asap_physical_operators::{ + operators::Operator, + physical_planner::{CompiledPhysicalDag, Source}, + runtime::{Limits, RunContext, Scope}, + values::Batch, +}; +use futures::{executor::block_on, StreamExt}; +use std::collections::BTreeMap; + +pub fn execute( + plan: &CompiledPhysicalDag, + inputs: BTreeMap, + scope: Scope, +) -> Vec> { + let sources = inputs + .into_iter() + .map(|(id, batch)| { + ( + id, + Box::new(Operator::source(batch.schema().clone(), vec![batch]).unwrap()) + as Source<'_>, + ) + }) + .collect(); + let dag = plan.instantiate(sources).unwrap(); + block_on(async { + let streams = dag + .execute( + plan.roots(), + RunContext::new(scope, Limits::default()).unwrap(), + ) + .unwrap(); + futures::future::join_all(streams.into_iter().map(|mut stream| async move { + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push((*batch.unwrap()).clone()); + } + batches + })) + .await + }) +} diff --git a/crates/integration-tests/tests/promql_to_post_asap.rs b/crates/integration-tests/tests/promql_to_post_asap.rs index 5f763fe4..70a4feed 100644 --- a/crates/integration-tests/tests/promql_to_post_asap.rs +++ b/crates/integration-tests/tests/promql_to_post_asap.rs @@ -831,7 +831,7 @@ fn execute_topk_reference(plan: &SummaryNode) -> Vec<(String, f64)> { } #[test] -fn planner_topk_reference_execution_matches_ground_truth() { +fn planner_heap_topk_reference_execution_matches_ground_truth() { // Pin numeric results independently of the emitted IR: swapping weights, // losing identity, changing the window, or dropping k changes the answer. for (query, expected) in [ @@ -865,8 +865,11 @@ fn planner_topk_reference_execution_matches_ground_truth() { &EqualSplitAllocator, &SeparatedTopK, ); - let candidates = strategy.replacements(&TargetSubDAG::new(&pre)); - assert!(!candidates.is_empty(), "no plan for {query}"); + // This reference executor consumes keyed heap updates. The inventory + // also contains maintained exact values followed by sort/limit; those + // have a different execution contract and must not enter this fixture. + let candidates: Vec<_> = strategy.replacements(&TargetSubDAG::new(&pre)).into_iter().filter(|candidate| matches!(&candidate.replacement, Replacement::Summary(plan) if matches!(plan.expr, SummaryExpr::SummaryEstimate { query: SketchQuery::TopK { .. }, .. }))).collect(); + assert!(!candidates.is_empty(), "no heap candidate for {query}"); for candidate in candidates { let Replacement::Summary(plan) = candidate.replacement else { panic!("expected summary plan for {query}") diff --git a/crates/integration-tests/tests/sql_to_physical.rs b/crates/integration-tests/tests/sql_to_physical.rs new file mode 100644 index 00000000..22cb92a4 --- /dev/null +++ b/crates/integration-tests/tests/sql_to_physical.rs @@ -0,0 +1,165 @@ +//! SQL frontend, candidate selection, physical compilation and fresh-run execution. +use asap_aware_mapping::{search_workload, DefaultCostModel}; +use asap_frontend_sql::{lower_sql, SqlCatalog}; +use asap_physical_operators::{ + physical_planner::{compile, InputContract, Source}, + runtime::{Limits, RunContext, Scope}, + sources::{DataSources, MemorySource}, + values::{Batch, Value}, +}; +use asap_types::{ + post_asap::{compile_executable_dag, ExecutableOperatorPayload, SummaryFamilyType}, + pre_asap::{Column, DataType, QueryExpr, Schema}, + types::AccuracyTarget, +}; +use futures::StreamExt; +use std::{collections::BTreeMap, rc::Rc, sync::Arc}; + +/// SQL filtering and grouped aggregation survive logical/physical lowering; +/// rebinding the compiled DAG runs against new data rather than cached results. +#[tokio::test] +async fn sql_filter_grouped_sum_executes_and_rebinds() { + let catalog = SqlCatalog::new().with_table( + "metrics", + Schema::new(vec![ + Column::new("service", DataType::Utf8, false), + Column::new("value", DataType::Float64, true), + ]), + ); + for query in [ + "SELECT service, SUM(value) AS total FROM metrics WHERE value > 1 GROUP BY service", + "SELECT service, SUM(value) AS total FROM metrics GROUP BY service", + ] { + let logical = Rc::new( + lower_sql(query, &catalog, AccuracyTarget::Exact) + .await + .unwrap(), + ); + let space = search_workload(vec![("sql", logical)]); + let selected = space + .global_selection(&DefaultCostModel) + .assemble_selected_dag(&space.roots[0].1) + .unwrap() + .unwrap(); + let dag = compile_executable_dag(&selected).unwrap(); + let scan = dag + .nodes + .iter() + .find(|node| { + matches!( + &node.payload, + ExecutableOperatorPayload::Fallback { + expression: QueryExpr::Scan { .. } + } + ) + }) + .expect("raw SQL scan"); + let schema = Arc::new(scan.output_schema.clone()); + assert!(schema + .fields + .iter() + .all(|field| matches!(field.dtype, SummaryFamilyType::Plain(_)))); + let plan = compile( + &dag, + BTreeMap::from([(u64::from(scan.id.0), InputContract::bounded(schema.clone()))]), + &[u64::from(dag.root.0)], + ) + .unwrap(); + for multiplier in [1., 2.] { + let rows = [ + ("api", Some(2.)), + ("api", Some(3.)), + ("api", None), + ("batch", Some(4.)), + ("batch", Some(1.)), + ] + .into_iter() + .map(|(service, value)| { + schema + .fields + .iter() + .map(|field| match field.name.as_str() { + "service" => Value::Utf8(service.into()), + "value" => { + value.map_or(Value::Null, |value| Value::Float64(value * multiplier)) + } + _ => panic!("unexpected field {field:?}"), + }) + .collect() + }) + .collect(); + let ExecutableOperatorPayload::Fallback { expression } = &scan.payload else { + unreachable!() + }; + let QueryExpr::Scan { source, .. } = expression else { + unreachable!() + }; + let mut sources = DataSources::default(); + sources + .register( + source.clone(), + Arc::new( + MemorySource::new( + schema.clone(), + vec![Batch::try_new(schema.clone(), rows).unwrap()], + ) + .unwrap(), + ), + ) + .unwrap(); + let bound = plan + .instantiate(BTreeMap::from([( + u64::from(scan.id.0), + Box::new(sources.bind(expression).unwrap()) as Source<'_>, + )])) + .unwrap(); + let mut stream = bound + .execute( + plan.roots(), + RunContext::new( + Scope::Query { + evaluation_time_ms: 300_000, + revision: 1, + }, + Limits::default(), + ) + .unwrap(), + ) + .unwrap() + .remove(0); + let mut batches = Vec::new(); + while let Some(batch) = stream.next().await { + batches.push(batch.unwrap()); + } + let mut actual: Vec<_> = batches + .iter() + .flat_map(|batch| batch.rows()) + .map(|row| { + let Value::Utf8(service) = &row[0] else { + panic!("missing service") + }; + let Value::Float64(value) = row[1] else { + panic!("missing sum") + }; + (service.to_string(), value) + }) + .collect(); + actual.sort_by(|a, b| a.0.cmp(&b.0)); + assert_eq!( + actual, + vec![ + ("api".into(), 5. * multiplier), + ( + "batch".into(), + 4. * multiplier + + if query.contains("WHERE") && multiplier == 1. { + 0. + } else { + multiplier + } + ) + ] + ); + } + } +} diff --git a/crates/integration-tests/tests/summary_maintenance_lifecycle_e2e.rs b/crates/integration-tests/tests/summary_maintenance_lifecycle_e2e.rs index 186dd04b..61003a50 100644 --- a/crates/integration-tests/tests/summary_maintenance_lifecycle_e2e.rs +++ b/crates/integration-tests/tests/summary_maintenance_lifecycle_e2e.rs @@ -119,52 +119,7 @@ fn dashboard_workload() -> PlanningWorkload { #[test] fn promql_dashboard_materializes_continuous_summary_with_explained_rejections() { let workload = dashboard_workload(); - workload.validate().unwrap(); - - let lowered = lower_promql_workload(&workload, 0) - .expect("valid PromQL workload") - .into_iter() - .next() - .expect("one normalized workload entry"); - let root = Rc::new(lowered); - let strategies = asap_aware_mapping::default_strategies_with(&FullyCostedRuntime); - let space = search_workload_with(vec![("dashboard", Rc::clone(&root))], &strategies); - let target = Rc::clone(&space.roots[0].1); - let capabilities = SummaryMaintenanceLifecycleCapabilities { - supports_ephemeral: true, - supports_prepared: false, - supports_shared: false, - supports_continuously_maintained: true, - }; - - let selection = global_selection_with_summary_maintenance_lifecycles( - &space, - WorkloadDemand { - workload: &workload.query_workload, - data_workload: workload.data_workload.as_ref(), - entry_indices: &[1], - }, - NOW_MS, - Some(Horizon(100.0)), - capabilities, - &FullyCostedRuntime, - ) - .unwrap(); - let plan = assemble_selected_dag_with_summary_maintenance_lifecycles( - &selection, - &target, - WorkloadDemand::new_with_data( - &workload.query_workload, - workload.data_workload.as_ref().unwrap(), - &[1], - ), - NOW_MS, - Some(Horizon(100.0)), - capabilities, - &FullyCostedRuntime, - ) - .unwrap() - .expect("selected summary plan"); + let plan = selected_plan(&workload); assert!(!plan.selected_raw_recompute); assert_eq!(plan.expected_reads, Some(100.0)); @@ -231,3 +186,706 @@ fn promql_dashboard_materializes_continuous_summary_with_explained_rejections() "continuously_maintained" ); } + +fn selected_plan( + workload: &PlanningWorkload, +) -> asap_aware_mapping::SummaryMaintenanceLifecyclePlan { + selected_plan_with_model(workload, &FullyCostedRuntime) +} + +fn selected_plan_with_model( + workload: &PlanningWorkload, + model: &dyn CostModel, +) -> asap_aware_mapping::SummaryMaintenanceLifecyclePlan { + selected_plan_with_horizon(workload, model, Horizon(100.)) +} + +fn selected_plan_with_horizon( + workload: &PlanningWorkload, + model: &dyn CostModel, + horizon: Horizon, +) -> asap_aware_mapping::SummaryMaintenanceLifecyclePlan { + workload.validate().unwrap(); + + let lowered = lower_promql_workload(workload, 0) + .expect("valid PromQL workload") + .into_iter() + .next() + .expect("one normalized workload entry"); + let root = Rc::new(lowered); + let strategies = asap_aware_mapping::default_strategies_with(model); + let space = search_workload_with(vec![("dashboard", Rc::clone(&root))], &strategies); + let target = Rc::clone(&space.roots[0].1); + let capabilities = SummaryMaintenanceLifecycleCapabilities { + supports_ephemeral: true, + supports_prepared: false, + supports_shared: false, + supports_continuously_maintained: true, + }; + + let selection = global_selection_with_summary_maintenance_lifecycles( + &space, + WorkloadDemand { + workload: &workload.query_workload, + data_workload: workload.data_workload.as_ref(), + entry_indices: &[1], + }, + NOW_MS, + Some(horizon), + capabilities, + model, + ) + .unwrap(); + assemble_selected_dag_with_summary_maintenance_lifecycles( + &selection, + &target, + WorkloadDemand::new_with_data( + &workload.query_workload, + workload.data_workload.as_ref().unwrap(), + &[1], + ), + NOW_MS, + Some(horizon), + capabilities, + model, + ) + .unwrap() + .expect("selected summary plan") +} + +mod physical_common; + +/// A selected continuous lifecycle supplies a materialization boundary; its +/// maintenance and query DAGs execute the selected KLL computation in fresh runs. +#[test] +fn continuous_lifecycle_compiles_and_executes_spatial_kll() { + use asap_physical_operators::{ + physical_planner::{compile_candidate, InputContract}, + runtime::Scope, + values::{Batch, Value}, + }; + use asap_types::{ + post_asap::{compile_executable_dag, ExecutableOperatorPayload, SummaryFamilyType}, + pre_asap::DataType, + }; + use std::{collections::BTreeMap, sync::Arc}; + let mut workload = dashboard_workload(); + workload.query_workload.query_batch.as_mut().unwrap()[0].query = + Query("quantile(0.99, latency)".into()); + workload.query_workload.repeating_queries.as_mut().unwrap()[0].query = + Query("quantile(0.99, latency)".into()); + let selected = selected_plan(&workload); + assert_eq!( + selected.deployments[0] + .summary_maintenance_lifecycle_guarantee + .as_ref() + .unwrap() + .summary_maintenance_lifecycle, + SummaryMaintenanceLifecycle::ContinuouslyMaintained + ); + let dag = compile_executable_dag(&selected.root).unwrap(); + let build = dag + .nodes + .iter() + .find(|node| matches!(node.payload, ExecutableOperatorPayload::SummaryAgg { .. })) + .unwrap(); + let input = dag + .edges + .iter() + .find(|edge| edge.consumer == build.id) + .unwrap() + .producer; + let raw = dag.nodes.iter().find(|node| node.id == input).unwrap(); + let schema = Arc::new(raw.output_schema.clone()); + let candidate = compile_candidate( + &dag, + BTreeMap::from([(u64::from(input.0), InputContract::bounded(schema.clone()))]), + &[u64::from(dag.root.0)], + &[u64::from(build.id.0)], + ) + .unwrap(); + + // A continuous input without a finite pane boundary cannot implement this + // blocking builder. Retain lifecycle ownership in the candidate payload; + // only the legal bounded request candidate reaches workload pricing. + let mut unbounded = InputContract::bounded(schema.clone()); + unbounded.properties.boundedness = asap_physical_operators::plan::Boundedness::Unbounded; + let rejected = compile_candidate( + &dag, + BTreeMap::from([(u64::from(input.0), unbounded)]), + &[u64::from(dag.root.0)], + &[u64::from(build.id.0)], + ); + assert!(rejected.is_err()); + let request = compile_candidate( + &dag, + BTreeMap::from([(u64::from(input.0), InputContract::bounded(schema.clone()))]), + &[u64::from(dag.root.0)], + &[], + ) + .unwrap(); + let mut priced = 0; + let feedback = asap_physical_operators::physical_planner::select_candidate( + vec![ + rejected.map(|candidate| { + ( + SummaryMaintenanceLifecycle::ContinuouslyMaintained, + candidate, + ) + }), + Ok((SummaryMaintenanceLifecycle::Ephemeral, request)), + ], + |_| { + priced += 1; + Ok(Some( + asap_physical_operators::physical_planner::CandidateCost { + workload_scope: "dashboard".into(), + horizon_seconds: 100., + total_cost: 1000., + }, + )) + }, + ) + .unwrap(); + assert_eq!(priced, 1); + assert_eq!(feedback.candidate.0, SummaryMaintenanceLifecycle::Ephemeral); + for revision in [1, 2] { + let rows = (1..=100) + .map(|value| { + schema + .fields + .iter() + .map(|field| match field.dtype { + SummaryFamilyType::Plain(DataType::Float64) => { + Value::Float64(f64::from(value)) + } + SummaryFamilyType::Plain(DataType::Timestamp) => Value::Timestamp(300_000), + _ => panic!("unexpected field {field:?}"), + }) + .collect() + }) + .collect(); + let raw_batch = Batch::try_new(schema.clone(), rows).unwrap(); + let direct = physical_common::execute( + &feedback.candidate.1.query, + BTreeMap::from([(u64::from(input.0), raw_batch.clone())]), + Scope::Query { + evaluation_time_ms: 300_000, + revision, + }, + ); + let state = physical_common::execute( + candidate.precompute.as_ref().unwrap(), + BTreeMap::from([(u64::from(input.0), raw_batch)]), + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 300_000, + revision, + }, + ); + let result = physical_common::execute( + &candidate.query, + BTreeMap::from([(u64::from(build.id.0), state[0][0].clone())]), + Scope::Query { + evaluation_time_ms: 300_000, + revision, + }, + ); + let values: Vec<_> = result[0] + .iter() + .flat_map(|batch| batch.rows()) + .flat_map(|row| row.iter()) + .filter_map(|value| { + if let Value::Float64(value) = value { + Some(*value) + } else { + None + } + }) + .collect(); + let direct_values: Vec<_> = direct[0] + .iter() + .flat_map(|batch| batch.rows()) + .flat_map(|row| row.iter()) + .filter_map(|value| { + if let Value::Float64(value) = value { + Some(*value) + } else { + None + } + }) + .collect(); + assert_eq!( + values, direct_values, + "maintenance and request candidates preserve the same population" + ); + assert_eq!(values.len(), 1); + assert!( + (98. ..=100.).contains(&values[0]), + "p99 rank must reflect the supplied population" + ); + } +} + +struct SlidingPaneModel; +impl CostModel for SlidingPaneModel { + fn raw_query_recompute_total_cost( + &self, + target: &asap_types::pre_asap::QueryExpr, + reads: f64, + ) -> Option { + let _ = (target, reads); + Some(Cost(100_000.0)) + } + fn rank_candidates( + &self, + intent: &AggIntent, + candidates: &[asap_types::post_asap::SketchAlgorithm], + ) -> Vec { + FullyCostedRuntime.rank_candidates(intent, candidates) + } + fn summary_maintenance_lifecycle_cost_inputs( + &self, + summary: &SummaryNode, + ) -> SummaryMaintenanceLifecycleCostInputs { + let mut costs = FullyCostedRuntime.summary_maintenance_lifecycle_cost_inputs(summary); + // Controlled workload evidence makes repeated raw construction more + // expensive than retaining and updating the same temporal population. + costs.build_cost = Some(Cost(1000.)); + costs + } + fn summary_maintenance_capabilities( + &self, + summary: &SummaryNode, + ) -> SummaryMaintenanceCapabilities { + FullyCostedRuntime.summary_maintenance_capabilities(summary) + } + fn complete_summary_candidate_estimate( + &self, + _root: &SummaryNode, + _target: Option<&asap_types::pre_asap::QueryExpr>, + deployments: &[asap_aware_mapping::cost_model::CostedSummaryDeployment<'_>], + _horizon: Option, + _reads: Option, + _accuracy: &[AccuracyTarget], + ) -> Option { + Some(asap_aware_mapping::CompleteSummaryCandidateEstimate { + cost: Cost( + deployments + .iter() + .map(|deployment| deployment.selected_cost.0) + .sum(), + ), + physical_plan_id: Some("bounded-sliding-pane-evidence".into()), + window_frameworks: deployments + .iter() + .map(|deployment| { + (deployment.guarantee.summary_maintenance_lifecycle + == SummaryMaintenanceLifecycle::ContinuouslyMaintained) + .then_some(asap_types::post_asap::SummaryWindowFramework::Sliding) + }) + .collect(), + window_accuracy_guarantee: None, + }) + } +} + +/// Workload and optimizer-selected lifecycle generate both physical DAGs. +/// No computational operators or graph edges are constructed by this fixture. +#[test] +fn selected_temporal_lifecycle_compiles_panes_and_executes() { + use asap_physical_operators::{ + operators::Operator, + physical_planner::{ + compile_temporal_pane_candidate, InputContract, Source, TemporalEntityIdentity, + TemporalPaneMaintenance, + }, + runtime::{Limits, RunContext, Scope}, + summary_kernels::datasketches_kll::DatasketchesKLLAccumulator, + values::{Batch, Value}, + }; + use asap_types::{ + post_asap::{ + compile_executable_dag, plan_pane_phase, ExecutableOperatorPayload, SummaryFamilyType, + SummaryWindowFramework, + }, + pre_asap::DataType, + workload::TimestampMs, + }; + use futures::{executor::block_on, StreamExt}; + use std::{collections::BTreeMap, sync::Arc}; + for quantile in [0.5, 0.99] { + let mut workload = dashboard_workload(); + let query = Query(format!( + "quantile_over_time({quantile}, latency{{job=\"api\"}}[5m])" + )); + workload.query_workload.query_batch.as_mut().unwrap()[0].query = query.clone(); + workload.query_workload.repeating_queries.as_mut().unwrap()[0].query = query; + + workload + .data_workload + .as_mut() + .unwrap() + .data_ingestion_interval + .value = Some(DurationMs(60_000)); + workload.query_workload.repeating_queries.as_mut().unwrap()[0].demand = + RepeatedDemand::FixedIntervalAt { + interval: RepetitionInterval(60_000), + evaluation_phase: TimestampMs(300_000), + }; + let plan = selected_plan_with_horizon(&workload, &SlidingPaneModel, Horizon(1000.)); + assert!(!plan.selected_raw_recompute); + assert_eq!(plan.deployments.len(), 1); + let deployment = &plan.deployments[0]; + assert_eq!( + deployment.selected_window_framework, + Some(SummaryWindowFramework::Sliding) + ); + let dag = compile_executable_dag(&plan.root).unwrap(); + let build = dag + .nodes + .iter() + .find(|node| matches!(node.payload, ExecutableOperatorPayload::SummaryAgg { .. })) + .unwrap(); + assert_eq!(build.id, deployment.post_asap_node_id); + let raw = dag + .nodes + .iter() + .find(|node| matches!(node.payload, ExecutableOperatorPayload::Fallback { .. })) + .unwrap(); + let schema = Arc::new(raw.output_schema.clone()); + let width = workload + .data_workload + .as_ref() + .unwrap() + .data_ingestion_interval + .value + .unwrap() + .0; + let layout = plan_pane_phase( + &workload.query_workload.repeating_queries.as_ref().unwrap()[0].demand, + width, + ) + .unwrap(); + let maintenance = TemporalPaneMaintenance { + summary_node: u64::from(build.id.0), + lifecycle: deployment + .summary_maintenance_lifecycle_guarantee + .clone() + .unwrap(), + framework: deployment.selected_window_framework.clone().unwrap(), + layout, + // The memory source has exactly the declared label columns; a + // schemaless deployment must resolve all entity keys first. + entity_identity: TemporalEntityIdentity::Columns( + schema + .fields + .iter() + .enumerate() + .filter(|(_, field)| field.name == "job") + .map(|(index, _)| index) + .collect(), + ), + }; + let candidate = compile_temporal_pane_candidate( + &dag, + BTreeMap::from([(u64::from(raw.id.0), InputContract::bounded(schema.clone()))]), + &[u64::from(dag.root.0)], + &maintenance, + ) + .unwrap(); + assert_eq!(candidate.window_width_ms, 300_000); + assert_eq!(candidate.pane_inputs.len(), 5); + assert_ne!( + candidate.physical.precompute.as_ref().unwrap().roots()[0], + maintenance.summary_node, + "a one-minute pane is not the logical five-minute summary output" + ); + let source_opens = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let mut stored = Vec::new(); + for pane in 0..6 { + let rows = (0..20) + .flat_map(|sample| { + ["api", "batch"] + .into_iter() + .map(move |entity| (sample, entity)) + }) + .map(|(sample, entity)| { + schema + .fields + .iter() + .map(|field| match field.dtype { + SummaryFamilyType::Plain(DataType::Timestamp) => { + Value::Timestamp(pane * 60_000 + (sample + 1) * 3000) + } + SummaryFamilyType::Plain(DataType::Float64) => Value::Float64( + (pane * 20 + sample) as f64 + + if entity == "batch" { 100_000. } else { 0. }, + ), + SummaryFamilyType::Plain(DataType::Utf8) => Value::Utf8(entity.into()), + _ => panic!("unexpected raw field {field:?}"), + }) + .collect() + }) + .collect(); + let result = physical_common::execute( + candidate.physical.precompute.as_ref().unwrap(), + BTreeMap::from([( + u64::from(raw.id.0), + Batch::try_new(schema.clone(), rows).unwrap(), + )]), + Scope::Ingestion { + window_start_ms: pane * 60_000, + window_end_ms: (pane + 1) * 60_000, + revision: 1, + }, + ); + let batch = &result[0][0]; + let rows = batch + .rows() + .iter() + .map(|row| { + row.iter() + .map(|value| match value { + Value::Summary { family, state } => Value::Summary { + family: family.clone(), + state: Arc::new( + DatasketchesKLLAccumulator::from_msgpack_bytes( + &state.serialize_to_bytes(), + ) + .unwrap(), + ), + }, + value => value.clone(), + }) + .collect() + }) + .collect(); + stored.push(Batch::try_new(batch.schema().clone(), rows).unwrap()); + } + for offset in [0, 1] { + let inputs: BTreeMap<_, _> = candidate + .pane_inputs + .iter() + .enumerate() + .map(|(index, &id)| (id, stored[offset + index].clone())) + .collect(); + let scope = Scope::Query { + evaluation_time_ms: (5 + offset as i64) * 60_000, + revision: 1, + }; + let result = + physical_common::execute(&candidate.physical.query, inputs.clone(), scope.clone()); + let rows = result[0][0].rows(); + assert_eq!(rows.len(), 1); + assert!(rows[0] + .iter() + .any(|value| matches!(value, Value::Utf8(label) if label.as_ref() == "api"))); + let Value::Float64(value) = rows[0][1] else { + panic!("missing p99") + }; + assert!( + (value - ((if quantile == 0.5 { 50 } else { 99 }) + offset * 20) as f64).abs() + <= 1. + ); + assert!( + matches!(rows[0][0], Value::Timestamp(timestamp) if timestamp == (5 + offset as i64) * 60_000) + ); + let sources = |inputs: BTreeMap| -> BTreeMap> { + inputs + .into_iter() + .map(|(id, batch)| { + ( + id, + Box::new(TemporalCountingSource { + operator: Operator::source(batch.schema().clone(), vec![batch]) + .unwrap(), + opens: source_opens.clone(), + }) as Source<'static>, + ) + }) + .collect() + }; + let bound = candidate + .physical + .query + .instantiate(sources(inputs.clone())) + .unwrap(); + let population = block_on( + bound + .execute( + &[candidate.merged_state], + RunContext::new(scope.clone(), Limits::default()).unwrap(), + ) + .unwrap() + .remove(0) + .collect::>(), + ); + let Value::Summary { state, .. } = population[0].as_ref().unwrap().rows()[0] + .iter() + .find(|value| matches!(value, Value::Summary { .. })) + .unwrap() + else { + panic!("missing merged state") + }; + assert_eq!( + state + .as_any() + .downcast_ref::() + .unwrap() + .inner + .count(), + 100 + ); + let mut missing = inputs.clone(); + missing.remove(&candidate.pane_inputs[0]); + assert!(candidate + .physical + .query + .instantiate(sources(missing)) + .is_err()); + let mut duplicate = inputs.clone(); + duplicate.insert(candidate.pane_inputs[1], stored[offset].clone()); + let bad = candidate + .physical + .query + .instantiate(sources(duplicate)) + .unwrap(); + let errors = block_on( + bad.execute( + candidate.physical.query.roots(), + RunContext::new(scope.clone(), Limits::default()).unwrap(), + ) + .unwrap() + .remove(0) + .collect::>(), + ); + assert!( + errors.iter().any(Result::is_err), + "duplicate pane must not be merged twice" + ); + let mut duplicate_entity = inputs.clone(); + let pane = &stored[offset]; + duplicate_entity.insert( + candidate.pane_inputs[0], + Batch::try_new( + pane.schema().clone(), + vec![pane.rows()[0].clone(), pane.rows()[0].clone()], + ) + .unwrap(), + ); + let bad = candidate + .physical + .query + .instantiate(sources(duplicate_entity)) + .unwrap(); + let errors = block_on( + bad.execute( + candidate.physical.query.roots(), + RunContext::new(scope.clone(), Limits::default()).unwrap(), + ) + .unwrap() + .remove(0) + .collect::>(), + ); + assert!( + errors.iter().any(Result::is_err), + "duplicate snapshots within a pane must fail" + ); + let bound = candidate + .physical + .query + .instantiate(sources(inputs)) + .unwrap(); + source_opens.store(0, std::sync::atomic::Ordering::SeqCst); + assert!(bound + .execute( + candidate.physical.query.roots(), + RunContext::new( + Scope::Query { + evaluation_time_ms: 330_000, + revision: 1 + }, + Limits::default() + ) + .unwrap() + ) + .is_err()); + assert_eq!( + source_opens.load(std::sync::atomic::Ordering::SeqCst), + 0, + "invalid phase must fail before opening readers" + ); + } + // A selected framework cannot be silently replaced by physical planning. + let mut wrong_framework = maintenance.clone(); + wrong_framework.framework = SummaryWindowFramework::ExponentialHistogram; + let mut wrong_identity = maintenance.clone(); + wrong_identity.entity_identity = TemporalEntityIdentity::SingleEntity; + let mut unknown_phase = maintenance.clone(); + unknown_phase.layout.pane_origin_ms = None; + let mut partial_panes = maintenance.clone(); + partial_panes.layout.pane_width_ms = 90_000; + let mut wrong_lifecycle = maintenance.clone(); + wrong_lifecycle.lifecycle.summary_maintenance_lifecycle = + SummaryMaintenanceLifecycle::Ephemeral; + for unsupported in [ + wrong_framework, + wrong_identity, + unknown_phase, + partial_panes, + wrong_lifecycle, + ] { + assert!(compile_temporal_pane_candidate( + &dag, + BTreeMap::from([(u64::from(raw.id.0), InputContract::bounded(schema.clone()))]), + &[u64::from(dag.root.0)], + &unsupported + ) + .is_err()); + } + } +} + +struct TemporalCountingSource { + operator: asap_physical_operators::operators::Operator, + opens: std::sync::Arc, +} +impl + asap_physical_operators::plan::PhysicalOperator< + asap_physical_operators::values::Batch, + asap_physical_operators::values::Schema, + > for TemporalCountingSource +{ + fn name(&self) -> &str { + "TemporalCountingSource" + } + fn properties( + &self, + inputs: &[asap_physical_operators::plan::PlanProperties], + ) -> asap_physical_operators::plan::PlanProperties { + self.operator.properties(inputs) + } + fn input_schemas(&self) -> Vec { + self.operator.input_schemas() + } + fn output_schema(&self) -> asap_physical_operators::values::Schema { + self.operator.output_schema() + } + fn output_bytes(&self, batch: &asap_physical_operators::values::Batch) -> usize { + batch.bytes() + } + fn start<'a>( + &'a self, + inputs: Vec< + asap_physical_operators::runtime::Input<'a, asap_physical_operators::values::Batch>, + >, + context: asap_physical_operators::runtime::RunContext, + ) -> Result< + asap_physical_operators::runtime::OutputStream<'a, asap_physical_operators::values::Batch>, + asap_physical_operators::Error, + > { + self.opens.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + self.operator.start(inputs, context) + } +} diff --git a/crates/types/Cargo.toml b/crates/types/Cargo.toml index 6b353657..6cf3c2b3 100644 --- a/crates/types/Cargo.toml +++ b/crates/types/Cargo.toml @@ -17,3 +17,4 @@ edition = "2021" serde = { version = "1", features = ["derive", "rc"] } serde_json = "1" thiserror = "2" +sha2 = "0.10" diff --git a/crates/types/src/post_asap/maintained_population.rs b/crates/types/src/post_asap/maintained_population.rs index 30dacf82..3939c7c1 100644 --- a/crates/types/src/post_asap/maintained_population.rs +++ b/crates/types/src/post_asap/maintained_population.rs @@ -63,7 +63,7 @@ impl CurrentSeriesInput { }; if self.metric.is_empty() || *metric != self.metric - || schema.closed + || (schema.closed && !schema.has_promql_series_identity()) || schema.time_index.is_none() { return false; diff --git a/crates/types/src/post_asap/mod.rs b/crates/types/src/post_asap/mod.rs index edd5b182..00473211 100644 --- a/crates/types/src/post_asap/mod.rs +++ b/crates/types/src/post_asap/mod.rs @@ -35,6 +35,8 @@ pub mod guarantee; pub mod maintained_population; pub mod query_time; pub mod schema; +pub mod semantic_definition; +pub use semantic_definition::{LogicalDatasetIdentity, SummarySemanticFragment}; pub mod sketch; pub mod summary_maintenance; pub mod summary_maintenance_lifecycle; diff --git a/crates/types/src/post_asap/semantic_definition.rs b/crates/types/src/post_asap/semantic_definition.rs new file mode 100644 index 00000000..1337222f --- /dev/null +++ b/crates/types/src/post_asap/semantic_definition.rs @@ -0,0 +1,653 @@ +//! Persistable dependency closure using Planner's typed operation vocabulary. +//! Node hashes are local semantic references, not executable or deployed IDs. +use crate::post_asap::{ + EdgeRole, ExecutableDag, ExecutableOperatorPayload, PostAsapNodeId, SummarySchema, +}; +use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; +use std::collections::{BTreeMap, BTreeSet}; + +/// Stable identity of a logical input dataset, independent of its endpoint. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LogicalDatasetIdentity { + pub namespace: String, + pub dataset: String, +} + +impl LogicalDatasetIdentity { + pub fn validate(&self) -> Result<(), String> { + if self.namespace.trim().is_empty() || self.dataset.trim().is_empty() { + return Err("dataset namespace and identity must be nonempty".into()); + } + Ok(()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SummarySemanticFragment { + pub format_version: u32, + /// Version 1 fragments are unbound structural descriptions. Persisted, + /// dataset-bound descriptions use version 2. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub dataset_identity: Option, + pub output: String, + pub nodes: BTreeMap, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SemanticOperation { + // Wire ownership must be Send + Sync. These values are checked against the + // Planner types on export and on recovery; arbitrary JSON is not accepted. + pub operation: serde_json::Value, + pub output_schema: serde_json::Value, + pub inputs: Vec, + /// The direct input range is supplied by the stored record, not by a query lookback. + #[serde(default, skip_serializing_if = "std::ops::Not::not")] + pub record_range: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SemanticInput { + pub role: EdgeRole, + pub node: String, +} + +pub(crate) fn canonical_bytes(value: &impl Serialize) -> Result, String> { + fn canonical(value: serde_json::Value) -> serde_json::Value { + match value { + serde_json::Value::Object(values) => serde_json::Value::Object( + values + .into_iter() + .map(|(k, v)| (k, canonical(v))) + .collect::>() + .into_iter() + .collect(), + ), + serde_json::Value::Array(values) => { + serde_json::Value::Array(values.into_iter().map(canonical).collect()) + } + value => value, + } + } + serde_json::to_vec(&canonical( + serde_json::to_value(value).map_err(|e| e.to_string())?, + )) + .map_err(|e| e.to_string()) +} +fn hash(value: &impl Serialize) -> Result { + Ok(format!("{:x}", Sha256::digest(canonical_bytes(value)?))) +} +fn role(role: EdgeRole) -> u8 { + match role { + EdgeRole::Input => 0, + EdgeRole::Left => 1, + EdgeRole::Right => 2, + } +} + +impl SummarySemanticFragment { + pub fn from_stored_output(dag: &ExecutableDag, output: PostAsapNodeId) -> Result { + Self::export(dag, output, true) + } + + /// All source names in this DAG resolve within this logical dataset. + pub fn from_stored_output_in_dataset( + dag: &ExecutableDag, + output: PostAsapNodeId, + dataset: LogicalDatasetIdentity, + ) -> Result { + dataset.validate()?; + let mut fragment = Self::export(dag, output, true)?; + fragment.format_version = 2; + fragment.dataset_identity = Some(dataset); + fragment.validate()?; + Ok(fragment) + } + + pub fn from_dag(dag: &ExecutableDag, output: PostAsapNodeId) -> Result { + Self::export(dag, output, false) + } + + fn export( + dag: &ExecutableDag, + output: PostAsapNodeId, + parameterize_range: bool, + ) -> Result { + let mut included = BTreeSet::new(); + let mut pending = vec![output]; + while let Some(id) = pending.pop() { + if included.insert(id) { + pending.extend( + dag.edges + .iter() + .filter(|e| e.consumer == id) + .map(|e| e.producer), + ); + } + } + let dag = ExecutableDag { + nodes: dag + .nodes + .iter() + .filter(|n| included.contains(&n.id)) + .cloned() + .collect(), + edges: dag + .edges + .iter() + .filter(|e| included.contains(&e.consumer)) + .cloned() + .collect(), + root: output, + }; + dag.validate().map_err(|e| e.to_string())?; + if dag.nodes.len() > 4096 { + return Err("semantic fragment exceeds node budget".into()); + } + // Open PromQL entities carry all labels. Nullable label columns demanded + // only by a downstream consumer do not change a per-entity scalar state. + let mut dag = dag; + let sample_only = matches!( + &dag.nodes + .iter() + .find(|n| n.id == output) + .ok_or("missing output")? + .payload, + ExecutableOperatorPayload::SummaryAgg { + reduction: crate::pre_asap::Reduction::PerEntity, + grouping: crate::post_asap::GroupingStrategy::PerSubpopulationInstance, + input: crate::post_asap::SummaryUpdate { + item: None, + weight: crate::post_asap::SummaryInputExpr::Column( + crate::pre_asap::ColumnRef::SampleValue + ), + .. + }, + .. + } + ); + if parameterize_range && sample_only { + let direct: BTreeSet<_> = dag + .edges + .iter() + .filter(|e| e.consumer == output) + .map(|e| e.producer) + .collect(); + let mut normalized = false; + for node in &mut dag.nodes { + if !direct.contains(&node.id) { + continue; + } + if let ExecutableOperatorPayload::Fallback { expression } = &mut node.payload { + let source = match expression { + crate::pre_asap::QueryExpr::TimeRange { child, .. } => { + std::rc::Rc::make_mut(child) + } + other => other, + }; + if let crate::pre_asap::QueryExpr::Scan { + source: crate::pre_asap::Source::TimeSeries { .. }, + predicates, + schema, + } = source + { + if !schema.closed + && predicates.is_empty() + && schema.unique_keys.is_empty() + && schema + .columns + .iter() + .take_while(|c| { + !(c.nullable && c.dtype == crate::pre_asap::DataType::Utf8) + }) + .count() + + schema + .columns + .iter() + .rev() + .take_while(|c| { + c.nullable && c.dtype == crate::pre_asap::DataType::Utf8 + }) + .count() + == schema.columns.len() + { + schema.columns.retain(|c| { + !(c.nullable && c.dtype == crate::pre_asap::DataType::Utf8) + }); + node.output_schema.fields.retain(|c| { + !(c.nullable + && c.dtype + == crate::post_asap::SummaryFamilyType::Plain( + crate::pre_asap::DataType::Utf8, + )) + }); + normalized = true; + } + } + } + } + if normalized { + dag.nodes + .iter_mut() + .find(|n| n.id == output) + .unwrap() + .output_schema + .fields + .retain(|c| { + !(c.nullable + && c.dtype + == crate::post_asap::SummaryFamilyType::Plain( + crate::pre_asap::DataType::Utf8, + )) + }); + } + } + let nodes: BTreeMap<_, _> = dag.nodes.iter().map(|n| (n.id, n)).collect(); + let mut hashes: BTreeMap = BTreeMap::new(); + let mut result = Self { + format_version: 1, + dataset_identity: None, + output: String::new(), + nodes: BTreeMap::new(), + }; + let mut stack = vec![(output, false)]; + while let Some((id, finish)) = stack.pop() { + if hashes.contains_key(&id) { + continue; + } + let node = nodes.get(&id).ok_or("missing semantic output")?; + let edges: Vec<_> = dag.edges.iter().filter(|e| e.consumer == id).collect(); + if !finish { + stack.push((id, true)); + for edge in &edges { + stack.push((edge.producer, false)); + } + continue; + } + let mut inputs = edges + .iter() + .map(|e| SemanticInput { + role: e.role, + node: hashes[&e.producer].clone(), + }) + .collect::>(); + inputs.sort_by(|a, b| (role(a.role), &a.node).cmp(&(role(b.role), &b.node))); + let mut payload = node.payload.clone(); + let mut record_range = false; + if parameterize_range + && matches!( + nodes[&output].payload, + ExecutableOperatorPayload::SummaryAgg { .. } + ) + && dag + .edges + .iter() + .any(|e| e.consumer == output && e.producer == id) + { + if let ExecutableOperatorPayload::Fallback { + expression: crate::pre_asap::QueryExpr::TimeRange { child, .. }, + } = &payload + { + payload = ExecutableOperatorPayload::Fallback { + expression: child.as_ref().clone(), + }; + record_range = true; + } + } + if let ExecutableOperatorPayload::RelationalJoin { pruning, .. } = &mut payload { + *pruning = None; + } + let operation = SemanticOperation { + record_range, + operation: serde_json::to_value(&payload).map_err(|e| e.to_string())?, + output_schema: serde_json::to_value(&node.output_schema) + .map_err(|e| e.to_string())?, + inputs, + }; + let key = hash(&operation)?; + result.nodes.insert(key.clone(), operation); + hashes.insert(id, key); + } + result.output = hashes + .remove(&output) + .ok_or("missing semantic output hash")?; + result.validate()?; + Ok(result) + } + + pub fn validate(&self) -> Result<(), String> { + match (&self.dataset_identity, self.format_version) { + (None, 1) => (), + (Some(dataset), 2) => dataset.validate()?, + _ => return Err("semantic version and dataset binding disagree".into()), + } + if !matches!(self.format_version, 1 | 2) + || self.nodes.is_empty() + || self.nodes.len() > 4096 + || canonical_bytes(self)?.len() > 4 * 1024 * 1024 + { + return Err("unsupported semantic fragment version or size".into()); + } + for (key, node) in &self.nodes { + let payload: ExecutableOperatorPayload = + serde_json::from_value(node.operation.clone()).map_err(|e| e.to_string())?; + if node.record_range { + let root = self + .nodes + .get(&self.output) + .ok_or("missing semantic root")?; + let root_payload: ExecutableOperatorPayload = + serde_json::from_value(root.operation.clone()).map_err(|e| e.to_string())?; + if !matches!(payload, ExecutableOperatorPayload::Fallback { .. }) + || !matches!(root_payload, ExecutableOperatorPayload::SummaryAgg { .. }) + || !root.inputs.iter().any(|input| &input.node == key) + { + return Err("record range must belong to a direct summary input".into()); + } + } + let _: SummarySchema = + serde_json::from_value(node.output_schema.clone()).map_err(|e| e.to_string())?; + if hash(node)? != *key + || node + .inputs + .iter() + .any(|i| !self.nodes.contains_key(&i.node)) + { + return Err("semantic fragment hash or dependency mismatch".into()); + } + if node + .inputs + .windows(2) + .any(|p| (role(p[0].role), &p[0].node) > (role(p[1].role), &p[1].node)) + { + return Err("noncanonical semantic input order".into()); + } + } + let mut seen = BTreeSet::new(); + let mut stack = vec![self.output.as_str()]; + while let Some(id) = stack.pop() { + let node = self.nodes.get(id).ok_or("missing semantic fragment root")?; + if seen.insert(id) { + stack.extend(node.inputs.iter().map(|i| i.node.as_str())); + } + } + if seen.len() != self.nodes.len() { + return Err("unrelated semantic fragment nodes".into()); + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::post_asap::{compile_executable_dag, SummaryExpr, SummaryNode}; + use crate::pre_asap::{Column, DataType, QueryExpr, Schema, Source}; + use std::rc::Rc; + + fn fixture(metric: &str) -> ExecutableDag { + let scan = QueryExpr::Scan { + source: Source::TimeSeries { + metric: metric.into(), + }, + predicates: vec![], + schema: Schema::new(vec![Column::new("value", DataType::Float64, false)]), + }; + let schema = SummarySchema { + fields: vec![crate::post_asap::SummaryField { + name: "value".into(), + dtype: crate::post_asap::SummaryFamilyType::Plain(DataType::Float64), + nullable: false, + }], + time_index: None, + }; + compile_executable_dag(&Rc::new(SummaryNode { + expr: SummaryExpr::KeepPreAsap(Rc::new(scan)), + schema, + guarantee: None, + })) + .unwrap() + } + + // Equal source names in different datasets must not alias persisted meaning. + #[test] + fn dataset_identity_is_semantic_and_roundtrips() { + let dag = fixture("latency"); + let export = |namespace: &str| { + SummarySemanticFragment::from_stored_output_in_dataset( + &dag, + dag.root, + LogicalDatasetIdentity { + namespace: namespace.into(), + dataset: "requests".into(), + }, + ) + .unwrap() + }; + let a = export("tenant-a"); + assert_ne!( + canonical_bytes(&a).unwrap(), + canonical_bytes(&export("tenant-b")).unwrap() + ); + assert_eq!(a, export("tenant-a")); // No endpoint enters the semantic API. + let restored: SummarySemanticFragment = + serde_json::from_slice(&canonical_bytes(&a).unwrap()).unwrap(); + restored.validate().unwrap(); + assert_eq!(a, restored); + let mut bad = a.clone(); + bad.dataset_identity.as_mut().unwrap().namespace.clear(); + assert!(bad.validate().is_err()); + bad = a; + bad.dataset_identity = None; + assert!(bad.validate().is_err()); + } + + // Storage identity must ignore temporary identifiers and execution placement. + #[test] + fn identity_ignores_node_ids_and_phase() { + let dag = fixture("latency"); + let expected = SummarySemanticFragment::from_dag(&dag, dag.root).unwrap(); + let mut other = dag.clone(); + other.root = PostAsapNodeId(71); + other.nodes[0].id = other.root; + other.nodes[0].output_state.timing = crate::post_asap::ExecutionTiming::IngestionTime; + assert_eq!( + canonical_bytes(&expected).unwrap(), + canonical_bytes(&SummarySemanticFragment::from_dag(&other, other.root).unwrap()) + .unwrap() + ); + } + + // Source identity and supported semantic format survive restart independently. + #[test] + fn semantics_roundtrip_and_reject_unknown_version() { + let a = fixture("latency"); + let b = fixture("bytes"); + let a = SummarySemanticFragment::from_dag(&a, a.root).unwrap(); + let b = SummarySemanticFragment::from_dag(&b, b.root).unwrap(); + assert_ne!(canonical_bytes(&a).unwrap(), canonical_bytes(&b).unwrap()); + let mut restored: SummarySemanticFragment = + serde_json::from_slice(&canonical_bytes(&a).unwrap()).unwrap(); + restored.validate().unwrap(); + restored.format_version += 1; + assert!(restored.validate().is_err()); + } + // A transformed value cannot share state identity with its source column. + #[test] + fn value_expression_is_semantic_and_nonfinite_constants_are_rejected() { + use crate::pre_asap::{ProjectItem, ScalarValue}; + let original = fixture("latency"); + let expected = SummarySemanticFragment::from_dag(&original, original.root).unwrap(); + let mut transformed = original.clone(); + let ExecutableOperatorPayload::Fallback { expression } = &mut transformed.nodes[0].payload + else { + unreachable!() + }; + *expression = QueryExpr::Project { + cols: vec![ProjectItem { + alias: Some("value".into()), + expr: QueryExpr::FunctionCall { + name: "ln".into(), + args: vec![QueryExpr::Column(0)], + }, + }], + qualifier: None, + child: Rc::new(expression.clone()), + }; + let logged = SummarySemanticFragment::from_dag(&transformed, transformed.root).unwrap(); + assert_ne!( + canonical_bytes(&expected).unwrap(), + canonical_bytes(&logged).unwrap() + ); + let ExecutableOperatorPayload::Fallback { + expression: QueryExpr::Project { cols, .. }, + } = &mut transformed.nodes[0].payload + else { + unreachable!() + }; + cols[0].expr = QueryExpr::Literal(ScalarValue::Float64(f64::NAN)); + assert!(SummarySemanticFragment::from_dag(&transformed, transformed.root).is_err()); + } + + // Changing a downstream consumer cannot change the persisted input definition. + #[test] + fn only_output_dependency_closure_is_exported() { + let mut dag = fixture("latency"); + let stored = dag.root; + let mut consumer = dag.nodes[0].clone(); + consumer.id = PostAsapNodeId(9); + consumer.payload = ExecutableOperatorPayload::Value { + operation: crate::post_asap::ValueOperation::Project { + cols: vec![], + qualifier: None, + }, + }; + dag.edges.push(crate::post_asap::ExecutableDagEdge { + producer: stored, + consumer: consumer.id, + role: EdgeRole::Input, + intermediate_schema: dag.nodes[0].output_schema.clone(), + data_state: dag.nodes[0].output_state, + grouping: crate::post_asap::GroupingEdgeCompatibility::NotApplicable, + window: crate::post_asap::WindowEdgeCompatibility::NotApplicable, + }); + dag.root = consumer.id; + dag.nodes.push(consumer); + let before = SummarySemanticFragment::from_dag(&dag, stored).unwrap(); + dag.nodes.reverse(); + let after = SummarySemanticFragment::from_dag(&dag, stored).unwrap(); + assert_eq!(before, after); + assert_eq!(before.nodes.len(), 1); + } + // Query lookback does not become the identity of each stored input pane. + #[test] + fn stored_input_range_is_parameterized_but_logical_range_is_preserved() { + use crate::post_asap::*; + use crate::pre_asap::{ColumnRef, Reduction}; + let make = |seconds| { + let mut dag = fixture("latency"); + let ExecutableOperatorPayload::Fallback { expression } = &mut dag.nodes[0].payload + else { + unreachable!() + }; + *expression = QueryExpr::TimeRange { + range: std::time::Duration::from_secs(seconds), + child: Rc::new(expression.clone()), + }; + let mut output = dag.nodes[0].clone(); + output.id = PostAsapNodeId(1); + let family = SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum); + output.payload = ExecutableOperatorPayload::SummaryAgg { + family: family.clone(), + input: SummaryUpdate { + item: None, + weight: SummaryInputExpr::Column(ColumnRef::SampleValue), + weight_domain: Default::default(), + }, + reduction: Reduction::PerEntity, + grouping: Default::default(), + }; + output.output_schema.fields[0].dtype = family; + output.output_state.primitive = DataPrimitive::SummaryState; + dag.edges.push(ExecutableDagEdge { + producer: dag.root, + consumer: output.id, + role: EdgeRole::Input, + intermediate_schema: dag.nodes[0].output_schema.clone(), + data_state: dag.nodes[0].output_state, + grouping: GroupingEdgeCompatibility::NotApplicable, + window: WindowEdgeCompatibility::NotApplicable, + }); + dag.root = output.id; + dag.nodes.push(output); + dag + }; + let one = make(60); + let five = make(300); + assert_ne!( + SummarySemanticFragment::from_dag(&one, one.root).unwrap(), + SummarySemanticFragment::from_dag(&five, five.root).unwrap() + ); + assert_eq!( + SummarySemanticFragment::from_stored_output(&one, one.root).unwrap(), + SummarySemanticFragment::from_stored_output(&five, five.root).unwrap() + ); + // Open entities retain their full label identity; consumer-demanded + // optional labels do not change the per-entity stored computation. + let mut open = one.clone(); + let ExecutableOperatorPayload::Fallback { + expression: QueryExpr::TimeRange { child, .. }, + } = &mut open.nodes[0].payload + else { + unreachable!() + }; + let QueryExpr::Scan { schema, .. } = Rc::make_mut(child) else { + unreachable!() + }; + schema.closed = false; + let expected = SummarySemanticFragment::from_stored_output(&open, open.root).unwrap(); + let ExecutableOperatorPayload::Fallback { + expression: QueryExpr::TimeRange { child, .. }, + } = &mut open.nodes[0].payload + else { + unreachable!() + }; + let QueryExpr::Scan { schema, .. } = Rc::make_mut(child) else { + unreachable!() + }; + schema + .columns + .push(Column::new("job", DataType::Utf8, true)); + let label = SummaryField { + name: "job".into(), + dtype: SummaryFamilyType::Plain(DataType::Utf8), + nullable: true, + }; + for node in &mut open.nodes { + node.output_schema.fields.push(label.clone()); + } + open.edges[0].intermediate_schema.fields.push(label); + assert_eq!( + expected, + SummarySemanticFragment::from_stored_output(&open, open.root).unwrap() + ); + let mut forged = SummarySemanticFragment::from_stored_output(&one, one.root).unwrap(); + forged.nodes.values_mut().next().unwrap().operation = + serde_json::json!({"kind": "unknown"}); + assert!(forged.validate().is_err()); + } + // Persisted semantic format changes require an explicit migration/version review. + #[test] + fn semantic_format_v1_has_stable_wire_identity() { + let dag = fixture("latency"); + let exported = SummarySemanticFragment::from_dag(&dag, dag.root).unwrap(); + assert_eq!( + exported.output, + "488a0550f37763997397403ae5ed3588dee5d2a09fa7840f4110b775095fe594" + ); + } +} diff --git a/crates/types/src/pre_asap/query_expr.rs b/crates/types/src/pre_asap/query_expr.rs index dcb73bee..dec0107a 100644 --- a/crates/types/src/pre_asap/query_expr.rs +++ b/crates/types/src/pre_asap/query_expr.rs @@ -65,6 +65,8 @@ pub enum QueryExprError { /// used by `Project`'s own `output_schema` arm instead). #[error("a scalar expression has no row schema of its own")] ScalarHasNoRowSchema, + #[error("invalid per-series sample column: {0}")] + InvalidSampleColumn(String), } // ── Leaf / supporting types ─────────────────────────────────────────────────── @@ -1451,12 +1453,23 @@ impl QueryExpr { /// one value per series, so every label column of `input` is preserved and only /// the sample value is replaced — kept named `value` so the PromQL sample-value /// convention (and any outer `SampleValue` reference) still resolves it by name. -fn per_series_reduction_schema(input: &Schema, agg: &AggIntent) -> Schema { - let value_idx = input - .column_id("value") - .or_else(|| (0..input.columns.len()).find(|&i| Some(i) != input.time_index)); +fn per_series_reduction_schema(input: &Schema, agg: &AggIntent) -> Result { + let vi = if let Some(index) = agg.input_cols().first() { + *index + } else { + super::column_resolution::resolve_column_ref(&ColumnRef::SampleValue, input) + .map_err(|error| QueryExprError::InvalidSampleColumn(error.to_string()))? + }; + if !matches!( + input.columns.get(vi).map(|column| &column.dtype), + Some(DataType::Float64 | DataType::Int64) + ) { + return Err(QueryExprError::InvalidSampleColumn(format!( + "column {vi} is not numeric" + ))); + } let mut columns = input.columns.clone(); - if let Some(vi) = value_idx { + { let mut out = agg.output_column(&columns[vi]); out.name = "value".into(); // A per-series range reduction produces a PromQL sample value, which is @@ -1466,14 +1479,14 @@ fn per_series_reduction_schema(input: &Schema, agg: &AggIntent) -> Schema { out.dtype = DataType::Float64; columns[vi] = out; } - Schema { + Ok(Schema { columns, time_index: input.time_index, unique_keys: input.unique_keys.clone(), // Per-series reduction is label-preserving: it inherits its input's // completeness (an open scan stays open; a closed one stays closed). closed: input.closed, - } + }) } /// The output schema of an `Aggregate { reduction, measures }` over `in_schema` — @@ -1500,7 +1513,7 @@ pub fn aggregate_output_schema( 1, "a per-entity reduction is single-aggregate" ); - return Ok(per_series_reduction_schema(in_schema, &measures[0])); + return per_series_reduction_schema(in_schema, &measures[0]); } Reduction::Reduce(by) => by, }; @@ -1562,6 +1575,11 @@ pub fn aggregate_output_schema( .and_then(|id| in_schema.columns.get(*id)) .unwrap_or(&probe); let mut out = intent.output_column(in_col); + // A global extremum emits NULL for an empty input, even if its input + // column is non-nullable. Grouped extrema only emit existing groups. + if by.is_empty() && matches!(intent, AggIntent::Min { .. } | AggIntent::Max { .. }) { + out.nullable = true; + } if let Some((arg, _)) = intent .arg_selector_columns(in_schema) .map_err(QueryExprError::InvalidScalarSignature)? @@ -2281,6 +2299,26 @@ mod tests { assert_eq!(back, s); } + // Nested temporal aggregation must replace the sample, never the grouping label. + #[test] + fn temporal_reduction_of_grouped_sum_preserves_job() { + let input = Schema::new(vec![ + col("job", DataType::Utf8, true), + col("sum", DataType::Float64, false), + ]); + for aggregate in [ + AggIntent::Avg { col: None }, + AggIntent::Avg { col: Some(1) }, + AggIntent::Rate, + ] { + let output = + aggregate_output_schema(&input, &Reduction::PerEntity, &[aggregate], &[]).unwrap(); + assert_eq!(output.columns[0], input.columns[0]); + assert_eq!(output.columns[1].name, "value"); + assert_eq!(output.columns[1].dtype, DataType::Float64); + } + } + #[test] fn per_series_rate_preserves_labels() { // A per-series range reduction (`rate`) is label-preserving: it produces diff --git a/crates/types/src/pre_asap/schema.rs b/crates/types/src/pre_asap/schema.rs index 81b4ae1f..fbcb9d67 100644 --- a/crates/types/src/pre_asap/schema.rs +++ b/crates/types/src/pre_asap/schema.rs @@ -154,7 +154,21 @@ pub struct Schema { pub closed: bool, } +/// Reserved physical row column carrying canonical JSON of a complete PromQL +/// label map. `$` cannot occur in a user PromQL label name. +pub const PROMQL_SERIES_IDENTITY: &str = "$promql_series_identity"; + impl Schema { + pub fn has_promql_series_identity(&self) -> bool { + self.closed + && self.columns.iter().any(|column| { + column.name == PROMQL_SERIES_IDENTITY + && column.dtype == DataType::Utf8 + && !column.nullable + && column.table.is_none() + }) + } + /// Construct a `Schema` from columns alone — no time index, no /// unique-key constraint. Used by `Scan` over a tabular source /// when the catalog supplies no primary-key metadata. diff --git a/docs/design_docs/physical-planning-and-deployment.md b/docs/design_docs/physical-planning-and-deployment.md new file mode 100644 index 00000000..73180984 --- /dev/null +++ b/docs/design_docs/physical-planning-and-deployment.md @@ -0,0 +1,521 @@ +# Physical Planning, Summary Maintenance, and Deployment + +## 1. Architecture + +A Post-ASAP computation is progressively realized through four layers: + +```mermaid +flowchart LR + L["Logical Post-ASAP DAG
What computation?"] + M["Summary Maintenance Lifecycle
How is state maintained?"] + P["Physical DAG(s)
How is it executed?"] + D["Deployment Plan / DAG
How is it instantiated?"] + + L -->|"Summary Maintenance
Candidate Generation"| M + M -->|"Physical Plan
Compiler"| P + P -->|"Deployment Plan
Compiler"| D +``` + +| Layer | Defines | +| --- | --- | +| **Logical Post-ASAP DAG** | Computation semantics | +| **Summary Maintenance Lifecycle** | Build, retention, reuse, and window strategy | +| **Physical DAG(s)** | Supported physical candidates, executable operators and typed input boundaries | +| **Deployment Plan / DAG** | Selected candidate, concrete data/state bindings and operational lifecycle | + +ASAPPlanner owns the first three layers and the shared physical operator +implementation library. Deployment systems such as ASAPQuery and asap-fusion +own deployment compilation and operation. The lifecycle is a planning contract +associated with the logical DAG, not a separate computation IR. + +### Candidate generation and deployment selection + +Planner exposes the supported, semantically legal **physical plan candidates**. +It does not discard a computation family or materialization placement merely +because a deployment-independent cost estimate prefers another candidate. +Logical candidates are an internal search stage, not the deployment handoff. + +```text +Query semantics + accuracy and lifecycle requirements + ↓ Planner +Supported Physical DAG candidates + typed inputs/outputs + requirements + ↓ backend +Binding feasibility + runtime statistics + resource limits + ERP + ↓ backend deployment compiler +Selected PrecomputePlan + QueryPlan + StoredOutputReferences +``` + +Planner owns operators, dependencies, sharing, and each candidate's +materialization frontier. The backend rejects candidates it cannot realize and +prices feasible candidates over a comparable workload and time horizon. It binds +the selected candidate; it does not lower the logical computation again, exchange +operators, or move an operator across the selected frontier. A missing quote is +not a zero-cost implementation. ERP evidence cannot authorize an illegal rewrite. + +The candidate inventory must identify its supported search scope and budget. +If a configured exhaustive enumeration exceeds its budget, planning fails +explicitly instead of selecting from an undisclosed partial inventory. Reports +separate unsupported compilation, deployment infeasibility, missing evidence, +and a feasible candidate that loses on cost. Absence is not a cost comparison. + +For `sum by(job)(rate(m[1m]))`, Rate remains per series before grouped Sum. +When lifecycle requirements permit it, a candidate may finalize Rate and Sum +within a bounded precompute run and persist the grouped value. Another may leave +those operators in the query DAG. Storing a value requires its exact evaluation +window, revision, readiness and serving cadence to match the query contract. + +For instant-vector TopK, CMS/CountSketch with a candidate heap requires explicit +series identity and a supported latest-value input protocol. Appending historical +sample values does not preserve instant-vector semantics. Replacement, rank +decrease, expiry, grouping and the required approximation guarantee must be +validated before admitting that physical candidate. + +This is the target ownership contract. A backend path that still reconstructs +operators from logical candidates has not completed this integration. + +### Input semantics and summary semantics + +`source`, `filter`, `grouping` and `window` describe input-data semantics: +where records originate, which records qualify, how they are grouped and which +time interval applies. They are not a complete description of arbitrary summary +computation. In particular, the same four fields can summarize different value +expressions or produce different states. + +| Concern | Required semantic information | +| --- | --- | +| Input computation | Source identities and schemas, filters, joins/transforms and their order, or a reference to the canonical input sub-DAG | +| Values and grouping | Value expressions, item identities and weights where applicable, group keys and types, and operation-defined null/duplicate handling | +| Time | Time column and interpretation, interval bounds, evaluation alignment, and distinction between query range and maintained panes | +| Summary computation | Exact operation or sketch family, algorithm and parameters, and supported build/merge behavior | +| Output | State versus finalized value, output schema/type, and readout parameters when part of the output computation | + +For example, KLL over `latency_seconds` and KLL over `log(latency_seconds)` differ +even with identical source, filter, grouping and window. Likewise, weighted +frequency state needs both item and weight expressions. More complex inputs +must retain their computation DAG; four descriptive fields cannot replace it. + +The canonical selected computation is authoritative. These categories describe +what must be preserved, not a new flat IR or a second expression language. +Operator-defined behavior should be referenced through its canonical contract, +not independently configured in deployment metadata. Unsupported or unresolved +semantics cannot be treated as compatible. + +Logical planning defines the semantics; physical compilation realizes them as +operators and typed boundaries. Deployment binds concrete readers and state +records that satisfy those requirements. A stored summary definition records or +references the relevant semantics for compatibility checks. Matching a definition +alone does not establish actual window coverage, revision compatibility or +readiness; those require runtime checks. Physical location, encoding, scheduling +and retention are separate execution/deployment contracts. + +The persisted semantic format contains only the dependency closure of the +selected output. It excludes execution timing, temporary node IDs and deployment +bindings. For a stored raw-input aggregate, its direct input interval is the +record's `(start, end]` interval; a consuming query's lookback is not the identity +of each pane. Nested computations retain their own time semantics. Planner exports +this contract through `SummarySemanticFragment`; changing its semantic wire +vocabulary requires an explicit format-version review. + +### Dataset-bound semantic export + +Backend supplies a stable `LogicalDatasetIdentity { namespace, dataset }` before +Planner exports a persisted definition. Source names within the exported DAG are +resolved in that dataset. Tenant A's `KLL(latency)` and tenant B's `KLL(latency)` +therefore differ; relocating the same dataset to another endpoint does not. + +`SummarySemanticFragment::from_stored_output_in_dataset` exports version 2 with +this identity. Version 1 remains an unbound structural description; Backend's new +planning path uses version 2 for persisted outputs and validates the identity +against the installed input binding. Endpoint and replica information do not enter +semantic identity. Definitions do not authorize cross-deployment reads or adoption +of another plan version's state. + +### Running example + +Suppose p50 and p99 are requested over the same latency samples in a five-minute window, +and one Planner candidate uses KLL with `k=200`. Assume query windows align with one-minute +pane boundaries and that the selected parameters satisfy the required guarantees. +Operator names below are illustrative; the example defines the design, not a +claim that the entire deployment integration is implemented. + +The data source identifies where samples come from. Filters, grouping and the +window determine which samples enter each summary. Here `pane_duration: 1m` +means each stored pane covers one minute; the query range is five minutes. +Neither duration specifies how often maintenance runs or how long state is kept. + +The example evolves through the architecture as follows: + +```text +1. Logical Post-ASAP DAG + +raw latency + ↓ +KLLBuild(k=200) + ↓ +KLLMerge + ┌─┴─────┐ + ↓ ↓ + p50 p99 + + │ + │ Summary Maintenance Candidate Generation + ▼ + +2. Summary Maintenance Lifecycle + +KLLBuild(k=200) + strategy = continuously maintain + window = 1-minute panes + reuse = p50 + p99 + query = merge panes covering requested aligned 5 minutes + + │ + │ Physical Plan Compiler + ▼ + +3. Physical DAGs + +Precompute DAG: +RawInput + ↓ +NativeKllBuild(k=200) + ↓ +KllStateOutput + +Query DAG: +InputSlot[5 panes] + ↓ +NativeKllMerge(k=200) + ┌─┴────────┐ + ↓ ↓ + NativeP50 NativeP99 + + │ + │ Deployment Plan Compiler + ▼ + +4. Deployment Plan / DAG + +Precompute: +OTLP latency source + ↓ +run KLL build over each complete 1-minute input pane + ↓ +store as latency-kll-1m/ + +Query: +resolve five latency-kll-1m states + ↓ +execute query Physical DAG + ↓ +return p50 / p99 +``` + +Each stage adds a different class of decision while preserving the preceding +contracts. Here, continuous maintenance means recurring production of pane state; +the bounded build DAG does not itself implement an unbounded streaming window. + +## 2. Logical Post-ASAP DAG → Summary Maintenance Lifecycle + +The **Logical Post-ASAP DAG** defines computation semantics: + +```text +Scan(latency) + ↓ +KLLBuild(k=200) + ↓ +KLLMerge + ┌─┴────────────┐ + ↓ ↓ +Quantile(.5) Quantile(.99) +``` + +It establishes that KLL with `k=200` is used and that the merge is shared by the +two readouts. It does not determine when KLL states are built or retained. + +**Summary Maintenance Candidate Generation** enumerates legal lifecycle choices +using workload demand, window/freshness requirements and supported physical +implementations. Backend selection uses runtime feasibility and cost after +physical compilation. The following example follows one candidate. + +For the running example, assume it selects: + +```text +producer: KLLBuild(k=200) + +strategy: + continuously maintain + +window realization: + 1-minute panes + +query requirement: + combine panes covering the requested aligned 5-minute range + +reuse: + one merged state serves p50 and p99 +``` + +This produces the **Summary Maintenance Lifecycle**. + +The lifecycle specifies how the selected logical summary should be maintained, +but not its concrete operator implementation or storage location. + +Physical feasibility may feed back into selection. For example, if the required +pane-based maintenance cannot be implemented, this lifecycle candidate cannot be +selected. One-minute panes alone also cannot cover an arbitrarily phased query +window; that requires supported boundary handling or a different candidate. + +## 3. Summary Maintenance Lifecycle → Physical DAG + +The **Physical Plan Compiler** consumes both computation semantics and maintenance +requirements: + +```text +Logical Post-ASAP DAG ++ Summary Maintenance Lifecycle ++ physical capabilities + ↓ +Physical Plan Compiler + ↓ +Physical DAG(s) +``` + +For the running example, the lifecycle creates two execution boundaries. + +These two halves are named as `PhysicalCandidate` names them, `precompute` +and `query`. *Maintenance* stays the lifecycle's word (section 2): it covers +how state is built, retained, reused and scheduled. A precompute DAG is the +physical object that a maintenance lifecycle compiles to, so reusing +*maintenance* for it collapses two layers that the crates keep apart: +`asap-aware-mapping::summary_maintenance_*` owns the lifecycle, and +`asap-physical-operators::physical_planner` owns the DAGs. + +### Precompute Physical DAG + +```text +RawInputSlot( + window = 1m, + bounded = true +) + ↓ +NativeKllBuild(k=200) + ↓ +KllStateOutput(k=200) +``` + +This DAG implements construction of each maintained one-minute pane. Its input +contract requires all input samples matching the source, filters and group within that pane; the deployment supplies that +bounded input from its source integration. + +### Query Physical DAG + +```text +InputSlot( + k = 200, + coverage = requested aligned 5m +) + ↓ +NativeKllMerge(k=200) + ┌─┴──────────────────┐ + ↓ ↓ +NativeQuantile(.50) NativeQuantile(.99) +``` + +The Physical Plan Compiler chooses `NativeKllBuild`, `NativeKllMerge`, and the +physical quantile implementations, validates state compatibility, and preserves +the shared merge. It also resolves expressions, schemas, ordered dependencies +and execution properties. + +The resulting Physical DAGs know that compatible KLL states are required, but +do not know where those states are stored. + +For example: + +```text +InputSlot +``` + +is physical, while: + +```text +s3://.../latency-kll/12:01 +``` + +is deployment-specific. Placement and scheduling also remain outside the Physical +DAG. If the required behavior cannot be realized, physical compilation fails. + +### Physical candidates include precompute computation + +Materialization frontiers are Planner decisions. A candidate records both the +precompute Physical DAG and the query Physical DAG, with typed outputs connecting +them. The deployment compiler binds those outputs; it does not move operators. + +For `sum by(job)(rate(m[1m]))`, legal physical candidates can include: + +```text +Candidate A: + precompute: compatible per-series counter states → per-series Rate + materialized output: per-series rate values for window/evaluation/revision + query: stored per-series rate values → grouped Sum + +Candidate B: + precompute: compatible per-series counter states → per-series Rate → grouped Sum + materialized output: grouped values for window/evaluation/revision + query: stored grouped values → result +``` + +Both preserve reset-aware Rate before Sum. Summing raw counters before Rate is +not equivalent. The counter-state build may be another precompute DAG; typed +state inputs do not imply that a deployment can construct or bind those states. + +The shared library exposes `physical_planner::compile_candidates(...)` to lower +explicit frontier candidates to `PhysicalCandidate { precompute, query, +materialized_outputs }`. `select_candidate(...)` accepts deployment feasibility +and scoped complete-workload costs and chooses the lowest-cost feasible +candidate. Costs must describe the same workload and planning horizon; missing +feasibility is rejected before pricing. The optimizer supplies candidate +frontiers and cost evidence, including updates, retention, recurrence and sharing. +`enumerate_frontiers` constructs bounded, reachable antichain frontiers above explicit input boundaries, including query-only and fully precomputed results. It fails explicitly when the candidate budget is exceeded. Maintenance selection must still reject frontiers that violate window, freshness, or reuse requirements; deployment feasibility is checked before pricing. + +Physical compilation opens no readers. Bounded precompute outputs become typed +query inputs. Their source, filters, grouping, build window, evaluation time, readiness and +revision contracts must accompany the selected lifecycle and be checked during +deployment binding. Type compatibility alone does not establish reuse legality. + +The Planner integration test executes both candidates through the shared runtime +and reverses the selected frontier with two controlled cost fixtures. It also +rejects shadowed/duplicate boundaries and incomparable planning horizons. This +establishes Planner capability; it does not establish that ASAPQuery currently +supports persisting every scalar/result-output frontier. + +## 4. Physical DAG → Deployment Plan / DAG + +The **Deployment Plan Compiler** binds the Physical DAGs to the concrete deployment: + +```text +Physical DAGs ++ Summary Maintenance Lifecycle ++ deployment catalog/state ++ sources/materializations ++ operational policy + ↓ +Deployment Plan Compiler + ↓ +Deployment Plan / DAG +``` + +For the precompute DAG, it may produce: + +```text +Source: + RawInputSlot + → complete bounded panes from the OTLP latency source + +Schedule: + each 1-minute pane, once its completion requirements are met + +Execution: + RawInput → NativeKllBuild(k=200) + +Output: + KllStateOutput + → latency-kll-1m/ +``` + +For a query over `(12:00, 12:05]`, its input-binding rule resolves: + +```text +InputSlot[5 panes] + ├── latency-kll-1m/(12:00,12:01] + ├── latency-kll-1m/(12:01,12:02] + ├── latency-kll-1m/(12:02,12:03] + ├── latency-kll-1m/(12:03,12:04] + └── latency-kll-1m/(12:04,12:05] + ↓ + Query Physical DAG + ↓ + p50, p99 +``` + +The Deployment Plan Compiler establishes bindings and checks that their contracts +satisfy the physical inputs and selected lifecycle, including KLL parameters, +source, filters, grouping, window coverage and revision scope. The deployment engine +resolves request-specific states and checks their actual coverage, revisions and +readiness at execution time. A compiled plan cannot establish future readiness. + +The compiler does not replace `NativeKllMerge`, choose another sketch, or decide +to maintain different windows. Such changes require replanning. A Deployment +Plan / DAG is an operational instantiation, not another computation IR. + +## 5. Responsibility Boundary + +The complete example makes the ownership boundary explicit: + +| Stage | KLL example decision | +| --- | --- | +| **Logical Post-ASAP DAG** | Use `KLL(k=200)` with shared merge for p50/p99 | +| **Summary Maintenance Candidate Generation** | Maintain 1-minute panes and reuse them for aligned five-minute queries | +| **Summary Maintenance Lifecycle** | Record pane/window/freshness/reuse requirements | +| **Physical Plan Compiler** | Lower to native KLL build, merge, and readout operators | +| **Physical DAG** | Define precompute and query DAGs with typed input/output boundaries | +| **Deployment Plan Compiler** | Bind raw input and KLL state slots to concrete sources/materializations | +| **Deployment Plan / DAG** | Specify maintenance schedules, stored-pane resolution and query execution | + +```text +Logical: + "Use KLL for p50/p99." + +Lifecycle: + "Maintain reusable 1-minute KLL panes." + +Physical: + "Execute NativeKllBuild and + NativeKllMerge → {p50, p99}." + +Deployment: + "Read OTLP here, store panes here, + and bind these five panes for this aligned query." +``` + +The deployment engine executes the bound Physical DAGs through ASAPPlanner's +shared physical operator implementation library, `asap-physical-operators`, and +its DAG runtime. The merge executes once per run for both consumers. Execution +does not introduce additional planning decisions. + +Each maintained pane contributes its input samples once. A replacement snapshot +replaces that pane's state; query merging must not count both the old and new +snapshots as separate inputs. + +## 6. Executable acceptance coverage + +The tests cover optimizer-selected lifecycle execution and automatic temporal +pane compilation, alongside independent operator/runtime fixtures: + +| Test | Contract exercised | +| --- | --- | +| `summary_maintenance_lifecycle_e2e::selected_temporal_lifecycle_compiles_panes_and_executes` | PromQL p50/p99 workloads → selected continuous lifecycle and Sliding framework → automatically generated precompute/query DAGs → real codec round-trip → adjacent aligned windows; checks filters, entity identity, sample counts, missing/duplicate panes and phase rejection before opening readers | +| `summary_maintenance_lifecycle_e2e::continuous_lifecycle_compiles_and_executes_spatial_kll` | PromQL workload → selected continuous lifecycle → logical DAG → compiled precompute/query candidate → results in independent revisions; an unbounded candidate fails before pricing, and a bounded request candidate summarizes the same input samples | +| `kll_pane_execution::five_panes_roundtrip_and_shared_merge_runs_once` | Explicit one-minute precompute DAGs → real MessagePack state bytes → five required query inputs → shared native merge → p50/p99; counts every sample once, checks adjacent aligned windows and instruments one merge start per run | +| `kll_pane_execution::restored_panes_reject_corruption_parameters_schema_and_missing_binding` | Corrupt bytes, parameter relabelling, incompatible schemas and absent bindings fail explicitly | +| `precompute_candidates::grouped_rate_can_be_materialized_before_or_after_grouped_sum` | Cost changes select different legal precompute frontiers; both selected candidates execute with the same reset-sensitive result; uncompilable candidates are not priced | +| `sql_to_physical::sql_filter_grouped_sum_executes_and_rebinds` | SQL text → candidate search → physical compilation → shared Scan predicates and grouped summary execution; NULL samples are ignored and fresh bindings produce new results | + +`physical_planner::compile_temporal_pane_candidate` consumes the logical DAG, +selected lifecycle/framework and a generic pane/entity input contract. It +generates pane construction, scan predicates, ordered state slots, a shared +merge, quantile readouts and run-scoped timestamps. A physical pane output has +its own identity: one minute of state cannot masquerade as the logical +five-minute summary. The returned candidate retains the maintenance contract. + +This initial realization supports bounded, complete KLL panes with known phase +and resolved entity identity, for Sliding windows or a single Tumbling window. +Source capability evidence must declare all entity keys or isolate one entity; +usage-derived PromQL columns alone cannot establish that identity. Partial edge +panes, exponential histograms and cross-run delta accumulation require further +physical candidates and are rejected by this entry point. + +Physical execution checks pane timestamps and duplicate entity states. Concrete +stored identity, revisions, readiness and complete coverage of required input samples remain +deployment responsibilities. Real storage and HTTP execution belong to +deployment-repository E2E tests. diff --git a/docs/develop_docs/native-promql-inputs.md b/docs/develop_docs/native-promql-inputs.md new file mode 100644 index 00000000..aa1a52b6 --- /dev/null +++ b/docs/develop_docs/native-promql-inputs.md @@ -0,0 +1,55 @@ +# Native PromQL source rows + +Audience: source-adapter and physical-executor developers. + +A PromQL query only names some labels. Those columns cannot establish series +identity for Rate or TopK: two series with the same `job` may have different +unreferenced instance labels. + +`physical_planner::promql_rows::with_series_identity` resolves supported unary +PromQL computations to a bounded row representation before candidate search. +It appends `$promql_series_identity`, a non-null UTF-8 column containing the +canonical JSON encoding of the full label map. The name cannot collide with a +legal PromQL label. The resulting schema is closed over physical columns; the +label map remains dynamic and is not restricted to labels named in the query. +This realization rejects unsupported label rewriting, implicit vector matching, +and `without` operations rather than dropping hidden labels. + +Source adapters construct batches with `series_row`. Named label columns are +projections of the same complete identity; absent named labels project to empty +strings. `decode_series_identity` restores all labels on result conversion and +rejects noncanonical encodings. A query adapter must still apply the selected +operator's metric-name/result-label rules. Source selection, complete window +coverage and revision admission remain deployment responsibilities. + +Planner's maintained-population candidate recognizes this explicit identity +representation. Its TopK readout compiles automatically to `CurrentSeries`, +`Sort`, and `Limit`; deployment supplies the raw boundary or an already maintained +population boundary. Compilation does not open either source. + +The native `CurrentSeries` operator selects the latest sample per complete +identity in `(evaluation_time - lookback, evaluation_time]`. It removes stale +markers after selecting the latest sample, so an older value cannot reappear. +It rejects conflicting values at one series timestamp and emits the evaluation +timestamp. Each run builds a new snapshot; decreased values and expired series +cannot retain earlier heap weights. It reserves workspace and observes the +run's cancellation and byte budget. Precompute scopes must match the declared +lookback before any input is polled. + +CMS/CountSketch heap operators can consume this snapshot. CMS still requires +nonnegative weights; legal approximate TopK admission still requires the +Planner's accuracy/membership evidence. Executing a heap does not establish +that its result satisfies a query's accuracy requirements. + +Tests cover open-label Rate → CMS/CountSketch heaps, hidden-label round trips, +reset and zero-rate cases, snapshot replacement/decrease/expiry/staleness, +serialized physical recovery, and resource rejection. These are shared-library +tests, not proof of Backend candidate selection or durable deployment execution. + +Spatial heap candidates use the same complete series identity. Planner's +`current_series_topk_candidates` explores a CountSketch-with-heap realization +of canonical Sort/Limit under an explicit accuracy target. The physical graph +selects the latest eligible samples before building a fresh heap. A maintained +population boundary can supply that snapshot directly. Arbitrary signed metric +values do not authorize CMS; counter Rate's non-negative proof is separate. +These candidates still require membership/score evidence for deployment admission.