From 71add656719ddde7c0d23812021df8873bd639e2 Mon Sep 17 00:00:00 2001 From: m00dy Date: Sun, 2 Aug 2026 11:37:56 -0500 Subject: [PATCH 1/4] fix(encoding): bounds-check the length prefix in VariableFullZipDecoder::unzip (#8138) `parse_length` reads the length prefix out of the page buffer with `get_unchecked`, and the only thing between it and the end of the buffer is a `debug_assert!`: ```rust // Safety: Data should have at least bytes_per_length bytes remaining debug_assert!(databuf.len() >= bytes_per_length); let length = unsafe { Self::parse_length(databuf, in_bits_per_length) }; ``` There is no `[profile.release]` override in the workspace `Cargo.toml`, so `debug-assertions` defaults to false in release and that assertion is not present in the published wheels. The loop it sits in continues on `while !databuf.is_empty()`, so it enters the body with as little as one byte remaining. `parse_length` then reads up to eight. A page whose item walk ends with a partial trailing item therefore reads past the end of the buffer. Reproduced on x86-64 with `-Zsanitizer=address` on a release build, driving the real `VariableFullZipDecoder::new`: ``` ERROR: AddressSanitizer: heap-buffer-overflow READ of size 8 at 0x7b9989be1017 #0 ::new 0x7b9989be1017 is located 3 bytes after 4-byte region [0x...1010,0x...1014) ``` A well-formed control buffer is clean in the same run. The payload read one line below the call site is already bounds checked and panics on malformed input: ```rust unzipped_data.extend_from_slice(&databuf[..length as usize]); ``` So a truncated item already fails cleanly on the payload path. Only the length read was inconsistent. This makes the two match by using safe indexing in `parse_length`, which lets the `unsafe` block and the `debug_assert!` both go away. On valid input the behaviour is unchanged. On a truncated trailing item the result is the same clean panic the payload path already produces, rather than an out-of-bounds read. Two, per the contributing guide: - `variable_full_zip_wellformed_length_prefix` decodes a well-formed prefix - `variable_full_zip_truncated_length_prefix_is_rejected` is `#[should_panic]` and covers the case above Both pass, and the crate's existing suite is unaffected (520 passing before and after). I have not established that a `.lance` file produced by the writer can reach this state. Truncating a data file is rejected earlier by the I/O range check, and a sweep of in-place single-byte edits either read cleanly, were rejected by that same check, or panicked in safe code further along in decode. So I am not claiming this is reachable from a crafted dataset, and I am filing it as hardening rather than as a security report. The case for the change does not depend on that: an `unsafe` read whose only guard is compiled out of release builds is worth removing on its own, particularly when the adjacent read of the same buffer is already checked. Not part of this change, and not something I have shown to be a bug, but it looked odd while reading. The length is read using `in_bits_per_length` and the cursor is then advanced by `bytes_per_offset`, which comes from `out_bits_per_offset`: ```rust let length = ... parse_length(databuf, in_bits_per_length); databuf = &databuf[bytes_per_offset..]; ``` Those are equal in the common case, so this may well be deliberate. Flagging it only in case the asymmetry is unintentional. --------- Co-authored-by: Xuanwo --- .../src/encodings/logical/primitive.rs | 653 +++++++++++++++++- 1 file changed, 624 insertions(+), 29 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 9f6be591d13..4f0d4c2f11c 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -2222,7 +2222,7 @@ impl FullZipScheduler { num_rows, bits_per_offset, bits_per_offset, - ))) + )?)) } } } @@ -2629,7 +2629,7 @@ impl VariableFullZipDecoder { num_rows: u64, in_bits_per_length: u8, out_bits_per_offset: u8, - ) -> Self { + ) -> Result { let decompressor = match details.value_decompressor { PerValueDecompressor::Variable(ref d) => d.clone(), _ => unreachable!(), @@ -2674,9 +2674,9 @@ impl VariableFullZipDecoder { // - We could force each decode task to do a full unzip of all the data. Each decode task now // has to do more work but the work is all fused. // - We could just try doing this work on the decode thread and see if it is a problem. - decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows); + decoder.unzip(data, in_bits_per_length, out_bits_per_offset, num_rows)?; - decoder + Ok(decoder) } fn slice_batch_data_and_rebase_offsets_typed( @@ -2751,28 +2751,33 @@ impl VariableFullZipDecoder { } } - unsafe fn parse_length(data: &[u8], bits_per_offset: u8) -> u64 { - match bits_per_offset { - 8 => *data.get_unchecked(0) as u64, - 16 => u16::from_le_bytes([*data.get_unchecked(0), *data.get_unchecked(1)]) as u64, - 32 => u32::from_le_bytes([ - *data.get_unchecked(0), - *data.get_unchecked(1), - *data.get_unchecked(2), - *data.get_unchecked(3), - ]) as u64, - 64 => u64::from_le_bytes([ - *data.get_unchecked(0), - *data.get_unchecked(1), - *data.get_unchecked(2), - *data.get_unchecked(3), - *data.get_unchecked(4), - *data.get_unchecked(5), - *data.get_unchecked(6), - *data.get_unchecked(7), - ]), - _ => unreachable!(), + /// Reads a single length prefix from the front of `data`. + /// + /// The bytes come from the file. A page whose item walk ends with a partial + /// trailing item leaves fewer than `bits_per_offset / 8` bytes here, so this + /// is bounds checked and reports a corrupt file rather than reading past the + /// end of the buffer. + fn parse_length(data: &[u8], bits_per_offset: u8) -> Result { + let width = bits_per_offset as usize / 8; + if data.len() < width { + return Err(Error::corrupt_file_named( + "variable_full_zip", + format!( + "truncated length prefix: {} byte(s) remain in the page buffer but a \ + {}-bit length prefix requires {}", + data.len(), + bits_per_offset, + width + ), + )); } + Ok(match bits_per_offset { + 8 => data[0] as u64, + 16 => u16::from_le_bytes(data[..2].try_into().unwrap()) as u64, + 32 => u32::from_le_bytes(data[..4].try_into().unwrap()) as u64, + 64 => u64::from_le_bytes(data[..8].try_into().unwrap()), + _ => unreachable!(), + }) } fn unzip( @@ -2781,7 +2786,7 @@ impl VariableFullZipDecoder { in_bits_per_length: u8, out_bits_per_offset: u8, num_rows: u64, - ) { + ) -> Result<()> { // This undercounts if there are lists but, at this point, we don't really know how many items we have let mut rep = Vec::with_capacity(num_rows as usize); let mut def = Vec::with_capacity(num_rows as usize); @@ -2832,9 +2837,7 @@ impl VariableFullZipDecoder { if ctrl_desc.is_visible { visible_item_count += 1; if ctrl_desc.is_valid_item { - // Safety: Data should have at least bytes_per_length bytes remaining - debug_assert!(databuf.len() >= bytes_per_length); - let length = unsafe { Self::parse_length(databuf, in_bits_per_length) }; + let length = Self::parse_length(databuf, in_bits_per_length)?; match out_bits_per_offset { 32 => offsets_data .extend_from_slice(&(current_offset as u32).to_le_bytes()), @@ -2870,6 +2873,7 @@ impl VariableFullZipDecoder { self.def = ScalarBuffer::from(def); self.data = LanceBuffer::from(unzipped_data); self.offsets = LanceBuffer::from(offsets_data); + Ok(()) } } @@ -7414,4 +7418,595 @@ mod tests { check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) .await; } + #[tokio::test] + async fn test_complex_all_null_constant_def_round_trip() { + use arrow_array::ListArray; + + // Every row is a null list => constant def levels => a single RLE run, + // exercising the lazy run-form decode end to end. + let list_array = ListArray::from_iter_primitive::( + (0..5000).map(|_| None::>>), + ); + + let test_cases = TestCases::default().with_u32_structural_encodings(); + check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) + .await; + } + + fn encoded_u16_frame(levels: &[u16], run_length_width: RunLengthWidth) -> LanceBuffer { + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_slice(Arc::from(levels)), + bits_per_value: 16, + num_values: levels.len() as u64, + block_info: BlockInfo::new(), + }); + BlockCompressor::compress(&RleEncoder::with_run_length_width(run_length_width), block) + .unwrap() + } + + fn encoded_u16_runs(levels: &[u16], run_length_width: RunLengthWidth) -> RleRuns { + let frame = encoded_u16_frame(levels, run_length_width); + RleDecompressor::with_run_length_width(16, run_length_width) + .decode_u16_runs(frame, levels.len() as u64) + .unwrap() + } + + fn physical_levels(levels: &[u16]) -> LazyLevels { + LazyLevels::Runs(Arc::new(RunStorage::Physical( + encoded_u16_runs(levels, RunLengthWidth::U8).into_owned(), + ))) + } + + fn coalesced_levels(levels: &[u16]) -> LazyLevels { + let mut values = Vec::new(); + let mut ends = RunEndsBuilder::with_capacity(levels.len(), levels.len()); + for (index, &value) in levels.iter().enumerate() { + if values.last() == Some(&value) { + ends.set_last(index + 1).unwrap(); + } else { + values.push(value); + ends.push(index + 1).unwrap(); + } + } + LazyLevels::Runs(Arc::new(RunStorage::Coalesced { + values: values.into_boxed_slice(), + ends: ends.finish(), + })) + } + + #[test] + fn lazy_levels_runs_match_dense() { + // Runs: 3x2, 1x1, 3x3, 0x2 => [3,3,1,3,3,3,0,0] + let expanded: Vec = vec![3, 3, 1, 3, 3, 3, 0, 0]; + let coalesced = coalesced_levels(&expanded); + let physical = physical_levels(&expanded); + let dense = LazyLevels::Dense(ScalarBuffer::::from(expanded.clone())); + let n = expanded.len(); + + assert_eq!(coalesced.len(), n); + assert_eq!(physical.len(), n); + assert_eq!(dense.len(), n); + + // Rows begin at each `max_rep` (3) position; row `num_rows` maps to `len`. + let max_rep = 3u16; + let row_starts: Vec = (0..n).filter(|&i| expanded[i] == max_rep).collect(); + for target in 0..=row_starts.len() as u64 { + let want = row_starts.get(target as usize).copied().unwrap_or(n); + for runs in [&coalesced, &physical] { + let mut cursor = LevelCursor::default(); + assert_eq!( + runs.seek_row_start(&mut cursor, target, max_rep).unwrap(), + want, + "seek_row_start({target})" + ); + } + let mut c_dense = LevelCursor::default(); + assert_eq!( + dense.seek_row_start(&mut c_dense, target, max_rep).unwrap(), + want + ); + } + + // `count_le_cursor` (fresh cursor per range) and `extend_into` agree with + // the dense reference on every sub-range. + for start in 0..=n { + for end in start..=n { + for max in [0u16, 1, 2, 3] { + let want = expanded[start..end].iter().filter(|&&d| d <= max).count() as u64; + for runs in [&coalesced, &physical] { + let mut cursor = RunPosition::default(); + assert_eq!( + runs.count_le_cursor(&mut cursor, start..end, max).0, + want, + "count_le_cursor({start}..{end}, {max})" + ); + } + let mut d_cur = RunPosition::default(); + assert_eq!(dense.count_le_cursor(&mut d_cur, start..end, max).0, want); + } + for runs in [&coalesced, &physical] { + let mut got = Vec::new(); + runs.extend_into(start..end, RunPosition::default(), &mut got); + assert_eq!( + got, + expanded[start..end].to_vec(), + "extend_into({start}..{end})" + ); + } + let mut got_dense = Vec::new(); + dense.extend_into(start..end, RunPosition::default(), &mut got_dense); + assert_eq!(got_dense, expanded[start..end].to_vec()); + } + } + } + + #[test] + fn physical_run_hints_support_deferred_materialization() { + let expanded: Vec = vec![3, 3, 1, 1, 2, 2, 0, 0]; + let physical = physical_levels(&expanded); + let LazyLevels::Runs(runs) = &physical else { + panic!("expected physical runs"); + }; + let mut first_hint = RunPosition::default(); + runs.seek(&mut first_hint, 2); + let mut second_hint = RunPosition::default(); + runs.seek(&mut second_hint, 6); + + let mut second = Vec::new(); + physical.extend_into(6..8, second_hint, &mut second); + let mut first = Vec::new(); + physical.extend_into(2..4, first_hint, &mut first); + assert_eq!(second, expanded[6..8]); + assert_eq!(first, expanded[2..4]); + } + + /// Fuzz parity for the run-oriented complex-all-null drain: the cursor walk + /// over `LazyLevels` must yield the exact level slices and visible + /// count that a brute-force reference over the fully expanded levels does, for + /// dense, physical-run, and coalesced-run forms and arbitrarily shaped range requests. + mod complex_all_null_drain_parity { + use std::ops::Range; + + use arrow_buffer::ScalarBuffer; + use proptest::prelude::*; + + use super::super::{LazyLevels, LevelCursor, RunPosition}; + use super::{coalesced_levels, physical_levels}; + use crate::Result; + + #[derive(Debug, Clone)] + struct DrainInput { + max_rep: u16, + max_visible: u16, + rep: Option>, + def: Option>, + ranges: Vec>, + } + + fn dense_levels(levels: &[u16]) -> LazyLevels { + LazyLevels::Dense(ScalarBuffer::from(levels.to_vec())) + } + + fn rle_levels(levels: &[u16]) -> LazyLevels { + coalesced_levels(levels) + } + + fn seek( + rep: Option<&LazyLevels>, + cursor: &mut LevelCursor, + row: u64, + max_rep: u16, + ) -> Result { + match rep { + Some(rep) => rep.seek_row_start(cursor, row, max_rep), + None => { + cursor.row = row; + cursor.level = row as usize; + Ok(row as usize) + } + } + } + + /// Mirror of `ComplexAllNullPageDecoder::drain`, driving the real + /// `seek_row_start` / `count_le_cursor` with monotonic cursors. + fn simulate_drain( + rep: Option<&LazyLevels>, + def: Option<&LazyLevels>, + max_rep: u16, + max_visible: u16, + ranges: &[Range], + ) -> Result<(Vec>, u64)> { + let mut rep_cursor = LevelCursor::default(); + let mut def_run_cursor = RunPosition::default(); + let mut slices: Vec> = Vec::new(); + let mut visible = 0u64; + for range in ranges { + let level_start = seek(rep, &mut rep_cursor, range.start, max_rep)?; + let level_end = seek(rep, &mut rep_cursor, range.end, max_rep)?; + visible += match def { + Some(def) => { + def.count_le_cursor( + &mut def_run_cursor, + level_start..level_end, + max_visible, + ) + .0 + } + None => (level_end - level_start) as u64, + }; + match slices.last_mut() { + Some(last) if last.end == level_start => last.end = level_end, + _ => slices.push(level_start..level_end), + } + } + Ok((slices, visible)) + } + + /// Independent brute-force reference over fully expanded levels. + fn reference_drain( + rep: Option<&[u16]>, + def: Option<&[u16]>, + max_rep: u16, + max_visible: u16, + ranges: &[Range], + ) -> (Vec>, u64) { + let total_levels = rep + .map(|r| r.len()) + .or_else(|| def.map(|d| d.len())) + .unwrap_or(0); + // Level index where each row starts (or `total_levels` for the end row). + let row_starts: Vec = match rep { + Some(rep) => (0..rep.len()).filter(|&i| rep[i] == max_rep).collect(), + None => (0..total_levels).collect(), + }; + let level_of_row = |row: u64| { + row_starts + .get(row as usize) + .copied() + .unwrap_or(total_levels) + }; + + let mut slices: Vec> = Vec::new(); + let mut visible = 0u64; + for range in ranges { + let ls = level_of_row(range.start); + let le = level_of_row(range.end); + visible += match def { + Some(def) => def[ls..le].iter().filter(|&&d| d <= max_visible).count() as u64, + None => (le - ls) as u64, + }; + match slices.last_mut() { + Some(last) if last.end == ls => last.end = le, + _ => slices.push(ls..le), + } + } + (slices, visible) + } + + fn ranges_strategy(num_rows: u64) -> BoxedStrategy>> { + if num_rows == 0 { + return Just(Vec::new()).boxed(); + } + // (gap, len) pairs; a zero gap yields ranges adjacent in row space, + // which exercises the level-slice coalescing path. + proptest::collection::vec((0u64..=3, 1u64..=4), 0..=8) + .prop_map(move |pairs| { + let mut ranges = Vec::new(); + let mut pos = 0u64; + for (gap, len) in pairs { + pos = pos.saturating_add(gap); + if pos >= num_rows { + break; + } + let end = (pos + len).min(num_rows); + ranges.push(pos..end); + pos = end; + } + ranges + }) + .boxed() + } + + fn drain_input() -> impl Strategy { + ( + 1u16..=3, + 0u16..=3, + any::(), + any::(), + 1usize..=48, + ) + .prop_flat_map(|(max_rep, max_visible, has_rep, has_def, len)| { + // Complex-all-null always has definition levels when there is + // no repetition, so force `def` present in that case. + let has_def = has_def || !has_rep; + let rep = if has_rep { + proptest::collection::vec(0u16..=max_rep, len) + .prop_map(move |mut v| { + // Row 0 must start at a max-rep boundary. + v[0] = max_rep; + Some(v) + }) + .boxed() + } else { + Just(None).boxed() + }; + let def = if has_def { + proptest::collection::vec(0u16..=(max_visible + 2), len) + .prop_map(Some) + .boxed() + } else { + Just(None).boxed() + }; + (Just(max_rep), Just(max_visible), rep, def) + }) + .prop_flat_map(|(max_rep, max_visible, rep, def)| { + let num_rows = match &rep { + Some(rep) => rep.iter().filter(|&&v| v == max_rep).count() as u64, + None => def.as_ref().map(|d| d.len() as u64).unwrap_or(0), + }; + ranges_strategy(num_rows).prop_map(move |ranges| DrainInput { + max_rep, + max_visible, + rep: rep.clone(), + def: def.clone(), + ranges, + }) + }) + } + + proptest! { + #![proptest_config(ProptestConfig::with_cases(512))] + + #[test] + fn drain_matches_reference(input in drain_input()) { + let DrainInput { max_rep, max_visible, rep, def, ranges } = input; + + let reference = + reference_drain(rep.as_deref(), def.as_deref(), max_rep, max_visible, &ranges); + + let rep_dense = rep.as_deref().map(dense_levels); + let def_dense = def.as_deref().map(dense_levels); + let got_dense = + simulate_drain(rep_dense.as_ref(), def_dense.as_ref(), max_rep, max_visible, &ranges) + .unwrap(); + prop_assert_eq!(&got_dense, &reference, "dense form diverged from reference"); + + let rep_rle = rep.as_deref().map(rle_levels); + let def_rle = def.as_deref().map(rle_levels); + let got_rle = + simulate_drain(rep_rle.as_ref(), def_rle.as_ref(), max_rep, max_visible, &ranges) + .unwrap(); + prop_assert_eq!(&got_rle, &reference, "rle form diverged from reference"); + + let rep_physical = rep.as_deref().map(physical_levels); + let def_physical = def.as_deref().map(physical_levels); + let got_physical = + simulate_drain(rep_physical.as_ref(), def_physical.as_ref(), max_rep, max_visible, &ranges) + .unwrap(); + prop_assert_eq!(&got_physical, &reference, "physical form diverged from reference"); + } + } + } + + #[test] + fn lazy_levels_runs_are_compact() { + let single_run = |n: usize| { + let mut ends = RunEndsBuilder::with_capacity(n, 1); + ends.push(n).unwrap(); + LazyLevels::Runs(Arc::new(RunStorage::Coalesced { + values: vec![1u16].into_boxed_slice(), + ends: ends.finish(), + })) + }; + // Run-form footprint is independent of the logical length within an end width... + assert_eq!(single_run(100).deep_size(), single_run(10_000).deep_size()); + assert!(single_run(10_000_000).deep_size() < 100); + assert_eq!(single_run(10_000_000).len(), 10_000_000); + // ...while Dense pays 2 bytes per value. + assert_eq!( + LazyLevels::Dense(ScalarBuffer::::from(vec![1u16; 1000])).deep_size(), + 2000 + ); + } + + #[test] + fn lazy_levels_selects_smallest_representation() { + let runs = encoded_u16_runs(&[7u16; 10], RunLengthWidth::U8); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense); + + let equal_size: Vec = std::iter::repeat_n(0, 256) + .chain(std::iter::repeat_n(1, 100)) + .chain(std::iter::repeat_n(2, 100)) + .collect(); + let runs = encoded_u16_runs(&equal_size, RunLengthWidth::U8); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced); + + let moderate_runs: Vec = (0..250) + .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4)) + .collect(); + let runs = encoded_u16_runs(&moderate_runs, RunLengthWidth::U8); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical); + + let split_constant = vec![7u16; 5000]; + let runs = encoded_u16_runs(&split_constant, RunLengthWidth::U8); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced); + + let high_density: Vec = (0..70_000).map(|index| (index % 2) as u16).collect(); + let runs = encoded_u16_runs(&high_density, RunLengthWidth::U32); + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense); + } + + #[test] + fn physical_runs_detach_from_large_encoded_frame() { + let levels: Vec = (0..250) + .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4)) + .collect(); + let frame = encoded_u16_frame(&levels, RunLengthWidth::U8); + let frame_offset = 4096; + let mut allocation = vec![0; frame_offset + frame.len() + 1_000_000]; + allocation[frame_offset..frame_offset + frame.len()].copy_from_slice(frame.as_ref()); + let frame = LanceBuffer::from(allocation).slice_with_length(frame_offset, frame.len()); + let runs = RleDecompressor::with_run_length_width(16, RunLengthWidth::U8) + .decode_u16_runs(frame, levels.len() as u64) + .unwrap(); + + assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical); + let cached = LazyLevels::from_rle_runs(runs).unwrap(); + assert!( + matches!(cached, LazyLevels::Runs(ref runs) if matches!(runs.as_ref(), RunStorage::Physical(_))) + ); + assert_eq!(cached.len(), levels.len()); + assert!(cached.deep_size() < 4096); + } + + #[test] + fn complex_all_null_levels_reject_invalid_values_and_lengths() { + let invalid_levels = vec![0u16, 3]; + for levels in [ + LazyLevels::Dense(ScalarBuffer::from(invalid_levels.clone())), + physical_levels(&invalid_levels), + coalesced_levels(&invalid_levels), + ] { + let error = validate_complex_all_null_levels(&None, &Some(levels), 0, 2).unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!(error.to_string().contains("Invalid definition level 3")); + } + + let rep = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16; 2]))); + let def = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16]))); + let error = validate_complex_all_null_levels(&rep, &def, 0, 0).unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("repetition has 2, definition has 1") + ); + } + + #[test] + fn block_levels_reject_malformed_payload_size() { + let block = DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::from(vec![0]), + bits_per_value: 16, + num_values: 1, + block_info: BlockInfo::new(), + }); + let error = dense_levels_from_block(block, 1, "definition").unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("expected 2 bytes for 1 values, got 1") + ); + } + + #[test] + fn complex_all_null_level_codec_validates_rle_metadata() { + let encoding = pb21::CompressiveEncoding { + compression: Some(Compression::Rle(Box::new(pb21::Rle { + values: None, + run_lengths: Some(Box::new(ProtobufUtils21::flat(8, None))), + }))), + }; + + let error = LevelCodec::try_new(Some(&encoding), &DefaultDecompressionStrategy::default()) + .unwrap_err(); + assert!(matches!(error, lance_core::Error::InvalidInput { .. })); + assert!( + error + .to_string() + .contains("RLE compression missing values encoding") + ); + } + + // https://github.com/lance-format/lance/issues/6681 + #[tokio::test] + async fn test_sparse_boolean_list_roundtrip() { + use arrow_array::builder::{BooleanBuilder, ListBuilder}; + + let mut list_builder = ListBuilder::new(BooleanBuilder::new()); + for i in 0..1000i32 { + if i % 64 == 0 { + // Alternate true/false so the array is not constant (constant path avoids the bug). + list_builder.values().append_value(i % 128 == 0); + list_builder.append(true); + } else { + list_builder.append(false); + } + } + let list_array = Arc::new(list_builder.finish()); + + let test_cases = TestCases::default().with_structural_encodings(); + check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; + } + + fn truncated_tail_details() -> std::sync::Arc { + use crate::compression::VariablePerValueDecompressor; + use crate::encodings::physical::binary::VariableDecoder; + use crate::repdef::{ControlWordParser, DefinitionInterpretation}; + use std::sync::Arc; + Arc::new(super::FullZipDecodeDetails { + value_decompressor: super::PerValueDecompressor::Variable(Arc::new( + VariableDecoder::default(), + ) + as Arc), + def_meaning: vec![DefinitionInterpretation::NullableItem].into(), + ctrl_word_parser: ControlWordParser::new(0, 0), + max_rep: 0, + max_visible_def: 0, + }) + } + + fn decode_variable_full_zip( + buf: Vec, + bits_per_offset: u8, + ) -> lance_core::Result { + use std::collections::VecDeque; + let mut data = VecDeque::new(); + data.push_back(crate::buffer::LanceBuffer::from(buf)); + super::VariableFullZipDecoder::new( + truncated_tail_details(), + data, + 1, + bits_per_offset, + bits_per_offset, + ) + } + + /// A well-formed length prefix decodes without incident, for both widths. + #[test] + fn variable_full_zip_wellformed_length_prefix() { + assert!(decode_variable_full_zip(0u32.to_le_bytes().to_vec(), 32).is_ok()); + assert!(decode_variable_full_zip(0u64.to_le_bytes().to_vec(), 64).is_ok()); + } + + /// A page whose item walk ends with a partial length prefix must surface a + /// corrupt-file error rather than read past the end of the buffer. + /// + /// This asserts the error variant and message rather than merely expecting a + /// panic: before the length prefix was bounds checked, the read was + /// `get_unchecked` behind a `debug_assert!`, so a debug build panicked here + /// (which a `#[should_panic]` test would have accepted as a pass) while a + /// release build read up to 8 bytes out of a 4 byte allocation. + #[test] + fn variable_full_zip_truncated_length_prefix_is_corrupt_file() { + use lance_core::Error; + + for (bits, buf_len) in [(32u8, 3usize), (64u8, 4usize)] { + let err = decode_variable_full_zip(vec![0xAA; buf_len], bits) + .expect_err("a truncated length prefix must not decode"); + assert!( + matches!(err, Error::CorruptFile { .. }), + "expected CorruptFile for a {}-bit prefix with {} byte(s), got: {:?}", + bits, + buf_len, + err + ); + let msg = err.to_string(); + assert!( + msg.contains("truncated length prefix"), + "error should say what is wrong, got: {msg}" + ); + } + } } From 982747d8bde43375645d592bbd5d739954a91e3b Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 04:29:05 +0800 Subject: [PATCH 2/4] chore: trim unrelated v6.1 backport context --- .../src/encodings/logical/primitive.rs | 555 ------------------ 1 file changed, 555 deletions(-) diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 4f0d4c2f11c..9dee00445a8 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -7418,561 +7418,6 @@ mod tests { check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) .await; } - #[tokio::test] - async fn test_complex_all_null_constant_def_round_trip() { - use arrow_array::ListArray; - - // Every row is a null list => constant def levels => a single RLE run, - // exercising the lazy run-form decode end to end. - let list_array = ListArray::from_iter_primitive::( - (0..5000).map(|_| None::>>), - ); - - let test_cases = TestCases::default().with_u32_structural_encodings(); - check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) - .await; - } - - fn encoded_u16_frame(levels: &[u16], run_length_width: RunLengthWidth) -> LanceBuffer { - let block = DataBlock::FixedWidth(FixedWidthDataBlock { - data: LanceBuffer::reinterpret_slice(Arc::from(levels)), - bits_per_value: 16, - num_values: levels.len() as u64, - block_info: BlockInfo::new(), - }); - BlockCompressor::compress(&RleEncoder::with_run_length_width(run_length_width), block) - .unwrap() - } - - fn encoded_u16_runs(levels: &[u16], run_length_width: RunLengthWidth) -> RleRuns { - let frame = encoded_u16_frame(levels, run_length_width); - RleDecompressor::with_run_length_width(16, run_length_width) - .decode_u16_runs(frame, levels.len() as u64) - .unwrap() - } - - fn physical_levels(levels: &[u16]) -> LazyLevels { - LazyLevels::Runs(Arc::new(RunStorage::Physical( - encoded_u16_runs(levels, RunLengthWidth::U8).into_owned(), - ))) - } - - fn coalesced_levels(levels: &[u16]) -> LazyLevels { - let mut values = Vec::new(); - let mut ends = RunEndsBuilder::with_capacity(levels.len(), levels.len()); - for (index, &value) in levels.iter().enumerate() { - if values.last() == Some(&value) { - ends.set_last(index + 1).unwrap(); - } else { - values.push(value); - ends.push(index + 1).unwrap(); - } - } - LazyLevels::Runs(Arc::new(RunStorage::Coalesced { - values: values.into_boxed_slice(), - ends: ends.finish(), - })) - } - - #[test] - fn lazy_levels_runs_match_dense() { - // Runs: 3x2, 1x1, 3x3, 0x2 => [3,3,1,3,3,3,0,0] - let expanded: Vec = vec![3, 3, 1, 3, 3, 3, 0, 0]; - let coalesced = coalesced_levels(&expanded); - let physical = physical_levels(&expanded); - let dense = LazyLevels::Dense(ScalarBuffer::::from(expanded.clone())); - let n = expanded.len(); - - assert_eq!(coalesced.len(), n); - assert_eq!(physical.len(), n); - assert_eq!(dense.len(), n); - - // Rows begin at each `max_rep` (3) position; row `num_rows` maps to `len`. - let max_rep = 3u16; - let row_starts: Vec = (0..n).filter(|&i| expanded[i] == max_rep).collect(); - for target in 0..=row_starts.len() as u64 { - let want = row_starts.get(target as usize).copied().unwrap_or(n); - for runs in [&coalesced, &physical] { - let mut cursor = LevelCursor::default(); - assert_eq!( - runs.seek_row_start(&mut cursor, target, max_rep).unwrap(), - want, - "seek_row_start({target})" - ); - } - let mut c_dense = LevelCursor::default(); - assert_eq!( - dense.seek_row_start(&mut c_dense, target, max_rep).unwrap(), - want - ); - } - - // `count_le_cursor` (fresh cursor per range) and `extend_into` agree with - // the dense reference on every sub-range. - for start in 0..=n { - for end in start..=n { - for max in [0u16, 1, 2, 3] { - let want = expanded[start..end].iter().filter(|&&d| d <= max).count() as u64; - for runs in [&coalesced, &physical] { - let mut cursor = RunPosition::default(); - assert_eq!( - runs.count_le_cursor(&mut cursor, start..end, max).0, - want, - "count_le_cursor({start}..{end}, {max})" - ); - } - let mut d_cur = RunPosition::default(); - assert_eq!(dense.count_le_cursor(&mut d_cur, start..end, max).0, want); - } - for runs in [&coalesced, &physical] { - let mut got = Vec::new(); - runs.extend_into(start..end, RunPosition::default(), &mut got); - assert_eq!( - got, - expanded[start..end].to_vec(), - "extend_into({start}..{end})" - ); - } - let mut got_dense = Vec::new(); - dense.extend_into(start..end, RunPosition::default(), &mut got_dense); - assert_eq!(got_dense, expanded[start..end].to_vec()); - } - } - } - - #[test] - fn physical_run_hints_support_deferred_materialization() { - let expanded: Vec = vec![3, 3, 1, 1, 2, 2, 0, 0]; - let physical = physical_levels(&expanded); - let LazyLevels::Runs(runs) = &physical else { - panic!("expected physical runs"); - }; - let mut first_hint = RunPosition::default(); - runs.seek(&mut first_hint, 2); - let mut second_hint = RunPosition::default(); - runs.seek(&mut second_hint, 6); - - let mut second = Vec::new(); - physical.extend_into(6..8, second_hint, &mut second); - let mut first = Vec::new(); - physical.extend_into(2..4, first_hint, &mut first); - assert_eq!(second, expanded[6..8]); - assert_eq!(first, expanded[2..4]); - } - - /// Fuzz parity for the run-oriented complex-all-null drain: the cursor walk - /// over `LazyLevels` must yield the exact level slices and visible - /// count that a brute-force reference over the fully expanded levels does, for - /// dense, physical-run, and coalesced-run forms and arbitrarily shaped range requests. - mod complex_all_null_drain_parity { - use std::ops::Range; - - use arrow_buffer::ScalarBuffer; - use proptest::prelude::*; - - use super::super::{LazyLevels, LevelCursor, RunPosition}; - use super::{coalesced_levels, physical_levels}; - use crate::Result; - - #[derive(Debug, Clone)] - struct DrainInput { - max_rep: u16, - max_visible: u16, - rep: Option>, - def: Option>, - ranges: Vec>, - } - - fn dense_levels(levels: &[u16]) -> LazyLevels { - LazyLevels::Dense(ScalarBuffer::from(levels.to_vec())) - } - - fn rle_levels(levels: &[u16]) -> LazyLevels { - coalesced_levels(levels) - } - - fn seek( - rep: Option<&LazyLevels>, - cursor: &mut LevelCursor, - row: u64, - max_rep: u16, - ) -> Result { - match rep { - Some(rep) => rep.seek_row_start(cursor, row, max_rep), - None => { - cursor.row = row; - cursor.level = row as usize; - Ok(row as usize) - } - } - } - - /// Mirror of `ComplexAllNullPageDecoder::drain`, driving the real - /// `seek_row_start` / `count_le_cursor` with monotonic cursors. - fn simulate_drain( - rep: Option<&LazyLevels>, - def: Option<&LazyLevels>, - max_rep: u16, - max_visible: u16, - ranges: &[Range], - ) -> Result<(Vec>, u64)> { - let mut rep_cursor = LevelCursor::default(); - let mut def_run_cursor = RunPosition::default(); - let mut slices: Vec> = Vec::new(); - let mut visible = 0u64; - for range in ranges { - let level_start = seek(rep, &mut rep_cursor, range.start, max_rep)?; - let level_end = seek(rep, &mut rep_cursor, range.end, max_rep)?; - visible += match def { - Some(def) => { - def.count_le_cursor( - &mut def_run_cursor, - level_start..level_end, - max_visible, - ) - .0 - } - None => (level_end - level_start) as u64, - }; - match slices.last_mut() { - Some(last) if last.end == level_start => last.end = level_end, - _ => slices.push(level_start..level_end), - } - } - Ok((slices, visible)) - } - - /// Independent brute-force reference over fully expanded levels. - fn reference_drain( - rep: Option<&[u16]>, - def: Option<&[u16]>, - max_rep: u16, - max_visible: u16, - ranges: &[Range], - ) -> (Vec>, u64) { - let total_levels = rep - .map(|r| r.len()) - .or_else(|| def.map(|d| d.len())) - .unwrap_or(0); - // Level index where each row starts (or `total_levels` for the end row). - let row_starts: Vec = match rep { - Some(rep) => (0..rep.len()).filter(|&i| rep[i] == max_rep).collect(), - None => (0..total_levels).collect(), - }; - let level_of_row = |row: u64| { - row_starts - .get(row as usize) - .copied() - .unwrap_or(total_levels) - }; - - let mut slices: Vec> = Vec::new(); - let mut visible = 0u64; - for range in ranges { - let ls = level_of_row(range.start); - let le = level_of_row(range.end); - visible += match def { - Some(def) => def[ls..le].iter().filter(|&&d| d <= max_visible).count() as u64, - None => (le - ls) as u64, - }; - match slices.last_mut() { - Some(last) if last.end == ls => last.end = le, - _ => slices.push(ls..le), - } - } - (slices, visible) - } - - fn ranges_strategy(num_rows: u64) -> BoxedStrategy>> { - if num_rows == 0 { - return Just(Vec::new()).boxed(); - } - // (gap, len) pairs; a zero gap yields ranges adjacent in row space, - // which exercises the level-slice coalescing path. - proptest::collection::vec((0u64..=3, 1u64..=4), 0..=8) - .prop_map(move |pairs| { - let mut ranges = Vec::new(); - let mut pos = 0u64; - for (gap, len) in pairs { - pos = pos.saturating_add(gap); - if pos >= num_rows { - break; - } - let end = (pos + len).min(num_rows); - ranges.push(pos..end); - pos = end; - } - ranges - }) - .boxed() - } - - fn drain_input() -> impl Strategy { - ( - 1u16..=3, - 0u16..=3, - any::(), - any::(), - 1usize..=48, - ) - .prop_flat_map(|(max_rep, max_visible, has_rep, has_def, len)| { - // Complex-all-null always has definition levels when there is - // no repetition, so force `def` present in that case. - let has_def = has_def || !has_rep; - let rep = if has_rep { - proptest::collection::vec(0u16..=max_rep, len) - .prop_map(move |mut v| { - // Row 0 must start at a max-rep boundary. - v[0] = max_rep; - Some(v) - }) - .boxed() - } else { - Just(None).boxed() - }; - let def = if has_def { - proptest::collection::vec(0u16..=(max_visible + 2), len) - .prop_map(Some) - .boxed() - } else { - Just(None).boxed() - }; - (Just(max_rep), Just(max_visible), rep, def) - }) - .prop_flat_map(|(max_rep, max_visible, rep, def)| { - let num_rows = match &rep { - Some(rep) => rep.iter().filter(|&&v| v == max_rep).count() as u64, - None => def.as_ref().map(|d| d.len() as u64).unwrap_or(0), - }; - ranges_strategy(num_rows).prop_map(move |ranges| DrainInput { - max_rep, - max_visible, - rep: rep.clone(), - def: def.clone(), - ranges, - }) - }) - } - - proptest! { - #![proptest_config(ProptestConfig::with_cases(512))] - - #[test] - fn drain_matches_reference(input in drain_input()) { - let DrainInput { max_rep, max_visible, rep, def, ranges } = input; - - let reference = - reference_drain(rep.as_deref(), def.as_deref(), max_rep, max_visible, &ranges); - - let rep_dense = rep.as_deref().map(dense_levels); - let def_dense = def.as_deref().map(dense_levels); - let got_dense = - simulate_drain(rep_dense.as_ref(), def_dense.as_ref(), max_rep, max_visible, &ranges) - .unwrap(); - prop_assert_eq!(&got_dense, &reference, "dense form diverged from reference"); - - let rep_rle = rep.as_deref().map(rle_levels); - let def_rle = def.as_deref().map(rle_levels); - let got_rle = - simulate_drain(rep_rle.as_ref(), def_rle.as_ref(), max_rep, max_visible, &ranges) - .unwrap(); - prop_assert_eq!(&got_rle, &reference, "rle form diverged from reference"); - - let rep_physical = rep.as_deref().map(physical_levels); - let def_physical = def.as_deref().map(physical_levels); - let got_physical = - simulate_drain(rep_physical.as_ref(), def_physical.as_ref(), max_rep, max_visible, &ranges) - .unwrap(); - prop_assert_eq!(&got_physical, &reference, "physical form diverged from reference"); - } - } - } - - #[test] - fn lazy_levels_runs_are_compact() { - let single_run = |n: usize| { - let mut ends = RunEndsBuilder::with_capacity(n, 1); - ends.push(n).unwrap(); - LazyLevels::Runs(Arc::new(RunStorage::Coalesced { - values: vec![1u16].into_boxed_slice(), - ends: ends.finish(), - })) - }; - // Run-form footprint is independent of the logical length within an end width... - assert_eq!(single_run(100).deep_size(), single_run(10_000).deep_size()); - assert!(single_run(10_000_000).deep_size() < 100); - assert_eq!(single_run(10_000_000).len(), 10_000_000); - // ...while Dense pays 2 bytes per value. - assert_eq!( - LazyLevels::Dense(ScalarBuffer::::from(vec![1u16; 1000])).deep_size(), - 2000 - ); - } - - #[test] - fn lazy_levels_selects_smallest_representation() { - let runs = encoded_u16_runs(&[7u16; 10], RunLengthWidth::U8); - assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense); - - let equal_size: Vec = std::iter::repeat_n(0, 256) - .chain(std::iter::repeat_n(1, 100)) - .chain(std::iter::repeat_n(2, 100)) - .collect(); - let runs = encoded_u16_runs(&equal_size, RunLengthWidth::U8); - assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced); - - let moderate_runs: Vec = (0..250) - .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4)) - .collect(); - let runs = encoded_u16_runs(&moderate_runs, RunLengthWidth::U8); - assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical); - - let split_constant = vec![7u16; 5000]; - let runs = encoded_u16_runs(&split_constant, RunLengthWidth::U8); - assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Coalesced); - - let high_density: Vec = (0..70_000).map(|index| (index % 2) as u16).collect(); - let runs = encoded_u16_runs(&high_density, RunLengthWidth::U32); - assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Dense); - } - - #[test] - fn physical_runs_detach_from_large_encoded_frame() { - let levels: Vec = (0..250) - .flat_map(|run| std::iter::repeat_n((run % 2) as u16, 4)) - .collect(); - let frame = encoded_u16_frame(&levels, RunLengthWidth::U8); - let frame_offset = 4096; - let mut allocation = vec![0; frame_offset + frame.len() + 1_000_000]; - allocation[frame_offset..frame_offset + frame.len()].copy_from_slice(frame.as_ref()); - let frame = LanceBuffer::from(allocation).slice_with_length(frame_offset, frame.len()); - let runs = RleDecompressor::with_run_length_width(16, RunLengthWidth::U8) - .decode_u16_runs(frame, levels.len() as u64) - .unwrap(); - - assert_eq!(LazyLevels::select_plan(&runs), LevelPlan::Physical); - let cached = LazyLevels::from_rle_runs(runs).unwrap(); - assert!( - matches!(cached, LazyLevels::Runs(ref runs) if matches!(runs.as_ref(), RunStorage::Physical(_))) - ); - assert_eq!(cached.len(), levels.len()); - assert!(cached.deep_size() < 4096); - } - - #[test] - fn complex_all_null_levels_reject_invalid_values_and_lengths() { - let invalid_levels = vec![0u16, 3]; - for levels in [ - LazyLevels::Dense(ScalarBuffer::from(invalid_levels.clone())), - physical_levels(&invalid_levels), - coalesced_levels(&invalid_levels), - ] { - let error = validate_complex_all_null_levels(&None, &Some(levels), 0, 2).unwrap_err(); - assert!(matches!(error, lance_core::Error::InvalidInput { .. })); - assert!(error.to_string().contains("Invalid definition level 3")); - } - - let rep = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16; 2]))); - let def = Some(LazyLevels::Dense(ScalarBuffer::from(vec![0u16]))); - let error = validate_complex_all_null_levels(&rep, &def, 0, 0).unwrap_err(); - assert!(matches!(error, lance_core::Error::InvalidInput { .. })); - assert!( - error - .to_string() - .contains("repetition has 2, definition has 1") - ); - } - - #[test] - fn block_levels_reject_malformed_payload_size() { - let block = DataBlock::FixedWidth(FixedWidthDataBlock { - data: LanceBuffer::from(vec![0]), - bits_per_value: 16, - num_values: 1, - block_info: BlockInfo::new(), - }); - let error = dense_levels_from_block(block, 1, "definition").unwrap_err(); - assert!(matches!(error, lance_core::Error::InvalidInput { .. })); - assert!( - error - .to_string() - .contains("expected 2 bytes for 1 values, got 1") - ); - } - - #[test] - fn complex_all_null_level_codec_validates_rle_metadata() { - let encoding = pb21::CompressiveEncoding { - compression: Some(Compression::Rle(Box::new(pb21::Rle { - values: None, - run_lengths: Some(Box::new(ProtobufUtils21::flat(8, None))), - }))), - }; - - let error = LevelCodec::try_new(Some(&encoding), &DefaultDecompressionStrategy::default()) - .unwrap_err(); - assert!(matches!(error, lance_core::Error::InvalidInput { .. })); - assert!( - error - .to_string() - .contains("RLE compression missing values encoding") - ); - } - - // https://github.com/lance-format/lance/issues/6681 - #[tokio::test] - async fn test_sparse_boolean_list_roundtrip() { - use arrow_array::builder::{BooleanBuilder, ListBuilder}; - - let mut list_builder = ListBuilder::new(BooleanBuilder::new()); - for i in 0..1000i32 { - if i % 64 == 0 { - // Alternate true/false so the array is not constant (constant path avoids the bug). - list_builder.values().append_value(i % 128 == 0); - list_builder.append(true); - } else { - list_builder.append(false); - } - } - let list_array = Arc::new(list_builder.finish()); - - let test_cases = TestCases::default().with_structural_encodings(); - check_round_trip_encoding_of_data(vec![list_array], &test_cases, HashMap::new()).await; - } - - fn truncated_tail_details() -> std::sync::Arc { - use crate::compression::VariablePerValueDecompressor; - use crate::encodings::physical::binary::VariableDecoder; - use crate::repdef::{ControlWordParser, DefinitionInterpretation}; - use std::sync::Arc; - Arc::new(super::FullZipDecodeDetails { - value_decompressor: super::PerValueDecompressor::Variable(Arc::new( - VariableDecoder::default(), - ) - as Arc), - def_meaning: vec![DefinitionInterpretation::NullableItem].into(), - ctrl_word_parser: ControlWordParser::new(0, 0), - max_rep: 0, - max_visible_def: 0, - }) - } - - fn decode_variable_full_zip( - buf: Vec, - bits_per_offset: u8, - ) -> lance_core::Result { - use std::collections::VecDeque; - let mut data = VecDeque::new(); - data.push_back(crate::buffer::LanceBuffer::from(buf)); - super::VariableFullZipDecoder::new( - truncated_tail_details(), - data, - 1, - bits_per_offset, - bits_per_offset, - ) - } - /// A well-formed length prefix decodes without incident, for both widths. #[test] fn variable_full_zip_wellformed_length_prefix() { From 8954303454bb4cd194402c418859e8ca5612f486 Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 03:56:41 +0800 Subject: [PATCH 3/4] fix(encoding): reject out-of-bounds variable-width offsets on decode (#8144) --- rust/lance-encoding/src/data.rs | 439 +++++++++++++++++- .../src/encodings/physical/binary.rs | 403 +++++++++++++++- rust/lance-file/src/reader.rs | 172 +++++++ 3 files changed, 981 insertions(+), 33 deletions(-) diff --git a/rust/lance-encoding/src/data.rs b/rust/lance-encoding/src/data.rs index fa4ffe3021e..588e953f232 100644 --- a/rust/lance-encoding/src/data.rs +++ b/rust/lance-encoding/src/data.rs @@ -578,20 +578,249 @@ pub struct VariableWidthBlock { pub block_info: BlockInfo, } +/// Proof that a [`VariableWidthBlock`] satisfies the Arrow layout contract for +/// its target data type (offsets buffer long enough, offsets monotonic and +/// within the data buffer, values valid UTF-8 where required). +/// +/// Only [`VariableWidthBlock::validate_layout`] can construct it, which ties the +/// unchecked Arrow build below to an actual validation pass instead of a +/// caller-controlled flag. +struct ValidVariableWidthLayout; + impl VariableWidthBlock { - fn into_arrow(self, data_type: DataType, validate: bool) -> Result { + // The offsets buffer comes straight from file bytes, so an unchecked build would + // let a corrupt file smuggle out-of-bounds offsets into an Arrow array whose + // consumers then read (or crash on) memory outside the data buffer. This + // boundary therefore always validates the layout, ignoring the optional + // `validate` flag. Lance validates the common layouts itself (a branchless + // scan, measurably cheaper than Arrow's element-wise checked build) and only + // falls back to Arrow's checked build for the cold cases. + fn into_arrow(self, data_type: DataType, _validate: bool) -> Result { + let Some(expected_bits_per_offset) = Self::expected_bits_per_offset(&data_type) else { + // Not an [offsets, bytes] layout we know how to prove; let Arrow + // check it. + return self.into_arrow_checked(data_type); + }; + if self.bits_per_offset != expected_bits_per_offset { + return Err(self.layout_error( + &data_type, + format!( + "expected {}-bit offsets but got {}-bit offsets", + expected_bits_per_offset, self.bits_per_offset + ), + )); + } + if self.num_values == 0 { + // Cold path; Arrow handles the empty-offsets special cases. + return self.into_arrow_checked(data_type); + } + let proof = self.validate_layout(&data_type)?; + Ok(self.into_arrow_unchecked(data_type, proof)) + } + + /// The offset width Arrow mandates for `data_type`, or `None` if the type + /// does not use the `[offsets, bytes]` layout this block represents. + fn expected_bits_per_offset(data_type: &DataType) -> Option { + match data_type { + DataType::Binary | DataType::Utf8 => Some(32), + DataType::LargeBinary | DataType::LargeUtf8 => Some(64), + _ => None, + } + } + + fn layout_error(&self, data_type: &DataType, detail: impl std::fmt::Display) -> Error { + Self::format_layout_error( + data_type, + detail, + self.num_values, + self.bits_per_offset, + self.offsets.len(), + self.data.len(), + ) + } + + fn format_layout_error( + data_type: &DataType, + detail: impl std::fmt::Display, + num_values: u64, + bits_per_offset: u8, + offsets_size: usize, + data_size: usize, + ) -> Error { + Error::corrupt_file_named( + "variable width data block", + format!( + "invalid variable-width layout for {}: {} (num_values: {}, bits_per_offset: {}, \ + offsets buffer size: {} bytes, data buffer size: {} bytes)", + data_type, detail, num_values, bits_per_offset, offsets_size, data_size, + ), + ) + } + + fn validate_layout(&self, data_type: &DataType) -> Result { + let bytes_per_offset = (self.bits_per_offset / 8) as u64; + let required_bytes = self + .num_values + .checked_add(1) + .and_then(|num_offsets| num_offsets.checked_mul(bytes_per_offset)) + .ok_or_else(|| self.layout_error(data_type, "offsets buffer size overflows"))?; + if (self.offsets.len() as u64) < required_bytes { + return Err(self.layout_error( + data_type, + format!( + "offsets buffer must hold at least {} offsets ({} bytes)", + self.num_values + 1, + required_bytes + ), + )); + } + let validate_utf8 = matches!(data_type, DataType::Utf8 | DataType::LargeUtf8); + match self.bits_per_offset { + 32 => self.validate_offsets_and_values::(data_type, validate_utf8), + 64 => self.validate_offsets_and_values::(data_type, validate_utf8), + other => Err(self.layout_error( + data_type, + format!("unsupported offset width: {} bits", other), + )), + } + } + + fn validate_offsets_and_values( + &self, + data_type: &DataType, + validate_utf8: bool, + ) -> Result { + let num_offsets = self.num_values as usize + 1; + // Slice before borrowing: the buffer may carry padding that is not a + // multiple of the offset width. + let offsets = self + .offsets + .slice_with_length(0, num_offsets * std::mem::size_of::()); + let offsets = offsets.borrow_to_typed_slice::(); + let offsets: &[T] = offsets.as_ref(); + let data = self.data.as_ref(); + + // A monotonic sequence with a non-negative first offset and an + // in-bounds last offset is entirely within [0, data.len()], so the hot + // loop only proves monotonicity; everything else is O(1) at the ends. + // The `&=` accumulation keeps the loop branchless so it vectorizes. + let mut is_monotonic = true; + for window in offsets.windows(2) { + is_monotonic &= window[0] <= window[1]; + } + let first = offsets[0]; + let last = offsets[num_offsets - 1]; + let bounds_ok = + first >= T::usize_as(0) && last.to_usize().is_some_and(|last| last <= data.len()); + if !is_monotonic || !bounds_ok { + return Err(self.offset_violation_error::(data_type, offsets)); + } + + if validate_utf8 { + let (first, last) = (first.as_usize(), last.as_usize()); + let values = std::str::from_utf8(&data[first..last]) + .map_err(|utf8_err| self.layout_error(data_type, utf8_err))?; + let mut on_char_boundaries = true; + for &offset in offsets { + on_char_boundaries &= values.is_char_boundary(offset.as_usize() - first); + } + if !on_char_boundaries { + // Cold path: rescan to pinpoint the offending offset. + let position = offsets + .iter() + .position(|offset| !values.is_char_boundary(offset.as_usize() - first)) + .expect("the fast scan found a non-boundary offset"); + return Err(self.layout_error( + data_type, + format!("offset at position {position} splits a UTF-8 character"), + )); + } + } + + Ok(ValidVariableWidthLayout) + } + + /// Cold path: pinpoint the first offending offset for the error message. + fn offset_violation_error( + &self, + data_type: &DataType, + offsets: &[T], + ) -> Error { + let data_size = self.data.len(); + for (position, window) in offsets.windows(2).enumerate() { + if window[0] > window[1] { + return self.layout_error( + data_type, + format!( + "non-monotonic offset at position {}: {:?} > {:?}", + position, window[0], window[1] + ), + ); + } + } + for (position, offset) in offsets.iter().enumerate() { + match offset.to_usize() { + None => { + return self.layout_error( + data_type, + format!("negative offset at position {}: {:?}", position, offset), + ); + } + Some(offset) if offset > data_size => { + return self.layout_error( + data_type, + format!( + "offset at position {} out of bounds: {} > {}", + position, offset, data_size + ), + ); + } + Some(_) => {} + } + } + // The fast scan only fails when one of the loops above finds the + // culprit; reaching here would be a bug in the fast scan itself. + self.layout_error(data_type, "offsets failed validation") + } + + fn into_arrow_checked(self, data_type: DataType) -> Result { + let num_values = self.num_values; + let bits_per_offset = self.bits_per_offset; + let offsets_size = self.offsets.len(); + let data_size = self.data.len(); + let builder = self.into_arrow_builder(data_type.clone()); + builder.build().map_err(|arrow_err| { + Self::format_layout_error( + &data_type, + arrow_err, + num_values, + bits_per_offset, + offsets_size, + data_size, + ) + }) + } + + fn into_arrow_unchecked( + self, + data_type: DataType, + _proof: ValidVariableWidthLayout, + ) -> ArrayData { + let builder = self.into_arrow_builder(data_type); + // SAFETY: `_proof` witnesses that `validate_layout` proved this block + // satisfies the Arrow layout contract for `data_type`. + unsafe { builder.build_unchecked() } + } + + fn into_arrow_builder(self, data_type: DataType) -> ArrayDataBuilder { + let num_values = self.num_values; let data_buffer = self.data.into_buffer(); let offsets_buffer = self.offsets.into_buffer(); - let builder = ArrayDataBuilder::new(data_type) + ArrayDataBuilder::new(data_type) .add_buffer(offsets_buffer) .add_buffer(data_buffer) - .len(self.num_values as usize) - .null_count(0); - if validate { - Ok(builder.build()?) - } else { - Ok(unsafe { builder.build_unchecked() }) - } + .len(num_values as usize) + .null_count(0) } fn into_buffers(self) -> Vec { @@ -1624,19 +1853,25 @@ mod tests { use std::sync::Arc; use arrow_array::{ - ArrayRef, DictionaryArray, Int8Array, LargeBinaryArray, StringArray, UInt8Array, - UInt16Array, make_array, new_null_array, + ArrayRef, BinaryArray, BinaryViewArray, DictionaryArray, Int8Array, LargeBinaryArray, + LargeStringArray, StringArray, StringViewArray, UInt8Array, UInt16Array, make_array, + new_null_array, types::{Int8Type, Int32Type}, }; use arrow_buffer::{BooleanBuffer, NullBuffer}; use arrow_schema::{DataType, Field, Fields}; + use lance_core::Error; use lance_datagen::{ArrayGeneratorExt, DEFAULT_SEED, RowCount, array}; use rand::SeedableRng; + use rstest::rstest; use crate::buffer::LanceBuffer; - use super::{AllNullDataBlock, DataBlock}; + use super::{ + AllNullDataBlock, BlockInfo, DataBlock, DictionaryDataBlock, FixedWidthDataBlock, + VariableWidthBlock, + }; use arrow_array::Array; @@ -1999,4 +2234,184 @@ mod tests { let total_nulls_size_in_bytes = concatenated_array.nulls().unwrap().len().div_ceil(8); assert!(block.data_size() == (total_buffer_size + total_nulls_size_in_bytes) as u64); } + + #[test] + fn variable_width_rejects_out_of_bounds_offsets_without_optional_validation() { + let block = VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets: LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + bits_per_offset: 32, + num_values: 3, + block_info: BlockInfo::new(), + }; + + let error = block + .into_arrow(DataType::Binary, false) + .expect_err("out-of-bounds offsets must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + let message = error.to_string(); + assert!( + message.contains("100000") && message.contains("data buffer size: 14 bytes"), + "error must report the offending offset and the data buffer size: {message}" + ); + } + + #[rstest] + #[case::binary_i32_tail_out_of_bounds( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::utf8_i32_tail_out_of_bounds( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::large_binary_i64_tail_out_of_bounds( + DataType::LargeBinary, + LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 100_000]), + 64, + 3, + b"alphabetagamma".as_slice() + )] + #[case::large_utf8_i64_tail_out_of_bounds( + DataType::LargeUtf8, + LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 100_000]), + 64, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_negative_offset( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, -1, 9, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_non_monotonic_offsets( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 9, 5, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_interior_offset_out_of_bounds( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 100_000, 100_000, 14]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::binary_offsets_buffer_too_short( + DataType::Binary, + LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9]), + 32, + 3, + b"alphabetagamma".as_slice() + )] + #[case::utf8_invalid_byte_sequence( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2, 3]), + 32, + 3, + &[b'a', 0xFF, b'b'] + )] + #[case::utf8_offset_splits_multibyte_char( + DataType::Utf8, + LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2]), + 32, + 2, + "é".as_bytes() + )] + #[case::large_utf8_invalid_byte_sequence( + DataType::LargeUtf8, + LanceBuffer::reinterpret_vec(vec![0_i64, 1, 2, 3]), + 64, + 3, + &[b'a', 0xFF, b'b'] + )] + fn variable_width_rejects_malformed_layout( + #[case] data_type: DataType, + #[case] offsets: LanceBuffer, + #[case] bits_per_offset: u8, + #[case] num_values: u64, + #[case] data: &[u8], + ) { + let block = VariableWidthBlock { + data: LanceBuffer::copy_slice(data), + offsets, + bits_per_offset, + num_values, + block_info: BlockInfo::new(), + }; + + // The malformed layout must be rejected regardless of the optional + // `validate` flag: the flag selects extra validation, not the memory + // safety proof required to construct an Arrow array. + for validate in [false, true] { + let error = DataBlock::VariableWidth(block.clone()) + .into_arrow(data_type.clone(), validate) + .expect_err("malformed variable-width layout must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile with validate={validate}, got: {error}" + ); + } + } + + #[test] + fn dictionary_rejects_malformed_variable_width_values_without_optional_validation() { + let values = VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets: LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 100_000]), + bits_per_offset: 32, + num_values: 3, + block_info: BlockInfo::new(), + }; + let dictionary = DataBlock::Dictionary(DictionaryDataBlock { + indices: FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(vec![0_i32, 1, 2]), + bits_per_value: 32, + num_values: 3, + block_info: BlockInfo::new(), + }, + dictionary: Box::new(DataBlock::VariableWidth(values)), + }); + + let data_type = DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Binary)); + let error = dictionary + .into_arrow(data_type, false) + .expect_err("dictionary with out-of-bounds value offsets must be rejected"); + assert!( + matches!(error, Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + } + + #[rstest] + #[case::binary(Arc::new(BinaryArray::from_vec(vec![b"alpha", b"", b"gamma"])) as ArrayRef)] + #[case::large_binary( + Arc::new(LargeBinaryArray::from_vec(vec![b"alpha", b"", b"gamma"])) as ArrayRef + )] + #[case::utf8(Arc::new(StringArray::from(vec!["héllo", "", "world"])) as ArrayRef)] + #[case::large_utf8(Arc::new(LargeStringArray::from(vec!["héllo", "", "world"])) as ArrayRef)] + fn variable_width_valid_data_survives_mandatory_validation(#[case] array: ArrayRef) { + let block = DataBlock::from_array(array.clone()); + for validate in [false, true] { + let round_tripped = make_array( + block + .clone() + .into_arrow(array.data_type().clone(), validate) + .unwrap(), + ); + assert_eq!(&round_tripped, &array); + } + } } diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index d02cf2da693..03cd8eb3d75 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -288,52 +288,144 @@ impl BinaryMiniBlockDecompressor { } } +/// Cold path: pinpoint why the chunk-relative offsets of a binary mini-block +/// chunk failed validation. +fn chunk_offset_violation_error>(offsets: &[T], chunk_len: usize) -> Error { + let mut previous: u64 = offsets[0].into(); + for (position, &offset) in offsets.iter().enumerate().skip(1) { + let offset: u64 = offset.into(); + if offset < previous { + return Error::corrupt_file_named( + "binary mini-block", + format!( + "value offset at position {position} decreases: {offset} < {previous} \ + (chunk is {chunk_len} bytes)" + ), + ); + } + previous = offset; + } + Error::corrupt_file_named( + "binary mini-block", + format!("value offset {previous} is out of bounds for a chunk of {chunk_len} bytes"), + ) +} + impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { // decompress a MiniBlock of binary data, the num_values must be less than or equal // to the number of values this MiniBlock has, BinaryMiniBlock doesn't store `the number of values` // it has so assertion can not be done here and the caller of `decompress` must ensure // `num_values` <= number of values in the chunk. + // + // The chunk-relative value offsets at the front of the chunk come straight + // from the file and are used to slice the chunk buffer, so corrupt values + // must surface as a typed error instead of a panic or an out-of-bounds + // read. The monotonicity check rides along the existing rebase loop (the + // `&=` accumulation keeps it branchless) so validation adds no extra pass. fn decompress(&self, data: Vec, num_values: u64) -> Result { assert_eq!(data.len(), 1); let data = data.into_iter().next().unwrap(); - if self.bits_per_offset == 64 { - // offset and at least one value - assert!(data.len() >= 16); + let bytes_per_offset = self.bits_per_offset as usize / 8; + if !data.len().is_multiple_of(bytes_per_offset) { + return Err(Error::corrupt_file_named( + "binary mini-block", + format!( + "chunk size {} is not a multiple of the {}-byte offset width", + data.len(), + bytes_per_offset + ), + )); + } + let num_offsets = (num_values as usize).checked_add(1).ok_or_else(|| { + Error::corrupt_file_named( + "binary mini-block", + format!("cannot decode {num_values} values from a single chunk"), + ) + })?; + if data.len() / bytes_per_offset < num_offsets { + return Err(Error::corrupt_file_named( + "binary mini-block", + format!( + "chunk of {} bytes holds {} offsets but decoding {} values requires {}", + data.len(), + data.len() / bytes_per_offset, + num_values, + num_offsets + ), + )); + } + + // The value region must start past the offsets being decoded, otherwise + // the offset table itself aliases into the value bytes. A lower bound + // (not equality) because a prefix read of the chunk legitimately leaves + // unrequested offsets between the requested prefix and the values. + let min_value_region_start = num_offsets * bytes_per_offset; + let value_region_overlap_error = |first: u64| { + Error::corrupt_file_named( + "binary mini-block", + format!( + "value region starts at offset {first} which overlaps the {num_offsets} \ + requested offsets ({min_value_region_start} bytes)" + ), + ) + }; + if self.bits_per_offset == 64 { let offsets_buffer = data.borrow_to_typed_slice::(); - let offsets = offsets_buffer.as_ref(); + let offsets = &offsets_buffer.as_ref()[..num_offsets]; - let result_offsets = offsets[0..(num_values + 1) as usize] + let first = offsets[0]; + if first < min_value_region_start as u64 { + return Err(value_region_overlap_error(first)); + } + let mut previous = first; + let mut is_monotonic = true; + let result_offsets = offsets .iter() - .map(|offset| offset - offsets[0]) + .map(|&offset| { + is_monotonic &= previous <= offset; + previous = offset; + offset.wrapping_sub(first) + }) .collect::>(); + let last = offsets[num_offsets - 1]; + if !is_monotonic || last as usize > data.len() { + return Err(chunk_offset_violation_error(offsets, data.len())); + } Ok(DataBlock::VariableWidth(VariableWidthBlock { - data: LanceBuffer::from( - data[offsets[0] as usize..offsets[num_values as usize] as usize].to_vec(), - ), + data: LanceBuffer::from(data[first as usize..last as usize].to_vec()), offsets: LanceBuffer::reinterpret_vec(result_offsets), bits_per_offset: 64, num_values, block_info: BlockInfo::new(), })) } else { - // offset and at least one value - assert!(data.len() >= 8); - let offsets_buffer = data.borrow_to_typed_slice::(); - let offsets = offsets_buffer.as_ref(); + let offsets = &offsets_buffer.as_ref()[..num_offsets]; - let result_offsets = offsets[0..(num_values + 1) as usize] + let first = offsets[0]; + if (first as u64) < min_value_region_start as u64 { + return Err(value_region_overlap_error(first as u64)); + } + let mut previous = first; + let mut is_monotonic = true; + let result_offsets = offsets .iter() - .map(|offset| offset - offsets[0]) + .map(|&offset| { + is_monotonic &= previous <= offset; + previous = offset; + offset.wrapping_sub(first) + }) .collect::>(); + let last = offsets[num_offsets - 1]; + if !is_monotonic || last as usize > data.len() { + return Err(chunk_offset_violation_error(offsets, data.len())); + } Ok(DataBlock::VariableWidth(VariableWidthBlock { - data: LanceBuffer::from( - data[offsets[0] as usize..offsets[num_values as usize] as usize].to_vec(), - ), + data: LanceBuffer::from(data[first as usize..last as usize].to_vec()), offsets: LanceBuffer::reinterpret_vec(result_offsets), bits_per_offset: 32, num_values, @@ -462,18 +554,48 @@ impl BlockDecompressor for BinaryBlockDecompressor { // never be more than 255 and it's little endian so the last 3 bytes will always be 0. These will be the least // significant 3 bytes of the number of values in the old scheme. It's pretty unlikely these are all 0 (that would // mean there are at least 16M values in a single page) so we'll use this to determine if the old scheme is used. + // + // The header fields and the offsets themselves come straight from the file. + // The structural checks below (all O(1)) reject blocks whose regions do not + // line up; the offset *values* are validated later, by the mandatory layout + // validation in `VariableWidthBlock::into_arrow`, so they are not rescanned + // here. + if data.len() < 4 { + return Err(Error::corrupt_file_named( + "variable-width block", + format!( + "block of {} bytes is too small to hold a header", + data.len() + ), + )); + } let is_old_scheme = data[1] != 0 || data[2] != 0 || data[3] != 0; + let ensure_header = |header_len: usize| { + if data.len() < header_len { + return Err(Error::corrupt_file_named( + "variable-width block", + format!( + "block of {} bytes is too small for a {} byte header", + data.len(), + header_len + ), + )); + } + Ok(()) + }; let (bits_per_offset, bytes_start_offset, offset_start) = if is_old_scheme { // Old scheme let bits_per_offset = data[0]; match bits_per_offset { 32 => { + ensure_header(9)?; debug_assert_eq!(LittleEndian::read_u32(&data[1..5]), num_values as u32); let bytes_start_offset = LittleEndian::read_u32(&data[5..9]); - (bits_per_offset, bytes_start_offset as u64, 9) + (bits_per_offset, bytes_start_offset as u64, 9_u64) } 64 => { + ensure_header(17)?; debug_assert_eq!(LittleEndian::read_u64(&data[1..9]), num_values); let bytes_start_offset = LittleEndian::read_u64(&data[9..17]); (bits_per_offset, bytes_start_offset, 17) @@ -489,10 +611,12 @@ impl BlockDecompressor for BinaryBlockDecompressor { let bits_per_offset = LittleEndian::read_u32(&data[0..4]) as u8; match bits_per_offset { 32 => { + ensure_header(8)?; let bytes_start_offset = LittleEndian::read_u32(&data[4..8]); (bits_per_offset, bytes_start_offset as u64, 8) } 64 => { + ensure_header(16)?; let bytes_start_offset = LittleEndian::read_u64(&data[8..16]); (bits_per_offset, bytes_start_offset, 16) } @@ -504,9 +628,55 @@ impl BlockDecompressor for BinaryBlockDecompressor { } }; + // The offsets region sits between the header and `bytes_start_offset` + // and must hold exactly `num_values + 1` offsets starting at zero. + let expected_offsets_bytes = num_values + .checked_add(1) + .and_then(|num_offsets| num_offsets.checked_mul(bits_per_offset as u64 / 8)) + .ok_or_else(|| { + Error::corrupt_file_named( + "variable-width block", + format!("offsets region size overflows for {num_values} values"), + ) + })?; + if bytes_start_offset < offset_start || bytes_start_offset > data.len() as u64 { + return Err(Error::corrupt_file_named( + "variable-width block", + format!( + "bytes start offset {} is outside the block (header: {} bytes, block: {} bytes)", + bytes_start_offset, + offset_start, + data.len() + ), + )); + } + if bytes_start_offset - offset_start != expected_offsets_bytes { + return Err(Error::corrupt_file_named( + "variable-width block", + format!( + "expected {} offset bytes for {} values but found {}", + expected_offsets_bytes, + num_values, + bytes_start_offset - offset_start + ), + )); + } + // the next `bytes_start_offset - offset_start` stores the offsets. - let offsets = - data.slice_with_length(offset_start, bytes_start_offset as usize - offset_start); + let offsets = data.slice_with_length( + offset_start as usize, + (bytes_start_offset - offset_start) as usize, + ); + let first_offset = match bits_per_offset { + 32 => LittleEndian::read_u32(&offsets[0..4]) as u64, + _ => LittleEndian::read_u64(&offsets[0..8]), + }; + if first_offset != 0 { + return Err(Error::corrupt_file_named( + "variable-width block", + format!("first offset must be 0 but found {first_offset}"), + )); + } // the rest are the binary bytes. let data = data.slice_with_length( @@ -533,10 +703,12 @@ mod tests { use arrow_schema::{DataType, Field}; use crate::{ + buffer::LanceBuffer, constants::{ COMPRESSION_META_KEY, STRUCTURAL_ENCODING_FULLZIP, STRUCTURAL_ENCODING_META_KEY, STRUCTURAL_ENCODING_MINIBLOCK, }, + data::{BlockInfo, DataBlock, VariableWidthBlock}, testing::check_specific_random, }; use rstest::rstest; @@ -968,4 +1140,193 @@ mod tests { } } } + + #[test] + fn test_binary_miniblock_rejects_corrupt_offsets() { + use super::BinaryMiniBlockDecompressor; + use crate::compression::MiniBlockDecompressor; + use lance_core::Error; + + // Chunk layout mirrors the on-disk format for ["alpha", "beta", "gamma"]: + // LE u32 offsets [16, 21, 25, 30] followed by the value bytes, padded to + // a multiple of 8 bytes. + fn chunk_u32(offsets: &[u32], values: &[u8]) -> LanceBuffer { + let mut chunk = offsets + .iter() + .flat_map(|offset| offset.to_le_bytes()) + .collect::>(); + chunk.extend_from_slice(values); + chunk.resize(chunk.len().next_multiple_of(8), 0); + LanceBuffer::from(chunk) + } + + let decompressor = BinaryMiniBlockDecompressor::new(32); + + // The tail offset points past the end of the 32-byte chunk. + let err = decompressor + .decompress( + vec![chunk_u32(&[16, 21, 25, 100_000], b"alphabetagamma")], + 3, + ) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("out of bounds"), "{err}"); + + // Offsets go backwards, which would underflow the rebase subtraction. + let err = decompressor + .decompress(vec![chunk_u32(&[16, 25, 21, 30], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("decreases"), "{err}"); + + // The first offset points inside the offset table, which would alias + // the serialized offsets into the value bytes. + let err = decompressor + .decompress(vec![chunk_u32(&[0, 21, 25, 30], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("overlaps"), "{err}"); + + // The chunk stores fewer offsets than the requested value count needs. + let err = decompressor + .decompress(vec![chunk_u32(&[8, 8], &[])], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("requires 4"), "{err}"); + + // The chunk size is not a multiple of the offset width. + let err = decompressor + .decompress(vec![LanceBuffer::from(vec![0u8; 10])], 1) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("multiple"), "{err}"); + + // 64-bit offsets take the same validation path. + fn chunk_u64(offsets: &[u64], values: &[u8]) -> LanceBuffer { + let mut chunk = offsets + .iter() + .flat_map(|offset| offset.to_le_bytes()) + .collect::>(); + chunk.extend_from_slice(values); + chunk.resize(chunk.len().next_multiple_of(8), 0); + LanceBuffer::from(chunk) + } + let decompressor = BinaryMiniBlockDecompressor::new(64); + let err = decompressor + .decompress( + vec![chunk_u64(&[32, 37, 41, 100_000], b"alphabetagamma")], + 3, + ) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("out of bounds"), "{err}"); + let err = decompressor + .decompress(vec![chunk_u64(&[0, 37, 41, 46], b"alphabetagamma")], 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("overlaps"), "{err}"); + + // A valid chunk still decodes: offsets rebase to [0, 5, 9, 14]. + let decompressor = BinaryMiniBlockDecompressor::new(32); + let block = decompressor + .decompress(vec![chunk_u32(&[16, 21, 25, 30], b"alphabetagamma")], 3) + .unwrap(); + let DataBlock::VariableWidth(block) = block else { + panic!("expected a variable-width block"); + }; + assert_eq!(block.data.as_ref(), b"alphabetagamma"); + assert_eq!( + block.offsets, + LanceBuffer::reinterpret_vec(vec![0_u32, 5, 9, 14]) + ); + } + + fn encoded_binary_block(bits_per_offset: u8) -> Vec { + use crate::compression::BlockCompressor; + + let offsets = match bits_per_offset { + 32 => LanceBuffer::reinterpret_vec(vec![0_i32, 5, 9, 14]), + 64 => LanceBuffer::reinterpret_vec(vec![0_i64, 5, 9, 14]), + _ => unreachable!(), + }; + let block = DataBlock::VariableWidth(VariableWidthBlock { + data: LanceBuffer::copy_slice(b"alphabetagamma"), + offsets, + bits_per_offset, + num_values: 3, + block_info: BlockInfo::new(), + }); + BlockCompressor::compress(&super::VariableEncoder::default(), block) + .unwrap() + .as_ref() + .to_vec() + } + + /// The block decompressor only checks the block structure (all O(1)); bad + /// offset values inside a structurally-sound block are rejected by the + /// mandatory layout validation when the block is converted to Arrow. + #[rstest] + #[case::i32_tail_out_of_bounds(32, 3, 100_000, "out of bounds")] + #[case::i64_tail_out_of_bounds(64, 3, 15, "out of bounds")] + #[case::i32_non_monotonic(32, 2, 4, "non-monotonic")] + #[case::i64_non_monotonic(64, 2, 4, "non-monotonic")] + fn test_binary_block_bad_offsets_rejected_at_arrow_conversion( + #[case] bits_per_offset: u8, + #[case] mutated_offset_index: usize, + #[case] mutated_offset_value: u64, + #[case] expected_message: &str, + ) { + use crate::compression::BlockDecompressor; + use lance_core::Error; + + let mut encoded = encoded_binary_block(bits_per_offset); + let bytes_per_offset = (bits_per_offset / 8) as usize; + // The standard scheme header is two offset-width fields. + let mutated_offset_start = bytes_per_offset * (2 + mutated_offset_index); + encoded[mutated_offset_start..mutated_offset_start + bytes_per_offset] + .copy_from_slice(&mutated_offset_value.to_le_bytes()[..bytes_per_offset]); + + let block = super::BinaryBlockDecompressor::default() + .decompress(LanceBuffer::from(encoded), 3) + .unwrap(); + let data_type = match bits_per_offset { + 32 => DataType::Binary, + _ => DataType::LargeBinary, + }; + let err = block.into_arrow(data_type, false).unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains(expected_message), "{err}"); + } + + #[test] + fn test_binary_block_rejects_corrupt_structure() { + use crate::compression::BlockDecompressor; + use lance_core::Error; + + let decompressor = super::BinaryBlockDecompressor::default(); + + // The first offset must be zero. + let mut encoded = encoded_binary_block(32); + encoded[8..12].copy_from_slice(&5_u32.to_le_bytes()); + let err = decompressor + .decompress(LanceBuffer::from(encoded), 3) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("first offset"), "{err}"); + + // The offsets region must hold exactly num_values + 1 offsets. + let encoded = encoded_binary_block(32); + let err = decompressor + .decompress(LanceBuffer::from(encoded), 4) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("offset bytes"), "{err}"); + + // A block too small to hold its header is rejected, not a panic. + let err = decompressor + .decompress(LanceBuffer::from(vec![0_u8; 2]), 1) + .unwrap_err(); + assert!(matches!(err, Error::CorruptFile { .. }), "{err:?}"); + assert!(err.to_string().contains("too small"), "{err}"); + } } diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index f631c0a7892..288f9440126 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -1755,6 +1755,178 @@ mod tests { assert_eq!(remaining, 0); } + /// Writes `batch` to a fresh file, overwrites `patch` bytes at `patch_offset` + /// into the single occurrence of `pattern`, and reads the file back with the + /// default reader configuration. + async fn read_file_with_mutated_bytes( + version: LanceFileVersion, + batch: RecordBatch, + pattern: &[u8], + patch_offset: usize, + patch: &[u8], + ) -> lance_core::Result> { + let fs = FsFixture::default(); + let schema = batch.schema(); + write_lance_file( + RecordBatchIterator::new(vec![Ok(batch)], schema), + &fs, + ConcreteFileVersion::from(version), + FileWriterOptions::default(), + ) + .await; + + let mut bytes = fs + .object_store + .read_one_all(&fs.tmp_path) + .await + .unwrap() + .to_vec(); + let matches = bytes + .windows(pattern.len()) + .enumerate() + .filter_map(|(position, window)| (window == pattern).then_some(position)) + .collect::>(); + assert_eq!( + matches.len(), + 1, + "expected the byte pattern to appear exactly once in the file" + ); + let patch_start = matches[0] + patch_offset; + bytes[patch_start..patch_start + patch.len()].copy_from_slice(patch); + fs.object_store.put(&fs.tmp_path, &bytes).await.unwrap(); + + let file_scheduler = fs + .scheduler + .open_file(&fs.tmp_path, &CachedFileSize::unknown()) + .await + .unwrap(); + let file_reader = FileReader::try_open( + file_scheduler, + None, + Arc::::default(), + &test_cache(), + FileReaderOptions::default(), + ) + .await + .unwrap(); + file_reader + .read_stream( + lance_io::ReadBatchParams::RangeFull, + 1024, + 16, + FilterExpression::no_filter(), + ) + .await? + .try_collect::>() + .await + } + + /// A corrupt file whose variable-width offsets point outside the value bytes + /// must fail with a typed error under the default reader configuration + /// (`validate_on_decode` disabled) instead of materializing values outside + /// the data buffer. + /// + /// Uses a dictionary-encoded string column because its values page stores + /// the offsets verbatim, so flipping the tail offset in the file reaches the + /// Arrow conversion boundary without being rejected by an intermediate + /// decompressor. + #[rstest] + #[tokio::test] + async fn test_default_reader_rejects_out_of_bounds_variable_width_offsets( + #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)] + version: LanceFileVersion, + ) { + use arrow_array::{Array, DictionaryArray, Int32Array, StringArray}; + + let values = StringArray::from(vec!["alpha", "beta", "gamma"]); + let indices = Int32Array::from((0..300).map(|i| i % 3).collect::>()); + let dictionary = DictionaryArray::new(indices, Arc::new(values)); + let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "category", + dictionary.data_type().clone(), + false, + )])); + let batch = RecordBatch::try_new(arrow_schema, vec![Arc::new(dictionary)]).unwrap(); + + // The dictionary values page stores the value offsets as plain + // little-endian i32s ending with [5, 9, 14] (2.1 also stores the leading + // zero, 2.2+ omits it). If a future encoding change stops storing these + // offsets verbatim this lookup fails loudly and the test needs a new + // byte pattern. The patch rewrites the tail offset so it points far + // beyond the value bytes. + let offsets_tail_pattern = [5_i32, 9, 14] + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let error = read_file_with_mutated_bytes( + version, + batch, + &offsets_tail_pattern, + 8, + &100_000_i32.to_le_bytes(), + ) + .await + .expect_err("out-of-bounds offsets must fail the read"); + assert!( + matches!(error, lance_core::Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + assert!( + error.to_string().contains("out of bounds"), + "unexpected message: {error}" + ); + } + + /// Same contract as the test above, but for a plain (non-dictionary) string + /// column: the mini-block chunk stores chunk-relative value offsets that are + /// used to slice the chunk, so a corrupt tail offset must surface as a typed + /// error from the chunk decompressor instead of a panic in the decode task. + #[rstest] + #[tokio::test] + async fn test_default_reader_rejects_out_of_bounds_miniblock_offsets( + #[values(LanceFileVersion::V2_1, LanceFileVersion::V2_2, LanceFileVersion::V2_3)] + version: LanceFileVersion, + ) { + use arrow_array::StringArray; + + let arrow_schema = Arc::new(ArrowSchema::new(vec![Field::new( + "strings", + DataType::Utf8, + false, + )])); + let batch = RecordBatch::try_new( + arrow_schema, + vec![Arc::new(StringArray::from(vec!["alpha", "beta", "gamma"]))], + ) + .unwrap(); + + // For ["alpha", "beta", "gamma"] the chunk stores LE i32 offsets + // [16, 21, 25, 30] (chunk-relative: a 16-byte offsets region precedes + // the value bytes). The patch rewrites the tail offset to point far + // past the chunk. + let chunk_offsets_pattern = [16_i32, 21, 25, 30] + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect::>(); + let error = read_file_with_mutated_bytes( + version, + batch, + &chunk_offsets_pattern, + 12, + &100_000_i32.to_le_bytes(), + ) + .await + .expect_err("an out-of-bounds chunk offset must fail the read"); + assert!( + matches!(error, lance_core::Error::CorruptFile { .. }), + "expected CorruptFile, got: {error}" + ); + assert!( + error.to_string().contains("out of bounds"), + "unexpected message: {error}" + ); + } + #[tokio::test] async fn test_round_trip() { let fs = FsFixture::default(); From 7ebed3c1bd1f82f779ad38d13f29edca71b4bf4a Mon Sep 17 00:00:00 2001 From: Xuanwo Date: Mon, 3 Aug 2026 10:31:07 +0800 Subject: [PATCH 4/4] fix(encoding): adapt corruption backport to release v6.1 APIs --- rust/lance-encoding/src/data.rs | 11 +++-- .../src/encodings/logical/primitive.rs | 42 ++++++++++++++++--- .../src/encodings/physical/binary.rs | 28 +++++++------ rust/lance-file/src/reader.rs | 26 +++++++----- 4 files changed, 74 insertions(+), 33 deletions(-) diff --git a/rust/lance-encoding/src/data.rs b/rust/lance-encoding/src/data.rs index 588e953f232..417c026f0fa 100644 --- a/rust/lance-encoding/src/data.rs +++ b/rust/lance-encoding/src/data.rs @@ -587,6 +587,10 @@ pub struct VariableWidthBlock { /// caller-controlled flag. struct ValidVariableWidthLayout; +fn corrupt_file_named(name: &str, message: impl Into) -> Error { + Error::corrupt_file(name.into(), message) +} + impl VariableWidthBlock { // The offsets buffer comes straight from file bytes, so an unchecked build would // let a corrupt file smuggle out-of-bounds offsets into an Arrow array whose @@ -647,7 +651,7 @@ impl VariableWidthBlock { offsets_size: usize, data_size: usize, ) -> Error { - Error::corrupt_file_named( + corrupt_file_named( "variable width data block", format!( "invalid variable-width layout for {}: {} (num_values: {}, bits_per_offset: {}, \ @@ -1853,9 +1857,8 @@ mod tests { use std::sync::Arc; use arrow_array::{ - ArrayRef, BinaryArray, BinaryViewArray, DictionaryArray, Int8Array, LargeBinaryArray, - LargeStringArray, StringArray, StringViewArray, UInt8Array, UInt16Array, make_array, - new_null_array, + ArrayRef, BinaryArray, DictionaryArray, Int8Array, LargeBinaryArray, LargeStringArray, + StringArray, UInt8Array, UInt16Array, make_array, new_null_array, types::{Int8Type, Int32Type}, }; use arrow_buffer::{BooleanBuffer, NullBuffer}; diff --git a/rust/lance-encoding/src/encodings/logical/primitive.rs b/rust/lance-encoding/src/encodings/logical/primitive.rs index 9dee00445a8..7ba6ccde57a 100644 --- a/rust/lance-encoding/src/encodings/logical/primitive.rs +++ b/rust/lance-encoding/src/encodings/logical/primitive.rs @@ -2622,6 +2622,10 @@ struct VariableFullZipDecoder { num_rows: u64, } +fn corrupt_file_named(name: &str, message: impl Into) -> Error { + Error::corrupt_file(name.into(), message) +} + impl VariableFullZipDecoder { fn new( details: Arc, @@ -2760,7 +2764,7 @@ impl VariableFullZipDecoder { fn parse_length(data: &[u8], bits_per_offset: u8) -> Result { let width = bits_per_offset as usize / 8; if data.len() < width { - return Err(Error::corrupt_file_named( + return Err(corrupt_file_named( "variable_full_zip", format!( "truncated length prefix: {} byte(s) remain in the page buffer but a \ @@ -7418,11 +7422,37 @@ mod tests { check_round_trip_encoding_of_data(vec![Arc::new(list_array)], &test_cases, HashMap::new()) .await; } - /// A well-formed length prefix decodes without incident, for both widths. - #[test] - fn variable_full_zip_wellformed_length_prefix() { - assert!(decode_variable_full_zip(0u32.to_le_bytes().to_vec(), 32).is_ok()); - assert!(decode_variable_full_zip(0u64.to_le_bytes().to_vec(), 64).is_ok()); + fn truncated_tail_details() -> std::sync::Arc { + use crate::compression::VariablePerValueDecompressor; + use crate::encodings::physical::binary::VariableDecoder; + use crate::repdef::{ControlWordParser, DefinitionInterpretation}; + use std::sync::Arc; + Arc::new(super::FullZipDecodeDetails { + value_decompressor: super::PerValueDecompressor::Variable(Arc::new( + VariableDecoder::default(), + ) + as Arc), + def_meaning: vec![DefinitionInterpretation::NullableItem].into(), + ctrl_word_parser: ControlWordParser::new(0, 0), + max_rep: 0, + max_visible_def: 0, + }) + } + + fn decode_variable_full_zip( + buf: Vec, + bits_per_offset: u8, + ) -> lance_core::Result { + use std::collections::VecDeque; + let mut data = VecDeque::new(); + data.push_back(crate::buffer::LanceBuffer::from(buf)); + super::VariableFullZipDecoder::new( + truncated_tail_details(), + data, + 1, + bits_per_offset, + bits_per_offset, + ) } /// A page whose item walk ends with a partial length prefix must surface a diff --git a/rust/lance-encoding/src/encodings/physical/binary.rs b/rust/lance-encoding/src/encodings/physical/binary.rs index 03cd8eb3d75..ecb24ad6d0b 100644 --- a/rust/lance-encoding/src/encodings/physical/binary.rs +++ b/rust/lance-encoding/src/encodings/physical/binary.rs @@ -13,6 +13,10 @@ use arrow_array::OffsetSizeTrait; use byteorder::{ByteOrder, LittleEndian}; use core::panic; +fn corrupt_file_named(name: &str, message: impl Into) -> Error { + Error::corrupt_file(name.into(), message) +} + use crate::compression::{ BlockCompressor, BlockDecompressor, MiniBlockDecompressor, VariablePerValueDecompressor, }; @@ -295,7 +299,7 @@ fn chunk_offset_violation_error>(offsets: &[T], chunk_len: u for (position, &offset) in offsets.iter().enumerate().skip(1) { let offset: u64 = offset.into(); if offset < previous { - return Error::corrupt_file_named( + return corrupt_file_named( "binary mini-block", format!( "value offset at position {position} decreases: {offset} < {previous} \ @@ -305,7 +309,7 @@ fn chunk_offset_violation_error>(offsets: &[T], chunk_len: u } previous = offset; } - Error::corrupt_file_named( + corrupt_file_named( "binary mini-block", format!("value offset {previous} is out of bounds for a chunk of {chunk_len} bytes"), ) @@ -328,7 +332,7 @@ impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { let bytes_per_offset = self.bits_per_offset as usize / 8; if !data.len().is_multiple_of(bytes_per_offset) { - return Err(Error::corrupt_file_named( + return Err(corrupt_file_named( "binary mini-block", format!( "chunk size {} is not a multiple of the {}-byte offset width", @@ -338,13 +342,13 @@ impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { )); } let num_offsets = (num_values as usize).checked_add(1).ok_or_else(|| { - Error::corrupt_file_named( + corrupt_file_named( "binary mini-block", format!("cannot decode {num_values} values from a single chunk"), ) })?; if data.len() / bytes_per_offset < num_offsets { - return Err(Error::corrupt_file_named( + return Err(corrupt_file_named( "binary mini-block", format!( "chunk of {} bytes holds {} offsets but decoding {} values requires {}", @@ -362,7 +366,7 @@ impl MiniBlockDecompressor for BinaryMiniBlockDecompressor { // unrequested offsets between the requested prefix and the values. let min_value_region_start = num_offsets * bytes_per_offset; let value_region_overlap_error = |first: u64| { - Error::corrupt_file_named( + corrupt_file_named( "binary mini-block", format!( "value region starts at offset {first} which overlaps the {num_offsets} \ @@ -561,7 +565,7 @@ impl BlockDecompressor for BinaryBlockDecompressor { // validation in `VariableWidthBlock::into_arrow`, so they are not rescanned // here. if data.len() < 4 { - return Err(Error::corrupt_file_named( + return Err(corrupt_file_named( "variable-width block", format!( "block of {} bytes is too small to hold a header", @@ -573,7 +577,7 @@ impl BlockDecompressor for BinaryBlockDecompressor { let ensure_header = |header_len: usize| { if data.len() < header_len { - return Err(Error::corrupt_file_named( + return Err(corrupt_file_named( "variable-width block", format!( "block of {} bytes is too small for a {} byte header", @@ -634,13 +638,13 @@ impl BlockDecompressor for BinaryBlockDecompressor { .checked_add(1) .and_then(|num_offsets| num_offsets.checked_mul(bits_per_offset as u64 / 8)) .ok_or_else(|| { - Error::corrupt_file_named( + corrupt_file_named( "variable-width block", format!("offsets region size overflows for {num_values} values"), ) })?; if bytes_start_offset < offset_start || bytes_start_offset > data.len() as u64 { - return Err(Error::corrupt_file_named( + return Err(corrupt_file_named( "variable-width block", format!( "bytes start offset {} is outside the block (header: {} bytes, block: {} bytes)", @@ -651,7 +655,7 @@ impl BlockDecompressor for BinaryBlockDecompressor { )); } if bytes_start_offset - offset_start != expected_offsets_bytes { - return Err(Error::corrupt_file_named( + return Err(corrupt_file_named( "variable-width block", format!( "expected {} offset bytes for {} values but found {}", @@ -672,7 +676,7 @@ impl BlockDecompressor for BinaryBlockDecompressor { _ => LittleEndian::read_u64(&offsets[0..8]), }; if first_offset != 0 { - return Err(Error::corrupt_file_named( + return Err(corrupt_file_named( "variable-width block", format!("first offset must be 0 but found {first_offset}"), )); diff --git a/rust/lance-file/src/reader.rs b/rust/lance-file/src/reader.rs index 288f9440126..8c89a722428 100644 --- a/rust/lance-file/src/reader.rs +++ b/rust/lance-file/src/reader.rs @@ -1651,7 +1651,7 @@ mod tests { use std::{collections::BTreeMap, pin::Pin, sync::Arc}; use arrow_array::{ - RecordBatch, UInt32Array, + RecordBatch, RecordBatchIterator, UInt32Array, types::{Float64Type, Int32Type}, }; use arrow_schema::{DataType, Field, Fields, Schema as ArrowSchema}; @@ -1770,8 +1770,10 @@ mod tests { write_lance_file( RecordBatchIterator::new(vec![Ok(batch)], schema), &fs, - ConcreteFileVersion::from(version), - FileWriterOptions::default(), + FileWriterOptions { + format_version: Some(version), + ..Default::default() + }, ) .await; @@ -1867,13 +1869,14 @@ mod tests { ) .await .expect_err("out-of-bounds offsets must fail the read"); + let error_message = error.to_string(); assert!( - matches!(error, lance_core::Error::CorruptFile { .. }), - "expected CorruptFile, got: {error}" + error_message.contains("corrupt file"), + "expected a corruption error, got: {error_message}" ); assert!( - error.to_string().contains("out of bounds"), - "unexpected message: {error}" + error_message.contains("out of bounds"), + "unexpected message: {error_message}" ); } @@ -1917,13 +1920,14 @@ mod tests { ) .await .expect_err("an out-of-bounds chunk offset must fail the read"); + let error_message = error.to_string(); assert!( - matches!(error, lance_core::Error::CorruptFile { .. }), - "expected CorruptFile, got: {error}" + error_message.contains("corrupt file"), + "expected a corruption error, got: {error_message}" ); assert!( - error.to_string().contains("out of bounds"), - "unexpected message: {error}" + error_message.contains("out of bounds"), + "unexpected message: {error_message}" ); }