diff --git a/crates/asap-aware-mapping/src/pass/mod.rs b/crates/asap-aware-mapping/src/pass/mod.rs index e26998ace..beaf421eb 100644 --- a/crates/asap-aware-mapping/src/pass/mod.rs +++ b/crates/asap-aware-mapping/src/pass/mod.rs @@ -267,6 +267,31 @@ impl PlanOutput { nodes } + /// The workload as one physical ASAP DAG, with shared sub-DAGs exported + /// once: operator queries become its `roots` and standalone scalar + /// queries its `scalar_roots`, keyed by position in the batch. That + /// position is the entry index, because `check_contract` requires dense + /// entry indices. + /// + /// An empty workload is an error + /// ([`PhysicalASAPDAGValidationError::NoRoots`](asap_types::ir::physical_export::PhysicalASAPDAGValidationError::NoRoots)). + /// A workload of only literal scalars yields a DAG with no nodes: just + /// its `scalar_roots`. + pub fn execution_timed_dag( + &self, + ) -> Result< + asap_types::ir::physical_export::PhysicalASAPDAG, + crate::SummaryMaintenanceTimingError, + > { + let plans: Vec<_> = self + .plans + .iter() + .chain(&self.scalar_plans) + .map(|p| &p.plan) + .collect(); + crate::summary_maintenance_lifecycle::execution_timed_workload_dag(&self.roots(), &plans) + } + pub fn len(&self) -> usize { self.plans.len() + self.scalar_roots.len() } diff --git a/crates/asap-aware-mapping/src/summary_maintenance_lifecycle.rs b/crates/asap-aware-mapping/src/summary_maintenance_lifecycle.rs index e0bbfbeff..3123ba558 100644 --- a/crates/asap-aware-mapping/src/summary_maintenance_lifecycle.rs +++ b/crates/asap-aware-mapping/src/summary_maintenance_lifecycle.rs @@ -22,11 +22,11 @@ use std::collections::{HashMap, HashSet}; use std::rc::Rc; use asap_types::ir::physical_export::{ - compile_physical_asap_dag_with_node_ids, PhysicalASAPDAG, PhysicalASAPDAGValidationError, - PhysicalASAPNodeId, + compile_physical_asap_dag_with_node_ids, compile_physical_asap_workload_with_node_ids, + PhysicalASAPDAG, PhysicalASAPDAGValidationError, PhysicalASAPNodeId, }; use asap_types::ir::timing::{apply_lifecycle_timings, LifecycleAssignment, TimingMemo}; -use asap_types::ir::{ASAPOp, Operator, OperatorNode}; +use asap_types::ir::{ASAPOp, Operator, OperatorNode, QueryRoot}; use asap_types::post_asap::{ EvaluationSchedule, ExecutionDataStateError, ExecutionTiming, OutputRepresentation, ResultGuarantee, SummaryMaintenanceLifecycle, SummaryMaintenanceLifecycleGuarantee, @@ -247,64 +247,94 @@ impl SummaryMaintenanceLifecyclePlan { /// a `SummaryAgg` is one of its inputs. Timings already on the root are /// ignored. pub fn execution_timed_dag(&self) -> Result { - let mut memo = TimingMemo::new(); - let timed = apply_lifecycle_timings( - &self.root, - &LifecycleAssignment::default_maintained(), - &mut memo, - )?; - let compiled = compile_physical_asap_dag_with_node_ids(&timed)?; - let dag = compiled.dag; - for population in &standalone_populations(&self.root) { - let id = compiled - .node_ids - .node_id(memo.timed(population).expect("population was timed")) - .expect("collected population belongs to the compiled DAG"); - if !self - .deployments - .iter() - .any(|deployment| deployment.post_asap_node_id == id) - { + execution_timed_workload_dag(&[QueryRoot::Operator(Rc::clone(&self.root))], &[self]) + } +} + +/// One physical ASAP DAG for a workload: `roots` are the queries in +/// workload order, operator roots becoming `roots` and standalone scalars +/// `scalar_roots` of the export, with sub-DAGs shared between queries +/// exported once. `plans` are the lifecycle plans of every planned sub-DAG +/// `roots` reach; timing follows the selected lifecycles of their +/// deployments, as in [`SummaryMaintenanceLifecyclePlan::execution_timed_dag`]. +pub(crate) fn execution_timed_workload_dag( + roots: &[QueryRoot], + plans: &[&SummaryMaintenanceLifecyclePlan], +) -> Result { + // One memo, so a node shared by several roots is timed and exported once. + let mut memo = TimingMemo::new(); + let assignment = LifecycleAssignment::default_maintained(); + let mut timed = Vec::with_capacity(roots.len()); + for root in roots { + timed.push(match root { + QueryRoot::Operator(node) => { + QueryRoot::Operator(apply_lifecycle_timings(node, &assignment, &mut memo)?) + } + QueryRoot::Scalar(expr) => { + for node in expr.operator_refs() { + apply_lifecycle_timings(node, &assignment, &mut memo)?; + } + QueryRoot::Scalar(expr.map_operator_refs(&mut |node| { + Rc::clone(memo.timed(node).expect("scalar reference was timed")) + })) + } + }); + } + let compiled = compile_physical_asap_workload_with_node_ids(&timed)?; + let id_of = |node: &Rc| { + compiled + .node_ids + .node_id(memo.timed(node).expect("plan node was timed")) + .expect("timed plan node belongs to the compiled DAG") + }; + let deployments: Vec<_> = plans + .iter() + .flat_map(|plan| &plan.deployments) + .map(|deployment| (id_of(&deployment.summary), deployment)) + .collect(); + for plan in plans { + for population in &standalone_populations(&plan.root) { + let id = id_of(population); + if !deployments.iter().any(|(deployed, _)| *deployed == id) { return Err(SummaryMaintenanceTimingError::UnplannedMaintainedState(id)); } } - let mut pending = Vec::new(); - for deployment in &self.deployments { - let guarantee = deployment - .summary_maintenance_lifecycle_guarantee - .as_ref() - .ok_or(SummaryMaintenanceTimingError::UnselectedLifecycle( - deployment.post_asap_node_id, - ))?; - if guarantee.summary_maintenance_lifecycle != SummaryMaintenanceLifecycle::Ephemeral { - pending.push(deployment.post_asap_node_id); - } + } + let dag = compiled.dag; + let mut pending = Vec::new(); + for (id, deployment) in &deployments { + let guarantee = deployment + .summary_maintenance_lifecycle_guarantee + .as_ref() + .ok_or(SummaryMaintenanceTimingError::UnselectedLifecycle(*id))?; + if guarantee.summary_maintenance_lifecycle != SummaryMaintenanceLifecycle::Ephemeral { + pending.push(*id); } - let mut ingestion = HashSet::new(); - while let Some(id) = pending.pop() { - if ingestion.insert(id) { - pending.extend( - dag.edges - .iter() - .filter(|edge| edge.consumer == id) - .map(|edge| edge.producer), - ); - } + } + let mut ingestion = HashSet::new(); + while let Some(id) = pending.pop() { + if ingestion.insert(id) { + pending.extend( + dag.edges + .iter() + .filter(|edge| edge.consumer == id) + .map(|edge| edge.producer), + ); } - let phases = dag - .nodes - .iter() - .map(|node| { - let timing = if ingestion.contains(&node.id) { - ExecutionTiming::IngestionTime - } else { - ExecutionTiming::QueryTime - }; - (node.id, timing) - }) - .collect(); - Ok(dag.with_execution_phases(&phases)?) } + let phases = dag + .nodes + .iter() + .map(|node| { + let timing = if ingestion.contains(&node.id) { + ExecutionTiming::IngestionTime + } else { + ExecutionTiming::QueryTime + }; + (node.id, timing) + }) + .collect(); + Ok(dag.with_execution_phases(&phases)?) } /// Explicit association between a materialized target and the normalized diff --git a/crates/integration-tests/tests/operator_design_examples.rs b/crates/integration-tests/tests/operator_design_examples.rs index e5b2c0a19..c5f0c66ee 100644 --- a/crates/integration-tests/tests/operator_design_examples.rs +++ b/crates/integration-tests/tests/operator_design_examples.rs @@ -398,4 +398,22 @@ async fn batch_planning_replaces_and_shares_summary_operators() { rows ); } + // The batch exports as one physical DAG: a root per query and the shared + // SUM state once. + let workload_dag = output.execution_timed_dag().unwrap(); + assert_eq!(workload_dag.roots.len(), 2); + assert_ne!(workload_dag.roots[0], workload_dag.roots[1]); + assert_eq!( + workload_dag + .nodes + .iter() + .filter(|n| matches!( + n.payload, + asap_types::ir::physical_export::PhysicalASAPOperatorPayload::ASAP( + asap_types::ir::ASAPOp::SummaryAgg { .. } + ) + )) + .count(), + 1 + ); } diff --git a/crates/integration-tests/tests/summary_coverage_e2e.rs b/crates/integration-tests/tests/summary_coverage_e2e.rs new file mode 100644 index 000000000..978788ef1 --- /dev/null +++ b/crates/integration-tests/tests/summary_coverage_e2e.rs @@ -0,0 +1,275 @@ +//! Query string → planned summary state → `coverage()`, for the examples of +//! the ASAP primitive schema design doc (#573 §4.2.2). Each query is lowered +//! and searched, and the quantile's KLL `SummaryAgg` candidate in the search +//! space is checked: what its `definition` keeps and what its `selection` +//! takes, row by row of §4.2.2's "More examples". The tests read the +//! candidate list, not a cost model's selection, so they do not depend on +//! which candidate a cost model prefers. + +use std::ops::Bound; +use std::rc::Rc; + +use asap_aware_mapping::replacement::Replacement; +use asap_aware_mapping::search_workload; +use asap_frontend_sql::{lower_sql, SqlCatalog}; +use asap_integration_tests::fixtures::lower_promql; +use asap_types::ir::summary_coverage::{ColumnIdentity, Constraint, SelectionBox}; +use asap_types::ir::{ASAPOp, NonASAPOp, OperatorNode, Predicate, ScalarExpr}; +use asap_types::post_asap::{FieldDataType, SketchAlgorithm}; +use asap_types::pre_asap::expr_ir::ScalarValue; +use asap_types::pre_asap::schema::{DataType, Field, Schema}; +use asap_types::types::AccuracyTarget; + +fn accuracy() -> AccuracyTarget { + AccuracyTarget::EpsilonDelta { + epsilon: 0.01, + delta: 0.01, + } +} + +/// The single KLL `SummaryAgg` among the search space's candidates for `pre`. +fn kll_state(pre: Rc) -> Rc { + let space = search_workload(vec![("query", pre)]).unwrap(); + let mut states: Vec> = Vec::new(); + for group in space.target_subdag_candidates() { + for candidate in &group.candidates { + let Replacement::SubDAG(dag) = &candidate.replacement else { + continue; + }; + for node in OperatorNode::reachable(dag) { + let is_kll = matches!(node.asap(), Some(ASAPOp::SummaryAgg { family: FieldDataType::Sketch(kind, _), .. }) + if kind.algorithm() == &SketchAlgorithm::Kll); + if is_kll && !states.iter().any(|s| Rc::ptr_eq(s, &node)) { + states.push(node); + } + } + } + } + assert_eq!( + states.len(), + 1, + "expected one KLL candidate state: {states:#?}" + ); + states.pop().unwrap() +} + +async fn sql_state(sql: &str) -> Rc { + let catalog = SqlCatalog::new().with_table( + "t", + Schema::new(vec![ + Field::plain("job", DataType::Utf8, false), + Field::plain("region", DataType::Utf8, false), + Field::plain("latency", DataType::Float64, false), + ]), + ); + let pre = lower_sql(sql, &catalog, accuracy()) + .await + .unwrap_or_else(|e| panic!("lower failed for {sql:?}: {e}")); + kll_state(pre) +} + +fn promql_state(query: &str) -> Rc { + kll_state(lower_promql(query, accuracy()).expect("lowering failed")) +} + +fn column(table: Option<&str>, name: &str) -> ColumnIdentity { + ColumnIdentity { + table: table.map(str::to_string), + name: name.to_string(), + } +} + +fn utf8(value: &str) -> ScalarValue { + ScalarValue::Utf8(value.to_string()) +} + +/// The child of the definition's `SummaryAgg`. +fn definition_child(state: &OperatorNode) -> Rc { + state.coverage().unwrap().definition.children()[0].clone() +} + +/// The definition's `SummaryAgg` carries no filter of its own. +fn assert_definition_unfiltered(state: &OperatorNode) { + let definition = state.coverage().unwrap().definition; + assert!( + matches!( + definition.asap(), + Some(ASAPOp::SummaryAgg { filter: None, .. }) + ), + "{definition:#?}" + ); +} + +/// The definition keeps the whole sub-DAG below the state: nothing moved +/// into the selection, which is the one unrestricted box. +fn assert_nothing_selected(state: &OperatorNode) { + assert_eq!( + state.coverage().unwrap().selection, + vec![SelectionBox::default()] + ); + assert!(Rc::ptr_eq(&definition_child(state), state.children()[0])); +} + +#[tokio::test] +async fn sql_where_value_set_moves_to_selection() { + let state = + sql_state("SELECT approx_percentile_cont(latency, 0.99) FROM t WHERE region = 'us'").await; + let coverage = state.coverage().unwrap(); + assert_eq!( + coverage.selection, + vec![SelectionBox { + columns: [( + column(Some("t"), "region"), + Constraint::In(vec![utf8("us")]) + )] + .into(), + relative_time: None, + }] + ); + assert_definition_unfiltered(&state); + let child = definition_child(&state); + assert!( + matches!(child.non_asap(), Some(NonASAPOp::Scan { predicates, .. }) if predicates.is_empty()) + ); +} + +#[tokio::test] +async fn sql_where_range_moves_to_selection() { + let state = + sql_state("SELECT approx_percentile_cont(latency, 0.99) FROM t WHERE latency < 100").await; + assert_eq!( + state.coverage().unwrap().selection, + vec![SelectionBox { + columns: [( + column(Some("t"), "latency"), + Constraint::Interval { + lower: Bound::Unbounded, + upper: Bound::Excluded(ScalarValue::Float64(100.0)), + } + )] + .into(), + relative_time: None, + }] + ); + assert_definition_unfiltered(&state); + let child = definition_child(&state); + assert!( + matches!(child.non_asap(), Some(NonASAPOp::Scan { predicates, .. }) if predicates.is_empty()) + ); +} + +/// The worked example without its `FILTER` clause: the rename carries +/// `region = 'us'` into the selection as `r`, and the expression condition +/// stays in the definition's scan. +#[tokio::test] +async fn sql_renamed_column_moves_and_expression_stays() { + let state = sql_state( + "SELECT job, approx_percentile_cont(latency, 0.99) \ + FROM (SELECT job, region AS r, latency FROM t \ + WHERE region = 'us' AND latency * 2 > 10) \ + GROUP BY job", + ) + .await; + assert_eq!( + state.coverage().unwrap().selection, + vec![SelectionBox { + columns: [(column(None, "r"), Constraint::In(vec![utf8("us")]))].into(), + relative_time: None, + }] + ); + assert_definition_unfiltered(&state); + let project = definition_child(&state); + assert!(matches!( + project.non_asap(), + Some(NonASAPOp::Project { .. }) + )); + let scan = project.children()[0].clone(); + let Some(NonASAPOp::Scan { predicates, .. }) = scan.non_asap() else { + panic!("expected a scan, got {scan:#?}"); + }; + assert_eq!(predicates.len(), 1, "{predicates:#?}"); + assert!( + matches!(&predicates[0], Predicate(ScalarExpr::Compare { left, .. }) + if matches!(**left, ScalarExpr::Arithmetic { .. })), + "the expression condition should stay: {predicates:#?}" + ); +} + +#[test] +fn promql_offset_window_is_relative_time() { + let state = promql_state("quantile_over_time(0.99, m[1m] offset 2m)"); + assert_eq!( + state.coverage().unwrap().selection, + vec![SelectionBox { + columns: Default::default(), + relative_time: Some((Bound::Excluded(-180_000), Bound::Included(-120_000))), + }] + ); + assert!(matches!( + definition_child(&state).non_asap(), + Some(NonASAPOp::Scan { .. }) + )); +} + +/// The label matcher is below `rate`, which the walk does not pass: the +/// definition keeps all of it, the matcher included. +#[test] +fn promql_matcher_below_rate_stays_in_definition() { + let state = promql_state("quantile(0.99, rate(m{job=\"api\"}[5m]))"); + assert_nothing_selected(&state); + let scans: Vec<_> = OperatorNode::reachable(&definition_child(&state)) + .into_iter() + .filter_map(|node| match node.non_asap() { + Some(NonASAPOp::Scan { predicates, .. }) => Some(predicates.len()), + _ => None, + }) + .collect(); + assert_eq!(scans, vec![1], "the matcher stays in the scan"); +} + +/// A PromQL comparison lowers to a `Filter` above `rate`, so it picks the +/// rate outputs the state reads; the definition keeps the `rate` below it. +#[test] +fn promql_comparison_above_rate_moves_to_selection() { + let state = promql_state("quantile(0.99, rate(m[5m]) > 0)"); + let filter = state.children()[0]; + assert!(matches!(filter.non_asap(), Some(NonASAPOp::Filter { .. }))); + assert!(Rc::ptr_eq(&definition_child(&state), filter.children()[0])); + assert_eq!( + state.coverage().unwrap().selection, + vec![SelectionBox { + columns: [( + column(None, "value"), + Constraint::Interval { + lower: Bound::Excluded(ScalarValue::Float64(0.0)), + upper: Bound::Unbounded, + }, + )] + .into(), + relative_time: None, + }] + ); +} + +/// A vector/vector comparison is a `BinaryOp`, which the walk does not enter. +#[test] +fn promql_vector_comparison_stays_in_definition() { + let state = promql_state("quantile(0.99, rate(a[5m]) > rate(b[5m]))"); + assert_nothing_selected(&state); + assert!(matches!( + definition_child(&state).non_asap(), + Some(NonASAPOp::BinaryOp { .. }) + )); +} + +/// A `bool` comparison lowers to a `Project` (values become 0 or 1), which +/// drops no row, so it stays in the definition. +#[test] +fn promql_bool_comparison_stays_in_definition() { + let state = promql_state("quantile(0.99, rate(m[5m]) > bool 0)"); + assert_nothing_selected(&state); + assert!(matches!( + definition_child(&state).non_asap(), + Some(NonASAPOp::Project { .. }) + )); +} diff --git a/crates/planner/tests/summary_sharing.rs b/crates/planner/tests/summary_sharing.rs index 73a2c62ff..cdb80f045 100644 --- a/crates/planner/tests/summary_sharing.rs +++ b/crates/planner/tests/summary_sharing.rs @@ -595,6 +595,20 @@ async fn scalar_queries_read_planned_summaries() { let refs = output.scalar_roots[1].1.operator_refs(); assert_eq!(refs.len(), 1); assert!(Rc::ptr_eq(refs[0], &read.plan.root)); + + // The export keys each scalar query by its entry and points the scalar + // read at the node the operator query exports: one summary, read twice. + let dag = output.execution_timed_dag().unwrap(); + let keys: Vec<_> = dag.scalar_roots.iter().map(|(i, _)| *i).collect(); + assert_eq!(keys, vec![0, 1]); + assert_eq!(dag.roots.len(), 1); + assert_eq!(dag.scalar_roots[1].1.operator_refs(), vec![&dag.roots[0]]); + assert!(dag.nodes.iter().any(|node| matches!( + node.payload, + asap_types::ir::physical_export::PhysicalASAPOperatorPayload::ASAP( + ASAPOp::SummaryAgg { .. } + ) + ))); } /// A combination of approximate values may read summaries: each reference of