feat: Support spilling for hash aggregation

datafusion/core/src/physical_plan/aggregates/group_values/mod.rs

@@ -15,6 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
+use arrow::record_batch::RecordBatch;
 use arrow_array::{downcast_primitive, ArrayRef};
 use arrow_schema::SchemaRef;
 use datafusion_common::Result;
@@ -42,6 +43,19 @@ pub trait GroupValues: Send {
 
     /// Emits the group values
     fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>>;
+
+    /// Try to reserve the capacity that at least a single [`RecordBatch`] can be inserted. The
+    /// accumulator may reserve more space to speculatively avoid frequent re-allocations. After
+    /// calling try_reserve, capacity will be greater than or equal to self.len() + additional if
+    /// it returns Ok(()). Does nothing if capacity is already sufficient. This method preserves
+    /// the contents even if an error occurs.
+    fn try_reserve(
+        &mut self,
+        batch: &RecordBatch,
+    ) -> Result<(), hashbrown::TryReserveError>;
+
+    /// clear the contents and shrink the capacity
+    fn clear_shrink(&mut self, batch: &RecordBatch);
 }
 
 pub fn new_group_values(schema: SchemaRef) -> Result<Box<dyn GroupValues>> {

datafusion/core/src/physical_plan/aggregates/group_values/primitive.rs

@@ -20,6 +20,7 @@ use ahash::RandomState;
 use arrow::array::BooleanBufferBuilder;
 use arrow::buffer::NullBuffer;
 use arrow::datatypes::i256;
+use arrow::record_batch::RecordBatch;
 use arrow_array::cast::AsArray;
 use arrow_array::{ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, PrimitiveArray};
 use arrow_schema::DataType;
@@ -206,4 +207,30 @@ where
         };
         Ok(vec![Arc::new(array.with_data_type(self.data_type.clone()))])
     }
+
+    // FIXME: cannot return std::collections::TryReserveError because std::collections::TryReserveErrorKind
+    //   is unstable. For now, use hashbrown::TryReserveError instead.
+    fn try_reserve(
+        &mut self,
+        batch: &RecordBatch,
+    ) -> Result<(), hashbrown::TryReserveError> {
+        let additional = batch.num_rows();
+        self.values
+            .try_reserve(additional)
+            .map_err(|_| hashbrown::TryReserveError::CapacityOverflow)
+            .and({
+                let state = &self.random_state;
+                self.map.try_reserve(additional, |g| unsafe {
+                    self.values.get_unchecked(*g).hash(state)
+                })
+            })
+    }
+
+    fn clear_shrink(&mut self, batch: &RecordBatch) {
+        let count = batch.num_rows();
+        self.values.clear();
+        self.values.shrink_to(count);
+        self.map.clear();
+        self.map.shrink_to(count, |_| 0); // hasher does not matter since the map is cleared
+    }
 }

datafusion/core/src/physical_plan/aggregates/group_values/row.rs

@@ -17,6 +17,7 @@
 
 use crate::physical_plan::aggregates::group_values::GroupValues;
 use ahash::RandomState;
+use arrow::record_batch::RecordBatch;
 use arrow::row::{RowConverter, Rows, SortField};
 use arrow_array::ArrayRef;
 use arrow_schema::SchemaRef;
@@ -59,17 +60,19 @@ pub struct GroupValuesRows {
 
     /// Random state for creating hashes
     random_state: RandomState,
+
+    /// Schema fields for the row converter
+    fields: Vec<SortField>,
 }
 
 impl GroupValuesRows {
     pub fn try_new(schema: SchemaRef) -> Result<Self> {
-        let row_converter = RowConverter::new(
-            schema
-                .fields()
-                .iter()
-                .map(|f| SortField::new(f.data_type().clone()))
-                .collect(),
-        )?;
+        let fields: Vec<SortField> = schema
+            .fields()
+            .iter()
+            .map(|f| SortField::new(f.data_type().clone()))
+            .collect();
+        let row_converter = RowConverter::new(fields.clone())?;
 
         let map = RawTable::with_capacity(0);
         let group_values = row_converter.empty_rows(0, 0);
@@ -81,6 +84,7 @@ impl GroupValuesRows {
             group_values,
             hashes_buffer: Default::default(),
             random_state: Default::default(),
+            fields,
         })
     }
 }
@@ -181,4 +185,30 @@ impl GroupValues for GroupValuesRows {
             }
         })
     }
