martin-g commented on code in PR #21119:
URL: https://github.com/apache/datafusion/pull/21119#discussion_r2987078975
##########
datafusion/functions/src/string/split_part.rs:
##########
@@ -123,71 +126,50 @@ impl ScalarUDFImpl for SplitPartFunc {
// Unpack the ArrayRefs from the arguments
let n_array = as_int64_array(&args[2])?;
- let result = match (args[0].data_type(), args[1].data_type()) {
- (DataType::Utf8View, DataType::Utf8View) => {
- split_part_impl::<&StringViewArray, &StringViewArray, i32>(
- &args[0].as_string_view(),
- &args[1].as_string_view(),
- n_array,
- )
- }
- (DataType::Utf8View, DataType::Utf8) => {
- split_part_impl::<&StringViewArray, &GenericStringArray<i32>,
i32>(
- &args[0].as_string_view(),
- &args[1].as_string::<i32>(),
- n_array,
- )
- }
- (DataType::Utf8View, DataType::LargeUtf8) => {
- split_part_impl::<&StringViewArray, &GenericStringArray<i64>,
i32>(
- &args[0].as_string_view(),
- &args[1].as_string::<i64>(),
- n_array,
- )
- }
- (DataType::Utf8, DataType::Utf8View) => {
- split_part_impl::<&GenericStringArray<i32>, &StringViewArray,
i32>(
- &args[0].as_string::<i32>(),
- &args[1].as_string_view(),
- n_array,
- )
- }
- (DataType::LargeUtf8, DataType::Utf8View) => {
- split_part_impl::<&GenericStringArray<i64>, &StringViewArray,
i64>(
- &args[0].as_string::<i64>(),
- &args[1].as_string_view(),
- n_array,
- )
- }
- (DataType::Utf8, DataType::Utf8) => {
- split_part_impl::<&GenericStringArray<i32>,
&GenericStringArray<i32>, i32>(
- &args[0].as_string::<i32>(),
- &args[1].as_string::<i32>(),
- n_array,
- )
- }
- (DataType::LargeUtf8, DataType::LargeUtf8) => {
- split_part_impl::<&GenericStringArray<i64>,
&GenericStringArray<i64>, i64>(
- &args[0].as_string::<i64>(),
- &args[1].as_string::<i64>(),
- n_array,
- )
- }
- (DataType::Utf8, DataType::LargeUtf8) => {
- split_part_impl::<&GenericStringArray<i32>,
&GenericStringArray<i64>, i32>(
- &args[0].as_string::<i32>(),
- &args[1].as_string::<i64>(),
- n_array,
- )
- }
- (DataType::LargeUtf8, DataType::Utf8) => {
- split_part_impl::<&GenericStringArray<i64>,
&GenericStringArray<i32>, i64>(
- &args[0].as_string::<i64>(),
- &args[1].as_string::<i32>(),
- n_array,
- )
- }
- _ => exec_err!("Unsupported combination of argument types for
split_part"),
+
+ // Dispatch on delimiter type for a given string array and builder.
+ macro_rules! split_part_for_delimiter_type {
+ ($str_arr:expr, $builder:expr) => {
+ match args[1].data_type() {
+ DataType::Utf8View => split_part_impl(
+ $str_arr,
+ &args[1].as_string_view(),
+ n_array,
+ $builder,
+ ),
+ DataType::Utf8 => split_part_impl(
+ $str_arr,
+ &args[1].as_string::<i32>(),
+ n_array,
+ $builder,
+ ),
+ DataType::LargeUtf8 => split_part_impl(
+ $str_arr,
+ &args[1].as_string::<i64>(),
+ n_array,
+ $builder,
+ ),
+ other => {
+ exec_err!("Unsupported delimiter type {other:?} for
split_part")
+ }
+ }
+ };
+ }
+
+ let result = match args[0].data_type() {
+ DataType::Utf8View => split_part_for_delimiter_type!(
+ &args[0].as_string_view(),
+ StringViewBuilder::with_capacity(inferred_length)
+ ),
+ DataType::Utf8 => split_part_for_delimiter_type!(
+ &args[0].as_string::<i32>(),
+ GenericStringBuilder::<i32>::new()
+ ),
Review Comment:
The memory allocation could be optimized a bit:
```suggestion
DataType::Utf8 => {
let str_arr = &args[0].as_string::<i32>();
split_part_for_delimiter_type!(
str_arr,
GenericStringBuilder::<i32>::with_capacity(
inferred_length,
str_arr.value_data().len(),
)
)
},
```
Same for LargeUtf8
##########
datafusion/functions/benches/split_part.rs:
##########
@@ -73,303 +73,139 @@ fn gen_split_part_data(
}
}
-fn gen_positions(n_rows: usize, position: i64) -> ColumnarValue {
- let positions: Vec<i64> = vec![position; n_rows];
- ColumnarValue::Array(Arc::new(Int64Array::from(positions)) as ArrayRef)
+#[expect(clippy::too_many_arguments)]
+fn bench_split_part(
+ group: &mut criterion::BenchmarkGroup<'_,
criterion::measurement::WallTime>,
+ func: &ScalarUDF,
+ config_options: &Arc<ConfigOptions>,
+ name: &str,
+ tag: &str,
+ strings: ColumnarValue,
+ delimiters: ColumnarValue,
+ position: i64,
+) {
+ let positions: ColumnarValue =
+ ColumnarValue::Array(Arc::new(Int64Array::from(vec![position;
N_ROWS])));
+ let args = vec![strings, delimiters, positions];
+ let arg_fields: Vec<_> = args
+ .iter()
+ .enumerate()
+ .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(),
true).into())
+ .collect();
+ let return_type = match args[0].data_type() {
+ DataType::Utf8View => DataType::Utf8View,
+ _ => DataType::Utf8,
+ };
+ let return_field = Field::new("f", return_type, true).into();
+
+ group.bench_function(BenchmarkId::new(name, tag), |b| {
+ b.iter(|| {
+ black_box(
+ func.invoke_with_args(ScalarFunctionArgs {
+ args: args.clone(),
+ arg_fields: arg_fields.clone(),
+ number_rows: N_ROWS,
+ return_field: Arc::clone(&return_field),
+ config_options: Arc::clone(config_options),
+ })
+ .expect("split_part should work"),
+ )
+ })
+ });
}
fn criterion_benchmark(c: &mut Criterion) {
let split_part_func = split_part();
let config_options = Arc::new(ConfigOptions::default());
-
let mut group = c.benchmark_group("split_part");
- // Test different scenarios
- // Scenario 1: Single-char delimiter, first position (should be fastest
with optimization)
- {
- let (strings, delimiters) = gen_split_part_data(N_ROWS, 10, 8, ".",
false);
- let positions = gen_positions(N_ROWS, 1);
- let args = vec![strings, delimiters, positions];
- let arg_fields: Vec<_> = args
- .iter()
- .enumerate()
- .map(|(idx, arg)| {
- Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
- })
- .collect();
- let return_field = Field::new("f", DataType::Utf8, true).into();
-
- group.bench_function(BenchmarkId::new("single_char_delim",
"pos_first"), |b| {
- b.iter(|| {
- black_box(
- split_part_func
- .invoke_with_args(ScalarFunctionArgs {
- args: args.clone(),
- arg_fields: arg_fields.clone(),
- number_rows: N_ROWS,
- return_field: Arc::clone(&return_field),
- config_options: Arc::clone(&config_options),
- })
- .expect("split_part should work"),
- )
- })
- });
- }
-
- // Scenario 2: Single-char delimiter, middle position
+ // Utf8, single-char delimiter, first position
{
let (strings, delimiters) = gen_split_part_data(N_ROWS, 10, 8, ".",
false);
- let positions = gen_positions(N_ROWS, 5);
- let args = vec![strings, delimiters, positions];
- let arg_fields: Vec<_> = args
- .iter()
- .enumerate()
- .map(|(idx, arg)| {
- Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
- })
- .collect();
- let return_field = Field::new("f", DataType::Utf8, true).into();
-
- group.bench_function(BenchmarkId::new("single_char_delim",
"pos_middle"), |b| {
- b.iter(|| {
- black_box(
- split_part_func
- .invoke_with_args(ScalarFunctionArgs {
- args: args.clone(),
- arg_fields: arg_fields.clone(),
- number_rows: N_ROWS,
- return_field: Arc::clone(&return_field),
- config_options: Arc::clone(&config_options),
- })
- .expect("split_part should work"),
- )
- })
- });
+ bench_split_part(
+ &mut group,
+ &split_part_func,
+ &config_options,
+ "utf8_single_char",
+ "pos_first",
+ strings,
+ delimiters,
+ 1,
+ );
}
- // Scenario 3: Single-char delimiter, last position
+ // Utf8, single-char delimiter, middle position
{
let (strings, delimiters) = gen_split_part_data(N_ROWS, 10, 8, ".",
false);
- let positions = gen_positions(N_ROWS, 10);
- let args = vec![strings, delimiters, positions];
- let arg_fields: Vec<_> = args
- .iter()
- .enumerate()
- .map(|(idx, arg)| {
- Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
- })
- .collect();
- let return_field = Field::new("f", DataType::Utf8, true).into();
-
- group.bench_function(BenchmarkId::new("single_char_delim",
"pos_last"), |b| {
- b.iter(|| {
- black_box(
- split_part_func
- .invoke_with_args(ScalarFunctionArgs {
- args: args.clone(),
- arg_fields: arg_fields.clone(),
- number_rows: N_ROWS,
- return_field: Arc::clone(&return_field),
- config_options: Arc::clone(&config_options),
- })
- .expect("split_part should work"),
- )
- })
- });
+ bench_split_part(
+ &mut group,
+ &split_part_func,
+ &config_options,
+ "utf8_single_char",
+ "pos_middle",
+ strings,
+ delimiters,
+ 5,
+ );
}
- // Scenario 4: Single-char delimiter, negative position (last element)
+ // Utf8, single-char delimiter, negative position
{
let (strings, delimiters) = gen_split_part_data(N_ROWS, 10, 8, ".",
false);
- let positions = gen_positions(N_ROWS, -1);
- let args = vec![strings, delimiters, positions];
- let arg_fields: Vec<_> = args
- .iter()
- .enumerate()
- .map(|(idx, arg)| {
- Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
- })
- .collect();
- let return_field = Field::new("f", DataType::Utf8, true).into();
-
- group.bench_function(
- BenchmarkId::new("single_char_delim", "pos_negative"),
- |b| {
- b.iter(|| {
- black_box(
- split_part_func
- .invoke_with_args(ScalarFunctionArgs {
- args: args.clone(),
- arg_fields: arg_fields.clone(),
- number_rows: N_ROWS,
- return_field: Arc::clone(&return_field),
- config_options: Arc::clone(&config_options),
- })
- .expect("split_part should work"),
- )
- })
- },
+ bench_split_part(
+ &mut group,
+ &split_part_func,
+ &config_options,
+ "utf8_single_char",
+ "pos_negative",
+ strings,
+ delimiters,
+ -1,
);
}
- // Scenario 5: Multi-char delimiter, first position
- {
- let (strings, delimiters) = gen_split_part_data(N_ROWS, 10, 8, "~@~",
false);
- let positions = gen_positions(N_ROWS, 1);
- let args = vec![strings, delimiters, positions];
- let arg_fields: Vec<_> = args
- .iter()
- .enumerate()
- .map(|(idx, arg)| {
- Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
- })
- .collect();
- let return_field = Field::new("f", DataType::Utf8, true).into();
-
- group.bench_function(BenchmarkId::new("multi_char_delim",
"pos_first"), |b| {
- b.iter(|| {
- black_box(
- split_part_func
- .invoke_with_args(ScalarFunctionArgs {
- args: args.clone(),
- arg_fields: arg_fields.clone(),
- number_rows: N_ROWS,
- return_field: Arc::clone(&return_field),
- config_options: Arc::clone(&config_options),
- })
- .expect("split_part should work"),
- )
- })
- });
- }
-
- // Scenario 6: Multi-char delimiter, middle position
+ // Utf8, multi-char delimiter, middle position
{
let (strings, delimiters) = gen_split_part_data(N_ROWS, 10, 8, "~@~",
false);
- let positions = gen_positions(N_ROWS, 5);
- let args = vec![strings, delimiters, positions];
- let arg_fields: Vec<_> = args
- .iter()
- .enumerate()
- .map(|(idx, arg)| {
- Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
- })
- .collect();
- let return_field = Field::new("f", DataType::Utf8, true).into();
-
- group.bench_function(BenchmarkId::new("multi_char_delim",
"pos_middle"), |b| {
- b.iter(|| {
- black_box(
- split_part_func
- .invoke_with_args(ScalarFunctionArgs {
- args: args.clone(),
- arg_fields: arg_fields.clone(),
- number_rows: N_ROWS,
- return_field: Arc::clone(&return_field),
- config_options: Arc::clone(&config_options),
- })
- .expect("split_part should work"),
- )
- })
- });
+ bench_split_part(
+ &mut group,
+ &split_part_func,
+ &config_options,
+ "utf8_multi_char",
+ "pos_middle",
+ strings,
+ delimiters,
+ 5,
+ );
}
- // Scenario 7: StringViewArray, single-char delimiter, first position
+ // Utf8View, single-char delimiter, first position
{
let (strings, delimiters) = gen_split_part_data(N_ROWS, 10, 8, ".",
true);
- let positions = gen_positions(N_ROWS, 1);
- let args = vec![strings, delimiters, positions];
- let arg_fields: Vec<_> = args
- .iter()
- .enumerate()
- .map(|(idx, arg)| {
- Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
- })
- .collect();
- let return_field = Field::new("f", DataType::Utf8, true).into();
-
- group.bench_function(
- BenchmarkId::new("string_view_single_char", "pos_first"),
- |b| {
- b.iter(|| {
- black_box(
- split_part_func
- .invoke_with_args(ScalarFunctionArgs {
- args: args.clone(),
- arg_fields: arg_fields.clone(),
- number_rows: N_ROWS,
- return_field: Arc::clone(&return_field),
- config_options: Arc::clone(&config_options),
- })
- .expect("split_part should work"),
- )
- })
- },
+ bench_split_part(
+ &mut group,
+ &split_part_func,
+ &config_options,
+ "utf8view_single_char",
+ "pos_first",
+ strings,
+ delimiters,
+ 1,
);
}
- // Scenario 8: Many parts (20), position near end - shows benefit of early
termination
+ // Utf8View, single-char delimiter, middle position, long parts
{
- let (strings, delimiters) = gen_split_part_data(N_ROWS, 20, 8, ".",
false);
- let positions = gen_positions(N_ROWS, 2);
- let args = vec![strings, delimiters, positions];
- let arg_fields: Vec<_> = args
- .iter()
- .enumerate()
- .map(|(idx, arg)| {
- Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
- })
- .collect();
- let return_field = Field::new("f", DataType::Utf8, true).into();
-
- group.bench_function(BenchmarkId::new("many_parts_20", "pos_second"),
|b| {
- b.iter(|| {
- black_box(
- split_part_func
- .invoke_with_args(ScalarFunctionArgs {
- args: args.clone(),
- arg_fields: arg_fields.clone(),
- number_rows: N_ROWS,
- return_field: Arc::clone(&return_field),
- config_options: Arc::clone(&config_options),
- })
- .expect("split_part should work"),
- )
- })
- });
- }
-
- // Scenario 9: Long strings with many parts - worst case for old
implementation
Review Comment:
Is this scenario still tested ?
It seems to be lost.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
