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/09/23 00:56:38 UTC
(datafusion) branch main updated: refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) (#25564)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new ebdb657777 refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) (#25564)
ebdb657777 is described below
commit ebdb657777d512b20ad40f50ce8fb85a39ee3175
Author: Yongting You <20...@gmail.com>
AuthorDate: Wed Sep 23 00:47:30 2026 +0000
refactor: Make `AggregateExec` state modeling and public updates safe (part 1: distinct limit) (#25564)
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->
- Part of https://github.com/apache/datafusion/issues/25563
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
Please explain the problem you are trying to solve in terms of the
user-visible
behavior, rather than the implementation.
For example, "The code in `foo.rs` doesn't handle nulls" is a symptom of
the
implementation. "COUNT(DISTINCT) returns wrong results when the column
contains
nulls" is the user-visible problem.
-->
See issue for the idea, this PR is the step 1 in the issue.
For backgrounds for distinct limit optimization (this PR targets), see
https://github.com/apache/datafusion/blob/1e09a2a3d76c03f469e80bed712a6bada45e77c8/datafusion/physical-plan/src/aggregates/hash_stream.rs#L77
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here, but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->
1. Refactor AggregateExec with enum state
2. Add safe API for this optimization, and use it in the optimizer.
(unsafe one will be able to be deleted after step2, now it's reused by
other path)
I'll mark the structural changes (only 2 files)
Most diff is mechanical changes like `aggregate_exec.aggr_expr` ->
`aggregate_exec.aggr_expr()`, caused by the refactor.
## What is the testing strategy for this PR?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
4. Serve as another way to document the expected behavior of the code
Briefly describe how this PR is tested, and point to the specific tests
you added. For example: 'This new feature is covered by the
`sqllogictest` cases added in `foo.slt`'.
If this PR does not add tests, explain why. For example, if the change
is already covered by existing tests, please mention it.
You should also check the `codecov` bot reply on this PR to confirm the
changed code is exercised.
-->
Existing tests.
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be
updated before approving the PR.
If there are any breaking changes to public APIs, please add the `api
change` label.
-->
no
---
.../src/limited_distinct_aggregation.rs | 24 +-
.../src/aggregates/aggregate_hash_table/common.rs | 10 +-
.../aggregate_hash_table/common_ordered.rs | 10 +-
.../aggregates/aggregate_hash_table/final_table.rs | 2 +-
.../src/aggregates/aggregate_hash_table/mod.rs | 2 +-
.../aggregate_hash_table/ordered_final_table.rs | 2 +-
.../aggregate_hash_table/ordered_partial_table.rs | 2 +-
.../aggregate_hash_table/ordered_single_table.rs | 2 +-
.../aggregate_hash_table/partial_reduce_table.rs | 2 +-
.../aggregate_hash_table/partial_table.rs | 2 +-
.../aggregate_hash_table/single_table.rs | 2 +-
.../src/aggregates/aggregate_stream.rs | 10 +-
.../src/aggregates/grouped_hash_stream.rs | 14 +-
.../src/aggregates/grouped_topk_stream.rs | 6 +-
.../physical-plan/src/aggregates/hash_stream.rs | 2 +-
datafusion/physical-plan/src/aggregates/mod.rs | 322 +++++++++++++--------
.../src/aggregates/ordered_single_stream.rs | 4 +-
.../physical-plan/src/aggregates/single_stream.rs | 4 +-
datafusion/physical-plan/src/aggregates/spill.rs | 6 +-
19 files changed, 259 insertions(+), 169 deletions(-)
diff --git a/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs b/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs
index 192a139f36..60833e0401 100644
--- a/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs
+++ b/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs
@@ -20,7 +20,7 @@
use std::sync::Arc;
-use datafusion_physical_plan::aggregates::{AggregateExec, LimitOptions};
+use datafusion_physical_plan::aggregates::AggregateExec;
use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
@@ -47,16 +47,10 @@ impl LimitedDistinctAggregation {
fn transform_agg(
aggr: &AggregateExec,
limit: usize,
- ) -> Option<Arc<dyn ExecutionPlan>> {
- // rules for transforming this Aggregate are held in this method
- if !aggr.is_unordered_unfiltered_group_by_distinct() {
- return None;
- }
-
- // We found what we want: clone, copy the limit down, and return modified node
- let new_aggr = aggr.with_new_limit_options(Some(LimitOptions::new(limit)));
-
- Some(Arc::new(new_aggr))
+ ) -> Option<Transformed<Arc<dyn ExecutionPlan>>> {
+ let new_aggr = aggr.clone().try_optimize_distinct_soft_limit(limit)?;
+ // An already limited aggregate still permits optimizing its partial child.
+ Some(new_aggr.update_data(|aggr| Arc::new(aggr) as Arc<dyn ExecutionPlan>))
}
/// transform_limit matches an `AggregateExec` as the child of a `LocalLimitExec`
@@ -119,14 +113,18 @@ impl LimitedDistinctAggregation {
Some(new_aggr) => {
match_aggr = plan;
found_match_aggr = true;
- return Ok(Transformed::yes(new_aggr));
+ return Ok(new_aggr);
}
}
}
rewrite_applicable = false;
Ok(Transformed::no(plan))
};
- let child = child.to_owned().transform_down(closure).data().ok()?;
+ let child = child.to_owned().transform_down(closure).ok()?;
+ if !child.transformed {
+ return None;
+ }
+ let child = child.data;
if is_global_limit {
return Some(Arc::new(GlobalLimitExec::new(
child,
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
index 0c5c4185e2..b25caf815e 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs
@@ -156,12 +156,12 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
let input_schema = agg.input().schema();
let metrics = AggregateTableMetrics::new(agg, partition);
let aggregate_arguments = aggregate_expressions(
- &agg.aggr_expr,
+ agg.aggr_expr(),
&agg.mode,
- agg.group_by.num_group_exprs(),
+ agg.group_by().num_group_exprs(),
)?;
let accumulators: Vec<_> = agg
- .aggr_expr
+ .aggr_expr()
.iter()
.zip(aggregate_arguments)
.zip(filters)
@@ -179,7 +179,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
})
.collect::<Result<_>>()?;
- let group_schema = agg.group_by.group_schema(&input_schema)?;
+ let group_schema = agg.group_by().group_schema(&input_schema)?;
let group_values = new_group_values(group_schema, &GroupOrdering::None)?;
Ok(Self {
@@ -192,7 +192,7 @@ impl<AggrMode> AggregateHashTable<AggrMode> {
state_schema,
batch_size,
state: AggregateHashTableState::Building(AggregateHashTableBuffer {
- group_by: Arc::clone(&agg.group_by),
+ group_by: Arc::clone(agg.group_by()),
group_values,
batch_group_indices: Default::default(),
accumulators,
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs
index 9476c0991a..fd9e92635f 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs
@@ -200,15 +200,15 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
);
let group_ordering = GroupOrdering::try_new(input_order_mode)?;
- let group_schema = agg.group_by.group_schema(input_schema)?;
+ let group_schema = agg.group_by().group_schema(input_schema)?;
let group_values = new_group_values(group_schema, &group_ordering)?;
let aggregate_arguments = aggregate_expressions(
- &agg.aggr_expr,
+ agg.aggr_expr(),
aggregate_mode,
- agg.group_by.num_group_exprs(),
+ agg.group_by().num_group_exprs(),
)?;
let accumulators = agg
- .aggr_expr
+ .aggr_expr()
.iter()
.zip(aggregate_arguments)
.zip(filters)
@@ -235,7 +235,7 @@ impl<AggrMode> OrderedAggregateTable<AggrMode> {
aggregate_accumulator_metrics: metrics.accumulator,
aggregate_submetrics: metrics.submetrics,
buffer: OrderedAggregateTableBuffer {
- group_by: Arc::clone(&agg.group_by),
+ group_by: Arc::clone(agg.group_by()),
group_ordering,
group_values,
group_indices: vec![],
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
index 307215cfd2..4efd2a1c3a 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs
@@ -46,7 +46,7 @@ impl AggregateHashTable<FinalMarker> {
output_schema,
Arc::clone(&agg.input().schema()),
batch_size,
- vec![None; agg.aggr_expr.len()],
+ vec![None; agg.aggr_expr().len()],
)
}
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs
index 005b2bac09..960b7498f1 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs
@@ -70,7 +70,7 @@ pub(super) struct AggregateTableMetrics {
impl AggregateTableMetrics {
pub(super) fn new(agg: &AggregateExec, partition: usize) -> Self {
let aggregate_labels = agg
- .aggr_expr
+ .aggr_expr()
.iter()
.map(|agg_expr| aggregate_metric_label(agg_expr))
.collect::<Vec<_>>();
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs
index d0d0c99bb5..e4b8118f94 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs
@@ -58,7 +58,7 @@ impl OrderedAggregateTable<FinalMarker> {
batch_size,
input_order_mode,
&AggregateMode::Final,
- vec![None; agg.aggr_expr.len()],
+ vec![None; agg.aggr_expr().len()],
metrics,
)
}
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs
index b53012d8bd..104855c5f0 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs
@@ -70,7 +70,7 @@ impl OrderedAggregateTable<PartialMarker> {
batch_size,
&agg.input_order_mode,
&AggregateMode::Partial,
- agg.filter_expr.iter().cloned().collect(),
+ agg.filter_expr().to_vec(),
metrics,
)
}
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs
index ce1ce647b4..dd266a68eb 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_single_table.rs
@@ -61,7 +61,7 @@ impl OrderedAggregateTable<SingleMarker> {
batch_size,
&agg.input_order_mode,
&agg.mode,
- agg.filter_expr.iter().cloned().collect(),
+ agg.filter_expr().to_vec(),
metrics,
)
}
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
index 2892d05933..fc59f7e92a 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs
@@ -40,7 +40,7 @@ impl AggregateHashTable<PartialReduceMarker> {
Arc::clone(&output_schema),
output_schema,
batch_size,
- vec![None; agg.aggr_expr.len()],
+ vec![None; agg.aggr_expr().len()],
)
}
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs
index 90f8193034..397b766f41 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs
@@ -51,7 +51,7 @@ impl AggregateHashTable<PartialMarker> {
Arc::clone(&output_schema),
output_schema,
batch_size,
- agg.filter_expr.iter().cloned().collect(),
+ agg.filter_expr().to_vec(),
)
}
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs
index 2d7dc2a63d..beaf098035 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs
@@ -45,7 +45,7 @@ impl AggregateHashTable<SingleMarker> {
output_schema,
state_schema,
batch_size,
- agg.filter_expr.iter().cloned().collect(),
+ agg.filter_expr().to_vec(),
)
}
diff --git a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs
index 3ce7915806..0da8f0944e 100644
--- a/datafusion/physical-plan/src/aggregates/aggregate_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs
@@ -295,23 +295,23 @@ impl AggregateStream {
partition: usize,
) -> Result<Self> {
let agg_schema = Arc::clone(&agg.schema);
- let agg_filter_expr = Arc::clone(&agg.filter_expr);
+ let agg_filter_expr = agg.clone_filter_exprs();
let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
let input = agg.input.execute(partition, Arc::clone(context))?;
- let aggregate_expressions = aggregate_expressions(&agg.aggr_expr, &agg.mode, 0)?;
+ let aggregate_expressions = aggregate_expressions(agg.aggr_expr(), &agg.mode, 0)?;
let filter_expressions = match agg.mode.input_mode() {
AggregateInputMode::Raw => agg_filter_expr,
- AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(),
+ AggregateInputMode::Partial => vec![None; agg.aggr_expr().len()].into(),
};
let aggregate_labels = agg
- .aggr_expr
+ .aggr_expr()
.iter()
.map(|agg_expr| aggregate_metric_label(agg_expr))
.collect::<Vec<_>>();
let accumulators = create_accumulators_with_metrics(
- &agg.aggr_expr,
+ agg.aggr_expr(),
&aggregate_sub_metrics(
&agg.metrics,
partition,
diff --git a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
index 042f4a9044..b9a6d2833a 100644
--- a/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs
@@ -400,15 +400,15 @@ impl GroupedHashAggregateStream {
) -> Result<Self> {
debug!("Creating GroupedHashAggregateStream");
let agg_schema = Arc::clone(&agg.schema);
- let agg_group_by = Arc::clone(&agg.group_by);
- let agg_filter_expr = Arc::clone(&agg.filter_expr);
+ let agg_group_by = Arc::clone(agg.group_by());
+ let agg_filter_expr = agg.clone_filter_exprs();
let batch_size = context.session_config().batch_size();
let input = agg.input.execute(partition, Arc::clone(context))?;
let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition);
let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition);
let aggregate_labels = agg
- .aggr_expr
+ .aggr_expr()
.iter()
.map(|agg_expr| aggregate_metric_label(agg_expr))
.collect::<Vec<_>>();
@@ -431,25 +431,25 @@ impl GroupedHashAggregateStream {
let timer = baseline_metrics.elapsed_compute().timer();
- let aggregate_exprs = Arc::clone(&agg.aggr_expr);
+ let aggregate_exprs = agg.clone_aggr_exprs();
// arguments for each aggregate, one vec of expressions per
// aggregate
let aggregate_arguments = aggregates::aggregate_expressions(
- &agg.aggr_expr,
+ agg.aggr_expr(),
&agg.mode,
agg_group_by.num_group_exprs(),
)?;
// arguments for aggregating spilled data is the same as the one for final aggregation
let merging_aggregate_arguments = aggregates::aggregate_expressions(
- &agg.aggr_expr,
+ agg.aggr_expr(),
&AggregateMode::Final,
agg_group_by.num_group_exprs(),
)?;
let filter_expressions = match agg.mode.input_mode() {
AggregateInputMode::Raw => agg_filter_expr,
- AggregateInputMode::Partial => vec![None; agg.aggr_expr.len()].into(),
+ AggregateInputMode::Partial => vec![None; agg.aggr_expr().len()].into(),
};
// Instantiate the accumulators
diff --git a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs
index 13ead73930..f4e719a9fc 100644
--- a/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs
@@ -71,19 +71,19 @@ impl GroupedTopKAggregateStream {
limit: usize,
) -> Result<Self> {
let agg_schema = Arc::clone(&aggr.schema);
- let group_by = Arc::clone(&aggr.group_by);
+ let group_by = Arc::clone(aggr.group_by());
let input = aggr.input.execute(partition, Arc::clone(context))?;
let baseline_metrics = BaselineMetrics::new(&aggr.metrics, partition);
let group_by_metrics = GroupByMetrics::new_topk(&aggr.metrics, partition);
let aggregate_argument_metrics = AggregateArgumentMetrics::new(
&aggr.metrics,
partition,
- aggr.aggr_expr
+ aggr.aggr_expr()
.iter()
.map(|agg_expr| aggregate_metric_label(agg_expr)),
);
let aggregate_arguments =
- aggregate_expressions(&aggr.aggr_expr, &aggr.mode, group_by.expr.len())?;
+ aggregate_expressions(aggr.aggr_expr(), &aggr.mode, group_by.expr.len())?;
let (expr, _) = &aggr.group_expr().expr()[0];
let kt = expr.data_type(&aggr.input().schema())?;
diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs
index 43257acc14..2ef7a6463f 100644
--- a/datafusion/physical-plan/src/aggregates/hash_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs
@@ -238,7 +238,7 @@ impl PartialHashAggregateStream {
Arc::clone(&schema),
batch_size,
)?;
- let skip_aggregation_probe = if agg.group_by.is_single() {
+ let skip_aggregation_probe = if agg.group_by().is_single() {
let options = &context.session_config().options().execution;
let probe_ratio_threshold =
options.skip_partial_aggregation_probe_ratio_threshold;
diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs
index 5fa350b55d..3f17cacb52 100644
--- a/datafusion/physical-plan/src/aggregates/mod.rs
+++ b/datafusion/physical-plan/src/aggregates/mod.rs
@@ -185,7 +185,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef};
use arrow::record_batch::RecordBatch;
use arrow_schema::FieldRef;
use datafusion_common::stats::Precision;
-use datafusion_common::tree_node::TreeNodeRecursion;
+use datafusion_common::tree_node::{Transformed, TreeNodeRecursion};
use datafusion_common::{
ColumnStatistics, Constraint, Constraints, Result, ScalarValue,
assert_eq_or_internal_err, internal_err, not_impl_err,
@@ -846,35 +846,30 @@ impl LimitOptions {
}
}
+/// Aggregation state, separating a DISTINCT soft limit from accumulators and filters.
+#[derive(Debug, Clone)]
+enum AggregateKind {
+ /// Ordinary aggregation, including the existing Top-K configuration.
+ General {
+ group_by: Arc<PhysicalGroupBy>,
+ aggr_expr: Arc<[Arc<AggregateFunctionExpr>]>,
+ filter_expr: Arc<[Option<Arc<dyn PhysicalExpr>>]>,
+ limit_options: Option<LimitOptions>,
+ },
+ /// `SELECT DISTINCT k FROM t LIMIT n`: eligible streams may stop after n groups.
+ /// Other streams consume all input; the parent LIMIT enforces the row count.
+ DistinctLimit {
+ group_by: Arc<PhysicalGroupBy>,
+ limit: usize,
+ },
+}
+
/// Hash aggregate execution plan
#[derive(Debug, Clone)]
pub struct AggregateExec {
/// Aggregation mode (full, partial)
mode: AggregateMode,
- /// Group by expressions
- /// [`Arc`] used for a cheap clone, which improves physical plan optimization performance.
- group_by: Arc<PhysicalGroupBy>,
- /// Aggregate expressions
- /// The same reason to [`Arc`] it as for [`Self::group_by`].
- aggr_expr: Arc<[Arc<AggregateFunctionExpr>]>,
- /// FILTER (WHERE clause) expression for each aggregate expression
- /// The same reason to [`Arc`] it as for [`Self::group_by`].
- filter_expr: Arc<[Option<Arc<dyn PhysicalExpr>>]>,
- /// Soft limit for best-effort, limit-based optimizations.
- ///
- /// `AggregateExec` has multiple stream implementations and not all of them
- /// support the optimization, so this is only a hint. The downstream limit
- /// operator enforces the exact row count either way.
- ///
- /// Supported by:
- /// - [`StreamType::GroupedPriorityQueue`]: retains only the best `limit`
- /// groups per partition (this stream is selected only when a limit is set)
- /// - [`StreamType::SingleHash`], [`StreamType::PartialHash`], [`StreamType::FinalHash`]
- /// and the legacy [`StreamType::GroupedHash`]: stop reading input once `limit` groups
- /// have been accumulated
- ///
- /// The remaining streams consume all input.
- limit_options: Option<LimitOptions>,
+ kind: AggregateKind,
/// Input plan, could be a partial aggregate or the input to the aggregate
pub input: Arc<dyn ExecutionPlan>,
/// Schema after the aggregate is applied. Contains the group by columns followed by the
@@ -906,6 +901,46 @@ pub struct AggregateExec {
}
impl AggregateExec {
+ /// Try to stop after enough distinct grouping keys have been accumulated.
+ ///
+ /// The caller must establish that discarding other groups is legal, for
+ /// example beneath `SELECT DISTINCT k FROM t LIMIT 10`. The parent limit
+ /// remains responsible for enforcing the exact number of rows.
+ ///
+ /// # Trigger conditions
+ ///
+ /// - Grouping with no aggregate expressions or aggregate filters.
+ /// - Ordering eligible under [`Self::is_unordered_unfiltered_group_by_distinct`],
+ /// including no existing Top-K direction.
+ /// - A positive limit that is tighter than any existing limit.
+ ///
+ /// # Consistency
+ ///
+ /// This is a safe, atomic optimization: it preserves the grouping keys and
+ /// leaves the aggregate in a consistent state. Ineligible aggregates return
+ /// `None`. Eligible aggregates are returned unchanged when the limit is zero
+ /// or no tighter than the existing limit.
+ pub fn try_optimize_distinct_soft_limit(
+ mut self,
+ limit: usize,
+ ) -> Option<Transformed<Self>> {
+ if !self.is_unordered_unfiltered_group_by_distinct() {
+ return None;
+ }
+ if limit == 0
+ || self
+ .limit_options()
+ .is_some_and(|existing| existing.limit <= limit)
+ {
+ return Some(Transformed::no(self));
+ }
+ self.kind = AggregateKind::DistinctLimit {
+ group_by: Arc::clone(self.group_by()),
+ limit,
+ };
+ Some(Transformed::yes(self))
+ }
+
/// Function used in `OptimizeAggregateOrder` optimizer rule,
/// where we need parts of the new value, others cloned from the old one
/// Rewrites aggregate exec with new aggregate expressions.
@@ -913,42 +948,31 @@ impl AggregateExec {
&self,
aggr_expr: impl Into<Arc<[Arc<AggregateFunctionExpr>]>>,
) -> Self {
- Self {
- aggr_expr: aggr_expr.into(),
- // clone the rest of the fields
- required_input_ordering: self.required_input_ordering.clone(),
- metrics: ExecutionPlanMetricsSet::new(),
- input_order_mode: self.input_order_mode.clone(),
- cache: Arc::clone(&self.cache),
- mode: self.mode,
- group_by: Arc::clone(&self.group_by),
- filter_expr: Arc::clone(&self.filter_expr),
- limit_options: self.limit_options,
- input: Arc::clone(&self.input),
- schema: Arc::clone(&self.schema),
- input_schema: Arc::clone(&self.input_schema),
- dynamic_filter: self.dynamic_filter.clone(),
+ let aggr_expr = aggr_expr.into();
+ let mut new = self.clone();
+ match &mut new.kind {
+ AggregateKind::General { aggr_expr: old, .. } => *old = aggr_expr,
+ AggregateKind::DistinctLimit { .. } if aggr_expr.is_empty() => {}
+ AggregateKind::DistinctLimit { group_by, .. } => {
+ // An accumulator rewrite cannot inherit DISTINCT's early stop.
+ new.kind = AggregateKind::General {
+ group_by: Arc::clone(group_by),
+ // The previous `DistinctLimit` type doesn't include filter
+ filter_expr: vec![None; aggr_expr.len()].into(),
+ aggr_expr,
+ limit_options: None,
+ };
+ }
}
+ new.metrics = ExecutionPlanMetricsSet::new();
+ new
}
/// Clone this exec, overriding only the limit hint.
pub fn with_new_limit_options(&self, limit_options: Option<LimitOptions>) -> Self {
- Self {
- limit_options,
- // clone the rest of the fields
- required_input_ordering: self.required_input_ordering.clone(),
- metrics: ExecutionPlanMetricsSet::new(),
- input_order_mode: self.input_order_mode.clone(),
- cache: Arc::clone(&self.cache),
- mode: self.mode,
- group_by: Arc::clone(&self.group_by),
- aggr_expr: Arc::clone(&self.aggr_expr),
- filter_expr: Arc::clone(&self.filter_expr),
- input: Arc::clone(&self.input),
- schema: Arc::clone(&self.schema),
- input_schema: Arc::clone(&self.input_schema),
- dynamic_filter: self.dynamic_filter.clone(),
- }
+ let mut new = self.clone().with_limit_options(limit_options);
+ new.metrics = ExecutionPlanMetricsSet::new();
+ new
}
pub fn cache(&self) -> &PlanProperties {
@@ -1092,15 +1116,17 @@ impl AggregateExec {
let mut exec = AggregateExec {
mode,
- group_by,
- aggr_expr: aggr_expr.into(),
- filter_expr,
+ kind: AggregateKind::General {
+ group_by,
+ aggr_expr: aggr_expr.into(),
+ filter_expr,
+ limit_options: None,
+ },
input,
schema,
input_schema,
metrics: ExecutionPlanMetricsSet::new(),
required_input_ordering,
- limit_options: None,
input_order_mode,
cache: Arc::new(cache),
dynamic_filter: None,
@@ -1118,33 +1144,98 @@ impl AggregateExec {
/// Set the limit options for this AggExec
pub fn with_limit_options(mut self, limit_options: Option<LimitOptions>) -> Self {
- self.limit_options = limit_options;
+ // Restoring an existing hint must not depend on input ordering: a later
+ // optimizer may have sorted the child since the hint was introduced.
+ if let Some(options) = limit_options
+ && options.descending.is_none()
+ && !self.group_by().is_true_no_grouping()
+ && self.aggr_expr().is_empty()
+ && self.filter_expr().is_empty()
+ {
+ self.kind = AggregateKind::DistinctLimit {
+ group_by: Arc::clone(self.group_by()),
+ limit: options.limit,
+ };
+ } else {
+ match &mut self.kind {
+ AggregateKind::General {
+ limit_options: old, ..
+ } => *old = limit_options,
+ AggregateKind::DistinctLimit { group_by, .. } => {
+ self.kind = AggregateKind::General {
+ group_by: Arc::clone(group_by),
+ aggr_expr: Arc::from([]),
+ filter_expr: Arc::from([]),
+ limit_options,
+ };
+ }
+ }
+ }
self
}
/// Get the limit options (if set)
pub fn limit_options(&self) -> Option<LimitOptions> {
- self.limit_options
+ match &self.kind {
+ AggregateKind::General { limit_options, .. } => *limit_options,
+ AggregateKind::DistinctLimit { limit, .. } => Some(LimitOptions::new(*limit)),
+ }
}
/// Grouping expressions
pub fn group_expr(&self) -> &PhysicalGroupBy {
- &self.group_by
+ self.group_by()
+ }
+
+ // Keep internal plan and stream clones cheap without exposing Arc in the public API.
+ fn group_by(&self) -> &Arc<PhysicalGroupBy> {
+ match &self.kind {
+ AggregateKind::General { group_by, .. }
+ | AggregateKind::DistinctLimit { group_by, .. } => group_by,
+ }
+ }
+
+ // Spill streams adapt a cloned descriptor to their partial-state input.
+ fn group_by_mut(&mut self) -> &mut Arc<PhysicalGroupBy> {
+ match &mut self.kind {
+ AggregateKind::General { group_by, .. }
+ | AggregateKind::DistinctLimit { group_by, .. } => group_by,
+ }
}
/// Grouping expressions as they occur in the output schema
pub fn output_group_expr(&self) -> Vec<Arc<dyn PhysicalExpr>> {
- self.group_by.output_exprs()
+ self.group_by().output_exprs()
}
/// Aggregate expressions
pub fn aggr_expr(&self) -> &[Arc<AggregateFunctionExpr>] {
- &self.aggr_expr
+ match &self.kind {
+ AggregateKind::General { aggr_expr, .. } => aggr_expr,
+ AggregateKind::DistinctLimit { .. } => &[],
+ }
}
/// FILTER (WHERE clause) expression for each aggregate expression
pub fn filter_expr(&self) -> &[Option<Arc<dyn PhysicalExpr>>] {
- &self.filter_expr
+ match &self.kind {
+ AggregateKind::General { filter_expr, .. } => filter_expr,
+ AggregateKind::DistinctLimit { .. } => &[],
+ }
+ }
+
+ fn clone_aggr_exprs(&self) -> Arc<[Arc<AggregateFunctionExpr>]> {
+ match &self.kind {
+ AggregateKind::General { aggr_expr, .. } => Arc::clone(aggr_expr),
+ AggregateKind::DistinctLimit { .. } => Arc::from([]),
+ }
+ }
+
+ fn clone_filter_exprs(&self) -> Arc<[Option<Arc<dyn PhysicalExpr>>]> {
+ match &self.kind {
+ AggregateKind::General { filter_expr, .. } => Arc::clone(filter_expr),
+ AggregateKind::DistinctLimit { .. } => Arc::from([]),
+ }
}
/// Returns the dynamic filter expression for this aggregate, if set.
@@ -1225,7 +1316,7 @@ impl AggregateExec {
partition: usize,
context: &Arc<TaskContext>,
) -> Result<StreamType> {
- if self.group_by.is_true_no_grouping() {
+ if self.group_by().is_true_no_grouping() {
return Ok(StreamType::AggregateStream(AggregateStream::new(
self, context, partition,
)?));
@@ -1248,7 +1339,7 @@ impl AggregateExec {
// grouped streams is correct in that case, because they treat the limit
// as a soft limit and the `LIMIT` above this aggregate still truncates
// the result.
- if let Some(config) = self.limit_options
+ if let Some(config) = self.limit_options()
&& !self.is_unordered_unfiltered_group_by_distinct()
&& (config.descending.is_some() || self.get_minmax_desc().is_some())
{
@@ -1276,7 +1367,7 @@ impl AggregateExec {
// consuming partial state must receive the expanded keys as a plain
// group by (see `PhysicalGroupBy::as_final`).
if self.mode.input_mode() == AggregateInputMode::Partial
- && !self.group_by.is_single()
+ && !self.group_by().is_single()
{
return internal_err!(
"Grouping sets must be expanded before {:?} aggregation, which consumes partial state",
@@ -1331,7 +1422,7 @@ impl AggregateExec {
/// Finds the DataType and SortDirection for this Aggregate, if there is one
pub fn get_minmax_desc(&self) -> Option<(FieldRef, bool)> {
- let agg_expr = self.aggr_expr.iter().exactly_one().ok()?;
+ let agg_expr = self.aggr_expr().iter().exactly_one().ok()?;
agg_expr.get_minmax_desc()
}
@@ -1533,7 +1624,7 @@ impl AggregateExec {
// self.schema: [<group by exprs>, <aggregate exprs>]
let mut column_statistics = Statistics::unknown_column(&self.schema());
- for (idx, (expr, _)) in self.group_by.expr.iter().enumerate() {
+ for (idx, (expr, _)) in self.group_by().expr.iter().enumerate() {
if let Some(col) = expr.downcast_ref::<Column>() {
let child_col_stats =
&child_statistics.column_statistics[col.index()];
@@ -1624,7 +1715,7 @@ impl AggregateExec {
/// `ROLLUP` and `CUBE` introduce alongside the non-empty ones.
fn output_rows_for_empty_input(&self, partition: Option<usize>) -> usize {
let empty_grouping_sets = self
- .group_by
+ .group_by()
.groups
.iter()
.filter(|nulls| nulls.iter().all(|is_null| *is_null))
@@ -1659,7 +1750,7 @@ impl AggregateExec {
let schema = self.schema();
for (idx, column_stats) in column_statistics
.iter_mut()
- .take(self.group_by.expr.len())
+ .take(self.group_by().expr.len())
.enumerate()
{
let typed_null = ScalarValue::try_from(schema.field(idx).data_type())
@@ -1685,10 +1776,10 @@ impl AggregateExec {
/// 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() {
+ 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 if self.group_by().expr.is_empty() {
+ Some(self.group_by().groups.len())
} else {
None
}
@@ -1701,12 +1792,12 @@ impl AggregateExec {
child_statistics: &Statistics,
partition: Option<usize>,
) -> Precision<usize> {
- let ndv = if !self.group_by.expr.is_empty() {
+ let ndv = if !self.group_by().expr.is_empty() {
self.compute_group_ndv(child_statistics)
} else {
None
};
- let limit = self.limit_options.as_ref().map(|lo| lo.limit);
+ let limit = self.limit_options().as_ref().map(|lo| lo.limit);
if let Some(&value) = child_statistics.num_rows.get_value() {
if value > 1 {
@@ -1725,7 +1816,7 @@ impl AggregateExec {
.num_rows
.map(|_| self.output_rows_for_empty_input(partition))
} else {
- let grouping_set_num = self.group_by.groups.len();
+ let grouping_set_num = self.group_by().groups.len();
let mut num_rows =
child_statistics.num_rows.map(|x| x * grouping_set_num);
if let Some(limit) = limit {
@@ -1762,9 +1853,9 @@ impl AggregateExec {
/// → total = 100 + 50 + 5,000 = 5,150
fn compute_group_ndv(&self, child_statistics: &Statistics) -> Option<usize> {
let mut total: usize = 0;
- for group_mask in &self.group_by.groups {
+ for group_mask in &self.group_by().groups {
let mut set_product: usize = 1;
- for (j, (expr, _)) in self.group_by.expr.iter().enumerate() {
+ for (j, (expr, _)) in self.group_by().expr.iter().enumerate() {
if group_mask[j] {
continue;
}
@@ -1787,7 +1878,7 @@ impl AggregateExec {
/// - If yes, init one inside `AggregateExec`'s `dynamic_filter` field.
/// - If not supported, `self.dynamic_filter` should be kept `None`
fn init_dynamic_filter(&mut self) {
- if (!self.group_by.is_empty()) || (self.mode != AggregateMode::Partial) {
+ if (!self.group_by().is_empty()) || (self.mode != AggregateMode::Partial) {
debug_assert!(
self.dynamic_filter.is_none(),
"The current operator node does not support dynamic filter"
@@ -1808,7 +1899,7 @@ impl AggregateExec {
// filter, and it's used to decide if this dynamic filter is able to get push
// through certain node during optimization.
let mut all_cols: Vec<Arc<dyn PhysicalExpr>> = Vec::new();
- for (i, aggr_expr) in self.aggr_expr.iter().enumerate() {
+ for (i, aggr_expr) in self.aggr_expr().iter().enumerate() {
// 1. Only `min` or `max` aggregate function
let fun_name = aggr_expr.fun().name();
// HACK: Should check the function type more precisely
@@ -1857,7 +1948,7 @@ impl AggregateExec {
.filter_map(|info| {
// This should always be true due to how the supported accumulators
// are constructed. See `init_dynamic_filter` for more details.
- if let [arg] = &self.aggr_expr[info.aggr_index].expressions().as_slice()
+ if let [arg] = &self.aggr_expr()[info.aggr_index].expressions().as_slice()
&& arg.is::<Column>()
{
return Some(Arc::clone(arg));
@@ -1913,14 +2004,14 @@ impl DisplayAs for AggregateExec {
};
write!(f, "AggregateExec: mode={:?}", self.mode)?;
- let g: Vec<String> = if self.group_by.is_single() {
- self.group_by
+ let g: Vec<String> = if self.group_by().is_single() {
+ self.group_by()
.expr
.iter()
.map(format_expr_with_alias)
.collect()
} else {
- self.group_by
+ self.group_by()
.groups
.iter()
.map(|group| {
@@ -1930,10 +2021,10 @@ impl DisplayAs for AggregateExec {
.map(|(idx, is_null)| {
if *is_null {
format_expr_with_alias(
- &self.group_by.null_expr[idx],
+ &self.group_by().null_expr[idx],
)
} else {
- format_expr_with_alias(&self.group_by.expr[idx])
+ format_expr_with_alias(&self.group_by().expr[idx])
}
})
.collect::<Vec<String>>()
@@ -1946,12 +2037,12 @@ impl DisplayAs for AggregateExec {
write!(f, ", gby=[{}]", g.join(", "))?;
let a: Vec<String> = self
- .aggr_expr
+ .aggr_expr()
.iter()
.map(|agg| format_aggregate_exec_expr(agg).to_string())
.collect();
write!(f, ", aggr=[{}]", a.join(", "))?;
- if let Some(config) = self.limit_options {
+ if let Some(config) = self.limit_options() {
write!(f, ", lim=[{}]", config.limit)?;
}
@@ -1970,14 +2061,14 @@ impl DisplayAs for AggregateExec {
}
};
- let g: Vec<String> = if self.group_by.is_single() {
- self.group_by
+ let g: Vec<String> = if self.group_by().is_single() {
+ self.group_by()
.expr
.iter()
.map(format_expr_with_alias)
.collect()
} else {
- self.group_by
+ self.group_by()
.groups
.iter()
.map(|group| {
@@ -1987,10 +2078,10 @@ impl DisplayAs for AggregateExec {
.map(|(idx, is_null)| {
if *is_null {
format_expr_with_alias(
- &self.group_by.null_expr[idx],
+ &self.group_by().null_expr[idx],
)
} else {
- format_expr_with_alias(&self.group_by.expr[idx])
+ format_expr_with_alias(&self.group_by().expr[idx])
}
})
.collect::<Vec<String>>()
@@ -2000,7 +2091,7 @@ impl DisplayAs for AggregateExec {
.collect()
};
let a: Vec<String> = self
- .aggr_expr
+ .aggr_expr()
.iter()
.map(|agg| format_tree_aggregate_expr(agg).to_string())
.collect();
@@ -2011,7 +2102,7 @@ impl DisplayAs for AggregateExec {
if !a.is_empty() {
writeln!(f, "aggr={}", a.join(", "))?;
}
- if let Some(config) = self.limit_options {
+ if let Some(config) = self.limit_options() {
writeln!(f, "limit={}", config.limit)?;
}
}
@@ -2070,7 +2161,7 @@ impl ExecutionPlan for AggregateExec {
vec![Distribution::UnspecifiedDistribution]
}
AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned => {
- vec![Distribution::KeyPartitioned(self.group_by.input_exprs())]
+ vec![Distribution::KeyPartitioned(self.group_by().input_exprs())]
}
AggregateMode::Final | AggregateMode::Single => {
vec![Distribution::SinglePartition]
@@ -2114,14 +2205,14 @@ impl ExecutionPlan for AggregateExec {
ChildrenPropertiesMode::Recompute => {
let mut me = AggregateExec::try_new_with_schema(
self.mode,
- Arc::clone(&self.group_by),
- self.aggr_expr.to_vec(),
- Arc::clone(&self.filter_expr),
+ Arc::clone(self.group_by()),
+ self.aggr_expr().to_vec(),
+ self.clone_filter_exprs(),
Arc::clone(&children[0]),
Arc::clone(&self.input_schema),
Arc::clone(&self.schema),
)?;
- me.limit_options = self.limit_options;
+ me = me.with_limit_options(self.limit_options());
me.dynamic_filter.clone_from(&self.dynamic_filter);
Ok(Arc::new(me))
}
@@ -2142,15 +2233,15 @@ impl ExecutionPlan for AggregateExec {
&self,
f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>,
) -> Result<TreeNodeRecursion> {
- let group_by = self.group_by.input_exprs();
- let aggregates = self.aggr_expr.iter().flat_map(|aggr| {
+ let group_by = self.group_by().input_exprs();
+ let aggregates = self.aggr_expr().iter().flat_map(|aggr| {
let expressions = aggr.all_expressions();
expressions
.args
.into_iter()
.chain(expressions.order_by_exprs)
});
- let filters = self.filter_expr.iter().flatten().cloned();
+ let filters = self.filter_expr().iter().flatten().cloned();
let dynamic_filter = self.dynamic_filter.iter().map(|dynamic_filter| {
Arc::<DynamicFilterPhysicalExpr>::clone(&dynamic_filter.filter)
as Arc<dyn PhysicalExpr>
@@ -2240,12 +2331,12 @@ impl ExecutionPlan for AggregateExec {
// that are plain input columns can be mapped; a grouping-set null mask
// marks grouping columns that are not available in that set.
let column_mapping: HashMap<usize, usize> = self
- .group_by
+ .group_by()
.expr()
.iter()
.enumerate()
.filter(|(idx, _)| {
- self.group_by
+ self.group_by()
.groups()
.iter()
.all(|null_mask| null_mask.get(*idx) != Some(&true))
@@ -2261,9 +2352,9 @@ impl ExecutionPlan for AggregateExec {
// emit a row even when their input is empty. Parent filters therefore
// cannot be pushed below them, including filters without column
// references.
- let may_emit_on_empty_input = self.group_by.is_true_no_grouping()
+ let may_emit_on_empty_input = self.group_by().is_true_no_grouping()
|| self
- .group_by
+ .group_by()
.groups()
.iter()
.any(|null_mask| null_mask.iter().all(|is_null| *is_null));
@@ -2337,10 +2428,7 @@ impl ExecutionPlan for AggregateExec {
// round-trip gap.
let Self {
mode,
- group_by,
- aggr_expr,
- filter_expr,
- limit_options,
+ kind: _,
input,
// Derived at construction by `create_schema` from `input_schema`,
// `group_by`, `aggr_expr` and `mode`.
@@ -2357,6 +2445,10 @@ impl ExecutionPlan for AggregateExec {
dynamic_filter,
} = self;
+ let group_by = self.group_by();
+ let aggr_expr = self.aggr_expr();
+ let filter_expr = self.filter_expr();
+ let limit_options = self.limit_options();
let input = ctx.encode_child(input)?;
let group_expr =
ctx.encode_expressions(group_by.expr().iter().map(|(expr, _)| expr))?;
@@ -5394,7 +5486,7 @@ mod tests {
// `GroupValuesColumn`. If a `Null` `GroupColumn` is ever added, this
// assertion fires and the test stops covering the row-encoded path.
let group_schema = partial_reduce
- .group_by
+ .group_by()
.group_schema(&partial_reduce.schema())?;
assert!(
!group_values::multi_group_by::supported_schema(&group_schema),
diff --git a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs
index 02c3b2a368..97afab093a 100644
--- a/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/ordered_single_stream.rs
@@ -157,8 +157,8 @@ impl OrderedSingleAggregateStream {
let spill_metrics = SpillMetrics::new(&agg.metrics, partition);
let state_schema = Arc::new(create_schema(
input_schema.as_ref(),
- &agg.group_by,
- &agg.aggr_expr,
+ agg.group_by(),
+ agg.aggr_expr(),
AggregateMode::Partial,
)?);
diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs
index a185650e2f..288bf025e4 100644
--- a/datafusion/physical-plan/src/aggregates/single_stream.rs
+++ b/datafusion/physical-plan/src/aggregates/single_stream.rs
@@ -185,8 +185,8 @@ impl SingleHashAggregateStream {
let spill_metrics = SpillMetrics::new(&agg.metrics, partition);
let state_schema = Arc::new(create_schema(
input_schema.as_ref(),
- &agg.group_by,
- &agg.aggr_expr,
+ agg.group_by(),
+ agg.aggr_expr(),
AggregateMode::Partial,
)?);
diff --git a/datafusion/physical-plan/src/aggregates/spill.rs b/datafusion/physical-plan/src/aggregates/spill.rs
index 5ad9a200e4..dea4fe1efd 100644
--- a/datafusion/physical-plan/src/aggregates/spill.rs
+++ b/datafusion/physical-plan/src/aggregates/spill.rs
@@ -145,7 +145,7 @@ impl AggregateSpill {
replay_agg.input_order_mode = InputOrderMode::Sorted;
let group_schema = match agg.mode {
AggregateMode::Final | AggregateMode::FinalPartitioned => {
- agg.group_by.group_schema(spill_schema)?
+ agg.group_by().group_schema(spill_schema)?
}
AggregateMode::Single | AggregateMode::SinglePartitioned => {
replay_agg.mode = if agg.mode == AggregateMode::Single {
@@ -153,8 +153,8 @@ impl AggregateSpill {
} else {
AggregateMode::FinalPartitioned
};
- replay_agg.group_by = Arc::new(agg.group_by.as_final());
- agg.group_by.group_schema(&agg.input().schema())?
+ *replay_agg.group_by_mut() = Arc::new(agg.group_by().as_final());
+ agg.group_by().group_schema(&agg.input().schema())?
}
mode => {
return internal_err!("{label}: cannot replay aggregate mode {mode:?}");
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org