+
+    // FIXME: cannot return std::collections::TryReserveError because std::collections::TryReserveErrorKind
+    //   is unstable. For now, use hashbrown::TryReserveError instead.
+    fn try_reserve(
+        &mut self,
+        batch: &RecordBatch,
+    ) -> Result<(), hashbrown::TryReserveError> {
+        let additional = batch.num_rows();
+        // FIXME: there is no good way to try_reserve for self.row_converter self.group_values
+        self.map.try_reserve(additional, |(hash, _)| *hash).and(
+            self.hashes_buffer
+                .try_reserve(additional)
+                .map_err(|_| hashbrown::TryReserveError::CapacityOverflow),
+        )
+    }
+
+    fn clear_shrink(&mut self, batch: &RecordBatch) {
+        let count = batch.num_rows();
+        // FIXME: there is no good way to clear_shrink for self.row_converter self.group_values
+        self.row_converter = RowConverter::new(self.fields.clone()).unwrap();
+        self.group_values = self.row_converter.empty_rows(count, 0);
+        self.map.clear();
+        self.map.shrink_to(count, |_| 0); // hasher does not matter since the map is cleared
+        self.hashes_buffer.clear();
+        self.hashes_buffer.shrink_to(count);
+    }
 }

datafusion/core/src/physical_plan/aggregates/mod.rs

@@ -281,6 +281,8 @@ pub struct AggregateExec {
     /// Stores mode and output ordering information for the `AggregateExec`.
     aggregation_ordering: Option<AggregationOrdering>,
     required_input_ordering: Option<LexOrderingReq>,
+    /// Force spilling for debugging
+    force_spill: bool,
 }
 
 /// Calculates the working mode for `GROUP BY` queries.
@@ -669,9 +671,34 @@ impl AggregateExec {
             metrics: ExecutionPlanMetricsSet::new(),
             aggregation_ordering,
             required_input_ordering,
+            force_spill: false,
         })
     }
 
+    /// Only for testing. When `force_spill` is true, it spills every batch.
+    pub fn try_new_for_test(
+        mode: AggregateMode,
+        group_by: PhysicalGroupBy,
+        aggr_expr: Vec<Arc<dyn AggregateExpr>>,
+        filter_expr: Vec<Option<Arc<dyn PhysicalExpr>>>,
+        order_by_expr: Vec<Option<LexOrdering>>,
+        input: Arc<dyn ExecutionPlan>,
+        input_schema: SchemaRef,
+        force_spill: bool,
+    ) -> Result<Self> {
+        let mut exec = AggregateExec::try_new(
+            mode,
+            group_by,
+            aggr_expr,
+            filter_expr,
+            order_by_expr,
+            input,
+            input_schema,
+        )?;
+        exec.force_spill = force_spill;
+        Ok(exec)
+    }
+
     /// Aggregation mode (full, partial)
     pub fn mode(&self) -> &AggregateMode {
         &self.mode
@@ -1389,7 +1416,10 @@ mod tests {
         )
     }
 
