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

[zzylol]

Summary

Issue #228 asked a design question first: should QueryExpr::Concat be able
to assert a compound unique_keys override ((discriminator_col, inner_key)) when a caller can prove branch disjointness via a discriminator
column? The investigation (first commit on this PR) found that neither
current Concat-constructing call site — PromQL histogram_quantiles, SQL
ROLLUP/CUBE/GROUPING SETS lowering — pays for a redundant Dedup that
this would let it remove, and recommended deferring.

Update: the repo owner reviewed that finding and explicitly asked for
Option 1 (producer-supplied override) to be implemented anyway — a
deliberate "ship the extension point ahead of a proven call-site win"
decision, not a disagreement with the investigation. This PR now contains
that implementation, plus a round of review fixes (see "Review fixes"
below). Full writeup (investigation + reversal + safety argument + review
fixes): docs/design_docs/concat-unique-keys-decision.md.

What shipped

• QueryExpr::Concat gains an opt-in field:
  discriminator_unique_key: Option<ConcatDiscriminatorKey<C>>.
• ConcatDiscriminatorKey<C> — private discriminator/inner_key fields,
  buildable only via ConcatDiscriminatorKey::new(discriminator, inner_key).
  This is the enforcement mechanism for "the caller must explicitly name a
  discriminator column" — from other Rust code, there's 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 (see the
  Deserialize caveat under "Review fixes" below). Mirrors the existing
  GroupKeys private-fields-plus-smart-constructor shape already in this
  file.
• QueryExpr::concat(children) — the new ordinary constructor
  (discriminator_unique_key: None), now used everywhere in the tree in
  place of the bare QueryExpr::Concat { children } struct literal.
• QueryExpr::concat_with_discriminator(children, discriminator, inner_key)
  — the override constructor.
• output_schema()'s Concat arm: unchanged default when the field is
  None (unique_keys cleared unconditionally, exactly as before); 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 (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 is correct end-to-end for a future caller upstream
  of resolve_root.
• Every other Concat match/construction site 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) updated mechanically — most just bind/ignore the
  new field; the one place that rebuilds a Concat node (cse.rs's
  rebuild_children, part of CSE interning) threads
  discriminator_unique_key through unchanged rather than silently dropping
  it.

Not wired into any lowering call site (deliberately)

Per the explicit instruction accompanying this decision, histogram_quantiles
and SQL's 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 as before this
issue. Why: PromQL's φ discriminator is structurally available
(PromqlRelabel) but nothing downstream needs the resulting unique key yet;
SQL's natural discriminator (__grouping_id) is actively discarded today
because this front end rejects GROUPING() — wiring that one in means first
reopening that rejection, a separate, larger design decision. Both are noted
as future work at their call sites and in the design doc.

Safety argument (why the default is unaffected)

1. Every real construction path in the tree defaults to None —
   QueryExpr::concat hardcodes it, and no call site (including both real
   lowering call sites) uses concat_with_discriminator.
2. output_schema()'s None arm is textually the same
   clear-and-return the code already did. merge_drops_the_branches_unique_keys
   and merge_and_setop_agree_on_unique_keys pass unchanged.
3. #[serde(default)] on the field means a pre-Concat drops unique_keys unconditionally — no way for a producer to assert cross-branch disjointness #228 serialized Concat
   (missing the field) deserializes to None — same convention already used
   elsewhere on this enum (Aggregate.output_names, Scan.predicates).

Safety argument (why the discriminator must be caller-proven, not inferred)

ConcatDiscriminatorKey's private fields plus its single, explicit-argument
constructor mean output_schema() never has enough information to fabricate
a discriminator on its own — it can only read one a constructor already
supplied. Nothing stops a caller from asserting a wrong discriminator (the
type system enforces "you named a column," not "you were right about it") —
that obligation is documented on the type and is the same shape of
unverified claim QueryExpr::Dedup.cols already carries elsewhere in this
module.

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 this branch:

1. binder.rs's collect_referenced_columns didn't walk
   discriminator_unique_key's ColumnRefs. 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
   had only been updated with , .., unlike the analogous Dedup.cols case
   (which is walked). Concretely: a future concat_with_discriminator
   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.
   New 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
   (full resolve_root end-to-end).

2. Resolved ColumnIds in discriminator_unique_key could go stale after
   canonicalize() runs. resolve() resolves the key's ColumnRefs
   against children.first()'s output schema as it stood at that point;
   canonicalize() then runs afterward and can restructure that same
   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) — with no consistency check downstream, so a matching
   branch could silently produce a wrong unique_keys claim: a wrong query
   answer, not a missed optimization. Fixed in canon()
   (canonicalize.rs): snapshot the first branch's output schema before
   recursing into a Concat's children — exactly the schema resolve.rs
   resolved the key against — and afterward, drop the key (None) if that
   schema differs at all (full Schema equality, not just a
   column-count/type heuristic). The key is never re-derived by guessing at
   name or position; dropping is the only sound outcome once the schema has
   moved. Two new tests in canonicalize.rs cover both outcomes:
   concat_discriminator_key_survives_canonicalize_when_first_branch_is_unaffected
   and concat_discriminator_key_is_dropped_when_first_branch_gets_rewritten.

3. Documentation overclaimed the privacy guarantee. Both the doc comment
   on ConcatDiscriminatorKey and this PR's original "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 naming the discriminator"
   — 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 stating the Deserialize caveat explicitly, rather
   than adding speculative runtime hardening for a currently-unreachable
   path (per the review's own preferred option).

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 — 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 — the
  bare struct literal with discriminator_unique_key: None still drops
  unique_keys.
• no_way_to_fabricate_a_unique_key_without_naming_a_discriminator — misuse
  check: neither the ordinary concat builder nor an explicit None
  literal can produce a unique key without a call site literally naming a
  discriminator column via ConcatDiscriminatorKey::new.

Review-fix regression tests:

• crates/types/src/pre_asap/binder.rs:
  concat_discriminator_key_is_seeded_into_the_binder_schema.
• crates/types/src/pre_asap/resolve.rs:
  resolve_root_seeds_and_resolves_an_otherwise_unreferenced_discriminator_column.
• crates/types/src/pre_asap/canonicalize.rs:
  concat_discriminator_key_survives_canonicalize_when_first_branch_is_unaffected,
  concat_discriminator_key_is_dropped_when_first_branch_gets_rewritten.

Testing

• cargo build --workspace — clean, no warnings.
• cargo test --workspace — all green: 115 passed in asap-types's unit
  suite (108 before this PR, +7 new tests across both commits), 0 failures
  anywhere in the workspace.
• cargo clippy --workspace --all-targets -- -D warnings — clean.
• cargo fmt --all -- --check — clean.

Open questions for review

• Is the private-fields-plus-explicit-constructor shape on
  ConcatDiscriminatorKey sufficient API-level protection given the
  Deserialize caveat, or is that worth closing now even though nothing
  reaches it yet?
• Should histogram_quantiles actually get wired up to
  concat_with_discriminator in a follow-up, given the discriminator is
  already sitting right there? (Left as future work since nothing downstream
  reads it yet.)
• Same question for SQL, contingent on revisiting the GROUPING()
  rejection.
• Is dropping (never re-deriving) the discriminator key on any
  post-canonicalize schema change the right call, or should a future change
  attempt a narrower re-derivation for the specific rewrite shapes involved?

Ref #228

🤖 Generated with Claude Code