Skip to content

feat: support corr with a general bivariate aggregate type - #421

Merged
Selvomega merged 5 commits into
mainfrom
feat/two-input-aggregate-corr
Sep 16, 2026
Merged

Selvomega merged 5 commits into
mainfrom
feat/two-input-aggregate-corr

Conversation

@zzylol

@zzylol zzylol commented Sep 15, 2026

Copy link
Copy Markdown
Contributor

Why

Closes #414. DataFusion recognizes corr, but ASAPPlanner rejects it because its built-in value reducers carry only one input column.

What

Add the general AggIntent::Bivariate { op, left, right } shape, with BivariateAggOp::Correlation as its first operation. SQL correlation accepts expressions in either argument and returns nullable Float64.

How

Resolve both inputs to column IDs and expose them through input_cols() for dependency and operand-state checks. Project both SQL arguments to preserve casts and qualified join columns. Keep correlation exact, exclude it from scalar rollups/sketch selection, and allow physical hash-aggregate costing with provider-supplied state size.

Before this PR

SELECT corr(x, y * 2) AS r FROM a fails with unsupported aggregate: corr.

After this PR

The query lowers to a projection feeding Bivariate { op: Correlation, left: 0, right: 1 }, with a nullable numeric output. Future paired aggregates can add an operation without adding another argument-carrying IR variant.

Evidence

The regression test first reproduced unsupported aggregate: corr, then passed after implementation. SQL tests also preserve distinct inputs for corr(a.x, b.x) across a join and compile the retained exact query to an executable DAG. Screenshots and performance measurements: not applicable.

Verification

  • cargo +1.98.0 test --workspace --locked: 1,141 passed, no failures or ignored tests.
  • cargo +1.98.0 clippy --workspace --all-targets --all-features --locked -- -D warnings: passed.
  • cargo +1.98.0 fmt --all -- --check and git diff --check: passed.
  • Tests and Clippy used an isolated target directory with dev/test debug info disabled; runtime numerical evaluation was not exercised.

New coverage:

  • corr_result_is_nullable_float: reproduces the original error and checks output name/type/nullability.
  • corr_materializes_both_arguments: retains expressions, casts, and constants in either input.
  • corr_preserves_qualified_join_inputs: binds same-named columns to different join sides.
  • corr_coexists_with_grouping_having_and_other_measures: retains both arguments alongside grouping, HAVING, sorting, and another reducer.
  • corr_repeated_input_and_serialization: preserves repeated argument positions and round-trips the query.
  • corr_rejects_unrepresented_forms: rejects modifiers, window usage, and invalid arity.
  • corr_survives_exact_plan_compilation: preserves the bivariate intent through exact fallback and DAG compilation.
  • Core tests cover serialization/exactness, both usage-derived dependencies, qualified resolution/missing inputs, both operand-state checks, exclusion from sketch selection, and physical hash-aggregate costing.

Architectural decisions

Use a typed bivariate aggregate plus an operation enum, rather than a correlation-specific argument shape or opaque extension JSON. This makes both dependencies available to core binding and lets later operations reuse the representation. Preserve existing single-input APIs; dependency walkers use input_cols().

Post-ASAP mapping and current scope

The current mapping is an exact fallback:

AggIntent::Bivariate { op: Correlation, left, right }
    -> Implementation::PassThrough
    -> SummaryExpr::KeepPreAsap

PassThrough preserves the original aggregate for exact execution by the engine; it does not skip correlation. This PR tracks both input dependencies, checks both operand states, permits physical HashAggregate costing with provider-supplied state size, and excludes correlation from single-column sketches and scalar-result rollups.

There is no dedicated correlation summary implementation yet. The following remain unimplemented:

  • A maintained, updatable, mergeable correlation accumulator.
  • A correlation readout from that accumulator.
  • Query-time composition via ExactOperation::Aggregate over independently realized child plans; correlation is not enabled in that reducer path.

The executable-DAG test covers the retained KeepPreAsap fallback. It does not demonstrate a maintained correlation state or a dedicated post-ASAP correlation execution operator.

Limitations and follow-up

Only correlation is implemented initially. DISTINCT, aggregate FILTER/ORDER BY, explicit null treatment, and OVER remain unsupported. No maintained correlation accumulator or approximation is added. Numerical execution and pairwise null handling remain the executing engine's responsibility; this PR validates planning, not runtime numerical results.

Human review — do not complete with an agent

  • The MVP boundary is correct.
  • New conceptual layers or public interfaces are necessary.
  • The before/after description matches the intended product behavior.
  • Human reviewer:
  • Decision and rationale:

@zzylol zzylol changed the title feat: support corr with a general binary aggregate type feat: support corr with a general bivariate aggregate type Sep 15, 2026
@zzylol
zzylol requested a review from Selvomega September 15, 2026 13:45

