Skip to content

Commit eadbed5

Browse files
kumarUjjawalkumarUjjawal
authored
perf: Optimize initcap scalar performance (#19776)
## Which issue does this PR close? <!-- We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123. --> - Part of apache/datafusion-comet#2986. ## Rationale for this change - `initcap` uses `make_scalar_function` which converts scalar inputs to arrays. <!-- Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed. Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes. --> ## What changes are included in this PR? - Add scalar fast path for Utf8/LargeUtf8/Utf8View inputs - Reuse existing `initcap_string` helper for direct scalar processing <!-- There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR. --> ## Are these changes tested? Yes. Unit tests and sqllogictest pass. ## Benchmark Results | Type | Before | After | Speedup | |------|--------|-------|---------| | scalar_utf8 | 698 ns | 250 ns | **2.8x** | | scalar_utf8view | 729 ns | 248 ns | **2.9x** | Measured using: ```bash cargo bench -p datafusion-functions --bench initcap -- "scalar" ``` <!-- We typically require tests for all PRs in order to: 1. Prevent the code from being accidentally broken by subsequent changes 2. Serve as another way to document the expected behavior of the code If tests are not included in your PR, please explain why (for example, are they covered by existing tests)? --> ## Are there any user-facing changes? No <!-- If there are user-facing changes then we may require documentation to be updated before approving the PR. --> <!-- If there are any breaking changes to public APIs, please add the `api change` label. -->
1 parent 0808f3a commit eadbed5

File tree

2 files changed

+111
-49
lines changed

2 files changed

+111
-49
lines changed

‎datafusion/functions/benches/initcap.rs‎

Lines changed: 77 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,14 @@ use arrow::datatypes::{DataType, Field};
2222
use arrow::util::bench_util::{
2323
create_string_array_with_len, create_string_view_array_with_len,
2424
};
25-
use criterion::{Criterion, criterion_group, criterion_main};
25+
use criterion::{Criterion, SamplingMode, criterion_group, criterion_main};
26+
use datafusion_common::ScalarValue;
2627
use datafusion_common::config::ConfigOptions;
2728
use datafusion_expr::{ColumnarValue, ScalarFunctionArgs};
2829
use datafusion_functions::unicode;
2930
use std::hint::black_box;
3031
use std::sync::Arc;
32+
use std::time::Duration;
3133

3234
fn create_args<O: OffsetSizeTrait>(
3335
size: usize,
@@ -49,60 +51,87 @@ fn create_args<O: OffsetSizeTrait>(
4951

5052
fn criterion_benchmark(c: &mut Criterion) {
5153
let initcap = unicode::initcap();
52-
for size in [1024, 4096] {
53-
let args = create_args::<i32>(size, 8, true);
54-
let arg_fields = args
55-
.iter()
56-
.enumerate()
57-
.map(|(idx, arg)| {
58-
Field::new(format!("arg_{idx}"), arg.data_type(), true).into()
59-
})
60-
.collect::<Vec<_>>();
61-
let config_options = Arc::new(ConfigOptions::default());
62-
63-
c.bench_function(
64-
format!("initcap string view shorter than 12 [size={size}]").as_str(),
65-
|b| {
66-
b.iter(|| {
67-
black_box(initcap.invoke_with_args(ScalarFunctionArgs {
68-
args: args.clone(),
69-
arg_fields: arg_fields.clone(),
70-
number_rows: size,
71-
return_field: Field::new("f", DataType::Utf8View, true).into(),
72-
config_options: Arc::clone(&config_options),
73-
}))
74-
})
75-
},
76-
);
77-
78-
let args = create_args::<i32>(size, 16, true);
79-
c.bench_function(
80-
format!("initcap string view longer than 12 [size={size}]").as_str(),
81-
|b| {
82-
b.iter(|| {
83-
black_box(initcap.invoke_with_args(ScalarFunctionArgs {
84-
args: args.clone(),
85-
arg_fields: arg_fields.clone(),
86-
number_rows: size,
87-
return_field: Field::new("f", DataType::Utf8View, true).into(),
88-
config_options: Arc::clone(&config_options),
89-
}))
90-
})
91-
},
92-
);
93-
94-
let args = create_args::<i32>(size, 16, false);
95-
c.bench_function(format!("initcap string [size={size}]").as_str(), |b| {
54+
let config_options = Arc::new(ConfigOptions::default());
55+
56+
// Grouped benchmarks for array sizes - to compare with scalar performance
57+
for size in [1024, 4096, 8192] {
58+
let mut group = c.benchmark_group(format!("initcap size={size}"));
59+
group.sampling_mode(SamplingMode::Flat);
60+
group.sample_size(10);
61+
group.measurement_time(Duration::from_secs(10));
62+
63+
// Array benchmark - Utf8
64+
let array_args = create_args::<i32>(size, 16, false);
65+
let array_arg_fields = vec![Field::new("arg_0", DataType::Utf8, true).into()];
66+
let batch_len = size;
67+
68+
group.bench_function("array_utf8", |b| {
9669
b.iter(|| {
9770
black_box(initcap.invoke_with_args(ScalarFunctionArgs {
98-
args: args.clone(),
99-
arg_fields: arg_fields.clone(),
100-
number_rows: size,
71+
args: array_args.clone(),
72+
arg_fields: array_arg_fields.clone(),
73+
number_rows: batch_len,
10174
return_field: Field::new("f", DataType::Utf8, true).into(),
10275
config_options: Arc::clone(&config_options),
10376
}))
10477
})
10578
});
79+
80+
// Array benchmark - Utf8View
81+
let array_view_args = create_args::<i32>(size, 16, true);
82+
let array_view_arg_fields =
83+
vec![Field::new("arg_0", DataType::Utf8View, true).into()];
84+
85+
group.bench_function("array_utf8view", |b| {
86+
b.iter(|| {
87+
black_box(initcap.invoke_with_args(ScalarFunctionArgs {
88+
args: array_view_args.clone(),
89+
arg_fields: array_view_arg_fields.clone(),
90+
number_rows: batch_len,
91+
return_field: Field::new("f", DataType::Utf8View, true).into(),
92+
config_options: Arc::clone(&config_options),
93+
}))
94+
})
95+
});
96+
97+
// Scalar benchmark - Utf8 (the optimization we added)
98+
let scalar_args = vec![ColumnarValue::Scalar(ScalarValue::Utf8(Some(
99+
"hello world test string".to_string(),
100+
)))];
101+
let scalar_arg_fields = vec![Field::new("arg_0", DataType::Utf8, false).into()];
102+
103+
group.bench_function("scalar_utf8", |b| {
104+
b.iter(|| {
105+
black_box(initcap.invoke_with_args(ScalarFunctionArgs {
106+
args: scalar_args.clone(),
107+
arg_fields: scalar_arg_fields.clone(),
108+
number_rows: 1,
109+
return_field: Field::new("f", DataType::Utf8, false).into(),
110+
config_options: Arc::clone(&config_options),
111+
}))
112+
})
113+
});
114+
115+
// Scalar benchmark - Utf8View
116+
let scalar_view_args = vec![ColumnarValue::Scalar(ScalarValue::Utf8View(Some(
117+
"hello world test string".to_string(),
118+
)))];
119+
let scalar_view_arg_fields =
120+
vec![Field::new("arg_0", DataType::Utf8View, false).into()];
121+
122+
group.bench_function("scalar_utf8view", |b| {
123+
b.iter(|| {
124+
black_box(initcap.invoke_with_args(ScalarFunctionArgs {
125+
args: scalar_view_args.clone(),
126+
arg_fields: scalar_view_arg_fields.clone(),
127+
number_rows: 1,
128+
return_field: Field::new("f", DataType::Utf8View, false).into(),
129+
config_options: Arc::clone(&config_options),
130+
}))
131+
})
132+
});
133+
134+
group.finish();
106135
}
107136
}
108137

‎datafusion/functions/src/unicode/initcap.rs‎

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ use arrow::datatypes::DataType;
2626
use crate::utils::{make_scalar_function, utf8_to_str_type};
2727
use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
2828
use datafusion_common::types::logical_string;
29-
use datafusion_common::{Result, exec_err};
29+
use datafusion_common::{Result, ScalarValue, exec_err};
3030
use datafusion_expr::{
3131
Coercion, ColumnarValue, Documentation, ScalarUDFImpl, Signature, TypeSignatureClass,
3232
Volatility,
@@ -99,6 +99,39 @@ impl ScalarUDFImpl for InitcapFunc {
9999
&self,
100100
args: datafusion_expr::ScalarFunctionArgs,
101101
) -> Result<ColumnarValue> {
102+
let arg = &args.args[0];
103+
104+
// Scalar fast path - handle directly without array conversion
105+
if let ColumnarValue::Scalar(scalar) = arg {
106+
return match scalar {
107+
ScalarValue::Utf8(None)
108+
| ScalarValue::LargeUtf8(None)
109+
| ScalarValue::Utf8View(None) => Ok(arg.clone()),
110+
ScalarValue::Utf8(Some(s)) => {
111+
let mut result = String::new();
112+
initcap_string(s, &mut result);
113+
Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result))))
114+
}
115+
ScalarValue::LargeUtf8(Some(s)) => {
116+
let mut result = String::new();
117+
initcap_string(s, &mut result);
118+
Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(result))))
119+
}
120+
ScalarValue::Utf8View(Some(s)) => {
121+
let mut result = String::new();
122+
initcap_string(s, &mut result);
123+
Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(result))))
124+
}
125+
other => {
126+
exec_err!(
127+
"Unsupported data type {:?} for function `initcap`",
128+
other.data_type()
129+
)
130+
}
131+
};
132+
}
133+
134+
// Array path
102135
let args = &args.args;
103136
match args[0].data_type() {
104137
DataType::Utf8 => make_scalar_function(initcap::<i32>, vec![])(args),

0 commit comments

Comments
 (0)