draft.

Author: Rachelint

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

@@ -33,6 +33,8 @@ mod bytes_view;
 use bytes::GroupValuesByes;
 use datafusion_physical_expr::binary_map::OutputType;
 
+pub type GroupBlock = Vec<ArrayRef>;
+
 /// An interning store for group keys
 pub trait GroupValues: Send {
     /// Calculates the `groups` for each input row of `cols`
@@ -48,13 +50,13 @@ pub trait GroupValues: Send {
     fn len(&self) -> usize;
 
     /// Emits the group values
-    fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>>;
+    fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<GroupBlock>>;
 
     /// Clear the contents and shrink the capacity to the size of the batch (free up memory usage)
     fn clear_shrink(&mut self, batch: &RecordBatch);
 }
 
-pub fn new_group_values(schema: SchemaRef) -> Result<Box<dyn GroupValues>> {
+pub fn new_group_values(schema: SchemaRef, batch_size: usize) -> Result<Box<dyn GroupValues>> {
     if schema.fields.len() == 1 {
         let d = schema.fields[0].data_type();
 
@@ -92,5 +94,5 @@ pub fn new_group_values(schema: SchemaRef) -> Result<Box<dyn GroupValues>> {
         }
     }
 
-    Ok(Box::new(GroupValuesRows::try_new(schema)?))
+    Ok(Box::new(GroupValuesRows::try_new(schema, batch_size)?))
 }

datafusion/physical-plan/src/aggregates/group_values/row.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::aggregates::group_values::GroupValues;
+use crate::aggregates::group_values::{GroupBlock, GroupValues};
 use ahash::RandomState;
 use arrow::compute::cast;
 use arrow::record_batch::RecordBatch;
@@ -57,7 +57,7 @@ pub struct GroupValuesRows {
     /// important for multi-column group keys.
     ///
     /// [`Row`]: arrow::row::Row
-    group_values: Option<Rows>,
+    group_values: Vec<Rows>,
 
     /// reused buffer to store hashes
     hashes_buffer: Vec<u64>,
@@ -67,10 +67,12 @@ pub struct GroupValuesRows {
 
     /// Random state for creating hashes
     random_state: RandomState,
+
+    page_size: usize,
 }
 
 impl GroupValuesRows {
-    pub fn try_new(schema: SchemaRef) -> Result<Self> {
+    pub fn try_new(schema: SchemaRef, page_size: usize) -> Result<Self> {
         let row_converter = RowConverter::new(
             schema
                 .fields()
@@ -94,6 +96,7 @@ impl GroupValuesRows {
             hashes_buffer: Default::default(),
             rows_buffer,
             random_state: Default::default(),
+            page_size,
         })
     }
 }
@@ -179,7 +182,7 @@ impl GroupValues for GroupValuesRows {
             .unwrap_or(0)
     }
 
-    fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<ArrayRef>> {
+    fn emit(&mut self, emit_to: EmitTo) -> Result<Vec<GroupBlock>> {
         let mut group_values = self
             .group_values
             .take()

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

@@ -463,7 +463,7 @@ impl GroupedHashAggregateStream {
             ordering.as_slice(),
         )?;
 
-        let group_values = new_group_values(group_schema)?;
+        let group_values = new_group_values(group_schema, batch_size)?;
         timer.done();
 
         let exec_state = ExecutionState::ReadingInput;