fix: support COUNT(DISTINCT col1, col2, ...) multi-column distinct#25325
Conversation
- Allow COUNT to accept multiple arguments (len(inputs) >= 1) - Reject COUNT(col1, col2) without DISTINCT with syntax error - Skip GROUP BY optimization for multi-arg and tuple COUNT DISTINCT - Expand tuple args in distinct_agg to reuse multi-vector encoding path - Fix increaseRefCnt missing *plan.Expr_List handling (scan output columns) - Implement multi-vector encoding in batchFillArgs with length-prefix format - Update readStateArg to guard argTypes[0] access with usesOpaqueArgEncoding - Change makeCount/newCountColumnExec signatures to accept []types.Type - Add unit tests for multi-column distinct (distinct/NULL/multi-group/merge) - Update BVT test cases and expected results Co-Authored-By: Claude <noreply@anthropic.com>
Qodo reviews are paused for this user.Troubleshooting steps vary by plan Learn more → On a Teams plan? Using GitHub Enterprise Server, GitLab Self-Managed, or Bitbucket Data Center? |
XuPeng-SH
left a comment
XuPeng-SH
left a comment
There was a problem hiding this comment.
Request changes for one correctness blocker.
COUNT(DISTINCT (a, b)) is only normalized inside optimizeDistinctAgg, and only after the rule has already accepted the narrow single-aggregate shape. That leaves valid tuple syntax unexpanded when another aggregate is present in the same AGG node.
Example planner reproduction on this PR:
select count(distinct (n_nationkey, n_regionkey)), count(*) from nation;The resulting AGG node still contains count with args=1 and T_tuple for the first aggregate, while starcount is the second aggregate. This happens because distinct_agg.go returns before the tuple-expansion block when len(node.AggList) != 1.
That shape is not executable as a row-wise multi-column distinct key. The group operator evaluates aggregate args through MakeEvalVector, and Expr_List is handled by ListExpressionExecutor, which builds a vector from the list elements themselves (UnionOne(vec, 0) for each child) rather than producing one tuple value per input row. So an unexpanded tuple aggregate input is semantically wrong and can miscount or fail once BatchFill indexes it using input-row offsets.
Suggested fix:
- Do semantic normalization outside this optimization rule. Prefer flattening
COUNT(DISTINCT (a, b))into the same arg list asCOUNT(DISTINCT a, b)during binding or in a general aggregate-normalization pass that runs for every aggregate expression, independent oflen(node.AggList)and independent of whether the distinct-agg optimization fires. - Keep
optimizeDistinctAggonly as an optimization; it should not be required for correctness. - Add regression coverage for tuple syntax with more than one aggregate, for example:
select count(distinct (g, v)), count(*) from t_count_distinct_multi;select txt, count(distinct (g, v)), count(*) from t_count_distinct_multi group by txt;- compare
count(distinct (g, v))withcount(distinct g, v)in the same query.
The direct multi-arg executor path and local merge tests look directionally correct, but this tuple path still violates the PR's stated "tuple syntax, same semantics" contract.
Previously, COUNT(DISTINCT (a, b)) tuple args were only expanded inside optimizeDistinctAgg, which runs only when len(AggList)==1 and the single-aggregate shape is accepted. When another aggregate was present (e.g. count(distinct (a,b)), count(*)), the tuple was never expanded, causing incorrect results. Move the normalization to HavingBinder.BindAggFunc so that every aggregate expression with a T_tuple arg gets expanded to multi-arg form, independent of the optimization rule. Add regression tests: - tuple syntax with multiple aggregates - tuple syntax with GROUP BY - tuple vs multi-arg comparison in same query Co-Authored-By: Claude <noreply@anthropic.com>
XuPeng-SH
left a comment
XuPeng-SH
left a comment
There was a problem hiding this comment.
Requesting changes on current head 7580283.
The previous tuple-with-peer-aggregate blocker looks fixed, but the current implementation opens a broader COUNT argument-validation hole.
pkg/sql/plan/function/list_agg.go:33 now lets count bind with len(inputs) >= 1, while the non-DISTINCT multi-arg guard only exists in HavingBinder.BindAggFunc (pkg/sql/plan/having_binder.go:149). That does not cover tuple-as-one-arg and does not cover the window binder path. I verified these with a temporary planner test: each query below expected a bind error, but BuildPlan returned nil error on this PR:
select count((n_nationkey, n_regionkey)) from nation;
select count(n_nationkey, n_regionkey) over() from nation;
select count((n_nationkey, n_regionkey)) over() from nation;This is not only a parse/compat issue. If the non-DISTINCT aggregate tuple path reaches execution after HavingBinder expands the tuple, countColumnExec.BatchFill takes the non-DISTINCT branch and only reads vectors[0] (pkg/sql/colexec/aggexec/count2.go:169-175), so it can silently behave like COUNT(first_arg) instead of rejecting an invalid COUNT form.
Suggested fix direction:
- Put COUNT argument validation in one shared binder helper and call it from both aggregate and window binding paths before
bindFuncExprImplByAstExpr. - Only allow multiple COUNT arguments for aggregate
COUNT(DISTINCT expr, expr, ...). - Only expand
COUNT(DISTINCT (a, b))whenfuncName == "count" && astExpr.Type == tree.FUNC_TYPE_DISTINCT; do not expand tuple args for non-DISTINCT COUNT. - Window COUNT has DISTINCT disabled here already, so reject both
COUNT(a,b) OVER()andCOUNT((a,b)) OVER(). - Add regression coverage for the three negative queries above, alongside the existing positive tuple/peer-aggregate cases.
Also please add an executor round-trip test for multi-column COUNT(DISTINCT ...) through SaveIntermediateResult/UnmarshalFromReader. I locally checked that such a round trip currently works, but this PR changed writeStateArg/readStateArg opaque key encoding, so that path should be explicitly covered.
Introduce validateCountArgs() as a single validation helper, called from both HavingBinder.BindAggFunc and bindWindowFuncExpr before the bindFuncExprImplByAstExpr call. This closes two holes in the previous COUNT argument validation: 1. COUNT((a,b)) (tuple) — len(astExpr.Exprs)==1 so the old multi-arg guard missed it; the tuple was then expanded unconditionally. 2. Window COUNT(a,b) OVER() and COUNT((a,b)) OVER() — the window binder had no COUNT argument validation at all. The tuple-expansion block in BindAggFunc now additionally checks funcName=="count" && astExpr.Type==FUNC_TYPE_DISTINCT, so it never expands a non-DISTINCT tuple into multi-arg form. Fixes issues raised by XuPeng-SH's review on PR matrixorigin#25325. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
aunjgr
left a comment
aunjgr
left a comment
There was a problem hiding this comment.
Tuple expansion only fires in the optimization rule — count((a,b)) without DISTINCT still leaks through. Validation is scattered across HavingBinder, window binder, and the optimizer. Should have a single binding-time gate that rejects multi-arg / tuple non-DISTINCT COUNT uniformly.
…round-trip - BVT: count((g,v)), count(g,v) over(), count((g,v)) over() must fail with 'Incorrect arguments to COUNT' - UT: multi-column COUNT(DISTINCT) state survives SaveIntermediateResultOfChunk/UnmarshalFromReader round trip Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- having_binder: only expand a genuine Expr_List tuple, not a row subquery Expr_Sub that also carries Typ.Id == T_tuple, which nil-derefed GetList(). - list_agg: build a cast list matching arg count so COUNT(DISTINCT NULL, x) returns 0 instead of failing with 'cast types length not match args length'. - aggState: grow the distinct-key arena by enough to fit an oversized key (arenaskl.MaxNodeSize) instead of a fixed 512 KiB step, which could still ErrArenaFull for large multi-column keys. Tests: - plan: TestCountDistinctRowSubqueryNoPanic (BuildPlan must not panic). - aggexec: multi-large-key-arena-grow (two 700 KiB text cols). - bvt: COUNT(DISTINCT NULL, ...) positive cases. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
A multi-column row subquery — COUNT(DISTINCT (SELECT a, b ...)) — binds to an Expr_Sub whose Typ.Id is T_tuple but is not an Expr_List. Previously the guard merely skipped expansion to avoid a nil-deref, leaving a single T_tuple arg that downstream handled as an opaque NYI error and, in flatten paths, could collapse to the subquery's first column. Reject it explicitly during binding instead. Strengthen the planner test to assert the bind-time error (and nil plan) rather than only NotPanics, and add a BVT negative case. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…main # Conflicts: # pkg/sql/plan/build_test.go
XuPeng-SH
left a comment
XuPeng-SH
left a comment
There was a problem hiding this comment.
Re-reviewed current head 05e8dd1. The previous blockers are fixed in the right place:
- COUNT arg validation is now a shared binding-time gate for aggregate and window paths, so non-DISTINCT multi-arg/tuple COUNT is rejected before execution.
- COUNT(DISTINCT (a, b)) is normalized during aggregate binding instead of depending on the distinct-agg optimization rule, so tuple syntax works with peer aggregates and grouped aggregates.
- Multi-column row subquery input is rejected explicitly instead of leaking an Expr_Sub/T_tuple into execution.
- Executor state now treats multi-arg distinct keys as opaque encoded keys, with merge/marshal and large-key arena growth covered.
Local verification:
- go test ./pkg/sql/colexec/aggexec -run 'TestCount|TestPR25325Manual'
- go test ./pkg/sql/plan -run 'TestCountDistinctRowSubqueryRejected|TestPR25325ManualCountArgBindingClosure|TestRewriteCountNotNullColToStarcount'
I also added temporary local-only scratch cases for varchar boundary-collision keys like ('a','bc') vs ('ab','c') and for valid/invalid planner closure around aggregate/window COUNT. They passed and were removed before review submission.
CI is green on the current head. Non-blocking note: the PR is still BEHIND main, so it may need an update branch before merge depending on queue policy.
Merge Queue Status
This pull request spent 1 hour 37 minutes 2 seconds in the queue, with no time running CI. Waiting for
All conditions
ReasonThe merge conditions cannot be satisfied due to failing checks Failing checks:
HintYou may have to fix your CI before adding the pull request to the queue again. Requeued — the merge queue status continues in this comment ↓. |
Merge Queue Status
This pull request spent 23 seconds in the queue, including 1 second running CI. Required conditions to merge
|
5 participants
What type of PR is this?
Which issue(s) this PR fixes:
issue #25284
What this PR does / why we need it:
Support
COUNT(DISTINCT col1, col2, ...)multi-column distinct count, matching MySQL semantics:COUNT(DISTINCT col1, col2)— counts distinct non-NULL combinationsCOUNT(DISTINCT (col1, col2))— tuple syntax, same semanticsCOUNT(col1, col2)without DISTINCT — properly rejected with syntax errorChanges
list_agg.gohaving_binder.godistinct_agg.goutils.goincreaseRefCntmissing*plan.Expr_Listhandlingcount2.go,types.goaggState.gocount2_test.gofunc_aggr_count.test/result🤖 Generated with Claude Code