aggregate blocked impl

datafusion/common/src/utils/proxy.rs

@@ -17,6 +17,7 @@
 
 //! [`VecAllocExt`] to help tracking of memory allocations
 
+use std::collections::VecDeque;
 use hashbrown::hash_table::HashTable;
 use std::mem::size_of;
 
@@ -111,6 +112,41 @@ impl<T> VecAllocExt for Vec<T> {
     }
 }
 
+/// Extension trait for [`Vec`] to account for allocations.
+pub trait VecDequeAllocExt {
+    /// Return the amount of memory allocated by this VecDeque to store elements
+    /// (`size_of<T> * capacity`).
+    ///
+    /// Note this calculation is not recursive, and does not include any heap
+    /// allocations contained within the Vec's elements. Does not include the
+    /// size of `self`
+    ///
+    /// # Example:
+    /// ```
+    /// # use datafusion_common::utils::proxy::VecDequeAllocExt;
+    /// let mut vec = VecDeque::new();
+    /// // Push data into the vec and the accounting will be updated to reflect
+    /// // memory allocation
+    /// vec.push_back(1);
+    /// assert_eq!(vec.allocated_size(), 16); // space for 4 u32s
+    /// vec.push_back(1);
+    /// assert_eq!(vec.allocated_size(), 16); // no new allocation needed
+    ///
+    /// // push more data into the vec
+    /// for _ in 0..10 {
+    ///     vec.push_back(1);
+    /// }
+    /// assert_eq!(vec.allocated_size(), 64); // space for 64 now
+    /// ```
+    fn allocated_size(&self) -> usize;
+}
+
+impl<T> VecDequeAllocExt for VecDeque<T> {
+    fn allocated_size(&self) -> usize {
+        size_of::<T>() * self.capacity()
+    }
+}
+
 /// Extension trait for hash browns [`HashTable`] to account for allocations.
 pub trait HashTableAllocExt {
     /// Item type.

datafusion/core/src/physical_planner.rs

@@ -4576,7 +4576,7 @@ mod tests {
 
         let execution_plan = plan(&logical_plan).await?;
         let final_hash_agg = execution_plan
-            .downcast_ref::<AggregateExec>()
+            .downcast_ref::<BlockedAggregateExec>()
             .expect("hash aggregate");
         assert_eq!(
             "sum(aggregate_test_100.c2)",
@@ -4603,7 +4603,7 @@ mod tests {
 
         let execution_plan = plan(&logical_plan).await?;
         let final_hash_agg = execution_plan
-            .downcast_ref::<AggregateExec>()
+            .downcast_ref::<BlockedAggregateExec>()
             .expect("hash aggregate");
         assert_eq!(
             "sum(aggregate_test_100.c3)",
@@ -5610,6 +5610,7 @@ digraph {
     use crate::catalog::TableProvider;
     use datafusion_catalog::Session;
     use datafusion_expr::TableType;
+    use datafusion_physical_plan::aggregates_blocked::BlockedAggregateExec;
 
     /// A TableProvider that returns schemas for logical planning vs physical planning.
     /// Used to test schema mismatch error messages.

datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs

@@ -552,7 +552,10 @@ async fn verify_ordered_aggregate(frame: &DataFrame, expected_sort: bool) {
         type Node = Arc<dyn ExecutionPlan>;
 
         fn f_down(&mut self, node: &'n Self::Node) -> Result<TreeNodeRecursion> {
-            if let Some(exec) = node.downcast_ref::<AggregateExec>() {
+            if node.downcast_ref::<AggregateExec>().is_some() {
+                return plan_err!("should not get AggregateExec, should be migrated to BlockedAggregateExec");
+            }
+            if let Some(exec) = node.downcast_ref::<BlockedAggregateExec>() {
                 if self.expected_sort {
                     assert!(matches!(
                         exec.input_order_mode(),

datafusion/core/tests/physical_optimizer/aggregate_statistics.rs

@@ -48,6 +48,7 @@ use datafusion_physical_plan::ExecutionPlan;
 use datafusion_physical_plan::aggregates::AggregateExec;
 use datafusion_physical_plan::aggregates::AggregateMode;
 use datafusion_physical_plan::aggregates::PhysicalGroupBy;
+use datafusion_physical_plan::aggregates_blocked::BlockedAggregateExec;
 use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec;
 use datafusion_physical_plan::common;
 use datafusion_physical_plan::displayable;
@@ -74,7 +75,7 @@ fn mock_data() -> Result<Arc<DataSourceExec>> {
 
 /// Checks that the count optimization was applied and we still get the right result
 async fn assert_count_optim_success(
-    plan: AggregateExec,
+    plan: BlockedAggregateExec,
     agg: TestAggregate,
 ) -> Result<()> {
     let task_ctx = Arc::new(TaskContext::default());

datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs

@@ -39,6 +39,7 @@ use datafusion_physical_plan::{
     collect, displayable,
     limit::{GlobalLimitExec, LocalLimitExec},
 };
+use datafusion_physical_plan::aggregates_blocked::BlockedAggregateExec;
 
 async fn run_plan_and_format(plan: Arc<dyn ExecutionPlan>) -> Result<String> {
     let cfg = SessionConfig::new().with_target_partitions(1);
@@ -133,7 +134,12 @@ async fn limited_distinct_aggregate_stream_respects_soft_limit() -> Result<()> {
         plan: &Arc<dyn ExecutionPlan>,
         metrics: &mut Vec<AggregateRuntimeMetric>,
     ) {
-        if let Some(agg) = plan.downcast_ref::<AggregateExec>() {
+        assert!(
+            plan.downcast_ref::<AggregateExec>().is_none(),
+            "should not get AggregateExec, should be migrated to BlockedAggregateExec"
+        );
+
+        if let Some(agg) = plan.downcast_ref::<BlockedAggregateExec>() {
             let output_rows = agg
                 .metrics()
                 .and_then(|metrics| metrics.aggregate_by_name().output_rows())

datafusion/core/tests/sql/explain_analyze.rs

@@ -141,6 +141,7 @@ async fn explain_analyze_baseline_metrics() {
 
         plan.is::<sorts::sort::SortExec>()
             || plan.is::<physical_plan::aggregates::AggregateExec>()
+            || plan.is::<physical_plan::aggregates_blocked::BlockedAggregateExec>()
             || plan.is::<physical_plan::filter::FilterExec>()
             || plan.is::<physical_plan::limit::LocalLimitExec>()
             || plan.is::<physical_plan::projection::ProjectionExec>()

datafusion/expr-common/Cargo.toml

@@ -45,6 +45,7 @@ arrow = { workspace = true }
 datafusion-common = { workspace = true }
 indexmap = { workspace = true }
 itertools = { workspace = true }
+libc = "0.2.185"
 
 [dev-dependencies]
 insta = { workspace = true }