Skip to content

Compact run buffer slice values #10292

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Jefffrey wants to merge 4 commits into
base: main
Choose a base branch
from
Draft

Compact run buffer slice values #10292

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4b2da1a
Compact run buffer slice values
Jefffrey Jul 6, 2026
e504e68
dont need to adjust offset actually
Jefffrey Jul 6, 2026
5575583
apply slicing optimization to runarray too, fix tests
Jefffrey Jul 23, 2026
6c92882
Merge branch 'main' into compact_run_slice
Jefffrey Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
102 changes: 67 additions & 35 deletions arrow-array/src/array/run_array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,10 +345,50 @@ impl<R: RunEndIndexType> RunArray<R> {
///
/// - Specified slice (`offset` + `length`) exceeds existing length
pub fn slice(&self, offset: usize, length: usize) -> Self {
assert!(
offset.saturating_add(length) <= self.run_ends.len(),
"the length + offset of the sliced array cannot exceed the existing length"
);

if length == 0 {
return Self {
data_type: self.data_type.clone(),
run_ends: self.run_ends.slice(0, 0),
values: self.values.slice(0, 0),
};
}

// Compact the physical run ends & values to only what is strictly
// needed.

// tmp makes it easier to use the get physical index methods here
// SAFETY: new offset + slice is within bounds, and existing values are
// valid
let tmp = unsafe {
RunEndBuffer::new_unchecked(
self.run_ends.inner().clone(),
self.run_ends.offset() + offset,
length,
)
};
let start = tmp.get_start_physical_index();
let end = tmp.get_end_physical_index();

let values = self.values.slice(start, end - start + 1);
// SAFETY: existing values are valid, and values referenced by the logical
// offset are preserved
let run_ends = unsafe {
RunEndBuffer::new_unchecked(
tmp.inner().slice(start, end - start + 1),
tmp.offset(),
tmp.len(),
)
};

Self {
data_type: self.data_type.clone(),
run_ends: self.run_ends.slice(offset, length),
values: self.values.clone(),
run_ends,
values,
}
}
}
Expand Down Expand Up @@ -894,6 +934,7 @@ mod tests {
};
});
}

