This is an automated email from the ASF dual-hosted git repository.
alamb pushed a commit to branch 57_maintenance
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/57_maintenance by this push:
new cdca3e38e8 [57_maintenance] fix: account for offsets in runarrays when
using extend for MutableArrayData (#11815)
cdca3e38e8 is described below
commit cdca3e38e82860184f69c2c0f472bf18fb40a581
Author: Jeffrey Vo <[email protected]>
AuthorDate: Fri Oct 23 00:23:45 2026 +0900
[57_maintenance] fix: account for offsets in runarrays when using extend
for MutableArrayData (#11815)
- Backport https://github.com/apache/arrow-rs/pull/11546
- For inclusion in https://github.com/apache/arrow-rs/issues/11666
Cherry-pick of
https://github.com/apache/arrow-rs/commit/2e26ab4b0baa351ca59945a1988685b974861e63
---
arrow-data/src/transform/run.rs | 81 ++++++++++++++++++++++++++++++++++++-----
1 file changed, 72 insertions(+), 9 deletions(-)
diff --git a/arrow-data/src/transform/run.rs b/arrow-data/src/transform/run.rs
index 89daa00516..ff3bbd155e 100644
--- a/arrow-data/src/transform/run.rs
+++ b/arrow-data/src/transform/run.rs
@@ -16,7 +16,7 @@
// under the License.
use super::{_MutableArrayData, ArrayData, Extend};
-use arrow_buffer::{ArrowNativeType, Buffer, ToByteSlice};
+use arrow_buffer::{ArrowNativeType, ToByteSlice};
use arrow_schema::DataType;
use num_traits::CheckedAdd;
@@ -58,8 +58,9 @@ pub fn extend_nulls(mutable: &mut _MutableArrayData, len:
usize) {
macro_rules! extend_nulls_impl {
($run_end_type:ty) => {{
let last_run_end =
get_last_run_end::<$run_end_type>(&mutable.child_data[0]);
- let new_value = last_run_end
- .checked_add(<$run_end_type as ArrowNativeType>::usize_as(len))
+ let len = <$run_end_type as ArrowNativeType>::from_usize(len);
+ let new_value = len
+ .and_then(|len| last_run_end.checked_add(len))
.expect("run end overflow");
mutable.child_data[0]
.data
@@ -81,7 +82,7 @@ pub fn extend_nulls(mutable: &mut _MutableArrayData, len:
usize) {
/// Build run ends bytes and values range directly for batch processing
fn build_extend_arrays<T: ArrowNativeType + std::ops::Add<Output = T> +
CheckedAdd>(
- buffer: &Buffer,
+ typed_slice: &[T],
length: usize,
start: usize,
len: usize,
@@ -93,9 +94,6 @@ fn build_extend_arrays<T: ArrowNativeType +
std::ops::Add<Output = T> + CheckedA
let mut prev_end = 0;
let mut current_run_end = dest_last_run_end;
- // Convert buffer to typed slice once
- let typed_slice: &[T] = buffer.typed_data();
-
for i in 0..length {
if i < typed_slice.len() {
let run_end = typed_slice[i].to_usize().unwrap();
@@ -188,6 +186,8 @@ pub fn build_extend(array: &ArrayData) -> Extend<'_> {
// We need to analyze the source array's run structure
let source_run_ends = &array.child_data()[0];
+ let source_buffer_offset = source_run_ends.offset();
+ let source_buffer_len = source_run_ends.len();
let source_buffer = &source_run_ends.buffers()[0];
// Get the run end type from the mutable array
@@ -201,6 +201,8 @@ pub fn build_extend(array: &ArrayData) -> Extend<'_> {
// Build run ends and values indices directly for batch processing
macro_rules! build_and_process_impl {
($run_end_type:ty) => {{
+ let source_buffer = &source_buffer.typed_data()
+ [source_buffer_offset..source_buffer_offset +
source_buffer_len];
let dest_last_run_end =
get_last_run_end::<$run_end_type>(&mutable.child_data[0]);
let (run_ends_bytes, values_range) =
build_extend_arrays::<$run_end_type>(
@@ -511,7 +513,8 @@ mod tests {
#[test]
fn test_build_extend_arrays_int16() {
let buffer = Buffer::from_vec(vec![3i16, 5i16, 8i16]);
- let (run_ends_bytes, values_range) =
build_extend_arrays::<i16>(&buffer, 3, 2, 4, 0i16);
+ let (run_ends_bytes, values_range) =
+ build_extend_arrays::<i16>(buffer.typed_data(), 3, 2, 4, 0i16);
// Logical array: [A, A, A, B, B, C, C, C]
// Requesting indices 2-6 should give us:
@@ -533,7 +536,8 @@ mod tests {
#[test]
fn test_build_extend_arrays_int64() {
let buffer = Buffer::from_vec(vec![3i64, 5i64, 8i64]);
- let (run_ends_bytes, values_range) =
build_extend_arrays::<i64>(&buffer, 3, 2, 4, 0i64);
+ let (run_ends_bytes, values_range) =
+ build_extend_arrays::<i64>(buffer.typed_data(), 3, 2, 4, 0i64);
// Same logic as above but with i64
assert_eq!(run_ends_bytes.len(), 3 * std::mem::size_of::<i64>());
@@ -644,4 +648,63 @@ mod tests {
// This should cause overflow: (i32::MAX - 50) + 100 > i32::MAX
mutable.extend(0, 0, 100);
}
+
+ #[test]
+ #[should_panic(expected = "run end overflow")]
+ fn test_extend_nulls_overflow() {
+ let values = create_int32_array_data(vec![10]);
+ let array = create_run_array_data_int16(vec![100], values);
+
+ let mut mutable = MutableArrayData::new(vec![&array], true, 10);
+ mutable.extend_nulls(i16::MAX as usize + 1);
+ }
+
+ #[test]
+ fn test_extend_accounts_for_offset() {
+ let dt = DataType::RunEndEncoded(
+ Arc::new(Field::new("run_ends", DataType::Int32, false)),
+ Arc::new(Field::new("values", DataType::Float32, true)),
+ );
+
+ let runs = ArrayData::builder(DataType::Int32)
+ .len(2)
+ .offset(2)
+ .add_buffer(vec![-10_i32, -20, 1_i32, 4, 0, 100].into())
+ .build()
+ .unwrap();
+ let values = ArrayData::builder(DataType::Float32)
+ .len(2)
+ .add_buffer(vec![1.5, 2.5].into())
+ .build()
+ .unwrap();
+
+ let source_data = ArrayDataBuilder::new(dt.clone())
+ .len(3)
+ .child_data(vec![runs, values])
+ .build()
+ .unwrap();
+
+ let mut mutable = MutableArrayData::new(vec![&source_data], true, 10);
+ mutable.extend(0, 0, 1);
+ let data = mutable.freeze();
+
+ let expected = {
+ let runs = ArrayData::builder(DataType::Int32)
+ .len(1)
+ .add_buffer(vec![1_i32].into())
+ .build()
+ .unwrap();
+ let values = ArrayData::builder(DataType::Float32)
+ .len(1)
+ .add_buffer(vec![1.5].into())
+ .build()
+ .unwrap();
+ ArrayDataBuilder::new(dt)
+ .len(1)
+ .child_data(vec![runs, values])
+ .build()
+ .unwrap()
+ };
+ assert_eq!(data, expected);
+ }
}
