Skip to content

Commit

Permalink
fix: Support slicing RunEndBool arrays to 0 elements (#1511)
Browse files Browse the repository at this point in the history
  • Loading branch information
robert3005 authored Nov 29, 2024
1 parent 9a78dd4 commit 6696266
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 8 deletions.
48 changes: 42 additions & 6 deletions encodings/runend-bool/src/compute/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,16 +69,52 @@ impl TakeFn<RunEndBoolArray> for RunEndBoolEncoding {

impl SliceFn<RunEndBoolArray> for RunEndBoolEncoding {
fn slice(&self, array: &RunEndBoolArray, start: usize, stop: usize) -> VortexResult<ArrayData> {
let slice_begin = array.find_physical_index(start)?;
let slice_end = array.find_physical_index(stop)?;
let new_length = stop - start;

let (slice_begin, slice_end) = if new_length == 0 {
let ends_len = array.ends().len();
(ends_len, ends_len)
} else {
let physical_begin = array.find_physical_index(start)?;
let physical_end = array.find_physical_index(stop)?;
(physical_begin, physical_end + 1)
};

Ok(RunEndBoolArray::with_offset_and_size(
slice(array.ends(), slice_begin, slice_end + 1)?,
slice(array.ends(), slice_begin, slice_end)?,
value_at_index(slice_begin, array.start()),
array.validity().slice(slice_begin, slice_end + 1)?,
stop - start,
start + array.offset(),
array.validity().slice(start, stop)?,
new_length,
if new_length == 0 {
0
} else {
start + array.offset()
},
)?
.into_array())
}
}

#[cfg(test)]
mod tests {
use vortex_array::compute::slice;
use vortex_array::validity::Validity;
use vortex_array::{ArrayLen, IntoArrayData};

use crate::RunEndBoolArray;

#[test]
fn slice_at_end() {
let re_array =
RunEndBoolArray::try_new(vec![7_u64, 10].into_array(), false, Validity::NonNullable)
.unwrap();

assert_eq!(re_array.len(), 10);

let sliced_array = slice(&re_array, re_array.len(), re_array.len()).unwrap();
assert!(sliced_array.is_empty());

let re_slice = RunEndBoolArray::try_from(sliced_array).unwrap();
assert!(re_slice.ends().is_empty());
}
}
2 changes: 1 addition & 1 deletion encodings/runend/src/array.rs
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ impl RunEndArray {
);
}

if offset != 0 && !ends.is_empty() {
if offset != 0 {
let first_run_end: usize = scalar_at(&ends, 0)?.as_ref().try_into()?;
if first_run_end <= offset {
vortex_bail!("First run end {first_run_end} must be bigger than offset {offset}");
Expand Down
6 changes: 5 additions & 1 deletion encodings/runend/src/compute/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -132,7 +132,11 @@ impl SliceFn<RunEndArray> for RunEndEncoding {
slice(array.ends(), slice_begin, slice_end)?,
slice(array.values(), slice_begin, slice_end)?,
array.validity().slice(start, stop)?,
start + array.offset(),
if new_length == 0 {
0
} else {
start + array.offset()
},
new_length,
)?
.into_array())
Expand Down

0 comments on commit 6696266

Please sign in to comment.