Skip to content

Rewrite AVG(expr) --> SUM(expr) / COUNT(expr) when components can be shared - #25536

Open
wudidapaopao wants to merge 3 commits into
apache:mainfrom
wudidapaopao:avg-simplify-sum-count
Open

wudidapaopao wants to merge 3 commits into
apache:mainfrom
wudidapaopao:avg-simplify-sum-count

Conversation

@wudidapaopao

@wudidapaopao wudidapaopao commented Sep 20, 2026

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

Rationale for this change

AVG maintains both sum and count state. When the same aggregate node already computes a matching SUM or COUNT, those states are redundant.

What changes are included in this PR?

  • Add an aggregate decomposition hook and apply candidates only when deduplication does not increase the number of distinct aggregates.
  • Decompose eligible Float64 AVG expressions into shared SUM / COUNT components.
  • Keep the combined AVG accumulator when no component can be shared.
  • Reuse NameTracker for internal aggregate name conflicts.

For example, given a Float64 column:

SELECT SUM(x), AVG(x) FROM t;

the optimized plan is equivalent to:

Projection: sum_x, sum_x / CAST(count_x AS Float64)
  Aggregate: SUM(x) AS sum_x, COUNT(x) AS count_x

The existing SUM(x) is computed once and reused by AVG(x). In contrast, SELECT AVG(x) alone keeps the original combined AVG accumulator because decomposition would add an aggregate.

Decimal AVG is not included because its accumulation and result types differ from regular Decimal SUM. For example, for Decimal128(15, 2), regular SUM returns Decimal128(25, 2), while AVG uses a Decimal128(38, 2) internal sum and returns Decimal128(19, 6). Reusing the regular SUM directly could change scale and overflow behavior. There is not yet a satisfactory way to share these states while preserving those semantics, so this PR does not optimize Decimal AVG. As a result, the standard TPC-H Q1, which uses Decimal columns, does not benefit from this PR as initially expected.

What is the testing strategy for this PR?

  • Added avg_to_sum_count.slt covering shared SUM/COUNT, no-share cases, NULL/empty input, grouping sets, unsupported AVG forms, Decimal, and naming conflicts.
  • Added unit tests for decomposition profitability and UDAF forwarding.
  • Ran cargo fmt --all, full workspace Clippy with -D warnings, and the extended workspace test suite.

Release benchmark using TPC-H SF1 with the four Q1 Decimal inputs materialized as Float64 Parquet:

  • baseline median: 29.137 ms
  • this PR median: 25.978 ms
  • improvement: 10.84%

Are there any user-facing changes?

Adds an optional AggregateUDFImpl::decompose hook with a default no-op implementation. Existing UDAFs do not need changes.

Eligible Float64 AVG aggregates may use shared SUM/COUNT states. Query results and output schemas are unchanged.

@github-actions github-actions Bot added logical-expr Logical plan and expressions optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) substrait Changes to the substrait crate functions Changes to functions implementation labels Sep 20, 2026
@alamb

alamb commented Sep 20, 2026

Copy link
Copy Markdown
Contributor

run benchmark tpch

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark running (GKE) | trigger
Instance: c4a-highmem-16 (12 vCPU / 65 GiB) | Linux bench-c5749175701-2523-w9qgt 6.12.94+ #1 SMP Tue Aug 4 08:44:15 UTC 2026 aarch64 GNU/Linux

CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected

Comparing avg-simplify-sum-count (9acbb37) to c149764 (merge-base) diff

Run configuration
run benchmark tpch

Results will be posted here when complete


File an issue against this benchmark runner

@adriangbot

Copy link
Copy Markdown

🤖 Benchmark completed (GKE) | trigger

Instance: c4a-highmem-16 (12 vCPU / 65 GiB)

Comparing avg-simplify-sum-count (9acbb37) to c149764 (merge-base) diff

