Skip to content

Commit 6c7eb33

Browse files
authored
fix[vortex-array]: fix offset_within_chunk underflow on patches array (#7312)
slice_start_idx was missing normalization to absolute coordinates. Signed-off-by: Alfonso Subiotto Marques <alfonso.subiotto@polarsignals.com>
1 parent 951bdf9 commit 6c7eb33

File tree

1 file changed

+22
-6
lines changed

1 file changed

+22
-6
lines changed

vortex-array/src/patches.rs

Lines changed: 22 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -656,7 +656,7 @@ impl Patches {
656656
let values = self.values().slice(slice_start_idx..slice_end_idx)?;
657657
let indices = self.indices().slice(slice_start_idx..slice_end_idx)?;
658658

659-
let chunk_offsets = self
659+
let new_chunk_offsets = self
660660
.chunk_offsets
661661
.as_ref()
662662
.map(|chunk_offsets| -> VortexResult<ArrayRef> {
@@ -667,15 +667,17 @@ impl Patches {
667667
})
668668
.transpose()?;
669669

670-
let offset_within_chunk = chunk_offsets
670+
let offset_within_chunk = new_chunk_offsets
671671
.as_ref()
672-
.map(|chunk_offsets| -> VortexResult<usize> {
673-
let base_offset = chunk_offsets
672+
.map(|new_chunk_offsets| -> VortexResult<usize> {
673+
let new_chunk_base = new_chunk_offsets
674674
.scalar_at(0)?
675675
.as_primitive()
676676
.as_::<usize>()
677677
.ok_or_else(|| vortex_err!("chunk offset does not fit in usize"))?;
678-
Ok(slice_start_idx - base_offset)
678+
let parent_chunk_base = self.chunk_offset_at(0)?;
679+
let parent_within = self.offset_within_chunk.unwrap_or(0);
680+
Ok(parent_chunk_base + parent_within + slice_start_idx - new_chunk_base)
679681
})
680682
.transpose()?;
681683

@@ -684,7 +686,7 @@ impl Patches {
684686
offset: range.start + self.offset(),
685687
indices,
686688
values,
687-
chunk_offsets,
689+
chunk_offsets: new_chunk_offsets,
688690
offset_within_chunk,
689691
}))
690692
}
@@ -2151,6 +2153,20 @@ mod test {
21512153
);
21522154
}
21532155

2156+
#[test]
2157+
fn test_nested_slice_with_dropped_first_chunk() {
2158+
// PATCH_CHUNK_SIZE = 1024, so the two patches land in different chunks.
2159+
let indices = buffer![0u64, 1024].into_array();
2160+
let values = buffer![1i32, 2].into_array();
2161+
let chunk_offsets = buffer![0u64, 1].into_array();
2162+
let patches = Patches::new(2048, 0, indices, values, Some(chunk_offsets)).unwrap();
2163+
2164+
// Drop chunk 0, then re-slice the result.
2165+
let dropped_first = patches.slice(1024..2048).unwrap().unwrap();
2166+
let resliced = dropped_first.slice(0..1024).unwrap().unwrap();
2167+
assert_eq!(resliced.num_patches(), 1);
2168+
}
2169+
21542170
#[test]
21552171
fn test_index_larger_than_length() {
21562172
let chunk_offsets = buffer![0u64].into_array();

0 commit comments

Comments
 (0)