2010YOUY01 commented on code in PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#discussion_r4068256404
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -905,49 +900,74 @@ 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. Inapplicable requests are
+ /// no-ops; the existing configuration is returned unchanged.
+ pub fn try_optimize_distinct_soft_limit(
+ mut self,
+ limit: usize,
+ ) -> Result<Transformed<Self>> {
+ if limit == 0
+ || !self.is_unordered_unfiltered_group_by_distinct()
+ || self
+ .limit_options()
+ .is_some_and(|existing| existing.limit <= limit)
+ {
+ return Ok(Transformed::no(self));
+ }
+ self.kind = AggregateKind::DistinctLimit {
+ group_by: Arc::clone(self.group_by()),
+ limit,
+ };
+ Ok(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.
pub fn with_new_aggr_exprs(
&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),
+ filter_expr: vec![None; aggr_expr.len()].into(),
Review Comment:
The previous `AggregateKind` is `DistinctLimit`, which is for the query
pattern `select distinct k from t limit 10` after it gets optimized into an
aggregate, so it must have no aggregate functions or filters.
This transformation converts `DistinctLimit` -> `General`, so there is no
filter to init.
However, the deeper question seems to be: is it safe? I'm not 100% sure.
Before this refactor, this conversion happened implicitly and was easy to miss.
With the typed state, it now shows up obviously.
I'll add more comment for it.
--
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]