Run configuration
run benchmark tpch
CPU Details (lscpu)
Architecture:                            aarch64
CPU op-mode(s):                          64-bit
Byte Order:                              Little Endian
CPU(s):                                  16
On-line CPU(s) list:                     0-15
Vendor ID:                               ARM
Model name:                              Neoverse-V2
Model:                                   1
Thread(s) per core:                      1
Core(s) per cluster:                     16
Socket(s):                               -
Cluster(s):                              1
Stepping:                                r0p1
BogoMIPS:                                2000.00
Flags:                                   fp asimd evtstrm aes pmull sha1 sha2 crc32 atomics fphp asimdhp cpuid asimdrdm jscvt fcma lrcpc dcpop sha3 sm3 sm4 asimddp sha512 sve asimdfhm dit uscat ilrcpc flagm sb paca pacg dcpodp sve2 sveaes svepmull svebitperm svesha3 svesm4 flagm2 frint svei8mm svebf16 i8mm bf16 dgh rng bti
L1d cache:                               1 MiB (16 instances)
L1i cache:                               1 MiB (16 instances)
L2 cache:                                32 MiB (16 instances)
L3 cache:                                80 MiB (1 instance)
NUMA node(s):                            1
NUMA node0 CPU(s):                       0-15
Vulnerability Gather data sampling:      Not affected
Vulnerability Indirect target selection: Not affected
Vulnerability Itlb multihit:             Not affected
Vulnerability L1tf:                      Not affected
Vulnerability Mds:                       Not affected
Vulnerability Meltdown:                  Not affected
Vulnerability Mmio stale data:           Not affected
Vulnerability Reg file data sampling:    Not affected
Vulnerability Retbleed:                  Not affected
Vulnerability Spec rstack overflow:      Not affected
Vulnerability Spec store bypass:         Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:                Mitigation; __user pointer sanitization
Vulnerability Spectre v2:                Mitigation; CSV2, BHB
Vulnerability Srbds:                     Not affected
Vulnerability Tsa:                       Not affected
Vulnerability Tsx async abort:           Not affected
Vulnerability Vmscape:                   Not affected
Details

Comparing HEAD and avg-simplify-sum-count
--------------------
Benchmark tpch_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Query     ┃     HEAD ┃ avg-simplify-sum-count ┃    Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ QQuery 1  │ 40.49 ms │               40.39 ms │ no change │
│ QQuery 2  │ 18.79 ms │               18.90 ms │ no change │
│ QQuery 3  │ 28.11 ms │               27.97 ms │ no change │
│ QQuery 4  │ 17.37 ms │               17.14 ms │ no change │
│ QQuery 5  │ 34.84 ms │               35.03 ms │ no change │
│ QQuery 6  │ 16.42 ms │               16.42 ms │ no change │
│ QQuery 7  │ 41.03 ms │               41.32 ms │ no change │
│ QQuery 8  │ 40.90 ms │               40.81 ms │ no change │
│ QQuery 9  │ 49.53 ms │               49.68 ms │ no change │
│ QQuery 10 │ 41.94 ms │               41.92 ms │ no change │
│ QQuery 11 │ 13.23 ms │               13.22 ms │ no change │
│ QQuery 12 │ 24.65 ms │               24.36 ms │ no change │
│ QQuery 13 │ 39.26 ms │               38.97 ms │ no change │
│ QQuery 14 │ 24.40 ms │               24.59 ms │ no change │
│ QQuery 15 │ 31.02 ms │               30.91 ms │ no change │
│ QQuery 16 │ 13.67 ms │               13.74 ms │ no change │
│ QQuery 17 │ 72.45 ms │               72.71 ms │ no change │
│ QQuery 18 │ 60.37 ms │               59.99 ms │ no change │
│ QQuery 19 │ 33.34 ms │               33.28 ms │ no change │
│ QQuery 20 │ 32.20 ms │               31.76 ms │ no change │
│ QQuery 21 │ 56.26 ms │               55.09 ms │ no change │
│ QQuery 22 │ 13.88 ms │               13.55 ms │ no change │
└───────────┴──────────┴────────────────────────┴───────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ Benchmark Summary                     ┃          ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│ Total Time (HEAD)                     │ 744.16ms │
│ Total Time (avg-simplify-sum-count)   │ 741.76ms │
│ Average Time (HEAD)                   │  33.83ms │
│ Average Time (avg-simplify-sum-count) │  33.72ms │
│ Queries Faster                        │        0 │
│ Queries Slower                        │        0 │
│ Queries with No Change                │       22 │
│ Queries with Failure                  │        0 │
└───────────────────────────────────────┴──────────┘

