Skip to content

Commit 543b806

Browse files
mapleFUclaude
andauthored
bench(arrow): add fsl/map interleave benchmark (#10322)
Add benchmark test cases for FixedSizeList and Map array interleave along with helper functions to create test arrays. # Which issue does this PR close? - Closes #10321 . # Rationale for this change Add benchmark for fsl/map interleave # What changes are included in this PR? Add benchmark for fsl/map interleave # Are these changes tested? No # Are there any user-facing changes? No --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 970391f commit 543b806

2 files changed

Lines changed: 88 additions & 2 deletions

File tree

arrow/benches/interleave_kernels.rs

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -150,6 +150,12 @@ fn add_benchmark(c: &mut Criterion) {
150150
)
151151
};
152152

153+
let fsl_i64 = create_primitive_fixed_size_list_array::<Int64Type>(8192, 0.0, 0.0, 5);
154+
let fsl_i64_nulls = create_primitive_fixed_size_list_array::<Int64Type>(8192, 0.1, 0.1, 5);
155+
156+
let map_str_i64 = create_string_map_array::<Int64Type>(8192, 0.0, 5, 8);
157+
let map_str_i64_nulls = create_string_map_array::<Int64Type>(8192, 0.1, 5, 8);
158+
153159
let ree_run_ends = Int32Array::from_iter_values((1..=64).map(|i| i * 16));
154160
let ree_i64_values = create_primitive_array::<Int64Type>(64, 0.0);
155161
let ree_i64 = RunArray::<Int32Type>::try_new(&ree_run_ends, &ree_i64_values).unwrap();
@@ -183,6 +189,10 @@ fn add_benchmark(c: &mut Criterion) {
183189
("list_view<i64>(0.1,0.1,20)", &list_view_i64),
184190
("list_view<i64>(0.0,0.0,20)", &list_view_i64_no_nulls),
185191
("list_view_overlapping<i64>(80x,20)", &list_view_overlapping),
192+
("fixed_size_list<i64,5>(0.0,0.0)", &fsl_i64),
193+
("fixed_size_list<i64,5>(0.1,0.1)", &fsl_i64_nulls),
194+
("map<utf8,i64>(0.0,5,8)", &map_str_i64),
195+
("map<utf8,i64>(0.1,5,8)", &map_str_i64_nulls),
186196
("ree_i32<i64>(64 runs)", &ree_i64),
187197
("ree_i32<dict<u32,utf8>>(64 runs)", &ree_dict),
188198
];

arrow/src/util/bench_util.rs

Lines changed: 78 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,17 +20,19 @@
2020
use crate::array::*;
2121
use crate::datatypes::*;
2222
use crate::util::test_util::seedable_rng;
23-
use arrow_buffer::{Buffer, IntervalMonthDayNano};
23+
use arrow_buffer::{Buffer, IntervalMonthDayNano, NullBuffer};
24+
use arrow_schema::Field;
2425
use half::f16;
2526
use rand::Rng;
2627
use rand::SeedableRng;
2728
use rand::distr::uniform::SampleUniform;
2829
use rand::rng;
2930
use rand::{
30-
distr::{Alphanumeric, Distribution, StandardUniform},
31+
distr::{Alphanumeric, Distribution, SampleString, StandardUniform},
3132
prelude::StdRng,
3233
};
3334
use std::ops::Range;
35+
use std::sync::Arc;
3436

3537
/// Creates an random (but fixed-seeded) array of a given size and null density
3638
pub fn create_primitive_array<T>(size: usize, null_density: f32) -> PrimitiveArray<T>
@@ -871,3 +873,77 @@ pub fn create_f64_array_with_seed(size: usize, nan_density: f32, seed: u64) -> F
871873
})
872874
.collect()
873875
}
876+
877+
/// Create a FixedSizeList array of primitive values
878+
///
879+
/// Arguments:
880+
/// - `size`: number of fixed-size lists in the array
881+
/// - `null_density`: density of nulls in the fixed-size list array (row-level nulls)
882+
/// - `value_null_density`: density of nulls in the primitive values inside each list
883+
/// - `list_size`: fixed size of each list element
884+
pub fn create_primitive_fixed_size_list_array<T>(
885+
size: usize,
886+
null_density: f32,
887+
value_null_density: f32,
888+
list_size: i32,
889+
) -> FixedSizeListArray
890+
where
891+
T: ArrowPrimitiveType,
892+
StandardUniform: Distribution<T::Native>,
893+
{
894+
let mut rng = seedable_rng();
895+
let list_size_usize = usize::try_from(list_size).expect("list_size must be non-negative");
896+
let values: PrimitiveArray<T> = (0..size * list_size_usize)
897+
.map(|_| {
898+
if rng.random::<f32>() < value_null_density {
899+
None
900+
} else {
901+
Some(rng.random())
902+
}
903+
})
904+
.collect();
905+
let field = Arc::new(Field::new("item", T::DATA_TYPE, value_null_density > 0.0));
906+
let nulls = (null_density > 0.0).then(|| {
907+
NullBuffer::new(arrow_buffer::BooleanBuffer::collect_bool(size, |_| {
908+
rng.random::<f32>() >= null_density
909+
}))
910+
});
911+
FixedSizeListArray::new(field, list_size, Arc::new(values), nulls)
912+
}
913+
914+
/// Create a Map array with string keys and primitive values
915+
///
916+
/// Arguments:
917+
/// - `size`: number of map entries in the array
918+
/// - `null_density`: density of nulls in the map array (row-level nulls)
919+
/// - `max_map_size`: maximum number of key-value pairs per map entry
920+
/// (actual size is random between 0 and max_map_size)
921+
/// - `key_len`: length of each random string key
922+
pub fn create_string_map_array<T>(
923+
size: usize,
924+
null_density: f32,
925+
max_map_size: usize,
926+
key_len: usize,
927+
) -> MapArray
928+
where
929+
T: ArrowPrimitiveType,
930+
StandardUniform: Distribution<T::Native>,
931+
{
932+
let mut rng = seedable_rng();
933+
let mut builder = MapBuilder::new(None, StringBuilder::new(), PrimitiveBuilder::<T>::new());
934+
for _ in 0..size {
935+
if rng.random::<f32>() < null_density {
936+
builder.append(false).unwrap();
937+
} else {
938+
let n = rng.random_range(0..=max_map_size);
939+
for _ in 0..n {
940+
builder
941+
.keys()
942+
.append_value(Alphanumeric.sample_string(&mut rng, key_len));
943+
builder.values().append_value(rng.random());
944+
}
945+
builder.append(true).unwrap();
946+
}
947+
}
948+
builder.finish()
949+
}

0 commit comments

Comments
 (0)