diff --git a/datafusion/functions/benches/pad.rs b/datafusion/functions/benches/pad.rs index c9220d8d3fce1..f6b2ed7636bf8 100644 --- a/datafusion/functions/benches/pad.rs +++ b/datafusion/functions/benches/pad.rs @@ -15,20 +15,22 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, ArrowPrimitiveType, OffsetSizeTrait, PrimitiveArray}; +extern crate criterion; + +use arrow::array::{ArrowPrimitiveType, OffsetSizeTrait, PrimitiveArray}; use arrow::datatypes::{DataType, Field, Int64Type}; use arrow::util::bench_util::{ create_string_array_with_len, create_string_view_array_with_len, }; -use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; -use datafusion_common::DataFusionError; +use criterion::{Criterion, SamplingMode, criterion_group, criterion_main}; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; -use datafusion_functions::unicode::{lpad, rpad}; +use datafusion_functions::unicode; use rand::Rng; use rand::distr::{Distribution, Uniform}; use std::hint::black_box; use std::sync::Arc; +use std::time::Duration; struct Filter { dist: Dist, @@ -67,104 +69,260 @@ where .collect() } -fn create_args( +/// Create args for pad benchmark +fn create_pad_args( size: usize, str_len: usize, - force_view_types: bool, + target_len: usize, + use_string_view: bool, ) -> Vec { - let length_array = Arc::new(create_primitive_array::(size, 0.0, str_len)); - - if !force_view_types { - let string_array = - Arc::new(create_string_array_with_len::(size, 0.1, str_len)); - let fill_array = Arc::new(create_string_array_with_len::(size, 0.1, str_len)); + let length_array = + Arc::new(create_primitive_array::(size, 0.0, target_len)); + if use_string_view { + let string_array = create_string_view_array_with_len(size, 0.1, str_len, false); + let fill_array = create_string_view_array_with_len(size, 0.1, str_len, false); vec![ - ColumnarValue::Array(string_array), - ColumnarValue::Array(Arc::clone(&length_array) as ArrayRef), - ColumnarValue::Array(fill_array), + ColumnarValue::Array(Arc::new(string_array)), + ColumnarValue::Array(length_array), + ColumnarValue::Array(Arc::new(fill_array)), ] } else { - let string_array = - Arc::new(create_string_view_array_with_len(size, 0.1, str_len, false)); - let fill_array = - Arc::new(create_string_view_array_with_len(size, 0.1, str_len, false)); - + let string_array = create_string_array_with_len::(size, 0.1, str_len); + let fill_array = create_string_array_with_len::(size, 0.1, str_len); vec![ - ColumnarValue::Array(string_array), - ColumnarValue::Array(Arc::clone(&length_array) as ArrayRef), - ColumnarValue::Array(fill_array), + ColumnarValue::Array(Arc::new(string_array)), + ColumnarValue::Array(length_array), + ColumnarValue::Array(Arc::new(fill_array)), ] } } -#[expect(clippy::needless_pass_by_value)] -fn invoke_pad_with_args( - args: Vec, - number_rows: usize, - left_pad: bool, -) -> Result { - let arg_fields = args - .iter() - .enumerate() - .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into()) - .collect::>(); - let config_options = Arc::new(ConfigOptions::default()); - - let scalar_args = ScalarFunctionArgs { - args: args.clone(), - arg_fields, - number_rows, - return_field: Field::new("f", DataType::Utf8, true).into(), - config_options: Arc::clone(&config_options), - }; +fn criterion_benchmark(c: &mut Criterion) { + for size in [1024, 4096] { + let mut group = c.benchmark_group(format!("lpad size={size}")); + group.sampling_mode(SamplingMode::Flat); + group.sample_size(10); + group.measurement_time(Duration::from_secs(10)); - if left_pad { - lpad().invoke_with_args(scalar_args) - } else { - rpad().invoke_with_args(scalar_args) - } -} + // Utf8 type + let args = create_pad_args::(size, 5, 20, false); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect::>(); + let config_options = Arc::new(ConfigOptions::default()); -fn criterion_benchmark(c: &mut Criterion) { - for size in [1024, 2048] { - let mut group = c.benchmark_group("lpad function"); + group.bench_function( + format!("lpad utf8 [size={size}, str_len=5, target=20]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(unicode::lpad().invoke_with_args(ScalarFunctionArgs { + args: args_cloned, + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8, true).into(), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); - let args = create_args::(size, 32, false); + // StringView type + let args = create_pad_args::(size, 5, 20, true); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect::>(); - group.bench_function(BenchmarkId::new("utf8 type", size), |b| { - b.iter(|| black_box(invoke_pad_with_args(args.clone(), size, true).unwrap())) - }); + group.bench_function( + format!("lpad stringview [size={size}, str_len=5, target=20]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(unicode::lpad().invoke_with_args(ScalarFunctionArgs { + args: args_cloned, + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8View, true).into(), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); - let args = create_args::(size, 32, false); - group.bench_function(BenchmarkId::new("largeutf8 type", size), |b| { - b.iter(|| black_box(invoke_pad_with_args(args.clone(), size, true).unwrap())) - }); + // Utf8 type with longer strings + let args = create_pad_args::(size, 20, 50, false); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect::>(); - let args = create_args::(size, 32, true); - group.bench_function(BenchmarkId::new("stringview type", size), |b| { - b.iter(|| black_box(invoke_pad_with_args(args.clone(), size, true).unwrap())) - }); + group.bench_function( + format!("lpad utf8 [size={size}, str_len=20, target=50]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(unicode::lpad().invoke_with_args(ScalarFunctionArgs { + args: args_cloned, + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8, true).into(), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); + + // StringView type with longer strings + let args = create_pad_args::(size, 20, 50, true); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect::>(); + + group.bench_function( + format!("lpad stringview [size={size}, str_len=20, target=50]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(unicode::lpad().invoke_with_args(ScalarFunctionArgs { + args: args_cloned, + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8View, true).into(), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); group.finish(); + } + + for size in [1024, 4096] { + let mut group = c.benchmark_group(format!("rpad size={size}")); + group.sampling_mode(SamplingMode::Flat); + group.sample_size(10); + group.measurement_time(Duration::from_secs(10)); + + // Utf8 type + let args = create_pad_args::(size, 5, 20, false); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect::>(); + let config_options = Arc::new(ConfigOptions::default()); + + group.bench_function( + format!("rpad utf8 [size={size}, str_len=5, target=20]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(unicode::rpad().invoke_with_args(ScalarFunctionArgs { + args: args_cloned, + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8, true).into(), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); + + // StringView type + let args = create_pad_args::(size, 5, 20, true); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect::>(); + + group.bench_function( + format!("rpad stringview [size={size}, str_len=5, target=20]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(unicode::rpad().invoke_with_args(ScalarFunctionArgs { + args: args_cloned, + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8View, true).into(), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); - let mut group = c.benchmark_group("rpad function"); + // Utf8 type with longer strings + let args = create_pad_args::(size, 20, 50, false); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect::>(); - let args = create_args::(size, 32, false); - group.bench_function(BenchmarkId::new("utf8 type", size), |b| { - b.iter(|| black_box(invoke_pad_with_args(args.clone(), size, false).unwrap())) - }); + group.bench_function( + format!("rpad utf8 [size={size}, str_len=20, target=50]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(unicode::rpad().invoke_with_args(ScalarFunctionArgs { + args: args_cloned, + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8, true).into(), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); - let args = create_args::(size, 32, false); - group.bench_function(BenchmarkId::new("largeutf8 type", size), |b| { - b.iter(|| black_box(invoke_pad_with_args(args.clone(), size, false).unwrap())) - }); + // StringView type with longer strings + let args = create_pad_args::(size, 20, 50, true); + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect::>(); - // rpad for stringview type - let args = create_args::(size, 32, true); - group.bench_function(BenchmarkId::new("stringview type", size), |b| { - b.iter(|| black_box(invoke_pad_with_args(args.clone(), size, false).unwrap())) - }); + group.bench_function( + format!("rpad stringview [size={size}, str_len=20, target=50]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(unicode::rpad().invoke_with_args(ScalarFunctionArgs { + args: args_cloned, + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8View, true).into(), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); group.finish(); } diff --git a/datafusion/functions/src/unicode/lpad.rs b/datafusion/functions/src/unicode/lpad.rs index 527aaf1389aa6..a892c0adf58de 100644 --- a/datafusion/functions/src/unicode/lpad.rs +++ b/datafusion/functions/src/unicode/lpad.rs @@ -206,6 +206,8 @@ where { let array = if let Some(fill_array) = fill_array { let mut builder: GenericStringBuilder = GenericStringBuilder::new(); + let mut graphemes_buf = Vec::new(); + let mut fill_chars_buf = Vec::new(); for ((string, length), fill) in string_array .iter() @@ -223,16 +225,20 @@ where continue; } - let graphemes = string.graphemes(true).collect::>(); - let fill_chars = fill.chars().collect::>(); + // Reuse buffers by clearing and refilling + graphemes_buf.clear(); + graphemes_buf.extend(string.graphemes(true)); - if length < graphemes.len() { - builder.append_value(graphemes[..length].concat()); - } else if fill_chars.is_empty() { + fill_chars_buf.clear(); + fill_chars_buf.extend(fill.chars()); + + if length < graphemes_buf.len() { + builder.append_value(graphemes_buf[..length].concat()); + } else if fill_chars_buf.is_empty() { builder.append_value(string); } else { - for l in 0..length - graphemes.len() { - let c = *fill_chars.get(l % fill_chars.len()).unwrap(); + for l in 0..length - graphemes_buf.len() { + let c = *fill_chars_buf.get(l % fill_chars_buf.len()).unwrap(); builder.write_char(c)?; } builder.write_str(string)?; @@ -246,6 +252,7 @@ where builder.finish() } else { let mut builder: GenericStringBuilder = GenericStringBuilder::new(); + let mut graphemes_buf = Vec::new(); for (string, length) in string_array.iter().zip(length_array.iter()) { if let (Some(string), Some(length)) = (string, length) { @@ -259,11 +266,15 @@ where continue; } - let graphemes = string.graphemes(true).collect::>(); - if length < graphemes.len() { - builder.append_value(graphemes[..length].concat()); + // Reuse buffer by clearing and refilling + graphemes_buf.clear(); + graphemes_buf.extend(string.graphemes(true)); + + if length < graphemes_buf.len() { + builder.append_value(graphemes_buf[..length].concat()); } else { - builder.write_str(" ".repeat(length - graphemes.len()).as_str())?; + builder + .write_str(" ".repeat(length - graphemes_buf.len()).as_str())?; builder.write_str(string)?; builder.append_value(""); } diff --git a/datafusion/functions/src/unicode/rpad.rs b/datafusion/functions/src/unicode/rpad.rs index a3c2d501c9127..14f517faf8cf1 100644 --- a/datafusion/functions/src/unicode/rpad.rs +++ b/datafusion/functions/src/unicode/rpad.rs @@ -216,6 +216,8 @@ where StringArrayLen: OffsetSizeTrait, { let mut builder: GenericStringBuilder = GenericStringBuilder::new(); + let mut graphemes_buf = Vec::new(); + let mut fill_chars_buf = Vec::new(); match fill_array { None => { @@ -233,14 +235,17 @@ where if length == 0 { builder.append_value(""); } else { - let graphemes = - string.graphemes(true).collect::>(); - if length < graphemes.len() { - builder.append_value(graphemes[..length].concat()); + // Reuse buffer by clearing and refilling + graphemes_buf.clear(); + graphemes_buf.extend(string.graphemes(true)); + + if length < graphemes_buf.len() { + builder + .append_value(graphemes_buf[..length].concat()); } else { builder.write_str(string)?; builder.write_str( - &" ".repeat(length - graphemes.len()), + &" ".repeat(length - graphemes_buf.len()), )?; builder.append_value(""); } @@ -268,19 +273,26 @@ where ); } let length = if length < 0 { 0 } else { length as usize }; - let graphemes = - string.graphemes(true).collect::>(); + // Reuse buffer by clearing and refilling + graphemes_buf.clear(); + graphemes_buf.extend(string.graphemes(true)); - if length < graphemes.len() { - builder.append_value(graphemes[..length].concat()); + if length < graphemes_buf.len() { + builder + .append_value(graphemes_buf[..length].concat()); } else if fill.is_empty() { builder.append_value(string); } else { builder.write_str(string)?; - fill.chars() - .cycle() - .take(length - graphemes.len()) - .for_each(|ch| builder.write_char(ch).unwrap()); + // Reuse fill_chars_buf by clearing and refilling + fill_chars_buf.clear(); + fill_chars_buf.extend(fill.chars()); + for l in 0..length - graphemes_buf.len() { + let c = *fill_chars_buf + .get(l % fill_chars_buf.len()) + .unwrap(); + builder.write_char(c)?; + } builder.append_value(""); } }