Distribution per query (min / mean ±stddev / max):

Comparing HEAD and avg-simplify-sum-count
--------------------
Benchmark tpch_sf1.json
--------------------
┏━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━┓
┃ Query     ┃                           HEAD ┃         avg-simplify-sum-count ┃    Change ┃
┡━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━┩
│ QQuery 1  │ 40.49 / 41.53 ±1.21 / 43.85 ms │ 40.39 / 41.68 ±1.42 / 43.57 ms │ no change │
│ QQuery 2  │ 18.79 / 19.05 ±0.20 / 19.34 ms │ 18.90 / 19.27 ±0.37 / 19.89 ms │ no change │
│ QQuery 3  │ 28.11 / 28.61 ±0.31 / 28.97 ms │ 27.97 / 28.57 ±0.37 / 29.09 ms │ no change │
│ QQuery 4  │ 17.37 / 18.00 ±0.64 / 18.92 ms │ 17.14 / 17.37 ±0.14 / 17.51 ms │ no change │
│ QQuery 5  │ 34.84 / 35.02 ±0.12 / 35.21 ms │ 35.03 / 35.69 ±1.00 / 37.67 ms │ no change │
│ QQuery 6  │ 16.42 / 16.61 ±0.24 / 17.08 ms │ 16.42 / 16.48 ±0.05 / 16.54 ms │ no change │
│ QQuery 7  │ 41.03 / 41.61 ±0.37 / 42.16 ms │ 41.32 / 42.62 ±1.13 / 44.17 ms │ no change │
│ QQuery 8  │ 40.90 / 42.19 ±1.53 / 45.14 ms │ 40.81 / 42.16 ±1.07 / 43.85 ms │ no change │
│ QQuery 9  │ 49.53 / 50.43 ±0.73 / 51.53 ms │ 49.68 / 50.15 ±0.35 / 50.64 ms │ no change │
│ QQuery 10 │ 41.94 / 42.37 ±0.37 / 42.92 ms │ 41.92 / 42.55 ±0.96 / 44.45 ms │ no change │
│ QQuery 11 │ 13.23 / 14.02 ±0.57 / 14.97 ms │ 13.22 / 13.53 ±0.20 / 13.84 ms │ no change │
│ QQuery 12 │ 24.65 / 24.97 ±0.32 / 25.57 ms │ 24.36 / 24.56 ±0.23 / 24.99 ms │ no change │
│ QQuery 13 │ 39.26 / 41.72 ±2.47 / 46.17 ms │ 38.97 / 40.58 ±2.28 / 45.12 ms │ no change │
│ QQuery 14 │ 24.40 / 24.79 ±0.32 / 25.26 ms │ 24.59 / 24.89 ±0.29 / 25.35 ms │ no change │
│ QQuery 15 │ 31.02 / 31.19 ±0.13 / 31.37 ms │ 30.91 / 31.32 ±0.31 / 31.82 ms │ no change │
│ QQuery 16 │ 13.67 / 13.82 ±0.11 / 14.01 ms │ 13.74 / 14.05 ±0.36 / 14.74 ms │ no change │
│ QQuery 17 │ 72.45 / 74.00 ±1.55 / 76.30 ms │ 72.71 / 73.72 ±1.33 / 76.30 ms │ no change │
│ QQuery 18 │ 60.37 / 61.90 ±1.52 / 64.73 ms │ 59.99 / 61.31 ±1.12 / 62.50 ms │ no change │
│ QQuery 19 │ 33.34 / 34.00 ±0.78 / 35.43 ms │ 33.28 / 33.56 ±0.21 / 33.92 ms │ no change │
│ QQuery 20 │ 32.20 / 32.42 ±0.20 / 32.72 ms │ 31.76 / 32.39 ±0.67 / 33.70 ms │ no change │
│ QQuery 21 │ 56.26 / 56.98 ±0.66 / 58.12 ms │ 55.09 / 55.53 ±0.34 / 56.02 ms │ no change │
│ QQuery 22 │ 13.88 / 14.32 ±0.51 / 15.29 ms │ 13.55 / 13.80 ±0.17 / 13.95 ms │ no change │
└───────────┴────────────────────────────────┴────────────────────────────────┴───────────┘
┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━┓
┃ Benchmark Summary                     ┃          ┃
┡━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━┩
│ Total Time (HEAD)                     │ 759.57ms │
│ Total Time (avg-simplify-sum-count)   │ 755.78ms │
│ Average Time (HEAD)                   │  34.53ms │
│ Average Time (avg-simplify-sum-count) │  34.35ms │
│ Queries Faster                        │        0 │
│ Queries Slower                        │        0 │
│ Queries with No Change                │       22 │
│ Queries with Failure                  │        0 │
└───────────────────────────────────────┴──────────┘

