From 35595568330093c46617889e06a3808a6cc0a8b1 Mon Sep 17 00:00:00 2001 From: discord9 <55937128+discord9@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:37:19 +0800 Subject: [PATCH] fix: preserve window inputs during TopK aggregation Signed-off-by: discord9 <55937128+discord9@users.noreply.github.com> --- .../src/topk_aggregation.rs | 6 + .../test_files/aggregates_topk.slt | 115 ++++++++++++++++++ 2 files changed, 121 insertions(+) diff --git a/datafusion/physical-optimizer/src/topk_aggregation.rs b/datafusion/physical-optimizer/src/topk_aggregation.rs index 0eddb5d5507e4..1907ce053840a 100644 --- a/datafusion/physical-optimizer/src/topk_aggregation.rs +++ b/datafusion/physical-optimizer/src/topk_aggregation.rs @@ -30,6 +30,7 @@ use datafusion_physical_plan::aggregates::{AggregateExec, topk_types_supported}; use datafusion_physical_plan::execution_plan::CardinalityEffect; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use itertools::Itertools; /// An optimizer rule that passes a `limit` hint to aggregations if the whole result is not needed @@ -153,6 +154,11 @@ impl TopKAggregation { cur_col_name = src_col.name().to_string(); } } + } else if plan.downcast_ref::().is_some() + || plan.downcast_ref::().is_some() + { + // Window values can depend on rows removed by a bounded aggregate. + cardinality_preserved = false; } else { // or we continue down through types that don't reduce cardinality match plan.cardinality_effect() { diff --git a/datafusion/sqllogictest/test_files/aggregates_topk.slt b/datafusion/sqllogictest/test_files/aggregates_topk.slt index e2d453068adf7..164de44a98438 100644 --- a/datafusion/sqllogictest/test_files/aggregates_topk.slt +++ b/datafusion/sqllogictest/test_files/aggregates_topk.slt @@ -789,3 +789,118 @@ drop table topk_two_groups; statement ok drop table t0; + +# A window result depends on every group, so the TopK limit must not pass +# through either window implementation to the aggregate below it. +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +CREATE TABLE topk_window_groups(g int) AS VALUES (1), (2); + +statement ok +set datafusion.optimizer.enable_topk_aggregation = false; + +# Control: without the outer limit, both groups contribute to the window value. +query II +SELECT g, COUNT(*) OVER () AS n +FROM topk_window_groups +GROUP BY g +ORDER BY g ASC; +---- +1 2 +2 2 + +query II +SELECT g, COUNT(*) OVER () AS n +FROM topk_window_groups +GROUP BY g +ORDER BY g ASC +LIMIT 1; +---- +1 2 + +statement ok +set datafusion.optimizer.enable_topk_aggregation = true; + +query II +SELECT g, COUNT(*) OVER () AS n +FROM topk_window_groups +GROUP BY g +ORDER BY g ASC +LIMIT 1; +---- +1 2 + +# WindowAggExec must prevent a TopK aggregate limit. +query TT +EXPLAIN SELECT g, COUNT(*) OVER () AS n +FROM topk_window_groups +GROUP BY g +ORDER BY g ASC +LIMIT 1; +---- +logical_plan +01)Sort: topk_window_groups.g ASC NULLS LAST, fetch=1 +02)--Projection: topk_window_groups.g, count(Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS count(*) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS n +03)----WindowAggr: windowExpr=[[count(Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] +04)------Aggregate: groupBy=[[topk_window_groups.g]], aggr=[[]] +05)--------TableScan: topk_window_groups projection=[g] +physical_plan +01)ProjectionExec: expr=[g@0 as g, count(Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as n] +02)--SortExec: TopK(fetch=1), expr=[g@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----WindowAggExec: wdw=[count(Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "count(Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +04)------AggregateExec: mode=Single, gby=[g@0 as g], aggr=[] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +set datafusion.optimizer.enable_topk_aggregation = false; + +query II +SELECT g, COUNT(*) OVER (ORDER BY g DESC ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS n +FROM topk_window_groups +GROUP BY g +ORDER BY g ASC +LIMIT 1; +---- +1 2 + +statement ok +set datafusion.optimizer.enable_topk_aggregation = true; + +query II +SELECT g, COUNT(*) OVER (ORDER BY g DESC ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS n +FROM topk_window_groups +GROUP BY g +ORDER BY g ASC +LIMIT 1; +---- +1 2 + +# The preceding-row frame selects BoundedWindowAggExec, which must also stop TopK. +query TT +EXPLAIN SELECT g, COUNT(*) OVER (ORDER BY g DESC ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS n +FROM topk_window_groups +GROUP BY g +ORDER BY g ASC +LIMIT 1; +---- +logical_plan +01)Sort: topk_window_groups.g ASC NULLS LAST, fetch=1 +02)--Projection: topk_window_groups.g, count(Int64(1)) ORDER BY [topk_window_groups.g DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND CURRENT ROW AS count(*) ORDER BY [topk_window_groups.g DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND CURRENT ROW AS n +03)----WindowAggr: windowExpr=[[count(Int64(1)) ORDER BY [topk_window_groups.g DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND CURRENT ROW]] +04)------Aggregate: groupBy=[[topk_window_groups.g]], aggr=[[]] +05)--------TableScan: topk_window_groups projection=[g] +physical_plan +01)ProjectionExec: expr=[g@0 as g, count(Int64(1)) ORDER BY [topk_window_groups.g DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND CURRENT ROW@1 as n] +02)--SortExec: TopK(fetch=1), expr=[g@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----BoundedWindowAggExec: wdw=[count(Int64(1)) ORDER BY [topk_window_groups.g DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND CURRENT ROW: Field { "count(Int64(1)) ORDER BY [topk_window_groups.g DESC NULLS FIRST] ROWS BETWEEN 1 PRECEDING AND CURRENT ROW": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------SortExec: expr=[g@0 DESC], preserve_partitioning=[false] +05)--------AggregateExec: mode=Single, gby=[g@0 as g], aggr=[] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +drop table topk_window_groups; + +statement ok +set datafusion.execution.target_partitions = 4;