Skip to content

Commit 1a797f5

Browse files
committed
fix: validate deserialized patch chunk offsets
Check restored offsets against the sliced row count and host-resident first-chunk patch indices. Cover malformed offsets and valid short single-chunk slices. Remove the unused ALP import and apply the CI-pinned formatter. Signed-off-by: blaginin <[email protected]>
1 parent 50eeca5 commit 1a797f5

3 files changed

Lines changed: 98 additions & 6 deletions

File tree

‎encodings/alp/src/alp/mod.rs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,6 @@ mod tests {
3636
use vortex_array::arrays::PrimitiveArray;
3737
use vortex_array::arrays::slice::SliceKernel;
3838
use vortex_array::assert_arrays_eq;
39-
use vortex_array::builtins::ArrayBuiltins;
4039
use vortex_array::dtype::PType;
4140
use vortex_array::patches::PatchesMetadata;
4241
use vortex_array::test_harness::check_metadata;
@@ -71,7 +70,9 @@ mod tests {
7170
}
7271

7372
#[rstest]
74-
fn serde_roundtrip_sliced_chunked_patches(#[values(500, 1500)] start: usize) -> VortexResult<()> {
73+
fn serde_roundtrip_sliced_chunked_patches(
74+
#[values(500, 1500)] start: usize,
75+
) -> VortexResult<()> {
7576
let session = VortexSession::empty();
7677
crate::initialize(&session);
7778
let mut ctx = session.create_execution_ctx();

‎encodings/fastlanes/src/bitpacking/plugin/tests.rs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -490,7 +490,9 @@ fn serde_roundtrip_sliced_chunked_patches(
490490
expected_patches.offset_within_chunk()
491491
);
492492
let resliced = read_patches.slice(800..read.len(), &mut ctx)?.unwrap();
493-
let expected_resliced = expected_patches.slice(800..sliced.len(), &mut ctx)?.unwrap();
493+
let expected_resliced = expected_patches
494+
.slice(800..sliced.len(), &mut ctx)?
495+
.unwrap();
494496
assert_eq!(
495497
resliced.offset_within_chunk(),
496498
expected_resliced.offset_within_chunk()

‎vortex-array/src/patches.rs‎

Lines changed: 92 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -310,8 +310,35 @@ impl Patches {
310310
chunk_offsets,
311311
ctx,
312312
)?;
313-
if patches.chunk_offsets.is_some() {
314-
patches.offset_within_chunk = Some(usize::try_from(metadata.offset_within_chunk())?);
313+
if let Some(chunk_offsets) = &patches.chunk_offsets {
314+
let offset_within_chunk = usize::try_from(metadata.offset_within_chunk())?;
315+
let sliced_rows = patches.offset % PATCH_CHUNK_SIZE;
316+
vortex_ensure!(
317+
offset_within_chunk <= sliced_rows,
318+
"Patch offset within chunk exceeds the sliced row count"
319+
);
320+
321+
// A wrong offset can assign first-chunk patches to a later chunk during decoding.
322+
// Check the boundary against the indices, not just the original chunk's patch count.
323+
if chunk_offsets.len() > 1 && chunk_offsets.is_host() && patches.indices.is_host() {
324+
let first_chunk_end = patches
325+
.offset
326+
.checked_add(PATCH_CHUNK_SIZE - sliced_rows)
327+
.ok_or_else(|| vortex_err!("Patch chunk boundary overflows usize"))?;
328+
let retained =
329+
search_index_binary_search(&patches.indices, first_chunk_end, ctx)?.to_index();
330+
let expected_offset = patches
331+
.chunk_offset_at(1, ctx)?
332+
.checked_sub(patches.chunk_offset_at(0, ctx)?)
333+
.and_then(|count| count.checked_sub(retained))
334+
.ok_or_else(|| vortex_err!("Invalid patch chunk offsets"))?;
335+
vortex_ensure_eq!(
336+
offset_within_chunk,
337+
expected_offset,
338+
"Patch offset within chunk disagrees with the patch indices"
339+
);
340+
}
341+
patches.offset_within_chunk = Some(offset_within_chunk);
315342
}
316343
Ok(patches)
317344
}
@@ -1380,6 +1407,7 @@ fn take_indices_with_search_fn<
13801407

13811408
#[cfg(test)]
13821409
mod test {
1410+
use std::ops::Range;
13831411
use std::sync::LazyLock;
13841412

13851413
use rstest::rstest;
@@ -1427,7 +1455,68 @@ mod test {
14271455
)?;
14281456

14291457
assert_eq!(patches.offset_within_chunk(), chunked.then_some(0));
1430-
assert_eq!(patches.search_index(1029, &mut ctx)?, SearchResult::Found(1));
1458+
assert_eq!(
1459+
patches.search_index(1029, &mut ctx)?,
1460+
SearchResult::Found(1)
1461+
);
1462+
Ok(())
1463+
}
1464+
1465+
#[rstest]
1466+
fn test_from_metadata_rejects_inconsistent_chunk_offset(
1467+
#[values(0, 2, 3, 501, usize::MAX)] offset_within_chunk: usize,
1468+
) {
1469+
let mut ctx = SESSION.create_execution_ctx();
1470+
// The original first chunk had patches at 100 and 500; slicing at 500 removed one.
1471+
let metadata = PatchesMetadata::new(
1472+
2,
1473+
500,
1474+
PType::U32,
1475+
Some(2),
1476+
Some(PType::U32),
1477+
Some(offset_within_chunk),
1478+
);
1479+
let error = Patches::from_metadata(
1480+
1000,
1481+
&metadata,
1482+
buffer![500u32, 1200].into_array(),
1483+
buffer![11i32, 22].into_array(),
1484+
Some(buffer![0u32, 2].into_array()),
1485+
&mut ctx,
1486+
)
1487+
.expect_err("inconsistent patch offset must be rejected");
1488+
assert!(error.to_string().contains("Patch offset within chunk"));
1489+
}
1490+
1491+
#[rstest]
1492+
#[case::single_chunk(900..901)]
1493+
#[case::later_single_chunk(1500..1501)]
1494+
#[case::multiple_chunks(900..2050)]
1495+
fn test_from_metadata_short_slice(#[case] range: Range<usize>) -> VortexResult<()> {
1496+
let mut ctx = SESSION.create_execution_ctx();
1497+
let indices = PrimitiveArray::from_iter(0u32..3072).into_array();
1498+
let patches = Patches::new(
1499+
3072,
1500+
0,
1501+
indices.clone(),
1502+
indices,
1503+
Some(buffer![0u32, 1024, 2048].into_array()),
1504+
&mut ctx,
1505+
)?;
1506+
let sliced = patches.slice(range, &mut ctx)?.expect("slice has patches");
1507+
let metadata = sliced.to_metadata(sliced.array_len(), sliced.dtype())?;
1508+
let read = Patches::from_metadata(
1509+
sliced.array_len(),
1510+
&metadata,
1511+
sliced.indices().clone(),
1512+
sliced.values().clone(),
1513+
sliced.chunk_offsets().clone(),
1514+
&mut ctx,
1515+
)?;
1516+
1517+
assert_eq!(read.offset_within_chunk(), sliced.offset_within_chunk());
1518+
assert_eq!(read.search_index(0, &mut ctx)?, SearchResult::Found(0));
1519+
assert_arrays_eq!(read.values(), sliced.values(), &mut ctx);
14311520
Ok(())
14321521
}
14331522

0 commit comments

Comments
 (0)