draft.

Rachelint · Rachelint · commit 680bf3b42355 · 2024-08-11T06:59:27.000+08:00
diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs
@@ -33,6 +33,32 @@ mod bytes_view;
 use bytes::GroupValuesByes;
 use datafusion_physical_expr::binary_map::OutputType;
 
+pub type GroupBlock = Vec<ArrayRef>;
+
+const GROUP_IDX_HIGH_16_BITS_MASK: u64 = 0xffff000000000000;
+const GROUP_IDX_LOW_48_BITS_MASK: u64 = 0x0000ffffffffffff;
+
+pub struct GroupIdx(u64);
+
+impl GroupIdx {
+    pub fn new(block_id: u16, block_offset: u64) -> Self {
+        let group_idx_high_part = ((block_id as u64) << 48) & GROUP_IDX_HIGH_16_BITS_MASK;
+        let group_idx_low_part = block_offset & GROUP_IDX_LOW_48_BITS_MASK;
+
+        Self(group_idx_high_part | group_idx_low_part)
+    }
+
+    #[inline]
+    pub fn block_id(&self) -> usize {
+        ((self.0 & GROUP_IDX_HIGH_16_BITS_MASK) >> 48) as usize
+    }
+
+    #[inline]
+    pub fn block_offset(&self) -> usize {
+        (self.0 & GROUP_IDX_LOW_48_BITS_MASK) as usize
+    }
+}
+
 /// An interning store for group keys
 pub trait GroupValues: Send {
     /// Calculates the `groups` for each input row of `cols`
@@ -48,13 +74,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 +118,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)?))
 }
diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs
@@ -15,7 +15,9 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use crate::aggregates::group_values::GroupValues;
+use std::mem;
+
+use crate::aggregates::group_values::{GroupBlock, GroupIdx, GroupValues};
 use ahash::RandomState;
 use arrow::compute::cast;
 use arrow::record_batch::RecordBatch;
@@ -44,7 +46,7 @@ pub struct GroupValuesRows {
     ///
     /// keys: u64 hashes of the GroupValue
     /// values: (hash, group_index)
-    map: RawTable<(u64, usize)>,
+    map: RawTable<(u64, GroupIdx)>,
 
     /// The size of `map` in bytes
     map_size: usize,
@@ -57,7 +59,7 @@ pub struct GroupValuesRows {
     /// important for multi-column group keys.
     ///
     /// [`Row`]: arrow::row::Row
-    group_values: Option<Rows>,
+    group_values_blocks: Vec<Rows>,
 
     /// reused buffer to store hashes
     hashes_buffer: Vec<u64>,
@@ -67,10 +69,14 @@ pub struct GroupValuesRows {
 
     /// Random state for creating hashes
     random_state: RandomState,
+
+    max_block_size: usize,
+
+    cur_block_id: u16,
 }
 
 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()
@@ -90,27 +96,31 @@ impl GroupValuesRows {
             row_converter,
             map,
             map_size: 0,
-            group_values: None,
+            group_values_blocks: Vec::new(),
             hashes_buffer: Default::default(),
             rows_buffer,
             random_state: Default::default(),
+            max_block_size: page_size,
         })
     }
 }
 
 impl GroupValues for GroupValuesRows {
-    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<usize>) -> Result<()> {
+    fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec<GroupIdx>) -> Result<()> {
         // Convert the group keys into the row format
         let group_rows = &mut self.rows_buffer;
         group_rows.clear();
         self.row_converter.append(group_rows, cols)?;
         let n_rows = group_rows.num_rows();
 
-        let mut group_values = match self.group_values.take() {
-            Some(group_values) => group_values,
-            None => self.row_converter.empty_rows(0, 0),
+        if self.group_values_blocks.is_empty() {
+            // TODO: calc and use the capacity to init.
+            let block = self.row_converter.empty_rows(0, 0);
+            self.group_values_blocks.push(block);
         };
 
+        let mut group_values_blocks = mem::take(&mut self.group_values_blocks);
+
         // tracks to which group each of the input rows belongs
         groups.clear();
 
@@ -126,21 +136,38 @@ impl GroupValues for GroupValuesRows {
                 // hash doesn't match, so check the hash first with an integer
                 // comparison first avoid the more expensive comparison with
                 // group value. https://github.com/apache/datafusion/pull/11718
-                target_hash == *exist_hash
-                    // verify that the group that we are inserting with hash is
-                    // actually the same key value as the group in
-                    // existing_idx  (aka group_values @ row)
-                    && group_rows.row(row) == group_values.row(*group_idx)
+                if target_hash != *exist_hash {
+                    return false;
+                }
+
+                // verify that the group that we are inserting with hash is
+                // actually the same key value as the group in
+                // existing_idx  (aka group_values @ row)
+                let block_id = group_idx.block_id();
+                let block_offset = group_idx.block_offset();
+                let group_value = group_values_blocks[block_id].row(block_offset);
+                group_rows.row(row) == group_value
             });
 
             let group_idx = match entry {
                 // Existing group_index for this group value
                 Some((_hash, group_idx)) => *group_idx,
                 //  1.2 Need to create new entry for the group
                 None => {
+                    // Check if the block size has reached the limit, if so we switch to next block.
+                    let block_size =  group_values_blocks.last().unwrap().num_rows();
+                    if block_size == self.max_block_size {
+                        self.cur_block_id += 1;
+                        // TODO: calc and use the capacity to init.
+                        let block = self.row_converter.empty_rows(0, 0);
+                        self.group_values_blocks.push(block);
+                    }
+
                     // Add new entry to aggr_state and save newly created index
-                    let group_idx = group_values.num_rows();
-                    group_values.push(group_rows.row(row));
+                    let cur_group_values = self.group_values_blocks.last_mut().unwrap();
+                    let block_offset = group_values.num_rows();
+                    let group_idx = GroupIdx::new(self.cur_block_id, block_offset);
+                    cur_group_values.push(group_rows.row(row));
 
                     // for hasher function, use precomputed hash value
                     self.map.insert_accounted(
@@ -154,13 +181,13 @@ impl GroupValues for GroupValuesRows {
             groups.push(group_idx);
         }
 
-        self.group_values = Some(group_values);
+        self.group_values_blocks = group_values_blocks;
 
         Ok(())
     }
 
     fn size(&self) -> usize {
-        let group_values_size = self.group_values.as_ref().map(|v| v.size()).unwrap_or(0);
+        let group_values_size = self.group_values_blocks.as_ref().map(|v| v.size()).unwrap_or(0);
         self.row_converter.size()
             + group_values_size
             + self.map_size
@@ -173,15 +200,15 @@ impl GroupValues for GroupValuesRows {
     }
 
     fn len(&self) -> usize {
-        self.group_values
+        self.group_values_blocks
             .as_ref()
             .map(|group_values| group_values.num_rows())
             .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
+            .group_values_blocks
             .take()
             .expect("Can not emit from empty rows");
 
@@ -232,13 +259,13 @@ impl GroupValues for GroupValuesRows {
             }
         }
 
-        self.group_values = Some(group_values);
+        self.group_values_blocks = Some(group_values);
         Ok(output)
     }
 
     fn clear_shrink(&mut self, batch: &RecordBatch) {
         let count = batch.num_rows();
-        self.group_values = self.group_values.take().map(|mut rows| {
+        self.group_values_blocks = self.group_values_blocks.take().map(|mut rows| {
             rows.clear();
             rows
         });
diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/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;
