@@ -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) ]
13821409mod 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