reduce memory

Author: forsaken628

src/query/expression/src/aggregate/aggregate_hashtable.rs

@@ -189,8 +189,8 @@ impl AggregateHashTable {
         group_hash_columns(group_columns, &mut state.group_hashes);
 
         let new_group_count = if self.direct_append {
-            for idx in 0..row_count {
-                state.empty_vector[idx] = idx;
+            for i in 0..row_count {
+                state.empty_vector[i] = i.into();
             }
             self.payload.append_rows(state, row_count, group_columns);
             row_count

src/query/expression/src/aggregate/hash_index.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use super::payload_row::CompareState;
+use super::CompareItem;
 use super::PartitionedPayload;
 use super::ProbeState;
 use super::RowPtr;
@@ -151,17 +153,16 @@ impl HashIndex {
         row_count: usize,
         mut adapter: impl TableAdapter,
     ) -> usize {
-        for (i, row) in state.no_match_vector[..row_count].iter_mut().enumerate() {
-            *row = i;
+        for (i, item) in state.no_match_vector[..row_count].iter_mut().enumerate() {
+            let hash = state.group_hashes[i];
+            *item = CompareItem {
+                row: i.into(),
+                salt: Entry::hash_to_salt(hash),
+                slot: self.init_slot(hash),
+                row_ptr: RowPtr::null(),
+            };
         }
 
-        let mut slots = state.get_temp();
-        slots.extend(
-            state.group_hashes[..row_count]
-                .iter()
-                .map(|hash| self.init_slot(*hash)),
-        );
-
         let mut new_group_count = 0;
         let mut remaining_entries = row_count;
 
@@ -171,18 +172,21 @@ impl HashIndex {
             let mut no_match_count = 0;
 
             // 1. inject new_group_count, new_entry_count, need_compare_count, no_match_count
-            for row in state.no_match_vector[..remaining_entries].iter().copied() {
-                let slot = &mut slots[row];
-                let hash = state.group_hashes[row];
-
-                let is_new;
-                (*slot, is_new) = self.find_or_insert(*slot, Entry::hash_to_salt(hash));
+            for item in state.no_match_vector[..remaining_entries].iter_mut() {
+                let (slot, is_new) = self.find_or_insert(item.slot, item.salt);
+                item.slot = slot;
 
                 if is_new {
-                    state.empty_vector[new_entry_count] = row;
+                    state.empty_vector[new_entry_count] = item.row;
+                    state.slots[new_entry_count] = slot;
                     new_entry_count += 1;
                 } else {
-                    state.group_compare_vector[need_compare_count] = row;
+                    state.group_compare_vector[need_compare_count] = CompareItem {
+                        row: item.row,
+                        slot: item.slot,
+                        salt: item.salt,
+                        row_ptr: self.mut_entry(slot).get_pointer(),
+                    };
                     need_compare_count += 1;
                 }
             }
@@ -193,42 +197,45 @@ impl HashIndex {
 
                 adapter.append_rows(state, new_entry_count);
 
-                for row in state.empty_vector[..new_entry_count].iter().copied() {
-                    let entry = self.mut_entry(slots[row]);
+                for (i, row) in state.empty_vector[..new_entry_count]
+                    .iter()
+                    .copied()
+                    .enumerate()
+                {
+                    let entry = self.mut_entry(state.slots[i]);
                     entry.set_pointer(state.addresses[row]);
                     debug_assert_eq!(entry.get_pointer(), state.addresses[row]);
                 }
             }
 
             // 3. set address of compare vector
             if need_compare_count > 0 {
-                for row in state.group_compare_vector[..need_compare_count]
-                    .iter()
-                    .copied()
-                {
-                    let entry = self.mut_entry(slots[row]);
+                for item in &mut state.group_compare_vector[..need_compare_count] {
+                    let entry = self.mut_entry(item.slot);
 
                     debug_assert!(entry.is_occupied());
-                    debug_assert_eq!(entry.get_salt(), (state.group_hashes[row] >> 48) as u16);
-                    state.addresses[row] = entry.get_pointer();
+                    debug_assert_eq!(
+                        entry.get_salt(),
+                        (state.group_hashes[item.row] >> 48) as u16
+                    );
+                    item.row_ptr = entry.get_pointer();
+                    state.addresses[item.row] = item.row_ptr;
                 }
 
                 // 4. compare
                 no_match_count = adapter.compare(state, need_compare_count, no_match_count);
             }
 
             // 5. Linear probing, just increase iter_times
-            for row in state.no_match_vector[..no_match_count].iter().copied() {
-                let slot = &mut slots[row];
-                *slot += 1;
-                if *slot >= self.capacity {
-                    *slot = 0;
+            for item in &mut state.no_match_vector[..no_match_count] {
+                item.slot += 1;
+                if item.slot >= self.capacity {
+                    item.slot = 0;
                 }
             }
             remaining_entries = no_match_count;
         }
 
-        state.save_temp(slots);
         self.count += new_group_count;
 
         new_group_count
@@ -251,7 +258,12 @@ impl<'a> TableAdapter for AdapterImpl<'a> {
         need_compare_count: usize,
         no_match_count: usize,
     ) -> usize {
-        state.row_match_columns(
+        CompareState {
+            compare: &mut state.group_compare_vector,
+            matched: &mut state.match_vector,
+            no_matched: &mut state.no_match_vector,
+        }
+        .row_match_columns(
             self.group_columns,
             &self.payload.row_layout,
             (need_compare_count, no_match_count),
@@ -284,8 +296,10 @@ mod tests {
         }
 
         fn init_state(&self) -> ProbeState {
-            let mut state = ProbeState::default();
-            state.row_count = self.incoming.len();
+            let mut state = ProbeState {
+                row_count: self.incoming.len(),
+                ..Default::default()
+            };
 
             for (i, (_, hash)) in self.incoming.iter().enumerate() {
                 state.group_hashes[i] = *hash
@@ -323,12 +337,13 @@ mod tests {
 
     impl TableAdapter for &mut TestTableAdapter {
         fn append_rows(&mut self, state: &mut ProbeState, new_entry_count: usize) {
-            for row in state.empty_vector[..new_entry_count].iter().copied() {
-                let (key, hash) = self.incoming[row];
+            for row in state.empty_vector[..new_entry_count].iter() {
+                let row_index = row.to_index();
+                let (key, hash) = self.incoming[row_index];
                 let value = key + 20;
 
                 self.payload.push((key, hash, value));
-                state.addresses[row] = self.get_row_ptr(true, row);
+                state.addresses[*row] = self.get_row_ptr(true, row_index);
             }
         }
 
@@ -338,23 +353,18 @@ mod tests {
             need_compare_count: usize,
             mut no_match_count: usize,
         ) -> usize {
-            for row in state.group_compare_vector[..need_compare_count]
-                .iter()
-                .copied()
-            {
-                let incoming = self.incoming[row];
-
-                let row_ptr = state.addresses[row];
+            for item in &state.group_compare_vector[..need_compare_count] {
+                let incoming = self.incoming[item.row.to_index()];
 
-                let (key, hash, _) = self.get_payload(row_ptr);
+                let (key, hash, _) = self.get_payload(item.row_ptr);
 
                 const POINTER_MASK: u64 = 0x0000FFFFFFFFFFFF;
                 assert_eq!(incoming.1 | POINTER_MASK, hash | POINTER_MASK);
                 if incoming.0 == key {
                     continue;
                 }
 
-                state.no_match_vector[no_match_count] = row;
+                state.no_match_vector[no_match_count] = item.clone();
                 no_match_count += 1;
             }
 

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

@@ -39,8 +39,9 @@ use hash_index::Entry;
 pub use partitioned_payload::*;
 pub use payload::*;
 pub use payload_flush::*;
-pub use probe_state::*;
-use row_ptr::RowPtr;
+pub use probe_state::ProbeState;
+use probe_state::*;
+use row_ptr::*;
 
 // A batch size to probe, flush, repartition, etc.
 pub(crate) const BATCH_SIZE: usize = 2048;

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

@@ -122,7 +122,7 @@ impl PartitionedPayload {
                 let partition_idx = ((hash & self.mask_v) >> self.shift_v) as usize;
                 let (count, sel) = &mut state.partition_entries[partition_idx];
 
-                sel[*count] = row;
+                sel[*count as usize] = row;
                 *count += 1;
             }
 
@@ -133,7 +133,7 @@ impl PartitionedPayload {
             {
                 if *count > 0 {
                     payload.reserve_append_rows(
-                        &sel[..*count],
+                        &sel[..*count as _],
                         &state.group_hashes,
                         &mut state.addresses,
                         &mut state.page_index,
@@ -200,7 +200,7 @@ impl PartitionedPayload {
                     let (count, sel) = &state.partition_entries[partition];
                     if *count > 0 {
                         let payload = &mut self.payloads[partition];
-                        payload.copy_rows(&sel[..*count], &flush_state.addresses);
+                        payload.copy_rows(&sel[..*count as _], &flush_state.addresses);
                     }
                 }
             }
@@ -239,7 +239,7 @@ impl PartitionedPayload {
             let partition_idx = ((hash & self.mask_v) >> self.shift_v) as usize;
 
             let (count, sel) = &mut state.partition_entries[partition_idx];
-            sel[*count] = idx;
+            sel[*count as usize] = idx.into();
             *count += 1;
         }
         flush_state.flush_page_row = end;

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

@@ -24,6 +24,7 @@ use super::payload_row::rowformat_size;
 use super::payload_row::serialize_column_to_rowformat;
 use super::row_ptr::RowLayout;
 use super::row_ptr::RowPtr;
+use super::RowID;
 use crate::types::DataType;
 use crate::AggrState;
 use crate::AggregateFunctionRef;
@@ -197,17 +198,17 @@ impl Payload {
 
     pub(super) fn reserve_append_rows(
         &mut self,
-        select_vector: &[usize],
+        select_vector: &[RowID],
         group_hashes: &[u64; BATCH_SIZE],
         address: &mut [RowPtr; BATCH_SIZE],
         page_index: &mut [usize],
         group_columns: ProjectedBlock,
     ) {
         let tuple_size = self.tuple_size;
         let (mut page, mut page_index_value) = self.writable_page();
-        for idx in select_vector.iter().copied() {
-            address[idx] = page.data_ptr(page.rows, tuple_size);
-            page_index[idx] = page_index_value;
+        for row in select_vector {
+            address[*row] = page.data_ptr(page.rows, tuple_size);
+            page_index[*row] = page_index_value;
             page.rows += 1;
 
             if page.rows == page.capacity {
@@ -226,7 +227,7 @@ impl Payload {
 
     fn append_rows(
         &mut self,
-        select_vector: &[usize],
+        select_vector: &[RowID],
         group_hashes: &[u64; BATCH_SIZE],
         address: &mut [RowPtr; BATCH_SIZE],
         page_index: &mut [usize],
@@ -240,15 +241,16 @@ impl Payload {
                 if bitmap.null_count() == 0 || bitmap.null_count() == bitmap.len() {
                     let val: u8 = if bitmap.null_count() == 0 { 1 } else { 0 };
                     // faster path
-                    for idx in select_vector.iter().copied() {
+                    for row in select_vector {
                         unsafe {
-                            address[idx].write_u8(write_offset, val);
+                            address[*row].write_u8(write_offset, val);
                         }
                     }
                 } else {
-                    for idx in select_vector.iter().copied() {
+                    for row in select_vector {
                         unsafe {
-                            address[idx].write_u8(write_offset, bitmap.get_bit(idx) as u8);
+                            address[*row]
+                                .write_u8(write_offset, bitmap.get_bit(row.to_index()) as u8);
                         }
                     }
                 }
@@ -276,8 +278,8 @@ impl Payload {
 
         // write group hashes
         debug_assert!(write_offset == self.row_layout.hash_offset);
-        for idx in select_vector.iter().copied() {
-            address[idx].set_hash(&self.row_layout, group_hashes[idx]);
+        for row in select_vector {
+            address[*row].set_hash(&self.row_layout, group_hashes[*row]);
         }
 
         debug_assert!(write_offset + 8 == self.row_layout.state_offset);
@@ -293,7 +295,7 @@ impl Payload {
                 .iter()
                 .copied()
                 .enumerate()
-                .map(|(i, idx)| (idx, unsafe { place.add(padded_size * i) }))
+                .map(|(i, row)| (row.to_index(), unsafe { place.add(padded_size * i) }))
             {
                 let place = StateAddr::from(place);
                 address[idx].set_state_addr(&self.row_layout, &place);
@@ -336,7 +338,7 @@ impl Payload {
         }
     }
 
-    pub fn copy_rows(&mut self, select_vector: &[usize], address: &[RowPtr; BATCH_SIZE]) {
+    pub fn copy_rows(&mut self, select_vector: &[RowID], address: &[RowPtr; BATCH_SIZE]) {
         let tuple_size = self.tuple_size;
         let agg_len = self.aggrs.len();
         let (mut page, _) = self.writable_page();
@@ -395,7 +397,7 @@ impl Payload {
             let partition_idx = (hash % mods) as usize;
 
             let (count, sel) = &mut state.partition_entries[partition_idx];
-            sel[*count] = idx;
+            sel[*count as usize] = idx.into();
             *count += 1;
         }
         flush_state.flush_page_row = end;