This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-11928-5746b443f70b4e55676be59cea317501a94ec810 in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
commit 35b20e63ddf211d5d011ff28bea44ae400fbf5c2 Author: RIchard Baah <[email protected]> AuthorDate: Tue Oct 27 17:26:48 2026 +0000 refactor(bench) : reuse shared array builders across benches (#11928) # Which issue does this PR close? - Closes #11927 # Rationale for this change Several benchmarks define their own random array and index builders. Keeping these helpers in one place avoids differences in seeds, null handling, and generated values. # What changes are included in this PR? - Add shared index, decimal, and seeded fixed-size-binary builders to `bench_util`. - Update the Arrow and Parquet benchmarks to use the shared builders and remove their local copies. # Are these changes tested? No new tests. # Are there any significant user-facing changes? No. --- arrow/benches/cast_kernels.rs | 85 +++++++-------------- arrow/benches/coalesce_kernels.rs | 26 +------ arrow/benches/partition_kernels.rs | 22 ++---- arrow/benches/primitive_run_take.rs | 18 ----- arrow/benches/sort_kernel.rs | 65 +++++++--------- arrow/benches/take_kernels.rs | 136 +++++++++++++++------------------- arrow/src/util/bench_util.rs | 54 +++++++++++++- parquet/benches/parquet_round_trip.rs | 78 +++---------------- 8 files changed, 188 insertions(+), 296 deletions(-) diff --git a/arrow/benches/cast_kernels.rs b/arrow/benches/cast_kernels.rs index f29bcca011..e7f5c7d66c 100644 --- a/arrow/benches/cast_kernels.rs +++ b/arrow/benches/cast_kernels.rs @@ -19,7 +19,7 @@ extern crate criterion; use criterion::Criterion; use rand::RngExt; -use rand::distr::{Distribution, StandardUniform, Uniform}; +use rand::distr::Uniform; use std::hint; use chrono::{DateTime, NaiveDate, NaiveDateTime}; @@ -31,14 +31,6 @@ use arrow::datatypes::*; use arrow::util::bench_util::*; use arrow::util::test_util::seedable_rng; -fn build_array<T: ArrowPrimitiveType>(size: usize) -> ArrayRef -where - StandardUniform: Distribution<T::Native>, -{ - let array = create_primitive_array::<T>(size, 0.1); - Arc::new(array) -} - fn build_utf8_date_array(size: usize, with_nulls: bool) -> ArrayRef { use chrono::NaiveDate; @@ -81,38 +73,6 @@ fn build_utf8_date_time_array(size: usize, with_nulls: bool) -> ArrayRef { Arc::new(builder.finish()) } -fn build_decimal128_array(size: usize, precision: u8, scale: i8) -> ArrayRef { - let mut rng = seedable_rng(); - let mut builder = Decimal128Builder::with_capacity(size); - - for _ in 0..size { - builder.append_value(rng.random_range::<i128, _>(0..1000000000)); - } - Arc::new( - builder - .finish() - .with_precision_and_scale(precision, scale) - .unwrap(), - ) -} - -fn build_decimal256_array(size: usize, precision: u8, scale: i8) -> ArrayRef { - let mut rng = seedable_rng(); - let mut builder = Decimal256Builder::with_capacity(size); - let mut bytes = [0; 32]; - for _ in 0..size { - let num = rng.random_range::<i128, _>(0..1000000000); - bytes[0..16].clone_from_slice(&num.to_le_bytes()); - builder.append_value(i256::from_le_bytes(bytes)); - } - Arc::new( - builder - .finish() - .with_precision_and_scale(precision, scale) - .unwrap(), - ) -} - fn build_string_array(size: usize) -> ArrayRef { let mut builder = StringBuilder::new(); for v in 0..size { @@ -258,32 +218,43 @@ fn cast_array(array: &ArrayRef, to_type: DataType) { } fn add_benchmark(c: &mut Criterion) { - let i32_array = build_array::<Int32Type>(512); - let i32_array_8192 = build_array::<Int32Type>(8192); + let i32_array: ArrayRef = Arc::new(create_primitive_array::<Int32Type>(512, 0.1)); + let i32_array_8192: ArrayRef = Arc::new(create_primitive_array::<Int32Type>(8192, 0.1)); let bool_array: ArrayRef = Arc::new(create_boolean_array(8192, 0.1, 0.5)); let bool_array_no_nulls: ArrayRef = Arc::new(create_boolean_array(8192, 0.0, 0.5)); - let i64_array = build_array::<Int64Type>(512); + let i64_array: ArrayRef = Arc::new(create_primitive_array::<Int64Type>(512, 0.1)); let i64_within_1e6 = create_primitive_array_range::<Int64Type>(512, 0.1, -1_000_000..1_000_000); let mut values: Vec<Option<i64>> = i64_within_1e6.iter().collect(); *values.last_mut().unwrap() = Some(10_000_000_000_000_000); let i64_late_overflow = Arc::new(Int64Array::from(values)) as ArrayRef; let i64_within_1e6 = Arc::new(i64_within_1e6) as ArrayRef; - let f32_array = build_array::<Float32Type>(512); - let f32_utf8_array = cast(&build_array::<Float32Type>(512), &DataType::Utf8).unwrap(); - let i32_utf8_array = cast(&build_array::<Int32Type>(512), &DataType::Utf8).unwrap(); - - let f64_array = build_array::<Float64Type>(512); - let date64_array = build_array::<Date64Type>(512); - let date32_array = build_array::<Date32Type>(512); - let time32s_array = build_array::<Time32SecondType>(512); - let time64ns_array = build_array::<Time64NanosecondType>(512); - let time_ns_array = build_array::<TimestampNanosecondType>(512); - let time_ms_array = build_array::<TimestampMillisecondType>(512); + let f32_array: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(512, 0.1)); + let f32_utf8_array = cast( + &create_primitive_array::<Float32Type>(512, 0.1), + &DataType::Utf8, + ) + .unwrap(); + let i32_utf8_array = cast( + &create_primitive_array::<Int32Type>(512, 0.1), + &DataType::Utf8, + ) + .unwrap(); + + let f64_array: ArrayRef = Arc::new(create_primitive_array::<Float64Type>(512, 0.1)); + let date64_array: ArrayRef = Arc::new(create_primitive_array::<Date64Type>(512, 0.1)); + let date32_array: ArrayRef = Arc::new(create_primitive_array::<Date32Type>(512, 0.1)); + let time32s_array: ArrayRef = Arc::new(create_primitive_array::<Time32SecondType>(512, 0.1)); + let time64ns_array: ArrayRef = + Arc::new(create_primitive_array::<Time64NanosecondType>(512, 0.1)); + let time_ns_array: ArrayRef = + Arc::new(create_primitive_array::<TimestampNanosecondType>(512, 0.1)); + let time_ms_array: ArrayRef = + Arc::new(create_primitive_array::<TimestampMillisecondType>(512, 0.1)); let utf8_date_array = build_utf8_date_array(512, true); let utf8_date_time_array = build_utf8_date_time_array(512, true); - let decimal128_array = build_decimal128_array(8_000, 10, 3); - let decimal256_array = build_decimal256_array(8_000, 50, 3); + let decimal128_array: ArrayRef = Arc::new(create_decimal128_array(8_000, 10, 3)); + let decimal256_array: ArrayRef = Arc::new(create_decimal256_array(8_000, 50, 3)); let string_array = build_string_array(512); let wide_string_array = cast(&string_array, &DataType::LargeUtf8).unwrap(); diff --git a/arrow/benches/coalesce_kernels.rs b/arrow/benches/coalesce_kernels.rs index faaee8f6e6..8f3693c36f 100644 --- a/arrow/benches/coalesce_kernels.rs +++ b/arrow/benches/coalesce_kernels.rs @@ -22,8 +22,6 @@ use arrow::util::data_gen::create_random_array; use std::sync::Arc; use arrow::array::*; -use arrow_array::FixedSizeBinaryArray; -use arrow_array::builder::FixedSizeBinaryBuilder; use arrow_array::types::{Float64Type, Int32Type, TimestampNanosecondType}; use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use arrow_select::coalesce::BatchCoalescer; @@ -1502,10 +1500,10 @@ impl DataStreamBuilder { .collect::<DictionaryArray<Int32Type>>(), ) } - DataType::FixedSizeBinary(size) => Arc::new(create_fixed_size_binary_array( + DataType::FixedSizeBinary(size) => Arc::new(create_fsb_array_with_seed( self.batch_size, self.null_density, - *size, + *size as usize, seed, )), DataType::Boolean => Arc::new(create_boolean_array_with_seed( @@ -1536,23 +1534,3 @@ impl DataStreamBuilder { } } } - -fn create_fixed_size_binary_array( - batch_size: usize, - null_density: f32, - value_length: i32, - seed: u64, -) -> FixedSizeBinaryArray { - let size = value_length as usize; - let mut rng = StdRng::seed_from_u64(seed); - let mut builder = FixedSizeBinaryBuilder::with_capacity(batch_size, value_length); - for _ in 0..batch_size { - if rng.random::<f32>() < null_density { - builder.append_null(); - } else { - let value: Vec<u8> = (0..size).map(|_| rng.random::<u8>()).collect(); - builder.append_value(&value).unwrap(); - } - } - builder.finish() -} diff --git a/arrow/benches/partition_kernels.rs b/arrow/benches/partition_kernels.rs index 0d3052b261..8c908c3b82 100644 --- a/arrow/benches/partition_kernels.rs +++ b/arrow/benches/partition_kernels.rs @@ -25,19 +25,9 @@ use arrow::{ }; use arrow_ord::partition::partition; use criterion::Criterion; -use rand::distr::{Distribution, StandardUniform}; use std::hint; use std::sync::Arc; -fn create_array<T: ArrowPrimitiveType>(size: usize, with_nulls: bool) -> ArrayRef -where - StandardUniform: Distribution<T::Native>, -{ - let null_density = if with_nulls { 0.5 } else { 0.0 }; - let array = create_primitive_array::<T>(size, null_density); - Arc::new(array) -} - fn bench_partition(sorted_columns: &[ArrayRef]) { hint::black_box(partition(sorted_columns).unwrap().ranges()); } @@ -60,14 +50,16 @@ fn create_sorted_low_cardinality_data(length: usize) -> Vec<ArrayRef> { } fn create_sorted_float_data(pow: u32, with_nulls: bool) -> Vec<ArrayRef> { + let size = 2u64.pow(pow) as usize; + let null_density = if with_nulls { 0.5 } else { 0.0 }; lexsort( &[ SortColumn { - values: create_array::<Float64Type>(2u64.pow(pow) as usize, with_nulls), + values: Arc::new(create_primitive_array::<Float64Type>(size, null_density)), options: None, }, SortColumn { - values: create_array::<Float64Type>(2u64.pow(pow) as usize, with_nulls), + values: Arc::new(create_primitive_array::<Float64Type>(size, null_density)), options: None, }, ], @@ -77,14 +69,16 @@ fn create_sorted_float_data(pow: u32, with_nulls: bool) -> Vec<ArrayRef> { } fn create_sorted_data(pow: u32, with_nulls: bool) -> Vec<ArrayRef> { + let size = 2u64.pow(pow) as usize; + let null_density = if with_nulls { 0.5 } else { 0.0 }; lexsort( &[ SortColumn { - values: create_array::<UInt8Type>(2u64.pow(pow) as usize, with_nulls), + values: Arc::new(create_primitive_array::<UInt8Type>(size, null_density)), options: None, }, SortColumn { - values: create_array::<UInt8Type>(2u64.pow(pow) as usize, with_nulls), + values: Arc::new(create_primitive_array::<UInt8Type>(size, null_density)), options: None, }, ], diff --git a/arrow/benches/primitive_run_take.rs b/arrow/benches/primitive_run_take.rs index ab6514fd9b..96ff676ca9 100644 --- a/arrow/benches/primitive_run_take.rs +++ b/arrow/benches/primitive_run_take.rs @@ -15,30 +15,12 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::UInt32Builder; use arrow::compute::take; use arrow::datatypes::{Int32Type, Int64Type}; use arrow::util::bench_util::*; -use arrow::util::test_util::seedable_rng; -use arrow_array::UInt32Array; use criterion::{Criterion, criterion_group, criterion_main}; -use rand::RngExt; use std::hint; -fn create_random_index(size: usize, null_density: f32, max_value: usize) -> UInt32Array { - let mut rng = seedable_rng(); - let mut builder = UInt32Builder::with_capacity(size); - for _ in 0..size { - if rng.random::<f32>() < null_density { - builder.append_null(); - } else { - let value = rng.random_range::<u32, _>(0u32..max_value as u32); - builder.append_value(value); - } - } - builder.finish() -} - fn criterion_benchmark(c: &mut Criterion) { let mut group = c.benchmark_group("primitive_run_take"); diff --git a/arrow/benches/sort_kernel.rs b/arrow/benches/sort_kernel.rs index ec9b6fd736..b48a221bf7 100644 --- a/arrow/benches/sort_kernel.rs +++ b/arrow/benches/sort_kernel.rs @@ -28,19 +28,6 @@ use arrow::{array::*, datatypes::Float32Type}; use arrow_ord::rank::rank; use std::hint; -fn create_f32_array(size: usize, with_nulls: bool) -> ArrayRef { - let null_density = if with_nulls { 0.5 } else { 0.0 }; - let array = create_primitive_array::<Float32Type>(size, null_density); - Arc::new(array) -} - -fn create_bool_array(size: usize, with_nulls: bool) -> ArrayRef { - let null_density = if with_nulls { 0.5 } else { 0.0 }; - let true_density = 0.5; - let array = create_boolean_array(size, null_density, true_density); - Arc::new(array) -} - fn bench_sort(array: &dyn Array) { hint::black_box(sort(array, None).unwrap()); } @@ -89,13 +76,13 @@ fn add_benchmark(c: &mut Criterion) { b.iter(|| bench_sort_to_indices(&arr, None)) }); - let arr = create_f32_array(2_usize.pow(12), false); + let arr: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2_usize.pow(12), 0.0)); c.bench_function("sort f32 2^12", |b| b.iter(|| bench_sort(&arr))); c.bench_function("sort f32 to indices 2^12", |b| { b.iter(|| bench_sort_to_indices(&arr, None)) }); - let arr = create_f32_array(2usize.pow(12), true); + let arr: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.5)); c.bench_function("sort f32 nulls 2^12", |b| b.iter(|| bench_sort(&arr))); c.bench_function("sort f32 nulls to indices 2^12", |b| { b.iter(|| bench_sort_to_indices(&arr, None)) @@ -219,72 +206,72 @@ fn add_benchmark(c: &mut Criterion) { b.iter(|| bench_sort_to_indices(&run_encoded_array, None)) }); - let arr_a = create_f32_array(2usize.pow(10), false); - let arr_b = create_f32_array(2usize.pow(10), false); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(10), 0.0)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(10), 0.0)); c.bench_function("lexsort (f32, f32) 2^10", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, None)) }); - let arr_a = create_f32_array(2usize.pow(12), false); - let arr_b = create_f32_array(2usize.pow(12), false); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); c.bench_function("lexsort (f32, f32) 2^12", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, None)) }); - let arr_a = create_f32_array(2usize.pow(10), true); - let arr_b = create_f32_array(2usize.pow(10), true); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(10), 0.5)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(10), 0.5)); c.bench_function("lexsort (f32, f32) nulls 2^10", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, None)) }); - let arr_a = create_f32_array(2usize.pow(12), true); - let arr_b = create_f32_array(2usize.pow(12), true); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.5)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.5)); c.bench_function("lexsort (f32, f32) nulls 2^12", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, None)) }); - let arr_a = create_bool_array(2usize.pow(12), false); - let arr_b = create_bool_array(2usize.pow(12), false); + let arr_a: ArrayRef = Arc::new(create_boolean_array(2usize.pow(12), 0.0, 0.5)); + let arr_b: ArrayRef = Arc::new(create_boolean_array(2usize.pow(12), 0.0, 0.5)); c.bench_function("lexsort (bool, bool) 2^12", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, None)) }); - let arr_a = create_bool_array(2usize.pow(12), true); - let arr_b = create_bool_array(2usize.pow(12), true); + let arr_a: ArrayRef = Arc::new(create_boolean_array(2usize.pow(12), 0.5, 0.5)); + let arr_b: ArrayRef = Arc::new(create_boolean_array(2usize.pow(12), 0.5, 0.5)); c.bench_function("lexsort (bool, bool) nulls 2^12", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, None)) }); - let arr_a = create_f32_array(2usize.pow(12), false); - let arr_b = create_f32_array(2usize.pow(12), false); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); c.bench_function("lexsort (f32, f32) 2^12 limit 10", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, Some(10))) }); - let arr_a = create_f32_array(2usize.pow(12), false); - let arr_b = create_f32_array(2usize.pow(12), false); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); c.bench_function("lexsort (f32, f32) 2^12 limit 100", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, Some(100))) }); - let arr_a = create_f32_array(2usize.pow(12), false); - let arr_b = create_f32_array(2usize.pow(12), false); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); c.bench_function("lexsort (f32, f32) 2^12 limit 1000", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, Some(1000))) }); - let arr_a = create_f32_array(2usize.pow(12), false); - let arr_b = create_f32_array(2usize.pow(12), false); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); c.bench_function("lexsort (f32, f32) 2^12 limit 2^12", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, Some(2usize.pow(12)))) }); - let arr_a = create_f32_array(2usize.pow(12), true); - let arr_b = create_f32_array(2usize.pow(12), true); + let arr_a: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.5)); + let arr_b: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.5)); c.bench_function("lexsort (f32, f32) nulls 2^12 limit 10", |b| { b.iter(|| bench_lexsort(&arr_a, &arr_b, Some(10))) @@ -299,12 +286,12 @@ fn add_benchmark(c: &mut Criterion) { b.iter(|| bench_lexsort(&arr_a, &arr_b, Some(2usize.pow(12)))) }); - let arr = create_f32_array(2usize.pow(12), false); + let arr: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.0)); c.bench_function("rank f32 2^12", |b| { b.iter(|| hint::black_box(rank(&arr, None).unwrap())) }); - let arr = create_f32_array(2usize.pow(12), true); + let arr: ArrayRef = Arc::new(create_primitive_array::<Float32Type>(2usize.pow(12), 0.5)); c.bench_function("rank f32 nulls 2^12", |b| { b.iter(|| hint::black_box(rank(&arr, None).unwrap())) }); diff --git a/arrow/benches/take_kernels.rs b/arrow/benches/take_kernels.rs index 2e93028e01..b951b522ef 100644 --- a/arrow/benches/take_kernels.rs +++ b/arrow/benches/take_kernels.rs @@ -30,20 +30,6 @@ use arrow::{array::*, util::bench_util::*}; use std::hint; use std::sync::Arc; -fn create_random_index(size: usize, null_density: f32) -> UInt32Array { - let mut rng = seedable_rng(); - let mut builder = UInt32Builder::with_capacity(size); - for _ in 0..size { - if rng.random::<f32>() < null_density { - builder.append_null(); - } else { - let value = rng.random_range::<u32, _>(0u32..size as u32); - builder.append_value(value); - } - } - builder.finish() -} - fn bench_take(values: &dyn Array, indices: &UInt32Array) { hint::black_box(take(values, indices, None).unwrap()); } @@ -168,40 +154,40 @@ fn create_dense_union(size: usize) -> UnionArray { fn add_benchmark(c: &mut Criterion) { let values = create_primitive_array::<Int32Type>(512, 0.0); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take i32 512", |b| b.iter(|| bench_take(&values, &indices))); let values = create_primitive_array::<Int32Type>(1024, 0.0); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take i32 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take i32 null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_array::<Int32Type>(1024, 0.5); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take i32 null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take i32 null values null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_array::<Int32Type>(512, 0.0); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take check bounds i32 512", |b| { b.iter(|| bench_take_bounds_check(&values, &indices)) }); let values = create_primitive_array::<Int32Type>(1024, 0.0); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take check bounds i32 1024", |b| { b.iter(|| bench_take_bounds_check(&values, &indices)) }); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take check bounds i32 null indices 1024", |b| { b.iter(|| bench_take_bounds_check(&values, &indices)) }); @@ -216,212 +202,212 @@ fn add_benchmark(c: &mut Criterion) { ); let values = create_boolean_array(512, 0.0, 0.5); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take bool 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_boolean_array(1024, 0.0, 0.5); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take bool 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take bool null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_boolean_array(1024, 0.5, 0.5); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take bool null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_boolean_array(1024, 0.5, 0.5); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take bool null values null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_array::<i32>(512, 0.0); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take str 512", |b| b.iter(|| bench_take(&values, &indices))); let values = create_string_array::<i32>(1024, 0.0); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take str 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_array::<i32>(512, 0.0); - let indices = create_random_index(512, 0.5); + let indices = create_random_index(512, 0.5, 512); c.bench_function("take str null indices 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_array::<i32>(1024, 0.0); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take str null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_array::<i32>(1024, 0.5); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take str null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_array::<i32>(1024, 0.5); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take str null values null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_view_array(512, 0.0); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take stringview 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_view_array(1024, 0.0); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take stringview 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_view_array(512, 0.0); - let indices = create_random_index(512, 0.5); + let indices = create_random_index(512, 0.5, 512); c.bench_function("take stringview null indices 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_view_array(1024, 0.0); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take stringview null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_view_array(1024, 0.5); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take stringview null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_view_array(1024, 0.5); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take stringview null values null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_array::<i32, Int32Type>(512, 0.0, 0.0, 20); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take list i32 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_array::<i32, Int32Type>(1024, 0.0, 0.0, 20); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take list i32 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_array::<i32, Int32Type>(1024, 0.5, 0.0, 20); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take list i32 null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_array::<i32, Int32Type>(1024, 0.0, 0.0, 20); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take list i32 null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_array::<i32, Int32Type>(1024, 0.5, 0.5, 20); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take list i32 null values null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_view_array::<i32, Int32Type>(512, 0.0, 0.0, 20); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take listview i32 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_view_array::<i32, Int32Type>(1024, 0.0, 0.0, 20); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take listview i32 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_view_array::<i32, Int32Type>(1024, 0.5, 0.0, 20); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take listview i32 null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_view_array::<i32, Int32Type>(1024, 0.0, 0.0, 20); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take listview i32 null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_list_view_array::<i32, Int32Type>(1024, 0.5, 0.5, 20); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take listview i32 null values null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_run_array::<Int32Type, Int32Type>(1024, 512); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function( "take primitive run logical len: 1024, physical len: 512, indices: 1024", |b| b.iter(|| bench_take(&values, &indices)), ); let values = create_string_run_array(1024, 128); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function( "take string run logical len: 1024, physical len: 128, indices: 1024", |b| b.iter(|| bench_take(&values, &indices)), ); let values = create_string_run_array(1024, 512); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function( "take string run logical len: 1024, physical len: 512, indices: 1024", |b| b.iter(|| bench_take(&values, &indices)), ); let values = create_string_run_array(1024, 512); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function( "take string run logical len: 1024, physical len: 512, null indices: 1024", |b| b.iter(|| bench_take(&values, &indices)), ); let values = create_fsb_array(1024, 0.0, 12); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take fsb value len: 12, indices: 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_fsb_array(1024, 0.5, 12); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take fsb value len: 12, null values, indices: 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_fsb_array(1024, 0.0, 16); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take fsb value optimized len: 16, indices: 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_fsb_array(1024, 0.5, 16); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function( "take fsb value optimized len: 16, null values, indices: 1024", |b| b.iter(|| bench_take(&values, &indices)), @@ -435,7 +421,7 @@ fn add_benchmark(c: &mut Criterion) { DataType::Boolean, ]; let batch = make_record_batch(create_columns(&types, 1024, 0.0)); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take_record_batch 5 primitive cols no nulls 1024", |b| { b.iter(|| bench_take_record_batch(&batch, &indices)) }); @@ -449,7 +435,7 @@ fn add_benchmark(c: &mut Criterion) { DataType::FixedSizeBinary(16), ]; let batch = make_record_batch(create_columns(&types, 1024, 0.0)); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function( "take_record_batch 6 string/binary cols no nulls 1024", |b| b.iter(|| bench_take_record_batch(&batch, &indices)), @@ -465,7 +451,7 @@ fn add_benchmark(c: &mut Criterion) { DataType::Binary, ]; let batch = make_record_batch(create_columns(&types, 1024, 0.5)); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take_record_batch 7 mixed cols null values 1024", |b| { b.iter(|| bench_take_record_batch(&batch, &indices)) }); @@ -480,7 +466,7 @@ fn add_benchmark(c: &mut Criterion) { DataType::Binary, ]; let batch = make_record_batch(create_columns(&types, 1024, 0.5)); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function( "take_record_batch 7 mixed cols null values null indices 1024", |b| b.iter(|| bench_take_record_batch(&batch, &indices)), @@ -488,86 +474,86 @@ fn add_benchmark(c: &mut Criterion) { // FixedSizeList — list_size=8 (power-of-two) and list_size=22 (arbitrary/dynamic length) let values = create_primitive_fixed_size_list_array::<Int32Type>(1024, 0.0, 0.0, 8); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take fixed_size_list<i32>[8] 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_fixed_size_list_array::<Int32Type>(1024, 0.5, 0.0, 8); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take fixed_size_list<i32>[8] null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_fixed_size_list_array::<Int32Type>(1024, 0.0, 0.0, 8); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take fixed_size_list<i32>[8] null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_fixed_size_list_array::<Int32Type>(1024, 0.0, 0.0, 22); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take fixed_size_list<i32>[22] 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_fixed_size_list_array::<Int32Type>(1024, 0.5, 0.0, 22); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take fixed_size_list<i32>[22] null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_primitive_fixed_size_list_array::<Int32Type>(1024, 0.0, 0.0, 22); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take fixed_size_list<i32>[22] null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); // Map let values = create_string_map_array::<Int32Type>(512, 0.0, 10, 8); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take map<str, i32> 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_map_array::<Int32Type>(1024, 0.0, 10, 8); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take map<str, i32> 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_map_array::<Int32Type>(1024, 0.5, 10, 8); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take map<str, i32> null values 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_string_map_array::<Int32Type>(1024, 0.0, 10, 8); - let indices = create_random_index(1024, 0.5); + let indices = create_random_index(1024, 0.5, 1024); c.bench_function("take map<str, i32> null indices 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_sparse_union(512); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take sparse union 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_sparse_union(1024); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take sparse union 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_dense_union(512); - let indices = create_random_index(512, 0.0); + let indices = create_random_index(512, 0.0, 512); c.bench_function("take dense union 512", |b| { b.iter(|| bench_take(&values, &indices)) }); let values = create_dense_union(1024); - let indices = create_random_index(1024, 0.0); + let indices = create_random_index(1024, 0.0, 1024); c.bench_function("take dense union 1024", |b| { b.iter(|| bench_take(&values, &indices)) }); diff --git a/arrow/src/util/bench_util.rs b/arrow/src/util/bench_util.rs index ec757b1a07..7d0ca0b60f 100644 --- a/arrow/src/util/bench_util.rs +++ b/arrow/src/util/bench_util.rs @@ -20,7 +20,7 @@ use crate::array::*; use crate::datatypes::*; use crate::util::test_util::seedable_rng; -use arrow_buffer::{Buffer, IntervalMonthDayNano, NullBuffer}; +use arrow_buffer::{Buffer, IntervalMonthDayNano, NullBuffer, i256}; use arrow_schema::Field; use half::f16; use rand::RngExt; @@ -193,6 +193,45 @@ where .collect() } +/// Creates a random (but fixed-seeded) [`UInt32Array`] of indices in `0..max_value`. +/// +/// Intended for use as the `indices` argument to take kernels. `null_density` +/// controls the fraction of null indices. +pub fn create_random_index(size: usize, null_density: f32, max_value: usize) -> UInt32Array { + create_primitive_array_range::<UInt32Type>(size, null_density, 0..max_value as u32) +} + +/// Creates a random (but fixed-seeded) [`Decimal128Array`] with values drawn +/// from `0..1_000_000_000`. +pub fn create_decimal128_array(size: usize, precision: u8, scale: i8) -> Decimal128Array { + let mut rng = seedable_rng(); + let mut builder = Decimal128Builder::with_capacity(size); + for _ in 0..size { + builder.append_value(rng.random_range::<i128, _>(0..1_000_000_000)); + } + builder + .finish() + .with_precision_and_scale(precision, scale) + .unwrap() +} + +/// Creates a random (but fixed-seeded) [`Decimal256Array`] with values drawn +/// from `0..1_000_000_000` (sign-extended into i256). +pub fn create_decimal256_array(size: usize, precision: u8, scale: i8) -> Decimal256Array { + let mut rng = seedable_rng(); + let mut builder = Decimal256Builder::with_capacity(size); + let mut bytes = [0; 32]; + for _ in 0..size { + let num = rng.random_range::<i128, _>(0..1_000_000_000); + bytes[0..16].clone_from_slice(&num.to_le_bytes()); + builder.append_value(i256::from_le_bytes(bytes)); + } + builder + .finish() + .with_precision_and_scale(precision, scale) + .unwrap() +} + /// Creates a random (but fixed-seeded) string array of a given size and null density. /// /// Strings have a random length @@ -782,7 +821,18 @@ pub fn create_binary_array_with_len_range_and_prefix_and_seed<Offset: OffsetSize /// Creates an random (but fixed-seeded) array of a given size and null density pub fn create_fsb_array(size: usize, null_density: f32, value_len: usize) -> FixedSizeBinaryArray { - let rng = &mut seedable_rng(); + create_fsb_array_with_seed(size, null_density, value_len, 42) +} + +/// Creates a random [`FixedSizeBinaryArray`] of a given size and null density +/// generated using the provided seed. +pub fn create_fsb_array_with_seed( + size: usize, + null_density: f32, + value_len: usize, + seed: u64, +) -> FixedSizeBinaryArray { + let rng = &mut StdRng::seed_from_u64(seed); FixedSizeBinaryArray::try_from_sparse_iter_with_size( (0..size).map(|_| { diff --git a/parquet/benches/parquet_round_trip.rs b/parquet/benches/parquet_round_trip.rs index bf82a8a952..8601c6e4b7 100644 --- a/parquet/benches/parquet_round_trip.rs +++ b/parquet/benches/parquet_round_trip.rs @@ -18,21 +18,16 @@ use arrow::array::{ArrayRef, RecordBatch}; use arrow::datatypes::{DataType, Field, Float32Type, Float64Type, Int32Type, Int64Type, Schema}; use arrow::util::bench_util::{ - create_binary_array_with_len_range_and_prefix_and_seed, create_primitive_array_with_seed, - create_string_array_with_len_range_and_prefix_and_seed, + create_binary_array_with_len_range_and_prefix_and_seed, create_fsb_array_with_seed, + create_primitive_array_with_seed, create_string_array_with_len_range_and_prefix_and_seed, + create_string_view_array_with_len_range_and_seed, }; -use arrow_array::{FixedSizeBinaryArray, StringViewArray}; use bytes::Bytes; use criterion::{Criterion, criterion_group, criterion_main}; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; use parquet::basic::Encoding; use parquet::file::properties::WriterProperties; -use rand::{ - RngExt, SeedableRng, - distr::{Alphanumeric, StandardUniform}, - prelude::StdRng, -}; use std::sync::Arc; #[derive(Copy, Clone)] @@ -47,61 +42,6 @@ pub enum ColumnType { Double, } -// arrow::util::bench_util::create_fsb_array with a seed - -/// Creates a random (but fixed-seeded) array of fixed size with a given null density and length -fn create_fsb_array_with_seed( - size: usize, - null_density: f32, - fixed_len: i32, - seed: u64, -) -> FixedSizeBinaryArray { - let mut rng = StdRng::seed_from_u64(seed); - - let rng = &mut rng; - FixedSizeBinaryArray::try_from_sparse_iter_with_size( - (0..size).map(|_| { - if rng.random::<f32>() < null_density { - None - } else { - let value = rng - .sample_iter::<u8, _>(StandardUniform) - .take(fixed_len as usize) - .collect::<Vec<u8>>(); - Some(value) - } - }), - fixed_len, - ) - .unwrap() -} - -// arrow::util::bench_util::create_string_view_array_with_max_len with a seed - -/// Creates a random (but fixed-seeded) array of rand size with a given max size, null density and length -pub fn create_string_view_array_with_seed( - size: usize, - null_density: f32, - max_str_len: usize, - seed: u64, -) -> StringViewArray { - let mut rng = StdRng::seed_from_u64(seed); - - let rng = &mut rng; - (0..size) - .map(|_| { - if rng.random::<f32>() < null_density { - None - } else { - let str_len = rng.random_range(max_str_len / 2..max_str_len); - let value = rng.sample_iter(&Alphanumeric).take(str_len).collect(); - let value = String::from_utf8(value).unwrap(); - Some(value) - } - }) - .collect() -} - fn schema(column_type: ColumnType, num_columns: usize) -> Arc<Schema> { let field_type = match column_type { ColumnType::Binary(_) => DataType::Binary, @@ -161,10 +101,10 @@ fn create_batch( ColumnType::StringView(max_str_len) => { for i in 0..num_columns { let array_seed = seed * num_columns + i; - let array = create_string_view_array_with_seed( + let array = create_string_view_array_with_len_range_and_seed( num_rows, null_density, - max_str_len, + max_str_len / 2..max_str_len, array_seed as u64, ); arrays.push(Arc::new(array)); @@ -173,8 +113,12 @@ fn create_batch( ColumnType::FixedLen(size) => { for i in 0..num_columns { let array_seed = seed * num_columns + i; - let array = - create_fsb_array_with_seed(num_rows, null_density, size, array_seed as u64); + let array = create_fsb_array_with_seed( + num_rows, + null_density, + size as usize, + array_seed as u64, + ); arrays.push(Arc::new(array)); } }