Resource Usage

tpch — base (merge-base)

Metric Value
Wall time 5.0s
Peak memory 1.2 GiB
Avg memory 496.0 MiB
CPU user 21.2s
CPU sys 1.7s
Peak spill 0 B

tpch — branch

Metric Value
Wall time 5.0s
Peak memory 1.3 GiB
Avg memory 510.3 MiB
CPU user 21.1s
CPU sys 1.7s
Peak spill 0 B

File an issue against this benchmark runner

# Conflicts:
#	datafusion/sqllogictest/test_files/aggregates_simplify.slt
@wudidapaopao

wudidapaopao commented Sep 20, 2026

Copy link
Copy Markdown
Contributor Author

Thanks for running the benchmark. The result shows no improvement because the standard TPC-H schema uses Decimal columns, while this PR only decomposes Float64 AVG, so the new optimization is not triggered.

In a local release benchmark, I materialized the four Q1 Decimal input columns as Float64 Parquet. Q1 improved from a 29.137 ms median to 25.978 ms, or 10.84%.

@codecov-commenter

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 92.67139% with 31 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.40%. Comparing base (a9579bd) to head (749751c).
⚠️ Report is 2 commits behind head on main.

Files with missing lines Patch % Lines
datafusion/expr/src/utils.rs 88.97% 0 Missing and 14 partials ⚠️
...timizer/src/simplify_expressions/simplify_exprs.rs 84.78% 0 Missing and 7 partials ⚠️
datafusion/expr/src/udaf.rs 91.11% 2 Missing and 2 partials ⚠️
datafusion/functions-aggregate/src/average.rs 91.17% 1 Missing and 2 partials ⚠️
...rc/simplify_expressions/aggregate_decomposition.rs 98.24% 0 Missing and 3 partials ⚠️
Additional details and impacted files
@@           Coverage Diff            @@
##             main   #25536    +/-   ##
========================================
  Coverage   82.39%   82.40%            
========================================
  Files        1138     1139     +1     
  Lines      434623   434914   +291     
  Branches   434623   434914   +291     
========================================
+ Hits       358100   358371   +271     
- Misses      54856    54862     +6     
- Partials    21667    21681    +14     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
  • 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

@wudidapaopao wudidapaopao changed the title Optimize shared Float64 AVG aggregates Rewrite AVG(expr) --> SUM(expr) / COUNT(expr) when components can be shared Sep 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

functions Changes to functions implementation logical-expr Logical plan and expressions optimizer Optimizer rules sqllogictest SQL Logic Tests (.slt) substrait Changes to the substrait crate

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Convert AVG(col) to SUM(x) / COUNT(*)

4 participants