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

[crepererum]

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: change some hashbrown RawTable uses to HashTable #13514
• refactor: change some hashbrown RawTable uses to HashTable (round 2) #13524
• refactor: change some hashbrown RawTable uses to HashTable (round 3) #13658

[Dandandan]

Thanks @crepererum for the steady progress on this

[crepererum]

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:

1. data heap: a "heap" (which I guess is just a vector) to store the actual payload data
2. 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.
3. 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.

[alamb]

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.

[thinkharderdev]

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:

1. data heap: a "heap" (which I guess is just a vector) to store the actual payload data
2. 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.
3. 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.