make hash function in RepartitionExec configurable #17648
Conversation
| let mut bitmap = build_side.left_data.visited_indices_bitmap().lock(); | ||
| left_indices.iter().flatten().for_each(|x| { | ||
| bitmap.set_bit(x as usize, true); | ||
| bitmap.set_bit(usize::try_from(x).expect("should fit"), true); |
|
cc @rkrishn7 |
Co-authored-by: Jonathan Chen <chenleejonathan@gmail.com>
| } | ||
|
|
||
| // Create hash buffer and compute hashes using DataFusion's internal algorithm | ||
| let mut hashes_buffer = vec![0u64; array_len]; |
There was a problem hiding this comment.
I will flag that previously the same vec was re-used with a capacity bump followed by a clear. Now we're creating a new one for each batch. It's some more allocations, but we're also pre-allocating the entire size, etc. I'm not sure if this will be measurable or not.
There was a problem hiding this comment.
I think we should maintain that behavior if possible. If the primary goal here is to encapsulate re-partitioning logic in this module, can we add a function for this? Something like:
/// Calculates the partition used by the repartition operator for each row. All arrays should have the same length.
fn compute_partition_indices(buf: &mut Vec<u64>, arrays: &[ArrayRef], num_partitions: usize) -> Result<()> {
buf.resize(arrays[0].len(), 0);
create_hashes(arrays, REPARTITION_RANDOM_STATE, buf);
buf.iter_mut().for_each(|hash| *hash %= num_partitions as u64);
}That way the repartitioning code can simply utilize this with the once allocated vector. But the dynamic filter can use the UDF, which still utilizes this under the hood
There was a problem hiding this comment.
Yes agreed, I think that's better. I'll cook something up.
| // Convert all arguments to arrays | ||
| let arrays = ColumnarValue::values_to_arrays(&args.args)?; | ||
|
|
||
| // Check that all arrays have the same length |
There was a problem hiding this comment.
I think ColumnarValue::values_to_arrays does this check already
| } | ||
|
|
||
| // Create hash buffer and compute hashes using DataFusion's internal algorithm | ||
| let mut hashes_buffer = vec![0u64; array_len]; |
There was a problem hiding this comment.
I think we should maintain that behavior if possible. If the primary goal here is to encapsulate re-partitioning logic in this module, can we add a function for this? Something like:
/// Calculates the partition used by the repartition operator for each row. All arrays should have the same length.
fn compute_partition_indices(buf: &mut Vec<u64>, arrays: &[ArrayRef], num_partitions: usize) -> Result<()> {
buf.resize(arrays[0].len(), 0);
create_hashes(arrays, REPARTITION_RANDOM_STATE, buf);
buf.iter_mut().for_each(|hash| *hash %= num_partitions as u64);
}That way the repartitioning code can simply utilize this with the once allocated vector. But the dynamic filter can use the UDF, which still utilizes this under the hood
Needed immediately for some testing needs in #17632. But this is long term useful so that users can customize the hash function used so that e.g. it is consistent in a distributed environment.