Optimize ClickBench q17 aggregate limit

Author: Dandandan

benchmarks/src/bin/dfbench.rs

@@ -20,12 +20,7 @@ use datafusion::error::Result;
 
 use clap::{Parser, Subcommand};
 
-#[cfg(all(feature = "snmalloc", feature = "mimalloc"))]
-compile_error!(
-    "feature \"snmalloc\" and feature \"mimalloc\" cannot be enabled at the same time"
-);
-
-#[cfg(feature = "snmalloc")]
+#[cfg(all(feature = "snmalloc", not(feature = "mimalloc")))]
 #[global_allocator]
 static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;
 

benchmarks/src/bin/imdb.rs

@@ -21,12 +21,7 @@ use clap::{Parser, Subcommand};
 use datafusion::error::Result;
 use datafusion_benchmarks::imdb;
 
-#[cfg(all(feature = "snmalloc", feature = "mimalloc"))]
-compile_error!(
-    "feature \"snmalloc\" and feature \"mimalloc\" cannot be enabled at the same time"
-);
-
-#[cfg(feature = "snmalloc")]
+#[cfg(all(feature = "snmalloc", not(feature = "mimalloc")))]
 #[global_allocator]
 static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc;
 

datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs

@@ -28,11 +28,12 @@ use crate::physical_optimizer::test_utils::{
 use arrow::datatypes::DataType;
 use arrow::{compute::SortOptions, util::pretty::pretty_format_batches};
 use datafusion::prelude::SessionContext;
-use datafusion_common::Result;
+use datafusion_common::{Result, config::ConfigOptions};
 use datafusion_execution::config::SessionConfig;
 use datafusion_expr::Operator;
 use datafusion_physical_expr::expressions::{self, cast, col};
 use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr;
+use datafusion_physical_optimizer::PhysicalOptimizerRule;
 use datafusion_physical_plan::{
     ExecutionPlan,
     aggregates::{AggregateExec, AggregateMode},
@@ -332,7 +333,7 @@ fn test_has_aggregate_expression() -> Result<()> {
     let schema = source.schema();
     let agg = TestAggregate::new_count_star();
 
-    // `SELECT <aggregate with no expressions> FROM DataSourceExec LIMIT 10;`, Single AggregateExec
+    // `SELECT a, COUNT(*) FROM DataSourceExec GROUP BY a LIMIT 10;`, Single AggregateExec
     let single_agg = AggregateExec::try_new(
         AggregateMode::Single,
         build_group_by(&schema, vec!["a".to_string()]),
@@ -345,21 +346,82 @@ fn test_has_aggregate_expression() -> Result<()> {
         Arc::new(single_agg),
         10, // fetch
     );
-    // expected not to push the limit to the AggregateExec
+    // expected to push the limit to the AggregateExec
     let plan: Arc<dyn ExecutionPlan> = Arc::new(limit_exec);
     let formatted = get_optimized_plan(&plan)?;
     let actual = formatted.trim();
     assert_snapshot!(
         actual,
         @r"
     LocalLimitExec: fetch=10
-      AggregateExec: mode=Single, gby=[a@0 as a], aggr=[COUNT(*)]
+      AggregateExec: mode=Single, gby=[a@0 as a], aggr=[COUNT(*)], lim=[10]
         DataSourceExec: partitions=1, partition_sizes=[1]
     "
     );
     Ok(())
 }
 
+#[tokio::test]
+async fn test_partial_final_with_aggregate_expression() -> Result<()> {
+    let source = mock_data()?;
+    let schema = source.schema();
+    let agg = TestAggregate::new_count_star();
+
+    // `SELECT a, COUNT(*) FROM DataSourceExec GROUP BY a LIMIT 4;`,
+    // Partial/Final AggregateExec. Both stages can keep the same deterministic
+    // top-k group keys.
+    let partial_agg = AggregateExec::try_new(
+        AggregateMode::Partial,
+        build_group_by(&schema.clone(), vec!["a".to_string()]),
+        vec![Arc::new(agg.count_expr(&schema))], /* aggr_expr */
+        vec![None],                              /* filter_expr */
+        source,                                  /* input */
+        schema.clone(),                          /* input_schema */
+    )?;
+    let final_agg = AggregateExec::try_new(
+        AggregateMode::Final,
+        build_group_by(&schema.clone(), vec!["a".to_string()]),
+        vec![Arc::new(agg.count_expr(&schema))], /* aggr_expr */
+        vec![None],                              /* filter_expr */
+        Arc::new(partial_agg),                   /* input */
+        schema.clone(),                          /* input_schema */
+    )?;
+    let limit_exec = LocalLimitExec::new(
+        Arc::new(final_agg),
+        4, // fetch
+    );
+    let plan: Arc<dyn ExecutionPlan> = Arc::new(limit_exec);
+    let formatted = get_optimized_plan(&plan)?;
+    let actual = formatted.trim();
+    assert_snapshot!(
+        actual,
+        @r"
+    LocalLimitExec: fetch=4
+      AggregateExec: mode=Final, gby=[a@0 as a], aggr=[COUNT(*)], lim=[4]
+        AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[COUNT(*)], lim=[4]
+          DataSourceExec: partitions=1, partition_sizes=[1]
+    "
+    );
+    let optimized =
+        datafusion_physical_optimizer::limited_distinct_aggregation::LimitedDistinctAggregation::new()
+            .optimize(Arc::clone(&plan), &ConfigOptions::new())?;
+    let expected = run_plan_and_format(optimized).await?;
+    assert_snapshot!(
+        expected,
+        @r"
+    +---+----------+
+    | a | COUNT(*) |
+    +---+----------+
+    |   | 1        |
+    | 1 | 2        |
+    | 2 | 1        |
+    | 4 | 1        |
+    +---+----------+
+    "
+    );
+    Ok(())
+}
+
 #[test]
 fn test_has_filter() -> Result<()> {
     let source = mock_data()?;

datafusion/physical-optimizer/src/limited_distinct_aggregation.rs

@@ -15,12 +15,12 @@
 // 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
+//! A special-case optimizer rule that pushes limit into unordered grouped
+//! aggregation when the query only needs an arbitrary subset of groups.
 
 use std::sync::Arc;
 
-use datafusion_physical_plan::aggregates::{AggregateExec, LimitOptions};
+use datafusion_physical_plan::aggregates::{AggregateExec, AggregateMode, LimitOptions};
 use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec};
 use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties};
 
@@ -32,9 +32,10 @@ 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:
+/// groups to be produced 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;`
+/// - `SELECT l_orderkey, COUNT(*) FROM lineitem GROUP BY l_orderkey LIMIT 10;`
 #[derive(Debug)]
 pub struct LimitedDistinctAggregation {}
 
@@ -48,21 +49,43 @@ impl LimitedDistinctAggregation {
         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() {
+        if aggr.is_unordered_unfiltered_group_by_distinct() {
+            let new_aggr = aggr.with_new_limit_options(Some(LimitOptions::new(limit)));
+            return Some(Arc::new(new_aggr));
+        }
+
+        if !Self::can_limit_aggregate(aggr) {
             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))
     }
 
+    fn can_limit_aggregate(aggr: &AggregateExec) -> bool {
+        if !aggr.is_unordered_unfiltered_group_by() {
+            return false;
+        }
+        if aggr.aggr_expr().is_empty() {
+            return false;
+        }
+        if !aggr.group_expr().is_single() {
+            return false;
+        }
+        matches!(
+            aggr.mode(),
+            AggregateMode::Partial
+                | AggregateMode::Final
+                | AggregateMode::FinalPartitioned
+                | AggregateMode::Single
+                | AggregateMode::SinglePartitioned
+        )
+    }
+
     /// 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
+    /// there is a group by, but no sorting or 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;

datafusion/physical-plan/src/aggregates/mod.rs

@@ -915,6 +915,7 @@ impl AggregateExec {
 
         // grouping by an expression that has a sort/limit upstream
         if let Some(config) = self.limit_options
+            && config.descending().is_some()
             && !self.is_unordered_unfiltered_group_by_distinct()
         {
             return Ok(StreamType::GroupedPriorityQueue(
@@ -934,11 +935,9 @@ impl AggregateExec {
         agg_expr.get_minmax_desc()
     }
 
-    /// true, if this Aggregate has a group-by with no required or explicit ordering,
-    /// no filtering and no aggregate expressions
-    /// This method qualifies the use of the LimitedDistinctAggregation rewrite rule
-    /// on an AggregateExec.
-    pub fn is_unordered_unfiltered_group_by_distinct(&self) -> bool {
+    /// true if this Aggregate has a group-by with no required or explicit
+    /// ordering and no aggregate filters.
+    pub fn is_unordered_unfiltered_group_by(&self) -> bool {
         if self
             .limit_options()
             .and_then(|config| config.descending)
@@ -950,12 +949,7 @@ impl AggregateExec {
         if self.group_expr().is_empty() && !self.group_expr().has_grouping_set() {
             return false;
         }
-        // ensure there are no aggregate expressions
-        if !self.aggr_expr().is_empty() {
-            return false;
-        }
-        // ensure there are no filters on aggregate expressions; the above check
-        // may preclude this case
+        // ensure there are no filters on aggregate expressions
         if self.filter_expr().iter().any(|e| e.is_some()) {
             return false;
         }
@@ -974,6 +968,15 @@ impl AggregateExec {
         true
     }
 
+    /// true, if this Aggregate has a group-by with no required or explicit ordering,
+    /// no filtering and no aggregate expressions.
+    ///
+    /// This method qualifies the distinct-only use of the
+    /// LimitedDistinctAggregation rewrite rule on an AggregateExec.
+    pub fn is_unordered_unfiltered_group_by_distinct(&self) -> bool {
+        self.aggr_expr().is_empty() && self.is_unordered_unfiltered_group_by()
+    }
+
     /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
     pub fn compute_properties(
         input: &Arc<dyn ExecutionPlan>,