alamb commented on code in PR #25564:
URL: https://github.com/apache/datafusion/pull/25564#discussion_r4065957123
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -47,16 +47,15 @@ impl LimitedDistinctAggregation {
fn transform_agg(
aggr: &AggregateExec,
limit: usize,
- ) -> Option<Arc<dyn ExecutionPlan>> {
+ ) -> Option<Transformed<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))
+ let new_aggr =
aggr.clone().try_optimize_distinct_soft_limit(limit).ok()?;
Review Comment:
The pattern of ignoring Errs I think is bad in DataFusion because each `Err`
contains an allocated string (so this is inefficient).
Perhaps we could just change `try_optimize_distinct_soft_limit` to return an
Option rather than Result 🤔
Also, it would be great to avoid the need to clone the input(as this is a
deep clone, not just a an arc clone)
##########
datafusion/physical-optimizer/src/limited_distinct_aggregation.rs:
##########
@@ -119,14 +118,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()?;
Review Comment:
though here is an example of discarding an error in existing code too 😬
##########
datafusion/physical-plan/src/aggregates/mod.rs:
##########
@@ -845,35 +845,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,
Review Comment:
👍
##########
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:
Why are the filter expressions reset?
##########
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
Review Comment:
This seems to always return "Ok" -- why is it returning Result?
--
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]
