|
| 1 | +//! Pass 1 alternatives over SQL row sources compose, compile and execute. |
| 2 | +mod physical_common; |
| 3 | +use std::collections::BTreeMap; |
| 4 | +use std::rc::Rc; |
| 5 | + |
| 6 | +use asap_executor::values::Value; |
| 7 | +use asap_frontend_sql::{lower_sql, SqlCatalog}; |
| 8 | +use asap_logical_optimizer::pass1::logical_candidates::{ |
| 9 | + compose_logical_candidate, enumerate_choices, enumerate_local_logical_candidates, |
| 10 | + LocalLogicalCandidates, |
| 11 | +}; |
| 12 | +use asap_plan_selection::{plan_stages, PlanningModels}; |
| 13 | +use asap_types::ir::schema::{DataType, Field, Schema}; |
| 14 | +use asap_types::ir::{ASAPOp, Operator, OperatorNode, QueryRoot}; |
| 15 | +use asap_types::types::AccuracyTarget; |
| 16 | +use asap_types::workload::{ |
| 17 | + DataArrival, DataWorkload, Evidence, EvidenceSource, Predictability, QueryRecurrence, Rate, |
| 18 | + RootDemand, |
| 19 | +}; |
| 20 | + |
| 21 | +fn catalog() -> SqlCatalog { |
| 22 | + SqlCatalog::new().with_table( |
| 23 | + "flows", |
| 24 | + Schema::new(vec![ |
| 25 | + Field::plain("ts", DataType::Timestamp, false), |
| 26 | + Field::plain("src_ip", DataType::Utf8, false), |
| 27 | + ]), |
| 28 | + ) |
| 29 | +} |
| 30 | + |
| 31 | +fn target() -> AccuracyTarget { |
| 32 | + AccuracyTarget::EpsilonDelta { |
| 33 | + epsilon: 0.01, |
| 34 | + delta: 0.01, |
| 35 | + } |
| 36 | +} |
| 37 | + |
| 38 | +fn rows() -> Vec<Vec<Value>> { |
| 39 | + ["a", "a", "b", "a", "c", "b"] |
| 40 | + .into_iter() |
| 41 | + .map(|ip| vec![Value::Timestamp(0), Value::Utf8(ip.into())]) |
| 42 | + .collect() |
| 43 | +} |
| 44 | + |
| 45 | +async fn inventory(sql: &str) -> LocalLogicalCandidates<usize> { |
| 46 | + let root = lower_sql(sql, &catalog(), target()).await.unwrap(); |
| 47 | + enumerate_local_logical_candidates(vec![(0, QueryRoot::Operator(root))], &BTreeMap::new()) |
| 48 | + .unwrap() |
| 49 | +} |
| 50 | + |
| 51 | +fn has_estimate(root: &Rc<OperatorNode>) -> bool { |
| 52 | + OperatorNode::reachable(root) |
| 53 | + .iter() |
| 54 | + .any(|n| matches!(n.operator, Operator::ASAP(ASAPOp::SummaryEstimate { .. }))) |
| 55 | +} |
| 56 | + |
| 57 | +/// `COUNT(*) … GROUP BY src_ip` (#509 Example 2's inner query) reads no |
| 58 | +/// sample value. Every candidate composes and compiles; the exact ones, |
| 59 | +/// the exact `Count` accumulator included, count rows per `src_ip`. |
| 60 | +#[tokio::test] |
| 61 | +async fn sql_count_star_group_by_candidates_compose_and_execute() { |
| 62 | + let inventory = inventory("SELECT src_ip, COUNT(*) AS c FROM flows GROUP BY src_ip").await; |
| 63 | + let choices = enumerate_choices(&inventory, usize::MAX); |
| 64 | + assert!(choices.len() > 2, "exact and summary candidates"); |
| 65 | + let mut executed = 0; |
| 66 | + for choice in choices { |
| 67 | + let roots = compose_logical_candidate(&inventory, &choice) |
| 68 | + .unwrap_or_else(|e| panic!("{choice:?} composes: {e}")); |
| 69 | + let QueryRoot::Operator(root) = &roots[0].1 else { |
| 70 | + panic!("operator root") |
| 71 | + }; |
| 72 | + physical_common::compile_physical_asap_dag(root) |
| 73 | + .unwrap_or_else(|e| panic!("{choice:?} compiles: {e}")); |
| 74 | + if has_estimate(root) { |
| 75 | + continue; |
| 76 | + } |
| 77 | + let mut rows = physical_common::execute_raw_rows(root, rows()); |
| 78 | + rows.sort_by_key(|row| format!("{row:?}")); |
| 79 | + let counts: Vec<_> = rows |
| 80 | + .iter() |
| 81 | + .map(|row| match (&row[0], &row[1]) { |
| 82 | + (Value::Utf8(ip), Value::Int64(n)) => (ip.to_string(), *n), |
| 83 | + other => panic!("{choice:?}: unexpected row {other:?}"), |
| 84 | + }) |
| 85 | + .collect(); |
| 86 | + assert_eq!( |
| 87 | + counts, |
| 88 | + [("a".into(), 3), ("b".into(), 2), ("c".into(), 1)], |
| 89 | + "{choice:?}" |
| 90 | + ); |
| 91 | + executed += 1; |
| 92 | + } |
| 93 | + assert_eq!(executed, 2, "pass-through and the exact Count accumulator"); |
| 94 | +} |
| 95 | + |
| 96 | +fn declared<T>(value: T) -> Evidence<T> { |
| 97 | + Evidence { |
| 98 | + value: Some(value), |
| 99 | + source: EvidenceSource::Declared, |
| 100 | + ..Default::default() |
| 101 | + } |
| 102 | +} |
| 103 | + |
| 104 | +/// #509 Example 2's design queries (its integer Q3 keeps `COUNT(*) GROUP BY |
| 105 | +/// src_ip` as a target): every candidate builds through Stages 1 and 2. |
| 106 | +/// Stage 3 may still reject one, e.g. for accuracy. |
| 107 | +#[tokio::test] |
| 108 | +async fn example2_design_candidates_all_build() { |
| 109 | + const QUERIES: [&str; 3] = [ |
| 110 | + "SELECT COUNT(DISTINCT src_ip) FROM flows", |
| 111 | + "SELECT -SUM(p * LN(p)) FROM (SELECT COUNT(*) * 1.0 / SUM(COUNT(*)) OVER () AS p FROM flows GROUP BY src_ip)", |
| 112 | + "SELECT SQRT(SUM(c * c)) FROM (SELECT src_ip, COUNT(*) AS c FROM flows GROUP BY src_ip)", |
| 113 | + ]; |
| 114 | + let mut roots = vec![]; |
| 115 | + for (i, sql) in QUERIES.into_iter().enumerate() { |
| 116 | + let root = lower_sql(sql, &catalog(), target()).await.unwrap(); |
| 117 | + roots.push((i, QueryRoot::Operator(root))); |
| 118 | + } |
| 119 | + let demand = vec![ |
| 120 | + RootDemand { |
| 121 | + accuracy: Some(target()), |
| 122 | + recurrence: QueryRecurrence::OneTime { |
| 123 | + invocations: 1, |
| 124 | + execute_at: None, |
| 125 | + }, |
| 126 | + predictability: Predictability::default(), |
| 127 | + }; |
| 128 | + roots.len() |
| 129 | + ]; |
| 130 | + let data = DataWorkload { |
| 131 | + arrival: DataArrival::ContinuouslyIngesting, |
| 132 | + ingestion_rate: declared(Rate(100_000.0)), |
| 133 | + input_cardinality: declared(10_000_000), |
| 134 | + ..Default::default() |
| 135 | + }; |
| 136 | + let run = plan_stages(roots, &demand, &data, PlanningModels::builtin(), 4096).unwrap(); |
| 137 | + let enumeration = run.enumeration.unwrap(); |
| 138 | + let unbuilt: Vec<_> = enumeration |
| 139 | + .selection |
| 140 | + .rejected |
| 141 | + .iter() |
| 142 | + .filter(|r| r.reason.starts_with("Stage ")) |
| 143 | + .collect(); |
| 144 | + assert!(enumeration.candidates.len() > 1); |
| 145 | + assert!(unbuilt.is_empty(), "{unbuilt:#?}"); |
| 146 | +} |
0 commit comments