[Enhancement] Push down non-group-by aggregate through union all

[JasonQu1]

Why I'm doing:

Statistics collection queries may aggregate data above UNION ALL, causing a large amount of raw sampled rows to be exchanged and merged before local aggregation. This can increase network transfer, memory usage, and execution time for non-group-by aggregate workloads such as ANALYZE SAMPLE.

This PR pushes local split aggregates into UNION ALL branches when the physical plan is already a safe two-stage aggregate plan, so each branch can aggregate raw rows earlier and UNION ALL only needs to pass intermediate aggregate states.

What I'm doing:

Add a physical plan rewrite rule to push eligible non-group-by local split aggregates into UNION ALL branches, while preserving the upper merge stage to combine intermediate aggregate states.

Also add tests to verify both plan rewriting and result correctness.

Performance Test

Environment

Test environment:

• 3 BE nodes
• Test table row count: 200M
• Buckets: 48

Test table schema:

CREATE TABLE t_analyze_union_pushdown_200m_bigstr (
    id BIGINT,
    bucket_key INT,
    v_int INT,
    v_bigint BIGINT,
    v_decimal DECIMAL(18, 4),
    v_double DOUBLE,
    v_str VARCHAR(1000),
    v_high_card_str VARCHAR(2000),
    v_date DATE,
    v_array ARRAY<INT>
)
DUPLICATE KEY(id)
DISTRIBUTED BY HASH(id) BUCKETS 48;

Test 1: ANALYZE SAMPLE

This test runs sample statistics collection on the 200M-row table:

ANALYZE SAMPLE TABLE t_analyze_union_pushdown_200m_bigstr;

The result below is the average of 5 runs.

Feature	Avg Wall Time	Avg Max QueryPeakMemoryUsage
OFF	2.659s	687.6 MB
ON	2.330s	521.0 MB

Result:

• Wall time reduced by about 12.4%
• QueryPeakMemoryUsage reduced by about 24.2%

Test 2: Statistics-like SQL

This test uses a statistics-collection-like SQL on the same 200M-row table. The SQL contains multiple UNION ALL branches and aggregate functions commonly used by statistics collection, including COUNT, SUM, MIN, MAX, and HLL_RAW.

Default parallelism settings are used.

WITH base AS (
    SELECT *
    FROM t_analyze_union_pushdown_200m_bigstr
    WHERE id <= 50000000
    UNION ALL
    SELECT *
    FROM t_analyze_union_pushdown_200m_bigstr
    WHERE id > 50000000 AND id <= 100000000
    UNION ALL
    SELECT *
    FROM t_analyze_union_pushdown_200m_bigstr
    WHERE id > 100000000 AND id <= 150000000
    UNION ALL
    SELECT *
    FROM t_analyze_union_pushdown_200m_bigstr
    WHERE id > 150000000
)
SELECT
    COUNT(*),
    HEX(HLL_SERIALIZE(IFNULL(HLL_RAW(v_high_card_str), HLL_EMPTY()))),
    SUM(CHAR_LENGTH(v_high_card_str)),
    COUNT(v_high_card_str),
    MAX(LEFT(v_high_card_str, 200)),
    MIN(LEFT(v_high_card_str, 200)),
    HEX(HLL_SERIALIZE(IFNULL(HLL_RAW(CAST(v_bigint AS STRING)), HLL_EMPTY()))),
    SUM(v_bigint),
    MAX(v_bigint),
    MIN(v_bigint),
    HEX(HLL_SERIALIZE(IFNULL(HLL_RAW(v_str), HLL_EMPTY()))),
    SUM(v_int),
    MAX(v_int),
    MIN(v_int),
    SUM(v_decimal),
    MAX(v_decimal),
    MIN(v_decimal),
    SUM(v_double),
    MIN(v_date),
    MAX(v_date)
FROM base;

The result below is the average of 5 runs.

Feature	Avg Wall Time	Avg QueryPeakMemoryUsage
OFF	35.411s	19.817 GB
ON	29.764s	8.755 GB

