diff --git a/rust/lance-core/src/error.rs b/rust/lance-core/src/error.rs index a85f08cb741..e24bb1a0acc 100644 --- a/rust/lance-core/src/error.rs +++ b/rust/lance-core/src/error.rs @@ -440,6 +440,18 @@ impl Error { CorruptFileSnafu { path }.into_error(message.into().into()) } + /// Reports a corrupt file when the caller only has a logical/section name + /// rather than the real file path (for example, a decoder that validates an + /// in-memory buffer and does not know where it came from). + /// + /// `name` is carried in the `path` field of the resulting [`Error::CorruptFile`] + /// variant and is NOT a filesystem path; callers that have the real path should + /// use [`Self::corrupt_file`] instead. + #[track_caller] + pub fn corrupt_file_named(name: &str, message: impl Into) -> Self { + Self::corrupt_file(object_store::path::Path::from(name), message) + } + #[track_caller] pub fn invalid_input(message: impl Into) -> Self { InvalidInputSnafu.into_error(message.into().into()) diff --git a/rust/lance-encoding/benches/decoder.rs b/rust/lance-encoding/benches/decoder.rs index cc0404e1bb3..abfad51b1b4 100644 --- a/rust/lance-encoding/benches/decoder.rs +++ b/rust/lance-encoding/benches/decoder.rs @@ -1,14 +1,28 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright The Lance Authors -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, hint::black_box, sync::Arc}; use arrow_array::{RecordBatch, UInt32Array}; +#[cfg(feature = "bitpacking")] +use arrow_buffer::ArrowNativeType; use arrow_schema::{DataType, Field, Schema, TimeUnit}; use arrow_select::take::take; +#[cfg(feature = "bitpacking")] +use bytemuck::Pod; use criterion::{Criterion, criterion_group, criterion_main}; use futures::StreamExt; +#[cfg(feature = "bitpacking")] +use lance_bitpacking::BitPacking; use lance_core::cache::LanceCache; use lance_datagen::ArrayGeneratorExt; +#[cfg(feature = "bitpacking")] +use lance_encoding::buffer::LanceBuffer; +#[cfg(feature = "bitpacking")] +use lance_encoding::compression::BlockDecompressor; +#[cfg(feature = "bitpacking")] +use lance_encoding::data::{BlockInfo, DataBlock, FixedWidthDataBlock}; +#[cfg(feature = "bitpacking")] +use lance_encoding::encodings::physical::bitpacking::{ELEMS_PER_CHUNK, InlineBitpacking}; use lance_encoding::{ decoder::{ DecodeBatchScheduler, DecoderConfig, DecoderPlugins, FilterExpression, create_decode_stream, @@ -563,6 +577,136 @@ fn bench_decode_compressed_parallel(c: &mut Criterion) { } } +#[cfg(feature = "bitpacking")] +fn make_inline_bitpacking_chunk(bit_width: usize) -> LanceBuffer +where + T: ArrowNativeType + BitPacking + Pod, +{ + let value_range = 1_usize << bit_width; + let values: Vec = (0..ELEMS_PER_CHUNK as usize) + .map(|i| T::from_usize((i * 31 + 7) % value_range).unwrap()) + .collect(); + let packed_words = ELEMS_PER_CHUNK as usize * bit_width / (std::mem::size_of::() * 8); + + let mut chunk = Vec::with_capacity(1 + packed_words); + chunk.push(T::from_usize(bit_width).unwrap()); + let payload_start = chunk.len(); + chunk.resize(payload_start + packed_words, T::from_usize(0).unwrap()); + unsafe { + BitPacking::unchecked_pack(bit_width, &values, &mut chunk[payload_start..]); + } + + LanceBuffer::reinterpret_vec(chunk) +} + +#[cfg(feature = "bitpacking")] +fn read_little_endian_header(bytes: &[u8]) -> usize { + bytes[..std::mem::size_of::()] + .iter() + .enumerate() + .fold(0_u64, |value, (idx, byte)| { + value | ((*byte as u64) << (idx * 8)) + }) as usize +} + +#[cfg(feature = "bitpacking")] +fn legacy_copy_unchunk(data: LanceBuffer, num_values: u64) -> DataBlock +where + T: ArrowNativeType + BitPacking + Pod, +{ + assert!(data.len() >= std::mem::size_of::()); + assert!(num_values <= ELEMS_PER_CHUNK); + + let chunk_in_u8 = data.to_vec(); + let bit_width_value = read_little_endian_header::(&chunk_in_u8); + let chunk = bytemuck::cast_slice(&chunk_in_u8[std::mem::size_of::()..]); + assert!(std::mem::size_of_val(chunk) == bit_width_value * ELEMS_PER_CHUNK as usize / 8); + + let mut decompressed = vec![T::from_usize(0).unwrap(); ELEMS_PER_CHUNK as usize]; + unsafe { + BitPacking::unchecked_unpack(bit_width_value, chunk, &mut decompressed); + } + + decompressed.truncate(num_values as usize); + DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::reinterpret_vec(decompressed), + bits_per_value: (std::mem::size_of::() * 8) as u64, + num_values, + block_info: BlockInfo::new(), + }) +} + +#[cfg(feature = "bitpacking")] +fn typed_view_unchunk(buffer: LanceBuffer, uncompressed_bits: u64, num_values: u64) -> DataBlock { + InlineBitpacking::new(uncompressed_bits) + .decompress(buffer, num_values) + .unwrap() +} + +#[cfg(feature = "bitpacking")] +fn assert_same_fixed_width_payloads(legacy: &DataBlock, typed_view: &DataBlock) { + let legacy = legacy.as_fixed_width_ref().unwrap(); + let typed_view = typed_view.as_fixed_width_ref().unwrap(); + + assert_eq!(legacy.num_values, typed_view.num_values); + assert_eq!(legacy.bits_per_value, typed_view.bits_per_value); + assert_eq!(legacy.data.as_ref(), typed_view.data.as_ref()); +} + +#[cfg(feature = "bitpacking")] +fn bench_inline_bitpacking_case( + group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>, + name: &str, + bit_width: usize, +) where + T: ArrowNativeType + BitPacking + Pod, +{ + let buffer = make_inline_bitpacking_chunk::(bit_width); + let compressed_bytes = buffer.len() as u64; + let uncompressed_bits = (std::mem::size_of::() * 8) as u64; + group.throughput(criterion::Throughput::Bytes(compressed_bytes)); + + let legacy = legacy_copy_unchunk::(buffer.clone(), ELEMS_PER_CHUNK); + let typed_view = typed_view_unchunk(buffer.clone(), uncompressed_bits, ELEMS_PER_CHUNK); + assert_same_fixed_width_payloads(&legacy, &typed_view); + + group.bench_function(format!("{name}/legacy_copy/compressed_bytes"), |b| { + b.iter(|| { + let decoded = + legacy_copy_unchunk::(black_box(buffer.clone()), black_box(ELEMS_PER_CHUNK)); + let fixed = decoded.as_fixed_width().unwrap(); + black_box(fixed.data.as_ref()); + }) + }); + + group.bench_function(format!("{name}/typed_view/compressed_bytes"), |b| { + b.iter(|| { + let decoded = typed_view_unchunk( + black_box(buffer.clone()), + black_box(uncompressed_bits), + black_box(ELEMS_PER_CHUNK), + ); + let fixed = decoded.as_fixed_width().unwrap(); + black_box(fixed.data.as_ref()); + }) + }); +} + +#[cfg(feature = "bitpacking")] +fn bench_decode_inline_bitpacking_unchunk(c: &mut Criterion) { + let mut group = c.benchmark_group("decode_inline_bitpacking_unchunk"); + bench_inline_bitpacking_case::(&mut group, "u32_bw12_1024", 12); + bench_inline_bitpacking_case::(&mut group, "u64_bw23_1024", 23); + group.finish(); +} + +#[cfg(not(feature = "bitpacking"))] +fn bench_decode_inline_bitpacking_unchunk(c: &mut Criterion) { + let mut group = c.benchmark_group("decode_inline_bitpacking_unchunk"); + group.bench_function("bitpacking_feature_disabled", |b| b.iter(|| black_box(()))); + group.finish(); +} + #[cfg(target_os = "linux")] criterion_group!( name=benches; @@ -570,7 +714,7 @@ criterion_group!( .with_profiler(lance_testing::pprof::PProfProfiler::new(100, lance_testing::pprof::Output::Flamegraph(None))); targets = bench_decode, bench_decode_fsl, bench_decode_str_with_dict_encoding, bench_decode_packed_struct, bench_decode_str_with_fixed_size_binary_encoding, bench_decode_compressed, - bench_decode_compressed_parallel); + bench_decode_compressed_parallel, bench_decode_inline_bitpacking_unchunk); // Non-linux version does not support pprof. #[cfg(not(target_os = "linux"))] @@ -578,5 +722,5 @@ criterion_group!( name=benches; config = Criterion::default().significance_level(0.1).sample_size(10); targets = bench_decode, bench_decode_fsl, bench_decode_str_with_dict_encoding, bench_decode_packed_struct, - bench_decode_compressed, bench_decode_compressed_parallel); + bench_decode_compressed, bench_decode_compressed_parallel, bench_decode_inline_bitpacking_unchunk); criterion_main!(benches); diff --git a/rust/lance-encoding/src/encodings/physical/bitpacking.rs b/rust/lance-encoding/src/encodings/physical/bitpacking.rs index be0b747e7dc..8f41d47a4e2 100644 --- a/rust/lance-encoding/src/encodings/physical/bitpacking.rs +++ b/rust/lance-encoding/src/encodings/physical/bitpacking.rs @@ -18,7 +18,6 @@ use arrow_array::types::UInt64Type; use arrow_array::{Array, PrimitiveArray}; use arrow_buffer::ArrowNativeType; -use byteorder::{ByteOrder, LittleEndian}; use lance_bitpacking::BitPacking; use lance_core::{Error, Result}; @@ -33,10 +32,11 @@ use crate::encodings::logical::primitive::miniblock::{ use crate::format::pb21::CompressiveEncoding; use crate::format::{ProtobufUtils21, pb21}; use crate::statistics::{GetStat, Stat}; -use bytemuck::{AnyBitPattern, cast_slice}; +use bytemuck::Pod; -const LOG_ELEMS_PER_CHUNK: u8 = 10; -const ELEMS_PER_CHUNK: u64 = 1 << LOG_ELEMS_PER_CHUNK; +pub(crate) const LOG_ELEMS_PER_CHUNK: u8 = 10; +/// Number of values encoded in each inline bitpacking chunk. +pub const ELEMS_PER_CHUNK: u64 = 1 << LOG_ELEMS_PER_CHUNK; #[derive(Debug, Default)] pub struct InlineBitpacking { @@ -181,27 +181,85 @@ impl InlineBitpacking { ) } - fn unchunk( + fn unchunk( data: LanceBuffer, num_values: u64, ) -> Result { - // Ensure at least the header is present - assert!(data.len() >= std::mem::size_of::()); - assert!(num_values <= ELEMS_PER_CHUNK); - // This macro decompresses a chunk(1024 values) of bitpacked values. let uncompressed_bit_width = std::mem::size_of::() * 8; - let mut decompressed = vec![T::from_usize(0).unwrap(); ELEMS_PER_CHUNK as usize]; - - // Copy for memory alignment - let chunk_in_u8: Vec = data.to_vec(); - let bit_width_bytes = &chunk_in_u8[..std::mem::size_of::()]; - let bit_width_value = LittleEndian::read_uint(bit_width_bytes, std::mem::size_of::()); - let chunk = cast_slice(&chunk_in_u8[std::mem::size_of::()..]); - // The bit-packed chunk should have number of bytes (bit_width_value * ELEMS_PER_CHUNK / 8) - assert!(std::mem::size_of_val(chunk) == (bit_width_value * ELEMS_PER_CHUNK) as usize / 8); + let word_size = std::mem::size_of::(); + + if data.len() < word_size { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking chunk is too small for {}-byte header: {} bytes", + word_size, + data.len() + ), + )); + } + if !data.len().is_multiple_of(word_size) { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking chunk size must be a multiple of {} bytes, got {} bytes", + word_size, + data.len() + ), + )); + } + if num_values > ELEMS_PER_CHUNK { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking chunk has {} values, expected at most {}", + num_values, ELEMS_PER_CHUNK + ), + )); + } + + let chunk_words = data.borrow_to_typed_view::(); + let bit_width_value = chunk_words[0].as_usize(); + if bit_width_value > uncompressed_bit_width { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking width {} exceeds {}-bit values", + bit_width_value, uncompressed_bit_width + ), + )); + } + let chunk = &chunk_words[1..]; + // bit_width_value has already been verified to be <= uncompressed_bit_width + // (8/16/32/64), so bit_width_value * ELEMS_PER_CHUNK (1024) can never + // overflow usize on supported targets. Keep checked_mul as defense in depth. + let expected_num_bits = bit_width_value + .checked_mul(ELEMS_PER_CHUNK as usize) + .ok_or_else(|| { + Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking width {} overflows chunk bit count", + bit_width_value + ), + ) + })?; + let expected_num_bytes = expected_num_bits / 8; + let actual_num_bytes = std::mem::size_of_val(chunk); + if actual_num_bytes != expected_num_bytes { + return Err(Error::corrupt_file_named( + "inline_bitpacking", + format!( + "Inline bitpacking payload has {} bytes, expected {} bytes for bit width {}", + actual_num_bytes, expected_num_bytes, bit_width_value + ), + )); + } + + let mut decompressed = vec![T::default(); ELEMS_PER_CHUNK as usize]; unsafe { - BitPacking::unchecked_unpack(bit_width_value as usize, chunk, &mut decompressed); + BitPacking::unchecked_unpack(bit_width_value, chunk, &mut decompressed); } decompressed.truncate(num_values as usize); @@ -212,6 +270,17 @@ impl InlineBitpacking { block_info: BlockInfo::new(), })) } + + /// An empty fixed-width block, used for the `num_values == 0` short-circuit in + /// both decompressor entry points so empty blocks skip chunk validation entirely. + fn empty_block(&self) -> DataBlock { + DataBlock::FixedWidth(FixedWidthDataBlock { + data: LanceBuffer::empty(), + bits_per_value: self.uncompressed_bit_width, + num_values: 0, + block_info: BlockInfo::new(), + }) + } } impl MiniBlockCompressor for InlineBitpacking { @@ -243,12 +312,7 @@ impl MiniBlockDecompressor for InlineBitpacking { let data = data.into_iter().next().unwrap(); if num_values == 0 { // Empty mini-blocks have no inline bit-width header to decode. - return Ok(DataBlock::FixedWidth(FixedWidthDataBlock { - data: LanceBuffer::empty(), - bits_per_value: self.uncompressed_bit_width, - num_values: 0, - block_info: BlockInfo::new(), - })); + return Ok(self.empty_block()); } match self.uncompressed_bit_width { 8 => Self::unchunk::(data, num_values), @@ -262,6 +326,12 @@ impl MiniBlockDecompressor for InlineBitpacking { impl BlockDecompressor for InlineBitpacking { fn decompress(&self, data: LanceBuffer, num_values: u64) -> Result { + if num_values == 0 { + // Empty blocks carry no inline bit-width header to decode; avoid + // spurious "too small for header" corrupt-file errors and mirror + // the MiniBlockDecompressor path. See #7794. + return Ok(self.empty_block()); + } match self.uncompressed_bit_width { 8 => Self::unchunk::(data, num_values), 16 => Self::unchunk::(data, num_values), @@ -536,13 +606,16 @@ mod test { use std::{collections::HashMap, sync::Arc}; use arrow_array::{Array, Int8Array, Int64Array}; + use arrow_buffer::ArrowNativeType; use arrow_schema::DataType; + use bytemuck::Pod; + use lance_bitpacking::BitPacking; use rstest::rstest; use super::{ELEMS_PER_CHUNK, InlineBitpacking, bitpack_out_of_line, unpack_out_of_line}; use crate::{ buffer::LanceBuffer, - compression::MiniBlockDecompressor, + compression::{BlockDecompressor, MiniBlockDecompressor}, data::{BlockInfo, DataBlock, FixedWidthDataBlock}, testing::{TestCases, check_round_trip_encoding_of_data}, version::LanceFileVersion, @@ -567,6 +640,98 @@ mod test { assert_eq!(block.data.len(), 0); } + // Regression test for #7794: the block-level decompressor must short-circuit + // on num_values == 0 the same way the mini-block decompressor does, instead + // of reporting a spurious "too small for header" corrupt-file error. + #[rstest] + #[case::u8(8)] + #[case::u16(16)] + #[case::u32(32)] + #[case::u64(64)] + fn test_inline_bitpacking_decompress_empty_block(#[case] bit_width: u64) { + let decompressor = InlineBitpacking::new(bit_width); + let decompressed = + BlockDecompressor::decompress(&decompressor, LanceBuffer::empty(), 0).unwrap(); + + let DataBlock::FixedWidth(block) = decompressed else { + panic!("Expected FixedWidth block"); + }; + assert_eq!(block.bits_per_value, bit_width); + assert_eq!(block.num_values, 0); + assert_eq!(block.data.len(), 0); + } + + fn roundtrip_unchunk(values: &[T], bit_width: usize) + where + T: ArrowNativeType + BitPacking + Pod, + { + assert!(values.len() <= ELEMS_PER_CHUNK as usize); + let num_values = values.len() as u64; + + let mut padded = vec![T::from_usize(0).unwrap(); ELEMS_PER_CHUNK as usize]; + padded[..values.len()].copy_from_slice(values); + + let packed_words = ELEMS_PER_CHUNK as usize * bit_width / (std::mem::size_of::() * 8); + let mut chunk: Vec = Vec::with_capacity(1 + packed_words); + chunk.push(T::from_usize(bit_width).unwrap()); + let out_len = chunk.len(); + chunk.resize(out_len + packed_words, T::from_usize(0).unwrap()); + unsafe { + BitPacking::unchecked_pack(bit_width, &padded, &mut chunk[out_len..]); + } + + let data = LanceBuffer::reinterpret_vec(chunk); + let decoded = InlineBitpacking::unchunk::(data, num_values).unwrap(); + let DataBlock::FixedWidth(fixed) = decoded else { + panic!("expected FixedWidth DataBlock"); + }; + let decoded_values = fixed.data.borrow_to_typed_view::(); + assert_eq!(decoded_values.as_ref(), values); + } + + fn assert_corrupt_unchunk(data: LanceBuffer, num_values: u64, expected_message: &str) + where + T: ArrowNativeType + BitPacking + Pod, + { + let err = InlineBitpacking::unchunk::(data, num_values).unwrap_err(); + assert!(matches!(err, lance_core::Error::CorruptFile { .. })); + let err = err.to_string(); + assert!( + err.contains(expected_message), + "expected error containing {expected_message:?}, got {err:?}" + ); + } + + #[test] + fn unchunk_u32_bw12_tail() { + let values: Vec = (0..500).map(|i| ((i * 7) % (1 << 12)) as u32).collect(); + roundtrip_unchunk(&values, 12); + } + + #[test] + fn unchunk_u64_bw23_full() { + let values: Vec = (0..1024).map(|i| ((i * 3) % (1 << 23)) as u64).collect(); + roundtrip_unchunk(&values, 23); + } + + #[rstest] + #[case::too_small_header(LanceBuffer::from(vec![1, 2, 3]), 1, "too small")] + #[case::misaligned_chunk_size(LanceBuffer::from(vec![0, 0, 0, 0, 0]), 1, "multiple")] + #[case::too_many_values( + LanceBuffer::reinterpret_vec(vec![0_u32]), + ELEMS_PER_CHUNK + 1, + "expected at most" + )] + #[case::payload_size_mismatch(LanceBuffer::reinterpret_vec(vec![12_u32]), 1, "payload")] + #[case::invalid_bit_width(LanceBuffer::reinterpret_vec(vec![33_u32]), 1, "exceeds")] + fn unchunk_rejects( + #[case] data: LanceBuffer, + #[case] num_values: u64, + #[case] expected_message: &str, + ) { + assert_corrupt_unchunk::(data, num_values, expected_message); + } + #[test_log::test(tokio::test)] async fn test_miniblock_bitpack() { let test_cases = TestCases::default().with_min_file_version(LanceFileVersion::V2_1);