feat(query): update

src/query/expression/src/aggregate/mod.rs

@@ -25,7 +25,6 @@ mod payload_flush;
 mod payload_row;
 mod probe_state;
 
-use std::hash::Hasher;
 use std::sync::atomic::AtomicU64;
 use std::sync::Arc;
 
@@ -92,68 +91,3 @@ impl HashTableConfig {
         self
     }
 }
-
-pub struct PerfectHashBuilder;
-
-// NOTE: This is a dummy hasher that just returns the value passed to it.
-// This is only used for i8-i64, u8-u64, isize and usize keys.
-pub struct PerfectHash {
-    val: u64,
-}
-
-impl std::hash::BuildHasher for PerfectHashBuilder {
-    type Hasher = PerfectHash;
-    fn build_hasher(&self) -> PerfectHash {
-        PerfectHash { val: 0 }
-    }
-}
-
-impl Hasher for PerfectHash {
-    fn finish(&self) -> u64 {
-        self.val
-    }
-
-    fn write(&mut self, _bytes: &[u8]) {
-        unreachable!()
-    }
-
-    fn write_u8(&mut self, i: u8) {
-        self.val = i as u64;
-    }
-
-    fn write_u16(&mut self, i: u16) {
-        self.val = i as u64;
-    }
-
-    fn write_u32(&mut self, i: u32) {
-        self.val = i as u64;
-    }
-
-    fn write_u64(&mut self, i: u64) {
-        self.val = i;
-    }
-
-    fn write_usize(&mut self, i: usize) {
-        self.val = i as u64;
-    }
-
-    fn write_i8(&mut self, i: i8) {
-        self.val = i as u64;
-    }
-
-    fn write_i16(&mut self, i: i16) {
-        self.val = i as u64;
-    }
-
-    fn write_i32(&mut self, i: i32) {
-        self.val = i as u64;
-    }
-
-    fn write_i64(&mut self, i: i64) {
-        self.val = i as u64;
-    }
-
-    fn write_isize(&mut self, i: isize) {
-        self.val = i as u64;
-    }
-}

src/query/expression/src/aggregate/partitioned_payload.rs

@@ -39,7 +39,7 @@ pub struct PartitionedPayload {
     pub state_addr_offsets: Vec<usize>,
     pub state_layout: Option<Layout>,
 
-    arenas: Vec<Arc<Bump>>,
+    pub arenas: Vec<Arc<Bump>>,
 
     partition_count: u64,
     mask_v: u64,
@@ -166,7 +166,7 @@ impl PartitionedPayload {
         self.arenas.append(&mut other.arenas);
     }
 
-    fn combine_single(&mut self, mut other: Payload, state: &mut PayloadFlushState) {
+    pub fn combine_single(&mut self, mut other: Payload, state: &mut PayloadFlushState) {
         if other.len() == 0 {
             return;
         }

src/query/expression/src/aggregate/payload.rs

@@ -303,6 +303,7 @@ impl Payload {
         let mut page = self.writable_page();
         for i in 0..row_count {
             let index = select_vector[i];
+
             unsafe {
                 std::ptr::copy_nonoverlapping(
                     address[index],
@@ -335,9 +336,13 @@ impl Drop for Payload {
                     for page in self.pages.iter() {
                         for row in 0..page.rows {
                             unsafe {
-                                let state_addr = self.data_ptr(page, row).add(self.state_offset);
-                                aggr.drop_state(StateAddr::new(state_addr as usize + *addr_offset))
-                            };
+                                let state_place = StateAddr::new(core::ptr::read::<u64>(
+                                    self.data_ptr(page, row).add(self.state_offset) as _,
+                                )
+                                    as usize);
+
+                                aggr.drop_state(state_place.next(*addr_offset));
+                            }
                         }
                     }
                 }

src/query/service/src/pipelines/processors/transforms/aggregator/transform_partition_bucket.rs

@@ -443,17 +443,27 @@ impl<Method: HashMethodBounds, V: Copy + Send + Sync + 'static> Processor
             let aggrs = payloads[0].aggrs.clone();
 
             let mut payload_map = (0..self.max_partition_count).map(|_| vec![]).collect_vec();
-            for payload in payloads.into_iter() {
+
+            // All arenas should be kept in the bucket partition payload
+            let mut arenas = vec![];
+
+            for mut payload in payloads.into_iter() {
                 for (bucket, p) in payload.payloads.into_iter().enumerate() {
                     payload_map[bucket].push(p);
                 }
+                arenas.append(&mut payload.arenas);
             }
 
             for (bucket, mut payloads) in payload_map.into_iter().enumerate() {
                 let mut partition_payload =
                     PartitionedPayload::new(group_types.clone(), aggrs.clone(), 1);
 
-                partition_payload.payloads.append(payloads.as_mut());
+                for payload in payloads.drain(0..) {
+                    partition_payload.combine_single(payload, &mut self.flush_state);
+                }
+
+                partition_payload.arenas.extend_from_slice(&arenas);
+
                 self.buckets_blocks
                     .insert(bucket as isize, vec![DataBlock::empty_with_meta(
                         AggregateMeta::<Method, V>::create_agg_hashtable(partition_payload),