Result:

• Wall time reduced by about 15.9%
• QueryPeakMemoryUsage reduced by about 55.8%

What type of PR is this:

• BugFix
• Feature
• Enhancement
• Refactor
• UT
• Doc
• Tool

Does this PR entail a change in behavior?

• Yes, this PR will result in a change in behavior.
• No, this PR will not result in a change in behavior.

If yes, please specify the type of change:

• Interface/UI changes: syntax, type conversion, expression evaluation, display information
• Parameter changes: default values, similar parameters but with different default values
• Policy changes: use new policy to replace old one, functionality automatically enabled
• Feature removed
• Miscellaneous: upgrade & downgrade compatibility, etc.

Checklist:

• I have added test cases for my bug fix or my new feature
• This pr needs user documentation (for new or modified features or behaviors)
  • I have added documentation for my new feature or new function
  • This pr needs auto generate documentation
• This is a backport pr

Bugfix cherry-pick branch check:

• I have checked the version labels which the pr will be auto-backported to the target branch
  • 4.1
  • 4.0
  • 3.5

[eshishki]

Drive-by review — I'm not a maintainer, just a passer-by who found this interesting. Take or leave anything that doesn't fit.

Overall the rewrite looks correct. The property that makes it safe is the all-or-nothing decomposition through UNION ALL: in PushDownAggregateCollector.visitLogicalUnion the allChildRewriteContext.size() != collectors.size() check pushes the partial aggregate into every branch or none, so the COUNT→SUM / HLL_RAW→HLL_UNION decomposition can never be applied to only some branches (which would mis-count). The blocking rules reinforce this:

• UNION DISTINCT — the non-group-by aggregate is not pushed into the branches; it stays above the union, so dedup still happens before aggregation.
• JOIN — the aggregate is not pushed below the join; under a UNION ALL branch that contains a join, that branch yields no push-down context and the all-or-nothing guard then abandons the whole push-down.

Both are the conservative-correct choice.

On the PR description: it says the whitelist "adds COUNT and HLL_RAW", but the change actually enables the full NON_GROUP_BY_WHITE_FNS set for the non-group-by case — MAX/MIN/SUM/HLL_UNION/BITMAP_UNION/PERCENTILE_UNION were unreachable for non-group-by before (the removed groupingKeys.isEmpty() short-circuit) and are now pushed too; only COUNT and HLL_RAW need the final-stage remap. Could you spell out the complete supported set in the description, and which ones need a combiner remap? It makes the actual blast radius clear for reviewers and future readers.

A few things worth considering:

1. Gate the non-group-by push-down on input size. For the non-group-by path checkStatistics has no group-by columns to evaluate, so it returns true unconditionally — every branch gets an aggregate node even when it scans almost nothing, which is pure overhead. There's already a precedent for the floor in the same method (the medium-cardinality branch ends with statistics.getOutputRowCount() >= SMALL_SCALE_ROWS_LIMIT), and the estimate is already computed here via estimatorStats(), so an explicit non-group-by arm with a row-count floor is a couple of lines.

