[58_maintenance] fix: account for offsets in runarrays when using extend for MutableArrayData

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,45 @@ 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)
+            .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();
+        assert_eq!(data, source_data.slice(0, 1));
+    }
 }