Move HashJoin from RawTable to HashTable (#14904)

Dandandan · web-flow · commit fc2fbb3d6b3a · 2025-02-27T07:27:54.000-05:00
* Move to HashTable

* Add back in comment

* Add TODO for switching to HashTable::allocation_size
diff --git a/datafusion/physical-plan/src/joins/hash_join.rs b/datafusion/physical-plan/src/joins/hash_join.rs
@@ -1656,7 +1656,7 @@ mod tests {
     use datafusion_expr::Operator;
     use datafusion_physical_expr::expressions::{BinaryExpr, Literal};
     use datafusion_physical_expr::PhysicalExpr;
-    use hashbrown::raw::RawTable;
+    use hashbrown::HashTable;
     use rstest::*;
     use rstest_reuse::*;
 
@@ -3295,7 +3295,7 @@ mod tests {
 
     #[test]
     fn join_with_hash_collision() -> Result<()> {
-        let mut hashmap_left = RawTable::with_capacity(2);
+        let mut hashmap_left = HashTable::with_capacity(2);
         let left = build_table_i32(
             ("a", &vec![10, 20]),
             ("x", &vec![100, 200]),
@@ -3311,8 +3311,8 @@ mod tests {
         )?;
 
         // Create hash collisions (same hashes)
-        hashmap_left.insert(hashes[0], (hashes[0], 1), |(h, _)| *h);
-        hashmap_left.insert(hashes[1], (hashes[1], 1), |(h, _)| *h);
+        hashmap_left.insert_unique(hashes[0], (hashes[0], 1), |(h, _)| *h);
+        hashmap_left.insert_unique(hashes[1], (hashes[1], 1), |(h, _)| *h);
 
         let next = vec![2, 0];
 
diff --git a/datafusion/physical-plan/src/joins/stream_join_utils.rs b/datafusion/physical-plan/src/joins/stream_join_utils.rs
@@ -32,6 +32,7 @@ use arrow::array::{
 use arrow::compute::concat_batches;
 use arrow::datatypes::{ArrowNativeType, Schema, SchemaRef};
 use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode};
+use datafusion_common::utils::memory::estimate_memory_size;
 use datafusion_common::{
     arrow_datafusion_err, DataFusionError, HashSet, JoinSide, Result, ScalarValue,
 };
@@ -42,7 +43,7 @@ use datafusion_physical_expr::utils::collect_columns;
 use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr};
 
 use datafusion_physical_expr_common::sort_expr::LexOrdering;
-use hashbrown::raw::RawTable;
+use hashbrown::HashTable;
 
 /// Implementation of `JoinHashMapType` for `PruningJoinHashMap`.
 impl JoinHashMapType for PruningJoinHashMap {
@@ -54,12 +55,12 @@ impl JoinHashMapType for PruningJoinHashMap {
     }
 
     /// Get mutable references to the hash map and the next.
-    fn get_mut(&mut self) -> (&mut RawTable<(u64, u64)>, &mut Self::NextType) {
+    fn get_mut(&mut self) -> (&mut HashTable<(u64, u64)>, &mut Self::NextType) {
         (&mut self.map, &mut self.next)
     }
 
     /// Get a reference to the hash map.
-    fn get_map(&self) -> &RawTable<(u64, u64)> {
+    fn get_map(&self) -> &HashTable<(u64, u64)> {
         &self.map
     }
 
@@ -106,7 +107,7 @@ impl JoinHashMapType for PruningJoinHashMap {
 /// ```
 pub struct PruningJoinHashMap {
     /// Stores hash value to last row index
-    pub map: RawTable<(u64, u64)>,
+    pub map: HashTable<(u64, u64)>,
     /// Stores indices in chained list data structure
     pub next: VecDeque<u64>,
 }
@@ -122,7 +123,7 @@ impl PruningJoinHashMap {
     /// A new instance of `PruningJoinHashMap`.
     pub(crate) fn with_capacity(capacity: usize) -> Self {
         PruningJoinHashMap {
-            map: RawTable::with_capacity(capacity),
+            map: HashTable::with_capacity(capacity),
             next: VecDeque::with_capacity(capacity),
         }
     }
@@ -155,7 +156,11 @@ impl PruningJoinHashMap {
     /// # Returns
     /// The size of the hash map in bytes.
     pub(crate) fn size(&self) -> usize {
-        self.map.allocation_info().1.size() + self.next.capacity() * size_of::<u64>()
+        let fixed_size = size_of::<PruningJoinHashMap>();
+
+        // TODO: switch to using [HashTable::allocation_size] when available after upgrading hashbrown to 0.15
+        estimate_memory_size::<(u64, u64)>(self.map.capacity(), fixed_size).unwrap()
+            + self.next.capacity() * size_of::<u64>()
     }
 
     /// Removes hash values from the map and the list based on the given pruning
@@ -177,20 +182,20 @@ impl PruningJoinHashMap {
         self.next.drain(0..prune_length);
 
         // Calculate the keys that should be removed from the map.
-        let removable_keys = unsafe {
-            self.map
-                .iter()
-                .map(|bucket| bucket.as_ref())
-                .filter_map(|(hash, tail_index)| {
-                    (*tail_index < prune_length as u64 + deleting_offset).then_some(*hash)
-                })
-                .collect::<Vec<_>>()
-        };
+        let removable_keys = self
+            .map
+            .iter()
+            .filter_map(|(hash, tail_index)| {
+                (*tail_index < prune_length as u64 + deleting_offset).then_some(*hash)
+            })
+            .collect::<Vec<_>>();
 
         // Remove the keys from the map.
         removable_keys.into_iter().for_each(|hash_value| {
             self.map
-                .remove_entry(hash_value, |(hash, _)| hash_value == *hash);
+                .find_entry(hash_value, |(hash, _)| hash_value == *hash)
+                .unwrap()
+                .remove();
         });
 
         // Shrink the map if necessary.
@@ -1094,7 +1099,7 @@ pub mod tests {
         let deleted_part = 3 * data_size / 4;
         // Add elements to the JoinHashMap
         for hash_value in 0..data_size {
-            join_hash_map.map.insert(
+            join_hash_map.map.insert_unique(
                 hash_value,
                 (hash_value, hash_value),
                 |(hash, _)| *hash,
@@ -1108,7 +1113,9 @@ pub mod tests {
         for hash_value in 0..deleted_part {
             join_hash_map
                 .map
-                .remove_entry(hash_value, |(hash, _)| hash_value == *hash);
+                .find_entry(hash_value, |(hash, _)| hash_value == *hash)
+                .unwrap()
+                .remove();
         }
 
         assert_eq!(join_hash_map.map.len(), (data_size - deleted_part) as usize);
diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs
@@ -54,12 +54,13 @@ use datafusion_physical_expr::utils::{collect_columns, merge_vectors};
 use datafusion_physical_expr::{
     LexOrdering, PhysicalExpr, PhysicalExprRef, PhysicalSortExpr,
 };
+use hashbrown::hash_table::Entry::{Occupied, Vacant};
+use hashbrown::HashTable;
 
 use crate::joins::SharedBitmapBuilder;
 use crate::projection::ProjectionExec;
 use futures::future::{BoxFuture, Shared};
 use futures::{ready, FutureExt};
-use hashbrown::raw::RawTable;
 use parking_lot::Mutex;
 
 /// Maps a `u64` hash value based on the build side ["on" values] to a list of indices with this key's value.
@@ -126,20 +127,20 @@ use parking_lot::Mutex;
 /// ```
 pub struct JoinHashMap {
     // Stores hash value to last row index
-    map: RawTable<(u64, u64)>,
+    map: HashTable<(u64, u64)>,
     // Stores indices in chained list data structure
     next: Vec<u64>,
 }
 
 impl JoinHashMap {
     #[cfg(test)]
-    pub(crate) fn new(map: RawTable<(u64, u64)>, next: Vec<u64>) -> Self {
+    pub(crate) fn new(map: HashTable<(u64, u64)>, next: Vec<u64>) -> Self {
         Self { map, next }
     }
 
     pub(crate) fn with_capacity(capacity: usize) -> Self {
         JoinHashMap {
-            map: RawTable::with_capacity(capacity),
+            map: HashTable::with_capacity(capacity),
             next: vec![0; capacity],
         }
     }
@@ -199,9 +200,9 @@ pub trait JoinHashMapType {
     /// Extend with zero
     fn extend_zero(&mut self, len: usize);
     /// Returns mutable references to the hash map and the next.
-    fn get_mut(&mut self) -> (&mut RawTable<(u64, u64)>, &mut Self::NextType);
+    fn get_mut(&mut self) -> (&mut HashTable<(u64, u64)>, &mut Self::NextType);
     /// Returns a reference to the hash map.
-    fn get_map(&self) -> &RawTable<(u64, u64)>;
+    fn get_map(&self) -> &HashTable<(u64, u64)>;
     /// Returns a reference to the next.
     fn get_list(&self) -> &Self::NextType;
 
@@ -212,24 +213,28 @@ pub trait JoinHashMapType {
         deleted_offset: usize,
     ) {
         let (mut_map, mut_list) = self.get_mut();
-        for (row, hash_value) in iter {
-            let item = mut_map.get_mut(*hash_value, |(hash, _)| *hash_value == *hash);
-            if let Some((_, index)) = item {
-                // Already exists: add index to next array
-                let prev_index = *index;
-                // Store new value inside hashmap
-                *index = (row + 1) as u64;
-                // Update chained Vec at `row` with previous value
-                mut_list[row - deleted_offset] = prev_index;
-            } else {
-                mut_map.insert(
-                    *hash_value,
-                    // store the value + 1 as 0 value reserved for end of list
-                    (*hash_value, (row + 1) as u64),
-                    |(hash, _)| *hash,
-                );
-                // chained list at `row` is already initialized with 0
-                // meaning end of list
+        for (row, &hash_value) in iter {
+            let entry = mut_map.entry(
+                hash_value,
+                |&(hash, _)| hash_value == hash,
+                |&(hash, _)| hash,
+            );
+
+            match entry {
+                Occupied(mut occupied_entry) => {
+                    // Already exists: add index to next array
+                    let (_, index) = occupied_entry.get_mut();
+                    let prev_index = *index;
+                    // Store new value inside hashmap
+                    *index = (row + 1) as u64;
+                    // Update chained Vec at `row` with previous value
+                    mut_list[row - deleted_offset] = prev_index;
+                }
+                Vacant(vacant_entry) => {
+                    vacant_entry.insert((hash_value, (row + 1) as u64));
+                    // chained list at `row` is already initialized with 0
+                    // meaning end of list
+                }
             }
         }
     }
@@ -251,7 +256,7 @@ pub trait JoinHashMapType {
         for (row_idx, hash_value) in iter {
             // Get the hash and find it in the index
             if let Some((_, index)) =
-                hash_map.get(*hash_value, |(hash, _)| *hash_value == *hash)
+                hash_map.find(*hash_value, |(hash, _)| *hash_value == *hash)
             {
                 let mut i = *index - 1;
                 loop {
@@ -299,7 +304,7 @@ pub trait JoinHashMapType {
 
         let mut remaining_output = limit;
 
-        let hash_map: &RawTable<(u64, u64)> = self.get_map();
+        let hash_map: &HashTable<(u64, u64)> = self.get_map();
         let next_chain = self.get_list();
 
         // Calculate initial `hash_values` index before iterating
@@ -330,7 +335,7 @@ pub trait JoinHashMapType {
         let mut row_idx = to_skip;
         for hash_value in &hash_values[to_skip..] {
             if let Some((_, index)) =
-                hash_map.get(*hash_value, |(hash, _)| *hash_value == *hash)
+                hash_map.find(*hash_value, |(hash, _)| *hash_value == *hash)
             {
                 chain_traverse!(
                     input_indices,
@@ -358,12 +363,12 @@ impl JoinHashMapType for JoinHashMap {
     fn extend_zero(&mut self, _: usize) {}
 
     /// Get mutable references to the hash map and the next.
-    fn get_mut(&mut self) -> (&mut RawTable<(u64, u64)>, &mut Self::NextType) {
+    fn get_mut(&mut self) -> (&mut HashTable<(u64, u64)>, &mut Self::NextType) {
         (&mut self.map, &mut self.next)
     }
 
     /// Get a reference to the hash map.
-    fn get_map(&self) -> &RawTable<(u64, u64)> {
+    fn get_map(&self) -> &HashTable<(u64, u64)> {
         &self.map
     }
 
