From 2b5debcd102756cf6b8a51d924dc967cd260df74 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 26 Aug 2026 13:39:49 -0600 Subject: [PATCH 1/3] docs(query_expr): investigate #228, recommend deferring Concat unique_keys override MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Investigated whether any current Concat call site (PromQL histogram_quantiles, SQL ROLLUP/CUBE/GROUPING SETS lowering) pays for a redundant Dedup that a discriminator-based unique_keys override could prove unnecessary. Finding: neither call site emits a Dedup (or equivalent) after its Concat today, and no consumer of Schema::unique_keys (CSE's share_common_subtrees, asap_aware_mapping::rollup's is_legal_rollup_source) is exercised by either Concat in any current test or workload. SQL grouping-set lowering also actively discards the one natural discriminator (__grouping_id) today, since this front end rejects GROUPING(). Recommendation: defer Option 1 (producer-supplied unique_keys override) until a real consumer exists; building it now would be speculative surface area with no call site to justify or exercise it. No behavior change — Concat still drops unique_keys unconditionally, per merge_drops_the_branches_unique_keys / merge_and_setop_agree_on_unique_keys (both still passing, unchanged). Ref #228 Co-Authored-By: Claude Sonnet 5 --- .../concat-unique-keys-decision.md | 108 ++++++++++++++++++ 1 file changed, 108 insertions(+) create mode 100644 docs/design_docs/concat-unique-keys-decision.md 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 00000000..f8ef0d15 --- /dev/null +++ b/docs/design_docs/concat-unique-keys-decision.md @@ -0,0 +1,108 @@ +# `Concat` and `unique_keys`: defer the discriminator override (issue #228) + +## 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: defer (Option 2 for now — no code change) + +Neither `histogram_quantiles` nor `ROLLUP`/`CUBE`/`GROUPING SETS` lowering +emits a `Dedup`, or anything playing that role, after its `Concat` today. +There is nothing redundant in the tree for a discriminator-based override to +remove. Building the producer-supplied-override machinery (Option 1) now +would be pure speculative surface area — new trust-boundary code +(`output_schema()` accepting a caller-asserted key that must actually be +sound, forever) with no current call site to exercise or justify it, exactly +what the issue's own "before implementing either" section says to avoid. + +This defers, it does not foreclose. If a future change adds a real consumer +— e.g. a canonicalization rule that drops a provably-redundant `Dedup` based +on a child's `unique_keys`, or a workload-level CSE scenario that actually +puts two identical `histogram_quantiles`/grouping-set `Concat`s in front of +`share_common_subtrees` — Option 1 (producer-supplied override, additive, +`output_schema()`'s default unchanged for ordinary callers) is still the +right shape for it then, for the reasons #228 already gives: it's strictly +additive, and the alternative of changing `Concat`'s default is unsound (a +key unique within one branch is not unique across the union unless the +branches are provably disjoint, and nothing about matching schemas or +matching per-branch keys establishes that on its own). + +For SQL grouping sets specifically, note that the natural discriminator +(`__grouping_id`) is not just unwired but actively discarded on purpose today +because `GROUPING()` itself is rejected — so revisiting this decision for SQL +means revisiting that rejection first, which is its own design question, not +a corollary of #228. + +## No code change + +This document is the entire change for #228: an investigation write-up, with +no implementation. `merge_drops_the_branches_unique_keys` and +`merge_and_setop_agree_on_unique_keys` continue to describe `Concat`'s only +behavior — dropping `unique_keys` unconditionally — unchanged. From b1740ab346ffc40e150376f620f460ef67e250f4 Mon Sep 17 00:00:00 2001 From: zz_y Date: Wed, 26 Aug 2026 13:56:00 -0600 Subject: [PATCH 2/3] feat(query_expr): add Concat discriminator unique_keys override (#228) Per explicit user direction overriding the investigation's own defer recommendation: build Option 1 (producer-supplied unique_keys override) ahead of a proven call-site win, as a deliberate 'ship the extension point now' decision. - QueryExpr::Concat gains an opt-in discriminator_unique_key: Option< ConcatDiscriminatorKey> field. ConcatDiscriminatorKey's fields (discriminator, inner_key) are private; the only constructor, ConcatDiscriminatorKey::new, requires the discriminator column to be named explicitly by the caller -- nothing infers or defaults one. - QueryExpr::concat(children) is the new ordinary constructor (discriminator_unique_key: None), replacing the bare struct literal at every call site in the tree. - QueryExpr::concat_with_discriminator(children, discriminator, inner_key) is the override constructor. - output_schema()'s Concat arm: unchanged default (drop unique_keys unconditionally) when the field is None; when Some, asserts (discriminator, inner_key) as the sole unique key, trusting the caller's claim without verifying it. - resolve.rs resolves a pre-bind discriminator key into its post-bind ColumnId equivalent against the first resolved branch's schema, so the feature is correct end-to-end for a future caller upstream of resolve_root. - Every other Concat match/construction site across the tree (canonicalize, cse, binder, dag_export, asap-aware-mapping's replacement/explanation, and test/tooling AST walkers) updated mechanically; cse.rs's rebuild_children (CSE interning) threads the field through unchanged rather than dropping it. Not wired into any real lowering call site, per instruction: neither histogram_quantiles nor SQL ROLLUP/CUBE/GROUPING SETS lowering calls concat_with_discriminator. Both still call the plain concat() builder -- byte-for-byte the same output_schema() behavior as before. SQL's natural discriminator (__grouping_id) is still discarded because this front end rejects GROUPING(); PromQL's phi discriminator is structurally available but nothing downstream needs the resulting key yet. Both noted as future work. Tests (crates/types/src/pre_asap/query_expr.rs): - merge_drops_the_branches_unique_keys / merge_and_setop_agree_on_unique_keys: unchanged, still pass. - discriminator_override_produces_a_compound_unique_key: new, concat_with_discriminator on branches individually deduplicated on the same column yields unique_keys == [[discriminator, inner_key]]. - ordinary_concat_struct_literal_still_drops_unique_keys_by_default: new, the bare struct literal with discriminator_unique_key: None still drops unique_keys. - no_way_to_fabricate_a_unique_key_without_naming_a_discriminator: new misuse check -- neither the ordinary builder nor an explicit None literal can produce a unique key without a call site literally naming a discriminator column via ConcatDiscriminatorKey::new. docs/design_docs/concat-unique-keys-decision.md updated to record the investigation (unchanged), the reversal, and the safety argument for shipping the override unused. cargo build --workspace, cargo test --workspace, and cargo clippy --workspace --tests all clean. Ref #228 Co-Authored-By: Claude Sonnet 5 --- crates/asap-aware-mapping/src/explanation.rs | 2 +- crates/asap-aware-mapping/src/replacement.rs | 5 +- crates/devtools/src/bin/variant_coverage.rs | 2 +- crates/frontend-promql/src/promql.rs | 8 +- .../awesome_prometheus_alerts.rs | 2 +- .../tests/promql_conformance.rs | 2 +- .../frontend-promql/tests/promql_lowering.rs | 6 +- crates/frontend-sql/src/sql/mod.rs | 9 +- .../synthetic_packet_trace.rs | 2 +- crates/frontend-sql/tests/netflow/netflow.rs | 2 +- crates/frontend-sql/tests/sql_lowering.rs | 2 +- crates/types/src/dag_export.rs | 29 +-- crates/types/src/pre_asap/binder.rs | 4 +- crates/types/src/pre_asap/canonicalize.rs | 2 +- crates/types/src/pre_asap/cse.rs | 15 +- crates/types/src/pre_asap/query_expr.rs | 233 +++++++++++++++++- crates/types/src/pre_asap/resolve.rs | 40 ++- .../concat-unique-keys-decision.md | 179 +++++++++++--- 18 files changed, 453 insertions(+), 91 deletions(-) diff --git a/crates/asap-aware-mapping/src/explanation.rs b/crates/asap-aware-mapping/src/explanation.rs index bcfee70a..026df916 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 e18a9a50..2edf92e4 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 a1b3b823..0cd66ee2 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 690ecbfc..f43a00df 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 b616d75e..254297de 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 698a1bdb..89426ba8 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 192118d5..f59f133c 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 1a650b04..538df705 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 91e08e08..5831aafb 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 2cf0cb78..09d742b8 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 06beba53..0b99cd41 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 252d319a..6d6eda23 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 97061e99..e7b0a159 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,7 @@ 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, .. } => children.iter().for_each(|c| walk(c, out)), QE::SetOp { left, right, .. } => { walk(left, out); walk(right, out); diff --git a/crates/types/src/pre_asap/canonicalize.rs b/crates/types/src/pre_asap/canonicalize.rs index 6dda2ec0..b1ae3b71 100644 --- a/crates/types/src/pre_asap/canonicalize.rs +++ b/crates/types/src/pre_asap/canonicalize.rs @@ -100,7 +100,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)] } diff --git a/crates/types/src/pre_asap/cse.rs b/crates/types/src/pre_asap/cse.rs index 593cfea6..9c689fbb 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 9f346f71..60462dca 100644 --- a/crates/types/src/pre_asap/query_expr.rs +++ b/crates/types/src/pre_asap/query_expr.rs @@ -561,6 +561,68 @@ 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): 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. +/// +/// # 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 +784,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 +986,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 +1241,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 +1836,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 +1853,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 +1869,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 4b8e0211..225bac41 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, diff --git a/docs/design_docs/concat-unique-keys-decision.md b/docs/design_docs/concat-unique-keys-decision.md index f8ef0d15..37e2b444 100644 --- a/docs/design_docs/concat-unique-keys-decision.md +++ b/docs/design_docs/concat-unique-keys-decision.md @@ -1,4 +1,14 @@ -# `Concat` and `unique_keys`: defer the discriminator override (issue #228) +# `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 @@ -71,38 +81,135 @@ Both current `Concat`-constructing call sites, and every consumer of perfectly-preserved `unique_keys` on these `Concat` nodes would not, by itself, delete any node from any plan that exists today. -## Decision: defer (Option 2 for now — no code change) - -Neither `histogram_quantiles` nor `ROLLUP`/`CUBE`/`GROUPING SETS` lowering -emits a `Dedup`, or anything playing that role, after its `Concat` today. -There is nothing redundant in the tree for a discriminator-based override to -remove. Building the producer-supplied-override machinery (Option 1) now -would be pure speculative surface area — new trust-boundary code -(`output_schema()` accepting a caller-asserted key that must actually be -sound, forever) with no current call site to exercise or justify it, exactly -what the issue's own "before implementing either" section says to avoid. - -This defers, it does not foreclose. If a future change adds a real consumer -— e.g. a canonicalization rule that drops a provably-redundant `Dedup` based -on a child's `unique_keys`, or a workload-level CSE scenario that actually -puts two identical `histogram_quantiles`/grouping-set `Concat`s in front of -`share_common_subtrees` — Option 1 (producer-supplied override, additive, -`output_schema()`'s default unchanged for ordinary callers) is still the -right shape for it then, for the reasons #228 already gives: it's strictly -additive, and the alternative of changing `Concat`'s default is unsound (a -key unique within one branch is not unique across the union unless the -branches are provably disjoint, and nothing about matching schemas or -matching per-branch keys establishes that on its own). - -For SQL grouping sets specifically, note that the natural discriminator -(`__grouping_id`) is not just unwired but actively discarded on purpose today -because `GROUPING()` itself is rejected — so revisiting this decision for SQL -means revisiting that rejection first, which is its own design question, not -a corollary of #228. - -## No code change - -This document is the entire change for #228: an investigation write-up, with -no implementation. `merge_drops_the_branches_unique_keys` and -`merge_and_setop_agree_on_unique_keys` continue to describe `Concat`'s only -behavior — dropping `unique_keys` unconditionally — unchanged. +## 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) — 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. +- `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: []`. + +### 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. + +## 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. From acfd2fe22ee7abaa3ebc61922969b4cf4ef2adac Mon Sep 17 00:00:00 2001 From: zz_y Date: Thu, 27 Aug 2026 10:53:41 -0600 Subject: [PATCH 3/3] fix(query_expr): close two Concat discriminator correctness gaps + doc caveat (#228) Review fixes on the Concat discriminator_unique_key override: 1. binder.rs's collect_referenced_columns didn't walk discriminator_unique_key's ColumnRefs -- the Concat arm was updated with ', ..' only, unlike the analogous Dedup.cols case (which is walked). This function seeds every name a query references into the Binder's usage-derived fallback schema; a discriminator column not otherwise referenced anywhere else in the tree, over a schema-less leaf Scan, would be absent from that fallback schema, 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. New regression tests: binder.rs's concat_discriminator_key_is_seeded_into_the_binder_schema and resolve.rs's resolve_root_seeds_and_resolves_an_otherwise_unreferenced_discriminator_column (end-to-end through resolve_root). 2. Resolved ColumnIds in discriminator_unique_key could go stale after canonicalize() runs. resolve() resolves the key's ColumnRefs against children.first()'s pre-canonicalize output schema; canonicalize() then runs afterward and can restructure that same branch (try_promote_heavy_hitter, try_rewrite_rownumber_topk both replace a Limit{Sort{Aggregate}}/Filter{...} shape with a differently-shaped Aggregate, anywhere within the branch), changing its column count/order with no consistency check downstream. Fixed in canon() (canonicalize.rs): snapshot the first branch's output schema before recursing into a Concat's children (exactly what resolve.rs resolved against), and after recursing, drop the key (never re-derive it by guessing) if the schema differs at all -- erring conservatively, since a wrong unique_keys claim is a wrong query answer, not a missed optimization. Two new tests in canonicalize.rs pin both outcomes: the key survives an untouched branch, and is dropped when the branch matches the heavy-hitter promotion trigger. 3. Documentation accuracy: the doc comment on ConcatDiscriminatorKey and the design doc's safety argument overclaimed "there is no path to a non-empty unique_keys claim that doesn't go through a call site literally naming the discriminator" -- true for other Rust code, but #[derive(Deserialize)] is same-module generated code that builds the struct directly from arbitrary field values, bypassing new() entirely. Currently unreachable (the only whole-QueryExpr deserialization call site in the repo is a same-file unit test), but the claim was factually incomplete as stated. Fixed by scoping both claims to 'other Rust code' and stating the Deserialize caveat explicitly, rather than adding speculative runtime hardening for a currently-unreachable path. docs/design_docs/concat-unique-keys-decision.md updated with a 'Review fixes' section covering all three, and both doc-accuracy fixes applied inline. cargo build --workspace, cargo test --workspace (115 passed, 0 failed), cargo clippy --workspace --all-targets -- -D warnings, and cargo fmt --all -- --check all clean. Ref #228 Co-Authored-By: Claude Sonnet 5 --- crates/types/src/pre_asap/binder.rs | 42 ++++++- crates/types/src/pre_asap/canonicalize.rs | 106 ++++++++++++++++++ crates/types/src/pre_asap/query_expr.rs | 18 ++- crates/types/src/pre_asap/resolve.rs | 41 +++++++ .../concat-unique-keys-decision.md | 93 ++++++++++++++- 5 files changed, 293 insertions(+), 7 deletions(-) diff --git a/crates/types/src/pre_asap/binder.rs b/crates/types/src/pre_asap/binder.rs index e7b0a159..221aebdb 100644 --- a/crates/types/src/pre_asap/binder.rs +++ b/crates/types/src/pre_asap/binder.rs @@ -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 b1ae3b71..092b93bb 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 @@ -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/query_expr.rs b/crates/types/src/pre_asap/query_expr.rs index 60462dca..819898ef 100644 --- a/crates/types/src/pre_asap/query_expr.rs +++ b/crates/types/src/pre_asap/query_expr.rs @@ -564,9 +564,21 @@ 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): 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. +/// 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 /// diff --git a/crates/types/src/pre_asap/resolve.rs b/crates/types/src/pre_asap/resolve.rs index 225bac41..4d80fb44 100644 --- a/crates/types/src/pre_asap/resolve.rs +++ b/crates/types/src/pre_asap/resolve.rs @@ -649,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 index 37e2b444..6b790c1e 100644 --- a/docs/design_docs/concat-unique-keys-decision.md +++ b/docs/design_docs/concat-unique-keys-decision.md @@ -102,9 +102,14 @@ for why it's fine to ship unused. - `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) — 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. + 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 @@ -190,6 +195,20 @@ file for its own `by`/`without` invariant. Concretely, this means: 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` — @@ -205,6 +224,74 @@ file for its own `by`/`without` invariant. Concretely, this means: - `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