diff --git a/arrow/benches/coalesce_kernels.rs b/arrow/benches/coalesce_kernels.rs index da62fa390c34..d43001a51449 100644 --- a/arrow/benches/coalesce_kernels.rs +++ b/arrow/benches/coalesce_kernels.rs @@ -18,6 +18,7 @@ //! Benchmarks for the `coalesce` kernels in Arrow. use arrow::util::bench_util::*; +use arrow::util::data_gen::create_random_array; use std::sync::Arc; use arrow::array::*; @@ -105,6 +106,29 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { // TODO model other dictionary types here (FixedSizeBinary for example) ])); + let single_dictionary_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + )])); + + let single_run_end_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::RunEndEncoded( + Arc::new(Field::new( + Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME, + DataType::Int32, + false, + )), + Arc::new(Field::new( + Field::REE_VALUES_FIELD_DEFAULT_NAME, + DataType::Int32, + true, + )), + ), + true, + )])); + let single_fsb16_schema = SchemaRef::new(Schema::new(vec![Field::new( "value", DataType::FixedSizeBinary(16), @@ -320,6 +344,18 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { } .build(); + FilterBenchmarkBuilder { + c, + name: "single_dict", + batch_size, + num_output_batches: 25, + null_density, + selectivity, + max_string_len: 30, + schema: &single_dictionary_schema, + } + .build(); + FilterBenchmarkBuilder { c, name: "single_fsb16", @@ -356,6 +392,18 @@ fn add_all_filter_benchmarks(c: &mut Criterion) { } .build(); + FilterBenchmarkBuilder { + c, + name: "single_ree", + batch_size, + num_output_batches: 25, + null_density, + selectivity, + max_string_len: 0, + schema: &single_run_end_schema, + } + .build(); + FilterBenchmarkBuilder { c, name: "single_fsl4", @@ -470,6 +518,29 @@ fn add_all_gather_benchmarks(c: &mut Criterion, operation: GatherOperation) { Field::new("float_val2", DataType::Float64, true), ])); + let single_dictionary_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + true, + )])); + + let single_run_end_schema = SchemaRef::new(Schema::new(vec![Field::new( + "value", + DataType::RunEndEncoded( + Arc::new(Field::new( + Field::REE_RUN_ENDS_FIELD_DEFAULT_NAME, + DataType::Int32, + false, + )), + Arc::new(Field::new( + Field::REE_VALUES_FIELD_DEFAULT_NAME, + DataType::Int32, + true, + )), + ), + true, + )])); + let single_fsb16_schema = SchemaRef::new(Schema::new(vec![Field::new( "value", DataType::FixedSizeBinary(16), @@ -568,6 +639,12 @@ fn add_all_gather_benchmarks(c: &mut Criterion, operation: GatherOperation) { max_string_len: 30, schema: &mixed_dict_schema, }, + GatherBenchmarkScenario { + name: "single_dict", + num_output_batches: 25, + max_string_len: 30, + schema: &single_dictionary_schema, + }, GatherBenchmarkScenario { name: "single_fsb16", num_output_batches: 50, @@ -586,6 +663,12 @@ fn add_all_gather_benchmarks(c: &mut Criterion, operation: GatherOperation) { max_string_len: 0, schema: &single_boolean_schema, }, + GatherBenchmarkScenario { + name: "single_ree", + num_output_batches: 25, + max_string_len: 0, + schema: &single_run_end_schema, + }, GatherBenchmarkScenario { name: "single_fsl4", num_output_batches: 50, @@ -1385,6 +1468,9 @@ impl DataStreamBuilder { 10, seed, )), + DataType::RunEndEncoded(_, _) => { + create_random_array(field, self.batch_size, self.null_density, 0.5).unwrap() + } _ => panic!("Unsupported data type: {field:?}"), } }