x

Author: forsaken628

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

@@ -28,6 +28,28 @@ use crate::ProjectedBlock;
 use crate::StatesLayout;
 use crate::BATCH_SIZE;
 
+#[derive(Debug, Clone, Copy)]
+struct PartitionMask {
+    mask: u64,
+    shift: u64,
+}
+
+impl PartitionMask {
+    fn new(partition_count: u64) -> Self {
+        let radix_bits = partition_count.trailing_zeros() as u64;
+        debug_assert_eq!(1 << radix_bits, partition_count);
+
+        let shift = 48 - radix_bits;
+        let mask = ((1 << radix_bits) - 1) << shift;
+
+        Self { mask, shift }
+    }
+
+    pub fn index(&self, hash: u64) -> usize {
+        ((hash & self.mask) >> self.shift) as _
+    }
+}
+
 pub struct PartitionedPayload {
     pub payloads: Vec<Payload>,
     pub group_types: Vec<DataType>,
@@ -37,9 +59,7 @@ pub struct PartitionedPayload {
 
     pub arenas: Vec<Arc<Bump>>,
 
-    partition_count: u64,
-    mask_v: u64,
-    shift_v: u64,
+    partition_mask: PartitionMask,
 }
 
 unsafe impl Send for PartitionedPayload {}
@@ -52,9 +72,6 @@ impl PartitionedPayload {
         partition_count: u64,
         arenas: Vec<Arc<Bump>>,
     ) -> Self {
-        let radix_bits = partition_count.trailing_zeros() as u64;
-        debug_assert_eq!(1 << radix_bits, partition_count);
-
         let states_layout = if !aggrs.is_empty() {
             Some(get_states_layout(&aggrs).unwrap())
         } else {
@@ -72,7 +89,7 @@ impl PartitionedPayload {
             })
             .collect_vec();
 
-        let offsets = RowLayout {
+        let row_layout = RowLayout {
             states_layout,
             ..payloads[0].row_layout.clone()
         };
@@ -81,12 +98,10 @@ impl PartitionedPayload {
             payloads,
             group_types,
             aggrs,
-            row_layout: offsets,
-            partition_count,
+            row_layout,
 
             arenas,
-            mask_v: mask(radix_bits),
-            shift_v: shift(radix_bits),
+            partition_mask: PartitionMask::new(partition_count),
         }
     }
 
@@ -119,7 +134,7 @@ impl PartitionedPayload {
             state.reset_partitions(self.partition_count());
             for &row in &state.empty_vector[..new_group_rows] {
                 let hash = state.group_hashes[row];
-                let partition_idx = ((hash & self.mask_v) >> self.shift_v) as usize;
+                let partition_idx = self.partition_mask.index(hash);
                 let (count, sel) = &mut state.partition_entries[partition_idx];
 
                 sel[*count as usize] = row;
@@ -149,19 +164,27 @@ impl PartitionedPayload {
             return self;
         }
 
-        let mut new_partition_payload = PartitionedPayload::new(
-            self.group_types.clone(),
-            self.aggrs.clone(),
-            new_partition_count as u64,
-            self.arenas.clone(),
-        );
+        let PartitionedPayload {
+            payloads,
+            group_types,
+            aggrs,
+            arenas,
+            ..
+        } = self;
+
+        let mut new_partition_payload =
+            PartitionedPayload::new(group_types, aggrs, new_partition_count as u64, arenas);
+
+        state.clear();
+        for payload in payloads.into_iter() {
+            new_partition_payload.combine_single(payload, state, None)
+        }
 
-        new_partition_payload.combine(self, state);
         new_partition_payload
     }
 
     pub fn combine(&mut self, other: PartitionedPayload, state: &mut PayloadFlushState) {
-        if other.partition_count == self.partition_count {
+        if other.partition_count() == self.partition_count() {
             for (l, r) in self.payloads.iter_mut().zip(other.payloads.into_iter()) {
                 l.combine(r);
             }
@@ -184,7 +207,7 @@ impl PartitionedPayload {
             return;
         }
 
-        if self.partition_count == 1 {
+        if self.partition_count() == 1 {
             self.payloads[0].combine(other);
         } else {
             flush_state.clear();
@@ -194,13 +217,19 @@ impl PartitionedPayload {
                 // copy rows
                 let state = &*flush_state.probe_state;
 
-                for partition in (0..self.partition_count as usize)
-                    .filter(|x| only_bucket.is_none() || only_bucket == Some(*x))
-                {
-                    let (count, sel) = &state.partition_entries[partition];
-                    if *count > 0 {
-                        let payload = &mut self.payloads[partition];
-                        payload.copy_rows(&sel[..*count as _], &flush_state.addresses);
+                match only_bucket {
+                    Some(i) => {
+                        let (count, sel) = &state.partition_entries[i];
+                        self.payloads[i].copy_rows(&sel[..*count as _], &flush_state.addresses);
+                    }
+                    None => {
+                        for ((count, sel), payload) in
+                            state.partition_entries.iter().zip(self.payloads.iter_mut())
+                        {
+                            if *count > 0 {
+                                payload.copy_rows(&sel[..*count as _], &flush_state.addresses);
+                            }
+                        }
                     }
                 }
             }
@@ -236,7 +265,7 @@ impl PartitionedPayload {
             flush_state.addresses[idx] = row_ptr;
 
             let hash = row_ptr.hash(&self.row_layout);
-            let partition_idx = ((hash & self.mask_v) >> self.shift_v) as usize;
+            let partition_idx = self.partition_mask.index(hash);
 
             let (count, sel) = &mut state.partition_entries[partition_idx];
             sel[*count as usize] = idx.into();
@@ -253,7 +282,7 @@ impl PartitionedPayload {
 
     #[inline]
     pub fn partition_count(&self) -> usize {
-        self.partition_count as usize
+        self.payloads.len()
     }
 
     #[allow(dead_code)]
@@ -266,13 +295,3 @@ impl PartitionedPayload {
         self.payloads.iter().map(|x| x.memory_size()).sum()
     }
 }
-
-#[inline]
-fn shift(radix_bits: u64) -> u64 {
-    48 - radix_bits
-}
-
-#[inline]
-fn mask(radix_bits: u64) -> u64 {
-    ((1 << radix_bits) - 1) << shift(radix_bits)
-}