From 668468808f6dbb2d206e854251956912fc7e5302 Mon Sep 17 00:00:00 2001 From: gonappuccino Date: Fri, 30 Oct 2026 18:30:20 -0400 Subject: [PATCH 1/2] fix: avoid lpad/rpad panic on out-of-range length arrays --- datafusion/functions/src/unicode/common.rs | 5 +- datafusion/functions/src/unicode/lpad.rs | 51 ++++++++++++++++++- datafusion/functions/src/unicode/rpad.rs | 39 +++++++++++++- .../test_files/string/string_literal.slt | 22 ++++++++ 4 files changed, 112 insertions(+), 5 deletions(-) diff --git a/datafusion/functions/src/unicode/common.rs b/datafusion/functions/src/unicode/common.rs index f42b871a45e81..1d63eab1aa637 100644 --- a/datafusion/functions/src/unicode/common.rs +++ b/datafusion/functions/src/unicode/common.rs @@ -44,7 +44,10 @@ pub(crate) fn pad_data_capacity(length_array: &Int64Array) -> usize { length_array .iter() .flatten() - .fold(0, |acc, len| acc.saturating_add(len as usize)) + // Out-of-range lengths produce "", an error or null, so add no bytes + .filter(|len| (0..=i32::MAX as i64).contains(len)) + .fold(0usize, |acc, len| acc.saturating_add(len as usize)) + .min(i32::MAX as usize) } /// A trait for `left` and `right` byte slicing operations diff --git a/datafusion/functions/src/unicode/lpad.rs b/datafusion/functions/src/unicode/lpad.rs index 75a645329aac9..92e0e15587281 100644 --- a/datafusion/functions/src/unicode/lpad.rs +++ b/datafusion/functions/src/unicode/lpad.rs @@ -81,10 +81,12 @@ mod tests { use crate::unicode::lpad::LPadFunc; use crate::utils::test::test_function; - use arrow::array::{Array, LargeStringArray, StringArray}; + use std::sync::Arc; + + use arrow::array::{Array, Int64Array, LargeStringArray, StringArray}; use arrow::datatypes::DataType::{LargeUtf8, Utf8}; - use datafusion_common::{Result, ScalarValue}; + use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; macro_rules! test_lpad { @@ -345,4 +347,49 @@ mod tests { Ok(()) } + + fn array_args(lengths: Vec) -> Vec { + vec![ + ColumnarValue::Array(Arc::new(StringArray::from(vec!["hi"; lengths.len()]))), + ColumnarValue::Array(Arc::new(Int64Array::from(lengths))), + ] + } + + #[test] + fn test_array_length_out_of_range() -> Result<()> { + test_function!( + LPadFunc::new(), + array_args(vec![-1]), + Ok(Some("")), + &str, + Utf8, + StringArray + ); + test_function!( + LPadFunc::new(), + array_args(vec![1 << 62]), + exec_err!( + "lpad requested length {} too large, maximum allowed length is {}", + 1i64 << 62, + i32::MAX + ), + &str, + Utf8, + StringArray + ); + test_function!( + LPadFunc::new(), + array_args(vec![i64::MAX, i64::MAX]), + exec_err!( + "lpad requested length {} too large, maximum allowed length is {}", + i64::MAX, + i32::MAX + ), + &str, + Utf8, + StringArray + ); + + Ok(()) + } } diff --git a/datafusion/functions/src/unicode/rpad.rs b/datafusion/functions/src/unicode/rpad.rs index 0717dd1392ed9..66ba77814b612 100644 --- a/datafusion/functions/src/unicode/rpad.rs +++ b/datafusion/functions/src/unicode/rpad.rs @@ -78,10 +78,12 @@ impl ScalarUDFImpl for RPadFunc { #[cfg(test)] mod tests { - use arrow::array::{Array, StringArray}; + use std::sync::Arc; + + use arrow::array::{Array, Int64Array, StringArray}; use arrow::datatypes::DataType::Utf8; - use datafusion_common::{Result, ScalarValue}; + use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; use crate::unicode::rpad::RPadFunc; @@ -292,4 +294,37 @@ mod tests { Ok(()) } + + fn array_args(lengths: Vec) -> Vec { + vec![ + ColumnarValue::Array(Arc::new(StringArray::from(vec!["hi"; lengths.len()]))), + ColumnarValue::Array(Arc::new(Int64Array::from(lengths))), + ] + } + + #[test] + fn test_array_length_out_of_range() -> Result<()> { + test_function!( + RPadFunc::new(), + array_args(vec![-1]), + Ok(Some("")), + &str, + Utf8, + StringArray + ); + test_function!( + RPadFunc::new(), + array_args(vec![i64::MAX]), + exec_err!( + "rpad requested length {} too large, maximum allowed length is {}", + i64::MAX, + i32::MAX + ), + &str, + Utf8, + StringArray + ); + + Ok(()) + } } diff --git a/datafusion/sqllogictest/test_files/string/string_literal.slt b/datafusion/sqllogictest/test_files/string/string_literal.slt index 690c6492a463e..aa6b9d0023b73 100644 --- a/datafusion/sqllogictest/test_files/string/string_literal.slt +++ b/datafusion/sqllogictest/test_files/string/string_literal.slt @@ -397,6 +397,17 @@ true true true query error lpad requested length 2147483648 too large SELECT lpad('hi', column1) FROM (VALUES (2147483648), (1)) AS t(column1); +# a negative length column is treated as 0 on the array path +query T +SELECT lpad('hi', column1, 'x') FROM (VALUES (-1), (3)) AS t(column1); +---- +(empty) +xhi + +# a length column far above i32::MAX is an error on the array path +query error lpad requested length 9223372036854775807 too large +SELECT lpad('hi', column1) FROM (VALUES (9223372036854775807), (1)) AS t(column1); + # a large scalar target length skips the scalar fast path query I SELECT character_length(lpad('x', 16385, 'a')); @@ -817,6 +828,17 @@ true true true query error rpad requested length 2147483648 too large SELECT rpad('hi', column1) FROM (VALUES (2147483648), (1)) AS t(column1); +# a negative length column is treated as 0 on the array path +query T +SELECT rpad('hi', column1, 'x') FROM (VALUES (-1), (3)) AS t(column1); +---- +(empty) +hix + +# a length column far above i32::MAX is an error on the array path +query error rpad requested length 9223372036854775807 too large +SELECT rpad('hi', column1) FROM (VALUES (9223372036854775807), (1)) AS t(column1); + # a large scalar target length skips the scalar fast path query I SELECT character_length(rpad('x', 16385, 'a')); From 46c2b047e1f1e794cc3ea95f09fa6da2933adaaf Mon Sep 17 00:00:00 2001 From: gonappuccino Date: Thu, 5 Nov 2026 23:18:35 -0500 Subject: [PATCH 2/2] test: check pad_data_capacity directly and clarify comments --- datafusion/functions/src/unicode/common.rs | 29 +++++++++++- datafusion/functions/src/unicode/lpad.rs | 51 +--------------------- datafusion/functions/src/unicode/rpad.rs | 39 +---------------- 3 files changed, 32 insertions(+), 87 deletions(-) diff --git a/datafusion/functions/src/unicode/common.rs b/datafusion/functions/src/unicode/common.rs index 1d63eab1aa637..859374cdba3bb 100644 --- a/datafusion/functions/src/unicode/common.rs +++ b/datafusion/functions/src/unicode/common.rs @@ -44,9 +44,12 @@ pub(crate) fn pad_data_capacity(length_array: &Int64Array) -> usize { length_array .iter() .flatten() - // Out-of-range lengths produce "", an error or null, so add no bytes + // Negative lengths produce "" and lengths above i32::MAX are an error, + // so neither adds bytes .filter(|len| (0..=i32::MAX as i64).contains(len)) .fold(0usize, |acc, len| acc.saturating_add(len as usize)) + // Many in-range lengths can still add up to more than one buffer + // should reserve up front .min(i32::MAX as usize) } @@ -232,3 +235,27 @@ fn general_left_right_view( Ok(Arc::new(array) as ArrayRef) } } + +#[cfg(test)] +mod tests { + use arrow::array::Int64Array; + + use super::pad_data_capacity; + + #[test] + fn test_pad_data_capacity() { + let capacity = + |lengths: Vec>| pad_data_capacity(&Int64Array::from(lengths)); + assert_eq!(capacity(vec![]), 0); + // Negative lengths produce "" and lengths above i32::MAX are an error + assert_eq!(capacity(vec![Some(3), None, Some(-1), Some(i64::MIN)]), 3); + assert_eq!( + capacity(vec![Some(3), Some(i32::MAX as i64 + 1), Some(i64::MAX)]), + 3 + ); + assert_eq!( + capacity(vec![Some(i32::MAX as i64), Some(1)]), + i32::MAX as usize + ); + } +} diff --git a/datafusion/functions/src/unicode/lpad.rs b/datafusion/functions/src/unicode/lpad.rs index 92e0e15587281..75a645329aac9 100644 --- a/datafusion/functions/src/unicode/lpad.rs +++ b/datafusion/functions/src/unicode/lpad.rs @@ -81,12 +81,10 @@ mod tests { use crate::unicode::lpad::LPadFunc; use crate::utils::test::test_function; - use std::sync::Arc; - - use arrow::array::{Array, Int64Array, LargeStringArray, StringArray}; + use arrow::array::{Array, LargeStringArray, StringArray}; use arrow::datatypes::DataType::{LargeUtf8, Utf8}; - use datafusion_common::{Result, ScalarValue, exec_err}; + use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; macro_rules! test_lpad { @@ -347,49 +345,4 @@ mod tests { Ok(()) } - - fn array_args(lengths: Vec) -> Vec { - vec![ - ColumnarValue::Array(Arc::new(StringArray::from(vec!["hi"; lengths.len()]))), - ColumnarValue::Array(Arc::new(Int64Array::from(lengths))), - ] - } - - #[test] - fn test_array_length_out_of_range() -> Result<()> { - test_function!( - LPadFunc::new(), - array_args(vec![-1]), - Ok(Some("")), - &str, - Utf8, - StringArray - ); - test_function!( - LPadFunc::new(), - array_args(vec![1 << 62]), - exec_err!( - "lpad requested length {} too large, maximum allowed length is {}", - 1i64 << 62, - i32::MAX - ), - &str, - Utf8, - StringArray - ); - test_function!( - LPadFunc::new(), - array_args(vec![i64::MAX, i64::MAX]), - exec_err!( - "lpad requested length {} too large, maximum allowed length is {}", - i64::MAX, - i32::MAX - ), - &str, - Utf8, - StringArray - ); - - Ok(()) - } } diff --git a/datafusion/functions/src/unicode/rpad.rs b/datafusion/functions/src/unicode/rpad.rs index 66ba77814b612..0717dd1392ed9 100644 --- a/datafusion/functions/src/unicode/rpad.rs +++ b/datafusion/functions/src/unicode/rpad.rs @@ -78,12 +78,10 @@ impl ScalarUDFImpl for RPadFunc { #[cfg(test)] mod tests { - use std::sync::Arc; - - use arrow::array::{Array, Int64Array, StringArray}; + use arrow::array::{Array, StringArray}; use arrow::datatypes::DataType::Utf8; - use datafusion_common::{Result, ScalarValue, exec_err}; + use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; use crate::unicode::rpad::RPadFunc; @@ -294,37 +292,4 @@ mod tests { Ok(()) } - - fn array_args(lengths: Vec) -> Vec { - vec![ - ColumnarValue::Array(Arc::new(StringArray::from(vec!["hi"; lengths.len()]))), - ColumnarValue::Array(Arc::new(Int64Array::from(lengths))), - ] - } - - #[test] - fn test_array_length_out_of_range() -> Result<()> { - test_function!( - RPadFunc::new(), - array_args(vec![-1]), - Ok(Some("")), - &str, - Utf8, - StringArray - ); - test_function!( - RPadFunc::new(), - array_args(vec![i64::MAX]), - exec_err!( - "rpad requested length {} too large, maximum allowed length is {}", - i64::MAX, - i32::MAX - ), - &str, - Utf8, - StringArray - ); - - Ok(()) - } }