diff --git a/Cargo.lock b/Cargo.lock index 7f0a37554..a3fe727ee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -385,6 +385,7 @@ dependencies = [ "asap-types", "asap_sketchlib 0.3.0 (git+https://github.com/ProjectASAP/asap_sketchlib?rev=5f03ccbd798ed5fec62bdd839bcb331123cab369)", "futures", + "rmp-serde", "serde", "serde_json", "thiserror 2.0.18", diff --git a/crates/asap-physical-operators/Cargo.toml b/crates/asap-physical-operators/Cargo.toml index 74c172e30..cec452e5a 100644 --- a/crates/asap-physical-operators/Cargo.toml +++ b/crates/asap-physical-operators/Cargo.toml @@ -14,5 +14,6 @@ thiserror = "2" [dev-dependencies] +rmp-serde = "1" asap-aware-mapping = { path = "../asap-aware-mapping" } asap-frontend-promql = { path = "../frontend-promql" } diff --git a/crates/asap-physical-operators/src/expressions/mod.rs b/crates/asap-physical-operators/src/expressions/mod.rs index 2916cca75..a16a3bde2 100644 --- a/crates/asap-physical-operators/src/expressions/mod.rs +++ b/crates/asap-physical-operators/src/expressions/mod.rs @@ -23,6 +23,18 @@ pub enum Expression { labels: Vec, without: bool, }, + /// One label of a label map; an absent label reads as empty, as in PromQL. + Label { + column: usize, + name: String, + }, + /// Canonical encoding of a label map less `excluding`, identical to + /// `promql_rows::encode_series_identity`. + LabelIdentity { + column: usize, + #[serde(default, skip_serializing_if = "Vec::is_empty")] + excluding: Vec, + }, Literal { value: Value, dtype: DataType, @@ -119,6 +131,21 @@ impl Expression { } Ok((expected, false)) } + Label { column, .. } | LabelIdentity { column, .. } => { + if plain(input, *column)? + != ( + &DataType::Map { + key: Box::new(DataType::Utf8), + value: Box::new(DataType::Utf8), + value_nullable: false, + }, + false, + ) + { + return Err(invalid("label read requires a non-null Utf8 map")); + } + Ok((DataType::Utf8, false)) + } Column(i) => { let (t, n) = plain(input, *i)?; Ok((t.clone(), n)) @@ -251,6 +278,41 @@ impl Expression { } } Planner(expression) => expression.evaluate(row)?, + Label { column, name } => { + let Value::Map(entries) = &row[*column] else { + return Err(invalid("label read requires a map")); + }; + let mut found = None; + for (key, value) in entries.iter() { + let (Value::Utf8(key), Value::Utf8(value)) = (key, value) else { + return Err(invalid("label read requires Utf8 entries")); + }; + if key.as_ref() == name.as_str() && found.replace(value.clone()).is_some() { + return Err(invalid("duplicate label name")); + } + } + Value::Utf8(found.unwrap_or_else(|| "".into())) + } + LabelIdentity { column, excluding } => { + let Value::Map(entries) = &row[*column] else { + return Err(invalid("label identity requires a map")); + }; + let mut labels = std::collections::BTreeMap::new(); + for (key, value) in entries.iter() { + let (Value::Utf8(key), Value::Utf8(value)) = (key, value) else { + return Err(invalid("label identity requires Utf8 entries")); + }; + if excluding.iter().any(|label| label.as_str() == key.as_ref()) { + continue; + } + if labels.insert(key.to_string(), value.to_string()).is_some() { + return Err(invalid("duplicate label name")); + } + } + Value::Utf8( + crate::physical_planner::promql_rows::encode_series_identity(&labels)?.into(), + ) + } Column(i) => row[*i].clone(), Literal { value, .. } => value.clone(), Negate(v) => match v.evaluate(row)? { diff --git a/crates/asap-physical-operators/src/physical_planner/precompute.rs b/crates/asap-physical-operators/src/physical_planner/precompute.rs index 063bd0b82..4fa40c682 100644 --- a/crates/asap-physical-operators/src/physical_planner/precompute.rs +++ b/crates/asap-physical-operators/src/physical_planner/precompute.rs @@ -1,4 +1,5 @@ //! Compile immutable summary-input computation with explicit population and pane identity. +use super::promql_rows::SERIES_IDENTITY_COLUMN as SERIES_IDENTITY; use super::*; use planner_types::{ post_asap::{ExecutionTiming, GroupingStrategy, SummarySchema}, @@ -34,6 +35,93 @@ pub fn population_schema(family: SummaryFamilyType) -> Schema { }) } +/// Raw sample rows at a precompute boundary. `$population` holds the series' +/// complete label set, so it is the complete source identity of per-series +/// summaries; `$timestamp` is the sample time and `value` a finite sample +/// (stale markers are not samples). Rows are what the boundary's source scan +/// selected; the deployment decides which rows and panes they are. Label sets +/// must be canonical (sorted, unique, no empty values), since they are the +/// population identity: build rows with [`raw_sample_row`]. +pub fn raw_sample_schema() -> Schema { + let mut schema = (*population_schema(SummaryFamilyType::Plain(DataType::Float64))).clone(); + schema.fields[1].name = "$timestamp".into(); + Arc::new(schema) +} + +/// A raw sample row whose label set is sorted, unique and omits empty values, +/// so one series always has one population identity. +pub fn raw_sample_row( + labels: &BTreeMap, + timestamp_ms: i64, + value: f64, +) -> Vec { + use crate::values::Value; + vec![ + Value::Map( + labels + .iter() + .filter(|(_, v)| !v.is_empty()) + .map(|(k, v)| { + ( + Value::Utf8(k.as_str().into()), + Value::Utf8(v.as_str().into()), + ) + }) + .collect::>() + .into(), + ), + Value::Timestamp(timestamp_ms), + Value::Float64(value), + ] +} + +/// Input contract of a precompute boundary: raw sample rows for a raw time +/// series scan, otherwise the stored population of its summary state. +pub fn boundary_schema(node: &PostAsapDagNode) -> Result { + let Payload::Fallback { expression } = &node.payload else { + return source_schema(&node.output_schema); + }; + let scan = match expression { + planner_types::pre_asap::QueryExpr::TimeRange { child, .. } => child.as_ref(), + expression => expression, + }; + if !matches!( + scan, + planner_types::pre_asap::QueryExpr::Scan { + source: planner_types::pre_asap::Source::TimeSeries { .. }, + .. + } + ) { + return source_schema(&node.output_schema); + } + let logical = &node.output_schema; + // Labels may be absent from a series; its label map then omits them. + let valid = logical + .fields + .iter() + .enumerate() + .all(|(i, field)| match &field.dtype { + SummaryFamilyType::Plain(DataType::Timestamp) => { + Some(i) == logical.time_index && !field.nullable + } + SummaryFamilyType::Plain(DataType::Float64) => field.name == "value" && !field.nullable, + SummaryFamilyType::Plain(DataType::Utf8) => true, + _ => false, + }) + && !logical + .fields + .iter() + .any(|f| f.name.starts_with('$') && f.name != SERIES_IDENTITY) + && logical.time_index.is_some() + && logical.fields.iter().filter(|f| f.name == "value").count() == 1; + if !valid { + return Err(invalid( + "raw sample boundary requires labels, a timestamp and one Float64 value", + )); + } + Ok(raw_sample_schema()) +} + /// Validate the adapter layout during installed-plan recovery without lowering operators. pub fn source_schema(logical: &SummarySchema) -> Result { let states = logical @@ -132,7 +220,7 @@ pub fn compile( for id in ordered { let node = nodes[&id]; if frontier.contains(&id) { - let schema = source_schema(&node.output_schema)?; + let schema = boundary_schema(node)?; sources.insert(id, InputContract::bounded(schema.clone())); outputs.insert(id, schema); continue; @@ -281,15 +369,48 @@ fn fragment( let [input] = schemas else { return Err(invalid("summary update requires one input")); }; - if update.item.is_some() - || !matches!(grouping, GroupingStrategy::PerSubpopulationInstance) - { + // Item identities resolve against the complete label set of raw + // samples; finalized readouts carry no such identity. + let raw = *input == raw_sample_schema(); + // A unit-frequency summary (HLL) observes each raw sample value. + let unit_frequency = raw + && crate::capability::is_unit_sample_frequency(update) + && matches!(family, SummaryFamilyType::Sketch(kind, _) if !matches!( + kind.algorithm(), + planner_types::post_asap::SketchAlgorithm::Cms + | planner_types::post_asap::SketchAlgorithm::CountSketch + | planner_types::post_asap::SketchAlgorithm::CmsWithHeap + | planner_types::post_asap::SketchAlgorithm::CountSketchWithHeap + )); + let keyed = update.item.is_some() && !unit_frequency; + if (keyed && !raw) || !matches!(grouping, GroupingStrategy::PerSubpopulationInstance) { return Err(invalid( "precompute keyed/shared update needs its dedicated physical candidate", )); } crate::capability::validate_summary_kernel(family, update, grouping) .map_err(Error::Invalid)?; + if raw + && matches!( + update.weight_domain, + planner_types::post_asap::WeightDomain::NonNegative { + proof: planner_types::post_asap::NonNegativeWeightProof::ResetAwareCounterDerivative + } + ) + { + return Err(invalid( + "a counter-derivative weight cannot be read from raw cumulative samples", + )); + } + if keyed + && 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")); + } let labels = match reduction { PlannerReduction::PerEntity => Expression::Column(0), PlannerReduction::Reduce(keys) => Expression::LabelSet { @@ -302,8 +423,11 @@ fn fragment( .output_schema .fields .get(*key) + // A raw label map omits absent labels; the + // series identity is not one of its labels. .filter(|field| { - !field.nullable + (raw || !field.nullable) + && field.name != SERIES_IDENTITY && field.dtype == SummaryFamilyType::Plain(DataType::Utf8) }) .map(|f| f.name.clone()) @@ -316,6 +440,7 @@ fn fragment( }, }; let weight = match &update.weight { + _ if unit_frequency => Expression::Column(2), SummaryInputExpr::Constant(value) => Expression::Literal { value: crate::values::Value::Float64(*value), dtype: DataType::Float64, @@ -334,20 +459,45 @@ fn fragment( )) } }; - let project = Operator::project( - input.clone(), - vec![ - ("$population".into(), labels), - ("$window_end".into(), Expression::Column(1)), - ("value".into(), Expression::FiniteFloat64(Box::new(weight))), - ], - )? - .with_output_schema(population_schema(SummaryFamilyType::Plain( - DataType::Float64, - )))?; + let mut columns = vec![ + ("$population".into(), labels), + ("$window_end".into(), Expression::Column(1)), + ("value".into(), Expression::FiniteFloat64(Box::new(weight))), + ]; + let mut fields = population_schema(SummaryFamilyType::Plain(DataType::Float64)) + .fields + .clone(); + if keyed { + let mut items = Vec::new(); + raw_items( + update.item.as_ref().expect("keyed item"), + &parents[0].output_schema, + &mut items, + )?; + for (index, (expression, dtype)) in items.into_iter().enumerate() { + let name = format!("$item{index}"); + fields.push(planner_types::post_asap::SummaryField { + name: name.clone(), + dtype: SummaryFamilyType::Plain(dtype), + nullable: false, + }); + columns.push((name, expression)); + } + } + let item_columns = (3..fields.len()).collect::>(); + let project = Operator::project(input.clone(), columns)?.with_output_schema( + Arc::new(SummarySchema { + fields, + time_index: Some(1), + }), + )?; let projected = project.schema(); let project = add(vec![0], project)?; - let build = Operator::summary_build(projected, family.clone(), 2, Some(1), vec![0])?; + let build = if keyed { + Operator::keyed_summary_build(projected, family.clone(), 2, item_columns, vec![0])? + } else { + Operator::summary_build(projected, family.clone(), 2, Some(1), vec![0])? + }; let built = build.schema(); let build = add(vec![project], build)?; add( @@ -382,3 +532,78 @@ fn fragment( }; CompiledPhysicalDag::from_operators(sources, operators, vec![root]) } + +/// Resolve keyed item identities over raw sample rows: labels (absent labels +/// read as empty, as in PromQL), the sample value, or the canonical encoding +/// of the label set less excluded labels. +fn raw_items( + expr: &SummaryInputExpr, + scan: &SummarySchema, + items: &mut Vec<(Expression, DataType)>, +) -> Result<(), Error> { + // Open PromQL scans need not list every label, so any name that is not + // another scan column (value, time, series identity) reads as a label. + let label = |column: &ColumnRef| match column { + ColumnRef::Named(name) | ColumnRef::Qualified { name, .. } + if !name.starts_with('$') + && scan.fields.iter().all(|f| { + &f.name != name || f.dtype == SummaryFamilyType::Plain(DataType::Utf8) + }) => + { + Some(name.clone()) + } + _ => None, + }; + let identity = |excluding: Vec| { + ( + Expression::LabelIdentity { + column: 0, + excluding, + }, + DataType::Utf8, + ) + }; + match expr { + SummaryInputExpr::Column(ColumnRef::SampleValue) => { + items.push((Expression::Column(2), DataType::Float64)) + } + SummaryInputExpr::Column(ColumnRef::Named(name) | ColumnRef::Qualified { name, .. }) + if name == "value" => + { + items.push((Expression::Column(2), DataType::Float64)) + } + SummaryInputExpr::Column(ColumnRef::Named(name) | ColumnRef::Qualified { name, .. }) + if name == SERIES_IDENTITY => + { + items.push(identity(vec![])) + } + SummaryInputExpr::Column(column) if label(column).is_some() => items.push(( + Expression::Label { + column: 0, + name: label(column).expect("resolved label"), + }, + DataType::Utf8, + )), + SummaryInputExpr::EntityIdentity( + planner_types::post_asap::EntityIdentity::PromqlLabelSet { excluding }, + ) => items.push(identity( + excluding + .iter() + .map(|column| { + label(column).ok_or_else(|| invalid("excluded identity label is not a label")) + }) + .collect::>()?, + )), + SummaryInputExpr::Tuple(parts) if !parts.is_empty() => { + for part in parts { + raw_items(part, scan, items)?; + } + } + _ => { + return Err(invalid( + "keyed summary item does not resolve over raw samples", + )) + } + } + Ok(()) +} 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 index 5a78fc489..bb0906911 100644 --- a/crates/asap-physical-operators/src/summary_kernels/count_min_sketch.rs +++ b/crates/asap-physical-operators/src/summary_kernels/count_min_sketch.rs @@ -28,6 +28,9 @@ impl AggregateCore for CountMinSketchAccumulator { 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, KernelError> { let other = other 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 index da07aeeeb..f2d8786c2 100644 --- 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 @@ -80,6 +80,9 @@ impl AggregateCore for CountMinSketchWithHeapAccumulator { 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, diff --git a/crates/asap-physical-operators/src/summary_kernels/count_sketch.rs b/crates/asap-physical-operators/src/summary_kernels/count_sketch.rs index 2728b5f29..c3aaa40c9 100644 --- a/crates/asap-physical-operators/src/summary_kernels/count_sketch.rs +++ b/crates/asap-physical-operators/src/summary_kernels/count_sketch.rs @@ -34,6 +34,9 @@ impl AggregateCore for CountSketchAccumulator { 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, 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 index a735e4bd9..9f15e5042 100644 --- 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 @@ -68,6 +68,9 @@ impl AggregateCore for CountSketchWithHeapAccumulator { 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, diff --git a/crates/asap-physical-operators/src/summary_kernels/datasketches_kll.rs b/crates/asap-physical-operators/src/summary_kernels/datasketches_kll.rs index 0f740a9dc..3417bb8cc 100644 --- a/crates/asap-physical-operators/src/summary_kernels/datasketches_kll.rs +++ b/crates/asap-physical-operators/src/summary_kernels/datasketches_kll.rs @@ -46,6 +46,9 @@ impl AggregateCore for 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, KernelError> { let other = other diff --git a/crates/asap-physical-operators/src/summary_kernels/dd_sketch.rs b/crates/asap-physical-operators/src/summary_kernels/dd_sketch.rs index 1d6be98cf..dcda7ad47 100644 --- a/crates/asap-physical-operators/src/summary_kernels/dd_sketch.rs +++ b/crates/asap-physical-operators/src/summary_kernels/dd_sketch.rs @@ -24,6 +24,9 @@ impl AggregateCore for DDSketchAccumulator { 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, KernelError> { let other = other diff --git a/crates/asap-physical-operators/src/summary_kernels/exact.rs b/crates/asap-physical-operators/src/summary_kernels/exact.rs index cb9d60995..a211b7f32 100644 --- a/crates/asap-physical-operators/src/summary_kernels/exact.rs +++ b/crates/asap-physical-operators/src/summary_kernels/exact.rs @@ -19,13 +19,50 @@ enum ScalarState { /// Both the family and population layout survive persistence. Sharing counter /// arithmetic never authorizes a Rate state to answer an Increase readout. +/// +/// Deserialization validates the payload against its declared family, so +/// deployments can persist this state with any serde format without mirroring +/// its shape. #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(try_from = "ExactPayload")] pub struct ExactAccumulator { family: SummaryFamilyType, scalar: ScalarState, keyed: Option>, } +#[derive(Deserialize)] +struct ExactPayload { + family: SummaryFamilyType, + scalar: ScalarState, + keyed: Option>, +} + +impl TryFrom for ExactAccumulator { + type Error = String; + + fn try_from(payload: ExactPayload) -> Result { + let empty = Self::new(payload.family, payload.keyed.is_some())?; + let same_kind = |state: &ScalarState| { + std::mem::discriminant(state) == std::mem::discriminant(&empty.scalar) + }; + if !same_kind(&payload.scalar) + || payload + .keyed + .iter() + .flat_map(HashMap::values) + .any(|s| !same_kind(s)) + { + return Err("exact payload differs from its declared Planner family".into()); + } + Ok(Self { + scalar: payload.scalar, + keyed: payload.keyed, + ..empty + }) + } +} + /// Planned readout of an exact summary. `lookback_ms` is the logical PromQL /// counter window; the evaluation range is resolved from it at run time. #[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)] @@ -214,6 +251,9 @@ impl AggregateCore for ExactAccumulator { 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() @@ -235,3 +275,68 @@ impl AggregateCore for ExactAccumulator { }) } } + +#[cfg(test)] +mod tests { + use super::*; + use serde::Serialize; + + #[derive(Serialize)] + struct Payload { + family: SummaryFamilyType, + scalar: ScalarState, + keyed: Option>, + } + + fn decode(payload: &Payload) -> Result { + rmp_serde::from_slice(&rmp_serde::to_vec_named(payload).unwrap()) + } + + fn sum() -> SummaryFamilyType { + SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum) + } + + // A persisted exact state decodes back to the same family, layout and readout. + #[test] + fn serialized_state_round_trips() { + let mut state = ExactAccumulator::new(sum(), true).unwrap(); + let key = KeyByLabelValues::new_with_labels(vec!["a".into()]); + state.update(Some(&key), 2.5, 0); + let bytes = rmp_serde::to_vec_named(&state).unwrap(); + let restored: ExactAccumulator = rmp_serde::from_slice(&bytes).unwrap(); + assert_eq!(restored.family(), &sum()); + assert_eq!( + restored.readout(Statistic::Sum, None, Some(&key)).unwrap(), + Some(2.5) + ); + } + + // Decoding rejects population states that differ from the declared family. + #[test] + fn decode_rejects_state_of_another_family() { + let scalar = Payload { + family: sum(), + scalar: ScalarState::Count(3), + keyed: None, + }; + assert!(decode(&scalar).is_err()); + let key = KeyByLabelValues::new_with_labels(vec!["a".into()]); + let keyed = Payload { + family: sum(), + scalar: ScalarState::Sum(0.0), + keyed: Some(HashMap::from([(key, ScalarState::Max(Some(1.0)))])), + }; + assert!(decode(&keyed).is_err()); + } + + // Decoding rejects families that have no exact Planner state. + #[test] + fn decode_rejects_unsupported_family() { + let payload = Payload { + family: SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Count), + scalar: ScalarState::Sum(0.0), + keyed: None, + }; + assert!(decode(&payload).is_err()); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/hll_sketch.rs b/crates/asap-physical-operators/src/summary_kernels/hll_sketch.rs index 089261c55..f675da230 100644 --- a/crates/asap-physical-operators/src/summary_kernels/hll_sketch.rs +++ b/crates/asap-physical-operators/src/summary_kernels/hll_sketch.rs @@ -24,6 +24,9 @@ impl AggregateCore for HllSketchAccumulator { 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, KernelError> { let other = other diff --git a/crates/asap-physical-operators/src/summary_kernels/hydra_kll.rs b/crates/asap-physical-operators/src/summary_kernels/hydra_kll.rs index c0347e69f..8d6f6a795 100644 --- a/crates/asap-physical-operators/src/summary_kernels/hydra_kll.rs +++ b/crates/asap-physical-operators/src/summary_kernels/hydra_kll.rs @@ -31,6 +31,9 @@ impl AggregateCore for 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, diff --git a/crates/asap-physical-operators/src/summary_kernels/increase.rs b/crates/asap-physical-operators/src/summary_kernels/increase.rs index f92ee582c..63095601d 100644 --- a/crates/asap-physical-operators/src/summary_kernels/increase.rs +++ b/crates/asap-physical-operators/src/summary_kernels/increase.rs @@ -99,6 +99,9 @@ impl AggregateCore for 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, diff --git a/crates/asap-physical-operators/src/summary_kernels/traits.rs b/crates/asap-physical-operators/src/summary_kernels/traits.rs index 7d866077e..3902a296c 100644 --- a/crates/asap-physical-operators/src/summary_kernels/traits.rs +++ b/crates/asap-physical-operators/src/summary_kernels/traits.rs @@ -13,6 +13,10 @@ pub trait AggregateCore: Send + Sync { fn as_any(&self) -> &dyn std::any::Any; + /// Mutable downcast, so a deployment can apply ingest deltas and merges to + /// a cached state in place instead of copying it for every frame. + fn as_any_mut(&mut self) -> &mut dyn std::any::Any; + /// Merge with a state of the same family and shape, leaving both inputs unchanged. fn merge_with(&self, other: &dyn AggregateCore) -> Result, KernelError>; @@ -33,3 +37,25 @@ impl Clone for Box { self.clone_boxed_core() } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::summary_kernels::DDSketchAccumulator; + + // A state behind a trait object can be updated in place through the mutable downcast. + #[test] + fn mutable_downcast_updates_in_place() { + let mut state: Box = Box::new(DDSketchAccumulator::new(0.01)); + let dd = state + .as_any_mut() + .downcast_mut::() + .unwrap(); + dd.inner.update(3.0); + let count = SketchQuery::PointCount { + key: planner_types::pre_asap::ColumnRef::SampleValue, + value: None, + }; + assert_eq!(state.estimate(&count).unwrap(), 1.0); + } +} diff --git a/crates/asap-physical-operators/src/summary_kernels/univmon.rs b/crates/asap-physical-operators/src/summary_kernels/univmon.rs index bffd8afaf..6d42c336f 100644 --- a/crates/asap-physical-operators/src/summary_kernels/univmon.rs +++ b/crates/asap-physical-operators/src/summary_kernels/univmon.rs @@ -2,9 +2,25 @@ use crate::AggregateCore; use asap_sketchlib::{DataInput, UnivMon}; +use planner_types::{post_asap::SketchQuery, pre_asap::ColumnRef}; type Error = Box; +fn check_dimensions( + heap_size: usize, + rows: usize, + cols: usize, + layers: usize, +) -> Result<(), Error> { + 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(()) +} + #[derive(Debug, Clone)] pub struct UnivMonAccumulator { inner: UnivMon, @@ -17,17 +33,33 @@ 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")?; + check_dimensions(heap_size, rows, cols, layers)?; Ok(Self { inner: UnivMon::init_univmon(heap_size, rows, cols, layers), }) } + /// Adopt a decoded sketch, e.g. one a deployment restored from its stored + /// bytes. Rejects dimensions [`Self::new`] rejects and terminal-mode + /// sketches, which do not accept this accumulator's updates. + pub fn from_sketch(sketch: UnivMon) -> Result { + check_dimensions( + sketch.heap_size, + sketch.sketch_row, + sketch.sketch_col, + sketch.layer_size, + )?; + if !sketch.accepts_standard_updates() { + return Err("terminal-mode UnivMon state cannot accept standard updates".into()); + } + Ok(Self { inner: sketch }) + } + + /// The underlying sketch, so a deployment can encode it with sketchlib's codec. + pub fn sketch(&self) -> &UnivMon { + &self.inner + } + /// 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() { @@ -96,6 +128,9 @@ impl AggregateCore for UnivMonAccumulator { 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 @@ -106,4 +141,79 @@ impl AggregateCore for UnivMonAccumulator { merged.merge_in_place(other)?; Ok(Box::new(merged)) } + + /// Sample count (a bare `PointCount`), distinct count, L2 norm and entropy + /// of the sample-value frequencies. + fn estimate(&self, query: &SketchQuery) -> Result { + Ok(match query { + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + } => self.inner.calc_l1(), + SketchQuery::Cardinality => self.inner.calc_card(), + SketchQuery::FrequencyL2 => self.inner.calc_l2(), + SketchQuery::FrequencyEntropy => self.inner.calc_entropy(), + other => return Err(format!("UnivMon does not answer {other:?}").into()), + }) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn count() -> SketchQuery { + SketchQuery::PointCount { + key: ColumnRef::SampleValue, + value: None, + } + } + + // Count, distinct, L2 and entropy readouts count each non-NaN sample once; + // signed zero is one identity. + #[test] + fn frequency_readouts() { + 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 read = |query| state.estimate(&query).unwrap(); + assert_eq!(read(count()), 4.0); + assert!((read(SketchQuery::Cardinality) - 2.0).abs() < 0.01); + assert!((read(SketchQuery::FrequencyL2) - 8.0f64.sqrt()).abs() < 0.01); + assert!((read(SketchQuery::FrequencyEntropy) - 1.0).abs() < 0.01); + assert!(state.estimate(&SketchQuery::Quantile { q: 0.5 }).is_err()); + } + + // A sketch taken out and adopted back answers the same readouts. + #[test] + fn adopted_sketch_keeps_readouts() { + let mut state = UnivMonAccumulator::new(32, 5, 1024, 4).unwrap(); + for value in [1.0, 2.0, 2.0] { + state.insert_sample(value).unwrap(); + } + let adopted = UnivMonAccumulator::from_sketch(state.sketch().clone()).unwrap(); + assert_eq!(adopted.dimensions(), state.dimensions()); + for query in [ + count(), + SketchQuery::Cardinality, + SketchQuery::FrequencyEntropy, + ] { + assert_eq!( + adopted.estimate(&query).unwrap(), + state.estimate(&query).unwrap() + ); + } + } + + // Adoption rejects terminal-mode sketches and dimensions `new` rejects. + #[test] + fn adoption_rejects_terminal_or_invalid_sketches() { + let mut terminal = UnivMon::init_univmon(4, 3, 16, 2); + terminal.fast_insert(&DataInput::U64(1), 1); + assert!(UnivMonAccumulator::from_sketch(terminal.clone()).is_err()); + terminal.free(); + assert!(UnivMonAccumulator::from_sketch(terminal).is_ok()); + assert!(UnivMonAccumulator::from_sketch(UnivMon::init_univmon(4, 21, 16, 2)).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 index 6f0085327..7c1122305 100644 --- a/crates/asap-physical-operators/src/summary_kernels/weighted_frequency.rs +++ b/crates/asap-physical-operators/src/summary_kernels/weighted_frequency.rs @@ -75,10 +75,12 @@ impl WeightedFrequency { Ok((algorithm, width as usize, depth as usize, capacity as usize)) } - pub(crate) fn algorithm(&self) -> FrequencyAlgorithm { + /// Algorithm and `(width, depth, capacity)` shape, so a deployment can check + /// a decoded state against its declared family. + pub fn algorithm(&self) -> FrequencyAlgorithm { self.inner.algorithm() } - pub(crate) fn shape(&self) -> (usize, usize, usize) { + pub fn shape(&self) -> (usize, usize, usize) { self.inner.shape() } pub fn new( @@ -114,6 +116,9 @@ impl AggregateCore for 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, diff --git a/crates/integration-tests/tests/precompute_raw_samples.rs b/crates/integration-tests/tests/precompute_raw_samples.rs new file mode 100644 index 000000000..ade52ec8d --- /dev/null +++ b/crates/integration-tests/tests/precompute_raw_samples.rs @@ -0,0 +1,610 @@ +//! Planner-selected summaries over raw samples compile as precompute graphs +//! and produce the same estimates as feeding their kernel sample by sample. +use std::{collections::BTreeMap, collections::BTreeSet, rc::Rc, sync::Arc}; + +use asap_aware_mapping::cost_model::DefaultCostModel; +use asap_aware_mapping::{ + search_workload, Replacement, ReplacementStrategy, ReplacementSubDAG, SketchAlgorithmStrategy, + TargetSubDAG, +}; +use asap_integration_tests::fixtures::lower_promql; +use asap_physical_operators::{ + factory::create_planner_accumulator, + operators::Operator, + physical_planner::{precompute, Source}, + runtime::{Limits, RunContext, Scope}, + summary_kernels::{exact::ExactAccumulator, weighted_frequency::WeightedFrequency}, + values::{Batch, Value}, + AggregateCore, KeyByLabelValues, Statistic, +}; +use asap_types::post_asap::{ + compile_post_asap_dag, EntityIdentity, ExactKind, PostAsapDag, PostAsapOperatorPayload, + SketchAlgorithm, SketchQuery, SummaryFamilyType, SummaryInputExpr, SummaryNode, SummaryUpdate, +}; +use asap_types::pre_asap::{expr_ir::ColumnRef, query_expr::Reduction}; +use asap_types::types::AccuracyTarget; +use futures::{executor::block_on, StreamExt}; + +type Series = BTreeMap; + +/// A series label set; `None` omits the label. An empty value is present +/// in the input but is not part of the series identity. +fn series(service: Option<&str>, instance: &str) -> Series { + [ + ("__name__", Some("m")), + ("service", service), + ("instance", Some(instance)), + ] + .into_iter() + .filter_map(|(k, v)| Some((k.to_owned(), v?.to_owned()))) + .collect() +} + +fn canonical(labels: &Series) -> Series { + labels + .iter() + .filter(|(_, v)| !v.is_empty()) + .map(|(k, v)| (k.clone(), v.clone())) + .collect() +} + +/// Every Planner candidate for `query`: the searched selection plus each +/// summary replacement of the root. +fn candidates(query: &str, accuracy: AccuracyTarget) -> Vec> { + let root = Rc::new(lower_promql(query, accuracy).expect("lowering failed")); + let mut result = SketchAlgorithmStrategy::default_cost_model() + .replacements(&TargetSubDAG::new(&root)) + .into_iter() + .filter_map(|candidate| match candidate { + ReplacementSubDAG { + replacement: Replacement::Summary(node), + .. + } => Some(node), + _ => None, + }) + .collect::>(); + let space = search_workload(vec![("query", root)]); + if let Ok(Some(selected)) = space + .global_selection(&DefaultCostModel) + .assemble_selected_dag(&space.roots[0].1) + { + result.push(selected); + } + result +} + +/// Raw-input summary nodes: `(dag, raw source id, summary id)`. +fn raw_summaries(dag: &PostAsapDag) -> Vec<(u64, u64)> { + dag.nodes + .iter() + .filter(|node| matches!(node.payload, PostAsapOperatorPayload::SummaryAgg { .. })) + .filter_map(|node| { + let inputs = dag + .edges + .iter() + .filter(|edge| edge.consumer == node.id) + .collect::>(); + let [edge] = inputs.as_slice() else { + return None; + }; + let source = dag.nodes.iter().find(|n| n.id == edge.producer)?; + matches!(source.payload, PostAsapOperatorPayload::Fallback { .. }) + .then_some((u64::from(source.id.0), u64::from(node.id.0))) + }) + .collect() +} + +fn samples() -> Vec<(Series, i64, f64)> { + let mut rows = Vec::new(); + let series_set = [ + (Some("a"), "1"), + (Some("a"), "2"), + (Some("b"), "1"), + (None, "3"), + (Some("b"), ""), + ]; + for (index, (service, instance)) in series_set.iter().enumerate() { + for step in 1..=5i64 { + let value = (index as f64 + 1.0) * step as f64 + (step % 2) as f64; + rows.push((series(*service, instance), step * 1000, value)); + } + } + rows +} + +fn execute( + dag: &PostAsapDag, + source: u64, + root: u64, + rows: &[(Series, i64, f64)], +) -> Vec<(Series, Arc)> { + let program = precompute::compile(dag, &[source], &[root]).unwrap_or_else(|error| { + panic!( + "raw summary {root} does not compile: {error}; source {:?}", + dag.nodes + .iter() + .find(|n| u64::from(n.id.0) == source) + .map(|n| (&n.output_schema, &n.payload)) + ); + }); + let program = serde_json::from_slice::< + asap_physical_operators::physical_planner::CompiledPhysicalDag, + >(&serde_json::to_vec(&program).unwrap()) + .unwrap(); + let schema = precompute::raw_sample_schema(); + let batch = Batch::try_new( + schema.clone(), + rows.iter() + .map(|(labels, time, value)| precompute::raw_sample_row(labels, *time, *value)) + .collect(), + ) + .unwrap(); + let sources = BTreeMap::from([( + source, + Box::new(Operator::source(schema, vec![batch]).unwrap()) as Source<'_>, + )]); + let graph = program.instantiate(sources).unwrap(); + let context = RunContext::new( + Scope::Ingestion { + window_start_ms: 0, + window_end_ms: 6000, + revision: 1, + }, + Limits::default(), + ) + .unwrap(); + block_on(async { + let mut stream = graph.execute(program.roots(), context).unwrap().remove(0); + let mut result = Vec::new(); + while let Some(batch) = stream.next().await { + for row in batch.unwrap().rows() { + let [Value::Map(labels), Value::Timestamp(6000), Value::Summary { state, .. }] = + row.as_slice() + else { + panic!("unexpected population row {row:?}"); + }; + let labels = labels + .iter() + .map(|(k, v)| match (k, v) { + (Value::Utf8(k), Value::Utf8(v)) => (k.to_string(), v.to_string()), + _ => panic!("non-label population entry"), + }) + .collect(); + result.push((labels, state.clone())); + } + } + result + }) +} + +/// PromQL grouping of a canonical label set: `by` keeps the named labels, +/// `without` drops them and `__name__`. +fn population(reduction: &Reduction, dag_labels: &[String], labels: &Series) -> Series { + let labels = canonical(labels); + match reduction { + Reduction::PerEntity => labels, + Reduction::Reduce(keys) => labels + .into_iter() + .filter(|(k, _)| { + if keys.is_without() { + k != "__name__" && !dag_labels.contains(k) + } else { + dag_labels.contains(k) + } + }) + .collect(), + } +} + +/// Item identity used by a keyed summary: labels by name, the sample value, +/// or the canonical label-set identity. +fn item(expr: &SummaryInputExpr, labels: &Series, value: f64, out: &mut Vec) { + match expr { + SummaryInputExpr::Column(ColumnRef::SampleValue) => out.push(Value::Float64(value)), + SummaryInputExpr::Column(ColumnRef::Named(name)) if name == "value" => { + out.push(Value::Float64(value)) + } + SummaryInputExpr::Column(ColumnRef::Named(name)) => out.push(Value::Utf8( + canonical(labels) + .get(name) + .cloned() + .unwrap_or_default() + .into(), + )), + SummaryInputExpr::EntityIdentity(EntityIdentity::PromqlLabelSet { excluding }) => { + let mut identity = canonical(labels); + for column in excluding { + let ColumnRef::Named(name) = column else { + panic!("unsupported fixture exclusion {column:?}") + }; + identity.remove(name); + } + out.push(Value::Utf8( + serde_json::to_string(&identity).unwrap().into(), + )) + } + SummaryInputExpr::Tuple(items) => items.iter().for_each(|i| item(i, labels, value, out)), + other => panic!("unsupported fixture item {other:?}"), + } +} + +fn weight(update: &SummaryUpdate, value: f64) -> f64 { + match update.weight { + SummaryInputExpr::Constant(weight) => weight, + _ => value, + } +} + +/// Estimates that identify a state's content for comparison. +fn readouts(state: &dyn AggregateCore, family: &SummaryFamilyType) -> Vec { + if let Some(exact) = state.as_any().downcast_ref::() { + let SummaryFamilyType::ExactAggregate(kind, _) = family else { + unreachable!() + }; + let statistic = match kind { + ExactKind::Sum => Statistic::Sum, + ExactKind::Count => Statistic::Count, + ExactKind::Min => Statistic::Min, + ExactKind::Max => Statistic::Max, + ExactKind::Rate => Statistic::Rate, + ExactKind::Increase => Statistic::Increase, + other => panic!("unexpected exact kind {other:?}"), + }; + return vec![exact + .readout(statistic, None, None::<&KeyByLabelValues>) + .unwrap() + .unwrap()]; + } + let SummaryFamilyType::Sketch(kind, _) = family else { + panic!("sketch state for exact family") + }; + match kind.algorithm() { + SketchAlgorithm::Kll | SketchAlgorithm::DDSketch => [0.1, 0.5, 0.9] + .into_iter() + .map(|q| state.estimate(&SketchQuery::Quantile { q }).unwrap()) + .collect(), + SketchAlgorithm::Hll => vec![state.estimate(&SketchQuery::Cardinality).unwrap()], + other => panic!("unexpected unkeyed sketch {other:?}"), + } +} + +/// Compile one raw-input summary, execute it over `rows`, and compare each +/// population with its kernel fed sample by sample. Returns the family label, +/// or the family when it has no native state. +fn check( + query: &str, + dag: &PostAsapDag, + source: u64, + root: u64, + rows: &[(Series, i64, f64)], +) -> Result { + let node = dag + .nodes + .iter() + .find(|n| u64::from(n.id.0) == root) + .unwrap(); + let PostAsapOperatorPayload::SummaryAgg { + family, + input, + reduction, + grouping, + } = &node.payload + else { + unreachable!() + }; + let source_node = dag + .nodes + .iter() + .find(|n| u64::from(n.id.0) == source) + .unwrap(); + let keys = match reduction { + Reduction::Reduce(keys) => keys + .keys() + .iter() + .map(|i| source_node.output_schema.fields[*i].name.clone()) + .collect(), + Reduction::PerEntity => vec![], + }; + if asap_physical_operators::capability::validate_native_family(family).is_err() { + // Families without a native state (e.g. plain CMS, UnivMon) + // are outside precompute execution; their compile must fail. + assert!(precompute::compile(dag, &[source], &[root]).is_err()); + return Err(format!("{family:?}")); + } + let actual = execute(dag, source, root, rows); + let label = match family { + SummaryFamilyType::ExactAggregate(kind, _) => format!("{kind:?}"), + SummaryFamilyType::Sketch(kind, _) => format!("{:?}", kind.algorithm()), + other => format!("{other:?}"), + }; + if let SummaryFamilyType::Sketch(kind, _) = family { + if let (Some(keyed), false) = (&input.item, kind.algorithm() == &SketchAlgorithm::Hll) { + // Keyed heaps: every item's estimated weight is its exact + // total at this scale (no collisions in the fixture). + let mut expected = BTreeMap::>::new(); + for (labels, _, value) in rows { + let mut items = Vec::new(); + item(keyed, labels, *value, &mut items); + *expected + .entry(population(reduction, &keys, labels)) + .or_default() + .entry(format!("{items:?}")) + .or_default() += weight(input, *value); + } + assert_eq!(actual.len(), expected.len(), "{query}"); + for (labels, state) in &actual { + let heap = state.as_any().downcast_ref::().unwrap(); + let got = heap + .rows(usize::MAX >> 1) + .into_iter() + .map(|mut row| { + let Some(Value::Float64(score)) = row.pop() else { + panic!("heap score") + }; + (format!("{row:?}"), score) + }) + .collect::>(); + assert_eq!(&got, &expected[labels], "{query}"); + } + return Ok(label); + } + } + let mut expected = + BTreeMap::>::new(); + for (labels, time, value) in rows { + let updater = expected + .entry(population(reduction, &keys, labels)) + .or_insert_with(|| create_planner_accumulator(family, input, grouping).unwrap()); + let unit = input.item.is_some(); + updater.update_single(if unit { *value } else { weight(input, *value) }, *time); + } + assert_eq!(actual.len(), expected.len(), "{query}"); + for (labels, state) in actual { + let reference = expected[&labels].snapshot_accumulator(); + assert_eq!( + readouts(state.as_ref(), family), + readouts(reference.as_ref(), family), + "{query}: {labels:?}" + ); + } + Ok(label) +} + +// Every raw-input summary selected by Planner compiles over raw sample rows, +// and each population's estimates equal feeding its kernel sample by sample. +#[test] +fn raw_sample_summaries_compile_and_match_their_kernels() { + let exact = AccuracyTarget::Exact; + let sketch = AccuracyTarget::Epsilon(0.02); + let queries = [ + ("sum_over_time(m[5m])", &exact), + ("count_over_time(m[5m])", &exact), + ("min_over_time(m[5m])", &exact), + ("max_over_time(m[5m])", &exact), + ("rate(m[5m])", &exact), + ("increase(m[5m])", &exact), + ("sum by (service) (sum_over_time(m[5m]))", &exact), + ("sum by (service) (rate(m[5m]))", &exact), + ("topk(2, sum_over_time(m[5m]))", &exact), + ("quantile_over_time(0.9, m[5m])", &sketch), + ("sum by (service) (quantile_over_time(0.9, m[5m]))", &sketch), + ("quantile by (service) (0.9, m)", &sketch), + ("distinct_over_time(m[5m])", &sketch), + ("count(m)", &sketch), + ("topk(2, m)", &sketch), + ( + "topk(2, sum by (service) (count_over_time(m[5m])))", + &sketch, + ), + ("topk(2, sum_over_time(m[5m]))", &sketch), + ("topk by (service) (2, sum_over_time(m[5m]))", &sketch), + ]; + let rows = samples(); + let mut families = BTreeSet::new(); + let mut unsupported = BTreeSet::new(); + let mut checked = BTreeMap::new(); + for (query, accuracy) in queries { + for candidate in candidates(query, accuracy.clone()) { + let dag = compile_post_asap_dag(&candidate).unwrap(); + for (source, root) in raw_summaries(&dag) { + match check(query, &dag, source, root, &rows) { + Ok(family) => { + families.insert(family); + *checked.entry(query).or_insert(0) += 1; + } + Err(family) => { + unsupported.insert(family); + } + } + } + } + } + println!("checked {checked:?}; families {families:?}; without native state {unsupported:?}"); + for query in [ + "topk by (service) (2, sum_over_time(m[5m]))", + "quantile by (service) (0.9, m)", + "distinct_over_time(m[5m])", + ] { + assert!( + checked.contains_key(query), + "{query} has no checked raw summary" + ); + } + for family in [ + "Sum", + "Count", + "Min", + "Max", + "Rate", + "Increase", + "Kll", + "DDSketch", + "Hll", + "CountSketchWithHeap", + ] { + assert!( + families.contains(family), + "no {family} fixture: {families:?}" + ); + } +} + +/// Replace the raw summary of `sum by (service) (sum_over_time(m[5m]))` with +/// another update, keeping its raw input and reduction. +fn grouped_raw_summary(family: SummaryFamilyType, input: SummaryUpdate) -> (PostAsapDag, u64, u64) { + let candidate = candidates( + "sum by (service) (sum_over_time(m[5m]))", + AccuracyTarget::Exact, + ) + .pop() + .unwrap(); + let mut dag = compile_post_asap_dag(&candidate).unwrap(); + let (source, root) = raw_summaries(&dag)[0]; + let node = dag + .nodes + .iter_mut() + .find(|n| u64::from(n.id.0) == root) + .unwrap(); + let PostAsapOperatorPayload::SummaryAgg { + family: old, + input: update, + .. + } = &mut node.payload + else { + unreachable!() + }; + for field in &mut node.output_schema.fields { + if field.dtype == *old { + field.dtype = family.clone(); + } + } + *old = family; + *update = input; + let schema = node.output_schema.clone(); + for edge in dag + .edges + .iter_mut() + .filter(|e| u64::from(e.producer.0) == root) + { + edge.intermediate_schema = schema.clone(); + } + (dag, source, root) +} + +// Keyed heaps over raw samples resolve items from the series label set and +// estimate each item's exact total; invalid weight contracts do not compile. +#[test] +fn raw_sample_heaps_resolve_items_from_labels() { + use asap_types::post_asap::{ + GroupingStrategy::PerSubpopulationInstance, NonNegativeWeightProof, SketchKind, + SketchParams, WeightDomain, + }; + let heap = |algorithm, params| { + SummaryFamilyType::Sketch(SketchKind::new(algorithm, params), PerSubpopulationInstance) + }; + let cms = heap( + SketchAlgorithm::CmsWithHeap, + SketchParams::CmsWithHeap { + width: 64, + depth: 3, + heap_size: 8, + }, + ); + let count_sketch = heap( + SketchAlgorithm::CountSketchWithHeap, + SketchParams::CountSketchWithHeap { + width: 64, + depth: 3, + heap_size: 8, + }, + ); + let identity = SummaryUpdate { + item: Some(SummaryInputExpr::EntityIdentity( + EntityIdentity::PromqlLabelSet { excluding: vec![] }, + )), + weight: SummaryInputExpr::Constant(1.0), + weight_domain: WeightDomain::NonNegative { + proof: NonNegativeWeightProof::UnitCount, + }, + }; + let by_instance = SummaryUpdate { + item: Some(SummaryInputExpr::Tuple(vec![SummaryInputExpr::Column( + ColumnRef::Named("instance".into()), + )])), + weight: SummaryInputExpr::Column(ColumnRef::SampleValue), + weight_domain: WeightDomain::UnknownOrSigned, + }; + let rows = samples(); + for (family, input) in [ + (cms.clone(), identity.clone()), + (count_sketch.clone(), identity), + (count_sketch, by_instance.clone()), + ] { + let (dag, source, root) = grouped_raw_summary(family, input); + assert!(check("heap", &dag, source, root, &rows).is_ok()); + } + let signed_cms = grouped_raw_summary(cms.clone(), by_instance.clone()); + assert!(precompute::compile(&signed_cms.0, &[signed_cms.1], &[signed_cms.2]).is_err()); + let derivative = grouped_raw_summary( + cms, + SummaryUpdate { + weight_domain: WeightDomain::NonNegative { + proof: NonNegativeWeightProof::ResetAwareCounterDerivative, + }, + ..by_instance.clone() + }, + ); + assert!( + precompute::compile(&derivative.0, &[derivative.1], &[derivative.2]).is_err(), + "raw samples are cumulative counters, not their derivative" + ); + // A scan's time column is not a label; it cannot silently read as empty. + let time_item = grouped_raw_summary( + heap( + SketchAlgorithm::CountSketchWithHeap, + SketchParams::CountSketchWithHeap { + width: 64, + depth: 3, + heap_size: 8, + }, + ), + SummaryUpdate { + item: Some(SummaryInputExpr::Column(ColumnRef::Named("ts".into()))), + ..by_instance + }, + ); + assert!(precompute::compile(&time_item.0, &[time_item.1], &[time_item.2]).is_err()); +} + +// `without` grouping over raw samples drops the listed labels and `__name__`. +#[test] +fn raw_sample_without_grouping_drops_labels_and_name() { + use asap_types::pre_asap::query_expr::GroupKeys; + let family = + SummaryFamilyType::ExactAggregate(ExactKind::Sum, asap_types::post_asap::ExactParams::Sum); + let (mut dag, source, root) = + grouped_raw_summary(family, SummaryUpdate::column(ColumnRef::SampleValue)); + let service = dag + .nodes + .iter() + .find(|n| u64::from(n.id.0) == source) + .unwrap() + .output_schema + .fields + .iter() + .position(|f| f.name == "service") + .unwrap(); + let node = dag + .nodes + .iter_mut() + .find(|n| u64::from(n.id.0) == root) + .unwrap(); + let PostAsapOperatorPayload::SummaryAgg { reduction, .. } = &mut node.payload else { + unreachable!() + }; + *reduction = Reduction::Reduce(GroupKeys::without(vec![service])); + assert_eq!( + check("without", &dag, source, root, &samples()), + Ok("Sum".into()) + ); +}