Rewrite AVG(expr) --> SUM(expr) / COUNT(expr) when components can be shared - #25536
wudidapaopao wants to merge 4 commits into
Conversation
github-actions
Bot
added
logical-expr
add0efd to
9acbb37
Compare
|
run benchmark tpch |
|
🤖 Benchmark running (GKE) | trigger CPU Details (lscpu)Comparing avg-simplify-sum-count (9acbb37) to c149764 (merge-base) diff Run configurationrun benchmark tpchResults will be posted here when complete File an issue against this benchmark runner |
|
🤖 Benchmark completed (GKE) | trigger Instance: Comparing avg-simplify-sum-count (9acbb37) to c149764 (merge-base) diff Run configurationrun benchmark tpchCPU Details (lscpu)Details
Resource Usagetpch — base (merge-base)
tpch — branch
File an issue against this benchmark runner |
# Conflicts: # datafusion/sqllogictest/test_files/aggregates_simplify.slt
# Conflicts: # datafusion/expr/src/udaf.rs
|
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 |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #25536 +/- ##
========================================
Coverage 82.42% 82.43%
========================================
Files 1138 1139 +1
Lines 435501 435803 +302
Branches 435501 435803 +302
========================================
+ Hits 358955 359238 +283
- Misses 54841 54843 +2
- Partials 21705 21722 +17 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
wudidapaopao
changed the title
AVG(expr) --> SUM(expr) / COUNT(expr) when components can be shared
# Conflicts: # datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs # datafusion/substrait/src/logical_plan/consumer/utils.rs
Could you please provide some instructions on how to reproduce these numbers? |
alamb
left a comment
alamb
left a comment
There was a problem hiding this comment.
Thank you @wudidapaopao -- this is a neat idea. I left some comments. Let me know what you think
| /// Returns this aggregate function's candidate decomposition, if any. | ||
| /// | ||
| /// See [`AggregateUDFImpl::decompose`] for more details. | ||
| pub fn decompose( |
There was a problem hiding this comment.
I wonder if you considered using the existing simplify method:
If you changed the avg udf to simplify to sum/count the existing common subexpr eliminate path probably will already avoid the recomputation.
Also it woudl allow us to delete the actual AVG accumulators (rather than having a special case like this) 🤔
There was a problem hiding this comment.
Thanks, I considered using simplify. I think we should retain the AVG accumulator and only decompose AVG when its generated SUM or COUNT can be shared. If implemented in simplify, every AVG would be unconditionally rewritten into SUM/COUNT.
Benchmark: 20 million random non-null Int64 rows, single-threaded.
| Scenario | SQL | Before decomposition | After decomposition | Change |
|---|---|---|---|---|
| No sharing | SELECT AVG(x) FROM t |
9.06 ms | 10.72 ms | 18.29% slower |
| One reusable SUM | SELECT SUM(CAST(x AS DOUBLE)), AVG(x) FROM t |
11.72 ms | 10.79 ms | 7.95% faster |
| Three reusable SUMs | SELECT SUM(CAST(x AS DOUBLE)), AVG(x), SUM(CAST(y AS DOUBLE)), AVG(y), SUM(CAST(z AS DOUBLE)), AVG(z) FROM t |
32.91 ms | 27.85 ms | 15.38% faster |
There was a problem hiding this comment.
If implemented in simplify, every AVG would be unconditionally rewritten into SUM/COUNT.
Given the internal avg implementation basically has a sum and count accumulator, I am surprised at these numbers. Can you profile them and find out why there is a performance difference?
There was a problem hiding this comment.
Thanks for pointing this out. I found that COUNT(*) materializes a full Int64Array for each batch. I will optimize this in a separate PR, then continue this PR.
| }; | ||
|
|
||
| let rewrote_aggs = rewrite_multiple_linear_aggregates(&mut aggr_expr)?; | ||
| let rewrote_linear = rewrite_multiple_linear_aggregates(&mut aggr_expr)?; |
There was a problem hiding this comment.
I think our CSE pass already does something like this 🤔
There was a problem hiding this comment.
Yes, the existing CSE already deduplicates repeated SUM or COUNT expressions. rewrite_shared_aggregate_components does not replace that deduplication. It determines whether decomposing AVG enables SUM or COUNT sharing without increasing the total number of distinct aggregate expressions, and only then applies the decomposition. The resulting duplicate aggregates are still deduplicated by the existing CSE path.
I converted the four non-nullable import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq
source = "..."
destination = "..."
table = pq.read_table(source)
for name in [
"l_quantity",
"l_extendedprice",
"l_discount",
"l_tax",
]:
index = table.schema.get_field_index(name)
table = table.set_column(
index,
pa.field(name, pa.float64(), nullable=False),
pc.cast(table.column(index), pa.float64()),
)
pq.write_table(
table,
destination,
compression="zstd",
compression_level=1,
row_group_size=131072,
) |
4 participants
Which issue does this PR close?
Rationale for this change
AVGmaintains both sum and count state. When the same aggregate node already computes a matchingSUMorCOUNT, those states are redundant.What changes are included in this PR?
AVGexpressions into sharedSUM/COUNTcomponents.AVGaccumulator when no component can be shared.NameTrackerfor internal aggregate name conflicts.For example, given a Float64 column:
the optimized plan is equivalent to:
The existing
SUM(x)is computed once and reused byAVG(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 returnsDecimal128(25, 2), while AVG uses aDecimal128(38, 2)internal sum and returnsDecimal128(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 standardTPC-H Q1, which uses Decimal columns, does not benefit from this PR as initially expected.What is the testing strategy for this PR?
avg_to_sum_count.sltcovering shared SUM/COUNT, no-share cases, NULL/empty input, grouping sets, unsupported AVG forms, Decimal, and naming conflicts.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:
Are there any user-facing changes?
Adds an optional
AggregateUDFImpl::decomposehook 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.