diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 2e85adab10b..bffa254835d 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -2650,15 +2650,40 @@ fn build_chunk_index( // and an invalid count must not reach the final-chunk subtraction. let value_counts = analyze_value_counts(words, items_in_page)?; - // Each chunk stores `(divided_bytes + 1) * MINIBLOCK_ALIGNMENT` bytes, so the - // deltas are the chunk sizes and their grand total is the data buffer size. - let byte_starts = PrefixSums::from_deltas( + // Each chunk stores `(divided_bytes + 1) * MINIBLOCK_ALIGNMENT` bytes. Validate + // that the metadata covers exactly the recorded data buffer before using its + // total to select the prefix-sum width. + let chunk_sizes = || { words .iter() - .map(|word| ((word >> 4) as u64 + 1) * MINIBLOCK_ALIGNMENT as u64), - num_chunks, - data_buf_size, - ); + .map(|word| ((word >> 4) as u64 + 1) * MINIBLOCK_ALIGNMENT as u64) + }; + let chunk_data_size = chunk_sizes().enumerate().try_fold( + 0u64, + |accumulated_size, (chunk_index, chunk_size)| { + accumulated_size.checked_add(chunk_size).ok_or_else(|| { + Error::corrupt_file_named( + "miniblock_metadata", + format!( + "chunk byte sizes overflow at chunk_index={chunk_index}: \ + accumulated_size={accumulated_size}, chunk_size={chunk_size}, \ + num_chunks={num_chunks}" + ), + ) + }) + }, + )?; + if chunk_data_size != data_buf_size { + return Err(Error::corrupt_file_named( + "miniblock_metadata", + format!( + "chunk byte sizes do not match the data buffer: \ + chunk_data_size={chunk_data_size}, data_buf_size={data_buf_size}, \ + num_chunks={num_chunks}" + ), + )); + } + let byte_starts = PrefixSums::from_deltas(chunk_sizes(), num_chunks, chunk_data_size); // Nested pages track rows via the repetition index and keep leaf item counts // separately; flat pages have row == value index, so value counts are rows. @@ -7385,6 +7410,7 @@ mod tests { }; use arrow_buffer::{BooleanBuffer, NullBuffer, ScalarBuffer}; use arrow_schema::{DataType, Field as ArrowField}; + use lance_core::Error; use std::collections::HashMap; use std::{collections::VecDeque, sync::Arc}; @@ -8715,6 +8741,24 @@ mod tests { assert_eq!(total_items, items_in_page); } + #[rstest] + #[case::metadata_longer_than_data(8)] + #[case::metadata_shorter_than_data(24)] + fn test_rejects_mismatched_miniblock_data_size(#[case] data_buf_size: u64) { + let (words, _) = words_from(&[(1, 8), (0, 8)]); + + let error = build_chunk_index(&words, 3, 0, data_buf_size, None, 0).unwrap_err(); + + assert!(matches!(error, Error::CorruptFile { .. })); + let message = error.to_string(); + assert!(message.contains("chunk_data_size=16"), "{message}"); + assert!( + message.contains(&format!("data_buf_size={data_buf_size}")), + "{message}" + ); + assert!(message.contains("num_chunks=2"), "{message}"); + } + #[test] fn test_nested_detection_and_axes() { // Repetition index (stride 2): three chunks holding 5, 4, 3 rows, no trailers.