feat(query_expr): Concat discriminator unique_keys override (#228)

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);
             }

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);
                     }

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));
         }

crates/frontend-promql/src/promql.rs

@@ -742,7 +742,13 @@ fn walk_histogram_quantiles(call: &Call) -> Result<Unresolved> {
             })
         })
         .collect::<Result<Vec<_>>>()?;
-    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`

crates/frontend-promql/tests/observability/awesome_prometheus_alerts.rs

@@ -83,7 +83,7 @@ fn intents(e: &QueryExpr) -> Vec<AggIntent> {
                 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)
             }

crates/frontend-promql/tests/promql_conformance.rs

@@ -92,7 +92,7 @@ fn collect(e: &QueryExpr, out: &mut Vec<AggIntent>) {
             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)
         }

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<Vec<String>> = 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 {

crates/frontend-sql/src/sql/mod.rs

@@ -865,7 +865,14 @@ impl<'a> SqlLowerer<'a> {
             })
             .collect::<Result<Vec<_>, 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<Unresolved, LoweringError> {

crates/frontend-sql/tests/data_quality_check/synthetic_packet_trace.rs

@@ -105,7 +105,7 @@ fn intents(e: &QueryExpr) -> Vec<AggIntent> {
                 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)
             }

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);
             }

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<QueryExpr> {
     fn find(qe: &QueryExpr) -> Option<&Vec<QueryExpr>> {
         match qe {
-            QueryExpr::Concat { children } => Some(children),
+            QueryExpr::Concat { children, .. } => Some(children),
             QueryExpr::Project { child, .. }
             | QueryExpr::Filter { child, .. }
             | QueryExpr::Sort { child, .. }

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<u32> = 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];

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<Stri
             | QE::PromqlSubquery { child, .. }
             | QE::TimeRange { child, .. }
             | QE::TimeShift { child, .. } => 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

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