alamb commented on code in PR #14575:
URL: https://github.com/apache/datafusion/pull/14575#discussion_r1950877830
##########
datafusion/functions/src/string/repeat.rs:
##########
@@ -130,22 +142,36 @@ fn repeat(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
-fn repeat_impl<'a, T, S>(string_array: S, number_array: &Int64Array) ->
Result<ArrayRef>
+fn repeat_impl<'a, T, S>(
+ string_array: S,
+ number_array: &Int64Array,
+ max_str_len: usize,
+) -> Result<ArrayRef>
where
T: OffsetSizeTrait,
S: StringArrayType<'a>,
{
let mut builder: GenericStringBuilder<T> = GenericStringBuilder::new();
- string_array
- .iter()
- .zip(number_array.iter())
- .for_each(|(string, number)| match (string, number) {
- (Some(string), Some(number)) if number >= 0 => {
- builder.append_value(string.repeat(number as usize))
+ string_array.iter().zip(number_array.iter()).try_for_each(
+ |(string, number)| -> Result<(), DataFusionError> {
+ match (string, number) {
+ (Some(string), Some(number)) if number >= 0 => {
+ if number as usize * string.len() > max_str_len {
+ return exec_err!(
+ "string size overflow on repeat, max size is {},
but got {}",
+ max_str_len,
+ number as usize * string.len()
+ );
+ } else {
+ builder.append_value(string.repeat(number as usize))
Review Comment:
this code checks for overflow on each individual element which is an
improvement.
We could also potentially ensure that the total length of the builder's data
after pushing the string would be iess than max length as well
Given it calls
[String::repeat](https://doc.rust-lang.org/std/string/struct.String.html#method.repeat)
which allocates and copies the string (before `append_value` copies it again)
I doubt checking for max sizes is going to make a difference
I am running some benchmarks to be sure, however.
##########
datafusion/functions/src/string/repeat.rs:
##########
@@ -130,22 +142,36 @@ fn repeat(args: &[ArrayRef]) -> Result<ArrayRef> {
}
}
-fn repeat_impl<'a, T, S>(string_array: S, number_array: &Int64Array) ->
Result<ArrayRef>
+fn repeat_impl<'a, T, S>(
+ string_array: S,
+ number_array: &Int64Array,
+ max_str_len: usize,
+) -> Result<ArrayRef>
where
T: OffsetSizeTrait,
S: StringArrayType<'a>,
{
let mut builder: GenericStringBuilder<T> = GenericStringBuilder::new();
- string_array
- .iter()
- .zip(number_array.iter())
- .for_each(|(string, number)| match (string, number) {
- (Some(string), Some(number)) if number >= 0 => {
- builder.append_value(string.repeat(number as usize))
+ string_array.iter().zip(number_array.iter()).try_for_each(
+ |(string, number)| -> Result<(), DataFusionError> {
+ match (string, number) {
+ (Some(string), Some(number)) if number >= 0 => {
+ if number as usize * string.len() > max_str_len {
+ return exec_err!(
+ "string size overflow on repeat, max size is {},
but got {}",
+ max_str_len,
+ number as usize * string.len()
+ );
+ } else {
+ builder.append_value(string.repeat(number as usize))
Review Comment:
I also filed a ticket to optimize this path some more
- https://github.com/apache/datafusion/issues/14610
--
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: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
