You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/06/10 16:54:02 UTC
(datafusion) branch gh-readonly-queue/main/pr-22780-a65ab339d907972a2b6a787173d99bd4b17b4824 created (now 82ea3cb268)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-22780-a65ab339d907972a2b6a787173d99bd4b17b4824
in repository https://gitbox.apache.org/repos/asf/datafusion.git
at 82ea3cb268 fix: Correct output-count stats for partitioned partial aggs (#22780)
This branch includes the following new commits:
new 82ea3cb268 fix: Correct output-count stats for partitioned partial aggs (#22780)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org
(datafusion) 01/01: fix: Correct output-count stats for partitioned partial aggs (#22780)
Posted by gi...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-22780-a65ab339d907972a2b6a787173d99bd4b17b4824
in repository https://gitbox.apache.org/repos/asf/datafusion.git
commit 82ea3cb268708b9a078de51e17109f761653e207
Author: Neil Conway <ne...@gmail.com>
AuthorDate: Wed Jun 10 12:50:16 2026 -0400
fix: Correct output-count stats for partitioned partial aggs (#22780)
## Which issue does this PR close?
- Closes #22779
## Rationale for this change
A partial aggregate with no group by expression emits one row per output
partition, even for input partitions that did not receive any rows. The
stats code gets this incorrect, and claims that the partial agg outputs
`Exact(0)` rows in this scenario. This is off by a factor of
`partition_count`, which can lead to suboptimal planning decisions
downstream.
Because `total_byte_size` is scaled from the output row count, this is
also improved: the byte-size estimate now scales with the corrected
output row count rather than the input row count.
## What changes are included in this PR?
* Pass the requested partition into `statistics_inner` and adjust
`statistics_inner` as described above
* Tighten optimizer metadata for grouping sets queries (minor
correctness fix)
* Add/extend existing unit tests
## Are these changes tested?
Yes; new tests added.
## Are there any user-facing changes?
No.
---
.../physical_optimizer/partition_statistics.rs | 37 ++++
datafusion/physical-plan/src/aggregates/mod.rs | 197 +++++++++++++++++++--
2 files changed, 215 insertions(+), 19 deletions(-)
diff --git a/datafusion/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs
index 4fba94ec3a..e1bf22201d 100644
--- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs
+++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs
@@ -982,6 +982,43 @@ mod test {
scan_schema.clone(),
)?);
+ let expect_partial_stat = Statistics {
+ num_rows: Precision::Exact(1),
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![ColumnStatistics::new_unknown()],
+ };
+ assert_eq!(
+ expect_partial_stat,
+ *agg_partial.partition_statistics(Some(0))?
+ );
+ assert_eq!(
+ expect_partial_stat,
+ *agg_partial.partition_statistics(Some(1))?
+ );
+
+ let expect_partial_overall_stat = Statistics {
+ num_rows: Precision::Exact(2),
+ total_byte_size: Precision::Absent,
+ column_statistics: vec![ColumnStatistics::new_unknown()],
+ };
+ assert_eq!(
+ expect_partial_overall_stat,
+ *agg_partial.partition_statistics(None)?
+ );
+
+ // Verify that the partial aggregate emits one accumulator-state row per
+ // output partition, even when the corresponding input partitions are empty.
+ let partitions = execute_stream_partitioned(
+ agg_partial.clone(),
+ Arc::new(TaskContext::default()),
+ )?;
+ assert_eq!(2, partitions.len());
+ for partition_stream in partitions {
+ let result: Vec<RecordBatch> = partition_stream.try_collect().await?;
+ let rows = result.iter().map(|batch| batch.num_rows()).sum::<usize>();
+ assert_eq!(1, rows);
+ }
+
let coalesce = Arc::new(CoalescePartitionsExec::new(agg_partial.clone()));
let agg_final = Arc::new(AggregateExec::try_new(
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index a5f1621812..940bdd41a8 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -873,6 +873,7 @@ impl AggregateExec {
&input,
Arc::clone(&schema),
&group_expr_mapping,
+ group_by.is_true_no_grouping(),
&mode,
&input_order_mode,
aggr_expr.as_ref(),
@@ -1115,6 +1116,7 @@ impl AggregateExec {
input: &Arc<dyn ExecutionPlan>,
schema: SchemaRef,
group_expr_mapping: &ProjectionMapping,
+ is_true_no_grouping: bool,
mode: &AggregateMode,
input_order_mode: &InputOrderMode,
aggr_exprs: &[Arc<AggregateFunctionExpr>],
@@ -1124,9 +1126,12 @@ impl AggregateExec {
.equivalence_properties()
.project(group_expr_mapping, schema);
- // If the group by is empty, then we ensure that the operator will produce
- // only one row, and mark the generated result as a constant value.
- if group_expr_mapping.is_empty() {
+ // True no-group aggregates produce only one row in each output
+ // partition, so aggregate outputs are constants within the partition.
+ // Grouping sets with empty grouping expressions are not covered here:
+ // their output schema can include grouping-set columns before the
+ // aggregate columns, so this aggregate-column mapping does not apply.
+ if is_true_no_grouping {
let new_constants = aggr_exprs.iter().enumerate().map(|(idx, func)| {
let column = Arc::new(Column::new(func.name(), idx));
ConstExpr::from(column as Arc<dyn PhysicalExpr>)
@@ -1185,6 +1190,11 @@ impl AggregateExec {
/// Estimates output statistics for this aggregate node.
///
+ /// For aggregations without group-by expressions, row count follows the
+ /// number of logical aggregate rows and the aggregate output mode. True
+ /// no-group aggregates have one logical row; empty grouping sets have one
+ /// logical row per grouping-set occurrence.
+ ///
/// For grouped aggregations with known input row count > 1, the output row
/// count is estimated as:
///
@@ -1222,7 +1232,11 @@ impl AggregateExec {
/// - Per-set products are summed across all grouping sets
/// - Requires NDV stats for ALL active group-by columns; if any lacks stats,
/// falls back to `input_rows` (or `Absent` if that is also unknown)
- fn statistics_inner(&self, child_statistics: &Statistics) -> Result<Statistics> {
+ fn statistics_inner(
+ &self,
+ child_statistics: &Statistics,
+ partition: Option<usize>,
+ ) -> Result<Statistics> {
// TODO stats: group expressions:
// - once expressions will be able to compute their own stats, use it here
// - case where we group by on a column for which with have the `distinct` stat
@@ -1246,20 +1260,18 @@ impl AggregateExec {
column_statistics
};
- match self.mode {
- AggregateMode::Final | AggregateMode::FinalPartitioned
- if self.group_by.expr.is_empty() =>
- {
+ match self.exact_output_rows_without_group_exprs(partition) {
+ Some(output_rows) => {
let total_byte_size =
- Self::calculate_scaled_byte_size(child_statistics, 1);
+ Self::calculate_scaled_byte_size(child_statistics, output_rows);
Ok(Statistics {
- num_rows: Precision::Exact(1),
+ num_rows: Precision::Exact(output_rows),
column_statistics,
total_byte_size,
})
}
- _ => {
+ None => {
let num_rows = self.estimate_num_rows(child_statistics);
let total_byte_size = num_rows
@@ -1280,6 +1292,51 @@ impl AggregateExec {
}
}
+ /// Exact physical output row count for aggregates without group-by
+ /// expressions.
+ ///
+ /// `partition` follows [`ExecutionPlan::partition_statistics`]: `Some(_)`
+ /// requests one output partition, while `None` requests the entire plan.
+ /// Partial-state output contains the logical rows in each output partition;
+ /// final-value output contains the global logical rows once.
+ /// This mirrors execution, where partial aggregation without group-by
+ /// expressions emits its logical rows from every output partition, including
+ /// empty input partitions.
+ ///
+ /// Returns `None` when grouping expressions are present and grouped
+ /// cardinality estimation should be used instead.
+ fn exact_output_rows_without_group_exprs(
+ &self,
+ partition: Option<usize>,
+ ) -> Option<usize> {
+ let logical_rows = self.logical_rows_without_group_exprs()?;
+
+ Some(match (self.mode.output_mode(), partition) {
+ (AggregateOutputMode::Final, _) => logical_rows,
+ (AggregateOutputMode::Partial, Some(_)) => logical_rows,
+ (AggregateOutputMode::Partial, None) => {
+ logical_rows * self.cache.output_partitioning().partition_count()
+ }
+ })
+ }
+
+ /// Exact number of logical aggregate rows for aggregates without group-by
+ /// expressions.
+ ///
+ /// A true no-group aggregate has one logical aggregate row. Empty grouping
+ /// sets have one logical aggregate row per grouping-set occurrence, even
+ /// when there are duplicate empty grouping sets. Returns `None` when there
+ /// are grouping expressions.
+ fn logical_rows_without_group_exprs(&self) -> Option<usize> {
+ if self.group_by.is_true_no_grouping() {
+ Some(1)
+ } else if self.group_by.expr.is_empty() {
+ Some(self.group_by.groups.len())
+ } else {
+ None
+ }
+ }
+
/// Estimates the output row count for grouped aggregations, combining NDV,
/// input row count, and TopK limit into a single [`Precision<usize>`].
fn estimate_num_rows(&self, child_statistics: &Statistics) -> Precision<usize> {
@@ -1708,7 +1765,9 @@ impl ExecutionPlan for AggregateExec {
fn partition_statistics(&self, partition: Option<usize>) -> Result<Arc<Statistics>> {
let child_statistics = self.input().partition_statistics(partition)?;
- Ok(Arc::new(self.statistics_inner(&child_statistics)?))
+ Ok(Arc::new(
+ self.statistics_inner(&child_statistics, partition)?,
+ ))
}
fn cardinality_effect(&self) -> CardinalityEffect {
@@ -4514,6 +4573,26 @@ mod tests {
let stats_zero = agg_zero.partition_statistics(None)?;
assert_eq!(stats_zero.total_byte_size, Precision::Absent);
+ let single_input =
+ Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc<dyn ExecutionPlan>;
+ let single_agg_zero = AggregateExec::try_new(
+ AggregateMode::Single,
+ PhysicalGroupBy::default(),
+ vec![count_a_aggregate(&schema)?],
+ vec![None],
+ single_input,
+ Arc::clone(&schema),
+ )?;
+ assert_eq!(
+ single_agg_zero
+ .properties()
+ .output_partitioning()
+ .partition_count(),
+ 1
+ );
+ let single_stats_zero = single_agg_zero.partition_statistics(None)?;
+ assert_eq!(single_stats_zero.num_rows, Precision::Exact(1));
+
Ok(())
}
@@ -4522,19 +4601,39 @@ mod tests {
stats: Statistics,
group_by: PhysicalGroupBy,
limit: Option<LimitOptions>,
+ ) -> Result<AggregateExec> {
+ build_test_aggregate_with_mode(
+ schema,
+ stats,
+ group_by,
+ limit,
+ AggregateMode::Final,
+ )
+ }
+
+ fn count_a_aggregate(schema: &SchemaRef) -> Result<Arc<AggregateFunctionExpr>> {
+ Ok(Arc::new(
+ AggregateExprBuilder::new(count_udaf(), vec![col("a", schema)?])
+ .schema(Arc::clone(schema))
+ .alias("COUNT(a)")
+ .build()?,
+ ))
+ }
+
+ fn build_test_aggregate_with_mode(
+ schema: &SchemaRef,
+ stats: Statistics,
+ group_by: PhysicalGroupBy,
+ limit: Option<LimitOptions>,
+ mode: AggregateMode,
) -> Result<AggregateExec> {
let input = Arc::new(StatisticsExec::new(stats, (**schema).clone()))
as Arc<dyn ExecutionPlan>;
let mut agg = AggregateExec::try_new(
- AggregateMode::Final,
+ mode,
group_by,
- vec![Arc::new(
- AggregateExprBuilder::new(count_udaf(), vec![col("a", schema)?])
- .schema(Arc::clone(schema))
- .alias("COUNT(a)")
- .build()?,
- )],
+ vec![count_a_aggregate(schema)?],
vec![None],
input,
Arc::clone(schema),
@@ -4969,6 +5068,66 @@ mod tests {
Ok(())
}
+ #[tokio::test]
+ async fn test_aggregate_stats_duplicate_empty_grouping_sets() -> Result<()> {
+ let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)]));
+
+ let duplicate_empty_grouping_sets =
+ PhysicalGroupBy::new(vec![], vec![], vec![vec![], vec![]], true);
+
+ let single_input =
+ Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc<dyn ExecutionPlan>;
+ let single_agg = AggregateExec::try_new(
+ AggregateMode::Single,
+ duplicate_empty_grouping_sets.clone(),
+ vec![count_a_aggregate(&schema)?],
+ vec![None],
+ single_input,
+ Arc::clone(&schema),
+ )?;
+ assert_eq!(
+ single_agg.partition_statistics(None)?.num_rows,
+ Precision::Exact(2)
+ );
+
+ let partial_input =
+ Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(2))
+ as Arc<dyn ExecutionPlan>;
+ let partial_agg = Arc::new(AggregateExec::try_new(
+ AggregateMode::Partial,
+ duplicate_empty_grouping_sets,
+ vec![count_a_aggregate(&schema)?],
+ vec![None],
+ partial_input,
+ Arc::clone(&schema),
+ )?);
+
+ assert_eq!(
+ partial_agg
+ .properties()
+ .output_partitioning()
+ .partition_count(),
+ 2
+ );
+ let task_ctx = Arc::new(TaskContext::default());
+ for partition in 0..2 {
+ assert_eq!(
+ partial_agg.partition_statistics(Some(partition))?.num_rows,
+ Precision::Exact(2)
+ );
+ let result =
+ collect(partial_agg.execute(partition, Arc::clone(&task_ctx))?).await?;
+ assert_eq!(result.iter().map(RecordBatch::num_rows).sum::<usize>(), 2);
+ }
+
+ assert_eq!(
+ partial_agg.partition_statistics(None)?.num_rows,
+ Precision::Exact(4)
+ );
+
+ Ok(())
+ }
+
#[test]
fn test_aggregate_stats_non_column_expr_bails_out() -> Result<()> {
use datafusion_common::ColumnStatistics;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org