@Selvomega Selvomega left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

As described in the comment. I might directly use AggIntent::PearsonCorr
if you still think current design good, you can just merge it

Comment thread crates/asap-aware-mapping/src/query_physical_lowering.rs Outdated
@zzylol

zzylol commented Sep 15, 2026

Copy link
Copy Markdown
Contributor Author

As described in the comment. I might directly use AggIntent::PearsonCorr if you still think current design good, you can just merge it

I agree using AggIntent::PearsonCorr is better. Would you like to take the change?

Selvomega and others added 2 commits September 15, 2026 17:30
`AggIntent::Bivariate { op, left, right }` placed a category alongside
specific instances such as `Avg` and `Variance`, and `BivariateAggOp`
carried exactly one operation. Correlation is the only two-input
aggregate, so name the variant after it and drop the operation enum.
`covar` and the `regr_*` family can each add a variant when implemented,
following the `StdDev` / `Variance` precedent.

The variant now sits next to `Variance` rather than ahead of the
data-model-agnostic banner. Serialized `kind` becomes `pearson_corr` and
the `op` field is gone. `input_cols()` is unchanged: dependency walkers
still need a generic accessor, and `input_col()` still returns `None` so
single-column consumers cannot pick up half of the pair.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Two conflicts, both from main landing after this branch was cut.

`tests/sql_lowering.rs`: each side appended a test function at the same
place — `corr_result_is_nullable_float` here, `composite_distinct_is_rejected`
on main (#419). They are unrelated, so both are kept.

`tests/data_quality_check/tpch_deequ.rs`: no textual conflict, since this
branch never touched the file. #419 added the corpus ratchet and pinned
U-P4d and U-P4l as rejected-because-`corr`. Both lower now, so the ratchet
moves: 47 -> 49 lowered, P2b alone rejected.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@Selvomega
Selvomega merged commit 8b90f04 into main Sep 16, 2026
3 checks passed
Selvomega added a commit that referenced this pull request Sep 16, 2026
…UNT(DISTINCT)

Two changes, in dependency order. The second is only expressible because of
the first.

`AggIntent` asserted that an aggregate reads one column: `input_col() ->
Option<ColumnId>` was how every consumer reached an input. #421 added
`PearsonCorr { left, right }` and had to make it return `None` from that
accessor *defensively*, so no single-column consumer could pick up half of the
pair. The accessor had become a trap each new multi-column intent must
remember to opt out of, and the cost of forgetting is a silently wrong number
rather than a compile error.

That is also why `COUNT(DISTINCT a, b)` had nowhere to land. It counts
distinct *tuples*; SQL lowering dropped every argument after the first,
reporting single-column cardinality where the query asked for tuple
cardinality. #419 stopped the miscount by rejecting the call outright, which
left TPC-H/Deequ U-P2b and nine synthetic-packet-trace flow counts permanently
unlowerable.

Before this PR:

    SELECT count(DISTINCT l_orderkey, l_linenumber) * 1.0 / count(*)
    FROM lineitem
    -> lowering failed: unsupported aggregate: multi-column COUNT(DISTINCT)

After this PR:

    Aggregate { measures: [Cardinality { cols: [0, 1], .. }], .. }

What changed:

- `input_col()` is removed. `input_cols() -> Vec<C>` is the only column
  accessor, so no consumer can ask for "the" input column of an aggregate that
  reads two. An intent may now declare whatever arity its semantics need
  without a defensive opt-out, and `PearsonCorr`'s opt-out is gone. The
  genuinely unary reducers keep `col: Option<C>` — this relaxes what consumers
  assume, it does not widen every variant.
- `AggIntent::Cardinality` takes `cols: Vec<C>`. One entry is
  `COUNT(DISTINCT col)`, several count distinct tuples, empty keeps the PromQL
  "the sample value" convention. It is the first variable-arity intent, and was
  not expressible while the accessor split existed.
- Realization is the single-column one with a wider item: a tuple becomes a
  `SummaryInputExpr::Tuple`, which HLL/Theta/KMV hash as one value. UnivMon is
  withheld from a tuple — it estimates frequency moments over a single value
  stream — and the value-frequency rule refuses a multi-column intent directly,
  so that invariant does not rest on the candidate table alone.
- SQL lowering resolves every `COUNT(DISTINCT ...)` argument as a grouping key,
  so a qualifier survives a join; an expression argument is rejected rather
  than reduced over a probe column.

Serialization of `Cardinality` changes from `col` to `cols`; a payload that
omits the field still reads as the implicit input.

Both corpus ratchets move in this commit: tpch_deequ 49 -> 50 lowered with no
rejections, synthetic_packet_trace 61 -> 70.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
zzylol added a commit that referenced this pull request Sep 17, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature] ASAPPlanner cannot parse DataFusion corr function

2 participants