#[test]
fn test_run_array() {
// Construct a value array
Expand Down Expand Up @@ -1220,7 +1261,6 @@ mod tests {
PrimitiveRunBuilder::<Int16Type, Int32Type>::with_capacity(input_array.len());
builder.extend(input_array.iter().copied());
let run_array = builder.finish();
let physical_values_array = run_array.values().as_primitive::<Int32Type>();

// test for all slice lengths.
for slice_len in 1..=total_len {
Expand All @@ -1234,15 +1274,11 @@ mod tests {
// test for offset = 0 and slice length = slice_len
// slice the input array using which the run array was built.
let sliced_input_array = &input_array[0..slice_len];

// slice the run array
let sliced_run_array: RunArray<Int16Type> =
run_array.slice(0, slice_len).into_data().into();

// Get physical indices.
let sliced_run_array = run_array.slice(0, slice_len);
let physical_indices = sliced_run_array
.get_physical_indices(&logical_indices)
.unwrap();
let physical_values_array = sliced_run_array.values().as_primitive::<Int32Type>();

compare_logical_and_physical_indices(
&logical_indices,
Expand All @@ -1254,17 +1290,11 @@ mod tests {
// test for offset = total_len - slice_len and slice length = slice_len
// slice the input array using which the run array was built.
let sliced_input_array = &input_array[total_len - slice_len..total_len];

// slice the run array
let sliced_run_array: RunArray<Int16Type> = run_array
.slice(total_len - slice_len, slice_len)
.into_data()
.into();

// Get physical indices
let sliced_run_array = run_array.slice(total_len - slice_len, slice_len);
let physical_indices = sliced_run_array
.get_physical_indices(&logical_indices)
.unwrap();
let physical_values_array = sliced_run_array.values().as_primitive::<Int32Type>();

compare_logical_and_physical_indices(
&logical_indices,
Expand All @@ -1291,8 +1321,11 @@ mod tests {
let slices = [(0, 12), (0, 2), (2, 5), (3, 0), (3, 3), (3, 4), (4, 8)];
for (offset, length) in slices {
let a = array.slice(offset, length);
let n = a.logical_nulls().unwrap();
let n = n.into_iter().collect::<Vec<_>>();
let n = if let Some(nulls) = a.logical_nulls() {
nulls.into_iter().collect::<Vec<_>>()
} else {
vec![true; length]
};
assert_eq!(&n, &expected[offset..offset + length], "{offset} {length}");
}
}
Expand Down Expand Up @@ -1364,33 +1397,32 @@ mod tests {

#[test]
fn test_run_array_values_slice() {

@Jefffrey Jefffrey Jul 23, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can reference this test to see how the physical values change underneath as we slice

// 0, 0, 1, 1, 1, 2...2 (15 2s)
let run_ends: PrimitiveArray<Int32Type> = vec![2, 5, 20].into();
// 0, 0, 1, 1, 1, 2
let run_ends: PrimitiveArray<Int32Type> = vec![2, 5, 6].into();
let values: PrimitiveArray<Int32Type> = vec![0, 1, 2].into();
let array = RunArray::<Int32Type>::try_new(&run_ends, &values).unwrap();

let slice = array.slice(1, 4); // 0 | 1, 1, 1 |
let slice = array.slice(1, 4); // 0, 1, 1, 1 |
// logical indices: 1, 2, 3, 4
// physical indices: 0, 1, 1, 1
// values at 0 is 0
// values at 1 is 1
// values slice should be [0, 1]
assert_eq!(slice.get_start_physical_index(), 0);
assert_eq!(slice.get_end_physical_index(), 1);

let values_slice = slice.values_slice();
let values_slice = values_slice.as_primitive::<Int32Type>();
assert_eq!(values_slice.values(), &[0, 1]);
assert_eq!(slice.run_ends.len(), 4);
assert_eq!(slice.run_ends.offset(), 1);
assert_eq!(slice.run_ends.values(), &[2, 5]);
assert_eq!(slice.values().as_ref(), &Int32Array::from(vec![0, 1]));

let slice2 = array.slice(2, 3); // 1, 1, 1
// logical indices: 2, 3, 4
// physical indices: 1, 1, 1
assert_eq!(slice2.get_start_physical_index(), 1);
assert_eq!(slice2.get_end_physical_index(), 1);

let values_slice2 = slice2.values_slice();
let values_slice2 = values_slice2.as_primitive::<Int32Type>();
assert_eq!(values_slice2.values(), &[1]);
// physical indices: 0, 0, 0
assert_eq!(slice2.get_start_physical_index(), 0);
assert_eq!(slice2.get_end_physical_index(), 0);

assert_eq!(slice2.run_ends.len(), 3);
assert_eq!(slice2.run_ends.offset(), 2);
assert_eq!(slice2.run_ends.values(), &[5]);
assert_eq!(slice2.values().as_ref(), &Int32Array::from(vec![1]));
}

#[test]
Expand Down
72 changes: 70 additions & 2 deletions arrow-buffer/src/buffer/run.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -277,9 +277,28 @@ where
logical_offset.saturating_add(logical_length) <= self.logical_length,
"the length + offset of the sliced RunEndBuffer cannot exceed the existing length"
);
Self {
let logical_offset = self.logical_offset + logical_offset;

@Jefffrey Jefffrey Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

cc @Rich-T-kid this is what i was thinking on my previous comments on your PR

essentially we can still reduce how many physical values we hold on to, without needing to readjust all of the values (still keeping it zero-copy)

we'd need to extend it to runarray too since need to ensure the values array is similarly compacted (might need a bit of plumbing to make it work) and as noted in the issue callsites may need to be adjusted

lemme know your thoughts


if logical_length == 0 {
return Self {
run_ends: self.run_ends.slice(0, 0),
logical_length: 0,
logical_offset: 0,
};
}

// tmp makes it easier to use the get physical index methods here
let tmp = Self {
run_ends: self.run_ends.clone(),
logical_offset: self.logical_offset + logical_offset,
logical_offset,
logical_length,
};
let start = tmp.get_start_physical_index();
let end = tmp.get_end_physical_index();

Self {
run_ends: self.run_ends.slice(start, end - start + 1),
logical_offset,
logical_length,
}
}
Expand Down Expand Up @@ -421,4 +440,53 @@ mod tests {
let sliced_values: Vec<i32> = sliced.sliced_values().collect();
assert_eq!(sliced_values, &[2]);
}

#[test]
fn test_compact_slicing() {
// {} represents a logical slice within physically expanded []
// [{A, A, A, A, B, B, B, B, C, C, C, C}]
let buffer = RunEndBuffer::<i32>::new(vec![4, 8, 12].into(), 0, 12);

// zero slice
let slice = buffer.slice(2, 0);
assert_eq!(slice.len(), 0);
assert_eq!(slice.offset(), 0);
assert!(slice.values().is_empty());

// compact start only
// [B, B, B, B, B, B, {B, B, C, C, C, C}]
let slice = buffer.slice(6, 6);
assert_eq!(slice.len(), 6);
assert_eq!(slice.offset(), 6);
Comment on lines +457 to +460

@Jefffrey Jefffrey Jul 6, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

i realized my initial version of the PR incorrectly adjusted the offset when that wasnt needed; by removing the first physical value its as if the first run was consumed by the following run since we always list run lengths; that is

[4, 8] -> AAAABBBB
[8] -> {BBBB}BBBB
  • where {BBBB} used to be a separate run but is now folded into the first run; this shouldnt matter as our logical offset points beyond it anyway

assert_eq!(slice.values(), &[8, 12]);
let values = slice.sliced_values().collect::<Vec<_>>();
assert_eq!(values, &[2, 6]);

// compact end only
// [A, {A, A, A, B, B, B}, B]
let slice = buffer.slice(1, 6);
assert_eq!(slice.len(), 6);
assert_eq!(slice.offset(), 1);
assert_eq!(slice.values(), &[4, 8]);
let values = slice.sliced_values().collect::<Vec<_>>();
assert_eq!(values, &[3, 6]);

// compact both
// [B, B, B, B, B, {B, B}, B]
let slice = buffer.slice(5, 2);
assert_eq!(slice.len(), 2);
assert_eq!(slice.offset(), 5);
assert_eq!(slice.values(), &[8]);
let values = slice.sliced_values().collect::<Vec<_>>();
assert_eq!(values, &[2]);

// no compaction
// [A, {A, A, A, B, B, B, B, C, C, C}, C]
let slice = buffer.slice(1, 10);
assert_eq!(slice.len(), 10);
assert_eq!(slice.offset(), 1);
assert_eq!(slice.values(), &[4, 8, 12]);
let values = slice.sliced_values().collect::<Vec<_>>();
assert_eq!(values, &[3, 7, 10]);
}
}
Loading