diff --git a/crates/asap-aware-mapping/src/explanation.rs b/crates/asap-aware-mapping/src/explanation.rs index bcfee70..026df91 100644 --- a/crates/asap-aware-mapping/src/explanation.rs +++ b/crates/asap-aware-mapping/src/explanation.rs @@ -438,7 +438,7 @@ fn visit_children( | SQLWindowFunc { child, .. } | Sort { child, .. } | Limit { child, .. } => visit(child, format!("{label} > child"), locations), - Concat { children } => { + Concat { children, .. } => { for (i, c) in children.iter().enumerate() { visit_children(c, &format!("{label} > concat[{i}]"), locations); } diff --git a/crates/asap-aware-mapping/src/replacement.rs b/crates/asap-aware-mapping/src/replacement.rs index e18a9a5..2edf92e 100644 --- a/crates/asap-aware-mapping/src/replacement.rs +++ b/crates/asap-aware-mapping/src/replacement.rs @@ -2355,6 +2355,7 @@ fn direct_child_counts(node: &QueryExpr) -> Vec<(*const QueryExpr, usize)> { } Concat { children: concat_children, + .. } => { for c in concat_children { collect(c, children); @@ -2808,7 +2809,7 @@ fn walk_children( | SQLWindowFunc { child, .. } | Sort { child, .. } | Limit { child, .. } => walk(child, order, nodes, counts), - Concat { children } => { + Concat { children, .. } => { for c in children { walk_children(c, order, nodes, counts); } @@ -3718,7 +3719,7 @@ mod tests { | SQLWindowFunc { child, .. } | Sort { child, .. } | Limit { child, .. } => walk(child, counts), - Concat { children } => { + Concat { children, .. } => { for c in children { walk_children(c, counts); } diff --git a/crates/devtools/src/bin/variant_coverage.rs b/crates/devtools/src/bin/variant_coverage.rs index a1b3b82..0cd66ee 100644 --- a/crates/devtools/src/bin/variant_coverage.rs +++ b/crates/devtools/src/bin/variant_coverage.rs @@ -88,7 +88,7 @@ fn walk(e: &QueryExpr, seen: &mut BTreeSet<&'static str>) { seen.insert("Dedup"); walk(child, seen); } - QueryExpr::Concat { children } => { + QueryExpr::Concat { children, .. } => { seen.insert("Concat"); children.iter().for_each(|c| walk(c, seen)); } diff --git a/crates/frontend-promql/src/promql.rs b/crates/frontend-promql/src/promql.rs index 690ecbf..f43a00d 100644 --- a/crates/frontend-promql/src/promql.rs +++ b/crates/frontend-promql/src/promql.rs @@ -742,7 +742,13 @@ fn walk_histogram_quantiles(call: &Call) -> Result { }) }) .collect::>>()?; - Ok(Unresolved::Concat { children: branches }) + // No discriminator asserted here today (issue #228): the φ value each + // branch carries via `PromqlRelabel` *is* structurally a distinct + // per-branch discriminator, but nothing downstream currently needs the + // resulting compound unique key — see + // `docs/design_docs/concat-unique-keys-decision.md`. `Unresolved::concat` + // keeps `output_schema`'s default (drop `unique_keys` entirely). + Ok(Unresolved::concat(branches)) } /// Prometheus's `labels.FormatOpenMetricsFloat` — how `histogram_quantiles` diff --git a/crates/frontend-promql/tests/observability/awesome_prometheus_alerts.rs b/crates/frontend-promql/tests/observability/awesome_prometheus_alerts.rs index b616d75..254297d 100644 --- a/crates/frontend-promql/tests/observability/awesome_prometheus_alerts.rs +++ b/crates/frontend-promql/tests/observability/awesome_prometheus_alerts.rs @@ -83,7 +83,7 @@ fn intents(e: &QueryExpr) -> Vec { go(lhs, out); go(rhs, out); } - QueryExpr::Concat { children } => children.iter().for_each(|c| go(c, out)), + QueryExpr::Concat { children, .. } => children.iter().for_each(|c| go(c, out)), QueryExpr::PromqlVectorFromScalar(inner) | QueryExpr::PromqlScalarFromVector(inner) => { go(inner, out) } diff --git a/crates/frontend-promql/tests/promql_conformance.rs b/crates/frontend-promql/tests/promql_conformance.rs index 698a1bd..89426ba 100644 --- a/crates/frontend-promql/tests/promql_conformance.rs +++ b/crates/frontend-promql/tests/promql_conformance.rs @@ -92,7 +92,7 @@ fn collect(e: &QueryExpr, out: &mut Vec) { collect(left, out); collect(right, out); } - QueryExpr::Concat { children } => children.iter().for_each(|c| collect(c, out)), + QueryExpr::Concat { children, .. } => children.iter().for_each(|c| collect(c, out)), QueryExpr::PromqlVectorFromScalar(inner) | QueryExpr::PromqlScalarFromVector(inner) => { collect(inner, out) } diff --git a/crates/frontend-promql/tests/promql_lowering.rs b/crates/frontend-promql/tests/promql_lowering.rs index 192118d..f59f133 100644 --- a/crates/frontend-promql/tests/promql_lowering.rs +++ b/crates/frontend-promql/tests/promql_lowering.rs @@ -905,7 +905,7 @@ fn topk_over_bare_selector_ranks_raw_samples() { /// The `(label value, intent)` of each `histogram_quantiles` branch. fn quantile_branches(q: &QueryExpr) -> Vec<(String, AggIntent)> { - let QueryExpr::Concat { children } = q else { + let QueryExpr::Concat { children, .. } = q else { panic!("expected a Concat at the root, got {q:?}"); }; children @@ -963,7 +963,7 @@ fn histogram_quantiles_branches_are_union_compatible() { // `Concat` derives its schema from the first child, so every branch must // agree on column names — the φ lives in the label, not the column name. let q = lower(r#"histogram_quantiles(testhistogram3, "q", 0.5, 0.9)"#); - let QueryExpr::Concat { children } = &q else { + let QueryExpr::Concat { children, .. } = &q else { panic!("expected Concat"); }; let shapes: Vec> = children @@ -989,7 +989,7 @@ fn histogram_quantiles_branches_are_union_compatible() { #[test] fn histogram_quantiles_uses_the_given_label_name() { let q = lower(r#"histogram_quantiles(h, "phi", 0.5)"#); - let QueryExpr::Concat { children } = &q else { + let QueryExpr::Concat { children, .. } = &q else { panic!("expected Concat"); }; let QueryExpr::PromqlRelabel { dst, .. } = &children[0] else { diff --git a/crates/frontend-sql/src/sql/mod.rs b/crates/frontend-sql/src/sql/mod.rs index 1a650b0..538df70 100644 --- a/crates/frontend-sql/src/sql/mod.rs +++ b/crates/frontend-sql/src/sql/mod.rs @@ -865,7 +865,14 @@ impl<'a> SqlLowerer<'a> { }) .collect::, LoweringError>>()?; - Ok(Unresolved::Concat { children: branches }) + // No discriminator asserted here today (issue #228): DataFusion's own + // `__grouping_id` would be the natural one, but this front end + // already discards it (see above — `GROUPING()` itself is rejected), + // so there is no distinct-per-branch column available to name yet. + // `Unresolved::concat` keeps `output_schema`'s default (drop + // `unique_keys` entirely). See + // `docs/design_docs/concat-unique-keys-decision.md`. + Ok(Unresolved::concat(branches)) } fn lower_sort(&self, sort: &logical_expr::Sort) -> Result { diff --git a/crates/frontend-sql/tests/data_quality_check/synthetic_packet_trace.rs b/crates/frontend-sql/tests/data_quality_check/synthetic_packet_trace.rs index 91e08e0..5831aaf 100644 --- a/crates/frontend-sql/tests/data_quality_check/synthetic_packet_trace.rs +++ b/crates/frontend-sql/tests/data_quality_check/synthetic_packet_trace.rs @@ -105,7 +105,7 @@ fn intents(e: &QueryExpr) -> Vec { go(lhs, out); go(rhs, out); } - QueryExpr::Concat { children } => children.iter().for_each(|c| go(c, out)), + QueryExpr::Concat { children, .. } => children.iter().for_each(|c| go(c, out)), QueryExpr::PromqlVectorFromScalar(inner) | QueryExpr::PromqlScalarFromVector(inner) => { go(inner, out) } diff --git a/crates/frontend-sql/tests/netflow/netflow.rs b/crates/frontend-sql/tests/netflow/netflow.rs index 2cf0cb7..09d742b 100644 --- a/crates/frontend-sql/tests/netflow/netflow.rs +++ b/crates/frontend-sql/tests/netflow/netflow.rs @@ -290,7 +290,7 @@ fn visit(qe: &QueryExpr, f: &mut impl FnMut(&QueryExpr)) { visit(lhs, f); visit(rhs, f); } - QueryExpr::Concat { children } => { + QueryExpr::Concat { children, .. } => { for child in children { visit(child, f); } diff --git a/crates/frontend-sql/tests/sql_lowering.rs b/crates/frontend-sql/tests/sql_lowering.rs index 06beba5..0b99cd4 100644 --- a/crates/frontend-sql/tests/sql_lowering.rs +++ b/crates/frontend-sql/tests/sql_lowering.rs @@ -1292,7 +1292,7 @@ async fn a_shared_expression_is_materialized_once() { fn merge_branches(qe: &QueryExpr) -> &Vec { fn find(qe: &QueryExpr) -> Option<&Vec> { match qe { - QueryExpr::Concat { children } => Some(children), + QueryExpr::Concat { children, .. } => Some(children), QueryExpr::Project { child, .. } | QueryExpr::Filter { child, .. } | QueryExpr::Sort { child, .. } diff --git a/crates/types/src/dag_export.rs b/crates/types/src/dag_export.rs index 252d319..6d6eda2 100644 --- a/crates/types/src/dag_export.rs +++ b/crates/types/src/dag_export.rs @@ -1011,21 +1011,18 @@ fn build_no_recheck( vec![c], ) } - QueryExpr::Concat { children } => { + QueryExpr::Concat { + children, + discriminator_unique_key, + } => { let ids: Vec = children .iter() .map(|c| build(c, nodes, cache, find_winner)) .collect(); let label = format!("Concat({} branches)", ids.len()); - push_node( - nodes, - expr, - cache, - "Concat", - label, - serde_json::json!({}), - ids, - ) + let detail = + serde_json::json!({ "discriminator_unique_key": discriminator_unique_key }); + push_node(nodes, expr, cache, "Concat", label, detail, ids) } QueryExpr::Join { kind, @@ -1294,13 +1291,11 @@ mod tests { #[test] fn merge_keeps_every_branch_as_a_child() { - let expr = QueryExpr::Concat { - children: vec![ - scan("a", value_col()), - scan("b", value_col()), - scan("c", value_col()), - ], - }; + let expr = QueryExpr::concat(vec![ + scan("a", value_col()), + scan("b", value_col()), + scan("c", value_col()), + ]); let graph = export(&expr); assert_eq!(graph.nodes.len(), 4, "3 branches + the Concat node"); let merge = &graph.nodes[graph.root as usize]; diff --git a/crates/types/src/pre_asap/binder.rs b/crates/types/src/pre_asap/binder.rs index 97061e9..221aebd 100644 --- a/crates/types/src/pre_asap/binder.rs +++ b/crates/types/src/pre_asap/binder.rs @@ -163,7 +163,7 @@ fn leftmost_scan_name(tree: &UnresolvedQueryExpr) -> Option<&str> { | QE::TimeRange { child, .. } | QE::TimeShift { child, .. } | QE::SQLWindowFunc { child, .. } => leftmost_scan_name(child), - QE::Concat { children } => children.first().and_then(leftmost_scan_name), + QE::Concat { children, .. } => children.first().and_then(leftmost_scan_name), QE::Join { left, .. } | QE::SetOp { left, .. } | QE::BinaryOp { lhs: left, .. } => { leftmost_scan_name(left) } @@ -310,7 +310,24 @@ pub(crate) fn collect_referenced_columns(tree: &UnresolvedQueryExpr) -> Vec walk(child, out), - QE::Concat { children } => children.iter().for_each(|c| walk(c, out)), + QE::Concat { + children, + discriminator_unique_key, + } => { + // Same treatment as `Dedup.cols` above: an own-field + // `ColumnRef` must be seeded here too, or a discriminator + // column that isn't otherwise referenced anywhere else in + // the tree (plausible — a raw usage-derived label, not one a + // `Project`/relabel freshly created) is absent from the + // Binder's usage-derived fallback schema, and + // `resolve.rs`'s later `resolve_column_ref` call fails with + // `NotFound` for a column the caller correctly named. + if let Some(key) = discriminator_unique_key { + push_ref_name(key.discriminator(), out); + key.inner_key().iter().for_each(|c| push_ref_name(c, out)); + } + children.iter().for_each(|c| walk(c, out)); + } QE::SetOp { left, right, .. } => { walk(left, out); walk(right, out); @@ -390,6 +407,29 @@ mod tests { assert!(schema.column_id("host").is_some()); } + /// Issue #228 review: a `Concat`'s `discriminator_unique_key` columns — + /// even one referenced nowhere else in the tree — must be seeded into + /// the usage-derived fallback schema, exactly like `Dedup.cols`, or + /// `resolve.rs`'s later `resolve_column_ref` fails `NotFound` for a + /// column the caller correctly named. + #[test] + fn concat_discriminator_key_is_seeded_into_the_binder_schema() { + let tree = UnresolvedQueryExpr::concat_with_discriminator( + vec![src("m")], + ColumnRef::Named("phi".into()), + vec![ColumnRef::Named("host".into())], + ); + let schema = Binder::new().bind(&tree); + assert!( + schema.column_id("phi").is_some(), + "discriminator column must be seeded" + ); + assert!( + schema.column_id("host").is_some(), + "inner_key column must be seeded" + ); + } + #[test] fn inherited_names_are_seeded_alongside_referenced() { // A `BinaryOp` side re-binds against its own sub-tree, but must still see diff --git a/crates/types/src/pre_asap/canonicalize.rs b/crates/types/src/pre_asap/canonicalize.rs index 6dda2ec..092b93b 100644 --- a/crates/types/src/pre_asap/canonicalize.rs +++ b/crates/types/src/pre_asap/canonicalize.rs @@ -38,11 +38,52 @@ pub fn canonicalize(mut expr: QueryExpr) -> QueryExpr { } fn canon(expr: &mut QueryExpr) { + // A `Concat` asserting a caller-proven `discriminator_unique_key` (issue + // #228) had that key's `ColumnId`s resolved, in `resolve.rs`, against + // exactly the first branch's output schema *as it stood before this + // pass ran*. `try_promote_heavy_hitter`/`try_rewrite_rownumber_topk` + // below can restructure that branch (anywhere within it — not only at + // its own top level, since the same recursive walk can rewrite a node + // nested under a pass-through wrapper too) into a shape with a + // different output schema, which would leave those `ColumnId`s + // pointing at the wrong column, or out of bounds, of the + // post-canonicalize schema. Snapshot the schema the discriminator key + // was actually resolved against, right here, before recursing into the + // children — this is the exact tree state `resolve.rs` saw. + let discriminator_branch_schema_before = match expr { + QueryExpr::Concat { + children, + discriminator_unique_key: Some(_), + } => children.first().and_then(|c| c.output_schema().ok()), + _ => None, + }; + // Bottom-up: canonicalize every child before matching at this node, so an // inner heavy-hitter is promoted before an enclosing rewrite inspects it. for child in children_mut(expr) { canon(child); } + + // 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: the two rewrites above don't preserve + // column identity in a way that's safe to infer). A wrong `unique_keys` + // claim is a wrong query answer, not a missed optimization — see + // `ConcatDiscriminatorKey`'s soundness doc — so this errs conservatively: + // any difference at all (not just a column-count/type change) drops the + // key, including the schema becoming undecidable in either direction. + if let QueryExpr::Concat { + children, + discriminator_unique_key: key @ Some(_), + } = expr + { + let discriminator_branch_schema_after = + children.first().and_then(|c| c.output_schema().ok()); + if discriminator_branch_schema_before != discriminator_branch_schema_after { + *key = None; + } + } + // Local rewrites chain: a `ROW_NUMBER()`-partitioned top-k rewrites to a // `Limit{Sort}`, which the heavy-hitter rule may then promote to an // `Aggregate([TopK])`. Each rule strictly simplifies the node, so applying @@ -100,7 +141,7 @@ fn children_mut(expr: &mut QueryExpr) -> Vec<&mut QueryExpr> { | PromqlInfoEnrich { child, .. } | Sort { child, .. } | Limit { child, .. } => vec![rc_mut(child)], - Concat { children } => children.iter_mut().collect(), + Concat { children, .. } => children.iter_mut().collect(), Join { left, right, .. } | SetOp { left, right, .. } => { vec![rc_mut(left), rc_mut(right)] } @@ -398,6 +439,71 @@ mod tests { assert_eq!(once, twice, "canonicalize must be idempotent"); } + // ── Concat's discriminator_unique_key vs. canonicalize (issue #228 review) ── + // + // `resolve.rs` resolves `discriminator_unique_key`'s `ColumnId`s against + // the first branch's *pre-canonicalize* output schema. If canonicalize + // then restructures that branch (heavy-hitter promotion, the + // `ROW_NUMBER()` top-k rewrite), those `ColumnId`s can end up pointing at + // the wrong column — or out of bounds — of the new schema. The two tests + // below pin the fix: the key is dropped whenever the branch's schema + // actually changed, and survives untouched otherwise. Never guessed at. + + #[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 = QueryExpr::concat_with_discriminator( + vec![count_by_service(), count_by_service()], + /* discriminator */ 0, + /* inner_key */ vec![1], + ); + let QueryExpr::Concat { + discriminator_unique_key, + .. + } = canonicalize(q) + else { + panic!("expected Concat"); + }; + assert!( + discriminator_unique_key.is_some(), + "an untouched first branch's discriminator key must survive canonicalize" + ); + } + + #[test] + fn concat_discriminator_key_is_dropped_when_first_branch_gets_rewritten() { + // The first branch is exactly the heavy-hitter promotion trigger — + // `Limit{Sort{Aggregate([Count])}}`, with an empty (global) + // `partition_by` — so canonicalize rewrites it in place to + // `Aggregate{TopK}`, whose own output is a single column, not the + // original two (`[service, count]`). A discriminator key resolved + // against the original 2-column shape (`discriminator` = `service` + // at index 0, `inner_key` = `count` at index 1) must not silently + // survive pointing at the new 1-column schema. + let promotable_branch = limit(5, 0, sort(desc(1), count_by_service())); + let q = QueryExpr::concat_with_discriminator( + vec![promotable_branch, count_by_service()], + /* discriminator */ 0, + /* inner_key */ vec![1], + ); + let QueryExpr::Concat { + children, + discriminator_unique_key, + } = canonicalize(q) + 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" + ); + } + #[test] fn does_not_promote_ascending_sort() { // Ascending = bottom-k: the shared `is_frequency_heavy_hitter` rule diff --git a/crates/types/src/pre_asap/cse.rs b/crates/types/src/pre_asap/cse.rs index 593cfea..9c689fb 100644 --- a/crates/types/src/pre_asap/cse.rs +++ b/crates/types/src/pre_asap/cse.rs @@ -301,8 +301,11 @@ pub fn structural_hash(node: &QueryExpr, cache: &mut HashCache) -> u64 { hash_own_fields(&mut hasher, &("Dedup", cols)); child_hash(child, cache).hash(&mut hasher); } - Concat { children } => { - "Concat".hash(&mut hasher); + Concat { + children, + discriminator_unique_key, + } => { + hash_own_fields(&mut hasher, &("Concat", discriminator_unique_key)); for c in children { // Stored by value, not `Rc` — see `rebuild_children`'s // `intern_owned` use for this variant — so there's no @@ -505,7 +508,7 @@ fn count_unique(node: &QueryExpr, seen: &mut std::collections::HashSet<*const Qu // this variant), so a branch has no `Rc` identity of its own to // dedup on at this position; still recurse into each in case an // `Rc`-shared descendant appears further down. - Concat { children } => children.iter().map(|c| count_unique(c, seen)).sum(), + Concat { children, .. } => children.iter().map(|c| count_unique(c, seen)).sum(), Join { left, right, .. } | SetOp { left, right, .. } => { visit(left, seen) + visit(right, seen) } @@ -620,11 +623,15 @@ fn rebuild_children(table: &mut InternTable, expr: QueryExpr) -> QueryExpr { cols, child: intern_child(table, child), }, - Concat { children } => Concat { + Concat { + children, + discriminator_unique_key, + } => Concat { children: children .into_iter() .map(|c| intern_owned(table, c)) .collect(), + discriminator_unique_key, }, Join { kind, diff --git a/crates/types/src/pre_asap/query_expr.rs b/crates/types/src/pre_asap/query_expr.rs index 9f346f7..819898e 100644 --- a/crates/types/src/pre_asap/query_expr.rs +++ b/crates/types/src/pre_asap/query_expr.rs @@ -561,6 +561,80 @@ impl Reduction { } } +/// A caller-proven compound unique key for a [`QueryExpr::Concat`] (issue +/// #228) — built only via [`QueryExpr::concat_with_discriminator`] / +/// [`ConcatDiscriminatorKey::new`], never by naming `discriminator` directly +/// in a struct literal (both fields are private): from *other Rust code*, +/// the only way to end up with one of these is to hand over a specific +/// column as the discriminator, by name, at the call site. +/// +/// Caveat: this is a Rust-API-level guarantee, not a data-level one. The +/// derived `Deserialize` impl below is same-module generated code, so it +/// builds a `ConcatDiscriminatorKey` directly from field values found in +/// arbitrary input, bypassing `new()` and its "name the discriminator" +/// requirement entirely — untrusted JSON can put in place any `discriminator` +/// / `inner_key` an attacker likes. This is not a reachable concern today: +/// the only place a whole `QueryExpr` is ever deserialized in this repo is a +/// same-file unit test, not an external input path. If `QueryExpr` (or a +/// subtree of it) ever gains a real from-external-input deserialization call +/// site, this guarantee would need revisiting there (a custom `Deserialize` +/// impl, or a post-deserialize validation pass) — nothing here does that yet. +/// +/// # Soundness +/// +/// `Concat`'s default (see its own doc) is to drop `unique_keys` +/// unconditionally, because a key unique **within** one branch is not unique +/// **across** the concatenation unless the branches' value sets for that key +/// are provably disjoint — nothing about matching schemas or matching +/// per-branch keys establishes that on its own. Two different branches can +/// trivially emit the same `inner_key` value (e.g. two PromQL +/// `histogram_quantiles` branches keyed on `(host, le)` can both produce a +/// `(host, le)` pair for different φ). +/// +/// Prepending `discriminator` is what restores it: if `discriminator`'s +/// value is **guaranteed to differ per branch** — a literal the producer +/// just tagged the branch with (PromQL φ riding along via +/// [`QueryExpr::PromqlRelabel`], a Postgres-style synthetic `GROUPING()` id +/// for `ROLLUP`/`CUBE`, …), never something inferred structurally from the +/// branches' own data — then `discriminator` alone partitions rows into +/// disjoint sets independent of what the branches actually contain, so +/// `(discriminator, inner_key)` is sound regardless of whether `inner_key` +/// values repeat across branches. +/// +/// This is a **caller-proven claim, not something `Concat` can verify**: +/// nothing stops a caller from asserting a discriminator that in fact +/// repeats across branches, in which case the resulting `unique_keys` claim +/// is simply wrong — `output_schema` trusts it without checking. The +/// obligation is on the constructor call site, exactly as it is on +/// [`QueryExpr::Dedup`]'s `cols` or any other unverified `unique_keys` +/// producer in this module. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(bound(serialize = "C: ColState", deserialize = "C: ColState"))] +pub struct ConcatDiscriminatorKey { + discriminator: C, + inner_key: Vec, +} + +impl ConcatDiscriminatorKey { + /// The only constructor — `discriminator` must be named explicitly by + /// the caller. See the type's doc for the soundness obligation this + /// puts on that caller. + pub fn new(discriminator: C, inner_key: Vec) -> Self { + Self { + discriminator, + inner_key, + } + } + + pub fn discriminator(&self) -> &C { + &self.discriminator + } + + pub fn inner_key(&self) -> &[C] { + &self.inner_key + } +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(bound(serialize = "C: ColState", deserialize = "C: ColState"))] pub enum QueryExpr { @@ -722,11 +796,25 @@ pub enum QueryExpr { /// project the branches into a common shape first. /// /// A row may appear in several branches, so no branch's unique key survives - /// the union — `unique_keys` is dropped, as in `SetOp`. + /// the union — `unique_keys` is dropped, as in `SetOp`. **Unless** the + /// constructor asserted `discriminator_unique_key` (issue #228, + /// [`QueryExpr::concat_with_discriminator`]): a caller-proven claim that + /// one column's value is guaranteed distinct per branch, which makes + /// `(discriminator, inner_key)` a sound compound unique key regardless of + /// whether `inner_key` alone repeats across branches. `None` — every + /// ordinary construction path, including the plain struct literal and + /// [`QueryExpr::concat`] — reproduces the old, unconditional-drop + /// behavior exactly; see [`ConcatDiscriminatorKey`]'s doc for the + /// soundness argument and the obligation this puts on whoever asserts it. /// /// Empty children is an error ([`QueryExprError::EmptyConcat`]), not an /// empty relation: there would be no schema to derive. - Concat { children: Vec> }, + Concat { + children: Vec>, + /// See the field-level doc above and [`ConcatDiscriminatorKey`]. + #[serde(default)] + discriminator_unique_key: Option>, + }, /// Logical join. Post-ASAP binding picks the physical alternative. Join { @@ -910,6 +998,37 @@ impl QueryExpr { QueryExpr::PromqlScalarBridge(Rc::new(QueryExpr::Literal(ScalarValue::Float64(v)))) } + /// Build an ordinary [`Concat`](Self::Concat) — the ordinary/default + /// construction path every call site should prefer over the bare struct + /// literal: `output_schema` drops `unique_keys` unconditionally, exactly + /// as before issue #228. Use + /// [`concat_with_discriminator`](Self::concat_with_discriminator) instead + /// when the caller can prove branch disjointness via a discriminator + /// column. + pub fn concat(children: Vec>) -> Self { + QueryExpr::Concat { + children, + discriminator_unique_key: None, + } + } + + /// Build a [`Concat`](Self::Concat) whose output schema carries the + /// caller-proven compound unique key `(discriminator, inner_key)` (issue + /// #228). See [`ConcatDiscriminatorKey`]'s doc for the soundness + /// argument and the obligation this puts on the caller — + /// `output_schema` trusts this claim without verifying it: nothing here + /// checks that `discriminator`'s value is actually distinct per branch. + pub fn concat_with_discriminator( + children: Vec>, + discriminator: C, + inner_key: Vec, + ) -> Self { + QueryExpr::Concat { + children, + discriminator_unique_key: Some(ConcatDiscriminatorKey::new(discriminator, inner_key)), + } + } + /// The value of a [`PromqlScalarBridge`](Self::PromqlScalarBridge) leaf /// wrapping a plain `Literal(ScalarValue::Float64(_))` — every one a /// front end constructs today (see [`promql_scalar`](Self::promql_scalar)). @@ -1134,13 +1253,26 @@ impl QueryExpr { // ⊕ — the branches are union-compatible by construction, so the // output shape is the first child's. A row can appear in more than // one branch, so no key of one branch is a key of the union: drop - // unique_keys, exactly as `SetOp` does. - QueryExpr::Concat { children } => { + // unique_keys, exactly as `SetOp` does — unless the constructor + // asserted `discriminator_unique_key` (issue #228), in which case + // `(discriminator, inner_key)` becomes the sole unique key. That + // assertion is trusted verbatim here, never checked: see + // `ConcatDiscriminatorKey`'s doc for the soundness argument and + // whose obligation it is. + QueryExpr::Concat { + children, + discriminator_unique_key, + } => { let mut s = children .first() .ok_or(QueryExprError::EmptyConcat) .and_then(|c| c.output_schema())?; s.unique_keys.clear(); + if let Some(key) = discriminator_unique_key { + let mut compound = vec![*key.discriminator()]; + compound.extend(key.inner_key().iter().copied()); + s.add_unique_key(compound); + } Ok(s) } // Set operations are union-compatible: both sides share the left's @@ -1716,9 +1848,7 @@ mod tests { "a Dedup branch does have a unique key on its own" ); - let merged = QueryExpr::Concat { - children: vec![branch(), branch()], - }; + let merged = QueryExpr::concat(vec![branch(), branch()]); let schema = merged.output_schema().unwrap(); assert!( schema.unique_keys.is_empty(), @@ -1735,9 +1865,7 @@ mod tests { cols: vec![0], child: Rc::new(scan(vec![col("k", DataType::Utf8, false)], None, vec![])), }; - let merged = QueryExpr::Concat { - children: vec![branch(), branch()], - }; + let merged = QueryExpr::concat(vec![branch(), branch()]); let setop = QueryExpr::SetOp { kind: SetOpKind::Union, all: true, @@ -1753,11 +1881,106 @@ mod tests { #[test] fn an_empty_merge_has_no_schema() { assert!(matches!( - QueryExpr::Concat { children: vec![] }.output_schema(), + QueryExpr::concat(vec![]).output_schema(), Err(QueryExprError::EmptyConcat) )); } + /// Issue #228: a `Concat` built via `concat_with_discriminator` gets a + /// sound compound `(discriminator, inner_key)` unique key, even though + /// each branch's own `inner_key` alone repeats across branches (exactly + /// the shape `merge_drops_the_branches_unique_keys` shows is unsafe + /// *without* a discriminator). + #[test] + fn discriminator_override_produces_a_compound_unique_key() { + // Two branches, each individually deduplicated on column 0 (`k`) — + // but, per `merge_drops_the_branches_unique_keys`, that alone proves + // nothing about the union. Column 1 (`branch_id`) stands in for a + // discriminator the constructor has separately proven distinct per + // branch (PromQL φ, a synthetic `GROUPING()` id, ...) — this + // schema-level test only checks the shape `output_schema` derives + // from asserting one, not how a real caller proves distinctness. + let branch = || QueryExpr::Dedup { + cols: vec![0], + child: Rc::new(scan( + vec![ + col("k", DataType::Utf8, false), + col("branch_id", DataType::Int64, false), + ], + None, + vec![], + )), + }; + let merged = QueryExpr::concat_with_discriminator( + vec![branch(), branch()], + /* discriminator */ 1, + /* inner_key */ vec![0], + ); + let schema = merged.output_schema().unwrap(); + assert_eq!( + schema.unique_keys, + vec![vec![1, 0]], + "(discriminator, inner_key) is the sole asserted unique key" + ); + assert_eq!( + schema.columns.len(), + 2, + "column shape is still the first branch's" + ); + } + + /// The override is opt-in: building a `Concat` without asserting a + /// discriminator — via the plain struct literal, exactly like every call + /// site before issue #228 — still drops `unique_keys` by default, + /// unchanged. + #[test] + fn ordinary_concat_struct_literal_still_drops_unique_keys_by_default() { + let branch = || QueryExpr::Dedup { + cols: vec![0], + child: Rc::new(scan(vec![col("k", DataType::Utf8, false)], None, vec![])), + }; + let merged = QueryExpr::Concat { + children: vec![branch(), branch()], + discriminator_unique_key: None, + }; + assert!(merged.output_schema().unwrap().unique_keys.is_empty()); + } + + /// Misuse check (issue #228): there is no way to end up with a + /// discriminator-backed unique key without a call site literally naming + /// a column as the discriminator. Neither the ordinary `concat` + /// constructor nor a bare struct literal with `discriminator_unique_key: + /// None` can be coaxed into fabricating one — the only path that + /// produces `Some` is `concat_with_discriminator` / + /// `ConcatDiscriminatorKey::new`, both of which require `discriminator` + /// as an explicit, named argument. + #[test] + fn no_way_to_fabricate_a_unique_key_without_naming_a_discriminator() { + let branch = || QueryExpr::Dedup { + cols: vec![0], + child: Rc::new(scan(vec![col("k", DataType::Utf8, false)], None, vec![])), + }; + // The ordinary builder. + assert_eq!( + QueryExpr::concat(vec![branch(), branch()]) + .output_schema() + .unwrap() + .unique_keys, + Vec::>::new() + ); + // The bare struct literal, explicitly opting out. + assert_eq!( + QueryExpr::Concat { + children: vec![branch(), branch()], + discriminator_unique_key: None, + } + .output_schema() + .unwrap() + .unique_keys, + Vec::>::new() + ); + } + #[test] fn project_retypes_and_renames_per_item() { let child = scan( diff --git a/crates/types/src/pre_asap/resolve.rs b/crates/types/src/pre_asap/resolve.rs index 4b8e021..4d80fb4 100644 --- a/crates/types/src/pre_asap/resolve.rs +++ b/crates/types/src/pre_asap/resolve.rs @@ -59,8 +59,8 @@ use super::column_resolution::{ }; use super::expr_ir::ColumnRef; use super::query_expr::{ - aggregate_output_schema, GroupKeys, Predicate, ProjectItem, QueryExprError, Reduction, - ResolvedQueryExpr, SortKey, UnresolvedQueryExpr, + aggregate_output_schema, ConcatDiscriminatorKey, GroupKeys, Predicate, ProjectItem, + QueryExprError, Reduction, ResolvedQueryExpr, SortKey, UnresolvedQueryExpr, }; use super::schema::{ColumnId, Schema}; @@ -242,12 +242,40 @@ fn resolve( } } - QE::Concat { children } => QE::Concat { - children: children + QE::Concat { + children, + discriminator_unique_key, + } => { + let children: Vec<_> = children .iter() .map(|c| resolve(c, fallback)) - .collect::, _>>()?, - }, + .collect::, _>>()?; + // No front end asserts this today (issue #228 shipped the + // extension point ahead of a wired call site) — resolved here + // regardless, against the first resolved branch's own output + // schema, exactly the schema `output_schema`'s `Concat` arm + // derives the merged schema from, so a future direct + // `concat_with_discriminator` caller upstream of `resolve_root` + // gets a correctly positional `ConcatDiscriminatorKey` out the + // other side. + let discriminator_unique_key = discriminator_unique_key + .as_ref() + .map(|key| -> Result<_, ResolveTreeError> { + let schema = children + .first() + .ok_or(QueryExprError::EmptyConcat)? + .output_schema()?; + Ok(ConcatDiscriminatorKey::new( + resolve_column_ref(key.discriminator(), &schema)?, + resolve_column_refs(key.inner_key(), &schema)?, + )) + }) + .transpose()?; + QE::Concat { + children, + discriminator_unique_key, + } + } QE::Join { kind, @@ -621,4 +649,45 @@ mod tests { lhs.output_schema().unwrap() ); } + + /// Issue #228 review, end-to-end: `resolve_root` over a `Concat` whose + /// discriminator column is referenced *nowhere else* in the tree, with a + /// schema-less (usage-derived) leaf `Scan` in the first branch — exactly + /// the scenario the review flagged. Before the `binder.rs` fix, the + /// Binder's fallback schema wouldn't contain `phi` at all, and this + /// `resolve_column_ref` call would fail `NotFound` for a column the + /// caller correctly named. It must resolve cleanly, and the resolved + /// `ConcatDiscriminatorKey` must carry the *positional* `ColumnId`s of + /// the branch's own (usage-derived) schema. + #[test] + fn resolve_root_seeds_and_resolves_an_otherwise_unreferenced_discriminator_column() { + let branch = || UnresolvedQueryExpr::Scan { + source: Source::TimeSeries { metric: "m".into() }, + predicates: vec![], + schema: None, + }; + let unresolved = UnresolvedQueryExpr::concat_with_discriminator( + vec![branch(), branch()], + ColumnRef::Named("phi".into()), + vec![ColumnRef::Named("host".into())], + ); + + let resolved = resolve_root(&unresolved).expect("resolves"); + let QueryExpr::Concat { + children, + discriminator_unique_key, + } = &resolved + else { + panic!("expected a resolved Concat, got {resolved:?}"); + }; + let schema = children[0].output_schema().unwrap(); + let key = discriminator_unique_key + .as_ref() + .expect("discriminator key survives resolution"); + assert_eq!(*key.discriminator(), schema.column_id("phi").unwrap()); + assert_eq!( + key.inner_key().to_vec(), + vec![schema.column_id("host").unwrap()] + ); + } } diff --git a/docs/design_docs/concat-unique-keys-decision.md b/docs/design_docs/concat-unique-keys-decision.md new file mode 100644 index 0000000..6b790c1 --- /dev/null +++ b/docs/design_docs/concat-unique-keys-decision.md @@ -0,0 +1,302 @@ +# `Concat` and `unique_keys`: the discriminator override (issue #228) + +> **Update**: the investigation below found no current call site paying for a +> redundant `Dedup` that this override would remove, and the original version +> of this document recommended deferring Option 1 on that basis. The repo +> owner reviewed that finding and explicitly asked for Option 1 to be built +> anyway — a deliberate "ship the extension point ahead of a proven call-site +> win" call, not a disagreement with the investigation. See "Decision" below +> for what actually shipped and the safety argument for why it's safe to ship +> unused. The investigation section is otherwise unchanged from the original +> write-up, since nothing about it stopped being true. + +## Context + +[`QueryExpr::Concat`](../../crates/types/src/pre_asap/query_expr.rs) (the +n-ary exact `UNION ALL` node, renamed from `Merge` in #226) always drops +`unique_keys` on its output — `merge_drops_the_branches_unique_keys` and +`merge_and_setop_agree_on_unique_keys` pin this down. Issue #228 asks whether +a `Concat`'s *constructor* should be able to assert a compound unique key +(`(discriminator_col, inner_key)`) when it can prove branch disjointness via +an explicit discriminator column (e.g. PromQL `histogram_quantiles`'s +per-branch φ, or a Postgres-style `GROUPING()` id for `ROLLUP`/`CUBE`/ +`GROUPING SETS`), and lists two options: + +1. **Producer-supplied override**: let the code building a `Concat` assert + its own `unique_keys` when it can prove branch disjointness. +2. **Don't preserve it in the IR at all**: a consumer that needs the + guarantee re-establishes it with an explicit `Dedup`. + +Per the issue's own instructions, the decisive question is empirical: does +any current call site actually pay for a redundant `Dedup`-equivalent today +that a discriminator-based unique key would let it drop? + +## What was checked + +Both current `Concat`-constructing call sites, and every consumer of +`Schema::unique_keys` in the tree: + +- **PromQL `histogram_quantiles`** — + [`walk_histogram_quantiles`](../../crates/frontend-promql/src/promql.rs). + Each branch is `PromqlRelabel { dst: label, value: Literal(φᵢ), child: Aggregate{…} }` + — the discriminator (φ, formatted the way `open_metrics_float` renders it) + really is a distinct literal per branch, so the "prove disjointness via a + tagged discriminator" premise the issue describes does hold structurally + here. The function returns `Unresolved::Concat { children: branches }` + directly — **no `Dedup` or dedup-equivalent node follows it**, in this + function or in any caller (`walk_call` returns its result unmodified). +- **SQL `ROLLUP`/`CUBE`/`GROUPING SETS`** — + [`lower_grouping_sets`](../../crates/frontend-sql/src/sql/mod.rs). Each + level's branch is `Project { …, child: Aggregate{…} }`, reinstating omitted + keys as typed `NULL`s, and the levels are `Concat`ed. The function returns + `Unresolved::Concat { children: branches }` directly — **no `Dedup` follows + it here either.** Notably, this lowering *already discards* DataFusion's + `__grouping_id` discriminator column on purpose (see the comment at the + bottom of `lower_grouping_sets`'s doc comment): "it only exists to tell a + subtotal's `NULL` apart from a data `NULL`, which is observable solely + through `GROUPING(col)` — an aggregate this front end rejects." So today + there is no discriminator column even available to ride along at this call + site; wiring one in would mean *first* deciding to stop rejecting + `GROUPING()` and surfacing `__grouping_id` as real IR — a separate, larger + change outside #228's scope, not a small addition to this lowering. +- **Every other `QueryExpr::Concat { … }` construction site** in the repo is + a test/tooling AST match (`promql_lowering.rs`, `promql_conformance.rs`, + `sql_lowering.rs`, `dag_export.rs`, `variant_coverage.rs`, netflow/synthetic + test fixtures) — none of them builds a fresh `Concat` with a `Dedup` on top + that this feature could remove. +- **Every consumer of `Schema::unique_keys`** in the tree, to check for a + cost beyond "a literal `Dedup` node": `pre_asap::cse::share_common_subtrees` + (gates CSE producer-sharing on `Schema::has_unique_key()`) and + `asap_aware_mapping::rollup::is_legal_rollup_source` (gates rollup-source + legality the same way, on an *`Aggregate`'s* own output schema). Neither + case is exercised by a `histogram_quantiles` or `ROLLUP`/`CUBE`/ + `GROUPING SETS` `Concat` in any current test, workload, or call site: no + test constructs a workload with two structurally-identical + `histogram_quantiles`/grouping-set queries for CSE to attempt to share, and + no call site aggregates further on top of a `Concat`'s output in a way that + would ask `is_legal_rollup_source` about it. +- No canonicalization/optimization pass in the repo removes a `Dedup` (or + anything else) on the strength of a child's `unique_keys` — there is no + such rewrite rule today, in `canonicalize.rs` or elsewhere — so even a + perfectly-preserved `unique_keys` on these `Concat` nodes would not, by + itself, delete any node from any plan that exists today. + +## Decision: build Option 1 anyway, unwired (explicit override of the defer) + +The investigation's conclusion stands: neither `histogram_quantiles` nor +`ROLLUP`/`CUBE`/`GROUPING SETS` lowering emits a `Dedup` (or anything playing +that role) after its `Concat` today, so there is nothing redundant in the +tree for a discriminator-based override to remove *right now*. On review, +the decision was made to build the extension point anyway, ahead of a proven +call-site win, rather than wait for one. That is a legitimate call to make +differently from the investigation's own recommendation — "no current +payoff" is a statement about today's call sites, not about whether the +shape is safe to add — and the rest of this section is the safety argument +for why it's fine to ship unused. + +### What shipped + +`QueryExpr::Concat` gained an opt-in field, `discriminator_unique_key: Option>` +(`crates/types/src/pre_asap/query_expr.rs`), plus: + +- `ConcatDiscriminatorKey` — a small struct with **private** `discriminator: C` / + `inner_key: Vec` fields, buildable only via `ConcatDiscriminatorKey::new(discriminator, inner_key)`. + Privacy is the enforcement mechanism for "the caller must explicitly name a + discriminator column" (see "Safety argument" below) — from other Rust code, + there is no path to a non-empty `unique_keys` claim that doesn't go through + a call site literally writing out which column it's proving is distinct. + Caveat: this is an API-level guarantee, not a data-level one — the derived + `Deserialize` impl bypasses `new()` entirely and can build one directly from + arbitrary field values in untrusted JSON. Not a reachable concern today (see + "Safety argument" below for why), but stated precisely rather than + overclaimed. +- `QueryExpr::concat(children)` — the ordinary constructor (`discriminator_unique_key: None`), + meant to replace the bare `QueryExpr::Concat { children }` struct literal + everywhere in the tree so a future field addition doesn't force every call + site to re-litigate this choice. +- `QueryExpr::concat_with_discriminator(children, discriminator, inner_key)` — + the override constructor. +- `output_schema()`'s `Concat` arm: unchanged default (`unique_keys` cleared + unconditionally) when `discriminator_unique_key` is `None`; when `Some`, + adds `(discriminator, inner_key)` as the sole unique key, trusting the + caller's claim without checking it. +- `resolve.rs`'s `Concat` arm resolves a pre-bind (`ColumnRef`) discriminator + key into its post-bind (`ColumnId`) equivalent against the first resolved + branch's own output schema — the same schema `output_schema()` derives the + merged shape from — so the feature works correctly end-to-end for a future + caller upstream of `resolve_root`, even though no such caller exists yet. +- Every other match/construction site touching `Concat` across the tree + (`canonicalize.rs`, `cse.rs`, `binder.rs`, `dag_export.rs`, + `asap-aware-mapping`'s `replacement.rs`/`explanation.rs`, and every + test/tooling AST walker) was mechanically updated to bind or ignore the new + field — most just added `, ..`; the two places that *rebuild* a `Concat` + node (`cse.rs`'s `rebuild_children`, part of CSE interning) thread + `discriminator_unique_key` through unchanged rather than dropping it. + +### Not wired into any lowering call site (deliberately) + +Per the explicit instruction accompanying this decision, `histogram_quantiles` +and `lower_grouping_sets` were **not** changed to call +`concat_with_discriminator` — both still call the plain `concat(children)` +builder, byte-for-byte the same `output_schema()` behavior they had before +this issue. The investigation's own findings are exactly why: `histogram_quantiles` +does have a structurally-available discriminator (φ, via `PromqlRelabel`) but +nothing downstream needs the resulting unique key yet, and SQL's natural +discriminator (`__grouping_id`) is actively discarded today because this +front end rejects `GROUPING()` — wiring that one in is a separate, +larger change (reopening that rejection) outside this issue's scope. Both +remain noted as future work at their call sites (see the comments added +there) and are not attempted here. + +### Safety argument: why the default is unaffected + +Three independent things hold `discriminator_unique_key: None` as the +observable behavior for every caller that doesn't ask for the override: + +1. **Every real construction path defaults to `None`.** `QueryExpr::concat` + hardcodes it; every call site in the tree (including both real lowering + call sites) uses `concat`, not `concat_with_discriminator`, so nothing in + the current tree can produce `Some` at all. +2. **`output_schema()`'s branch on the field is additive.** The `None` arm is + textually the same clear-and-return the code already did — `s.unique_keys.clear(); ... Ok(s)` + — with the `Some` branch reached only when the field is populated. This is + exactly what `merge_drops_the_branches_unique_keys` and + `merge_and_setop_agree_on_unique_keys` assert, and both pass unchanged. +3. **`Serialize`/`Deserialize` back-compat.** `#[serde(default)]` on the field + means a pre-#228 serialized `Concat` (missing the field entirely) + deserializes to `None`, matching its old unconditional-drop behavior. No + test in the repo constructs raw JSON for a `Concat` node to check this + directly, but the field shape follows the same `#[serde(default)]` + convention already used elsewhere on this enum (e.g. `Aggregate.output_names`, + `Scan.predicates`) for exactly this reason. + +### Safety argument: why the discriminator must be caller-proven, not inferred + +`ConcatDiscriminatorKey`'s fields are private; the only constructor, +`ConcatDiscriminatorKey::new(discriminator, inner_key)`, takes `discriminator` +as a required, explicitly-named argument — there is no default, no inference +from the branches' schemas, and no way to derive one structurally (e.g. "the +first column all branches disagree on"). This mirrors the same +private-field-plus-smart-constructor shape `GroupKeys` already uses in this +file for its own `by`/`without` invariant. Concretely, this means: + +- `output_schema()` never has enough information to fabricate a discriminator + on its own — it can only read one that a constructor already supplied. +- Nothing prevents a caller from asserting a **wrong** discriminator (one + that isn't actually distinct per branch) — the type system enforces "you + named a column," not "you were right about it." That obligation is + documented on `ConcatDiscriminatorKey` itself and is the same shape of + unverified claim `QueryExpr::Dedup.cols` already carries elsewhere in this + module (nothing checks a `Dedup`'s `cols` are actually a real key of its + child either). +- The `no_way_to_fabricate_a_unique_key_without_naming_a_discriminator` test + (in `query_expr.rs`) checks the two "how would you accidentally get + `Some`?" shapes concretely: the ordinary `concat` builder, and a bare + struct literal with `discriminator_unique_key: None` — both still produce + `unique_keys: []`. + +**Caveat, stated precisely (review fix, see "Review fixes" below): this is a +guarantee against *other Rust code*, not against arbitrary data.** +`#[derive(Deserialize)]` on `ConcatDiscriminatorKey` generates same-module +code that builds one directly from whatever `discriminator`/`inner_key` +values are present in the input, bypassing `new()` and its +naming-requirement entirely. A `Concat` node deserialized from untrusted JSON +could therefore carry a `discriminator_unique_key` nobody's Rust call site +ever named. This is not a reachable concern in the current repo — the only +place a whole `QueryExpr` is deserialized at all is a same-file unit test — +so no runtime hardening (a custom `Deserialize` impl, a post-deserialize +validation pass) was added for it; if `QueryExpr` ever gains a real +external-input deserialization path, this is the guarantee that would need +revisiting there. + +### Tests added (`crates/types/src/pre_asap/query_expr.rs`) + +- `merge_drops_the_branches_unique_keys` / `merge_and_setop_agree_on_unique_keys` — + unchanged, still pass (default behavior untouched). +- `discriminator_override_produces_a_compound_unique_key` — `concat_with_discriminator` + on two branches individually deduplicated on the same column (the exact + "looks safe but isn't" shape `merge_drops_the_branches_unique_keys` warns + about) yields `unique_keys == [[discriminator, inner_key…]]`. +- `ordinary_concat_struct_literal_still_drops_unique_keys_by_default` — the + bare struct literal (`discriminator_unique_key: None`) still drops + `unique_keys`, confirming the field addition didn't change the literal + construction path's behavior. +- `no_way_to_fabricate_a_unique_key_without_naming_a_discriminator` — the + misuse check described above. + +## Review fixes + +A code review of the initial implementation found two real correctness gaps +in the untested `resolve()`/`canonicalize()` path, plus a documentation +accuracy issue. All three are fixed on the same PR: + +1. **`binder.rs`'s `collect_referenced_columns` didn't walk + `discriminator_unique_key`'s `ColumnRef`s.** This function seeds every + name a query references into the Binder's usage-derived fallback schema, + which a schema-less `Scan` leaf (PromQL) falls back to. The `Concat` arm + was updated with `, ..` only, unlike the analogous `Dedup.cols` case + (which *is* walked, `push_ref_name`-style). Concretely: a future + `concat_with_discriminator(branches, discriminator_col, inner_key)` call + over an open query, where the discriminator column isn't otherwise + referenced anywhere else in the tree, with a schema-less leaf `Scan` in + the first branch — the Binder's fallback schema wouldn't contain the + discriminator name, and `resolve.rs`'s later `resolve_column_ref` call + would fail `NotFound` for a column the caller correctly named. Fixed: + the `Concat` arm now pushes `key.discriminator()` and every + `key.inner_key()` column into the walk, mirroring `Dedup.cols` exactly. + +2. **Resolved `ColumnId`s in `discriminator_unique_key` could go stale after + `canonicalize()` runs.** `resolve_root_with_inherited` calls `resolve()` + first — which resolves the key's `ColumnRef`s into `ColumnId`s against + `children.first()`'s output schema *as it stood at that point* — then + `canonicalize()` runs afterward and can restructure that same first + branch: `try_promote_heavy_hitter` and `try_rewrite_rownumber_topk` both + replace a `Limit{Sort{Aggregate}}`/`Filter{...}` shape with a + differently-shaped `Aggregate`, anywhere within the branch (not only at + its own top level — the walk is recursive), potentially changing its + column count/order. `output_schema()` read the previously-resolved + `ColumnId`s with no consistency check, so a future branch matching one of + these rewrite triggers could silently produce a wrong `unique_keys` claim + — a wrong query answer, not a missed optimization (per `cse.rs`'s own + module doc). Fixed in `canon()` (`canonicalize.rs`): before recursing + into a `Concat`'s children, if `discriminator_unique_key` is `Some`, + snapshot `children.first()`'s output schema — exactly the schema + `resolve.rs` resolved the key against. After the children have been + canonicalized, re-derive that schema and compare by full equality + (`Schema` is `PartialEq`/`Eq`); any difference at all — not just a + column-count/type change, since a same-shaped-but-different schema is + just as unsafe to trust positionally — drops the key (`None`) rather + than risk keeping a `ColumnId` that now points at the wrong column or is + out of bounds. The key is never *re-derived* by guessing at name or + position: the two rewrites don't preserve column identity in a way + that's safe to infer, so dropping is the only sound outcome once the + schema has moved. Two new tests in `canonicalize.rs`'s test module cover + both outcomes: `concat_discriminator_key_survives_canonicalize_when_first_branch_is_unaffected` + (an untouched branch keeps its key) and + `concat_discriminator_key_is_dropped_when_first_branch_gets_rewritten` + (a branch matching the heavy-hitter promotion trigger — 2 columns + collapsing to 1 — drops its key, never keeps a wrong one). + +3. **Documentation overclaimed the privacy guarantee.** Both the doc comment + on `ConcatDiscriminatorKey` and this document's "Safety argument" section + said flatly "there is no path to a non-empty `unique_keys` claim that + doesn't go through a call site literally writing out which column it's + proving is distinct" — true for other *Rust code*, but not for + `#[derive(Deserialize)]`, which is same-module generated code that builds + the struct directly from field values in arbitrary JSON, bypassing + `new()` entirely. Currently unreachable (the only place a whole + `QueryExpr` is deserialized in the repo is a same-file unit test, not an + external-input path), but the claim was factually incomplete as stated. + Fixed by narrowing both to "from other Rust code" and adding the + `Deserialize` caveat explicitly, rather than adding speculative runtime + hardening for a currently-unreachable path (per the review's preferred + option). + +## Future work (explicitly out of scope here) + +- Wiring `concat_with_discriminator` into `histogram_quantiles` (discriminator + readily available; no current downstream consumer). +- Reopening SQL's rejection of `GROUPING()` so `lower_grouping_sets` has a + real discriminator (`__grouping_id`) to assert — a separate design decision. +- Any canonicalization rule or CSE/rollup scenario that would actually *read* + a `Concat`'s asserted `unique_keys` for the first time in the current tree.