alamb commented on code in PR #26069:
URL: https://github.com/apache/datafusion/pull/26069#discussion_r4194420295
##########
datafusion/sqllogictest/test_files/aggregate.slt:
##########
@@ -7639,9 +7639,8 @@ logical_plan
03)----TableScan: distinct_limit_combine projection=[a, b]
physical_plan
01)GlobalLimitExec: skip=0, fetch=2
-02)--AggregateExec: mode=Final, gby=[b@0 as b, a@1 as a], aggr=[], lim=[2]
-03)----AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[]
-04)------DataSourceExec: partitions=1, partition_sizes=[1]
+02)--AggregateExec: mode=Single, gby=[b@1 as b, a@0 as a], aggr=[], lim=[2]
Review Comment:
this does seem a better plan (use only the final single partition execution
plan)
##########
datafusion/sqllogictest/test_files/aggregate.slt:
##########
@@ -7553,7 +7553,7 @@ physical_plan
02)--GlobalLimitExec: skip=0, fetch=3
03)----AggregateExec: mode=Final, gby=[c2@0 as c2, c3@1 as c3, __grouping_id@2
as __grouping_id], aggr=[], lim=[3]
04)------CoalescePartitionsExec
-05)--------AggregateExec: mode=Partial, gby=[(NULL as c2, NULL as c3), (c2@0
as c2, NULL as c3), (c2@0 as c2, c3@1 as c3)], aggr=[]
+05)--------AggregateExec: mode=Partial, gby=[(NULL as c2, NULL as c3), (c2@0
as c2, NULL as c3), (c2@0 as c2, c3@1 as c3)], aggr=[], lim=[3]
Review Comment:
This looks correct to me -- to push the limit down into each partial
aggregate / partition
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -44,95 +86,70 @@ impl LimitedDistinctAggregation {
Self {}
}
- fn transform_agg(
- aggr: &AggregateExec,
- limit: usize,
- ) -> 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>))
- }
+ /// Rewrite a limit and its immediately adjacent final/partial aggregate
pair.
+ fn transform_limit(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ // Step 1: Identify the plan shape,
+ //
+ // Limit
+ // Aggregate(final)
+ // Aggregate(partial)
+
+ // Check the current plan is limit, and extract limit value, input
plan.
+ let (limit, input) = match (
+ plan.downcast_ref::<LocalLimitExec>(),
+ plan.downcast_ref::<GlobalLimitExec>(),
+ ) {
+ (Some(local), _) => (local.fetch(), local.input()),
+ (_, Some(global)) => match global.fetch() {
+ Some(fetch) => (global.skip() + fetch, global.input()),
+ None => return Ok(Transformed::no(plan)),
+ },
+ _ => return Ok(Transformed::no(plan)),
+ };
- /// transform_limit matches an `AggregateExec` as the child of a
`LocalLimitExec`
- /// or `GlobalLimitExec` and pushes the limit into the aggregation as a
soft limit when
- /// there is a group by, but no sorting, no aggregate expressions, and no
filters in the
- /// aggregation
- fn transform_limit(plan: Arc<dyn ExecutionPlan>) -> Option<Arc<dyn
ExecutionPlan>> {
- let limit: usize;
- let mut global_fetch: Option<usize> = None;
- let mut global_skip: usize = 0;
- let children: Vec<Arc<dyn ExecutionPlan>>;
- let mut is_global_limit = false;
- if let Some(local_limit) = plan.downcast_ref::<LocalLimitExec>() {
- limit = local_limit.fetch();
- children = local_limit.children().into_iter().cloned().collect();
- } else {
- let global_limit = plan.downcast_ref::<GlobalLimitExec>()?;
- global_fetch = global_limit.fetch();
- global_fetch?;
- global_skip = global_limit.skip();
- // the aggregate must read at least fetch+skip number of rows
- limit = global_fetch.unwrap() + global_skip;
- children = global_limit.children().into_iter().cloned().collect();
- is_global_limit = true
- }
- let child = children.iter().exactly_one().ok()?;
- // ensure there is no output ordering; can this rule be relaxed?
- if plan.output_ordering().is_some() {
- return None;
- }
- // ensure no ordering is required on the input
- if plan.required_input_ordering()[0].is_some() {
- return None;
+ if plan.output_ordering().is_some() ||
plan.required_input_ordering()[0].is_some()
+ {
+ return Ok(Transformed::no(plan));
}
- // if found_match_aggr is true, match_aggr holds a parent aggregation
whose group_by
- // must match that of a child aggregation in order to rewrite the
child aggregation
- let mut match_aggr: Arc<dyn ExecutionPlan> = plan;
- let mut found_match_aggr = false;
-
- let mut rewrite_applicable = true;
- let closure = |plan: Arc<dyn ExecutionPlan>| {
- if !rewrite_applicable {
- return Ok(Transformed::no(plan));
- }
- if let Some(aggr) = plan.downcast_ref::<AggregateExec>() {
- if found_match_aggr
- && let Some(parent_aggr) =
match_aggr.downcast_ref::<AggregateExec>()
- && !parent_aggr.group_expr().eq(aggr.group_expr())
- {
- // a partial and final aggregation with different
groupings disqualifies
- // rewriting the child aggregation
- rewrite_applicable = false;
- return Ok(Transformed::no(plan));
- }
- // either we run into an Aggregate and transform it, or
disable the rewrite
- // for subsequent children
- match Self::transform_agg(aggr, limit) {
- None => {}
- Some(new_aggr) => {
- match_aggr = plan;
- found_match_aggr = true;
- return Ok(new_aggr);
- }
- }
- }
- rewrite_applicable = false;
- Ok(Transformed::no(plan))
+ let Some(final_agg) = input.downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(plan));
};
- let child = child.to_owned().transform_down(closure).ok()?;
- if !child.transformed {
- return None;
+ let Some(partial_agg) =
final_agg.input().downcast_ref::<AggregateExec>() else {
+ return Ok(Transformed::no(plan));
+ };
+
+ // Step 2: Verify partial/final aggregate is compatible, then apply
optimization
+ match (final_agg.mode(), partial_agg.mode()) {
+ (
+ AggregateMode::Final | AggregateMode::FinalPartitioned,
+ AggregateMode::Partial,
+ ) if final_agg.group_expr() ==
&partial_agg.group_expr().as_final() => {}
Review Comment:
Maybe this check could be a method on AggregateExec -- mostly so that it
could be given a name and better documented
Something like:
```rust
if final_agg.matches_partial(partial_agg) {
...
}
```
(not sure if that is a good name for it)
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -15,26 +15,68 @@
// specific language governing permissions and limitations
// under the License.
-//! A special-case optimizer rule that pushes limit into a grouped aggregation
-//! which has no aggregate expressions or sorting requirements
+//! Stop unordered DISTINCT aggregation once enough groups have been found.
use std::sync::Arc;
-use datafusion_physical_plan::aggregates::AggregateExec;
+use datafusion_physical_plan::aggregates::{AggregateExec, AggregateMode};
use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
-use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
+use datafusion_physical_plan::{
+ ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties,
+ ReplaceChildrenOptions,
+};
use datafusion_common::Result;
use datafusion_common::config::ConfigOptions;
use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
use crate::PhysicalOptimizerRule;
-use itertools::Itertools;
-/// An optimizer rule that passes a `limit` hint into grouped aggregations
which don't require all
-/// rows in the group to be processed for correctness. Example queries fitting
this description are:
-/// - `SELECT distinct l_orderkey FROM lineitem LIMIT 10;`
-/// - `SELECT l_orderkey FROM lineitem GROUP BY l_orderkey LIMIT 10;`
+/// Pushes a soft limit into unordered DISTINCT aggregations. For
+/// `SELECT DISTINCT a FROM t LIMIT 10`, any ten distinct values are
sufficient:
+/// later rows cannot change the groups already found. The same reasoning
applies
+/// to `GROUP BY a` without aggregate expressions. Aggregates such as `SUM`
need
+/// every row in each group.
+///
+/// Each partial stage can also stop after finding ten distinct keys: those
keys
+/// remain distinct in the final stage, even if other partitions contribute
+/// duplicates.
+///
+/// Before:
+///
+/// ```txt
+/// Limit(10)
+/// Aggregate(Final, a)
+/// Aggregate(Partial, a)
+/// Scan
+/// ```
+///
+/// After:
+///
+/// ```txt
+/// Limit(10)
+/// Aggregate(Final, a, soft_limit=10)
+/// Aggregate(Partial, a, soft_limit=10)
+/// Scan
+/// ```
+///
+/// # Invariants before and after the rewrite
Review Comment:
I am not sure I would use the term "invariant" here though it is technically
accurate
This comment mostly explains the effect of this optimizer rule (what it
does).
I normally think of Invariant as a property that will not change as any
operation is applied to it. I don't think there is any reason to require that
some future optimizer rule preserves the same property
For example, what if a future rule (or a user defined rule) has some
additional special operator that requires the full intermediate results (aka
undoes the soft limit on the Partial AggregateExec)? I don't see any reason to
try and prevent that 🤔
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -44,95 +86,70 @@ impl LimitedDistinctAggregation {
Self {}
}
- fn transform_agg(
- aggr: &AggregateExec,
- limit: usize,
- ) -> 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>))
- }
+ /// Rewrite a limit and its immediately adjacent final/partial aggregate
pair.
+ fn transform_limit(
+ plan: Arc<dyn ExecutionPlan>,
+ ) -> Result<Transformed<Arc<dyn ExecutionPlan>>> {
+ // Step 1: Identify the plan shape,
Review Comment:
thank you -- these comments really make the code much easier to follow
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
