You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/06/20 19:08:28 UTC
(datafusion) branch main updated: Return errors on string builder offset overflow in `replace` and `initcap` (#22990)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion.git
The following commit(s) were added to refs/heads/main by this push:
new 2e8a1d9685 Return errors on string builder offset overflow in `replace` and `initcap` (#22990)
2e8a1d9685 is described below
commit 2e8a1d968565c94b1e133caf144b4231ef4ef38a
Author: kosiew <ko...@gmail.com>
AuthorDate: Sun Jun 21 02:56:50 2026 +0800
Return errors on string builder offset overflow in `replace` and `initcap` (#22990)
## Which issue does this PR close?
* Part of #22688
## Rationale for this change
`GenericStringArrayBuilder` only exposed infallible append APIs that
panic when string offsets exceed the underlying offset type limits.
String functions such as `replace`, `replace_view`, and the generic
`initcap` path relied on these APIs, meaning extreme output sizes could
panic instead of returning a recoverable `DataFusionError`.
This change introduces fallible builder APIs and migrates selected
string UDFs to use them so offset overflow is reported as an error
rather than causing a panic.
## What changes are included in this PR?
* Add overflow-checked helper functions to `GenericStringArrayBuilder`:
* `try_offset`
* `try_push_offset_for_len`
* `try_append_bytes`
* Add fallible append APIs:
* `try_append_value`
* `try_append_placeholder`
* `try_append_byte_map`
* `try_append_with`
* Introduce a shared overflow error path that returns a
`DataFusionError` instead of panicking.
* Keep existing infallible append APIs for compatibility while
documenting that new overflow-sensitive call sites should prefer the
`try_*` variants.
* Refactor `replace` and `replace_view` to share a generic
`replace_arrays` implementation.
* Change `apply_replace` to return `Result<()>` and propagate errors
from builder operations.
* Update `replace`/`replace_view` to use the new fallible builder APIs
and thread errors with `?`.
* Update the generic `Utf8`/`LargeUtf8` path in `initcap` to use
`try_append_placeholder` and `try_append_value`.
* Add rollback handling in `try_append_with` so builder state is
restored if offset validation fails.
## Are these changes tested?
Yes.
Added tests in `datafusion/functions/src/strings.rs`:
* `generic_string_builder_try_append_success_path`
* `generic_string_builder_mixed_append_success_path`
* `generic_string_builder_try_offset_overflow`
* `generic_string_builder_try_append_bytes_overflow`
Existing `replace` and `initcap` tests remain in place and the migrated
code paths continue to be exercised by those test suites.
## Are there any user-facing changes?
Yes.
For extreme string outputs that exceed the offset limits of the
underlying string array type, affected functions now return a
`DataFusionError` instead of panicking. Normal behavior and results are
otherwise unchanged.
## LLM-generated code disclosure
This PR includes LLM-generated code and comments. All LLM-generated
content has been manually reviewed.
---
datafusion/functions/src/string/replace.rs | 87 ++++-------
datafusion/functions/src/strings.rs | 221 ++++++++++++++++++++++++++--
datafusion/functions/src/unicode/initcap.rs | 6 +-
3 files changed, 242 insertions(+), 72 deletions(-)
diff --git a/datafusion/functions/src/string/replace.rs b/datafusion/functions/src/string/replace.rs
index 28f81769f5..a2fda21461 100644
--- a/datafusion/functions/src/string/replace.rs
+++ b/datafusion/functions/src/string/replace.rs
@@ -17,13 +17,11 @@
use std::sync::Arc;
-use arrow::array::{Array, ArrayRef, OffsetSizeTrait};
+use arrow::array::{ArrayRef, OffsetSizeTrait, StringArrayType};
use arrow::buffer::NullBuffer;
use arrow::datatypes::DataType;
-use crate::strings::{
- BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringWriter,
-};
+use crate::strings::{GenericStringArrayBuilder, StringWriter};
use crate::utils::{make_scalar_function, utf8_to_str_type};
use datafusion_common::cast::{as_generic_string_array, as_string_view_array};
use datafusion_common::types::logical_string;
@@ -164,40 +162,7 @@ fn replace_view(args: &[ArrayRef]) -> Result<ArrayRef> {
let from_array = as_string_view_array(&args[1])?;
let to_array = as_string_view_array(&args[2])?;
- let len = string_array.len();
- let mut builder = GenericStringArrayBuilder::<i32>::with_capacity(len, 0);
- let nulls = NullBuffer::union_many([
- string_array.nulls(),
- from_array.nulls(),
- to_array.nulls(),
- ]);
-
- // Hoist the nulls.is_some() check out of the loop. LLVM does not always
- // unswitch this loop on its own (the Utf8View body is large enough to
- // exceed its cost-benefit threshold).
- if let Some(nulls_ref) = nulls.as_ref() {
- for i in 0..len {
- if nulls_ref.is_null(i) {
- builder.append_placeholder();
- continue;
- }
- // SAFETY: union of input nulls is non-null at i, so each input is too.
- let string = unsafe { string_array.value_unchecked(i) };
- let from = unsafe { from_array.value_unchecked(i) };
- let to = unsafe { to_array.value_unchecked(i) };
- apply_replace(&mut builder, string, from, to);
- }
- } else {
- for i in 0..len {
- // SAFETY: i < len, and no input has a null buffer.
- let string = unsafe { string_array.value_unchecked(i) };
- let from = unsafe { from_array.value_unchecked(i) };
- let to = unsafe { to_array.value_unchecked(i) };
- apply_replace(&mut builder, string, from, to);
- }
- }
-
- Ok(Arc::new(builder.finish(nulls)?) as ArrayRef)
+ replace_arrays::<_, i32>(string_array, from_array, to_array)
}
/// Replaces all occurrences in string of substring from with substring to.
@@ -207,28 +172,39 @@ fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
let from_array = as_generic_string_array::<T>(&args[1])?;
let to_array = as_generic_string_array::<T>(&args[2])?;
+ replace_arrays::<_, T>(string_array, from_array, to_array)
+}
+
+fn replace_arrays<'a, S, O>(
+ string_array: S,
+ from_array: S,
+ to_array: S,
+) -> Result<ArrayRef>
+where
+ S: StringArrayType<'a> + Copy,
+ O: OffsetSizeTrait,
+{
let len = string_array.len();
- let mut builder = GenericStringArrayBuilder::<T>::with_capacity(len, 0);
+ let mut builder = GenericStringArrayBuilder::<O>::with_capacity(len, 0);
let nulls = NullBuffer::union_many([
string_array.nulls(),
from_array.nulls(),
to_array.nulls(),
]);
- // Hoist the nulls.is_some() check out of the loop. LLVM unswitches this
- // automatically today, but kept explicit so the no-nulls fast path is not
- // contingent on the optimizer's cost heuristic.
+ // Hoist the nulls.is_some() check out of the loop so the no-nulls fast
+ // path does not depend on LLVM loop-unswitching heuristics.
if let Some(nulls_ref) = nulls.as_ref() {
for i in 0..len {
if nulls_ref.is_null(i) {
- builder.append_placeholder();
+ builder.try_append_placeholder()?;
continue;
}
// SAFETY: union of input nulls is non-null at i, so each input is too.
let string = unsafe { string_array.value_unchecked(i) };
let from = unsafe { from_array.value_unchecked(i) };
let to = unsafe { to_array.value_unchecked(i) };
- apply_replace(&mut builder, string, from, to);
+ apply_replace(&mut builder, string, from, to)?;
}
} else {
for i in 0..len {
@@ -236,7 +212,7 @@ fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
let string = unsafe { string_array.value_unchecked(i) };
let from = unsafe { from_array.value_unchecked(i) };
let to = unsafe { to_array.value_unchecked(i) };
- apply_replace(&mut builder, string, from, to);
+ apply_replace(&mut builder, string, from, to)?;
}
}
@@ -244,12 +220,12 @@ fn replace<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
}
#[inline]
-fn apply_replace<B: BulkNullStringArrayBuilder>(
- builder: &mut B,
+fn apply_replace<O: OffsetSizeTrait>(
+ builder: &mut GenericStringArrayBuilder<O>,
string: &str,
from: &str,
to: &str,
-) {
+) -> Result<()> {
// Hot path: single ASCII byte → single ASCII byte. An ASCII byte (< 0x80)
// cannot appear inside a multi-byte UTF-8 sequence, so any multi-byte
// sequences in `string` pass through unchanged and output stays valid
@@ -259,21 +235,19 @@ fn apply_replace<B: BulkNullStringArrayBuilder>(
&& to_byte.is_ascii()
{
// SAFETY: see the contract above.
- unsafe {
- builder.append_byte_map(string.as_bytes(), |b| {
+ return unsafe {
+ builder.try_append_byte_map(string.as_bytes(), |b| {
if b == from_byte { to_byte } else { b }
- });
- }
- return;
+ })
+ };
}
if from.is_empty() {
// PostgreSQL returns the input unchanged when `from` is empty (#22253).
- builder.append_value(string);
- return;
+ return builder.try_append_value(string);
}
- builder.append_with(|w| replace_into_writer(w, string, from, to));
+ builder.try_append_with(|w| replace_into_writer(w, string, from, to))
}
#[inline]
@@ -291,6 +265,7 @@ fn replace_into_writer<W: StringWriter>(w: &mut W, string: &str, from: &str, to:
mod tests {
use super::*;
use crate::utils::test::test_function;
+ use arrow::array::Array;
use arrow::array::LargeStringArray;
use arrow::array::StringArray;
use arrow::datatypes::DataType::{LargeUtf8, Utf8};
diff --git a/datafusion/functions/src/strings.rs b/datafusion/functions/src/strings.rs
index 144d567f5b..0ef27ad8bb 100644
--- a/datafusion/functions/src/strings.rs
+++ b/datafusion/functions/src/strings.rs
@@ -19,7 +19,7 @@ use std::marker::PhantomData;
use std::mem::size_of;
use std::sync::Arc;
-use datafusion_common::{Result, exec_datafusion_err, internal_err};
+use datafusion_common::{DataFusionError, Result, exec_datafusion_err, internal_err};
use arrow::array::{
Array, ArrayAccessor, ArrayDataBuilder, ArrayRef, BinaryArray, ByteView,
@@ -453,6 +453,20 @@ pub(crate) struct GenericStringArrayBuilder<O: OffsetSizeTrait> {
_phantom: PhantomData<O>,
}
+fn offset_overflow_error<O: OffsetSizeTrait>() -> DataFusionError {
+ exec_datafusion_err!(
+ "byte array offset overflow: output size exceeds {} bytes",
+ O::MAX_OFFSET
+ )
+}
+
+fn try_offset<O: OffsetSizeTrait>(len: usize) -> Result<O> {
+ if len > O::MAX_OFFSET {
+ return Err(offset_overflow_error::<O>());
+ }
+ Ok(O::usize_as(len))
+}
+
impl<O: OffsetSizeTrait> GenericStringArrayBuilder<O> {
pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self {
let capacity = item_capacity
@@ -470,30 +484,107 @@ impl<O: OffsetSizeTrait> GenericStringArrayBuilder<O> {
}
}
+ #[inline]
+ fn try_push_offset_for_len(&mut self, len: usize) -> Result<()> {
+ let next_offset = try_offset::<O>(len)?;
+ self.offsets_buffer.push(next_offset);
+ Ok(())
+ }
+
+ #[inline]
+ fn try_append_bytes<F>(&mut self, additional_len: usize, append: F) -> Result<()>
+ where
+ F: FnOnce(&mut MutableBuffer),
+ {
+ let next_len = self
+ .value_buffer
+ .len()
+ .checked_add(additional_len)
+ .ok_or_else(offset_overflow_error::<O>)?;
+ let next_offset = try_offset::<O>(next_len)?;
+ append(&mut self.value_buffer);
+ debug_assert_eq!(self.value_buffer.len(), next_len);
+ self.offsets_buffer.push(next_offset);
+ Ok(())
+ }
+
+ /// Fallible variant of [`Self::append_value`].
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the cumulative byte length exceeds this builder's
+ /// offset type limit.
+ #[inline]
+ pub fn try_append_value(&mut self, value: &str) -> Result<()> {
+ self.try_append_bytes(value.len(), |value_buffer| {
+ value_buffer.extend_from_slice(value.as_bytes());
+ })
+ }
+
+ /// Fallible variant of [`Self::append_placeholder`].
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the current cumulative byte length exceeds this
+ /// builder's offset type limit.
+ #[inline]
+ pub fn try_append_placeholder(&mut self) -> Result<()> {
+ self.try_push_offset_for_len(self.value_buffer.len())?;
+ self.placeholder_count += 1;
+ Ok(())
+ }
+
/// See [`BulkNullStringArrayBuilder::append_value`].
///
+ /// Note: new call sites that need recoverable overflow handling should
+ /// prefer [`Self::try_append_value`].
+ ///
/// # Panics
///
/// Panics if the cumulative byte length exceeds `O::MAX`.
#[inline]
pub fn append_value(&mut self, value: &str) {
- self.value_buffer.extend_from_slice(value.as_bytes());
- let next_offset =
- O::from_usize(self.value_buffer.len()).expect("byte array offset overflow");
- self.offsets_buffer.push(next_offset);
+ self.try_append_value(value)
+ .expect("byte array offset overflow");
}
/// See [`BulkNullStringArrayBuilder::append_placeholder`].
+ ///
+ /// Note: new call sites that need recoverable overflow handling should
+ /// prefer [`Self::try_append_placeholder`].
#[inline]
pub fn append_placeholder(&mut self) {
- let next_offset =
- O::from_usize(self.value_buffer.len()).expect("byte array offset overflow");
- self.offsets_buffer.push(next_offset);
- self.placeholder_count += 1;
+ self.try_append_placeholder()
+ .expect("byte array offset overflow");
+ }
+
+ /// Fallible variant of [`Self::append_byte_map`].
+ ///
+ /// # Safety
+ ///
+ /// The bytes produced by applying `map` to each byte of `src`, in order,
+ /// must form valid UTF-8.
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the cumulative byte length exceeds this builder's
+ /// offset type limit.
+ #[inline]
+ pub unsafe fn try_append_byte_map<F: FnMut(u8) -> u8>(
+ &mut self,
+ src: &[u8],
+ mut map: F,
+ ) -> Result<()> {
+ self.try_append_bytes(src.len(), |value_buffer| {
+ value_buffer.extend(src.iter().map(|&b| map(b)));
+ })
}
/// See [`BulkNullStringArrayBuilder::append_byte_map`].
///
+ /// Note: new call sites that need recoverable overflow handling should
+ /// prefer [`Self::try_append_byte_map`].
+ ///
/// # Safety
///
/// The bytes produced by applying `map` to each byte of `src`, in order,
@@ -503,15 +594,46 @@ impl<O: OffsetSizeTrait> GenericStringArrayBuilder<O> {
///
/// Panics if the cumulative byte length exceeds `O::MAX`.
#[inline]
- pub unsafe fn append_byte_map<F: FnMut(u8) -> u8>(&mut self, src: &[u8], mut map: F) {
- self.value_buffer.extend(src.iter().map(|&b| map(b)));
- let next_offset =
- O::from_usize(self.value_buffer.len()).expect("byte array offset overflow");
+ pub unsafe fn append_byte_map<F: FnMut(u8) -> u8>(&mut self, src: &[u8], map: F) {
+ // SAFETY: caller upholds this method's UTF-8 contract.
+ unsafe { self.try_append_byte_map(src, map) }
+ .expect("byte array offset overflow");
+ }
+
+ /// Fallible variant of [`Self::append_with`].
+ ///
+ /// # Errors
+ ///
+ /// Returns an error if the cumulative byte length exceeds this builder's
+ /// offset type limit.
+ #[inline]
+ pub fn try_append_with<F>(&mut self, f: F) -> Result<()>
+ where
+ F: FnOnce(&mut GenericStringWriter<'_>),
+ {
+ let old_len = self.value_buffer.len();
+ let mut writer = GenericStringWriter {
+ value_buffer: &mut self.value_buffer,
+ };
+ f(&mut writer);
+ let next_offset = match try_offset::<O>(self.value_buffer.len()) {
+ Ok(offset) => offset,
+ Err(e) => {
+ // SAFETY: `old_len` was the initialized length before `f` wrote to
+ // this owned buffer, so shrinking back preserves initialized data.
+ unsafe { self.value_buffer.set_len(old_len) };
+ return Err(e);
+ }
+ };
self.offsets_buffer.push(next_offset);
+ Ok(())
}
/// See [`BulkNullStringArrayBuilder::append_with`].
///
+ /// Note: new call sites that need recoverable overflow handling should
+ /// prefer [`Self::try_append_with`].
+ ///
/// # Panics
///
/// Panics if the cumulative byte length exceeds `O::MAX`.
@@ -520,6 +642,9 @@ impl<O: OffsetSizeTrait> GenericStringArrayBuilder<O> {
where
F: FnOnce(&mut GenericStringWriter<'_>),
{
+ // Do not delegate to `try_append_with`: it rolls back value_buffer on
+ // overflow before returning Err, which would change this infallible
+ // method's state if its panic is caught.
let mut writer = GenericStringWriter {
value_buffer: &mut self.value_buffer,
};
@@ -1388,6 +1513,76 @@ mod tests {
assert_finish_errs_on_length_mismatch(StringViewArrayBuilder::with_capacity(2));
}
+ #[test]
+ fn generic_string_builder_try_append_success_path() {
+ let mut builder = GenericStringArrayBuilder::<i32>::with_capacity(4, 16);
+ builder.try_append_value("abc").unwrap();
+ builder.try_append_placeholder().unwrap();
+ // SAFETY: ASCII input and output.
+ unsafe {
+ builder
+ .try_append_byte_map(b"de", |b| b.to_ascii_uppercase())
+ .unwrap();
+ }
+ builder
+ .try_append_with(|w| {
+ w.write_str("f");
+ w.write_char('é');
+ })
+ .unwrap();
+
+ let nulls = Some(NullBuffer::from(vec![true, false, true, true]));
+ let array = builder.finish(nulls).unwrap();
+ assert_eq!(
+ &array,
+ &StringArray::from(vec![Some("abc"), None, Some("DE"), Some("fé")])
+ );
+ }
+
+ #[test]
+ fn generic_string_builder_mixed_append_success_path() {
+ let mut builder = GenericStringArrayBuilder::<i32>::with_capacity(4, 16);
+ builder.append_value("ab");
+ builder.try_append_value("cd").unwrap();
+ // SAFETY: ASCII input and output.
+ unsafe {
+ builder.append_byte_map(b"ef", |b| b.to_ascii_uppercase());
+ builder
+ .try_append_byte_map(b"gh", |b| b.to_ascii_uppercase())
+ .unwrap();
+ }
+
+ let array = builder.finish(None).unwrap();
+ assert_eq!(
+ &array,
+ &StringArray::from(vec![Some("ab"), Some("cd"), Some("EF"), Some("GH")])
+ );
+ }
+
+ #[test]
+ fn generic_string_builder_try_offset_overflow() {
+ let err = try_offset::<i32>(i32::MAX as usize + 1)
+ .unwrap_err()
+ .to_string();
+ assert!(
+ err.contains("byte array offset overflow"),
+ "unexpected error: {err}"
+ );
+ }
+
+ #[test]
+ fn generic_string_builder_try_append_bytes_overflow() {
+ let mut builder = GenericStringArrayBuilder::<i32>::with_capacity(0, 0);
+ let err = builder
+ .try_append_bytes(i32::MAX as usize + 1, |_| unreachable!())
+ .unwrap_err()
+ .to_string();
+ assert!(
+ err.contains("byte array offset overflow"),
+ "unexpected error: {err}"
+ );
+ }
+
#[test]
#[cfg(debug_assertions)]
#[should_panic(expected = "placeholder rows")]
diff --git a/datafusion/functions/src/unicode/initcap.rs b/datafusion/functions/src/unicode/initcap.rs
index 711b2c49b0..9192f23844 100644
--- a/datafusion/functions/src/unicode/initcap.rs
+++ b/datafusion/functions/src/unicode/initcap.rs
@@ -166,12 +166,12 @@ fn initcap<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
if let Some(ref n) = nulls {
for i in 0..len {
if n.is_null(i) {
- builder.append_placeholder();
+ builder.try_append_placeholder()?;
} else {
// SAFETY: not null per check above.
let s = unsafe { string_array.value_unchecked(i) };
initcap_string(s, &mut container);
- builder.append_value(&container);
+ builder.try_append_value(&container)?;
}
}
} else {
@@ -179,7 +179,7 @@ fn initcap<T: OffsetSizeTrait>(args: &[ArrayRef]) -> Result<ArrayRef> {
// SAFETY: no null buffer means every index is valid.
let s = unsafe { string_array.value_unchecked(i) };
initcap_string(s, &mut container);
- builder.append_value(&container);
+ builder.try_append_value(&container)?;
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org