2. Make the whitelist and the partial→final remap a single source of truth. NON_GROUP_BY_WHITE_FNS (collector) and the COUNT→SUM / HLL_RAW→HLL_UNION mapping (PushDownAggregateRewriter.genAggregation) live apart, and nothing keeps them in sync. The self-mergeable functions (MAX/MIN/SUM/*_UNION) happen to need no remap, but a future edit that adds e.g. AVG to the whitelist without a matching final-stage combiner would silently produce wrong results. Consider a single Map<fnName, mergeFnName> — whitelist = keySet(), remap = lookup — where self-mergeable functions map to themselves. That also makes adding safe functions trivial (e.g. ANY_VALUE → ANY_VALUE).

3. (Desirable) Extend to Iceberg/MOR tables. Tables with equality deletes are rewritten by IcebergEqualityDeleteRewriteRule into a UNION ALL whose delete-applied branch is always a LEFT ANTI JOIN. Today, because non-group-by is blocked at joins, the all-or-nothing guard aborts push-down for these tables — the aggregate stays on top of the union and every row funnels through the exchange:

Aggregate [count(*), sum(v), hll_raw(v)]      ← single node, ALL rows funnel here
  └─ UNION ALL
       ├─ IcebergScan            (no-delete data files)
       └─ Project
            └─ LEFT ANTI JOIN
                 ├─ IcebergScan   (with-delete data files)
                 └─ IcebergEqualityDeleteScan

What we'd want is the same decomposition this PR already does for plain UNION ALL, with the branch-2 partial landed above the anti-join (never below it):

Aggregate [sum(cnt), sum(s), hll_union(h)]    ← final combiner, 1 row per branch in
  └─ UNION ALL
       ├─ Aggregate [count(*), sum(v), hll_raw(v)]   ← partial, on the scan
       │    └─ IcebergScan            (no-delete data files)
       └─ Aggregate [count(*), sum(v), hll_raw(v)]   ← partial, ABOVE the anti-join
            └─ Project
                 └─ LEFT ANTI JOIN
                      ├─ IcebergScan   (with-delete data files)
                      └─ IcebergEqualityDeleteScan

Landing the partial above a non-inflating anti/semi join is always correct (it aggregates exactly what the branch feeds the union), and the "land above a join" machinery already exists for the small-broadcast-join case. The push-vs-not decision uses the same row-count floor as #1, just read at the join's output rather than a scan — the only caveat being that the anti-join's output cardinality is a weaker estimate. For this PR, a regression test asserting that a non-group-by aggregate over a UNION ALL with a join branch does not push down (and returns correct results) would lock in the current behavior.

[CelerData-Reviewer]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Bravo.

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[Wenjun7J]

LGTM

[Copilot]

Pull request overview

This PR enhances FE optimizer aggregate push-down so non-group-by aggregates over UNION ALL can be partially computed in each branch and merged above the union, reducing raw-row exchange pressure for sampled-statistics style queries.

Changes:

• Adds non-group-by aggregate eligibility/remapping for COUNT and HLL_RAW, plus self-mergeable aggregate functions.
• Blocks unsafe non-group-by push-down through joins and UNION DISTINCT.
• Adds regression tests for union-all push-down, join/distinct guards, filter/project edge cases, and local/global/auto modes.

Reviewed changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
PushDownAggregateCollector.java	Extends collection/eligibility logic for non-group-by aggregate push-down and adds safety/statistics guards.
PushDownAggregateRewriter.java	Rewrites final-stage aggregate functions and handles zero-argument COUNT(*) safely.
AggregatePushDownTest.java	Adds optimizer plan tests covering the new push-down behavior and regressions.

[CelerData-Reviewer]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: eff7e0d956

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[CelerData-Reviewer]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f868405349

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[JasonQu1]

@satanson Submitted to the community edition, PTAL.

[github-actions]

[Java-Extensions Incremental Coverage Report]

✅ pass : 0 / 0 (0%)

[github-actions]

[FE Incremental Coverage Report]

✅ pass : 201 / 224 (89.73%)

file detail

	path	covered_line	new_line	coverage	not_covered_line_detail
🔵	com/starrocks/sql/optimizer/rule/tree/PushDownNonGroupedAggregateBelowUnion.java	199	222	89.64%	[145, 146, 155, 230, 295, 317, 318, 319, 321, 322, 323, 324, 325, 326, 342, 343, 344, 345, 347, 348, 349, 350, 351]
🔵	com/starrocks/sql/optimizer/QueryOptimizer.java	1	1	100.00%	[]
🔵	com/starrocks/common/Config.java	1	1	100.00%	[]

[github-actions]

[BE Incremental Coverage Report]

✅ pass : 0 / 0 (0%)

[github-actions]

@Mergifyio backport branch-4.1

[mergify]

backport branch-4.1

✅ Backports have been created

Details

• #75624 [Enhancement] Push down non-group-by aggregate through union all (backport #73930) has been created for branch branch-4.1