diff --git a/crates/types/src/ir/canonicalize.rs b/crates/types/src/ir/canonicalize.rs new file mode 100644 index 000000000..d9a0034a4 --- /dev/null +++ b/crates/types/src/ir/canonicalize.rs @@ -0,0 +1,1077 @@ +//! Post-lowering canonicalization of the operator DAG. +//! +//! Erases *structural* differences between semantically identical queries so +//! a post-ASAP binding rule matching on the intent algebra sees one canonical +//! spelling regardless of source language (issue #34). +//! +//! ## Heavy-hitter promotion +//! +//! An additive-ranked "order by the aggregate, take the top k" is a +//! heavy-hitter represented by [`AggIntent::TopK`]. Front ends may emit it as +//! an ordinary `Limit { Sort { … Aggregate } }`; this pass promotes that shape +//! to the canonical +//! +//! ```text +//! Aggregate { reduction: Reduce(), measures: [TopK{k}], +//! child: Aggregate { measures: [Count | Sum], … } } +//! ``` +//! +//! Count supplies unit weights and Sum supplies value weights. Because the +//! match is positional, aliases do not affect it. Other ranked expressions +//! retain Sort + Limit. +//! +//! ## Subquery lowering +//! +//! EXISTS/NOT EXISTS and positive IN filter conjuncts may use semi/anti joins. +//! Scalar subqueries remain explicit: a cross join does not preserve their +//! zero-row NULL or multiple-row error semantics. All scalar plan references +//! participate in DAG traversal and canonicalization. + +use std::collections::HashMap; +use std::rc::Rc; + +use super::node::{Operator, OperatorNode}; +use super::non_asap::NonASAPOp; +use super::scalar::{ExprSemantics, Predicate, ProjectItem, ScalarExpr, SortKey}; +use crate::ir::operator_properties::{JoinKind, Reduction}; +use crate::ir::SchemaDerivationError; +use crate::pre_asap::agg_intent::{topk, AggIntent}; +use crate::pre_asap::expr_ir::{CompareOpKind, ScalarValue}; +use crate::types::AccuracyTarget; + +/// Rewrite the DAG under `root` into its canonical form (bottom-up). +/// Idempotent: an already-canonical DAG comes back as the same `Rc`. Only +/// nodes that change (or whose inputs change) are rebuilt; every untouched +/// sub-DAG keeps its pointer identity, and a shared sub-DAG that is rewritten +/// stays shared. +pub fn canonicalize(root: Rc) -> Result, SchemaDerivationError> { + canon(&root, &mut HashMap::new()) +} + +fn canon( + node: &Rc, + memo: &mut HashMap<*const OperatorNode, Rc>, +) -> Result, SchemaDerivationError> { + if let Some(done) = memo.get(&Rc::as_ptr(node)) { + return Ok(Rc::clone(done)); + } + + // A `Concat` asserting a caller-proven `discriminator_unique_key` (issue + // #228) had that key's `ColumnId`s resolved against exactly the first + // branch's output schema *as it stood before this pass ran*. The rewrites + // below can restructure that branch (anywhere within it) into a shape + // with a different output schema, which would leave those `ColumnId`s + // pointing at the wrong column, or out of bounds. Snapshot the schema the + // key was resolved against before recursing into the children. + let discriminator_branch_schema_before = match &node.operator { + Operator::NonASAP(NonASAPOp::Concat { + children, + discriminator_unique_key: Some(_), + }) => children.first().map(|c| c.schema.clone()), + _ => None, + }; + + // Bottom-up: canonicalize every operator input before matching at this + // node, so an inner heavy-hitter is promoted before an enclosing rewrite + // inspects it. + let mut rebuilt: Vec<(*const OperatorNode, Rc)> = Vec::new(); + let mut changed = false; + for child in operator_children(&node.operator) { + let new = canon(child, memo)?; + changed |= !Rc::ptr_eq(&new, child); + rebuilt.push((Rc::as_ptr(child), new)); + } + let mut current = if changed { + // `map_children` also visits the operator nodes referenced from + // scalar expressions; those are not in `rebuilt` and pass through + // unchanged. (A node that is both an operator input and a scalar + // reference is one shared node, so it takes its canonical form in + // both places.) + let rebuilt_child = |c: &Rc| { + rebuilt + .iter() + .find(|(ptr, _)| *ptr == Rc::as_ptr(c)) + .map_or_else(|| Rc::clone(c), |(_, new)| Rc::clone(new)) + }; + Rc::new(node.map_children(rebuilt_child)?) + } else { + Rc::clone(node) + }; + + // If the first branch's output schema moved out from under the asserted + // key, the key can no longer be trusted — drop it (never re-derive it by + // guessing at name/position). A wrong `unique_keys` claim is a wrong + // query answer, not a missed optimization, so any difference at all + // drops the key. + if let Operator::NonASAP(NonASAPOp::Concat { + children, + discriminator_unique_key: Some(_), + }) = ¤t.operator + { + let after = children.first().map(|c| &c.schema); + if discriminator_branch_schema_before.as_ref() != after { + current = OperatorNode::new_shared(crate::ir::Operator::NonASAP(NonASAPOp::Concat { + children: children.clone(), + discriminator_unique_key: None, + }))?; + } + } + + let current = apply_local_rules(current, memo)?; + + memo.insert(Rc::as_ptr(node), Rc::clone(¤t)); + Ok(current) +} + +type Memo = HashMap<*const OperatorNode, Rc>; + +/// Apply the local rewrite rules at `node` (whose inputs are already +/// canonical) until none matches. The rules chain: a `ROW_NUMBER()`- +/// partitioned top-k rewrites to a `Limit{Sort}`, which the heavy-hitter +/// rule may then promote to an `Aggregate([TopK])`; a `Filter` with several +/// subquery conjuncts sheds one per round. Each rule strictly simplifies the +/// node (one fewer idiom, or one fewer subquery reference), so the loop +/// terminates. +fn apply_local_rules( + mut current: Rc, + memo: &mut Memo, +) -> Result, SchemaDerivationError> { + loop { + let next = if let Some(next) = try_promote_additive_top_ranking(¤t)? { + next + } else if let Some(next) = try_lower_subquery_conjunct(¤t, memo)? { + next + } else { + break; + }; + current = next; + } + Ok(current) +} + +/// The direct **operator** inputs of a node — the relational skeleton only. +/// Operator nodes referenced from a scalar position (`ScalarSubquery`, +/// `Exists`, …) are not visited here: a subquery that the lowering rules +/// lift into a join is canonicalized at that point, and one they leave in +/// place (`NOT IN`, an `EXISTS` outside a `Filter` conjunct) stays as the +/// front end emitted it. +fn operator_children(op: &Operator) -> Vec<&Rc> { + op.children() +} + +/// Recognise an additive-ranked +/// `Limit { Sort { [Project] Aggregate([Count | Sum]) } }` and rewrite it to +/// the canonical heavy-hitter `Aggregate([TopK])` over the explicit inner +/// aggregate. Returns `None` when the shape does not match. +fn try_promote_additive_top_ranking( + node: &OperatorNode, +) -> Result>, SchemaDerivationError> { + // Limit k, no offset (an OFFSET means "not the top k"). + let Some(NonASAPOp::Limit { + n: Some(k), + offset: 0, + partition_by: limit_partition, + child, + }) = node.non_asap() + else { + return Ok(None); + }; + // A single ordering key on a column. + let Some(NonASAPOp::Sort { + keys, + partition_by, + child: sort_child, + }) = child.non_asap() + else { + return Ok(None); + }; + // A per-group `Limit` must agree with its `Sort`'s partition: the + // ranking's partition is what the outer `TopK` groups by. + if !limit_partition.is_empty() && limit_partition != partition_by { + return Ok(None); + } + let [SortKey { + expr: ScalarExpr::Column(sort_col), + ascending, + .. + }] = keys.as_slice() + else { + return Ok(None); + }; + + // The ordered relation is an `Aggregate`, optionally behind a passthrough + // projection (a bare-column SELECT list). Map the sort key through the + // projection to the aggregate's own output column. + let (agg_node, ranked_col) = match sort_child.non_asap() { + Some(NonASAPOp::Project { cols, child, .. }) => { + let Some(ProjectItem { + expr: ScalarExpr::Column(underlying), + .. + }) = cols.get(*sort_col) + else { + return Ok(None); + }; + (child, *underlying) + } + _ => (sort_child, *sort_col), + }; + + // Exactly one aggregate, ranked by *its* output column — the measure sits + // at index `by.len()` (after the group keys). A `PerEntity` reduction has + // no `by` to rank a measure against, so it is a non-match. + let Some(NonASAPOp::Aggregate { + reduction, + measures, + child: aggregate_child, + .. + }) = agg_node.non_asap() + else { + return Ok(None); + }; + let Reduction::Reduce(by) = reduction else { + return Ok(None); + }; + let [ranked_agg] = measures.as_slice() else { + return Ok(None); + }; + if ranked_col != by.len() { + return Ok(None); + } + // The heavy-hitter decision — descending, over a measure with a realised + // heavy-hitter sketch — is the shared rule both front ends consult (issue + // #38). An ascending additive-ranked limit (bottom-k) stays generic. + if !topk::Ranking::from_aggregate(ranked_agg).is_supported(!ascending) { + return Ok(None); + } + // A direct Sum is a stream of additive observation weights. A Sum over a + // derived child such as Rate/Increase still needs exact reset-aware + // values to rerank sketch candidates, and the post-ASAP IR has no + // candidate-sidecar + exact-rerank node, so that shape keeps Sort + Limit. + if matches!(ranked_agg, AggIntent::Sum { .. }) + && matches!( + aggregate_child.non_asap(), + Some(NonASAPOp::Aggregate { .. }) + ) + { + return Ok(None); + } + // Count ranks unit updates; a direct Sum ranks weighted updates. + let accuracy = match ranked_agg { + AggIntent::Count { accuracy } => accuracy.clone(), + AggIntent::Sum { .. } => AccuracyTarget::Exact, + _ => unreachable!("additive ranking gate admitted a non-additive measure"), + }; + + // Outer heavy-hitter `TopK`, grouped by the ranking's partition (empty for + // a global `ORDER BY … LIMIT k`; the `by` labels for a partitioned `topk + // by`), over the unchanged inner additive aggregate. + OperatorNode::new_shared(crate::ir::Operator::NonASAP(NonASAPOp::Aggregate { + reduction: Reduction::by(partition_by.to_vec()), + measures: vec![AggIntent::TopK { k: *k, accuracy }], + output_names: Vec::new(), + filters: vec![], + having: None, + child: Rc::clone(agg_node), + })) + .map(Some) +} + +// ROW_NUMBER filters retain the window output. Eliminating it without a +// consumer-aware rewrite drops a visible column and invalidates outer scopes. + +/// `Predicate(true)`: the unconditional join predicate the SQL front end +/// emits for an uncorrelated `EXISTS` and for a `CROSS JOIN`. +fn always_true() -> Predicate { + Predicate(ScalarExpr::Literal(ScalarValue::Boolean(true))) +} + +/// Whether `conjunct` is one this pass lowers to a semi/anti join. +fn is_join_conjunct(conjunct: &ScalarExpr) -> bool { + match conjunct { + ScalarExpr::Exists { .. } => true, + // An `IN` whose probe expression itself reads a scalar subquery is + // lowered only after that subquery has been joined in by + // `try_lower_scalar_subquery` (a `Join` predicate is not a place that + // rule looks). `NOT IN` is never lowered — see the module docs. + ScalarExpr::InSubquery { + expr, + negated: false, + .. + } => find_scalar_subquery(expr).is_none(), + _ => false, + } +} + +/// Lower one `[NOT] EXISTS (s)` / `x IN (s)` conjunct of a `Filter` to the +/// semi-/anti-join the SQL front end used to emit directly. The remaining +/// conjuncts stay in an outer `Filter` over the join: a semi/anti join's +/// output schema is the left's, so their column ids are unchanged. One +/// conjunct per call; the fixpoint loop picks up the next. +fn try_lower_subquery_conjunct( + node: &OperatorNode, + memo: &mut Memo, +) -> Result>, SchemaDerivationError> { + let Some(NonASAPOp::Filter { + pred: Predicate(pred), + child, + }) = node.non_asap() + else { + return Ok(None); + }; + let conjuncts = pred.conjuncts(); + let Some(idx) = conjuncts.iter().position(is_join_conjunct) else { + return Ok(None); + }; + let left_width = child.schema.fields.len(); + let (kind, subquery, join_pred) = match &conjuncts[idx] { + // Uncorrelated by construction (the IR's `Exists` carries no outer + // column references), so the join condition is unconditionally true. + ScalarExpr::Exists { subquery, negated } => { + let kind = if *negated { + JoinKind::Anti + } else { + JoinKind::Semi + }; + (kind, subquery, always_true()) + } + // `x = `, which sits right after the + // left's columns in the `left ++ right` scope the predicate resolves + // against. + ScalarExpr::InSubquery { expr, subquery, .. } => ( + JoinKind::Semi, + subquery, + Predicate(ScalarExpr::Compare { + left: expr.clone(), + op: CompareOpKind::Eq, + right: Box::new(ScalarExpr::Column(left_width)), + semantics: ExprSemantics::Sql, + }), + ), + _ => unreachable!("`is_join_conjunct` admitted a non-subquery conjunct"), + }; + let join = OperatorNode::new_shared(crate::ir::Operator::NonASAP(NonASAPOp::Join { + kind, + pred: join_pred, + left: Rc::clone(child), + right: canon(subquery, memo)?, + }))?; + let mut rest: Vec = conjuncts + .iter() + .enumerate() + .filter(|(i, _)| *i != idx) + .map(|(_, c)| c.clone()) + .collect(); + let out = match rest.len() { + 0 => join, + 1 => filter(rest.remove(0), join)?, + _ => filter(ScalarExpr::BoolAnd(rest), join)?, + }; + Ok(Some(out)) +} + +/// Lower one scalar subquery read by a `Project` item or a `Filter` +/// predicate: the owner reads it through a cross join against the subquery, +/// whose single column is appended after the left's (`Column(|left|)`), and +/// every occurrence of that subquery node in the owner is replaced by that +/// column reference. One subquery node per call; the fixpoint loop handles +/// the rest, each getting its own cross join further out (so earlier column +/// ids are never shifted). For a `Filter` the output schema is restored to +/// the left's columns by a positional `Project` over the result. +/// +/// Not representable in the IR, and therefore not checked here: SQL raises +/// an error when a scalar subquery yields more than one row (the cross join +/// would duplicate the left's rows instead), and yields NULL when it yields +/// none (the cross join yields no rows instead). +fn filter( + pred: ScalarExpr, + child: Rc, +) -> Result, SchemaDerivationError> { + OperatorNode::new_shared(crate::ir::Operator::NonASAP(NonASAPOp::Filter { + pred: Predicate(pred), + child, + })) +} + +/// The first `ScalarSubquery` node read by `expr` (pre-order over its scalar +/// children; referenced operator subgraphs are their own scope and are not +/// entered). +fn find_scalar_subquery(expr: &ScalarExpr) -> Option<&Rc> { + if let ScalarExpr::ScalarSubquery(node) = expr { + return Some(node); + } + expr.children().into_iter().find_map(find_scalar_subquery) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ir::operator_properties::WindowFuncKind; + use crate::ir::operator_properties::{ + ConcatDiscriminatorKey, GroupKeys, Source, WindowFrame, WindowFrameBound, + WindowFrameOffset, WindowFrameUnits, + }; + use crate::pre_asap::schema::{DataType, Field, Schema}; + + fn node(op: NonASAPOp) -> Rc { + Rc::new(OperatorNode::new(Operator::NonASAP(op)).expect("fixture derives a schema")) + } + + fn scan() -> Rc { + node(NonASAPOp::Scan { + source: Source::TimeSeries { metric: "m".into() }, + predicates: vec![], + schema: Schema::with_time_index( + vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("service", DataType::Utf8, false), + Field::plain("value", DataType::Float64, false), + ], + 0, + vec![], + ), + }) + } + + fn aggregate( + reduction: Reduction, + agg: AggIntent, + child: Rc, + ) -> Rc { + node(NonASAPOp::Aggregate { + reduction, + measures: vec![agg], + output_names: vec![], + filters: vec![], + having: None, + child, + }) + } + + fn count() -> AggIntent { + AggIntent::Count { + accuracy: AccuracyTarget::Exact, + } + } + + /// `Aggregate{ by: [1], [Count] }` over the scan — output cols `[service, count]`. + fn count_by_service() -> Rc { + aggregate(Reduction::by(vec![1]), count(), scan()) + } + + fn key(col: usize, ascending: bool) -> Vec { + vec![SortKey { + expr: ScalarExpr::Column(col), + ascending, + nulls_first: false, + }] + } + + fn desc(col: usize) -> Vec { + key(col, false) + } + + fn limit(n: usize, offset: usize, child: Rc) -> Rc { + node(NonASAPOp::Limit { + n: Some(n), + offset, + partition_by: GroupKeys::none(), + child, + }) + } + + fn sort(keys: Vec, child: Rc) -> Rc { + node(NonASAPOp::Sort { + keys, + partition_by: GroupKeys::none(), + child, + }) + } + + fn passthrough_project(child: Rc) -> Rc { + node(NonASAPOp::Project { + cols: vec![ + ProjectItem { + alias: None, + expr: ScalarExpr::Column(0), + }, + ProjectItem { + alias: Some("c".into()), + expr: ScalarExpr::Column(1), + }, + ], + qualifier: None, + child, + }) + } + + fn concat( + children: Vec>, + key: Option, + ) -> Rc { + node(NonASAPOp::Concat { + children, + discriminator_unique_key: key, + }) + } + + fn measures(n: &OperatorNode) -> &[AggIntent] { + match n.non_asap() { + Some(NonASAPOp::Aggregate { measures, .. }) => measures, + _ => &[], + } + } + + fn is_topk_over_count(n: &OperatorNode) -> bool { + let Some(NonASAPOp::Aggregate { + measures, child, .. + }) = n.non_asap() + else { + return false; + }; + matches!(measures.as_slice(), [AggIntent::TopK { k: 5, .. }]) + && matches!(self::measures(child), [AggIntent::Count { .. }]) + } + + #[test] + fn promotes_count_ranked_limit_sort() { + // Limit 5 { Sort DESC by count-col (1) { Aggregate[Count] by [1] } }. + let q = limit(5, 0, sort(desc(1), count_by_service())); + assert!(is_topk_over_count(&canonicalize(q).unwrap())); + } + + #[test] + fn promotes_through_a_passthrough_projection() { + // …with a `SELECT service, count` projection between the Sort and the Agg. + let q = limit(5, 0, sort(desc(1), passthrough_project(count_by_service()))); + assert!(is_topk_over_count(&canonicalize(q).unwrap())); + } + + #[test] + fn promoted_topk_reuses_the_inner_aggregate_node() { + // The inner aggregate is untouched, so the rewrite shares it rather + // than copying it. + let agg = count_by_service(); + let out = canonicalize(limit(5, 0, sort(desc(1), Rc::clone(&agg)))).unwrap(); + let Some(NonASAPOp::Aggregate { child, .. }) = out.non_asap() else { + panic!("expected TopK aggregate"); + }; + assert!(Rc::ptr_eq(child, &agg)); + } + + #[test] + fn is_idempotent() { + let q = limit(5, 0, sort(desc(1), count_by_service())); + let once = canonicalize(q).unwrap(); + let twice = canonicalize(Rc::clone(&once)).unwrap(); + assert!(Rc::ptr_eq(&once, &twice), "canonicalize must be idempotent"); + } + + #[test] + fn untouched_dag_is_returned_pointer_equal() { + // Nothing here matches a rewrite: a Concat of two projections over + // one shared aggregate. The root (and everything under it) must come + // back as the same `Rc`. + let agg = count_by_service(); + let q = concat( + vec![ + passthrough_project(Rc::clone(&agg)), + passthrough_project(Rc::clone(&agg)), + ], + None, + ); + let out = canonicalize(Rc::clone(&q)).unwrap(); + assert!(Rc::ptr_eq(&out, &q)); + } + + #[test] + fn rewritten_shared_subtree_stays_shared() { + // One promotable sub-DAG referenced twice is rewritten once. + let branch = limit(5, 0, sort(desc(1), count_by_service())); + let q = concat(vec![Rc::clone(&branch), Rc::clone(&branch)], None); + let out = canonicalize(q).unwrap(); + let Some(NonASAPOp::Concat { children, .. }) = out.non_asap() else { + panic!("expected Concat"); + }; + assert!(is_topk_over_count(&children[0])); + assert!(Rc::ptr_eq(&children[0], &children[1])); + } + + // ── Concat's discriminator_unique_key vs. canonicalize (issue #228) ── + // + // `discriminator_unique_key`'s `ColumnId`s were resolved against the + // first branch's *pre-canonicalize* output schema. The key is dropped + // whenever that branch's schema actually changed, and survives untouched + // otherwise. Never guessed at. + + fn discriminator_key(n: &OperatorNode) -> &Option { + match n.non_asap() { + Some(NonASAPOp::Concat { + discriminator_unique_key, + .. + }) => discriminator_unique_key, + _ => panic!("expected Concat"), + } + } + + #[test] + fn concat_discriminator_key_survives_canonicalize_when_first_branch_is_unaffected() { + // A plain `Aggregate` first branch matches neither rewrite trigger, + // so its schema is identical before and after canonicalize. + let q = concat( + vec![count_by_service(), count_by_service()], + Some(ConcatDiscriminatorKey::new(0, vec![1])), + ); + let out = canonicalize(Rc::clone(&q)).unwrap(); + assert!( + discriminator_key(&out).is_some(), + "an untouched first branch's discriminator key must survive canonicalize" + ); + assert!(Rc::ptr_eq(&out, &q)); + } + + #[test] + fn concat_discriminator_key_is_dropped_when_first_branch_gets_rewritten() { + // The first branch is exactly the heavy-hitter promotion trigger, so + // canonicalize rewrites it to `Aggregate{TopK}`, whose own output is + // a single column, not the original two (`[service, count]`). A key + // resolved against the 2-column shape must not survive pointing at + // the new 1-column schema. + let promotable_branch = limit(5, 0, sort(desc(1), count_by_service())); + let q = concat( + vec![promotable_branch, count_by_service()], + Some(ConcatDiscriminatorKey::new(0, vec![1])), + ); + let out = canonicalize(q).unwrap(); + let Some(NonASAPOp::Concat { + children, + discriminator_unique_key, + }) = out.non_asap() + else { + panic!("expected Concat"); + }; + assert!( + is_topk_over_count(&children[0]), + "the first branch is still promoted normally" + ); + assert!( + discriminator_unique_key.is_none(), + "a stale discriminator key must be dropped, never silently kept wrong" + ); + assert!( + out.schema.unique_keys.is_empty(), + "the dropped key leaves the schema" + ); + } + + #[test] + fn does_not_promote_ascending_sort() { + // Ascending = bottom-k: the Top-K ranking rule rejects it (needs + // descending), so it stays a generic Sort+Limit (issue #38). + let q = limit(5, 0, sort(key(1, true), count_by_service())); + assert!(!is_topk_over_count(&canonicalize(q).unwrap())); + } + + #[test] + fn does_not_promote_with_offset() { + let q = limit(5, 2, sort(desc(1), count_by_service())); + assert!(!is_topk_over_count(&canonicalize(q).unwrap())); + } + + #[test] + fn does_not_promote_ranking_by_a_group_key() { + // DESC by col 0 (the `service` group key), not the count → not a + // frequency heavy-hitter. + let q = limit(5, 0, sort(desc(0), count_by_service())); + assert!(!is_topk_over_count(&canonicalize(q).unwrap())); + } + + #[test] + fn does_not_promote_when_limit_partition_disagrees_with_sort() { + // A per-group Limit partitioned differently from its Sort is not the + // top-k shape. + let q = node(NonASAPOp::Limit { + n: Some(5), + offset: 0, + partition_by: GroupKeys::by(vec![0]), + child: sort(desc(1), count_by_service()), + }); + assert!(!is_topk_over_count(&canonicalize(q).unwrap())); + } + + #[test] + fn promotes_sum_ranked_limit_sort_as_weighted_heavy_hitter() { + let sum = aggregate(Reduction::by(vec![1]), AggIntent::Sum { col: None }, scan()); + let out = canonicalize(limit(5, 0, sort(desc(1), sum))).unwrap(); + let Some(NonASAPOp::Aggregate { + measures, child, .. + }) = out.non_asap() + else { + panic!("expected weighted TopK aggregate"); + }; + assert!(matches!( + measures.as_slice(), + [AggIntent::TopK { k: 5, .. }] + )); + assert!(matches!(self::measures(child), [AggIntent::Sum { .. }])); + } + + #[test] + fn keeps_sum_over_counter_reduction_as_exact_value_ranking() { + for counter in [AggIntent::Rate, AggIntent::Increase] { + let derived = aggregate(Reduction::PerEntity, counter, scan()); + let sum = aggregate( + Reduction::by(vec![1]), + AggIntent::Sum { col: None }, + derived, + ); + let out = canonicalize(limit(5, 0, sort(desc(1), sum))).unwrap(); + let Some(NonASAPOp::Limit { child, .. }) = out.non_asap() else { + panic!("expected Limit, got {out:?}"); + }; + let Some(NonASAPOp::Sort { child, .. }) = child.non_asap() else { + panic!("expected Sort under the Limit"); + }; + let Some(NonASAPOp::Aggregate { + measures, child, .. + }) = child.non_asap() + else { + panic!("expected Aggregate under the Sort"); + }; + assert!(matches!(measures.as_slice(), [AggIntent::Sum { .. }])); + assert!(matches!( + child.non_asap(), + Some(NonASAPOp::Aggregate { .. }) + )); + } + } + + // ── ROW_NUMBER() partitioned top-k (issue #24) ────────────────────────── + + /// A scan with `[ts, service, region, value]`. + fn scan4() -> Rc { + node(NonASAPOp::Scan { + source: Source::TimeSeries { metric: "m".into() }, + predicates: vec![], + schema: Schema::with_time_index( + vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("service", DataType::Utf8, false), + Field::plain("region", DataType::Utf8, false), + Field::plain("value", DataType::Float64, false), + ], + 0, + vec![], + ), + }) + } + + /// `Aggregate{ by: [1,2] (service, region), [agg] }` — output `[service, + /// region, ]` (3 cols), so a ROW_NUMBER over it appends `rn` at index 3. + fn grouped(agg: AggIntent) -> Rc { + aggregate(Reduction::by(vec![1, 2]), agg, scan4()) + } + + /// `ROW_NUMBER` ignores its frame clause; any concrete frame works. + fn rownumber_frame() -> WindowFrame { + WindowFrame { + units: WindowFrameUnits::Rows, + start_bound: WindowFrameBound::Preceding(WindowFrameOffset::Scalar(ScalarValue::Null)), + end_bound: WindowFrameBound::Following(WindowFrameOffset::Scalar(ScalarValue::Null)), + } + } + + /// `SQLWindowFunc{ RowNumber, PARTITION BY region(2), ORDER BY col(2) DESC } { agg }`. + fn rownumber_window(agg: Rc) -> Rc { + node(NonASAPOp::SQLWindowFunc { + func: WindowFuncKind::RowNumber, + args: vec![], + partition_by: GroupKeys::by(vec![2]), // region + order_by: vec![SortKey { + expr: ScalarExpr::Column(2), // the aggregate output column + ascending: false, + nulls_first: true, + }], + frame: Some(rownumber_frame()), + output_name: "rn".into(), + child: agg, + }) + } + + /// `Filter{ col <= 5 } { child }`. + fn filter_le_5(col: usize, child: Rc) -> Rc { + node(NonASAPOp::Filter { + pred: Predicate(ScalarExpr::Compare { + left: Box::new(ScalarExpr::Column(col)), + op: CompareOpKind::Le, + right: Box::new(ScalarExpr::Literal(ScalarValue::Int64(5))), + semantics: ExprSemantics::Sql, + }), + child, + }) + } + + /// `Filter{ rn(3) <= 5 } { ROW_NUMBER window { agg } }`. + fn rownumber_topk(agg: Rc) -> Rc { + filter_le_5(3, rownumber_window(agg)) + } + + #[test] + fn rownumber_count_topk_becomes_a_partitioned_heavy_hitter() { + let original = rownumber_topk(grouped(count())); + let out = canonicalize(Rc::clone(&original)).unwrap(); + assert_eq!(out.schema, original.schema); + assert!( + matches!(out.non_asap(),Some(NonASAPOp::Filter { child,.. }) if matches!(child.non_asap(),Some(NonASAPOp::SQLWindowFunc { .. }))) + ); + assert_idempotent(&out); + } + + #[test] + fn rownumber_avg_topk_becomes_a_partitioned_sort_limit() { + let original = rownumber_topk(grouped(AggIntent::Avg { col: None })); + let out = canonicalize(Rc::clone(&original)).unwrap(); + assert_eq!(out.schema, original.schema); + assert!( + matches!(out.non_asap(),Some(NonASAPOp::Filter { child,.. }) if matches!(child.non_asap(),Some(NonASAPOp::SQLWindowFunc { .. }))) + ); + assert_idempotent(&out); + } + + #[test] + fn filter_on_a_non_rownumber_column_is_left_alone() { + // `WHERE service_len <= 5` (col 0, not the rn window column) must not + // be mistaken for a top-k. + let q = filter_le_5(0, rownumber_window(grouped(count()))); + let out = canonicalize(Rc::clone(&q)).unwrap(); + assert!(Rc::ptr_eq(&out, &q), "left as the same Filter"); + } + + // ── Subquery lowering ─────────────────────────────────────────────────── + + fn filter_of(pred: ScalarExpr, child: Rc) -> Rc { + node(NonASAPOp::Filter { + pred: Predicate(pred), + child, + }) + } + + /// `SELECT service FROM scan` — a one-column subquery. + /// Scalar reads retain cardinality/null semantics and a shared producer. + #[test] + fn scalar_subqueries_remain_explicit_and_shared() { + let sub = one_column_subquery(); + let root = node(NonASAPOp::Project { + child: scan(), + qualifier: None, + cols: vec![ + ProjectItem { + alias: Some("a".into()), + expr: ScalarExpr::ScalarSubquery(Rc::clone(&sub)), + }, + ProjectItem { + alias: Some("b".into()), + expr: ScalarExpr::ScalarSubquery(Rc::clone(&sub)), + }, + ], + }); + let out = canonicalize(root).unwrap(); + let NonASAPOp::Project { cols, child, .. } = out.expect_non_asap() else { + panic!() + }; + assert!(matches!(child.expect_non_asap(), NonASAPOp::Scan { .. })); + for col in cols { + assert!(matches!(&col.expr,ScalarExpr::ScalarSubquery(node) if Rc::ptr_eq(node,&sub))); + } + assert!(out.schema.fields.iter().all(|f| f.nullable)); + assert_idempotent(&out); + } + + fn one_column_subquery() -> Rc { + node(NonASAPOp::Scan { + source: Source::Table { + table_ref: "sub".into(), + }, + predicates: vec![], + schema: Schema::new(vec![Field::plain("value", DataType::Float64, false)]), + }) + } + + fn exists(subquery: Rc, negated: bool) -> ScalarExpr { + ScalarExpr::Exists { subquery, negated } + } + + fn in_subquery(expr: ScalarExpr, subquery: Rc, negated: bool) -> ScalarExpr { + ScalarExpr::InSubquery { + expr: Box::new(expr), + subquery, + negated, + } + } + + /// `value(2) > 1`. + fn value_gt_1() -> ScalarExpr { + ScalarExpr::Compare { + left: Box::new(ScalarExpr::Column(2)), + op: CompareOpKind::Gt, + right: Box::new(ScalarExpr::Literal(ScalarValue::Int64(1))), + semantics: ExprSemantics::Sql, + } + } + + fn literal_true() -> ScalarExpr { + ScalarExpr::Literal(ScalarValue::Boolean(true)) + } + + fn join_parts( + n: &OperatorNode, + ) -> (JoinKind, &ScalarExpr, &Rc, &Rc) { + match n.non_asap() { + Some(NonASAPOp::Join { + kind, + pred: Predicate(pred), + left, + right, + }) => (kind.clone(), pred, left, right), + _ => panic!("expected a Join, got {n:?}"), + } + } + + fn assert_idempotent(once: &Rc) { + let twice = canonicalize(Rc::clone(once)).unwrap(); + assert!(Rc::ptr_eq(once, &twice), "canonicalize must be idempotent"); + } + + #[test] + fn exists_filter_becomes_semi_join() { + let (left, sub) = (scan(), one_column_subquery()); + let q = filter_of(exists(Rc::clone(&sub), false), Rc::clone(&left)); + let out = canonicalize(q).unwrap(); + let (kind, pred, l, r) = join_parts(&out); + assert_eq!(kind, JoinKind::Semi); + assert_eq!(*pred, literal_true()); + assert!(Rc::ptr_eq(l, &left) && Rc::ptr_eq(r, &sub)); + assert_eq!( + out.schema.fields, left.schema.fields, + "a semi join outputs the left's columns" + ); + assert_idempotent(&out); + } + + #[test] + fn not_exists_becomes_anti_join() { + let (left, sub) = (scan(), one_column_subquery()); + let q = filter_of(exists(Rc::clone(&sub), true), Rc::clone(&left)); + let out = canonicalize(q).unwrap(); + let (kind, pred, l, r) = join_parts(&out); + assert_eq!(kind, JoinKind::Anti); + assert_eq!(*pred, literal_true()); + assert!(Rc::ptr_eq(l, &left) && Rc::ptr_eq(r, &sub)); + assert_idempotent(&out); + } + + #[test] + fn in_subquery_becomes_semi_join_on_the_subquery_column() { + // `WHERE service IN (SELECT service …)` over a 3-column left: the + // subquery's column is `Column(3)` in the `left ++ right` scope. + let (left, sub) = (scan(), one_column_subquery()); + let q = filter_of( + in_subquery(ScalarExpr::Column(1), Rc::clone(&sub), false), + Rc::clone(&left), + ); + let out = canonicalize(q).unwrap(); + let (kind, pred, l, r) = join_parts(&out); + assert_eq!(kind, JoinKind::Semi); + assert_eq!( + *pred, + ScalarExpr::Compare { + left: Box::new(ScalarExpr::Column(1)), + op: CompareOpKind::Eq, + right: Box::new(ScalarExpr::Column(3)), + semantics: ExprSemantics::Sql, + } + ); + assert!(Rc::ptr_eq(l, &left) && Rc::ptr_eq(r, &sub)); + assert_eq!( + out.schema.fields, left.schema.fields, + "a semi join outputs the left's columns" + ); + assert_idempotent(&out); + } + + #[test] + fn exists_with_other_conjuncts_keeps_an_outer_filter() { + // `WHERE value > 1 AND EXISTS (…)` → Filter{ value > 1 }{ Semi }. + let (left, sub) = (scan(), one_column_subquery()); + let q = filter_of( + ScalarExpr::BoolAnd(vec![value_gt_1(), exists(Rc::clone(&sub), false)]), + Rc::clone(&left), + ); + let out = canonicalize(q).unwrap(); + let Some(NonASAPOp::Filter { + pred: Predicate(pred), + child, + }) = out.non_asap() + else { + panic!("expected an outer Filter, got {out:?}"); + }; + assert_eq!(*pred, value_gt_1()); + let (kind, _, l, r) = join_parts(child); + assert_eq!(kind, JoinKind::Semi); + assert!(Rc::ptr_eq(l, &left) && Rc::ptr_eq(r, &sub)); + assert_idempotent(&out); + } + + #[test] + fn two_subquery_conjuncts_become_nested_joins() { + // `WHERE EXISTS (a) AND service NOT EXISTS (b) AND value > 1` sheds + // one conjunct per round: Filter{ value > 1 }{ Anti{ Semi{ l, a }, b } }. + let (left, a, b) = (scan(), one_column_subquery(), one_column_subquery()); + let q = filter_of( + ScalarExpr::BoolAnd(vec![ + exists(Rc::clone(&a), false), + exists(Rc::clone(&b), true), + value_gt_1(), + ]), + Rc::clone(&left), + ); + let out = canonicalize(q).unwrap(); + let Some(NonASAPOp::Filter { + pred: Predicate(pred), + child, + }) = out.non_asap() + else { + panic!("expected an outer Filter, got {out:?}"); + }; + assert_eq!(*pred, value_gt_1()); + let (kind, _, inner, r) = join_parts(child); + assert_eq!(kind, JoinKind::Anti); + assert!(Rc::ptr_eq(r, &b)); + let (kind, _, l, r) = join_parts(inner); + assert_eq!(kind, JoinKind::Semi); + assert!(Rc::ptr_eq(l, &left) && Rc::ptr_eq(r, &a)); + assert_idempotent(&out); + } + + #[test] + fn not_in_subquery_is_left_alone() { + let q = filter_of( + in_subquery(ScalarExpr::Column(1), one_column_subquery(), true), + scan(), + ); + let out = canonicalize(Rc::clone(&q)).unwrap(); + assert!(Rc::ptr_eq(&out, &q), "NOT IN keeps its Filter"); + assert_idempotent(&out); + } + + #[test] + fn lifted_subquery_is_canonicalized() { + // The subquery is itself a promotable heavy-hitter; once lifted into + // the join it is canonical, so a second pass finds nothing to do. + let sub = limit(5, 0, sort(desc(1), count_by_service())); + let q = filter_of(exists(sub, false), scan()); + let out = canonicalize(q).unwrap(); + let (_, _, _, r) = join_parts(&out); + assert!(is_topk_over_count(r)); + assert_idempotent(&out); + } +} diff --git a/crates/types/src/ir/cse.rs b/crates/types/src/ir/cse.rs new file mode 100644 index 000000000..eb6fc4bf7 --- /dev/null +++ b/crates/types/src/ir/cse.rs @@ -0,0 +1,791 @@ +//! Structural common-subexpression elimination over the unified operator IR: +//! bottom-up hash-consing of [`OperatorNode`] DAGs across a workload's roots. +//! +//! CSE only runs on already-bound, already-canonicalized plans — structural +//! matching is meaningless before canonicalization has converged +//! semantically-equivalent queries onto one shape. [`share_common_sub_dags`] +//! is the single entry point, run once per workload batch (a batch of one +//! still deduplicates a query's own repeated sub-DAGs, see below). +//! +//! ## Algorithm: classic hash-consing / value-numbering +//! +//! Bottom-up: every child is interned before its parent, so two parents whose +//! children were independently deduplicated down to the same `Rc`s are +//! structurally identical iff their own fields also match, without re-walking +//! the sub-DAGs. "Child" means everything [`OperatorNode::children`] returns: +//! the operator inputs *and* the operator nodes a scalar expression reads +//! (`PromqlScalarFromVector`, `ScalarSubquery`, `Exists`, `InSubquery`), so a +//! vector read by `scalar(v)` in two queries is shared like any other input. +//! The scalar expressions themselves stay opaque data on their owning node. +//! +//! ## Correctness: hash is a filter, `PartialEq` is the decision +//! +//! This is the one non-negotiable rule. A **false positive** here — two +//! sub-DAGs wrongly judged shareable — is a wrong query answer, not a missed +//! optimization: two different queries would read each other's data. +//! [`structural_hash`] (SipHash over a canonical serialization, no +//! collision-freedom guarantee) may only narrow the candidate set within one +//! bucket; the typed equality check on that bucket ([`same_node`]) is what +//! actually decides sharing, every time, no exceptions for "the hash probably +//! didn't collide." Equality is intentionally conservative: it recognizes +//! *exact* structural matches only, never "a stricter-accuracy summary could +//! also answer a looser request" (that subsumption question belongs to the +//! ASAP matcher, not here). +//! +//! ## Legality +//! +//! Structural equality is necessary but not sufficient. A non-ASAP node is +//! only ever *returned* as a match for another when its output has a provable +//! unique key (`Schema::has_unique_key()`): a producer's output can only be +//! shared across consumers when its row identity is stable across reads, so +//! an ungrouped aggregate, a `without(..)` grouping, a `Concat`/`SetOp` that +//! drops its keys, … is always inserted fresh even when it is structurally +//! identical to something already interned. An ASAP node (summary state and +//! its evaluations) has no such gate: equal operator, schema and guarantee make +//! it shareable, exactly as post-ASAP sharing decided before this IR. +//! +//! ## Single-query CSE falls out for free +//! +//! A repeated sub-expression within *one* query (the same grouped aggregate on +//! both `BinaryOp` branches) is deduplicated by the same bottom-up interning — +//! a workload of size one still interns bottom-up within that one DAG. + +use std::collections::hash_map::DefaultHasher; +use std::collections::HashMap; +use std::hash::{Hash, Hasher}; +use std::rc::Rc; + +use super::node::{Operator, OperatorNode}; +use super::non_asap::NonASAPOp; +use crate::pre_asap::schema::Schema; + +/// [`structural_hash`]'s memoization cache: an already-hashed node's `Rc` +/// pointer to its hash. A fresh cache is always correct; what matters is +/// letting it persist across every node of one bottom-up pass rather than +/// starting a new one per call. The caller must keep every cached node alive +/// for the cache's lifetime, or a reused address would alias a stale entry. +pub type HashCache = HashMap<*const OperatorNode, u64>; + +/// A constant stand-in for every child position. Substituting it before +/// serializing or comparing a node leaves exactly the node's own fields. +fn placeholder() -> Rc { + Rc::new(OperatorNode::with_schema( + Operator::NonASAP(NonASAPOp::Values { + rows: vec![], + schema: Schema::lifted(vec![], None), + }), + Schema::lifted(vec![], None), + )) +} + +/// The operator with every child — operator inputs and the operator nodes +/// referenced from its scalar expressions alike — replaced by +/// [`placeholder`]. What remains is the node's own data: variant tag, scalar +/// expressions (with their operator references blanked), parameters. +fn own_fields(node: &OperatorNode) -> Operator { + let placeholder = placeholder(); + node.operator.map_children(|_| Rc::clone(&placeholder)) +} + +/// Coarse structural hash used only to bucket [`InternTable::intern`]'s +/// candidate search — never the sharing decision ([`same_node`] is). +/// +/// `OperatorNode` carries `f64`s (`ScalarValue::Float64`, quantile targets, +/// `ResultGuarantee` bounds, …), so it cannot derive `std::hash::Hash`. The +/// hash is SipHash over two parts: +/// +/// 1. the canonical JSON of [`own_fields`] plus `result_kind`, `schema`, +/// `guarantee` and `timing` — every field `PartialEq` compares except the +/// children. A scalar expression is serialized as data with each operator +/// node it reads replaced by a constant placeholder, so a reference to an +/// interned sub-DAG contributes nothing of its own here; +/// 2. for every child in [`OperatorNode::children`] order (operator inputs, +/// then scalar-referenced nodes), the child's own `structural_hash`, +/// memoized in `cache` by `Rc` pointer identity. +/// +/// Part 2 is what makes equal sub-DAGs hash equal whether they are reached +/// through an operator input or through a `scalar(v)`, and what keeps the +/// pass linear: a node is generally a DAG, and re-serializing a shared +/// descendant once per parent would cost `O(sub-DAG)` per node instead of +/// `O(1)` beyond the children's already-known hashes. A non-finite `f64` +/// serializes as `null`, merely widening one (still equality-checked) bucket. +pub fn structural_hash(node: &OperatorNode, cache: &mut HashCache) -> u64 { + fn child_hash(child: &Rc, cache: &mut HashCache) -> u64 { + let ptr = Rc::as_ptr(child); + if let Some(&h) = cache.get(&ptr) { + return h; + } + let h = structural_hash(child, cache); + cache.insert(ptr, h); + h + } + + let mut hasher = DefaultHasher::new(); + let own = ( + own_fields(node), + node.result_kind, + &node.schema, + &node.guarantee, + node.timing, + &node.coverage, + ); + serde_json::to_string(&own) + .unwrap_or_default() + .hash(&mut hasher); + for child in node.children() { + child_hash(child, cache).hash(&mut hasher); + } + hasher.finish() +} + +/// Numeric `PartialEq` alone conflates signed zeros. The serialized check is +/// additional evidence, never a replacement for typed equality (JSON maps +/// non-finite floats to `null`). Used for the guarantee, whose bounds are +/// floats a shared node must preserve bit-for-bit. +fn same_value(left: &T, right: &T) -> bool { + left == right + && match (serde_json::to_string(left), serde_json::to_string(right)) { + (Ok(left), Ok(right)) => left == right, + _ => false, + } +} + +/// Memo of child-pair comparisons already decided by [`same_node`], keyed by +/// pointer pair. Only interned (table-owned, hence alive) nodes are keys. +type EqMemo = HashMap<(*const OperatorNode, *const OperatorNode), bool>; + +/// The sharing decision: typed equality of two nodes. +/// +/// `OperatorNode`'s derived `PartialEq` would recurse into children by value +/// even when both sides hold the same `Rc` (`OperatorNode` is not `Eq`, so +/// `Rc` gets no pointer shortcut), expanding a shared diamond once per path. +/// Children are therefore compared by pointer first; only when the pointers +/// differ (an equal child that was not legal to share) are the values +/// compared, memoized per pair so a diamond is still walked once. +fn same_node(left: &OperatorNode, right: &OperatorNode, memo: &mut EqMemo) -> bool { + let (lc, rc) = (left.children(), right.children()); + if lc.len() != rc.len() { + return false; + } + let children_equal = lc.iter().zip(&rc).all(|(a, b)| { + if Rc::ptr_eq(a, b) { + return true; + } + let key = (Rc::as_ptr(a), Rc::as_ptr(b)); + if let Some(&eq) = memo.get(&key) { + return eq; + } + let eq = same_node(a, b, memo); + memo.insert(key, eq); + eq + }); + children_equal + && left.result_kind == right.result_kind + && left.schema == right.schema + && left.timing == right.timing + && left.coverage == right.coverage + && same_value(&left.guarantee, &right.guarantee) + && same_value(&own_fields(left), &own_fields(right)) +} + +/// Bottom-up hash-consing table: structurally-equal, sharing-legal nodes +/// collapse onto one `Rc`. +/// +/// `buckets` is keyed by [`structural_hash`] — a coarse candidate filter +/// only. Every entry within one bucket is a full node kept around for the +/// [`same_node`] comparison that actually decides a match; a hash collision +/// between structurally different nodes just means a harmless linear scan of +/// a few extra candidates. +struct InternTable { + buckets: HashMap>>, + /// Persisted for the table's whole lifetime so hashing is `O(1)` per node + /// beyond its children; every cached node is owned by `buckets`. + hash_cache: HashCache, + eq_memo: EqMemo, +} + +impl InternTable { + fn new() -> Self { + Self { + buckets: HashMap::new(), + hash_cache: HashMap::new(), + eq_memo: HashMap::new(), + } + } + + /// Intern one node whose children are already interned: look it up by + /// [`structural_hash`], confirm with [`same_node`], and — only when + /// sharing is legal (module doc, "Legality") — return the existing `Rc` + /// instead of allocating a new one. + fn intern(&mut self, node: OperatorNode) -> Rc { + let hash = structural_hash(&node, &mut self.hash_cache); + // A node that is not legal to share is never *returned* as a match + // for something else; it still occupies a fresh slot in the bucket + // (harmless: later scans require legality of the new node too). + let reusable = node.is_asap() || node.schema.has_unique_key(); + let bucket = self.buckets.entry(hash).or_default(); + if reusable { + if let Some(existing) = bucket + .iter() + .find(|candidate| same_node(candidate, &node, &mut self.eq_memo)) + { + return Rc::clone(existing); + } + } + let rc = Rc::new(node); + bucket.push(Rc::clone(&rc)); + rc + } +} + +/// Count of *unique* nodes reachable from `root` (pointer identity, +/// following [`OperatorNode::children`]): the real size of the DAG, not a +/// tree-walk count that re-counts a shared descendant once per parent. +pub fn dag_node_count(root: &Rc) -> usize { + OperatorNode::reachable(root).len() +} + +/// Input pointer → (input `Rc`, interned result). The input `Rc` is retained +/// so its address cannot be freed and reused by a fresh allocation while the +/// memo still maps it. +type Visited = HashMap<*const OperatorNode, (Rc, Rc)>; + +/// Intern `node`'s children (recursively), then `node` itself. The rebuilt +/// node keeps `node`'s retained schema, result kind, guarantee and timing: +/// every child is replaced by an equal node, so each derived property stays +/// valid, and the result is `PartialEq`-equal to the input. +fn intern_bottom_up( + table: &mut InternTable, + visited: &mut Visited, + node: &Rc, +) -> Rc { + if let Some((_, interned)) = visited.get(&Rc::as_ptr(node)) { + return Rc::clone(interned); + } + let operator = node + .operator + .map_children(|child| intern_bottom_up(table, visited, child)); + let rebuilt = OperatorNode { + operator, + result_kind: node.result_kind, + schema: node.schema.clone(), + guarantee: node.guarantee.clone(), + timing: node.timing, + coverage: node.coverage.clone(), + }; + let interned = table.intern(rebuilt); + visited.insert(Rc::as_ptr(node), (Rc::clone(node), Rc::clone(&interned))); + interned +} + +/// Share structurally-identical, sharing-legal sub-DAGs across a workload's +/// roots (or within one root). Every root's *value* is unchanged +/// (`PartialEq`-equal to its input) — only its internal `Rc` structure may +/// now alias another root's, or another part of its own DAG. A node already +/// reached through two paths is visited once. +/// +/// `Id` is caller-chosen — a workload entry's key, an index, a query name. +pub fn share_common_sub_dags( + roots: Vec<(Id, Rc)>, +) -> Vec<(Id, Rc)> { + let mut table = InternTable::new(); + let mut visited = Visited::new(); + roots + .into_iter() + .map(|(id, root)| (id, intern_bottom_up(&mut table, &mut visited, &root))) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::ir::asap::ASAPOp; + use crate::ir::operator_properties::{BinaryOpKind, GroupKeys, Reduction, Source}; + use crate::ir::BinaryOperator; + use crate::ir::ScalarExpr; + use crate::post_asap::guarantee::ResultGuarantee; + use crate::post_asap::sketch::{ + GroupingStrategy, SketchAlgorithm, SketchKind, SketchParams, SummaryUpdate, + }; + use crate::pre_asap::agg_intent::AggIntent; + use crate::pre_asap::expr_ir::{ColumnRef, CompareOpKind}; + use crate::pre_asap::schema::{DataType, Field, FieldDataType, Schema}; + + use crate::types::AccuracyTarget; + + fn node(op: NonASAPOp) -> Rc { + OperatorNode::new_shared(crate::ir::Operator::NonASAP(op)).unwrap() + } + + /// `[ts, service, value, latency]`, no unique key. + fn scan() -> Rc { + node(NonASAPOp::Scan { + source: Source::TimeSeries { metric: "m".into() }, + predicates: vec![], + schema: Schema::with_time_index( + vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("service", DataType::Utf8, false), + Field::plain("value", DataType::Float64, false), + Field::plain("latency", DataType::Float64, false), + ], + 0, + vec![], + ), + }) + } + + fn quantile_agg(by: Vec, col: Option, q: f64) -> Rc { + node(NonASAPOp::Aggregate { + reduction: Reduction::by(by), + measures: vec![AggIntent::Quantile { + col, + q, + accuracy: AccuracyTarget::Exact, + }], + output_names: vec![], + filters: vec![], + having: None, + child: scan(), + }) + } + + fn compare(lhs: Rc, rhs: Rc) -> Rc { + node(NonASAPOp::BinaryOp { + operator: BinaryOperator { + checked_relative_division: false, + checked_finite_division: false, + kind: BinaryOpKind::Compare(CompareOpKind::Eq), + vector_match: None, + }, + return_bool: false, + lhs, + rhs, + }) + } + + fn two_roots(a: Rc, b: Rc) -> (Rc, Rc) { + let shared = share_common_sub_dags(vec![("a", a), ("b", b)]); + let [(_, ra), (_, rb)] = shared.as_slice() else { + panic!("expected 2 roots"); + }; + (Rc::clone(ra), Rc::clone(rb)) + } + + #[test] + fn distinct_column_quantiles_do_not_merge() { + // Grouped (unique key present) so only the differing `col` blocks it. + let (ra, rb) = two_roots( + quantile_agg(vec![1], Some(2), 0.5), + quantile_agg(vec![1], Some(3), 0.5), + ); + assert!(!Rc::ptr_eq(&ra, &rb)); + assert_ne!(ra, rb); + } + + #[test] + fn no_unique_keys_means_no_merge_even_when_structurally_identical() { + let a = quantile_agg(vec![], Some(2), 0.9); + let b = quantile_agg(vec![], Some(2), 0.9); + assert_eq!(a, b, "fixture sanity: structurally equal"); + assert!( + !a.schema.has_unique_key(), + "fixture sanity: a global aggregate has no provable unique key" + ); + let (ra, rb) = two_roots(a, b); + assert!( + !Rc::ptr_eq(&ra, &rb), + "no unique key ⇒ never hoisted, even for an identical structural match" + ); + } + + #[test] + fn median_and_explicit_half_percentile_merge() { + // Two spellings that lower to the identical grouped `Quantile { q: 0.5 }`. + let (m, p) = two_roots( + quantile_agg(vec![1], Some(2), 0.5), + quantile_agg(vec![1], Some(2), 0.5), + ); + assert!(Rc::ptr_eq(&m, &p)); + } + + #[test] + fn single_query_shares_its_own_repeated_sub_dag() { + // One root with the same grouped aggregate on both branches, built as + // two separately-allocated sub-DAGs (no sharing yet). + let root = compare( + quantile_agg(vec![1], Some(2), 0.5), + quantile_agg(vec![1], Some(2), 0.5), + ); + let shared = share_common_sub_dags(vec![("q", root)]); + let [(_, root)] = shared.as_slice() else { + panic!("expected 1 root"); + }; + let Some(NonASAPOp::BinaryOp { lhs, rhs, .. }) = root.non_asap() else { + panic!("expected BinaryOp root, got {root:?}"); + }; + assert!(Rc::ptr_eq(lhs, rhs)); + } + + #[test] + fn shared_root_value_is_unchanged() { + let a = quantile_agg(vec![1], Some(2), 0.5) + .as_ref() + .clone() + .with_guarantee(Some(ResultGuarantee::exact("fixture"))); + let before = Rc::new(a); + let (ra, _) = two_roots(Rc::clone(&before), Rc::clone(&before)); + assert_eq!(ra.as_ref(), before.as_ref()); + assert!( + ra.guarantee.is_some(), + "retained properties survive the rebuild" + ); + } + + // ── scalar-referenced sub-DAGs ────────────────────────────────────── + + /// `vector(scalar(sum by (service) (up)))`. + fn scalar_of_vector() -> Rc { + let sum_up = node(NonASAPOp::Aggregate { + reduction: Reduction::by(vec![1]), + measures: vec![AggIntent::Sum { col: Some(2) }], + output_names: vec![], + filters: vec![], + having: None, + child: scan(), + }); + assert!(sum_up.schema.has_unique_key(), "fixture sanity"); + node(NonASAPOp::PromqlVectorFromScalar( + ScalarExpr::PromqlScalarFromVector(sum_up), + )) + } + + fn bridged_vector(root: &Rc) -> &Rc { + match root.non_asap() { + Some(NonASAPOp::PromqlVectorFromScalar(ScalarExpr::PromqlScalarFromVector(v))) => v, + other => panic!("expected vector(scalar(v)), got {other:?}"), + } + } + + #[test] + fn scalar_referenced_vector_is_shared_across_queries() { + let (ra, rb) = two_roots(scalar_of_vector(), scalar_of_vector()); + assert!( + Rc::ptr_eq(bridged_vector(&ra), bridged_vector(&rb)), + "the vector read by scalar(v) is a child and must be interned" + ); + assert!( + !Rc::ptr_eq(&ra, &rb), + "the scalar bridge itself has no unique key and stays separate" + ); + } + + #[test] + fn structural_hash_sees_through_a_scalar_reference() { + // Two equal bridges must hash equal whether or not their referenced + // vector is the same Rc — the reference contributes the vector's + // memoized hash, not its identity. + let a = scalar_of_vector(); + let b = scalar_of_vector(); + let mut cache = HashMap::new(); + assert_eq!( + structural_hash(&a, &mut cache), + structural_hash(&b, &mut cache) + ); + assert_eq!( + cache.len(), + 4, + "aggregate + scan cached once per root: {cache:?}" + ); + let other = node(NonASAPOp::PromqlVectorFromScalar( + ScalarExpr::PromqlScalarFromVector(quantile_agg(vec![1], Some(2), 0.5)), + )); + assert_ne!( + structural_hash(&a, &mut cache), + structural_hash(&other, &mut cache) + ); + } + + // ── ASAP nodes ────────────────────────────────────────────────────── + + fn summary_agg(alpha: f64, guarantee: Option) -> Rc { + let family = FieldDataType::Sketch( + SketchKind::new(SketchAlgorithm::DDSketch, SketchParams::DDSketch { alpha }), + GroupingStrategy::default(), + ); + let schema = Schema::lifted(vec![Field::new("state", family.clone(), false)], None); + assert!(!schema.has_unique_key(), "fixture sanity"); + Rc::new( + OperatorNode::with_schema( + Operator::ASAP(ASAPOp::SummaryAgg { + child: scan(), + family, + input: SummaryUpdate::column(ColumnRef::SampleValue), + reduction: Reduction::PerEntity, + grouping: GroupingStrategy::default(), + filter: None, + }), + schema, + ) + .with_guarantee(guarantee), + ) + } + + #[test] + fn asap_nodes_share_without_a_unique_key() { + let exact = || Some(ResultGuarantee::exact("fixture")); + let (ra, rb) = two_roots(summary_agg(0.01, exact()), summary_agg(0.01, exact())); + assert!(Rc::ptr_eq(&ra, &rb)); + assert!(ra.guarantee.is_some()); + } + + #[test] + fn asap_nodes_with_distinct_parameters_or_guarantees_are_not_shared() { + let exact = || Some(ResultGuarantee::exact("fixture")); + let (ra, rb) = two_roots(summary_agg(0.01, exact()), summary_agg(0.001, exact())); + assert!(!Rc::ptr_eq(&ra, &rb), "different sketch parameters"); + let (ra, rb) = two_roots(summary_agg(0.01, exact()), summary_agg(0.01, None)); + assert!( + !Rc::ptr_eq(&ra, &rb), + "an unknown guarantee never borrows an exact one" + ); + assert!(rb.guarantee.is_none()); + } + + #[test] + fn evaluations_share_their_producer_but_not_each_other() { + use crate::post_asap::sketch::SketchStatistic; + let evaluation = |q: f64| { + Rc::new(OperatorNode::with_schema( + Operator::ASAP(ASAPOp::SummaryEstimate { + summary_input: summary_agg(0.01, None), + query: SketchStatistic::Quantile { q }, + }), + Schema::lifted( + vec![Field::plain("quantile", DataType::Float64, false)], + None, + ), + )) + }; + let (p95, p99) = two_roots(evaluation(0.95), evaluation(0.99)); + let producer = |n: &Rc| Rc::clone(n.children()[0]); + assert!(!Rc::ptr_eq(&p95, &p99)); + assert!(Rc::ptr_eq(&producer(&p95), &producer(&p99))); + } + + // ── structural_hash (DAG-aware memoization) ───────────────────────── + + #[test] + fn structural_hash_is_stable_across_cache_states() { + let agg = quantile_agg(vec![1], Some(2), 0.5); + let mut cold = HashMap::new(); + let mut warm = HashMap::new(); + structural_hash(&scan(), &mut warm); + assert_eq!( + structural_hash(&agg, &mut cold), + structural_hash(&agg, &mut warm), + "hash must be independent of unrelated cache state" + ); + } + + #[test] + fn structural_hash_of_an_internally_shared_dag_matches_the_unshared_equivalent() { + let agg = quantile_agg(vec![1], Some(2), 0.5); + let shared_root = compare(Rc::clone(&agg), Rc::clone(&agg)); + let unshared_root = compare( + quantile_agg(vec![1], Some(2), 0.5), + quantile_agg(vec![1], Some(2), 0.5), + ); + assert_eq!( + structural_hash(&shared_root, &mut HashMap::new()), + structural_hash(&unshared_root, &mut HashMap::new()), + ); + } + + #[test] + fn structural_hash_memoizes_a_shared_descendant_exactly_once() { + let agg = quantile_agg(vec![1], Some(2), 0.5); + let root = compare(Rc::clone(&agg), Rc::clone(&agg)); + let mut cache = HashMap::new(); + structural_hash(&root, &mut cache); + assert_eq!( + cache.len(), + 2, + "one entry per unique node in the shared branch (Aggregate + Scan): {cache:?}" + ); + } + + // ── dag_node_count ─────────────────────────────────────────────────── + + #[test] + fn dag_node_count_is_the_naive_count_when_nothing_is_shared() { + assert_eq!(dag_node_count(&scan()), 1); + assert_eq!(dag_node_count(&quantile_agg(vec![1], Some(2), 0.5)), 2); + assert_eq!( + dag_node_count(&scalar_of_vector()), + 3, + "follows scalar references" + ); + } + + #[test] + fn dag_node_count_deduplicates_an_internally_shared_sub_dag() { + let root = compare( + quantile_agg(vec![1], Some(2), 0.5), + quantile_agg(vec![1], Some(2), 0.5), + ); + assert_eq!( + dag_node_count(&root), + 5, + "fixture sanity: nothing shared yet" + ); + let shared = share_common_sub_dags(vec![("q", root)]); + let [(_, root)] = shared.as_slice() else { + panic!("expected 1 root"); + }; + assert_eq!( + dag_node_count(root), + 3, + "BinaryOp + one Aggregate + its Scan" + ); + } + + #[test] + fn dag_node_count_deduplicates_across_two_workload_roots() { + let (ra, rb) = two_roots( + quantile_agg(vec![1], Some(2), 0.5), + quantile_agg(vec![1], Some(2), 0.5), + ); + assert!(Rc::ptr_eq(&ra, &rb), "fixture sanity: the two roots merged"); + assert_eq!(dag_node_count(&ra), 2); + assert_eq!(dag_node_count(&rb), 2); + } + + #[test] + fn dedup_gates_sharing_the_same_as_aggregate() { + // `Dedup { cols }` adds `cols` as a unique key, so two identical + // `Dedup`s merge even though their keyless `Scan`s could not. + let dedup = || { + node(NonASAPOp::Dedup { + cols: vec![1], + child: scan(), + }) + }; + let (ra, rb) = two_roots(dedup(), dedup()); + assert!(Rc::ptr_eq(&ra, &rb)); + } + + #[test] + fn group_keys_gate_still_prevented_when_partition_by_without_used() { + let without_agg = || { + node(NonASAPOp::Aggregate { + reduction: Reduction::Reduce(GroupKeys::without(vec![0])), + measures: vec![AggIntent::Count { + accuracy: AccuracyTarget::Exact, + }], + output_names: vec![], + filters: vec![], + having: None, + child: scan(), + }) + }; + let a = without_agg(); + assert!(!a.schema.has_unique_key()); + let (ra, rb) = two_roots(a, without_agg()); + assert!(!Rc::ptr_eq(&ra, &rb)); + } + + #[test] + fn already_shared_nodes_are_visited_once() { + // A diamond already present in the input stays one node and is not + // re-interned per path. + let agg = quantile_agg(vec![1], Some(2), 0.5); + let root = compare(Rc::clone(&agg), Rc::clone(&agg)); + let shared = share_common_sub_dags(vec![("q", root)]); + let Some(NonASAPOp::BinaryOp { lhs, rhs, .. }) = shared[0].1.non_asap() else { + panic!("expected BinaryOp root"); + }; + assert!(Rc::ptr_eq(lhs, rhs)); + assert_eq!(dag_node_count(&shared[0].1), 3); + } + + // Comparing a shareable node whose equal-but-unshareable children form a + // deep diamond must not expand the diamond once per path. The timeout is + // a coarse runaway guard, not a performance SLA. + #[test] + fn shared_diamond_does_not_expand_during_comparison() { + let (done, completion) = std::sync::mpsc::channel(); + let worker = std::thread::spawn(move || { + fn keyed_diamond() -> Rc { + // BinaryOp over a keyless scan has no unique key at any level, + // so none of the 24 levels is shareable; the `Dedup` on top is. + let mut current = scan(); + for _ in 0..24 { + current = compare(Rc::clone(¤t), current); + } + node(NonASAPOp::Dedup { + cols: vec![1], + child: current, + }) + } + let (ra, rb) = two_roots(keyed_diamond(), keyed_diamond()); + assert!(Rc::ptr_eq(&ra, &rb)); + done.send(()).unwrap(); + }); + completion + .recv_timeout(std::time::Duration::from_secs(5)) + .expect("comparison expanded the shared DAG"); + worker.join().unwrap(); + } + + /// A keyed (hence shareable) projection emitting the literal `value`. + fn keyed_literal(value: f64) -> Rc { + let keyed = node(NonASAPOp::Scan { + source: Source::TimeSeries { metric: "m".into() }, + predicates: vec![], + schema: Schema::with_time_index( + vec![ + Field::plain("ts", DataType::Timestamp, false), + Field::plain("service", DataType::Utf8, false), + ], + 0, + vec![vec![1]], + ), + }); + node(NonASAPOp::Project { + cols: vec![ + crate::ir::ProjectItem { + alias: None, + expr: ScalarExpr::Column(1), + }, + crate::ir::ProjectItem { + alias: Some("v".into()), + expr: ScalarExpr::literal_f64(value), + }, + ], + qualifier: None, + child: keyed, + }) + } + + /// Sharing preserves IEEE signed zero, and JSON's `null` encoding of + /// non-finite floats never becomes the equality decision. + #[test] + fn signed_zero_and_nonfinite_values_remain_distinct() { + assert!( + keyed_literal(0.0).schema.has_unique_key(), + "fixture is shareable" + ); + for (a, b) in [ + (0.0, -0.0), + (-0.0, 0.0), + (f64::INFINITY, f64::NEG_INFINITY), + (f64::NAN, f64::NAN), + ] { + let (ra, rb) = two_roots(keyed_literal(a), keyed_literal(b)); + assert!(!Rc::ptr_eq(&ra, &rb), "{a} and {b} must not be shared"); + } + let (ra, rb) = two_roots(keyed_literal(f64::INFINITY), keyed_literal(f64::INFINITY)); + assert!(Rc::ptr_eq(&ra, &rb)); + } +} diff --git a/crates/types/src/ir/export.rs b/crates/types/src/ir/export.rs new file mode 100644 index 000000000..276d8ee43 --- /dev/null +++ b/crates/types/src/ir/export.rs @@ -0,0 +1,364 @@ +//! Logical ASAP DAG transport (planner-layering stage 1), with no execution timing assigned. +//! +//! This representation preserves operator semantics and summary state types. +//! Timing is derived from materialization during physical planning; +//! physical implementation, materialization and retention remain downstream. +use std::collections::{HashMap, HashSet}; +use std::rc::Rc; + +use serde::{Deserialize, Serialize}; +use thiserror::Error; + +use super::wire::{grouping_compatibility, input_edges, payload_of}; +pub use super::wire::{ + EdgeRole, GroupingEdgeCompatibility, LogicalASAPNodeId, LogicalASAPOperatorPayload, + WireScalarExpr, +}; +use super::{ASAPOp, Operator, OperatorNode, OperatorResultKind, QueryRoot, SchemaDerivationError}; +use crate::post_asap::guarantee::ResultGuarantee; +use crate::pre_asap::{FieldDataType, Schema}; + +/// Independent envelope version: this replaces the older phase-assigned format. +pub const LOGICAL_ASAP_DAG_WIRE_VERSION: u32 = 1; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LogicalASAPDAGNode { + pub id: LogicalASAPNodeId, + pub payload: LogicalASAPOperatorPayload, + pub result_kind: OperatorResultKind, + pub output_schema: Schema, + pub guarantee: Option, + #[serde(default)] + pub coverage: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LogicalASAPDAGEdge { + pub producer: LogicalASAPNodeId, + pub consumer: LogicalASAPNodeId, + pub role: EdgeRole, + pub intermediate_schema: Schema, + pub grouping: GroupingEdgeCompatibility, +} + +/// Standalone scalars remain scalar roots rather than fabricated operator nodes. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub enum LogicalASAPQueryRoot { + Operator(LogicalASAPNodeId), + Scalar(WireScalarExpr), +} +impl LogicalASAPQueryRoot { + pub fn operator_refs(&self) -> Vec { + match self { + Self::Operator(id) => vec![*id], + Self::Scalar(expr) => expr.operator_refs(), + } + } +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LogicalASAPDAG { + pub nodes: Vec, + pub edges: Vec, + /// One root per query of the batch, in workload order. Queries that share + /// a sub-DAG reference the same exported nodes. + pub roots: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct LogicalASAPDAGDocument { + pub schema_version: u32, + pub dag: LogicalASAPDAG, +} + +#[derive(Debug, Clone, PartialEq, Eq, Error)] +pub enum LogicalASAPDAGValidationError { + #[error("unsupported logical ASAP DAG version {0}")] + UnsupportedVersion(u32), + #[error("duplicate logical node {0:?}")] + DuplicateNode(LogicalASAPNodeId), + #[error("missing logical node {0:?}")] + MissingNode(LogicalASAPNodeId), + #[error("edge schema differs from producer {0:?}")] + EdgeSchemaMismatch(LogicalASAPNodeId), + #[error("summary node {0:?} schema does not contain its declared family/grouping")] + SummarySchemaMismatch(LogicalASAPNodeId), + #[error("invalid summary coverage at {0:?}")] + InvalidCoverage(LogicalASAPNodeId), + #[error("logical ASAP DAG has no query roots")] + NoRoots, + #[error("logical ASAP DAG contains a cycle")] + Cycle, + #[error("unreachable logical node {0:?}")] + UnreachableNode(LogicalASAPNodeId), +} + +impl LogicalASAPDAGDocument { + pub fn new(dag: LogicalASAPDAG) -> Self { + Self { + schema_version: LOGICAL_ASAP_DAG_WIRE_VERSION, + dag, + } + } + + pub fn validate(&self) -> Result<(), LogicalASAPDAGValidationError> { + if self.schema_version != LOGICAL_ASAP_DAG_WIRE_VERSION { + return Err(LogicalASAPDAGValidationError::UnsupportedVersion( + self.schema_version, + )); + } + self.dag.validate() + } +} + +impl LogicalASAPDAG { + /// Transport integrity checks; full operator/scalar typing is checked on + /// the in-memory IR before compilation. + pub fn validate(&self) -> Result<(), LogicalASAPDAGValidationError> { + let mut nodes = HashMap::new(); + for node in &self.nodes { + if nodes.insert(node.id, node).is_some() { + return Err(LogicalASAPDAGValidationError::DuplicateNode(node.id)); + } + if let Some(coverage) = &node.coverage { + if node.result_kind != OperatorResultKind::State || coverage.validate().is_err() { + return Err(LogicalASAPDAGValidationError::InvalidCoverage(node.id)); + } + } + if let LogicalASAPOperatorPayload::SummaryAgg { + family, grouping, .. + } = &node.payload + { + if node.coverage.is_none() { + return Err(LogicalASAPDAGValidationError::InvalidCoverage(node.id)); + } + if !node.output_schema.fields.iter().any(|field| &field.dtype == family) + || node.output_schema.fields.iter().any(|field| matches!(&field.dtype, FieldDataType::Sketch(_, actual) if actual != grouping)) { + return Err(LogicalASAPDAGValidationError::SummarySchemaMismatch(node.id)); + } + } + } + if self.roots.is_empty() { + return Err(LogicalASAPDAGValidationError::NoRoots); + } + let roots: Vec<_> = self + .roots + .iter() + .flat_map(LogicalASAPQueryRoot::operator_refs) + .collect(); + for root in &roots { + if !nodes.contains_key(root) { + return Err(LogicalASAPDAGValidationError::MissingNode(*root)); + } + } + let mut inputs: HashMap<_, Vec<_>> = HashMap::new(); + for edge in &self.edges { + let producer = nodes + .get(&edge.producer) + .ok_or(LogicalASAPDAGValidationError::MissingNode(edge.producer))?; + if !nodes.contains_key(&edge.consumer) { + return Err(LogicalASAPDAGValidationError::MissingNode(edge.consumer)); + } + if edge.intermediate_schema != producer.output_schema { + return Err(LogicalASAPDAGValidationError::EdgeSchemaMismatch( + edge.producer, + )); + } + inputs.entry(edge.consumer).or_default().push(edge.producer); + } + for node in &self.nodes { + if matches!(node.payload, LogicalASAPOperatorPayload::SummaryMerge) { + let coverage = inputs + .get(&node.id) + .into_iter() + .flatten() + .map(|id| { + nodes[id] + .coverage + .clone() + .ok_or(LogicalASAPDAGValidationError::InvalidCoverage(node.id)) + }) + .collect::, _>>()?; + let merged = super::summary_coverage::SummaryCoverage::merge_disjoint(&coverage) + .map_err(|_| LogicalASAPDAGValidationError::InvalidCoverage(node.id))?; + if node.coverage.as_ref() != Some(&merged) { + return Err(LogicalASAPDAGValidationError::InvalidCoverage(node.id)); + } + } + } + fn visit( + id: LogicalASAPNodeId, + inputs: &HashMap>, + active: &mut HashSet, + done: &mut HashSet, + ) -> Result<(), LogicalASAPDAGValidationError> { + if done.contains(&id) { + return Ok(()); + } + if !active.insert(id) { + return Err(LogicalASAPDAGValidationError::Cycle); + } + for child in inputs.get(&id).into_iter().flatten() { + visit(*child, inputs, active, done)?; + } + active.remove(&id); + done.insert(id); + Ok(()) + } + let mut done = HashSet::new(); + for root in roots { + visit(root, &inputs, &mut HashSet::new(), &mut done)?; + } + if let Some(id) = nodes.keys().find(|id| !done.contains(id)) { + return Err(LogicalASAPDAGValidationError::UnreachableNode(*id)); + } + Ok(()) + } +} + +/// Compiler-local identity mapping; IDs are local to this logical export. +#[derive(Debug, Clone)] +pub struct LogicalASAPNodeIdentityMap { + nodes_by_id: Vec>, +} + +impl LogicalASAPNodeIdentityMap { + pub fn node_id(&self, node: &Rc) -> Option { + self.nodes_by_id + .iter() + .position(|candidate| Rc::ptr_eq(candidate, node)) + .map(|id| LogicalASAPNodeId(id as u32)) + } + pub fn operator_node(&self, id: LogicalASAPNodeId) -> Option<&Rc> { + self.nodes_by_id.get(id.0 as usize) + } +} + +#[derive(Debug, Clone)] +pub struct LogicalASAPDAGCompilation { + pub dag: LogicalASAPDAG, + pub node_ids: LogicalASAPNodeIdentityMap, +} + +pub fn compile_logical_asap_dag( + root: &Rc, +) -> Result { + Ok(compile_logical_asap_dag_with_node_ids(root)?.dag) +} + +pub fn compile_logical_asap_dag_with_node_ids( + root: &Rc, +) -> Result { + compile_logical_asap_query_with_node_ids(&QueryRoot::Operator(Rc::clone(root))) +} + +pub fn compile_logical_asap_query( + root: &QueryRoot, +) -> Result { + Ok(compile_logical_asap_query_with_node_ids(root)?.dag) +} + +pub fn compile_logical_asap_query_with_node_ids( + root: &QueryRoot, +) -> Result { + compile_logical_asap_workload_with_node_ids(std::slice::from_ref(root)) +} + +/// Export a batch of queries as one DAG with one root per query. +pub fn compile_logical_asap_workload( + roots: &[QueryRoot], +) -> Result { + Ok(compile_logical_asap_workload_with_node_ids(roots)?.dag) +} + +pub fn compile_logical_asap_workload_with_node_ids( + roots: &[QueryRoot], +) -> Result { + let mut exporter = Exporter::default(); + let mut exported = Vec::with_capacity(roots.len()); + for root in roots { + root.validate_structure()?; + exported.push(match root { + QueryRoot::Operator(node) => LogicalASAPQueryRoot::Operator(exporter.visit(node)), + QueryRoot::Scalar(expr) => { + for node in expr.operator_refs() { + exporter.visit(node); + } + LogicalASAPQueryRoot::Scalar(WireScalarExpr::from_expr(expr, &mut |n| { + exporter.ids[&Rc::as_ptr(n)] + })) + } + }); + } + let dag = LogicalASAPDAG { + nodes: exporter.nodes, + edges: exporter.edges, + roots: exported, + }; + Ok(LogicalASAPDAGCompilation { + dag, + node_ids: LogicalASAPNodeIdentityMap { + nodes_by_id: exporter.nodes_by_id, + }, + }) +} + +#[derive(Default)] +struct Exporter { + ids: HashMap<*const OperatorNode, LogicalASAPNodeId>, + nodes: Vec, + edges: Vec, + nodes_by_id: Vec>, +} + +impl Exporter { + fn visit(&mut self, node: &Rc) -> LogicalASAPNodeId { + if let Some(id) = self.ids.get(&Rc::as_ptr(node)) { + return *id; + } + let mut producers = Vec::new(); + for (child, role) in input_edges(&node.operator) { + producers.push((self.visit(child), child, role)); + } + let scalars = match &node.operator { + Operator::NonASAP(op) => op.scalar_exprs(), + Operator::ASAP(ASAPOp::SummaryAgg { + filter: Some(filter), + .. + }) => vec![&filter.0], + _ => vec![], + }; + for expr in scalars { + for referenced in expr.operator_refs() { + producers.push((self.visit(referenced), referenced, EdgeRole::ScalarRef)); + } + } + let id = LogicalASAPNodeId(self.nodes.len() as u32); + let payload = payload_of(&node.operator, &mut |n| self.ids[&Rc::as_ptr(n)]); + self.nodes.push(LogicalASAPDAGNode { + id, + payload, + result_kind: node.result_kind, + output_schema: node.schema.clone(), + guarantee: node.guarantee.clone(), + coverage: node.coverage.clone(), + }); + self.nodes_by_id.push(Rc::clone(node)); + self.ids.insert(Rc::as_ptr(node), id); + for (producer, child, role) in producers { + self.edges.push(LogicalASAPDAGEdge { + producer, + consumer: id, + role, + intermediate_schema: child.schema.clone(), + grouping: grouping_compatibility(&child.operator, &node.operator), + }); + } + id + } +} diff --git a/crates/types/src/ir/mod.rs b/crates/types/src/ir/mod.rs index fcb2eb49b..04b9f33e3 100644 --- a/crates/types/src/ir/mod.rs +++ b/crates/types/src/ir/mod.rs @@ -16,5 +16,9 @@ pub use non_asap::{BinaryOperator, NonASAPOp, TimeRangeKind}; pub use query::QueryRoot; pub use scalar::{ExprSemantics, Predicate, ProjectItem, ScalarExpr, SortKey}; +pub mod canonicalize; +pub mod cse; +pub mod export; /// Semantic observation coverage, separate from field layout and physical timing. pub mod summary_coverage; +mod wire; diff --git a/crates/types/src/ir/wire.rs b/crates/types/src/ir/wire.rs new file mode 100644 index 000000000..a36243c3c --- /dev/null +++ b/crates/types/src/ir/wire.rs @@ -0,0 +1,746 @@ +//! Flat operator/scalar payloads shared by logical transport. +use std::rc::Rc; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; + +use super::asap::ASAPOp; +use super::node::{Operator, OperatorNode}; +use super::non_asap::{BinaryOperator, NonASAPOp, TimeRangeKind}; +use super::scalar::{ExprSemantics, Predicate, ProjectItem, ScalarExpr, SortKey}; +use crate::ir::operator_properties::{ + ConcatDiscriminatorKey, GroupKeys, InfoMatcher, JoinKind, Reduction, RelationalSetOpKind, + SampleKind, Source, TimeShift, WindowFrame, WindowFuncKind, +}; +use crate::post_asap::maintained_population::{MaintainedPopulation, PopulationStatistic}; +use crate::post_asap::sketch::{GroupingStrategy, SketchStatistic, SummaryUpdate}; +use crate::pre_asap::agg_intent::AggIntent; +use crate::pre_asap::expr_ir::{ArithmeticOpKind, CompareOpKind, ScalarValue}; +use crate::pre_asap::schema::{ColumnId, DataType, FieldDataType, Schema}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum EdgeRole { + Input, + Left, + Right, + /// The consumer reads the producer from inside one of its scalar + /// expressions (`scalar(v)`, a scalar subquery, `EXISTS`, `IN`). + ScalarRef, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +pub enum GroupingEdgeCompatibility { + Identical, + ConsumerCoarsensProducer, + Incompatible, + NotApplicable, +} + +/// Stable identity of a node within one exported logical ASAP DAG. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)] +#[serde(transparent)] +pub struct LogicalASAPNodeId(pub u32); + +// ── Wire mirrors of the scalar language ────────────────────────────────── + +/// [`ScalarExpr`] with every operator reference replaced by the id of the +/// exported node (connected to the owner by an [`EdgeRole::ScalarRef`] edge). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub enum WireScalarExpr { + Column(ColumnId), + Literal(ScalarValue), + Negative { + expr: Box, + semantics: ExprSemantics, + }, + Compare { + left: Box, + op: CompareOpKind, + right: Box, + semantics: ExprSemantics, + }, + BoolAnd(Vec), + BoolOr(Vec), + Not(Box), + IsNull(Box), + IsNotNull(Box), + Cast { + expr: Box, + to: DataType, + try_cast: bool, + }, + InList { + expr: Box, + list: Vec, + negated: bool, + }, + FunctionCall { + name: String, + args: Vec, + }, + Arithmetic { + op: ArithmeticOpKind, + left: Box, + right: Box, + semantics: ExprSemantics, + }, + Case { + operand: Option>, + branches: Vec<(WireScalarExpr, WireScalarExpr)>, + else_expr: Option>, + }, + CurrentTimestamp, + EvalTimestamp, + PromqlScalarFromVector(LogicalASAPNodeId), + ScalarSubquery(LogicalASAPNodeId), + Exists { + subquery: LogicalASAPNodeId, + negated: bool, + }, + InSubquery { + expr: Box, + subquery: LogicalASAPNodeId, + negated: bool, + }, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct WirePredicate(pub WireScalarExpr); + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct WireProjectItem { + pub alias: Option, + pub expr: WireScalarExpr, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct WireSortKey { + pub expr: WireScalarExpr, + pub ascending: bool, + pub nulls_first: bool, +} + +impl WireScalarExpr { + /// Explicit producer IDs recursively referenced by this scalar tree. + pub fn operator_refs(&self) -> Vec { + fn collect(expr: &WireScalarExpr, out: &mut Vec) { + use WireScalarExpr::*; + match expr { + PromqlScalarFromVector(id) | ScalarSubquery(id) => out.push(*id), + Exists { subquery, .. } => out.push(*subquery), + InSubquery { expr, subquery, .. } => { + out.push(*subquery); + collect(expr, out); + } + Negative { expr, .. } + | Not(expr) + | IsNull(expr) + | IsNotNull(expr) + | Cast { expr, .. } => collect(expr, out), + Compare { left, right, .. } | Arithmetic { left, right, .. } => { + collect(left, out); + collect(right, out); + } + BoolAnd(args) | BoolOr(args) | FunctionCall { args, .. } => { + for expr in args { + collect(expr, out); + } + } + InList { expr, list, .. } => { + collect(expr, out); + for expr in list { + collect(expr, out); + } + } + Case { + operand, + branches, + else_expr, + } => { + if let Some(expr) = operand { + collect(expr, out); + } + for (when, then) in branches { + collect(when, out); + collect(then, out); + } + if let Some(expr) = else_expr { + collect(expr, out); + } + } + Column(_) | Literal(_) | CurrentTimestamp | EvalTimestamp => {} + } + } + let mut out = Vec::new(); + collect(self, &mut out); + out + } + + /// Mirror `expr`, resolving every operator reference through `id_of`. + pub fn from_expr( + expr: &ScalarExpr, + id_of: &mut impl FnMut(&Rc) -> LogicalASAPNodeId, + ) -> Self { + fn boxed( + e: &ScalarExpr, + id_of: &mut impl FnMut(&Rc) -> LogicalASAPNodeId, + ) -> Box { + Box::new(WireScalarExpr::from_expr(e, id_of)) + } + fn list( + es: &[ScalarExpr], + id_of: &mut impl FnMut(&Rc) -> LogicalASAPNodeId, + ) -> Vec { + es.iter() + .map(|e| WireScalarExpr::from_expr(e, id_of)) + .collect() + } + match expr { + ScalarExpr::Column(id) => WireScalarExpr::Column(*id), + ScalarExpr::Literal(v) => WireScalarExpr::Literal(v.clone()), + ScalarExpr::Negative { expr, semantics } => WireScalarExpr::Negative { + expr: boxed(expr, id_of), + semantics: *semantics, + }, + ScalarExpr::Compare { + left, + op, + right, + semantics, + } => WireScalarExpr::Compare { + left: boxed(left, id_of), + op: op.clone(), + right: boxed(right, id_of), + semantics: *semantics, + }, + ScalarExpr::BoolAnd(parts) => WireScalarExpr::BoolAnd(list(parts, id_of)), + ScalarExpr::BoolOr(parts) => WireScalarExpr::BoolOr(list(parts, id_of)), + ScalarExpr::Not(e) => WireScalarExpr::Not(boxed(e, id_of)), + ScalarExpr::IsNull(e) => WireScalarExpr::IsNull(boxed(e, id_of)), + ScalarExpr::IsNotNull(e) => WireScalarExpr::IsNotNull(boxed(e, id_of)), + ScalarExpr::Cast { expr, to, try_cast } => WireScalarExpr::Cast { + expr: boxed(expr, id_of), + to: to.clone(), + try_cast: *try_cast, + }, + ScalarExpr::InList { + expr, + list: items, + negated, + } => WireScalarExpr::InList { + expr: boxed(expr, id_of), + list: list(items, id_of), + negated: *negated, + }, + ScalarExpr::FunctionCall { name, args } => WireScalarExpr::FunctionCall { + name: name.clone(), + args: list(args, id_of), + }, + ScalarExpr::Arithmetic { + op, + left, + right, + semantics, + } => WireScalarExpr::Arithmetic { + op: op.clone(), + left: boxed(left, id_of), + right: boxed(right, id_of), + semantics: *semantics, + }, + ScalarExpr::Case { + operand, + branches, + else_expr, + } => WireScalarExpr::Case { + operand: operand.as_ref().map(|e| boxed(e, id_of)), + branches: branches + .iter() + .map(|(w, t)| (Self::from_expr(w, id_of), Self::from_expr(t, id_of))) + .collect(), + else_expr: else_expr.as_ref().map(|e| boxed(e, id_of)), + }, + ScalarExpr::CurrentTimestamp => WireScalarExpr::CurrentTimestamp, + ScalarExpr::EvalTimestamp => WireScalarExpr::EvalTimestamp, + ScalarExpr::PromqlScalarFromVector(node) => { + WireScalarExpr::PromqlScalarFromVector(id_of(node)) + } + ScalarExpr::ScalarSubquery(node) => WireScalarExpr::ScalarSubquery(id_of(node)), + ScalarExpr::Exists { subquery, negated } => WireScalarExpr::Exists { + subquery: id_of(subquery), + negated: *negated, + }, + ScalarExpr::InSubquery { + expr, + subquery, + negated, + } => WireScalarExpr::InSubquery { + expr: boxed(expr, id_of), + subquery: id_of(subquery), + negated: *negated, + }, + } + } +} + +impl WirePredicate { + fn from_pred( + p: &Predicate, + id_of: &mut impl FnMut(&Rc) -> LogicalASAPNodeId, + ) -> Self { + WirePredicate(WireScalarExpr::from_expr(&p.0, id_of)) + } +} + +impl WireSortKey { + fn from_keys( + keys: &[SortKey], + id_of: &mut impl FnMut(&Rc) -> LogicalASAPNodeId, + ) -> Vec { + keys.iter() + .map(|k| WireSortKey { + expr: WireScalarExpr::from_expr(&k.expr, id_of), + ascending: k.ascending, + nulls_first: k.nulls_first, + }) + .collect() + } +} + +// ── Wire mirror of the non-ASAP operator vocabulary ────────────────────── + +/// [`NonASAPOp`] without its child fields (children are edges) and with +/// every scalar expression mirrored as [`WireScalarExpr`]. Fields named +/// `kind` in the IR are renamed so they do not collide with the variant tag. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum NonASAPOpKind { + Scan { + source: Source, + #[serde(default)] + predicates: Vec, + schema: Schema, + }, + Values { + rows: Vec>, + schema: Schema, + }, + Filter { + pred: WirePredicate, + }, + Project { + cols: Vec, + #[serde(default)] + qualifier: Option, + }, + Aggregate { + reduction: Reduction, + measures: Vec, + #[serde(default)] + output_names: Vec, + #[serde(default)] + filters: Vec>, + #[serde(default)] + having: Option, + }, + Join { + join_kind: JoinKind, + pred: WirePredicate, + }, + SetOp { + set_kind: RelationalSetOpKind, + all: bool, + }, + Concat { + #[serde(default)] + discriminator_unique_key: Option, + }, + Dedup { + cols: Vec, + }, + Sort { + keys: Vec, + #[serde(default)] + partition_by: GroupKeys, + }, + Limit { + n: Option, + offset: usize, + #[serde(default)] + partition_by: GroupKeys, + }, + BinaryOp { + operator: BinaryOperator, + #[serde(default)] + return_bool: bool, + }, + #[serde(rename = "sql_window_func")] + SQLWindowFunc { + func: WindowFuncKind, + args: Vec, + partition_by: GroupKeys, + order_by: Vec, + #[serde(default)] + frame: Option, + output_name: String, + }, + TimeRange { + range: Duration, + range_kind: TimeRangeKind, + }, + TimeShift { + shift: TimeShift, + }, + PromqlVectorFromScalar { + expr: WireScalarExpr, + }, + PromqlRelabel { + dst: String, + value: WireScalarExpr, + }, + PromqlInfoEnrich { + #[serde(default)] + selector: Vec, + }, + PromqlSeriesSample { + #[serde(default)] + by: GroupKeys, + sample_kind: SampleKind, + }, + PromqlSubquery { + range: Duration, + #[serde(default)] + resolution: Option, + }, +} + +impl NonASAPOpKind { + /// Mirror `op`, resolving every operator node its scalar expressions + /// reference through `id_of`. + pub fn from_op( + op: &NonASAPOp, + id_of: &mut impl FnMut(&Rc) -> LogicalASAPNodeId, + ) -> Self { + use NonASAPOp as Op; + match op { + Op::Scan { + source, + predicates, + schema, + } => NonASAPOpKind::Scan { + source: source.clone(), + predicates: predicates + .iter() + .map(|p| WirePredicate::from_pred(p, id_of)) + .collect(), + schema: schema.clone(), + }, + Op::Values { rows, schema } => NonASAPOpKind::Values { + rows: rows + .iter() + .map(|row| { + row.iter() + .map(|e| WireScalarExpr::from_expr(e, id_of)) + .collect() + }) + .collect(), + schema: schema.clone(), + }, + Op::Filter { pred, .. } => NonASAPOpKind::Filter { + pred: WirePredicate::from_pred(pred, id_of), + }, + Op::Project { + cols, qualifier, .. + } => NonASAPOpKind::Project { + cols: cols + .iter() + .map(|ProjectItem { alias, expr }| WireProjectItem { + alias: alias.clone(), + expr: WireScalarExpr::from_expr(expr, id_of), + }) + .collect(), + qualifier: qualifier.clone(), + }, + Op::Aggregate { + reduction, + measures, + output_names, + filters, + having, + .. + } => NonASAPOpKind::Aggregate { + reduction: reduction.clone(), + measures: measures.clone(), + output_names: output_names.clone(), + filters: filters + .iter() + .map(|p| p.as_ref().map(|p| WirePredicate::from_pred(p, id_of))) + .collect(), + having: having.as_ref().map(|p| WirePredicate::from_pred(p, id_of)), + }, + Op::Join { kind, pred, .. } => NonASAPOpKind::Join { + join_kind: kind.clone(), + pred: WirePredicate::from_pred(pred, id_of), + }, + Op::SetOp { kind, all, .. } => NonASAPOpKind::SetOp { + set_kind: kind.clone(), + all: *all, + }, + Op::Concat { + discriminator_unique_key, + .. + } => NonASAPOpKind::Concat { + discriminator_unique_key: discriminator_unique_key.clone(), + }, + Op::Dedup { cols, .. } => NonASAPOpKind::Dedup { cols: cols.clone() }, + Op::Sort { + keys, partition_by, .. + } => NonASAPOpKind::Sort { + keys: WireSortKey::from_keys(keys, id_of), + partition_by: partition_by.clone(), + }, + Op::Limit { + n, + offset, + partition_by, + .. + } => NonASAPOpKind::Limit { + n: *n, + offset: *offset, + partition_by: partition_by.clone(), + }, + Op::BinaryOp { + operator, + return_bool, + .. + } => NonASAPOpKind::BinaryOp { + operator: operator.clone(), + return_bool: *return_bool, + }, + Op::SQLWindowFunc { + func, + args, + partition_by, + order_by, + frame, + output_name, + .. + } => NonASAPOpKind::SQLWindowFunc { + func: func.clone(), + args: args + .iter() + .map(|e| WireScalarExpr::from_expr(e, id_of)) + .collect(), + partition_by: partition_by.clone(), + order_by: WireSortKey::from_keys(order_by, id_of), + frame: frame.clone(), + output_name: output_name.clone(), + }, + Op::TimeRange { range, kind, .. } => NonASAPOpKind::TimeRange { + range: *range, + range_kind: *kind, + }, + Op::TimeShift { shift, .. } => NonASAPOpKind::TimeShift { shift: *shift }, + Op::PromqlVectorFromScalar(e) => NonASAPOpKind::PromqlVectorFromScalar { + expr: WireScalarExpr::from_expr(e, id_of), + }, + Op::PromqlRelabel { dst, value, .. } => NonASAPOpKind::PromqlRelabel { + dst: dst.clone(), + value: WireScalarExpr::from_expr(value, id_of), + }, + Op::PromqlInfoEnrich { selector, .. } => NonASAPOpKind::PromqlInfoEnrich { + selector: selector.clone(), + }, + Op::PromqlSeriesSample { by, kind, .. } => NonASAPOpKind::PromqlSeriesSample { + by: by.clone(), + sample_kind: *kind, + }, + Op::PromqlSubquery { + range, resolution, .. + } => NonASAPOpKind::PromqlSubquery { + range: *range, + resolution: *resolution, + }, + } + } +} + +// ── The exported DAG ───────────────────────────────────────────────────── + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)] +pub enum LogicalASAPOperatorPayload { + Relational { + operator: NonASAPOpKind, + }, + SummaryAgg { + family: FieldDataType, + input: SummaryUpdate, + reduction: Reduction, + grouping: GroupingStrategy, + #[serde(default)] + filter: Option, + }, + SummaryEstimate { + query: SketchStatistic, + }, + FinalizeExactAccumulator, + MaintainPopulation { + population: MaintainedPopulation, + }, + EvaluatePopulation { + evaluation: PopulationStatistic, + }, + SummaryMerge, + SummarySubtract, + SummaryDelete { + key: ColumnId, + }, + SummaryJoin { + key: ColumnId, + family: FieldDataType, + }, + Extension { + name: String, + }, +} + +/// The operator's own inputs with their edge roles, in field order. +pub(super) fn input_edges(operator: &Operator) -> Vec<(&Rc, EdgeRole)> { + match operator { + Operator::NonASAP(op) => match op { + NonASAPOp::Join { left, right, .. } + | NonASAPOp::SetOp { left, right, .. } + | NonASAPOp::BinaryOp { + lhs: left, + rhs: right, + .. + } => vec![(left, EdgeRole::Left), (right, EdgeRole::Right)], + NonASAPOp::Concat { children, .. } => { + children.iter().map(|c| (c, EdgeRole::Input)).collect() + } + NonASAPOp::Filter { child, .. } + | NonASAPOp::Project { child, .. } + | NonASAPOp::Aggregate { child, .. } + | NonASAPOp::Dedup { child, .. } + | NonASAPOp::Sort { child, .. } + | NonASAPOp::Limit { child, .. } + | NonASAPOp::SQLWindowFunc { child, .. } + | NonASAPOp::TimeRange { child, .. } + | NonASAPOp::TimeShift { child, .. } + | NonASAPOp::PromqlRelabel { child, .. } + | NonASAPOp::PromqlInfoEnrich { child, .. } + | NonASAPOp::PromqlSeriesSample { child, .. } + | NonASAPOp::PromqlSubquery { child, .. } => vec![(child, EdgeRole::Input)], + NonASAPOp::Scan { .. } + | NonASAPOp::Values { .. } + | NonASAPOp::PromqlVectorFromScalar(_) => vec![], + }, + Operator::ASAP(op) => match op { + ASAPOp::SummarySubtract { left, right } + | ASAPOp::SummaryJoin { + outer: left, + inner: right, + .. + } => vec![(left, EdgeRole::Left), (right, EdgeRole::Right)], + ASAPOp::SummaryMerge { children } => { + children.iter().map(|c| (c, EdgeRole::Input)).collect() + } + ASAPOp::SummaryAgg { child, .. } + | ASAPOp::FinalizeExactAccumulator { child } + | ASAPOp::MaintainPopulation { child, .. } + | ASAPOp::EvaluatePopulation { child, .. } + | ASAPOp::Extension { child, .. } => vec![(child, EdgeRole::Input)], + ASAPOp::SummaryEstimate { summary_input, .. } + | ASAPOp::SummaryDelete { summary_input, .. } => { + vec![(summary_input, EdgeRole::Input)] + } + }, + } +} + +pub(super) fn payload_of( + operator: &Operator, + id_of: &mut impl FnMut(&Rc) -> LogicalASAPNodeId, +) -> LogicalASAPOperatorPayload { + match operator { + Operator::NonASAP(op) => LogicalASAPOperatorPayload::Relational { + operator: NonASAPOpKind::from_op(op, id_of), + }, + Operator::ASAP(op) => match op { + ASAPOp::SummaryAgg { + family, + input, + reduction, + grouping, + filter, + .. + } => LogicalASAPOperatorPayload::SummaryAgg { + family: family.clone(), + input: input.clone(), + reduction: reduction.clone(), + grouping: grouping.clone(), + filter: filter.as_ref().map(|p| WirePredicate::from_pred(p, id_of)), + }, + ASAPOp::SummaryEstimate { query, .. } => LogicalASAPOperatorPayload::SummaryEstimate { + query: query.clone(), + }, + ASAPOp::FinalizeExactAccumulator { .. } => { + LogicalASAPOperatorPayload::FinalizeExactAccumulator + } + ASAPOp::MaintainPopulation { population, .. } => { + LogicalASAPOperatorPayload::MaintainPopulation { + population: population.clone(), + } + } + ASAPOp::EvaluatePopulation { evaluation, .. } => { + LogicalASAPOperatorPayload::EvaluatePopulation { + evaluation: evaluation.clone(), + } + } + ASAPOp::SummaryMerge { .. } => LogicalASAPOperatorPayload::SummaryMerge, + ASAPOp::SummarySubtract { .. } => LogicalASAPOperatorPayload::SummarySubtract, + ASAPOp::SummaryDelete { key, .. } => { + LogicalASAPOperatorPayload::SummaryDelete { key: *key } + } + ASAPOp::SummaryJoin { key, family, .. } => LogicalASAPOperatorPayload::SummaryJoin { + key: *key, + family: family.clone(), + }, + ASAPOp::Extension { name, .. } => { + LogicalASAPOperatorPayload::Extension { name: name.clone() } + } + }, + } +} + +/// Grouping compatibility between two `SummaryAgg`s by their reductions. +pub(super) fn grouping_compatibility( + producer: &Operator, + consumer: &Operator, +) -> GroupingEdgeCompatibility { + let ( + Operator::ASAP(ASAPOp::SummaryAgg { + reduction: producer, + .. + }), + Operator::ASAP(ASAPOp::SummaryAgg { + reduction: consumer, + .. + }), + ) = (producer, consumer) + else { + return GroupingEdgeCompatibility::NotApplicable; + }; + match (producer, consumer) { + (p, c) if p == c => GroupingEdgeCompatibility::Identical, + (Reduction::PerEntity, Reduction::Reduce(_)) => { + GroupingEdgeCompatibility::ConsumerCoarsensProducer + } + (Reduction::Reduce(p), Reduction::Reduce(c)) + if !p.is_without() && !c.is_without() && c.iter().all(|key| p.contains(key)) => + { + GroupingEdgeCompatibility::ConsumerCoarsensProducer + } + _ => GroupingEdgeCompatibility::Incompatible, + } +} diff --git a/crates/types/tests/logical_export.rs b/crates/types/tests/logical_export.rs new file mode 100644 index 000000000..95c812cec --- /dev/null +++ b/crates/types/tests/logical_export.rs @@ -0,0 +1,239 @@ +//! Logical transport must accept plans with no execution timing assigned, before materialization. +use asap_types::{ + ir::{NonASAPOp, Operator, OperatorNode}, + pre_asap::Schema, +}; + +#[test] +fn logical_export_accepts_unassigned_timing() { + let root = OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Values { + rows: vec![vec![]], + schema: Schema::lifted(vec![], None), + })) + .unwrap(); + assert!(root.timing.is_none()); + assert!(asap_types::ir::export::compile_logical_asap_dag(&root).is_ok()); +} + +use asap_types::{ + ir::export::{ + compile_logical_asap_dag_with_node_ids, EdgeRole, LogicalASAPDAGDocument, + LogicalASAPDAGValidationError, LogicalASAPNodeId, LogicalASAPOperatorPayload, + }, + ir::operator_properties::Reduction, + ir::{ASAPOp, OperatorResultKind, ProjectItem, ScalarExpr}, + post_asap::{GroupingStrategy, SketchAlgorithm, SketchKind, SketchParams, SummaryUpdate}, + pre_asap::{ColumnRef, DataType, Field, FieldDataType, ScalarValue, Source}, +}; +use std::rc::Rc; + +fn values() -> Rc { + OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Values { + rows: vec![vec![ScalarExpr::Literal(ScalarValue::Float64(1.0))]], + schema: Schema::lifted(vec![Field::plain("value", DataType::Float64, false)], None), + })) + .unwrap() +} + +/// Scalar subqueries contribute real edges; repeated references export one producer. +#[test] +fn scalar_dependencies_share_one_exported_producer() { + let child = values(); + let root = OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Project { + cols: vec![ProjectItem { + alias: Some("result".into()), + expr: ScalarExpr::ScalarSubquery(child.clone()), + }], + qualifier: None, + child: child.clone(), + })) + .unwrap(); + let compiled = compile_logical_asap_dag_with_node_ids(&root).unwrap(); + compiled.dag.validate().unwrap(); + assert_eq!(compiled.dag.nodes.len(), 2); + assert_eq!(compiled.dag.edges.len(), 2); + assert!(compiled + .dag + .edges + .iter() + .any(|edge| edge.role == EdgeRole::ScalarRef)); + let id = compiled.node_ids.node_id(&child).unwrap(); + assert!(Rc::ptr_eq( + compiled.node_ids.operator_node(id).unwrap(), + &child + )); + let document = LogicalASAPDAGDocument::new(compiled.dag); + let json = serde_json::to_string(&document).unwrap(); + for physical_metadata in [ + "output_state", + "data_state", + "timing", + "retention", + "window", + ] { + assert!( + !json.contains(physical_metadata), + "logical JSON contains {physical_metadata}" + ); + } + let decoded: LogicalASAPDAGDocument = serde_json::from_str(&json).unwrap(); + assert_eq!(document, decoded); + decoded.validate().unwrap(); +} + +/// Summary state identity survives logical merge export without a phase assignment. +#[test] +fn merged_summary_preserves_typed_state() { + let state = OperatorNode::new_shared(Operator::ASAP(ASAPOp::SummaryAgg { + child: values(), + family: FieldDataType::Sketch( + SketchKind::new(SketchAlgorithm::Kll, SketchParams::Kll { k: 200 }), + GroupingStrategy::default(), + ), + input: SummaryUpdate::column(ColumnRef::Named("value".into())), + reduction: Reduction::by(vec![]), + grouping: GroupingStrategy::default(), + filter: None, + })) + .unwrap(); + let root = OperatorNode::new_shared(Operator::ASAP(ASAPOp::SummaryMerge { + children: (0..2) + .map(|start| { + let coverage = asap_types::ir::summary_coverage::SummaryCoverage { + source: Source::Table { + table_ref: "values".into(), + }, + regions: vec![asap_types::ir::summary_coverage::CoverageRegion { + time_ms: Some(start..start + 1), + population: Default::default(), + }], + }; + std::rc::Rc::new((*state).clone().with_coverage(coverage).unwrap()) + }) + .collect(), + })) + .unwrap(); + let dag = asap_types::ir::export::compile_logical_asap_dag(&root).unwrap(); + dag.validate().unwrap(); + assert_eq!(dag.nodes.len(), 4); + let root_id = dag.roots[0].operator_refs()[0]; + let merged = &dag.nodes[root_id.0 as usize]; + assert_eq!(merged.result_kind, OperatorResultKind::State); + assert_eq!(merged.output_schema, root.schema); + assert_eq!(merged.coverage, root.coverage); + assert!(matches!( + merged.payload, + LogicalASAPOperatorPayload::SummaryMerge + )); + // Transport rejects a summary producer whose required coverage was dropped. + let mut stripped = dag.clone(); + let producer = stripped + .nodes + .iter_mut() + .find(|node| matches!(node.payload, LogicalASAPOperatorPayload::SummaryAgg { .. })) + .unwrap(); + producer.coverage = None; + let id = producer.id; + assert!(matches!( + stripped.validate(), + Err(LogicalASAPDAGValidationError::InvalidCoverage(bad)) if bad == id + )); +} + +/// Malformed wire graphs fail transport integrity checks rather than reaching execution. +#[test] +fn malformed_transport_is_rejected() { + let dag = asap_types::ir::export::compile_logical_asap_dag(&values()).unwrap(); + let mut document = LogicalASAPDAGDocument::new(dag.clone()); + document.schema_version = 99; + assert!(matches!( + document.validate(), + Err(LogicalASAPDAGValidationError::UnsupportedVersion(99)) + )); + let mut duplicate = dag.clone(); + duplicate.nodes.push(dag.nodes[0].clone()); + assert!(matches!( + duplicate.validate(), + Err(LogicalASAPDAGValidationError::DuplicateNode(_)) + )); + let mut missing = dag.clone(); + missing.roots = vec![asap_types::ir::export::LogicalASAPQueryRoot::Operator( + LogicalASAPNodeId(9), + )]; + assert!(matches!( + missing.validate(), + Err(LogicalASAPDAGValidationError::MissingNode(_)) + )); + let mut unreachable = dag.clone(); + let mut extra = dag.nodes[0].clone(); + extra.id = LogicalASAPNodeId(1); + unreachable.nodes.push(extra); + assert!(matches!( + unreachable.validate(), + Err(LogicalASAPDAGValidationError::UnreachableNode(_)) + )); + let mut json = serde_json::to_value(LogicalASAPDAGDocument::new(dag)).unwrap(); + json["dag"]["nodes"][0]["output_state"] = serde_json::json!({"timing":"ingestion_time"}); + assert!(serde_json::from_value::(json).is_err()); +} + +/// Standalone constants need no fake relation, while scalar subqueries retain their producer DAG. +#[test] +fn standalone_scalar_roots_roundtrip_without_synthetic_operators() { + use asap_types::ir::{export::compile_logical_asap_query, QueryRoot}; + for root in [ + QueryRoot::Scalar(ScalarExpr::literal_f64(42.0)), + QueryRoot::Scalar(ScalarExpr::ScalarSubquery(values())), + ] { + let expected = root.as_operator().is_some(); + assert!(!expected); + let dag = compile_logical_asap_query(&root).unwrap(); + dag.validate().unwrap(); + let expected_nodes = match root { + QueryRoot::Scalar(ScalarExpr::ScalarSubquery(_)) => 1, + _ => 0, + }; + assert_eq!(dag.nodes.len(), expected_nodes); + let document = LogicalASAPDAGDocument::new(dag); + let decoded: LogicalASAPDAGDocument = + serde_json::from_str(&serde_json::to_string(&document).unwrap()).unwrap(); + decoded.validate().unwrap(); + assert_eq!(document, decoded); + } +} + +/// A batch exports as one DAG: one root per query, shared producers exported once. +#[test] +fn batch_exports_one_root_per_query_and_shares_producers() { + use asap_types::ir::{export::compile_logical_asap_workload, QueryRoot}; + let shared = values(); + let project = |alias: &str| { + OperatorNode::new_shared(Operator::NonASAP(NonASAPOp::Project { + child: shared.clone(), + cols: vec![ProjectItem { + alias: Some(alias.into()), + expr: ScalarExpr::Column(0), + }], + qualifier: None, + })) + .unwrap() + }; + let dag = compile_logical_asap_workload(&[ + QueryRoot::Operator(project("a")), + QueryRoot::Operator(project("b")), + ]) + .unwrap(); + dag.validate().unwrap(); + assert_eq!(dag.roots.len(), 2); + assert_eq!( + dag.nodes.len(), + 3, + "the shared Values node is exported once" + ); + let mut empty = dag.clone(); + empty.roots.clear(); + assert!(matches!( + empty.validate(), + Err(LogicalASAPDAGValidationError::NoRoots) + )); +} diff --git a/docs/develop_docs/logical-asap-dag.md b/docs/develop_docs/logical-asap-dag.md new file mode 100644 index 000000000..99696b175 --- /dev/null +++ b/docs/develop_docs/logical-asap-dag.md @@ -0,0 +1,50 @@ +# Logical ASAP DAG transport + +This interface exports the unified operator/scalar IR at the logical stage of +[planner layering](../design_docs/proposals/planner-layering.md). It describes +what to compute, including committed summary families, before physical planning +chooses implementations and materialization. + +## Interface + +`asap_types::ir::export` exposes: + +- `compile_logical_asap_dag(&Rc)` for a flat `LogicalASAPDAG`. +- `compile_logical_asap_query(&QueryRoot)` for operator or standalone scalar roots. +- `compile_logical_asap_dag_with_node_ids(...)` for that DAG and a compiler-local + `LogicalASAPNodeIdentityMap` with `node_id` and `operator_node` lookups. +- `LogicalASAPDAGDocument::new(dag)` and `validate()` for the versioned transport + envelope. Logical wire version 1 is distinct from the older phase-assigned + post-ASAP format. + +The compiler first checks the in-memory DAG's structural contracts. Untimed +ordinary plans and summary plans are valid inputs. Export does not assess +accuracy against request requirements or select a physical plan. + +Each `LogicalASAPDAGNode` contains an ID, operator payload, result kind, output +schema and optional accuracy guarantee. Each `LogicalASAPDAGEdge` contains +producer/consumer IDs, input role, intermediate schema and grouping compatibility. +The DAG has one semantic operator or scalar root. A standalone scalar constant +needs no synthetic operator node. IDs are local to one export. + +Scalar expressions remain owned by their operators. Their wire representations +replace explicit operator references with IDs. `ScalarRef` edges record those +producer dependencies. A shared operator is exported once even when ordinary +inputs and scalar expressions both reference it. + +## Physical boundary + +Logical nodes and edges contain no execution state, assigned timing, storage tier, +retention or pane-alignment assertion. Physical planning chooses, for each eligible +sub-DAG, no materialization, query-time materialization, or ingestion-time +materialization. Execution timing follows that choice and its dependencies. + +The optional `OperatorNode.timing` field belongs to the common IR and may later +record a physical assignment; it is not part of logical transport. There is no +intermediate timed-DAG stage. Physical planning owns phase validation and any +splitting needed when shared consumers require incompatible execution contexts. + +Transport validation checks graph identity, connectivity, acyclicity, edge schemas +and declared summary family/grouping metadata. Full scalar/operator typing remains +an in-memory structural validation responsibility. Neither check proves runtime +capability, cost, response latency or accuracy feasibility.