fix: [9018]Fixed RunArray slice offsets(row, cast, eq)

manishkr · manishkr · commit 14f94b2f1ab4 · 2026-01-18T11:59:07.000+05:30
diff --git a/arrow-array/src/array/run_array.rs b/arrow-array/src/array/run_array.rs
@@ -141,6 +141,9 @@ impl<R: RunEndIndexType> RunArray<R> {
     ///
     /// [`values`]: Self::values
     pub fn values_slice(&self) -> ArrayRef {
+        if self.len() == 0 {
+            return self.values.slice(0, 0);
+        }
         let start = self.get_start_physical_index();
         let end = self.get_end_physical_index();
         self.values.slice(start, end - start + 1)
@@ -1173,4 +1176,56 @@ mod tests {
         let values_slice2 = values_slice2.as_primitive::<Int32Type>();
         assert_eq!(values_slice2.values(), &[1]);
     }
+
+    #[test]
+    fn test_run_array_eq_diff_physical_same_logical() {
+        let run_ends1 = Int32Array::from(vec![1, 3, 6]);
+        let values1 = StringArray::from(vec!["a", "b", "c"]);
+        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
+
+        let run_ends2 = Int32Array::from(vec![1, 2, 3, 4, 5, 6]);
+        let values2 = StringArray::from(vec!["a", "b", "b", "c", "c", "c"]);
+        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
+
+        assert_eq!(array1, array2);
+    }
+
+    #[test]
+    fn test_run_array_eq_sliced() {
+        let run_ends1 = Int32Array::from(vec![2, 5, 10]);
+        let values1 = StringArray::from(vec!["a", "b", "c"]);
+        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
+        // Logical: a, a, b, b, b, c, c, c, c, c
+
+        let slice1 = array1.slice(1, 6);
+        // Logical: a, b, b, b, c, c
+
+        let run_ends2 = Int32Array::from(vec![1, 4, 6]);
+        let values2 = StringArray::from(vec!["a", "b", "c"]);
+        let array2 = RunArray::<Int32Type>::try_new(&run_ends2, &values2).unwrap();
+        // Logical: a, b, b, b, c, c
+
+        assert_eq!(slice1, array2);
+
+        let slice2 = array1.slice(2, 3);
+        // Logical: b, b, b
+        let run_ends3 = Int32Array::from(vec![3]);
+        let values3 = StringArray::from(vec!["b"]);
+        let array3 = RunArray::<Int32Type>::try_new(&run_ends3, &values3).unwrap();
+        assert_eq!(slice2, array3);
+    }
+
+    #[test]
+    fn test_run_array_eq_sliced_different_offsets() {
+        let run_ends1 = Int32Array::from(vec![2, 5, 10]);
+        let values1 = StringArray::from(vec!["a", "b", "c"]);
+        let array1 = RunArray::<Int32Type>::try_new(&run_ends1, &values1).unwrap();
+
+        let slice1 = array1.slice(1, 4); // a, b, b, b
+        let slice2 = array1.slice(1, 4);
+        assert_eq!(slice1, slice2);
+
+        let slice3 = array1.slice(0, 4); // a, a, b, b
+        assert_ne!(slice1, slice3);
+    }
 }
diff --git a/arrow-buffer/src/buffer/run.rs b/arrow-buffer/src/buffer/run.rs
@@ -199,9 +199,14 @@ where
     pub fn sliced_values(&self) -> impl Iterator<Item = E> + '_ {
         let offset = self.logical_offset;
         let len = self.logical_length;
-        let start = self.get_start_physical_index();
-        let end = self.get_end_physical_index();
-        self.run_ends[start..=end].iter().map(move |&val| {
+        let physical_slice = if len == 0 {
+            &self.run_ends[0..0]
+        } else {
+            let start = self.get_start_physical_index();
+            let end = self.get_end_physical_index();
+            &self.run_ends[start..=end]
+        };
+        physical_slice.iter().map(move |&val| {
             let val = val.as_usize().saturating_sub(offset).min(len);
             E::from_usize(val).unwrap()
         })
diff --git a/arrow-cast/src/cast/mod.rs b/arrow-cast/src/cast/mod.rs
@@ -12451,4 +12451,32 @@ mod tests {
             assert_eq!(casted.as_ref(), &expected);
         }
     }
+
+    #[test]
+    fn test_cast_between_sliced_run_end_encoded() {
+        let run_ends = Int16Array::from(vec![2, 5, 8]);
+        let values = StringArray::from(vec!["a", "b", "c"]);
+
+        let ree_array = RunArray::<Int16Type>::try_new(&run_ends, &values).unwrap();
+        let ree_array = ree_array.slice(1, 2);
+        let array_ref = Arc::new(ree_array) as ArrayRef;
+
+        let target_type = DataType::RunEndEncoded(
+            Arc::new(Field::new("run_ends", DataType::Int64, false)),
+            Arc::new(Field::new("values", DataType::Utf8, true)),
+        );
+        let cast_options = CastOptions {
+            safe: false,
+            format_options: FormatOptions::default(),
+        };
+
+        let result = cast_with_options(&array_ref, &target_type, &cast_options).unwrap();
+        let run_array = result.as_run::<Int64Type>();
+        let run_array = run_array.downcast::<StringArray>().unwrap();
+
+        let expected = vec!["a", "b"];
+        let actual = run_array.into_iter().flatten().collect::<Vec<_>>();
+
+        assert_eq!(expected, actual);
+    }
 }
diff --git a/arrow-cast/src/cast/run_array.rs b/arrow-cast/src/cast/run_array.rs
@@ -32,17 +32,18 @@ pub(crate) fn run_end_encoded_cast<K: RunEndIndexType>(
                 .downcast_ref::<RunArray<K>>()
                 .ok_or_else(|| ArrowError::CastError("Expected RunArray".to_string()))?;
 
-            let values = run_array.values();
-
             match to_type {
                 // Stay as RunEndEncoded, cast only the values
                 DataType::RunEndEncoded(target_index_field, target_value_field) => {
-                    let cast_values =
-                        cast_with_options(values, target_value_field.data_type(), cast_options)?;
+                    let values = run_array.values_slice();
+                    let cast_values = cast_with_options(
+                        values.as_ref(),
+                        target_value_field.data_type(),
+                        cast_options,
+                    )?;
 
-                    let run_ends_array = PrimitiveArray::<K>::from_iter_values(
-                        run_array.run_ends().values().iter().copied(),
-                    );
+                    let run_ends_array =
+                        PrimitiveArray::<K>::from_iter_values(run_array.run_ends().sliced_values());
                     let cast_run_ends = cast_with_options(
                         &run_ends_array,
                         target_index_field.data_type(),
@@ -72,6 +73,7 @@ pub(crate) fn run_end_encoded_cast<K: RunEndIndexType>(
 
                 // Expand to logical form
                 _ => {
+                    let values = run_array.values();
                     let len = run_array.len();
                     let offset = run_array.offset();
                     let run_ends = run_array.run_ends().values();
diff --git a/arrow-data/src/equal/run.rs b/arrow-data/src/equal/run.rs
@@ -16,71 +16,118 @@
 // under the License.
 
 use crate::data::ArrayData;
+use arrow_buffer::ArrowNativeType;
+use arrow_schema::DataType;
+use num_traits::ToPrimitive;
 
 use super::equal_range;
 
-/// The current implementation of comparison of run array support physical comparison.
-/// Comparing run encoded array based on logical indices (`lhs_start`, `rhs_start`) will
-/// be time consuming as converting from logical index to physical index cannot be done
-/// in constant time. The current comparison compares the underlying physical arrays.
 pub(super) fn run_equal(
     lhs: &ArrayData,
     rhs: &ArrayData,
     lhs_start: usize,
     rhs_start: usize,
     len: usize,
 ) -> bool {
-    if lhs_start != 0
-        || rhs_start != 0
-        || (lhs.len() != len && rhs.len() != len)
-        || lhs.offset() > 0
-        || rhs.offset() > 0
-    {
-        unimplemented!("Logical comparison for run array not supported.")
+    let lhs_index_type = match lhs.data_type() {
+        DataType::RunEndEncoded(f, _) => f.data_type(),
+        _ => unreachable!(),
+    };
+
+    match lhs_index_type {
+        DataType::Int16 => run_equal_inner::<i16>(lhs, rhs, lhs_start, rhs_start, len),
+        DataType::Int32 => run_equal_inner::<i32>(lhs, rhs, lhs_start, rhs_start, len),
+        DataType::Int64 => run_equal_inner::<i64>(lhs, rhs, lhs_start, rhs_start, len),
+        _ => unreachable!(),
     }
+}
 
-    if lhs.len() != rhs.len() {
-        return false;
+fn run_equal_inner<T: ArrowNativeType + ToPrimitive>(
+    lhs: &ArrayData,
+    rhs: &ArrayData,
+    lhs_start: usize,
+    rhs_start: usize,
+    len: usize,
+) -> bool {
+    if len == 0 {
+        return true;
     }
+    // RunEndEncoded arrays are guaranteed to have at least 2 children [run_ends, values]
+    let lhs_run_ends_data = &lhs.child_data()[0];
+    let rhs_run_ends_data = &rhs.child_data()[0];
+    let lhs_values = &lhs.child_data()[1];
+    let rhs_values = &rhs.child_data()[1];
 
-    let lhs_child_data = lhs.child_data();
-    let lhs_run_ends_array = &lhs_child_data[0];
-    let lhs_values_array = &lhs_child_data[1];
+    let lhs_run_ends = &lhs_run_ends_data.buffers()[0].typed_data::<T>()
+        [lhs_run_ends_data.offset()..lhs_run_ends_data.offset() + lhs_run_ends_data.len()];
+    let rhs_run_ends = &rhs_run_ends_data.buffers()[0].typed_data::<T>()
+        [rhs_run_ends_data.offset()..rhs_run_ends_data.offset() + rhs_run_ends_data.len()];
 
-    let rhs_child_data = rhs.child_data();
-    let rhs_run_ends_array = &rhs_child_data[0];
-    let rhs_values_array = &rhs_child_data[1];
+    let lhs_abs_start = lhs.offset() + lhs_start;
+    let rhs_abs_start = rhs.offset() + rhs_start;
+    let lhs_abs_end = lhs_abs_start + len;
+    let rhs_abs_end = rhs_abs_start + len;
 
-    if lhs_run_ends_array.len() != rhs_run_ends_array.len() {
-        return false;
-    }
+    let l_start_phys = find_physical_index(lhs_run_ends, lhs_abs_start);
+    let r_start_phys = find_physical_index(rhs_run_ends, rhs_abs_start);
+
+    let l_end_phys = find_physical_index(lhs_run_ends, lhs_abs_end - 1);
+    let r_end_phys = find_physical_index(rhs_run_ends, rhs_abs_end - 1);
+
+    let mut l_phys = l_start_phys;
+    let mut r_phys = r_start_phys;
 
-    if lhs_values_array.len() != rhs_values_array.len() {
-        return false;
+    let l_runs = l_end_phys - l_start_phys + 1;
+    let r_runs = r_end_phys - r_start_phys + 1;
+
+    if l_runs == r_runs {
+        let physical_match = lhs_run_ends[l_start_phys..l_end_phys]
+            .iter()
+            .zip(&rhs_run_ends[r_start_phys..r_end_phys])
+            .all(|(l, r)| l.as_usize() - lhs_abs_start == r.as_usize() - rhs_abs_start);
+
+        if physical_match {
+            return equal_range(lhs_values, rhs_values, l_start_phys, r_start_phys, l_runs);
+        }
     }
 
-    // check run ends array are equal. The length of the physical array
-    // is used to validate the child arrays.
-    let run_ends_equal = equal_range(
-        lhs_run_ends_array,
-        rhs_run_ends_array,
-        lhs_start,
-        rhs_start,
-        lhs_run_ends_array.len(),
-    );
-
-    // if run ends array are not the same return early without validating
-    // values array.
-    if !run_ends_equal {
-        return false;
+    let mut processed = 0;
+    while processed < len {
+        if !equal_range(lhs_values, rhs_values, l_phys, r_phys, 1) {
+            return false;
+        }
+
+        let l_run_end = lhs_run_ends[l_phys].as_usize();
+        let r_run_end = rhs_run_ends[r_phys].as_usize();
+
+        let l_remaining = l_run_end - (lhs_abs_start + processed);
+        let r_remaining = r_run_end - (rhs_abs_start + processed);
+
+        let step = l_remaining.min(r_remaining).min(len - processed);
+        processed += step;
+
+        if processed < len {
+            if lhs_abs_start + processed == l_run_end {
+                l_phys += 1;
+            }
+            if rhs_abs_start + processed == r_run_end {
+                r_phys += 1;
+            }
+        }
     }
 
-    // check values array are equal
-    equal_range(
-        lhs_values_array,
-        rhs_values_array,
-        lhs_start,
-        rhs_start,
-        rhs_values_array.len(),
-    )
+    true
+}
+
+fn find_physical_index<T: ArrowNativeType + ToPrimitive>(
+    run_ends: &[T],
+    logical_index: usize,
+) -> usize {
+    if logical_index == 0 {
+        return 0;
+    }
+    match run_ends.binary_search_by(|val| val.as_usize().cmp(&logical_index)) {
+        Ok(idx) => idx + 1,
+        Err(idx) => idx,
+    }
 }
diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs
@@ -657,14 +657,14 @@ impl Codec {
             Codec::RunEndEncoded(converter) => {
                 let values = match array.data_type() {
                     DataType::RunEndEncoded(r, _) => match r.data_type() {
-                        DataType::Int16 => array.as_run::<Int16Type>().values(),
-                        DataType::Int32 => array.as_run::<Int32Type>().values(),
-                        DataType::Int64 => array.as_run::<Int64Type>().values(),
+                        DataType::Int16 => array.as_run::<Int16Type>().values_slice(),
+                        DataType::Int32 => array.as_run::<Int32Type>().values_slice(),
+                        DataType::Int64 => array.as_run::<Int64Type>().values_slice(),
                         _ => unreachable!("Unsupported run end index type: {r:?}"),
                     },
                     _ => unreachable!(),
                 };
-                let rows = converter.convert_columns(std::slice::from_ref(values))?;
+                let rows = converter.convert_columns(std::slice::from_ref(&values))?;
                 Ok(Encoder::RunEndEncoded(rows))
             }
             Codec::Union(converters, _) => {
diff --git a/arrow-row/src/run.rs b/arrow-row/src/run.rs
@@ -27,11 +27,11 @@ pub fn compute_lengths<R: RunEndIndexType>(
     rows: &Rows,
     array: &RunArray<R>,
 ) {
-    let run_ends = array.run_ends().values();
+    let run_ends = array.run_ends().sliced_values();
     let mut logical_start = 0;
 
     // Iterate over each run and apply the same length to all logical positions in the run
-    for (physical_idx, &run_end) in run_ends.iter().enumerate() {
+    for (physical_idx, run_end) in run_ends.enumerate() {
         let logical_end = run_end.as_usize();
         let row_len = rows.row_len(physical_idx);
         let encoded_len = variable::padded_length(Some(row_len));
@@ -55,14 +55,14 @@ pub fn encode<R: RunEndIndexType>(
     opts: SortOptions,
     array: &RunArray<R>,
 ) {
-    let run_ends = array.run_ends();
+    let run_ends = array.run_ends().sliced_values();
 
     let mut logical_idx = 0;
     let mut offset_idx = 1; // Skip first offset
 
     // Iterate over each run
-    for physical_idx in 0..run_ends.values().len() {
-        let run_end = run_ends.values()[physical_idx].as_usize();
+    for (physical_idx, run_end) in run_ends.enumerate() {
+        let run_end = run_end.as_usize();
 
         // Process all elements in this run
         while logical_idx < run_end && offset_idx < offsets.len() {
@@ -639,4 +639,15 @@ mod tests {
         let result_ree = arrays[0].as_run::<Int32Type>();
         assert_eq!(result_ree.len(), 0);
     }
+
+    #[test]
+    fn test_run_end_encoded_round_trip_sliced() {
+        let values = Int64Array::from(vec![100, 200, 100, 300]);
+        let run_ends = vec![2, 3, 5, 6];
+        let array: RunArray<Int16Type> =
+            RunArray::try_new(&PrimitiveArray::from(run_ends), &values).unwrap();
+        let array = array.slice(2, 3);
+
+        assert_roundtrip(&array, DataType::Int16, DataType::Int64, None);
+    }
 }
