Upgrade to hashbrown 0.16.1: migrate from hashbrown::raw::RawTable to hashbrown::hash_table::HashTable

What

Migrate from hashbrown::raw::RawTable to hashbrown::hash_table::HashTable.

Why

RawTable and raw_entry are removed in hashbrown 0.15.1. See #13256 . Personally I think it's the right thing to do because RawTable was a bit of a weird API.

How

First, we need to get some memory accounting for HashTable. This can be done in a similar way to

datafusion/datafusion/common/src/utils/proxy.rs

Lines 110 to 111 in e25f5e7

    
           /// Extension trait for hash browns [`RawTable`] to account for allocations. 
        
           pub trait RawTableAllocExt {

/// Extension trait for hash browns [`HashTable`] to account for allocations.
pub trait HashTableAllocExt {
    /// Item type.
    type T;

    /// Insert new element into table and increase
    /// `accounting` by any newly allocated bytes.
    ///
    /// Returns the bucket where the element was inserted.
    /// Note that allocation counts capacity, not size.
    ///
    /// # Example:
    /// ```
    /// TODO: rewrite this example!
    /// ```
    fn insert_accounted(
        &mut self,
        x: Self::T,
        hasher: impl Fn(&Self::T) -> u64,
        accounting: &mut usize,
    );
}

impl<T> HashTableAllocExt for HashTable<T> where T: Eq {
    type T = T;

    fn insert_accounted(
        &mut self,
        x: Self::T,
        hasher: impl Fn(&Self::T) -> u64,
        accounting: &mut usize,
    ) {
        let hash = hasher(&x);

        // NOTE: `find_entry` does NOT grow!
        match self.find_entry(hash, |y| y == &x) {
            Ok(_occupied) => {}
            Err(_absent) => {
                if self.len() == self.capacity() {
                    // need to request more memory
                    let bump_elements = self.capacity().max(16);
                    let bump_size = bump_elements * size_of::<T>();
                    *accounting = (*accounting).checked_add(bump_size).expect("overflow");

                    self.reserve(bump_elements, &hasher);
                }

                // still need to insert the element since first try failed
                self.entry(hash, |y| y == &x, hasher).insert(x);
            }
        }
    }
}

Then, migrate every RawTable to HashTable. Since RawTable is used all over the place, I think this should be split into multiple PRs. Luckily, HashTable is already available in hashbrown 0.14.x, so we can do this iteratively before upgrading to hashbrown 0.15.x.

Progress

refactor: migrate `GroupValuesRows` to `HashTable`

refactor: migrate `GroupValuesPrimitive` to `HashTable`

refactor: migrate `GroupValuesColumn` to `HashTable`

refactor: migrate `GroupValuesRows` to `HashTable`

refactor: migrate `GroupValuesPrimitive` to `HashTable`

refactor: migrate `GroupValuesColumn` to `HashTable`

refactor: change some `hashbrown` `RawTable` uses to `HashTable` (rou…

refactor: migrate `LinearSearch` to `HashTable`

refactor: migrate `LinearSearch` to `HashTable` (#13658)

@crepererum

Thanks @crepererum for the steady progress on this

TBH the last two batches are rather hard:

`hashbrown` 0.14 & allocation size

hashbrown 0.14 doesn't expose the allocation size for HashTable which we would need here:

datafusion/datafusion/physical-plan/src/joins/stream_join_utils.rs

Lines 152 to 158 in 47569b2

    
               /// Calculates the size of the `PruningJoinHashMap` in bytes. 
        
               /// 
        
               /// # 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>() 
        
               }

hashbrown 0.15 offers that, but then we have to roll the RawTable->HashTable change together with the upgrade, not as separate steps.

`ArrowHashTable`

This is a bit of a wild interface and I dunno if it's worth the 15% uplift TBH:

datafusion/datafusion/physical-plan/src/aggregates/topk/hash_table.rs

Lines 63 to 86 in 47569b2

    
           /// An interface to hide the generic type signature of TopKHashTable behind arrow arrays 
        
           pub trait ArrowHashTable { 
        
               fn set_batch(&mut self, ids: ArrayRef); 
        
               fn len(&self) -> usize; 
        
               // JUSTIFICATION 
        
               //  Benefit:  ~15% speedup + required to index into RawTable from binary heap 
        
               //  Soundness: the caller must provide valid indexes 
        
               unsafe fn update_heap_idx(&mut self, mapper: &[(usize, usize)]); 
        
               // JUSTIFICATION 
        
               //  Benefit:  ~15% speedup + required to index into RawTable from binary heap 
        
               //  Soundness: the caller must provide a valid index 
        
               unsafe fn heap_idx_at(&self, map_idx: usize) -> usize; 
        
               unsafe fn take_all(&mut self, indexes: Vec<usize>) -> ArrayRef; 
        
               // JUSTIFICATION 
        
               //  Benefit:  ~15% speedup + required to index into RawTable from binary heap 
        
               //  Soundness: the caller must provide valid indexes 
        
               unsafe fn find_or_insert( 
        
                   &mut self, 
        
                   row_idx: usize, 
        
                   replace_idx: usize, 
        
                   map: &mut Vec<(usize, usize)>, 
        
               ) -> (usize, bool); 
        
           }

The issue here is that HashTable doesn't offer an unsafe "directly address the bucket via index" kind of interface because it is fundamentally rather risky. I wonder if we should somewhat redesign this part. As far as I understand, this uses the following concepts:

data heap: a "heap" (which I guess is just a vector) to store the actual payload data
mutable slot: a slot that stores a mutable index to the data heap. I guess this is used to update the data pointer whenever a new/better value was found. The slots are referenced by a simple usize index.
key lookup: A way to lookup of key->mutable slot.

The current solution fuses 2 & 3 into a single RawTable w/ a LOT of unsafe. I think we could deconstruct that into a HashTable (for 3) + Vec (for 2). This can be done BEFORE the hashbrown 0.15 update.

@avantgardnerio

This is a bit of a wild interface and I dunno if it's worth the 15% uplift TBH:

I think this is something @avantgardnerio and @thinkharderdev may know more about. Perhaps they can offer some suggestions / thoughts about improving the performance or some benchmarks that we can use to see how much performance would be affected if we switched to the newer hashbrown version

@avantgardnerio

I think this is something @avantgardnerio and @thinkharderdev may know more about.

The main reason that 15% was so important was to keep it on par with the unoptimized version. If CPU is equal, then the optimizer rule was easy: just always use this rule.

If we are okay taking a performance hit, or if @crepererum 's idea keeps performance the same, then there's no reason to keep the unsafe code.

@avantgardnerio

TBH the last two batches are rather hard:

hashbrown 0.14 & allocation size

hashbrown 0.14 doesn't expose the allocation size for HashTable which we would need here:

datafusion/datafusion/physical-plan/src/joins/stream_join_utils.rs

Lines 152 to 158 in 47569b2

/// Calculates the size of the `PruningJoinHashMap` in bytes.

///

/// # 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>()

}

hashbrown 0.15 offers that, but then we have to roll the RawTable->HashTable change together with the upgrade, not as separate steps.

ArrowHashTable

This is a bit of a wild interface and I dunno if it's worth the 15% uplift TBH:

datafusion/datafusion/physical-plan/src/aggregates/topk/hash_table.rs

Lines 63 to 86 in 47569b2

/// An interface to hide the generic type signature of TopKHashTable behind arrow arrays

pub trait ArrowHashTable {

fn set_batch(&mut self, ids: ArrayRef);

fn len(&self) -> usize;

// JUSTIFICATION

// Benefit: ~15% speedup + required to index into RawTable from binary heap

// Soundness: the caller must provide valid indexes

unsafe fn update_heap_idx(&mut self, mapper: &[(usize, usize)]);

// JUSTIFICATION

// Benefit: ~15% speedup + required to index into RawTable from binary heap

// Soundness: the caller must provide a valid index

unsafe fn heap_idx_at(&self, map_idx: usize) -> usize;

unsafe fn take_all(&mut self, indexes: Vec<usize>) -> ArrayRef;

// JUSTIFICATION

// Benefit: ~15% speedup + required to index into RawTable from binary heap

// Soundness: the caller must provide valid indexes

unsafe fn find_or_insert(

&mut self,

row_idx: usize,

replace_idx: usize,

map: &mut Vec<(usize, usize)>,

) -> (usize, bool);

}

The issue here is that HashTable doesn't offer an unsafe "directly address the bucket via index" kind of interface because it is fundamentally rather risky. I wonder if we should somewhat redesign this part. As far as I understand, this uses the following concepts:

data heap: a "heap" (which I guess is just a vector) to store the actual payload data

mutable slot: a slot that stores a mutable index to the data heap. I guess this is used to update the data pointer whenever a new/better value was found. The slots are referenced by a simple usize index.

key lookup: A way to lookup of key->mutable slot.

The current solution fuses 2 & 3 into a single RawTable w/ a LOT of unsafe. I think we could deconstruct that into a HashTable (for 3) + Vec (for 2). This can be done BEFORE the hashbrown 0.15 update.

The basic idea is just a hash map where the values are sorted so you can use as a combo aggregation by key + topk heap. @avantgardnerio and I went through a couple iterations on the design and landed on this because it doesn't regress performance when the grouping key is low cardinality.

@crepererum

If we are okay taking a performance hit, or if @crepererum 's idea keeps performance the same, then there's no reason to keep the unsafe code.

My idea was to ideally not regress significantly, at least not without discussing it first.

refactor: migrate `LinearSearch` to `HashTable` (apache#13658)

Crosslinking original raw feature deprecation from hashbrown rust-lang/hashbrown#546

reg to PRs above I think it is agreed to replace RawTable with HashTable across DataFusion?

reg to PRs above I think it is agreed to replace RawTable with HashTable across DataFusion?

correct. As far as I can tell, this is the last usage of it:

datafusion/datafusion/physical-plan/src/aggregates/topk/hash_table.rs

Line 57 in 3fe4640

map: RawTable<HashTableItem<ID>>,

Refactor TopKHashTable to use HashTable API (#19464)

Update: when this is merged, we can upgrade to 0.16: #19464

Upgrade hashbrown to 0.16 (#19554)

Navigation Menu

Repository navigation

Upgrade to hashbrown 0.16.1: migrate from `hashbrown::raw::RawTable` to `hashbrown::hash_table::HashTable` #13433

What

Why

How

Progress

`hashbrown` 0.14 & allocation size

`ArrowHashTable`

`hashbrown` 0.14 & allocation size

`ArrowHashTable`

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

	/// Extension trait for hash browns [`RawTable`] to account for allocations.
	pub trait RawTableAllocExt {

	/// Calculates the size of the `PruningJoinHashMap` in bytes.
	///
	/// # 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>()
	}

	/// An interface to hide the generic type signature of TopKHashTable behind arrow arrays
	pub trait ArrowHashTable {
	fn set_batch(&mut self, ids: ArrayRef);
	fn len(&self) -> usize;
	// JUSTIFICATION
	// Benefit: ~15% speedup + required to index into RawTable from binary heap
	// Soundness: the caller must provide valid indexes
	unsafe fn update_heap_idx(&mut self, mapper: &[(usize, usize)]);
	// JUSTIFICATION
	// Benefit: ~15% speedup + required to index into RawTable from binary heap
	// Soundness: the caller must provide a valid index
	unsafe fn heap_idx_at(&self, map_idx: usize) -> usize;
	unsafe fn take_all(&mut self, indexes: Vec<usize>) -> ArrayRef;

	// JUSTIFICATION
	// Benefit: ~15% speedup + required to index into RawTable from binary heap
	// Soundness: the caller must provide valid indexes
	unsafe fn find_or_insert(
	&mut self,
	row_idx: usize,
	replace_idx: usize,
	map: &mut Vec<(usize, usize)>,
	) -> (usize, bool);
	}

Upgrade to hashbrown 0.16.1: migrate from hashbrown::raw::RawTable to hashbrown::hash_table::HashTable #13433

Description

What

Why

How

Progress

Activity

Dandandan commented on Dec 6, 2024

crepererum commented on Dec 9, 2024

hashbrown 0.14 & allocation size

ArrowHashTable

alamb commented on Dec 9, 2024

avantgardnerio commented on Dec 9, 2024

thinkharderdev commented on Dec 10, 2024

hashbrown 0.14 & allocation size

ArrowHashTable

crepererum commented on Dec 10, 2024

comphead commented on Jun 3, 2025

comphead commented on Jun 3, 2025

crepererum commented on Jun 12, 2025

Dandandan commented on Dec 29, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Upgrade to hashbrown 0.16.1: migrate from `hashbrown::raw::RawTable` to `hashbrown::hash_table::HashTable` #13433

`hashbrown` 0.14 & allocation size

`ArrowHashTable`

`hashbrown` 0.14 & allocation size

`ArrowHashTable`