This is an automated email from the ASF dual-hosted git repository.
Jefffrey pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/main by this push:
new fdeada16e9 Fix sums of sliced run-end encoded arrays (#11283)
fdeada16e9 is described below
commit fdeada16e916ad1c213b6765244222c2b558c9e1
Author: Langning Zhang <[email protected]>
AuthorDate: Tue Sep 29 17:13:45 2026 +0800
Fix sums of sliced run-end encoded arrays (#11283)
# Which issue does this PR close?
Closes #11282.
# Rationale for this change
REE slices with a nonzero offset can produce incorrect sums or incorrect
overflow results.
# What changes are included in this PR?
Compute run lengths from slice-relative endpoints, with regression
coverage for slices, nulls and overflow.
# Are these changes tested?
`cargo test --locked -p arrow-arith --all-features`, crate clippy,
workspace formatting and the affected file's spelling check pass.
# Are there any user-facing changes?
Correct sums and overflow detection for sliced REE arrays; no API
changes.
AI disclosure: Codex generated the fix, regression tests and PR text;
Claude Code assisted investigation.
---
arrow-arith/src/aggregate.rs | 112 +++++++++++++++++++++++++++++++++++++++----
1 file changed, 104 insertions(+), 8 deletions(-)
diff --git a/arrow-arith/src/aggregate.rs b/arrow-arith/src/aggregate.rs
index d059f2b287..99615f7601 100644
--- a/arrow-arith/src/aggregate.rs
+++ b/arrow-arith/src/aggregate.rs
@@ -700,10 +700,7 @@ mod ree {
where
F: FnMut(V::Native, V::Native, usize) -> Result<V::Native, E>,
{
- let run_ends = array.run_ends();
- let logical_start = run_ends.offset();
- let logical_end = run_ends.offset() + run_ends.len();
- let run_ends = run_ends.sliced_values();
+ let run_ends = array.run_ends().sliced_values();
let values_slice = array.run_array().values_slice();
let values = values_slice
@@ -717,7 +714,7 @@ mod ree {
let mut has_non_null_value = false;
for (run_end, value) in run_ends.zip(values) {
- let current_run_end = run_end.as_usize().clamp(logical_start,
logical_end);
+ let current_run_end = run_end.as_usize();
let run_length = current_run_end - prev_end;
if let Some(value) = value {
@@ -726,9 +723,6 @@ mod ree {
}
prev_end = current_run_end;
- if current_run_end == logical_end {
- break;
- }
}
Ok(if has_non_null_value { Some(acc) } else { None })
@@ -2083,6 +2077,108 @@ mod tests {
assert_eq!(result, Some(100));
}
+ #[test]
+ fn test_ree_sum_array_sliced_across_runs() {
+ let run_ends = Int16Array::from(vec![4, 8]);
+ let values = Int32Array::from(vec![10, 100]);
+ let array = RunArray::<Int16Type>::try_new(&run_ends,
&values).unwrap();
+ let sliced = array.slice(3, 4);
+ let typed_array = sliced.downcast::<Int32Array>().unwrap();
+
+ assert_eq!(sum_array::<Int32Type, _>(typed_array), Some(310));
+ }
+
+ #[test]
+ fn test_ree_sum_array_all_slices() {
+ fn check<I: RunEndIndexType>() {
+ let values = [
+ Some(2),
+ Some(2),
+ None,
+ None,
+ Some(-3),
+ Some(-3),
+ Some(-3),
+ Some(4),
+ ];
+ let run_array = make_run_array::<I, Int32Type, _>(&values);
+ let plain_array = Int32Array::from(values.to_vec());
+
+ for offset in 0..=values.len() {
+ for len in 0..=values.len() - offset {
+ let sliced = run_array.slice(offset, len);
+ let typed = sliced.downcast::<Int32Array>().unwrap();
+ let expected = plain_array.slice(offset, len);
+
+ assert_eq!(
+ sum_array::<Int32Type, _>(typed),
+ sum_array::<Int32Type, _>(&expected),
+ "offset={offset}, len={len}"
+ );
+ assert_eq!(
+ sum_array_checked::<Int32Type, _>(typed).unwrap(),
+ sum_array_checked::<Int32Type, _>(&expected).unwrap(),
+ "offset={offset}, len={len}"
+ );
+ }
+ }
+
+ let nested = run_array.slice(1, 7).slice(2, 4);
+ let typed = nested.downcast::<Int32Array>().unwrap();
+ let expected = plain_array.slice(1, 7).slice(2, 4);
+ assert_eq!(
+ sum_array::<Int32Type, _>(typed),
+ sum_array::<Int32Type, _>(&expected)
+ );
+ assert_eq!(
+ sum_array_checked::<Int32Type, _>(typed).unwrap(),
+ sum_array_checked::<Int32Type, _>(&expected).unwrap()
+ );
+ }
+
+ check::<Int16Type>();
+ check::<Int32Type>();
+ check::<Int64Type>();
+ }
+
+ #[test]
+ fn test_ree_sum_array_checked_sliced_overflow() {
+ let values = [50, 50, 50, 50, 1, 1, 1, 1];
+ let run_array = make_run_array::<Int16Type, Int8Type, _>(&values);
+ let sliced = run_array.slice(3, 4);
+ let typed = sliced.downcast::<Int8Array>().unwrap();
+ assert_eq!(sum_array_checked::<Int8Type, _>(typed).unwrap(), Some(53));
+
+ let values = [1, 1, 1, 1, 50, 50, 50, 50];
+ let run_array = make_run_array::<Int16Type, Int8Type, _>(&values);
+ let sliced = run_array.slice(3, 4);
+ let typed = sliced.downcast::<Int8Array>().unwrap();
+ assert!(sum_array_checked::<Int8Type, _>(typed).is_err());
+ }
+
+ #[test]
+ fn test_ree_sum_array_sliced_infinity() {
+ let values = [
+ 1.0,
+ 1.0,
+ 1.0,
+ 1.0,
+ 2.0,
+ 2.0,
+ 2.0,
+ 2.0,
+ f64::INFINITY,
+ f64::INFINITY,
+ f64::INFINITY,
+ f64::INFINITY,
+ ];
+ let run_array = make_run_array::<Int16Type, Float64Type, _>(&values);
+ let sliced = run_array.slice(5, 4);
+ let typed = sliced.downcast::<Float64Array>().unwrap();
+
+ assert_eq!(sum_array::<Float64Type, _>(typed), Some(f64::INFINITY));
+ }
+
#[test]
fn test_ree_min_max_array_basic() {
let run_array = make_run_array::<Int16Type, Int32Type, _>(&[30, 30,
10, 20, 20]);
