Repository: pola-rs/polars
File: crates/polars-ops/src/frame/join/hash_join/single_keys.rs (Rust)

Code before:
use super::*;

pub(crate) fn build_tables<T, I>(keys: Vec<I>) -> Vec<PlHashMap<T, Vec<IdxSize>>>
where
    T: Send + Hash + Eq + Sync + Copy + AsU64,
    I: IntoIterator<Item = T> + Send + Sync + Clone,
    // <I as IntoIterator>::IntoIter: TrustedLen,
{
    let n_partitions = _set_partition_size();

    // We will create a hashtable in every thread.
    // We use the hash to partition the keys to the matching hashtable.
    // Every thread traverses all keys/hashes and ignores the ones that doesn't fall in that partition.
    POOL.install(|| {
        (0..n_partitions)
            .into_par_iter()
            .map(|partition_no| {
                let partition_no = partition_no as u64;

                let mut hash_tbl: PlHashMap<T, Vec<IdxSize>> =
                    PlHashMap::with_capacity(_HASHMAP_INIT_SIZE);

                let n_partitions = n_partitions as u64;
                let mut offset = 0;
                for keys in &keys {
                    let keys = keys.clone().into_iter();
                    let len = keys.size_hint().1.unwrap() as IdxSize;

                    let mut cnt = 0;
                    keys.for_each(|k| {
                        let idx = cnt + offset;
                        cnt += 1;

                        if this_partition(k.as_u64(), partition_no, n_partitions) {
                            let entry = hash_tbl.entry(k);

                            match entry {
                                Entry::Vacant(entry) => {
                                    entry.insert(vec![idx]);
                                },
                                Entry::Occupied(mut entry) => {
                                    let v = entry.get_mut();
                                    v.push(idx);
                                },
                            }
                        }
                    });
                    offset += len;
                }
                hash_tbl
            })
            .collect()
    })
}

// we determine the offset so that we later know which index to store in the join tuples
pub(super) fn probe_to_offsets<T, I>(probe: &[I]) -> Vec<usize>
where
    I: IntoIterator<Item = T> + Clone,
    // <I as IntoIterator>::IntoIter: TrustedLen,
    T: Send + Hash + Eq + Sync + Copy + AsU64,
{
    probe
        .iter()
        .map(|ph| ph.clone().into_iter().size_hint().1.unwrap())
        .scan(0, |state, val| {
            let out = *state;
            *state += val;
            Some(out)
        })
        .collect()
}


Review comment (performance):
Don't know if the compiler is able to optimize the indexing aways, otherwise we could use `get_unchecked_release`. Those elide bound checks in a release build.

Code after:
                }
                partition_sizes
            })
            .collect();

        // Compute output offsets with a cumulative sum.
        let mut per_thread_partition_offsets = vec![0; n_partitions * n_threads + 1];
        let mut partition_offsets = vec![0; n_partitions + 1];
        let mut cum_offset = 0;
        for p in 0..n_partitions {
            partition_offsets[p] = cum_offset;
            for t in 0..n_threads {
                per_thread_partition_offsets[t * n_partitions + p] = cum_offset;
                cum_offset += per_thread_partition_sizes[t][p];
            }
        }
        let num_keys = cum_offset;
        per_thread_partition_offsets[n_threads * n_partitions] = num_keys;
        partition_offsets[n_partitions] = num_keys;

        // FIXME: we wouldn't need this if we changed our interface to split the
        // input in this function, instead of taking a vec of iterators.
        let mut per_thread_input_offsets = vec![0; n_partitions];
        cum_offset = 0;
        for t in 0..n_threads {
            per_thread_input_offsets[t] = cum_offset;
            for p in 0..n_partitions {
                cum_offset += per_thread_partition_sizes[t][p];
            }
        }

        // Scatter values into partitions.
        let mut scatter_keys: Vec<T> = Vec::with_capacity(num_keys);
        let mut scatter_idxs: Vec<IdxSize> = Vec::with_capacity(num_keys);
        let scatter_keys_ptr = unsafe { SyncPtr::new(scatter_keys.as_mut_ptr()) };
        let scatter_idxs_ptr = unsafe { SyncPtr::new(scatter_idxs.as_mut_ptr()) };
        keys.into_par_iter()
            .enumerate()
            .for_each(|(t, key_portion)| {
                let mut partition_offsets =
                    per_thread_partition_offsets[t * n_partitions..(t + 1) * n_partitions].to_vec();
                for (i, key) in key_portion.into_iter().enumerate() {
                    unsafe {
                        let p = hash_to_partition(key.as_u64(), n_partitions);
                        let off = partition_offsets.get_unchecked_mut(p);
                        *scatter_keys_ptr.get().add(*off) = key;
                        *scatter_idxs_ptr.get().add(*off) =
                            (per_thread_input_offsets[t] + i) as IdxSize;
                        *off += 1;
                    }
                }