diff --git a/datafusion/common/benches/with_hashes.rs b/datafusion/common/benches/with_hashes.rs index 7fc4a29fc55ec..c4fda45e3e204 100644 --- a/datafusion/common/benches/with_hashes.rs +++ b/datafusion/common/benches/with_hashes.rs @@ -160,6 +160,31 @@ fn criterion_benchmark(c: &mut Criterion) { do_hash_test(b, &arrays); }); } + // Apply additional null-density cases to all Int32 dictionary + // benchmarks, including any added in the future. + if let Some(dict) = array.as_any().downcast_ref::>() { + for null_percent in [50, 97] { + let nullable_array = with_null_keys(dict, null_percent as f32 / 100.0); + + c.bench_function( + &format!("{name}: single, {null_percent}% null keys"), + |b| { + do_hash_test(b, std::slice::from_ref(&nullable_array)); + }, + ); + c.bench_function( + &format!("{name}: multiple, {null_percent}% null keys"), + |b| { + let arrays = vec![ + nullable_array.clone(), + nullable_array.clone(), + nullable_array.clone(), + ]; + do_hash_test(b, &arrays); + }, + ); + } + } } } @@ -178,8 +203,13 @@ fn create_null_mask(len: usize) -> NullBuffer where StandardUniform: Distribution, { + create_null_mask_with_density(len, 0.03).expect("should be nulls in buffer") +} + +/// Returns a random null mask where each slot is null with probability +/// `null_density`, or `None` if no slot ended up null +fn create_null_mask_with_density(len: usize, null_density: f32) -> Option { let mut rng = make_rng(); - let null_density = 0.03; let mut builder = NullBufferBuilder::new(len); for _ in 0..len { if rng.random::() < null_density { @@ -188,7 +218,21 @@ where builder.append_non_null(); } } - builder.finish().expect("should be nulls in buffer") + builder.finish() +} + +/// Returns `dict` with a new random null mask on its keys, where each key is null +/// with probability `null_density`. Any existing key nulls are replaced, and the +/// key and value buffers are shared rather than copied +fn with_null_keys( + dict: &DictionaryArray, + null_density: f32, +) -> ArrayRef { + let keys = PrimitiveArray::::new( + dict.keys().values().clone(), + create_null_mask_with_density(dict.len(), null_density), + ); + Arc::new(DictionaryArray::new(keys, Arc::clone(dict.values()))) } // Returns a new array that is the same as array, but with nulls