-    async fn check_grouping_sets(input: Arc<dyn ExecutionPlan>) -> Result<()> {
+    async fn check_grouping_sets(
+        input: Arc<dyn ExecutionPlan>,
+        spill: bool,
+    ) -> Result<()> {
         let input_schema = input.schema();
 
         let grouping_set = PhysicalGroupBy {
@@ -1416,14 +1446,15 @@ mod tests {
 
         let task_ctx = Arc::new(TaskContext::default());
 
-        let partial_aggregate = Arc::new(AggregateExec::try_new(
+        let partial_aggregate = Arc::new(AggregateExec::try_new_for_test(
             AggregateMode::Partial,
             grouping_set.clone(),
             aggregates.clone(),
             vec![None],
             vec![None],
             input,
             input_schema.clone(),
+            spill,
         )?);
 
         let result =
@@ -1460,14 +1491,15 @@ mod tests {
 
         let final_grouping_set = PhysicalGroupBy::new_single(final_group);
 
-        let merged_aggregate = Arc::new(AggregateExec::try_new(
+        let merged_aggregate = Arc::new(AggregateExec::try_new_for_test(
             AggregateMode::Final,
             final_grouping_set,
             aggregates,
             vec![None],
             vec![None],
             merge,
             input_schema,
+            spill,
         )?);
 
         let result =
@@ -1505,7 +1537,7 @@ mod tests {
     }
 
     /// build the aggregates on the data from some_data() and check the results
-    async fn check_aggregates(input: Arc<dyn ExecutionPlan>) -> Result<()> {
+    async fn check_aggregates(input: Arc<dyn ExecutionPlan>, spill: bool) -> Result<()> {
         let input_schema = input.schema();
 
         let grouping_set = PhysicalGroupBy {
@@ -1522,14 +1554,15 @@ mod tests {
 
         let task_ctx = Arc::new(TaskContext::default());
 
-        let partial_aggregate = Arc::new(AggregateExec::try_new(
+        let partial_aggregate = Arc::new(AggregateExec::try_new_for_test(
             AggregateMode::Partial,
             grouping_set.clone(),
             aggregates.clone(),
             vec![None],
             vec![None],
             input,
             input_schema.clone(),
+            spill,
         )?);
 
         let result =
@@ -1556,14 +1589,15 @@ mod tests {
 
         let final_grouping_set = PhysicalGroupBy::new_single(final_group);
 
-        let merged_aggregate = Arc::new(AggregateExec::try_new(
+        let merged_aggregate = Arc::new(AggregateExec::try_new_for_test(
             AggregateMode::Final,
             final_grouping_set,
             aggregates,
             vec![None],
             vec![None],
             merge,
             input_schema,
+            spill,
         )?);
 
         let result =
@@ -1707,31 +1741,63 @@ mod tests {
         let input: Arc<dyn ExecutionPlan> =
             Arc::new(TestYieldingExec { yield_first: false });
 
-        check_aggregates(input).await
+        check_aggregates(input, false).await
     }
 
     #[tokio::test]
     async fn aggregate_grouping_sets_source_not_yielding() -> Result<()> {
         let input: Arc<dyn ExecutionPlan> =
             Arc::new(TestYieldingExec { yield_first: false });
 
-        check_grouping_sets(input).await
+        check_grouping_sets(input, false).await
     }
 
     #[tokio::test]
     async fn aggregate_source_with_yielding() -> Result<()> {
         let input: Arc<dyn ExecutionPlan> =
             Arc::new(TestYieldingExec { yield_first: true });
 
-        check_aggregates(input).await
+        check_aggregates(input, false).await
     }
 
     #[tokio::test]
     async fn aggregate_grouping_sets_with_yielding() -> Result<()> {
         let input: Arc<dyn ExecutionPlan> =
             Arc::new(TestYieldingExec { yield_first: true });
 
-        check_grouping_sets(input).await
+        check_grouping_sets(input, false).await
+    }
+
+    #[tokio::test]
+    async fn aggregate_source_not_yielding_with_spill() -> Result<()> {
+        let input: Arc<dyn ExecutionPlan> =
+            Arc::new(TestYieldingExec { yield_first: false });
+
+        check_aggregates(input, true).await
+    }
+
+    #[tokio::test]
+    async fn aggregate_grouping_sets_source_not_yielding_with_spill() -> Result<()> {
+        let input: Arc<dyn ExecutionPlan> =
+            Arc::new(TestYieldingExec { yield_first: false });
+
+        check_grouping_sets(input, true).await
+    }
+
+    #[tokio::test]
+    async fn aggregate_source_with_yielding_with_spill() -> Result<()> {
+        let input: Arc<dyn ExecutionPlan> =
+            Arc::new(TestYieldingExec { yield_first: true });
+
+        check_aggregates(input, true).await
+    }
+
+    #[tokio::test]
+    async fn aggregate_grouping_sets_with_yielding_with_spill() -> Result<()> {
+        let input: Arc<dyn ExecutionPlan> =
+            Arc::new(TestYieldingExec { yield_first: true });
+
+        check_grouping_sets(input, true).await
     }
 
     #[tokio::test]
@@ -1899,7 +1965,10 @@ mod tests {
     async fn run_first_last_multi_partitions() -> Result<()> {
         for use_coalesce_batches in [false, true] {
             for is_first_acc in [false, true] {
-                first_last_multi_partitions(use_coalesce_batches, is_first_acc).await?
+                for spill in [false, true] {
+                    first_last_multi_partitions(use_coalesce_batches, is_first_acc, spill)
+                        .await?
+                }
             }
         }
         Ok(())
@@ -1925,6 +1994,7 @@ mod tests {
     async fn first_last_multi_partitions(
         use_coalesce_batches: bool,
         is_first_acc: bool,
+        spill: bool,
     ) -> Result<()> {
         let task_ctx = Arc::new(TaskContext::default());
 
@@ -1969,14 +2039,15 @@ mod tests {
             schema.clone(),
             None,
         )?);
-        let aggregate_exec = Arc::new(AggregateExec::try_new(
+        let aggregate_exec = Arc::new(AggregateExec::try_new_for_test(
             AggregateMode::Partial,
             groups.clone(),
             aggregates.clone(),
             vec![None],
             vec![Some(ordering_req.clone())],
             memory_exec,
             schema.clone(),
+            spill,
         )?);
         let coalesce = if use_coalesce_batches {
             let coalesce = Arc::new(CoalescePartitionsExec::new(aggregate_exec));
@@ -1985,14 +2056,15 @@ mod tests {
             Arc::new(CoalescePartitionsExec::new(aggregate_exec))
                 as Arc<dyn ExecutionPlan>
         };
-        let aggregate_final = Arc::new(AggregateExec::try_new(
+        let aggregate_final = Arc::new(AggregateExec::try_new_for_test(
             AggregateMode::Final,
             groups,
             aggregates.clone(),
             vec![None],
             vec![Some(ordering_req)],
             coalesce,
             schema,
+            spill,
         )?) as Arc<dyn ExecutionPlan>;
 
         let result = crate::physical_plan::collect(aggregate_final, task_ctx).await?;