|
| 1 | +//! End-to-end SQL query-string → post-ASAP IR pin (issue #191). |
| 2 | +//! |
| 3 | +//! The SQL counterpart of `l4_binding.rs`: drives SQL text — `lower_sql` |
| 4 | +//! (text → pre-ASAP `QueryExpr`) → `asap_aware_mapping::implement_tree` |
| 5 | +//! (pre-ASAP → post-ASAP `SummaryExpr`) — and pins the resulting |
| 6 | +//! sketch-vs-exact-accumulator binding shape node by node, the way |
| 7 | +//! `l4_binding.rs` does for PromQL. |
| 8 | +//! |
| 9 | +//! ## A structural wrinkle PromQL doesn't have |
| 10 | +//! |
| 11 | +//! `lower_promql` returns a *bare* `QueryExpr::Aggregate` for a top-level |
| 12 | +//! aggregation (`sum by (job) (m)`, `quantile(0.99, …)`), so `implement_tree` |
| 13 | +//! can bind it directly at the tree root. `lower_sql` never does: DataFusion's |
| 14 | +//! planner always wraps even a single, unaliased aggregate in an identity |
| 15 | +//! `Project` (confirmed below), so a SQL tree's *root* is always `Project { |
| 16 | +//! child: Aggregate { .. } }`. `implement_tree` only fires |
| 17 | +//! `bind_summary_agg` when the node it's looking at is itself a bindable |
| 18 | +//! `QueryExpr::Aggregate` (see `bind.rs`'s module docs on the "logical parent |
| 19 | +//! subsumes bindable child" conservative fallback); a `Project` at the root |
| 20 | +//! is exactly such a logical parent, so feeding a raw `lower_sql` result |
| 21 | +//! straight into `implement_tree` always yields a whole-tree |
| 22 | +//! `SummaryExpr::Logical` — never a genuine sketch or accumulator binding. |
| 23 | +//! |
| 24 | +//! The tests below extract the inner `Aggregate` node the same way this |
| 25 | +//! crate's own `frontend-sql/tests/sql_lowering.rs` does (its |
| 26 | +//! `find_aggregate`/`find_aggregate_node` helpers) and hand that to |
| 27 | +//! `implement_tree` directly, which is the shape a future Project-elision |
| 28 | +//! rewrite (tracked with the rest of the post-ASAP rule engine, issues |
| 29 | +//! #6/#33) would present to this pass in production. |
| 30 | +
|
| 31 | +use asap_aware_mapping::implement_tree; |
| 32 | +use asap_frontend_sql::{lower_sql, SqlCatalog}; |
| 33 | +use asap_types::post_asap::{ |
| 34 | + ExactKind, ExactParams, SketchKind, SketchParams, SketchQuery, SummaryExpr, SummaryFamilyType, |
| 35 | + SummarySchema, |
| 36 | +}; |
| 37 | +use asap_types::pre_asap::expr_ir::ColumnRef; |
| 38 | +use asap_types::pre_asap::query_expr::{QueryExpr, Reduction}; |
| 39 | +use asap_types::pre_asap::schema::{Column, DataType, Schema}; |
| 40 | +use asap_types::types::AccuracyTarget; |
| 41 | + |
| 42 | +fn dtype<'a>(schema: &'a SummarySchema, name: &str) -> &'a SummaryFamilyType { |
| 43 | + &schema |
| 44 | + .fields |
| 45 | + .iter() |
| 46 | + .find(|f| f.name == name) |
| 47 | + .unwrap_or_else(|| panic!("no field {name:?} in {schema:?}")) |
| 48 | + .dtype |
| 49 | +} |
| 50 | + |
| 51 | +fn col(name: &str, dtype: DataType) -> Column { |
| 52 | + Column::new(name, dtype, false) |
| 53 | +} |
| 54 | + |
| 55 | +/// `metrics(ts, service, latency, bytes)` — mirrors |
| 56 | +/// `frontend-sql/tests/sql_lowering.rs`'s catalog. |
| 57 | +fn catalog() -> SqlCatalog { |
| 58 | + SqlCatalog::new().with_table( |
| 59 | + "metrics", |
| 60 | + Schema::with_time_index( |
| 61 | + vec![ |
| 62 | + col("ts", DataType::Timestamp), |
| 63 | + col("service", DataType::Utf8), |
| 64 | + col("latency", DataType::Float64), |
| 65 | + col("bytes", DataType::Int64), |
| 66 | + ], |
| 67 | + 0, |
| 68 | + vec![], |
| 69 | + ), |
| 70 | + ) |
| 71 | +} |
| 72 | + |
| 73 | +async fn lower(sql: &str, accuracy: AccuracyTarget) -> QueryExpr { |
| 74 | + lower_sql(sql, &catalog(), accuracy) |
| 75 | + .await |
| 76 | + .unwrap_or_else(|e| panic!("lower failed for {sql:?}: {e}")) |
| 77 | +} |
| 78 | + |
| 79 | +/// The `Aggregate` node beneath the identity `Project` DataFusion's planner |
| 80 | +/// always wraps a top-level aggregate in — see the module docs above. |
| 81 | +fn inner_aggregate(qe: &QueryExpr) -> &QueryExpr { |
| 82 | + match qe { |
| 83 | + QueryExpr::Project { child, .. } => inner_aggregate(child), |
| 84 | + QueryExpr::Aggregate { .. } => qe, |
| 85 | + other => panic!("expected a Project{{Aggregate}} shape, got {other:?}"), |
| 86 | + } |
| 87 | +} |
| 88 | + |
| 89 | +/// Sanity + documentation: feeding a raw `lower_sql` root straight into |
| 90 | +/// `implement_tree` never binds anything — the wrapping `Project` always |
| 91 | +/// subsumes the `Aggregate` beneath it into one logical passthrough. This is |
| 92 | +/// the "conservative fallback" `bind.rs`'s module docs describe, hitting |
| 93 | +/// unconditionally for SQL because of the Project DataFusion always inserts. |
| 94 | +#[tokio::test] |
| 95 | +async fn sql_full_query_root_stays_logical_under_the_identity_projection() { |
| 96 | + let l3 = lower( |
| 97 | + "SELECT approx_percentile_cont(latency, 0.99) FROM metrics", |
| 98 | + AccuracyTarget::Epsilon(0.01), |
| 99 | + ) |
| 100 | + .await; |
| 101 | + assert!( |
| 102 | + matches!(l3, QueryExpr::Project { .. }), |
| 103 | + "sanity: a SQL root is a Project, unlike lower_promql's bare Aggregate" |
| 104 | + ); |
| 105 | + let root = implement_tree(&l3).expect("binding failed"); |
| 106 | + assert!( |
| 107 | + matches!(root.expr, SummaryExpr::Logical(ref e) if **e == l3), |
| 108 | + "implement_tree does not look inside a Project to find a bindable \ |
| 109 | + Aggregate child, so the whole Project{{Aggregate}} tree stays logical" |
| 110 | + ); |
| 111 | +} |
| 112 | + |
| 113 | +/// `SELECT approx_percentile_cont(latency, 0.99) FROM metrics` at ε = 0.01, |
| 114 | +/// with the wrapping `Project` stripped (see module docs): |
| 115 | +/// |
| 116 | +/// ```text |
| 117 | +/// SummaryEstimate { query: Quantile{0.99} } → {…: Float64} |
| 118 | +/// └─ SummaryAgg { Kll{k:200}, col: metrics.latency } → {…: Sketch(Kll, {k:200})} |
| 119 | +/// └─ Logical(Scan) → {ts, service, latency, bytes} |
| 120 | +/// ``` |
| 121 | +/// |
| 122 | +/// The SQL counterpart of `l4_binding.rs`'s |
| 123 | +/// `promql_quantile_of_rate_binds_kll_over_rate_accumulator`: same intent |
| 124 | +/// (`Quantile`), same KLL sizing (k=200 from ε=0.01), but the summarised |
| 125 | +/// column is the intent's own *named* SQL column rather than PromQL's |
| 126 | +/// synthetic sample value. |
| 127 | +#[tokio::test] |
| 128 | +async fn sql_quantile_binds_kll_sketch_over_named_column() { |
| 129 | + let l3 = lower( |
| 130 | + "SELECT approx_percentile_cont(latency, 0.99) FROM metrics", |
| 131 | + AccuracyTarget::Epsilon(0.01), |
| 132 | + ) |
| 133 | + .await; |
| 134 | + let agg = inner_aggregate(&l3); |
| 135 | + let root = implement_tree(agg).expect("binding failed"); |
| 136 | + |
| 137 | + let SummaryExpr::SummaryEstimate { |
| 138 | + summary_input, |
| 139 | + query, |
| 140 | + } = &root.expr |
| 141 | + else { |
| 142 | + panic!("expected SummaryEstimate root, got {:?}", root.expr); |
| 143 | + }; |
| 144 | + assert!(matches!(query, SketchQuery::Quantile { q } if *q == 0.99)); |
| 145 | + assert_eq!( |
| 146 | + root.schema.fields.len(), |
| 147 | + 1, |
| 148 | + "no GROUP BY — a single output column" |
| 149 | + ); |
| 150 | + assert_eq!( |
| 151 | + root.schema.fields[0].dtype, |
| 152 | + SummaryFamilyType::Plain(DataType::Float64), |
| 153 | + "the summary-state type must not propagate past the estimate" |
| 154 | + ); |
| 155 | + |
| 156 | + let SummaryExpr::SummaryAgg { |
| 157 | + child, |
| 158 | + family, |
| 159 | + col, |
| 160 | + reduction, |
| 161 | + } = &summary_input.expr |
| 162 | + else { |
| 163 | + panic!("expected SummaryAgg, got {:?}", summary_input.expr); |
| 164 | + }; |
| 165 | + assert_eq!( |
| 166 | + family, |
| 167 | + &SummaryFamilyType::Sketch(SketchKind::Kll, SketchParams::Kll { k: 200 }) |
| 168 | + ); |
| 169 | + assert_eq!( |
| 170 | + col, |
| 171 | + &ColumnRef::Qualified { |
| 172 | + table: "metrics".into(), |
| 173 | + name: "latency".into(), |
| 174 | + }, |
| 175 | + "SQL binds the intent's own named input column, not a synthetic sample value" |
| 176 | + ); |
| 177 | + assert_eq!( |
| 178 | + reduction, |
| 179 | + &Reduction::by(vec![]), |
| 180 | + "global quantile — no GROUP BY, full reduction" |
| 181 | + ); |
| 182 | + assert_eq!( |
| 183 | + summary_input.schema.fields[0].dtype, |
| 184 | + SummaryFamilyType::Sketch(SketchKind::Kll, SketchParams::Kll { k: 200 }) |
| 185 | + ); |
| 186 | + |
| 187 | + let SummaryExpr::Logical(logical_leaf) = &child.expr else { |
| 188 | + panic!("expected Logical leaf, got {:?}", child.expr); |
| 189 | + }; |
| 190 | + assert!(matches!(logical_leaf.as_ref(), QueryExpr::Scan { .. })); |
| 191 | + assert!( |
| 192 | + child |
| 193 | + .schema |
| 194 | + .fields |
| 195 | + .iter() |
| 196 | + .all(|f| matches!(f.dtype, SummaryFamilyType::Plain(_))), |
| 197 | + "logical edges carry only plain columns" |
| 198 | + ); |
| 199 | +} |
| 200 | + |
| 201 | +/// `SELECT COUNT(DISTINCT service) FROM metrics` at ε = 0.01 lowers to |
| 202 | +/// `AggIntent::Cardinality` (`sql_lowering.rs::count_distinct_is_cardinality`) |
| 203 | +/// — unlike `Quantile`/`Count`/`TopK`, its preferred candidate is HLL, not |
| 204 | +/// KLL/CMS (`boundary::summary_candidates`), so this exercises a distinct |
| 205 | +/// branch of the sketch-vs-exact boundary than the quantile test above. |
| 206 | +#[tokio::test] |
| 207 | +async fn sql_count_distinct_binds_hll_sketch_over_named_column() { |
| 208 | + let l3 = lower( |
| 209 | + "SELECT COUNT(DISTINCT service) FROM metrics", |
| 210 | + AccuracyTarget::Epsilon(0.01), |
| 211 | + ) |
| 212 | + .await; |
| 213 | + let agg = inner_aggregate(&l3); |
| 214 | + let root = implement_tree(agg).expect("binding failed"); |
| 215 | + |
| 216 | + let SummaryExpr::SummaryEstimate { |
| 217 | + summary_input, |
| 218 | + query, |
| 219 | + } = &root.expr |
| 220 | + else { |
| 221 | + panic!("expected SummaryEstimate root, got {:?}", root.expr); |
| 222 | + }; |
| 223 | + assert!(matches!(query, SketchQuery::Cardinality)); |
| 224 | + assert_eq!( |
| 225 | + root.schema.fields[0].dtype, |
| 226 | + SummaryFamilyType::Plain(DataType::Int64), |
| 227 | + "COUNT(DISTINCT …) reads back out as an integer count" |
| 228 | + ); |
| 229 | + |
| 230 | + let SummaryExpr::SummaryAgg { |
| 231 | + family, |
| 232 | + col, |
| 233 | + reduction, |
| 234 | + .. |
| 235 | + } = &summary_input.expr |
| 236 | + else { |
| 237 | + panic!("expected SummaryAgg, got {:?}", summary_input.expr); |
| 238 | + }; |
| 239 | + assert_eq!( |
| 240 | + family, |
| 241 | + &SummaryFamilyType::Sketch(SketchKind::Hll, SketchParams::Hll { precision: 14 }) |
| 242 | + ); |
| 243 | + assert_eq!( |
| 244 | + col, |
| 245 | + &ColumnRef::Qualified { |
| 246 | + table: "metrics".into(), |
| 247 | + name: "service".into(), |
| 248 | + } |
| 249 | + ); |
| 250 | + assert_eq!(reduction, &Reduction::by(vec![])); |
| 251 | +} |
| 252 | + |
| 253 | +/// An exact workload binds zero sketches: `SUM(bytes) GROUP BY service` at |
| 254 | +/// `AccuracyTarget::Exact` still gets its mergeable exact accumulator, and |
| 255 | +/// `AVG(bytes)` (non-mergeable) stays a whole logical subtree untouched. SQL |
| 256 | +/// counterpart of `l4_binding.rs`'s |
| 257 | +/// `promql_exact_workload_binds_accumulators_not_sketches`. |
| 258 | +#[tokio::test] |
| 259 | +async fn sql_exact_workload_binds_accumulators_not_sketches() { |
| 260 | + let l3 = lower( |
| 261 | + "SELECT service, SUM(bytes) FROM metrics GROUP BY service", |
| 262 | + AccuracyTarget::Exact, |
| 263 | + ) |
| 264 | + .await; |
| 265 | + let agg = inner_aggregate(&l3); |
| 266 | + let root = implement_tree(agg).expect("binding failed"); |
| 267 | + let SummaryExpr::SummaryAgg { |
| 268 | + family, reduction, .. |
| 269 | + } = &root.expr |
| 270 | + else { |
| 271 | + panic!("expected SummaryAgg, got {:?}", root.expr); |
| 272 | + }; |
| 273 | + assert_eq!( |
| 274 | + family, |
| 275 | + &SummaryFamilyType::ExactAggregate(ExactKind::Sum, ExactParams::Sum) |
| 276 | + ); |
| 277 | + assert_eq!( |
| 278 | + reduction, |
| 279 | + &Reduction::by(vec![1]), |
| 280 | + "service is col 1 in [ts, service, latency, bytes]" |
| 281 | + ); |
| 282 | + assert_eq!( |
| 283 | + dtype(&root.schema, "service"), |
| 284 | + &SummaryFamilyType::Plain(DataType::Utf8), |
| 285 | + "group keys pass through verbatim" |
| 286 | + ); |
| 287 | + |
| 288 | + let l3 = lower("SELECT AVG(bytes) FROM metrics", AccuracyTarget::Exact).await; |
| 289 | + let agg = inner_aggregate(&l3); |
| 290 | + let root = implement_tree(agg).expect("binding failed"); |
| 291 | + assert!( |
| 292 | + matches!(root.expr, SummaryExpr::Logical(_)), |
| 293 | + "avg has no mergeable accumulator — stays logical" |
| 294 | + ); |
| 295 | +} |
0 commit comments